feat: remove page_size_variable mode

Author: niushengxiao

lightllm/common/basemodel/basemodel.py

@@ -687,7 +687,7 @@ 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.req_manager.alloc_token_indices(
+            mem_indexes = self.req_manager.alloc_mem_indices(
                 len(dummy_input_ids), b_req_idx, b_seq_len, b_ready_cache_len
             ).cuda()
             total_token_num = self.batch_max_tokens
@@ -764,7 +764,7 @@ def _autotune_warmup(self):
                 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_token_indices(
+                mem_indexes = self.req_manager.alloc_mem_indices(
                     len(dummy_input_ids), b_req_idx, b_seq_len, b_ready_cache_len
                 ).cuda()
                 model_input = ModelInput(

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.req_manager.alloc_token_indices(len(input_ids), b_req_idx, b_seq_len).cuda()
+            mem_indexes = model.req_manager.alloc_mem_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.req_manager.alloc_token_indices(len(input_ids), b_req_idx, b_seq_len).cuda()
+                mem_indexes = model.req_manager.alloc_mem_indices(len(input_ids), b_req_idx, b_seq_len).cuda()
 
                 micro_batch = ModelInput(
                     is_prefill=False,

lightllm/common/deepseek2_mem_manager.py

@@ -8,6 +8,7 @@
 from lightllm.common.kv_trans_kernel.kv_trans import kv_trans
 from lightllm.common.kv_trans_kernel.kv_trans_v2 import kv_trans_v2_for_d_node, kv_trans_v2_for_p_node
 from lightllm.distributed.pynccl import PyNcclCommunicator
+from lightllm.utils.envs_utils import get_page_size
 
 logger = init_logger(__name__)
 
@@ -20,7 +21,12 @@ def get_cell_size(self):
         return self.head_num * self.head_dim * self.layer_num * torch._utils._element_size(self.dtype)
 
     def _init_buffers(self, size, dtype, head_num, head_dim, layer_num):
-        self.kv_buffer = torch.empty((layer_num, size + 1, head_num, head_dim), dtype=dtype, device="cuda")
+        page_size = get_page_size()
+        self.kv_buffer = torch.empty(
+            (layer_num, (size // page_size + 1) * page_size, head_num, head_dim),
+            dtype=dtype,
+            device="cuda",
+        )
 
         # todo, etp or edp use the same work buffer here
         # also it can be used for any kernels for work buffer witout save info only

lightllm/common/deepseek2_paged_mem_manager.py

@@ -2,14 +2,15 @@
 import os
 import torch
 import torch.distributed as dist
+import triton
 from typing import List, Union
 from lightllm.server.pd_io_struct import KVMoveTask
 from lightllm.utils.log_utils import init_logger
 from lightllm.server.router.dynamic_prompt.shared_arr import SharedInt
 from lightllm.utils.profile_max_tokens import get_available_gpu_memory, get_total_gpu_memory
 from lightllm.common.kv_trans_kernel.kv_trans import kv_trans
 from lightllm.utils.dist_utils import get_current_rank_in_node
-from lightllm.utils.envs_utils import get_unique_server_name, get_env_start_args
+from lightllm.utils.envs_utils import get_unique_server_name, get_env_start_args, get_page_size
 from lightllm.distributed.pynccl import PyNcclCommunicator
 from lightllm.utils.dist_utils import get_current_device_id
 
@@ -81,7 +82,12 @@ def _init_buffers(self, size, dtype, head_num, head_dim, layer_num):
         # 分配，内部实际也没有管理，这个token是预留来对一些特殊的运行模式，如多dp下，overlap microbatch
         # 等模式下 padding 一些请求，使推理过程可以正常运行采用的，其索引值为size，存储在HOLD_TOKEN_MEMINDEX
         # 成员变量中，其与 req_manager 中的HOLD_REQUEST_ID具有类似的作用和意义。
-        self.kv_buffer = torch.empty((layer_num, size + 1, 2 * head_num, head_dim), dtype=dtype, device="cuda")
+        page_size = get_page_size()
+        self.kv_buffer = torch.empty(
+            (layer_num, (size // page_size + 1) * page_size, 2 * head_num, head_dim),
+            dtype=dtype,
+            device="cuda",
+        )
 
     def alloc_kv_move_buffer(self, max_req_total_len):
         """
@@ -244,6 +250,7 @@ def _free_buffers(self):
         self.kv_buffer = None
 
     def alloc(self, need_size) -> torch.Tensor:
+        assert need_size % get_page_size() == 0, "Need size must be a multiple of page size"
         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"
@@ -265,18 +272,25 @@ def free(self, free_index: Union[torch.Tensor, List[int]]):
         """
 
         end = self.mark_start
-        start = self.mark_start - len(free_index)
-        assert start >= 0, f"error free state start: {self.mark_start} free len {len(free_index)}"
+        page_size = get_page_size()
+        free_len = page_size * triton.cdiv(len(free_index), page_size)
+        start = self.mark_start - free_len
+        assert start >= 0, f"error free state start: {self.mark_start} free len {free_len}"
 
         if isinstance(free_index, list):
-            self.mem_state.numpy()[start:end] = free_index
+            free_index = torch.tensor(free_index)
+
+        # 从 gpu 到 cpu 的拷贝操作是流内阻塞操作
+        if page_size > 1:
+            base_free_index = free_index[free_index % page_size == 0]
+            token_idxs = base_free_index[:, None] + torch.arange(page_size)
+            self.mem_state[start:end] = token_idxs.flatten()
         else:
-            # 从 gpu 到 cpu 的拷贝操作是流内阻塞操作
             self.mem_state[start:end] = free_index
 
-        self.mark_start -= len(free_index)
+        self.mark_start -= free_len
 
-        self.can_use_mem_size += len(free_index)
+        self.can_use_mem_size += free_len
         self.shared_can_use_token_num.set_value(self.can_use_mem_size)
 
         if self.can_use_mem_size == len(self.mem_state):

lightllm/common/mem_manager.py

@@ -4,7 +4,6 @@
 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.paged_mem_manager import PagedMemoryManager
 from lightllm.utils.log_utils import init_logger
 
 logger = init_logger(__name__)
@@ -29,9 +28,6 @@ def select_mem_manager_class(mode):
     elif "export_fp8kv_calibration" in mode:
         memory_manager_class = ExportCalibrationMemoryManager
         logger.info("Using mode export fp8kv calibration")
-    elif "page_size_variable" in mode:
-        memory_manager_class = PagedMemoryManager
-        logger.info("Page size will be variable")
     else:
         memory_manager_class = MemoryManager
         logger.info("Model kv cache using mode normal")

lightllm/common/mem_utils.py

@@ -1,5 +1,6 @@
 import torch
 import collections
+import triton
 from lightllm.utils.log_utils import init_logger
 from .mem_manager import MemoryManager
 from typing import List, Optional
@@ -11,10 +12,6 @@
 logger = init_logger(__name__)
 
 
-def cdiv(a, b):
-    return (a + b - 1) // b
-
-
 class _ReqNode:
     def __init__(self, index):
         self.index = index
@@ -71,25 +68,60 @@ def __init__(self, max_request_num, max_sequence_length, mem_manager: MemoryMana
         self.max_request_num = max_request_num
         self.HOLD_REQUEST_ID = max_request_num
 
+    def calc_real_need_token_num(self, need_token_num, b_seq_len, b_ready_cache_len=None):
+        return max(need_token_num, self._get_need_paged_token_num(b_seq_len, b_ready_cache_len))
+
+    def alloc_mem_indices(self, need_size, b_req_idx=None, b_seq_len=None, b_ready_cache_len=None) -> torch.Tensor:
+        page_size = get_page_size()
+        if page_size > 1 and b_req_idx is not None and b_seq_len is not None:
+            return self._alloc_paged_mem_indices(b_req_idx, page_size, b_seq_len, b_ready_cache_len)
+        else:
+            return self.mem_manager.alloc(need_size)
+
+    def alloc(self):
+        return self.req_list.alloc()
+
+    def free(self, free_req_indexes: List[int], free_token_index):
+        for req_index in free_req_indexes:
+            self.req_list.free(req_index)
+
+        if self.req_list.is_all_free():
+            logger.debug(f"freed all request size {self.req_list.can_alloc_size}")
+        self.mem_manager.free(free_token_index)
+
+    def free_req(self, free_req_index: int):
+        self.req_list.free(free_req_index)
+        if self.req_list.is_all_free():
+            logger.debug(f"freed all request size {self.req_list.can_alloc_size}")
+        return
+
+    def free_token(self, free_token_index):
+        self.mem_manager.free(free_token_index)
+        return
+
+    def free_all(self):
+        self.req_list = _ReqLinkedList(self.max_request_num)
+        return
+
     def _expand_by_page_size(self, b_token_len, page_size):
-        # 将seq_len按page整数倍展开，例如seq_len = [9,9,9] -> page_len = [4,4,1,4,4,1,4,4,1], page_size = 4
-        b_page_len = cdiv(b_token_len, page_size)
+        # 将seq_len按page整数倍展开，例如seq_len = [9,9,9] -> p_token_len = [4,4,1,4,4,1,4,4,1], page_size = 4
+        b_page_len = triton.cdiv(b_token_len, page_size)
         need_pages_num = b_page_len.sum()
         p_token_len = torch.full((need_pages_num,), page_size, dtype=b_token_len.dtype, device=b_token_len.device)
         cumsum_pages = torch.cumsum(b_page_len, dim=0)
         last_page_positions = cumsum_pages - 1
         remainders = b_token_len - (b_page_len - 1) * page_size
         p_token_len[last_page_positions] = remainders
-        return need_pages_num, b_page_len, p_token_len
+        return need_pages_num, p_token_len
 
-    def _alloc_paged_token_indices(self, b_req_idx, page_size, b_seq_len, b_ready_cache_len):
+    def _alloc_paged_mem_indices(self, b_req_idx, page_size, b_seq_len, b_ready_cache_len):
         if b_ready_cache_len is not None:
             # prefill
             b_seq_len = b_seq_len.cpu()
             b_ready_cache_len = b_ready_cache_len.cpu()
 
             b_token_len = b_seq_len - b_ready_cache_len
-            total_pages_needed, b_page_len, p_token_len = self._expand_by_page_size(b_token_len, page_size)
+            total_pages_needed, p_token_len = self._expand_by_page_size(b_token_len, page_size)
             paged_token_idxs = self.mem_manager.alloc(total_pages_needed * page_size)
             pages = paged_token_idxs.view(-1, page_size)
             mask = torch.arange(page_size, device=p_token_len.device) < p_token_len.unsqueeze(1)
@@ -126,41 +158,6 @@ def _get_need_paged_token_num(self, b_seq_len, b_ready_cache_len=None):
             need_new_pages = mask.sum()
         return need_new_pages * page_size
 
-    def calc_real_need_token_num(self, need_token_num, b_seq_len, b_ready_cache_len=None):
-        return max(need_token_num, self._get_need_paged_token_num(b_seq_len, b_ready_cache_len))
-
-    def alloc_token_indices(self, need_size, b_req_idx, b_seq_len, b_ready_cache_len=None) -> torch.Tensor:
-        page_size = get_page_size()
-        if page_size > 1:
-            return self._alloc_paged_token_indices(b_req_idx, page_size, b_seq_len, b_ready_cache_len)
-        else:
-            return self.mem_manager.alloc(need_size)
-
-    def alloc(self):
-        return self.req_list.alloc()
-
-    def free(self, free_req_indexes: List[int], free_token_index):
-        for req_index in free_req_indexes:
-            self.req_list.free(req_index)
-
-        if self.req_list.is_all_free():
-            logger.debug(f"freed all request size {self.req_list.can_alloc_size}")
-        self.mem_manager.free(free_token_index)
-
-    def free_req(self, free_req_index: int):
-        self.req_list.free(free_req_index)
-        if self.req_list.is_all_free():
-            logger.debug(f"freed all request size {self.req_list.can_alloc_size}")
-        return
-
-    def free_token(self, free_token_index):
-        self.mem_manager.free(free_token_index)
-        return
-
-    def free_all(self):
-        self.req_list = _ReqLinkedList(self.max_request_num)
-        return
-
 
 class ReqSamplingParamsManager:
     """