Fix

Author: none

lightllm/server/multi_level_kv_cache/cpu_cache_client.py

@@ -213,7 +213,7 @@ def _create_shm_cpu_kv_cache(self):
 
     def _attach_shm_cpu_kv_cache(self):
         shm_ptr = attach_shm_kv_cache_ptr()
-        register_shm_ptr_to_pin(shm_ptr=shm_ptr, size=self.kv_cache_tensor_meta.calcu_size())
+        device_ptr = register_shm_ptr_to_pin(shm_ptr=shm_ptr, size=self.kv_cache_tensor_meta.calcu_size())
         shape = (
             self.kv_cache_tensor_meta.page_num,
             self.kv_cache_tensor_meta.layer_num,
@@ -223,7 +223,7 @@ def _attach_shm_cpu_kv_cache(self):
         )
         self.cpu_kv_cache_tensor = torch.empty(size=shape, dtype=torch.bfloat16, device="meta")
         # 将指针绑定到 tensor上，方便triton获取真实的地址。
-        self.cpu_kv_cache_tensor.data_ptr = lambda: shm_ptr
+        self.cpu_kv_cache_tensor.data_ptr = lambda: device_ptr
         return
 
 

lightllm/server/router/model_infer/mode_backend/base_backend.py

@@ -192,16 +192,16 @@ def init_model(self, kvargs):
         # 开启 mtp 模式，需要完成mtp model的初始化
         if self.args.mtp_mode:
             self.init_mtp_draft_model(kvargs)
+        
+        if self.args.enable_cpu_cache:
+            self.multi_level_cache_module = MultiLevelKvCacheModule(self)
 
         # 启动infer_loop_thread, 启动两个线程进行推理，对于具备双batch推理折叠得场景
         # 可以降低 cpu overhead，大幅提升gpu得使用率。
         self.infer_loop_thread = threading.Thread(target=self.infer_loop, daemon=True)
         self.infer_loop_thread.start()
         self.infer_loop_thread1 = threading.Thread(target=self.infer_loop, daemon=True)
         self.infer_loop_thread1.start()
-
-        if self.args.enable_cpu_cache:
-            self.multi_level_cache_module = MultiLevelKvCacheModule(self)
         return
 
     def init_custom(self):

lightllm/server/router/model_infer/mode_backend/multi_level_kv_cache.py

@@ -21,7 +21,10 @@ def __init__(self, backend):
         self.gloo_group = create_new_group_for_current_dp("gloo")
         self.filter_group = create_new_group_for_current_dp("gloo")
         self.sync_group = create_new_group_for_current_dp("nccl")
+        dist.barrier(group=self.sync_group)
         self.init_sync_group = create_new_group_for_current_dp("nccl")
+        dist.barrier(group=self.init_sync_group)
+
 
         self.cpu_cache_handle_queue: Deque[TransTask] = deque()
         self.cpu_cache_client = CpuKvCacheClient(init_shm_data=False)

lightllm/utils/kv_cache_utils.py

@@ -148,14 +148,18 @@ def calcu_size(self):
         return self.page_num * self.layer_num * self.token_page_size * self.num_heads * self.head_dim * self.item_size
 
 
-def register_shm_ptr_to_pin(shm_ptr: int, size: int):
+def register_shm_ptr_to_pin(shm_ptr: int, size: int) -> int:
     # 加载 CUDA 库
     cuda = ctypes.CDLL("/usr/local/cuda/targets/x86_64-linux/lib/libcudart.so")  # Linux 下的 CUDA 库路径
 
     # 定义 cudaHostRegister 函数的参数和返回类型
     cuda.cudaHostRegister.argtypes = [ctypes.c_void_p, ctypes.c_size_t, ctypes.c_uint]
     cuda.cudaHostRegister.restype = ctypes.c_int
 
+    # 定义 cudaHostGetDevicePointer 函数原型
+    cuda.cudaHostGetDevicePointer.argtypes = [ctypes.POINTER(ctypes.c_void_p), ctypes.c_void_p, ctypes.c_int]
+    cuda.cudaHostGetDevicePointer.restype = ctypes.c_int
+
     # 定义常量
     cudaHostRegisterDefault = 0  # 默认注册标志
 
@@ -166,4 +170,15 @@ def register_shm_ptr_to_pin(shm_ptr: int, size: int):
         raise Exception(f"Error registering host memory: {result}")
     else:
         logger.info("Host memory registered successfully.")
-    return
+
+    device_ptr = ctypes.c_void_p()  # 输出设备指针
+    host_ptr = ctypes.c_void_p(shm_ptr)  # 输入主机指针
+    
+    result = cuda.cudaHostGetDevicePointer(ctypes.byref(device_ptr), host_ptr, 0)
+    
+    if result != 0:
+        raise RuntimeError(f"cudaHostGetDevicePointer failed with error code {result}")
+    
+    logger.info(f"get Host memory registered Device ptr {device_ptr.value}")
+    
+    return device_ptr.value