Improve Jit (#457)

Author: Autumn1998

csrc/hybrid_ep/executor/executor.cu

@@ -3,7 +3,7 @@
 
 #include "executor.cuh"
 
-Executor::Executor(int local_rank, int node_rank) : local_rank(local_rank), node_rank(node_rank), kernel_cache(local_rank) {}  
+Executor::Executor(int local_rank, int node_rank, std::string base_path) : local_rank(local_rank), node_rank(node_rank), kernel_cache(local_rank, base_path) {}  
 
 std::tuple<torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor>
 Executor::metadata_preprocess_core(

csrc/hybrid_ep/executor/executor.cuh

@@ -7,13 +7,13 @@
 #include <torch/torch.h>
 
 #include "utils.cuh"
-#include "backend/hybrid_ep_backend.cuh"
+#include "hybrid_ep_backend.cuh"
 #include "jit/compiler.cuh"
 #include "extension/permute.cuh"
 
 class Executor {
 public:
-    Executor(int local_rank, int node_rank);
+    Executor(int local_rank, int node_rank, std::string base_path);
 
     struct DispatchArgs {
         // Input tensors

csrc/hybrid_ep/hybrid_ep.cu

@@ -2,9 +2,9 @@
 // SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved
 #include "hybrid_ep.cuh"
 
-HybridEPBuffer::HybridEPBuffer(BufferConfig config, int local_rank, int node_rank, int group_size)
+HybridEPBuffer::HybridEPBuffer(BufferConfig config, int local_rank, int node_rank, int group_size, std::string base_path)
     : buffer_config(config), local_rank(local_rank), node_rank(node_rank), group_size(group_size),
-    executor(local_rank, node_rank) {
+    executor(local_rank, node_rank, base_path) {
     if(group_size <= buffer_config.num_of_ranks_per_node) {
       // If used on only intra-node communication, the dispatch/combine can share same buffers.
       use_shared_buffer = true;
@@ -145,6 +145,7 @@ void HybridEPBuffer::allocate_buffer_for_dispatch() {
   // Allocate and initialize synchronization buffers
   if (local_rank == 0) {
     remote_allocator.allocate((void**)&dispatch_buffers.intra_node_write_completion_flags, sizeof(uint32_t));
+    CUDA_CHECK(cudaMemset(dispatch_buffers.intra_node_write_completion_flags, 0, sizeof(uint32_t)));
   }
 
   CUDA_CHECK(cudaMalloc((void**)&dispatch_buffers.expected_rdma_flag_value, sizeof(uint64_t)));
@@ -207,6 +208,7 @@ void HybridEPBuffer::allocate_buffer_for_combine() {
   // Allocate and initialize synchronization buffers
   if (local_rank == 0) {
     remote_allocator.allocate((void**)&combine_buffers.intra_node_write_completion_flags, sizeof(uint32_t));
+    CUDA_CHECK(cudaMemset(combine_buffers.intra_node_write_completion_flags, 0, sizeof(uint32_t)));
   }
 
   CUDA_CHECK(cudaMalloc((void**)&combine_buffers.expected_rdma_flag_value, sizeof(uint64_t)));

csrc/hybrid_ep/hybrid_ep.cuh

@@ -2,7 +2,7 @@
 // SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved
 #pragma once
 #include "config.cuh"
-#include "backend/hybrid_ep_backend.cuh"
+#include "hybrid_ep_backend.cuh"
 #include "allocator/allocator.cuh"
 #include "utils.cuh"
 #include "executor/executor.cuh"
@@ -11,10 +11,11 @@
 #include <torch/torch.h>
 #include <vector>
 #include <algorithm>
+#include <string>
 
 class HybridEPBuffer {
 public:
-  HybridEPBuffer(BufferConfig config, int local_rank, int node_rank, int group_size);
+  HybridEPBuffer(BufferConfig config, int local_rank, int node_rank, int group_size, std::string base_path);
   ~HybridEPBuffer();
   bool update_buffer(HybridEpConfigInstance config); // True means the buffer is reallocated.
 

csrc/hybrid_ep/jit/compiler.cu

@@ -11,7 +11,7 @@ inline std::string get_env(std::string name) {
     return std::string(env);
 }
 
-NVCCCompiler::NVCCCompiler() {
+NVCCCompiler::NVCCCompiler(std::string base_path): base_path(base_path) {
     nvcc_path = get_env("CUDA_HOME") + "/bin/nvcc";
 
     // Init the flags to compiler
@@ -21,10 +21,7 @@ NVCCCompiler::NVCCCompiler() {
             " -Xcompiler -fPIC -shared";
 
     // Add the include path of the hybrid-ep library
-    std::string base_path = BASE_PATH;
-
-    // Add the include path of the hybrid-ep library
-    include = " -I" + base_path + "/csrc/hybrid_ep" + " -I" + get_env("CUDA_HOME") + "/include";
+    include = " -I" + base_path + "/backend" + " -I" + get_env("CUDA_HOME") + "/include";
 
     // Add the library path of the hybrid-ep library
     library = "-L" + get_env("CUDA_HOME") + "/lib64 -lcudart";
@@ -40,8 +37,6 @@ NVCCCompiler::NVCCCompiler() {
 
 
 std::string NVCCCompiler::build(std::string code, std::string signature, int local_rank) {
-    std::string base_path = BASE_PATH;
-
     // Create the source directory
     std::string jit_dir = base_path + "/build/jit";
     std::filesystem::create_directories(jit_dir);
@@ -80,7 +75,7 @@ std::string NVCCCompiler::build(std::string code, std::string signature, int loc
     return output_path;
 }
 
-std::any NVCCCompiler::get_instance(std::string library_path) {
+std::any NVCCCompiler::get_instance(std::string library_path, std::string kernel_key) {
     // Open the compiled library with RTLD_GLOBAL for symbol visibility
     void* handle = dlopen(library_path.c_str(), RTLD_NOW | RTLD_GLOBAL);
     if (handle == nullptr) {
@@ -98,8 +93,15 @@ std::any NVCCCompiler::get_instance(std::string library_path) {
                                 library_path);
     }
 
-    // After using the library, clear the built library
-    remove(library_path.c_str());
+    // Unique the compiled lib from different rank
+    std::string unique_library_path = base_path + "/build/jit/" + kernel_key + ".so";
+    std::string unique_command = "mv " + library_path + " " + unique_library_path;
+    if(library_path != unique_library_path) {
+        auto ret = std::system(unique_command.c_str());
+        if (ret != 0) {
+            throw std::runtime_error("Failed to unique the library: " + unique_command);
+        }
+    }
 
     // Run the get_function_ptr, then we get the compiled template
     std::any func_ptr = get_ptr();
@@ -109,7 +111,7 @@ std::any NVCCCompiler::get_instance(std::string library_path) {
 
 std::string NVCCCompiler::get_metadata_preprocessing_code(HybridEpConfigInstance config) {
   return R"(
-        #include "backend/hybrid_ep_backend.cuh"
+        #include "hybrid_ep_backend.cuh"
         #include <any>
         
         extern "C" {
@@ -130,7 +132,7 @@ std::string NVCCCompiler::get_dispatch_code(HybridEpConfigInstance config) {
       (config.token_data_type == TOKEN_DATA_TYPE::UINT8) ? "uint8_t" : "uint16_t";
 
   return R"(
-        #include "backend/hybrid_ep_backend.cuh"
+        #include "hybrid_ep_backend.cuh"
         #include <any>
         
         extern "C" {
@@ -150,7 +152,7 @@ std::string NVCCCompiler::get_dispatch_code(HybridEpConfigInstance config) {
 
 std::string NVCCCompiler::get_combine_code(HybridEpConfigInstance config) {
   return R"(
-        #include "backend/hybrid_ep_backend.cuh"
+        #include "hybrid_ep_backend.cuh"
         #include <any>
 
         extern "C" {
@@ -171,7 +173,18 @@ std::string NVCCCompiler::get_combine_code(HybridEpConfigInstance config) {
       )";
 }
 
-
+KernelCache::KernelCache(int local_rank, std::string base_path): 
+local_rank(local_rank), base_path(base_path), nvcc_compiler(base_path) {
+    // Load all cached kernels from the cache directory
+    std::string cache_dir = base_path + "/build/jit";
+    std::filesystem::create_directories(cache_dir);
+    for (const auto& entry : std::filesystem::directory_iterator(cache_dir)) {
+        if (entry.path().extension() == ".so") {
+            std::string kernel_key = entry.path().stem().string();
+            kernel_cache[kernel_key] = nvcc_compiler.get_instance(entry.path().string(), kernel_key);
+        }
+    }
+}
 
 void KernelCache::run_proprecess_kernel(
     HybridEpConfigInstance config, 
@@ -202,7 +215,7 @@ void KernelCache::run_proprecess_kernel(
     if (it == kernel_cache.end()) {
         auto preprocessing_code = nvcc_compiler.get_metadata_preprocessing_code(config);
         auto preprocessing_path = nvcc_compiler.build(preprocessing_code, preprocess_kernel_key, local_rank);
-        kernel_cache[preprocess_kernel_key] = nvcc_compiler.get_instance(preprocessing_path);
+        kernel_cache[preprocess_kernel_key] = nvcc_compiler.get_instance(preprocessing_path, preprocess_kernel_key);
     }
     auto preprocessing_instance = kernel_cache[preprocess_kernel_key];
 
@@ -255,7 +268,7 @@ void KernelCache::run_dispatch_kernel(
         // JIT Compile the kernel
         auto dispatch_code = nvcc_compiler.get_dispatch_code(config);
         auto dispatch_path = nvcc_compiler.build(dispatch_code, dispatch_kernel_key, local_rank);
-        kernel_cache[dispatch_kernel_key] = nvcc_compiler.get_instance(dispatch_path);
+        kernel_cache[dispatch_kernel_key] = nvcc_compiler.get_instance(dispatch_path, dispatch_kernel_key);
     }
     auto dispatch_instance = kernel_cache[dispatch_kernel_key];
 
@@ -295,7 +308,7 @@ void KernelCache::run_combine_kernel(
         // JIT Compile the kernel
         auto combine_code = nvcc_compiler.get_combine_code(config);
         auto combine_path = nvcc_compiler.build(combine_code, combine_kernel_key, local_rank);
-        kernel_cache[combine_kernel_key] = nvcc_compiler.get_instance(combine_path);
+        kernel_cache[combine_kernel_key] = nvcc_compiler.get_instance(combine_path, combine_kernel_key);
     }
     auto combine_instance = kernel_cache[combine_kernel_key];
 

csrc/hybrid_ep/jit/compiler.cuh

@@ -17,13 +17,13 @@
 #include <iostream>
 
 #include "config.cuh"
-#include "backend/hybrid_ep_backend.cuh"
+#include "hybrid_ep_backend.cuh"
 #include "utils.cuh"
 
 class NVCCCompiler{
 public:
     // Init the flags required by nvcc compiler
-    NVCCCompiler();
+    NVCCCompiler(std::string base_path);
 
     // Generate the code for jit compile
     std::string get_metadata_preprocessing_code(HybridEpConfigInstance config);
@@ -45,12 +45,14 @@ public:
     * @brief Get the compiled function pointer from the compiled .so file
     *
     * @param library_path The path of the compiled .so file
+    * @param kernel_key The key of the kernel, used to cache the compiled function pointer
     * @return std::any The function pointer
     */
-    std::any get_instance(std::string library_path);
+    std::any get_instance(std::string library_path, std::string kernel_key);
 
 
 private:
+    std::string base_path;  // The path of the installed package
     std::string flags;      // The flags required by nvcc compiler, which contains the
     // base flags(-O3, -arch...), include files, library files
     std::string nvcc_path;  // The path of the nvcc compiler
@@ -60,8 +62,7 @@ private:
 
 class KernelCache{
 public:
-    KernelCache() = default;
-    KernelCache(int local_rank) : local_rank(local_rank) {}
+    KernelCache(int local_rank, std::string base_path);
 
     void run_proprecess_kernel(
         HybridEpConfigInstance config, 
@@ -94,5 +95,6 @@ public:
 private:
     NVCCCompiler nvcc_compiler;
     std::unordered_map<std::string, std::any> kernel_cache;
+    std::string base_path;  // The path of the installed package
     int local_rank;   // Used to generate the unique signature for each rank
 };

csrc/hybrid_ep/pybind_hybrid_ep.cu

@@ -101,7 +101,7 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
         });
 
     pybind11::class_<HybridEPBuffer>(m, "HybridEPBuffer")
-        .def(py::init<BufferConfig, int, int, int>())
+        .def(py::init<BufferConfig, int, int, int, std::string>())
         .def("update_buffer", &HybridEPBuffer::update_buffer, py::arg("config"))
         .def("exchange_ipc_address", &HybridEPBuffer::exchange_ipc_address)
         .def("metadata_preprocessing", &HybridEPBuffer::metadata_preprocessing,

csrc/hybrid_ep/backend/hybrid_ep_backend.cuh renamed to deep_ep/backend/hybrid_ep_backend.cuh

@@ -111,11 +111,19 @@ def __init__(
 
         # Create C++ buffer - this will allocate all buffers during construction
         self.runtime = hybrid_ep_cpp.HybridEPBuffer(
-            self.config, self.local_rank, self.node_rank, self.group_size
+            self.config, self.local_rank, self.node_rank, self.group_size, os.path.dirname(os.path.abspath(__file__))
         )
         # Exchange IPC addresses using C++ distributed communication
         self.runtime.exchange_ipc_address(self.group)
 
+    def empty_jit_cache(self):
+        '''
+        Clean the cached kernel files.
+        '''
+        jit_cache_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "build", "jit")
+        if os.path.exists(jit_cache_path):
+            shutil.rmtree(jit_cache_path)
+
     def update_template_config(
         self,
         hidden_dim: int = None,