feat: clean code

Author: niushengxiao

lightllm/common/basemodel/basemodel.py

@@ -61,8 +61,8 @@ def __init__(self, kvargs):
         self.finetune_config = kvargs.get("finetune_config", None)
         self.max_req_num = kvargs.get("max_req_num", 1000)
         self.max_seq_length = kvargs.get("max_seq_length", 1024 * 5)
-        # 一个可选的钩子函数，会在模型 warmup 之前被调用，用于做外部初始化等待（如 CPU KV Cache 注册完成）
-        self._pre_warmup_hook = kvargs.get("pre_warmup_hook", None)
+        # 用于做外部初始化等待（如 CPU KV Cache 注册完成）
+        self.waiting_hook = kvargs.get("waiting_hook", None)
         # is_token_healing 和 return_all_prompt_logics 是有排斥关系的两个模式，只能单独有一个生效
         # 主要是在prefill阶段返回多少个token的用于后续处理相关。
         self.is_token_healing = kvargs.get("is_token_healing", False)
@@ -108,19 +108,19 @@ def __init__(self, kvargs):
         self._init_inferstate_cls()
         self._autotune_warmup()
         self._init_padded_req()
-        # 在进入 autotune warmup 之前执行可选的预热钩子（例如等待 CPU KV Cache 注册完成）
-        if callable(self._pre_warmup_hook):
-            try:
-                self._pre_warmup_hook()
-            except Exception as e:
-                logger.exception(f"pre_warmup_hook failed: {e}")
-                raise
+        # wait必须在init cudagraph之前，避免错误捕获
+        self._run_waiting_hook()
         self._init_cudagraph()
         self._check_max_len_infer()
         torch.cuda.empty_cache()
         set_model_init_status(True)
         return
 
+    def _run_waiting_hook(self):
+        if self.waiting_hook is not None:
+            self.waiting_hook()
+        return
+
     def _init_config(self):
         with open(os.path.join(self.weight_dir_, "config.json"), "r") as json_file:
             self.config = json.load(json_file)

lightllm/common/basemodel/triton_kernel/kv_cache_offload.py

@@ -2,6 +2,7 @@
 
 import triton
 import triton.language as tl
+from lightllm.utils.dist_utils import get_current_rank_in_dp, get_dp_world_size
 
 
 @triton.jit
@@ -72,16 +73,13 @@ def offload_gpu_kv_to_cpu(
 ):
     """
     this function is used to offload GPU KV cache to CPU KV cache.
-    Supports tensor parallelism (TP > 1).
     Args:
         token_indexes: (token_num,)
         gpu_kv_cache: (layer_num, token_num, head_num, head_dim)
         cpu_kv_cache: (all_page_num, layer_num, token_block_size, head_num, head_dim)
         page_indexes: (page_num,)
         page_readies: (page_num,)
     """
-    from lightllm.utils.dist_utils import get_current_rank_in_dp, get_dp_world_size
-    
     token_block_size = cpu_kv_cache.shape[2]
     token_num = page_indexes.shape[0] * token_block_size
     assert token_indexes.shape[0] >= token_num
@@ -92,9 +90,15 @@ def offload_gpu_kv_to_cpu(
 
     # Calculate head offset for tensor parallelism
     tp_rank = get_current_rank_in_dp()
+    tp_num = get_dp_world_size()
     gpu_heads = gpu_kv_cache.shape[2]
     gpu_head_dim = gpu_kv_cache.shape[3]
-    cpu_head_offset = tp_rank * gpu_heads * gpu_head_dim
+    cpu_heads = cpu_kv_cache.shape[3]
+    factor = (tp_num * gpu_heads) // cpu_heads
+    cpu_head_offset = (tp_rank // factor) * gpu_heads * gpu_head_dim
+    if tp_rank % factor != 0:
+        # redundant kv does not need to offload
+        return
 
     grid = (page_num,)
     num_warps = 4
@@ -142,7 +146,6 @@ def _load_cpu_cache_to_gpu(
     page_indexes_ptr,
     layer_num,
     head_all_dim,
-    all_move_token_num,
     cpu_head_offset,
     BLOCK_HEAD_ALL_DIM: tl.constexpr,
     TOKEN_BLOCK: tl.constexpr,
@@ -152,17 +155,11 @@ def _load_cpu_cache_to_gpu(
     if cpu_page_index == -1:
         return
 
-    gpu_stride0 = tl.cast(gpu_stride0, dtype=tl.int64)
-    padded_size = TOKEN_BLOCK * tl.num_programs(0) - all_move_token_num
-    head_all_dim_range = tl.arange(0, BLOCK_HEAD_ALL_DIM)
     token_range = block_index * TOKEN_BLOCK + tl.arange(0, TOKEN_BLOCK)
-    token_range = token_range - padded_size
-
-    token_mask = token_range >= 0
+    token_indexes = tl.load(token_indexes_ptr + token_range).to(tl.int64)
+    head_all_dim_range = tl.arange(0, BLOCK_HEAD_ALL_DIM)
     head_dim_mask = head_all_dim_range < head_all_dim
 
-    token_indexes = tl.load(token_indexes_ptr + token_range, mask=token_mask, other=0).to(tl.int64)
-
     cpu_page_index = tl.load(page_indexes_ptr + block_index).to(tl.int64)
     for layer_index in range(layer_num):
         cpu_ptr = (
@@ -176,14 +173,14 @@ def _load_cpu_cache_to_gpu(
 
         gpu_ptr = (
             gpu_kv_cache_ptr
-            + layer_index * gpu_stride0
+            + layer_index.to(tl.int64) * gpu_stride0
             + token_indexes[:, None] * gpu_stride1
             + head_all_dim_range[None, :]
         )
         tl.store(
             gpu_ptr,
             cpu_data,
-            mask=token_mask[:, None] & head_dim_mask[None, :],
+            mask=head_dim_mask[None, :],
         )
     return
 
@@ -196,27 +193,28 @@ def load_cpu_kv_to_gpu(
     page_indexes: torch.Tensor,
 ):
     """
-    this function is used to load CPU KV cache to GPU KV cache.
-    Supports tensor parallelism (TP > 1).
+    this function is used to offload GPU KV cache to CPU KV cache.
     Args:
         mem_indexes: (token_num,)
         gpu_kv_cache: (layer_num, token_num, head_num, head_dim)
         cpu_kv_cache: (page_num, layer_num, token_block_size, head_num, head_dim)
         page_indexes: (page_num,)
     """
-    from lightllm.utils.dist_utils import get_current_rank_in_dp, get_dp_world_size
-    
     token_block_size = cpu_kv_cache.shape[2]
     token_num = page_indexes.shape[0] * token_block_size
     assert mem_indexes.shape[0] >= token_num
     page_num = page_indexes.shape[0]
+    assert len(mem_indexes) == page_num * token_block_size
     BLOCK_HEAD_ALL_DIM = triton.next_power_of_2(gpu_kv_cache.shape[-1] * gpu_kv_cache.shape[-2])
 
     # Calculate head offset for tensor parallelism
     tp_rank = get_current_rank_in_dp()
+    tp_num = get_dp_world_size()
     gpu_heads = gpu_kv_cache.shape[2]
     gpu_head_dim = gpu_kv_cache.shape[3]
-    cpu_head_offset = tp_rank * gpu_heads * gpu_head_dim
+    cpu_heads = cpu_kv_cache.shape[3]
+    factor = (tp_num * gpu_heads) // cpu_heads
+    cpu_head_offset = (tp_rank // factor) * gpu_heads * gpu_head_dim
 
     grid = (page_num,)
     num_warps = 1
@@ -237,7 +235,6 @@ def load_cpu_kv_to_gpu(
         page_indexes_ptr=page_indexes,
         layer_num=gpu_kv_cache.shape[0],
         head_all_dim=gpu_kv_cache.shape[-1] * gpu_kv_cache.shape[-2],
-        all_move_token_num=len(mem_indexes),
         cpu_head_offset=cpu_head_offset,
         BLOCK_HEAD_ALL_DIM=BLOCK_HEAD_ALL_DIM,
         TOKEN_BLOCK=token_block_size,

lightllm/server/api_cli.py

@@ -508,7 +508,7 @@ def make_argument_parser() -> argparse.ArgumentParser:
     parser.add_argument(
         "--enable_cpu_cache",
         action="store_true",
-        help="""enable cpu cache to store kv cache.""",
+        help="""enable cpu cache to store kv cache. prefer to use hugepages for better performance.""",
     )
     parser.add_argument(
         "--cpu_cache_storage_size",
@@ -519,7 +519,7 @@ def make_argument_parser() -> argparse.ArgumentParser:
     parser.add_argument(
         "--cpu_cache_token_page_size",
         type=int,
-        default=256,
+        default=64,
         help="""The token page size of cpu cache""",
     )
     parser.add_argument("--enable_disk_cache", action="store_true", help="""enable disk cache to store kv cache.""")

lightllm/server/core/objs/atomic_array_lock.py

@@ -2,6 +2,7 @@
 import atomics
 from multiprocessing import shared_memory
 from lightllm.utils.log_utils import init_logger
+from lightllm.utils.auto_shm_cleanup import register_posix_shm_for_cleanup
 
 logger = init_logger(__name__)
 
@@ -26,6 +27,7 @@ def __init__(self, lock_name: str, lock_num: int):
     def _init_shm(self):
         try:
             shm = shared_memory.SharedMemory(name=self.lock_name, create=True, size=self.dest_size)
+            register_posix_shm_for_cleanup(self.lock_name)
             logger.info(f"create lock shm {self.lock_name}")
         except:
             shm = shared_memory.SharedMemory(name=self.lock_name, create=False, size=self.dest_size)

lightllm/server/core/objs/atomic_lock.py

@@ -2,6 +2,7 @@
 import time
 from multiprocessing import shared_memory
 from lightllm.utils.log_utils import init_logger
+from lightllm.utils.auto_shm_cleanup import register_posix_shm_for_cleanup
 
 logger = init_logger(__name__)
 
@@ -25,6 +26,7 @@ def _init_shm(self):
         try:
             shm = shared_memory.SharedMemory(name=self.lock_name, create=True, size=self.dest_size)
             logger.info(f"create lock shm {self.lock_name}")
+            register_posix_shm_for_cleanup(self.lock_name)
         except:
             shm = shared_memory.SharedMemory(name=self.lock_name, create=False, size=self.dest_size)
             logger.info(f"link lock shm {self.lock_name}")

lightllm/server/core/objs/rpc_shm.py

@@ -5,6 +5,7 @@
 from typing import List
 from lightllm.utils.envs_utils import get_unique_server_name
 from lightllm.utils.log_utils import init_logger
+from lightllm.utils.auto_shm_cleanup import register_posix_shm_for_cleanup
 
 logger = init_logger(__name__)
 
@@ -20,6 +21,7 @@ def __init__(self):
     def create_or_link_shm(self):
         try:
             shm = shared_memory.SharedMemory(name=self.name, create=True, size=LIGHTLLM_RPC_BYTE_SIZE)
+            register_posix_shm_for_cleanup(self.name)
         except:
             shm = shared_memory.SharedMemory(name=self.name, create=False, size=LIGHTLLM_RPC_BYTE_SIZE)
 
@@ -57,6 +59,7 @@ def __init__(self):
     def create_or_link_shm(self):
         try:
             shm = shared_memory.SharedMemory(name=self.name, create=True, size=LIGHTLLM_RPC_RESULT_BYTE_SIZE)
+            register_posix_shm_for_cleanup(self.name)
         except:
             shm = shared_memory.SharedMemory(name=self.name, create=False, size=LIGHTLLM_RPC_RESULT_BYTE_SIZE)
 
@@ -99,6 +102,7 @@ def __init__(self, world_size):
     def create_or_link_shm(self):
         try:
             shm = shared_memory.SharedMemory(name=self.name, create=True, size=self.dest_size)
+            register_posix_shm_for_cleanup(self.name)
         except:
             shm = shared_memory.SharedMemory(name=self.name, create=False, size=self.dest_size)
 

lightllm/server/core/objs/shm_array.py

@@ -1,5 +1,6 @@
 import numpy as np
 from multiprocessing import shared_memory
+from lightllm.utils.auto_shm_cleanup import register_posix_shm_for_cleanup
 from lightllm.utils.log_utils import init_logger
 
 logger = init_logger(__name__)
@@ -18,6 +19,7 @@ def __init__(self, name, shape, dtype):
     def create_shm(self):
         try:
             shm = shared_memory.SharedMemory(name=self.name, create=True, size=self.dest_size)
+            register_posix_shm_for_cleanup(self.name)
         except:
             shm = shared_memory.SharedMemory(name=self.name, create=False, size=self.dest_size)
 
@@ -28,6 +30,7 @@ def create_shm(self):
             try:
                 shm = shared_memory.SharedMemory(name=self.name, create=True, size=self.dest_size)
                 logger.info(f"create shm {self.name}")
+                register_posix_shm_for_cleanup(self.name)
             except:
                 shm = shared_memory.SharedMemory(name=self.name, create=False, size=self.dest_size)
                 logger.info(f"link shm {self.name}")

lightllm/server/core/objs/shm_req_manager.py

@@ -2,6 +2,7 @@
 import numpy as np
 from lightllm.utils.envs_utils import get_unique_server_name
 from multiprocessing import shared_memory
+from lightllm.utils.auto_shm_cleanup import register_posix_shm_for_cleanup
 from lightllm.utils.log_utils import init_logger
 from .req import Req, ChunkedPrefillReq, TokenHealingReq
 from .shm_array import ShmArray
@@ -53,6 +54,7 @@ def _init_reqs_shm(self):
         shm_name = f"{get_unique_server_name()}_req_shm_total"
         try:
             shm = shared_memory.SharedMemory(name=shm_name, create=True, size=self.req_shm_byte_size)
+            register_posix_shm_for_cleanup(shm_name)
             logger.info(f"create lock shm {shm_name}")
         except:
             shm = shared_memory.SharedMemory(name=shm_name, create=False, size=self.req_shm_byte_size)

lightllm/server/multi_level_kv_cache/cpu_cache_client.py

@@ -12,7 +12,6 @@
     attach_shm_kv_cache_ptr,
     register_shm_ptr_to_pin,
 )
-from lightllm.utils.infer_utils import mark_start, mark_end
 
 logger = init_logger(__name__)
 
@@ -31,9 +30,9 @@ def __init__(self, init_shm_data: bool):
         self._create_cpu_status_list(init_shm_data)
         if init_shm_data:
             self._create_shm_cpu_kv_cache()
-            self.pin_reg_handle = None
+            self.attach_shm_handle = None
         else:
-            self.pin_reg_handle = self._attach_shm_cpu_kv_cache()
+            self.attach_shm_handle = self._attach_shm_cpu_kv_cache()
         return
 
     def get_one_empty_page(self, hash_key: int, disk_offload_enable: bool) -> Optional[int]:
@@ -215,7 +214,6 @@ def _create_shm_cpu_kv_cache(self):
 
     def _attach_shm_cpu_kv_cache(self):
         shm_ptr = attach_shm_kv_cache_ptr()
-        mark_start("blueswhen1")
         handle = register_shm_ptr_to_pin(shm_ptr=shm_ptr, size=self.kv_cache_tensor_meta.calcu_size())
         numpy_array = np.frombuffer(
             memoryview((ctypes.c_uint8 * self.kv_cache_tensor_meta.calcu_size()).from_address(shm_ptr)), dtype=np.uint8

lightllm/server/multi_level_kv_cache/manager.py

@@ -31,15 +31,6 @@ def __init__(
         self.send_to_router = context.socket(zmq.PUSH)
         self.send_to_router.connect(f"{args.zmq_mode}127.0.0.1:{args.router_port}")
         logger.info(f"send_to_router sendhwm {self.send_to_router.getsockopt(zmq.SNDHWM)}")
-
-        # 自动注册共享内存清理
-        try:
-            from lightllm.utils.auto_shm_cleanup import auto_register_cpu_cache
-
-            auto_register_cpu_cache()
-        except Exception as e:
-            logger.warning(f"Failed to register auto shm cleanup: {e}")
-
         self.cpu_cache_client = CpuKvCacheClient(init_shm_data=True)
         self.shm_req_manager = ShmReqManager()
         # 控制同时进行cpu cache 匹配操作的数量。