fix: fix the page not enough bug

Author: niushengxiao

lightllm/common/mem_manager.py

@@ -341,8 +341,17 @@ 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/page_size_variable_mem_manager.py

@@ -3,7 +3,9 @@
 from .mem_manager import MemoryManager
 from typing import List, Union
 from lightllm.utils.log_utils import init_logger
-from lightllm.utils.envs_utils import get_page_size
+from lightllm.utils.envs_utils import get_unique_server_name, get_page_size
+from lightllm.server.router.dynamic_prompt.shared_arr import SharedInt
+from lightllm.utils.dist_utils import get_current_rank_in_node
 
 
 def cdiv(a, b):
@@ -24,6 +26,12 @@ def __init__(self, size, dtype, head_num, head_dim, layer_num, always_copy=False
         self.mark_page_start = 0
         self.can_use_page_size = cdiv(self.size, page_size)
 
+        rank_in_node = get_current_rank_in_node()
+        self.shared_can_use_page_num = SharedInt(
+            f"{get_unique_server_name()}_mem_manger_can_use_page_num_{rank_in_node}"
+        )
+        self.shared_can_use_page_num.set_value(self.can_use_page_size)
+
     def _init_buffers(self, size, dtype, head_num, head_dim, layer_num):
         self.kv_buffer = torch.empty(
             (layer_num, cdiv(size, get_page_size()) * get_page_size(), 2 * head_num, head_dim),
@@ -141,6 +149,7 @@ def alloc(self, need_size, b_req_idx, b_seq_len, b_ready_cache_len=None, is_pref
         token_idxs = self.get_paged_token_indexs(b_req_idx, page_size, b_seq_len, b_ready_cache_len, is_prefill)
         self.can_use_mem_size -= need_size
         self.shared_can_use_token_num.set_value(self.can_use_mem_size)
+        self.shared_can_use_page_num.set_value(self.can_use_page_size)
         return token_idxs
 
     def free(self, free_index: Union[torch.Tensor, List[int]]):
@@ -154,12 +163,13 @@ def free(self, free_index: Union[torch.Tensor, List[int]]):
         if len(free_index) == 0:
             return
 
-        page_indices = free_index // page_size
-        unique_pages = torch.unique(page_indices)
-        for page_idx in sorted(unique_pages, reverse=True):  # 逆序放回，保持池的相对顺序
+        base_free_index = free_index[free_index % page_size == 0]
+        page_indices = base_free_index // page_size
+        for page_idx in sorted(page_indices, reverse=True):  # 逆序放回，保持池的相对顺序
             self.mark_page_start -= 1
             self.page_idx_pool[self.mark_page_start] = page_idx
             self.can_use_page_size += 1
+        self.shared_can_use_page_num.set_value(self.can_use_page_size)
 
         return
 
@@ -168,6 +178,7 @@ def free_all(self):
         page_size = get_page_size()
         self.mark_page_start = 0
         self.can_use_page_size = cdiv(self.size, page_size)
+        self.shared_can_use_page_num.set_value(self.can_use_page_size)
         self.page_idx_pool = torch.arange(
             0, cdiv(self.size, page_size), dtype=torch.int32, device="cpu", requires_grad=False, pin_memory=True
         )

lightllm/server/router/dynamic_prompt/paged_radix_cache.py

@@ -159,7 +159,7 @@ def __init__(self, unique_name, total_token_num, rank_in_node, mem_manager=None)
         )
         self.tree_total_tokens_num.arr[0] = 0
 
-    def _get_page_aligned_key(self, key, value=None):
+    def _get_page_aligned_key(self, key, value=None, free_truncated=False):
         aligned_len = len(key)
         if aligned_len == 0:
             return None, None
@@ -171,6 +171,13 @@ def _get_page_aligned_key(self, key, value=None):
                     aligned_len = aligned_len & ~self._page_size_mask
                 else:
                     aligned_len = (aligned_len // self.page_size) * self.page_size
+
+                # 释放被截断的部分
+                if free_truncated and aligned_len < len(key) and self.mem_manager is not None:
+                    truncated_value = value[aligned_len:] if value is not None else key[aligned_len:]
+                    if len(truncated_value) > 0:
+                        self.mem_manager.free(truncated_value)
+
                 return (
                     key[:aligned_len] if aligned_len > 0 else None,
                     value[:aligned_len] if value is not None and aligned_len > 0 else None,
@@ -182,7 +189,7 @@ def insert(self, key, value=None):
             value = key
 
         assert len(key) == len(value)  # and len(key) >= 1
-        key, value = self._get_page_aligned_key(key, value)
+        key, value = self._get_page_aligned_key(key, value, free_truncated=True)
         if key is None:
             return 0
         return self._insert_helper(self.root_node, key, value)
@@ -422,41 +429,3 @@ def release_mem(mem_index):
                 mem_index = torch.concat(release_mems)
                 self.mem_manager.free(mem_index)
         return
-
-
-class _RadixCacheReadOnlyClient:
-    """
-    router 端只读用的客户端，用于从共享内存中读取树结构中的信息，用于进行prompt cache 的调度估计。
-    """
-
-    def __init__(self, unique_name, total_token_num, rank_in_node):
-        self.refed_tokens_num = SharedArray(f"{unique_name}_refed_tokens_num_{rank_in_node}", (1,), dtype=np.int64)
-        self.tree_total_tokens_num = SharedArray(
-            f"{unique_name}_tree_total_tokens_num_{rank_in_node}", (1,), dtype=np.int64
-        )
-
-    def get_refed_tokens_num(self):
-        return self.refed_tokens_num.arr[0]
-
-    def get_tree_total_tokens_num(self):
-        return self.tree_total_tokens_num.arr[0]
-
-    def get_unrefed_tokens_num(self):
-        return self.tree_total_tokens_num.arr[0] - self.refed_tokens_num.arr[0]
-
-
-class RadixCacheReadOnlyClient:
-    def __init__(self, unique_name, total_token_num, node_world_size, dp_world_size):
-        self.dp_rank_clients: List[_RadixCacheReadOnlyClient] = [
-            _RadixCacheReadOnlyClient(unique_name, total_token_num, rank_in_node)
-            for rank_in_node in range(0, node_world_size, dp_world_size)
-        ]
-
-    def get_refed_tokens_num(self, dp_rank_in_node):
-        return self.dp_rank_clients[dp_rank_in_node].get_refed_tokens_num()
-
-    def get_tree_total_tokens_num(self, dp_rank_in_node):
-        return self.dp_rank_clients[dp_rank_in_node].get_tree_total_tokens_num()
-
-    def get_unrefed_tokens_num(self, dp_rank_in_node):
-        return self.dp_rank_clients[dp_rank_in_node].get_unrefed_tokens_num()

lightllm/server/router/req_queue/chunked_prefill/impl.py

@@ -3,6 +3,11 @@
 from ...batch import Batch, Req
 from lightllm.server.router.req_queue.base_queue import BaseQueue
 from lightllm.common.basemodel.infer_lock import g_router_lock
+from lightllm.utils.envs_utils import get_page_size
+
+
+def cdiv(a, b):
+    return (a + b - 1) // b
 
 
 class ChunkedPrefillQueue(BaseQueue):
@@ -21,8 +26,9 @@ def _init_cache_list(self, current_batch: Batch, is_busy):
         return
 
     # @calculate_time(show=True, min_cost_ms=0.1)
-    def _can_add_new_req(self, req: Req, is_busy, new_batch_first_router_need_tokens):
-        self.cache_len_list.append(req.get_tuple_tokens(is_busy, self.router_max_new_token_len))  # hard to analysis
+    def _can_add_new_req(self, req: Req, is_busy, new_batch_first_router_need_tokens, new_batch_prefill_need_pages):
+        token_infos = req.get_tuple_tokens(is_busy, self.router_max_new_token_len)
+        self.cache_len_list.append(token_infos)  # hard to analysis
         self.cache_len_list.sort(key=lambda x: -x[1])
 
         left_out_len_array = np.array([e[1] for e in self.cache_len_list])
@@ -42,16 +48,29 @@ def _can_add_new_req(self, req: Req, is_busy, new_batch_first_router_need_tokens
             new_batch_first_router_need_tokens += req.get_first_router_need_tokens()
             ok_prefill = new_batch_first_router_need_tokens <= self.batch_max_tokens
 
-            if ok_token_num and ok_req_num and ok_prefill:
+            # 检查page
+            ok_page_num = True
+            if "page_size_variable" in self.router.mode:
+                available_pages = self.router.read_only_statics_mem_manager.get_unrefed_page_num(self.dp_index)
+                page_size = get_page_size()
+                if self.router.radix_cache_client is not None:
+                    radix_cache = self.router.radix_cache_client
+                    available_pages += radix_cache.get_unrefed_tokens_num(self.dp_index) // page_size
+
+                new_batch_prefill_need_pages += cdiv(req.input_len + req.shm_cur_output_len, page_size)
+                decode_need_pages = cdiv((left_out_len_array * size_array).max(), page_size)
+                ok_page_num = new_batch_prefill_need_pages + decode_need_pages < available_pages
+
+            if ok_token_num and ok_req_num and ok_prefill and ok_page_num:
                 self.router.shared_token_load.set_estimated_peak_token_count(need_max_token_num, self.dp_index)
                 self.router.shared_token_load.set_dynamic_max_load(
                     (need_max_token_num + self.router.shared_token_load.get_frozened_token_count(self.dp_index))
                     / self.max_total_tokens,
                     self.dp_index,
                 )
-                return True, new_batch_first_router_need_tokens
+                return True, new_batch_first_router_need_tokens, new_batch_prefill_need_pages
             else:
-                return False, new_batch_first_router_need_tokens
+                return False, new_batch_first_router_need_tokens, new_batch_prefill_need_pages
 
     # @calculate_time(show=True, min_cost_ms=10)
     def generate_new_batch(self, current_batch: Batch):
@@ -77,15 +96,16 @@ def generate_new_batch(self, current_batch: Batch):
 
         waiting_queue = self.waiting_req_list
 
+        new_batch_prefill_need_pages = cdiv(new_batch_first_router_need_tokens, get_page_size())
         for req in waiting_queue:
             if req.is_aborted:
                 # 由于管理的复杂性，只有没有被调度运行过的请求可以因为abort直接在队列中忽略掉.
                 # 暂停的请求需要恢复后，由 router manager 部分来过滤。暂时保持这种处理方法, 否则会导致管理token的泄漏
                 aborted_count += 1
                 abort_req_list.append(req)
                 continue
-            ok_insert, new_batch_first_router_need_tokens = self._can_add_new_req(
-                req, is_busy, new_batch_first_router_need_tokens
+            ok_insert, new_batch_first_router_need_tokens, new_batch_prefill_need_pages = self._can_add_new_req(
+                req, is_busy, new_batch_first_router_need_tokens, new_batch_prefill_need_pages
             )
             if ok_insert:
                 can_run_list.append(req)