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/models/deepseek2/flashattention_infer_struct.py

@@ -16,6 +16,7 @@ class Deepseek2FlashAttentionStateInfo(Deepseek2InferStateInfo):
 
     def __init__(self):
         super().__init__()
+        self.page_size = get_page_size()
 
     @classmethod
     def get_page_table_buffer(cls, graph_max_batch_size: int, max_seq_len: int):
@@ -43,19 +44,18 @@ def init_some_extra_state(self, model, input_ids: torch.Tensor):
             self.cu_seqlens_q = self.b1_cu_q_seq_len
             self.cu_seqlens_k = self.b1_cu_kv_seq_len
             max_seq_len_k = self.max_kv_seq_len
-            page_size = get_page_size()
             if self.batch_size <= model.graph_max_batch_size and self.max_len_in_batch <= model.graph_max_len_in_batch:
-                length = cdiv(model.graph_max_len_in_batch, page_size)
+                length = cdiv(model.graph_max_len_in_batch, self.page_size)
                 page_buffer = Deepseek2FlashAttentionStateInfo.get_page_table_buffer(model.graph_max_batch_size, length)
                 self.page_table = page_buffer[self.microbatch_index][: self.batch_size * length].reshape(
                     self.batch_size, length
                 )
             else:
-                length = cdiv(self.max_len_in_batch, page_size)
+                length = cdiv(self.max_len_in_batch, self.page_size)
                 self.page_table = torch.empty((self.batch_size, length), dtype=torch.int32).to(input_ids.device)
 
             if "page_size_variable" in model.mode:
-                length = cdiv(max_seq_len_k, page_size)
+                length = cdiv(max_seq_len_k, self.page_size)
                 self.page_table[:, :length].copy_(model.req_manager.req_to_page_indexs[self.b_req_idx, :length])
                 self.page_table[:, length:].fill_(0)
             else:

lightllm/models/deepseek2/layer_infer/transformer_layer_infer.py

@@ -26,7 +26,7 @@
 from functools import partial
 from lightllm.models.llama.yarn_rotary_utils import get_deepseek_mscale
 from lightllm.distributed.communication_op import all_gather, all_gather_into_tensor, all_reduce, reduce_scatter_tensor
-from lightllm.utils.envs_utils import get_env_start_args, get_page_size
+from lightllm.utils.envs_utils import get_env_start_args
 from lightllm.utils.dist_utils import get_global_world_size
 from lightllm.utils.log_utils import init_logger
 from lightllm.utils.sgl_utils import flash_attn_varlen_func, flash_attn_with_kvcache, merge_state_v2
@@ -589,12 +589,11 @@ def _token_gqa_decode_attention_flashattention(
     def _token_gqa_decode_attention_flashattention_paged(
         self, q, infer_state: Deepseek2FlashInferStateInfo, layer_weight: Deepseek2TransformerLayerWeight, out=None
     ):
-        page_size = get_page_size()
         q_nope, q_rope = q[:, :, : -self.qk_rope_head_dim], q[:, :, -self.qk_rope_head_dim :]
         q_nope = layer_weight.k_b_proj_.bmm(q_nope.transpose(0, 1)).transpose(0, 1)
         kv = infer_state.mem_manager.kv_buffer[self.layer_num_]
-        k_rope = kv[:, :, -self.qk_rope_head_dim :].reshape(-1, page_size, 1, self.qk_rope_head_dim)
-        kv_nope = kv[:, :, : -self.qk_rope_head_dim].reshape(-1, page_size, 1, self.kv_lora_rank)
+        k_rope = kv[:, :, -self.qk_rope_head_dim :].reshape(-1, infer_state.page_size, 1, self.qk_rope_head_dim)
+        kv_nope = kv[:, :, : -self.qk_rope_head_dim].reshape(-1, infer_state.page_size, 1, self.kv_lora_rank)
         k_descale, v_descale = None, None
         o_tensor = flash_attn_with_kvcache(
             q=q_rope,
@@ -638,7 +637,6 @@ def _token_gqa_decode_attention_flashinfer(
     def _token_gqa_decode_attention_flashinfer_paged(
         self, q, infer_state: Deepseek2FlashInferStateInfo, layer_weight: Deepseek2TransformerLayerWeight, out=None
     ):
-        page_size = get_page_size()
         q_nope, q_rope = q[:, :, : -self.qk_rope_head_dim], q[:, :, -self.qk_rope_head_dim :]
         q_nope = layer_weight.k_b_proj_.bmm(q_nope.transpose(0, 1)).transpose(0, 1)
 
@@ -648,8 +646,8 @@ def _token_gqa_decode_attention_flashinfer_paged(
         infer_state.decode_wrapper.run(
             q_nope,
             q_rope,
-            kv[:, :, : -self.qk_rope_head_dim].reshape(-1, page_size, 1, self.kv_lora_rank),
-            kv[:, :, -self.qk_rope_head_dim :].reshape(-1, page_size, 1, self.qk_rope_head_dim),
+            kv[:, :, : -self.qk_rope_head_dim].reshape(-1, infer_state.page_size, 1, self.kv_lora_rank),
+            kv[:, :, -self.qk_rope_head_dim :].reshape(-1, infer_state.page_size, 1, self.qk_rope_head_dim),
             out=o_tensor,
             return_lse=False,
         )

lightllm/models/llama/flashattention_infer_struct.py

@@ -18,6 +18,7 @@ class FlashAttentionStateInfo(LlamaInferStateInfo):
 
     def __init__(self):
         super().__init__()
+        self.page_size = get_page_size()
 
     @classmethod
     def get_page_table_buffer(cls, graph_max_batch_size: int, max_seq_len: int):
@@ -32,7 +33,7 @@ def _init_flash_attention_state(self, model, input_ids: torch.Tensor):
         if self.is_prefill:
             self.cu_seqlens_q = self.b1_cu_q_seq_len.int()
             self.cu_seqlens_k = self.b1_cu_kv_seq_len.int()
-            length = cdiv(self.max_seq_len, get_page_size())
+            length = cdiv(self.max_seq_len, self.page_size)
             self.page_table = torch.empty((self.batch_size, length), dtype=torch.int32, device=input_ids.device)
             if "page_size_variable" in model.mode:
                 self.page_table.copy_(model.req_manager.req_to_page_indexs[self.b_req_idx, :length])
@@ -44,17 +45,16 @@ def _init_flash_attention_state(self, model, input_ids: torch.Tensor):
             self.cu_seqlens_k = self.b1_cu_kv_seq_len.int()
             max_seq_len_k = self.max_kv_seq_len
             if self.batch_size <= model.graph_max_batch_size and self.max_len_in_batch <= model.graph_max_len_in_batch:
-                page_size = get_page_size()
-                length = cdiv(model.graph_max_len_in_batch, page_size)
+                length = cdiv(model.graph_max_len_in_batch, self.page_size)
                 page_buffer = FlashAttentionStateInfo.get_page_table_buffer(model.graph_max_batch_size, length)
                 self.page_table = page_buffer[self.microbatch_index][: self.batch_size * length].reshape(
                     self.batch_size, length
                 )
             else:
-                length = cdiv(self.max_len_in_batch, get_page_size())
+                length = cdiv(self.max_len_in_batch, self.page_size)
                 self.page_table = torch.empty((self.batch_size, length), dtype=torch.int32, device=input_ids.device)
 
-            length = cdiv(max_seq_len_k, get_page_size())
+            length = cdiv(max_seq_len_k, self.page_size)
             if "page_size_variable" in model.mode:
                 self.page_table[:, :length].copy_(model.req_manager.req_to_page_indexs[self.b_req_idx, :length])
             else:

lightllm/models/llama/layer_infer/transformer_layer_infer.py

@@ -27,7 +27,7 @@
 from lightllm.models.llama.triton_kernel.ppl_quant_copy_kv import destindex_copy_dequantize_kv
 from lightllm.distributed.communication_op import all_gather_into_tensor, reduce_scatter_tensor
 from lightllm.utils.log_utils import init_logger
-from lightllm.utils.envs_utils import get_env_start_args, get_page_size
+from lightllm.utils.envs_utils import get_env_start_args
 from lightllm.utils.light_utils import HAS_LIGHTLLM_KERNEL, light_ops
 from lightllm.common.basemodel.triton_kernel.q_per_head_fp8_quant import q_per_head_fp8_quant
 from lightllm.utils.vllm_utils import HAS_VLLM, vllm_ops
@@ -291,9 +291,8 @@ def _paged_context_attention_flashinfer_kernel(
         self, q, kv, infer_state: LlamaFlashInferStateInfo, layer_weight, out=None
     ) -> torch.Tensor:
         o_tensor = self.alloc_tensor(q.shape, q.dtype) if out is None else out
-        page_size = get_page_size()
         kv = infer_state.mem_manager.kv_buffer[self.layer_num_].view(
-            -1, page_size, 2 * self.tp_k_head_num_, self.head_dim_
+            -1, infer_state.page_size, 2 * self.tp_k_head_num_, self.head_dim_
         )
         infer_state.prefill_wrapper.run(
             q.view(q.shape[0], -1, self.head_dim_),
@@ -356,13 +355,12 @@ def _context_attention_kernel_ppl_int8kv(
     def _paged_context_attention_flashattention(
         self, q, kv, infer_state: FlashAttentionStateInfo, layer_weight, out=None
     ):
-        page_size = get_page_size()
         cache_k = infer_state.mem_manager.kv_buffer[self.layer_num_][:, 0 : self.tp_k_head_num_, :].reshape(
-            -1, page_size, self.tp_k_head_num_, self.head_dim_
+            -1, infer_state.page_size, self.tp_k_head_num_, self.head_dim_
         )
         cache_v = infer_state.mem_manager.kv_buffer[self.layer_num_][
             :, self.tp_k_head_num_ : self.tp_k_head_num_ + self.tp_v_head_num_, :
-        ].reshape(-1, page_size, self.tp_v_head_num_, self.head_dim_)
+        ].reshape(-1, infer_state.page_size, self.tp_v_head_num_, self.head_dim_)
         q = q.reshape(-1, self.tp_q_head_num_, self.head_dim_)
         k_descale, v_descale = None, None  # disable quantization
         Lq = q.shape[-1]
@@ -622,9 +620,8 @@ def _paged_token_decode_attention_flashinfer(
         calcu_shape1 = (batch_size, self.tp_q_head_num_, self.head_dim_)
 
         o_tensor = self.alloc_tensor(q.shape, q.dtype) if out is None else out
-        page_size = get_page_size()
         kv = infer_state.mem_manager.kv_buffer[self.layer_num_].view(
-            -1, page_size, 2 * self.tp_k_head_num_, self.head_dim_
+            -1, infer_state.page_size, 2 * self.tp_k_head_num_, self.head_dim_
         )
         infer_state.decode_wrapper.run(
             q.view(calcu_shape1),
@@ -914,13 +911,12 @@ def _token_decode_attention_gqa_flashdecoding_vsm(
     def _paged_token_decode_attention_flashattention(
         self, q, infer_state: FlashAttentionStateInfo, layer_weight, out=None
     ):
-        page_size = get_page_size()
         cache_k = infer_state.mem_manager.kv_buffer[self.layer_num_][:, 0 : self.tp_k_head_num_, :].reshape(
-            -1, page_size, self.tp_k_head_num_, self.head_dim_
+            -1, infer_state.page_size, self.tp_k_head_num_, self.head_dim_
         )
         cache_v = infer_state.mem_manager.kv_buffer[self.layer_num_][
             :, self.tp_k_head_num_ : self.tp_k_head_num_ + self.tp_v_head_num_, :
-        ].reshape(-1, page_size, self.tp_v_head_num_, self.head_dim_)
+        ].reshape(-1, infer_state.page_size, self.tp_v_head_num_, self.head_dim_)
         q = q.reshape(-1, self.tp_q_head_num_, self.head_dim_)
         k_descale, v_descale = None, None  # disable quantization
         Lq = q.shape[-1]

lightllm/server/api_cli.py

@@ -179,7 +179,7 @@ def make_argument_parser() -> argparse.ArgumentParser:
         nargs="+",
         help="""Model mode: [triton_int8kv | ppl_int8kv | ppl_fp16 | triton_flashdecoding
                         | triton_gqa_attention | triton_gqa_flashdecoding | triton_fp8kv | offline_calibration_fp8kv
-                        | export_fp8kv_calibration
+                        | export_fp8kv_calibration | page_size_variable
                         triton_flashdecoding mode is for long context, current support llama llama2 qwen;
                         triton_gqa_attention and triton_gqa_flashdecoding is fast kernel for model which use GQA;
                         triton_int8kv mode use int8 to store kv cache, can increase token capacity, use triton kernel;
@@ -191,6 +191,8 @@ def make_argument_parser() -> argparse.ArgumentParser:
                         Calibration need to disable cudagraph and use fa3 or flashinfer backend.
                         ppl_int8kv mode use int8 to store kv cache, and use ppl fast kernel;
                         ppl_fp16 mode use ppl fast fp16 decode attention kernel;
+                        page_size_variable allow to use page size > 1, use PAGE_SIZE env to set page size,
+                        page_size_variable only support fa3 and flashinfer backend for now
                         you need to read source code to make sure the supported detail mode for all models""",
     )
     parser.add_argument(

lightllm/server/api_start.py

@@ -94,7 +94,7 @@ def normal_or_p_d_start(args):
 
     if args.graph_max_len_in_batch == 0:
         args.graph_max_len_in_batch = args.max_req_total_len
-    
+
     # mode setting check.
     if args.output_constraint_mode != "none":
         assert args.disable_dynamic_prompt_cache is False
@@ -126,6 +126,13 @@ def normal_or_p_d_start(args):
             "--enable_flashinfer_prefill and --enable_flashinfer_decode"
         )
         assert args.disable_cudagraph is True, "export_fp8kv_calibration mode need disable cudagraph"
+    if "page_size_variable" in args.mode:
+        assert args.enable_fa3 is True or (
+            args.enable_flashinfer_prefill is True and args.enable_flashinfer_decode is True
+        ), (
+            "page_size_variable mode need enable fa3 or flashinfer, add --enable_fa3 or "
+            "--enable_flashinfer_prefill and --enable_flashinfer_decode"
+        )
 
     # 部分模式还不能支持与高级动态调度算法协同，to do.
     if args.diverse_mode:

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()