feat: add alloc_paged_token_indices function in req_manager

Author: niushengxiao

lightllm/common/basemodel/basemodel.py

@@ -687,8 +687,8 @@ def _check_max_len_infer(self):
             b_seq_len = torch.ones(1, dtype=torch.int32, device="cuda")
             b_seq_len[:] = self.batch_max_tokens
             b_ready_cache_len = torch.zeros(1, dtype=torch.int32, device="cuda")
-            mem_indexes = self.mem_manager.alloc(
-                len(dummy_input_ids), b_req_idx, b_seq_len, b_ready_cache_len, True
+            mem_indexes = self.req_manager.alloc_paged_token_indices(
+                len(dummy_input_ids), b_req_idx, b_seq_len, b_ready_cache_len
             ).cuda()
             total_token_num = self.batch_max_tokens
             b_mtp_index = torch.zeros(1, dtype=torch.int32, device="cuda")
@@ -759,12 +759,14 @@ def _autotune_warmup(self):
                     0, 10000, (input_len,), dtype=torch.int32, device="cuda", generator=rand_gen
                 )
                 b_req_idx = torch.tensor([self.req_manager.alloc()], dtype=torch.int32, device="cuda")
-                mem_indexes = self.mem_manager.alloc(len(dummy_input_ids)).cuda()
                 b_seq_len = torch.ones(1, dtype=torch.int32, device="cuda")
                 b_seq_len[:] = input_len
                 b_ready_cache_len = torch.zeros(1, dtype=torch.int32, device="cuda")
                 total_token_num = input_len
                 b_mtp_index = torch.zeros(1, dtype=torch.int32, device="cuda")
+                mem_indexes = self.req_manager.alloc_paged_token_indices(
+                    len(dummy_input_ids), b_req_idx, b_seq_len, b_ready_cache_len
+                ).cuda()
                 model_input = ModelInput(
                     batch_size=1,
                     total_token_num=total_token_num,

lightllm/common/basemodel/cuda_graph.py

@@ -202,7 +202,7 @@ def warmup(self, model):
             b_seq_len = torch.empty(batch_size, dtype=torch.int32, device="cuda")
             b_seq_len.fill_(seq_len)
             b_mtp_index = torch.zeros(batch_size, dtype=torch.int32, device="cuda")
-            mem_indexes = model.mem_manager.alloc(len(input_ids), b_req_idx, b_seq_len).cuda()
+            mem_indexes = model.req_manager.alloc_paged_token_indices(len(input_ids), b_req_idx, b_seq_len).cuda()
 
             model_input = ModelInput(
                 batch_size=batch_size,
@@ -258,7 +258,7 @@ def warmup_overlap(self, model):
                 b_seq_len = torch.empty(batch_size, dtype=torch.int32, device="cuda")
                 b_seq_len.fill_(seq_len)
                 b_mtp_index = torch.zeros(batch_size, dtype=torch.int32, device="cuda")
-                mem_indexes = model.mem_manager.alloc(len(input_ids), b_req_idx, b_seq_len).cuda()
+                mem_indexes = model.req_manager.alloc_paged_token_indices(len(input_ids), b_req_idx, b_seq_len).cuda()
 
                 micro_batch = ModelInput(
                     is_prefill=False,

lightllm/common/deepseek2_page_size_variable_mem_manager.py renamed to lightllm/common/deepseek2_paged_mem_manager.py

@@ -1,7 +1,7 @@
 import torch
 import numpy as np
 from .deepseek2_mem_manager import Deepseek2MemoryManager
-from .page_size_variable_mem_manager import PageSizeVariableMemoryManager
+from .paged_mem_manager import PagedMemoryManager
 from lightllm.utils.log_utils import init_logger
 from lightllm.utils.envs_utils import get_page_size
 
@@ -13,7 +13,7 @@ def cdiv(a, b):
 logger = init_logger(__name__)
 
 
-class Deepseek2PageSizeVariableMemoryManager(PageSizeVariableMemoryManager, Deepseek2MemoryManager):
+class Deepseek2PagedMemoryManager(PagedMemoryManager, Deepseek2MemoryManager):
     def __init__(self, size, dtype, head_num, head_dim, layer_num, always_copy=False, mem_fraction=0.9):
         super().__init__(size, dtype, head_num, head_dim, layer_num, always_copy, mem_fraction)
 

lightllm/common/mem_manager.py

@@ -52,8 +52,6 @@ def __init__(self, size, dtype, head_num, head_dim, layer_num, always_copy=False
             layer_num,
         )
         self.HOLD_TOKEN_MEMINDEX = self.size
-        # MemoryManager也需要个引用备份，供内部使用
-        self.req_to_token_indexs = None
 
     def get_cell_size(self):
         return 2 * self.head_num * self.head_dim * self.layer_num * torch._utils._element_size(self.dtype)
@@ -245,9 +243,7 @@ def _write_kv_move_data_p2p(self, token_indexes: torch.Tensor, buffer_tensor: to
     def _free_buffers(self):
         self.kv_buffer = None
 
-    def alloc(
-        self, need_size, b_req_idx=None, b_seq_len=None, b_ready_cache_len=None, is_prefill=False
-    ) -> torch.Tensor:
+    def alloc(self, need_size) -> torch.Tensor:
         if need_size > self.mark_end - self.mark_start:
             logger.error(f"warn no enough cache need_size {need_size} left_size {self.can_use_mem_size}")
             assert False, "error alloc state"
@@ -261,9 +257,6 @@ def alloc(
         self.shared_can_use_token_num.set_value(self.can_use_mem_size)
         return ans
 
-    def set_prefix_cache_to_req(self, req_idx: int, start: int, end: int, values: torch.Tensor):
-        self.req_to_token_indexs[req_idx, start:end] = values
-
     def free(self, free_index: Union[torch.Tensor, List[int]]):
         """_summary_
 
@@ -342,17 +335,8 @@ def __init__(self) -> None:
             SharedInt(f"{get_unique_server_name()}_mem_manger_can_use_token_num_{rank_in_node}")
             for rank_in_node in range(0, self.node_world_size, self.dp_world_size)
         ]
-        self.shared_tp_info_pages = [
-            SharedInt(f"{get_unique_server_name()}_mem_manger_can_use_page_num_{rank_in_node}")
-            for rank_in_node in range(0, self.node_world_size, self.dp_world_size)
-        ]
 
     def get_unrefed_token_num(self, dp_rank_in_node: int):
         if self.is_multinode_tp:
             return self.shared_tp_infos[0].get_value()
         return self.shared_tp_infos[dp_rank_in_node].get_value()
-
-    def get_unrefed_page_num(self, dp_rank_in_node: int):
-        if self.is_multinode_tp:
-            return self.shared_tp_info_pages[0].get_value()
-        return self.shared_tp_info_pages[dp_rank_in_node].get_value()

lightllm/common/mem_utils.py

@@ -4,7 +4,7 @@
 from lightllm.common.export_calibration_mem_manager import ExportCalibrationMemoryManager
 from lightllm.common.ppl_int8kv_mem_manager import PPLINT8KVMemoryManager
 from lightllm.common.ppl_int4kv_mem_manager import PPLINT4KVMemoryManager
-from lightllm.common.page_size_variable_mem_manager import PageSizeVariableMemoryManager
+from lightllm.common.paged_mem_manager import PagedMemoryManager
 from lightllm.utils.log_utils import init_logger
 
 logger = init_logger(__name__)
@@ -30,7 +30,7 @@ def select_mem_manager_class(mode):
         memory_manager_class = ExportCalibrationMemoryManager
         logger.info("Using mode export fp8kv calibration")
     elif "page_size_variable" in mode:
-        memory_manager_class = PageSizeVariableMemoryManager
+        memory_manager_class = PagedMemoryManager
         logger.info("Page size will be variable")
     else:
         memory_manager_class = MemoryManager