feat: reduce startup time for cpu cache

Author: niushengxiao

lightllm/server/multi_level_kv_cache/cpu_cache_client.py

@@ -31,8 +31,7 @@ def __init__(self, only_create_meta_data: bool, init_shm_data: bool):
 
         if not only_create_meta_data:
             if init_shm_data:
-                self._create_shm_cpu_kv_cache()
-                self.attach_shm_handle = None
+                self.attach_shm_handle = self._create_shm_cpu_kv_cache()
             else:
                 self.attach_shm_handle = self._attach_shm_cpu_kv_cache()
         return
@@ -275,7 +274,7 @@ def _create_cpu_status_list(self, init_shm_data: bool):
         return
 
     def _create_shm_cpu_kv_cache(self):
-        shm_ptr = create_shm_kv_cache_ptr()
+        shm_ptr, prefault_handle = create_shm_kv_cache_ptr()
         numpy_array = np.frombuffer(
             memoryview((ctypes.c_uint8 * self.kv_cache_tensor_meta.calcu_size()).from_address(shm_ptr)), dtype=np.uint8
         )
@@ -290,7 +289,7 @@ def _create_shm_cpu_kv_cache(self):
         self.cpu_kv_cache_tensor = (
             torch.from_numpy(numpy_array).view(dtype=self.kv_cache_tensor_meta.data_type).view(shape)
         )
-        return
+        return prefault_handle
 
     def _attach_shm_cpu_kv_cache(self):
         shm_ptr = attach_shm_kv_cache_ptr()

lightllm/server/multi_level_kv_cache/manager.py

@@ -61,6 +61,9 @@ def __init__(
         return
 
     def cpu_cache_hanle_loop(self):
+        if self.cpu_cache_client.attach_shm_handle is not None:
+            self.cpu_cache_client.attach_shm_handle.wait()
+
         while True:
             try:
                 current_group_req = self.recv_queue.get()

lightllm/utils/kv_cache_utils.py

@@ -124,12 +124,16 @@ def calcu_cpu_cache_meta() -> "CpuKVCacheMeta":
 
 
 @lru_cache(maxsize=None)
-def create_shm_kv_cache_ptr() -> int:
+def create_shm_kv_cache_ptr() -> tuple[int, "AsyncPrefaultHandle"]:
     libc = ctypes.CDLL("/usr/lib/x86_64-linux-gnu/libc.so.6", use_errno=True)
     libc.shmget.argtypes = (ctypes.c_long, ctypes.c_size_t, ctypes.c_int)
     libc.shmget.restype = ctypes.c_int
     libc.shmat.argtypes = (ctypes.c_int, ctypes.c_void_p, ctypes.c_int)
     libc.shmat.restype = ctypes.c_void_p
+    libc.madvise.argtypes = (ctypes.c_void_p, ctypes.c_size_t, ctypes.c_int)
+    libc.madvise.restype = ctypes.c_int
+    libc.memset.argtypes = (ctypes.c_void_p, ctypes.c_int, ctypes.c_size_t)
+    libc.memset.restype = ctypes.c_void_p
 
     args = get_env_start_args()
     key = args.cpu_kv_cache_shm_id
@@ -189,7 +193,28 @@ def _get_default_hugepage_size() -> int:
         raise Exception("Error attaching shared memory")
     logger.info(f"Shared cpu kv cache tensor memory at address: {shm_addr}")
 
-    return shm_addr
+    # Best-effort memory prefaulting in background to speed up subsequent cudaHostRegister
+    def _prefault_memory():
+        if not use_hugetlb:
+            # MADV_HUGEPAGE only makes sense for regular pages with THP, not for SHM_HUGETLB
+            MADV_HUGEPAGE = 14
+            ret = libc.madvise(ctypes.c_void_p(shm_addr), ctypes.c_size_t(size_to_alloc), MADV_HUGEPAGE)
+            if ret != 0:
+                err = ctypes.get_errno()
+                logger.warning(f"madvise(MADV_HUGEPAGE) failed with errno {err}, continue without THP")
+            else:
+                logger.info("madvise(MADV_HUGEPAGE) succeeded (best-effort THP)")
+        else:
+            logger.debug("Skipping MADV_HUGEPAGE for SHM_HUGETLB mode (already using hugepages)")
+
+        # memset触发页面提前分配
+        libc.memset(ctypes.c_void_p(shm_addr), 0, ctypes.c_size_t(size_to_alloc))
+        logger.info("prefaulted shared memory pages successfully")
+
+    th = threading.Thread(target=_prefault_memory, name="cpu_cache_prefault", daemon=True)
+    th.start()
+
+    return shm_addr, AsyncPrefaultHandle(thread=th)
 
 
 @dataclasses.dataclass
@@ -306,6 +331,18 @@ def wait(self):
         return
 
 
+class AsyncPrefaultHandle:
+    """A handle for async memory prefaulting."""
+
+    def __init__(self, thread: Optional[threading.Thread]):
+        self.thread = thread
+
+    def wait(self):
+        if self.thread is not None and self.thread.is_alive():
+            self.thread.join()
+        return
+
+
 @lru_cache(maxsize=None)
 def attach_shm_kv_cache_ptr() -> int:
     libc = ctypes.CDLL("/usr/lib/x86_64-linux-gnu/libc.so.6", use_errno=True)