feat: support page size variable for deepseek2

Author: niushengxiao

lightllm/common/deepseek2_page_size_variable_mem_manager.py

@@ -0,0 +1,25 @@
+import torch
+import numpy as np
+from .deepseek2_mem_manager import Deepseek2MemoryManager
+from .page_size_variable_mem_manager import PageSizeVariableMemoryManager
+from lightllm.utils.log_utils import init_logger
+from lightllm.utils.envs_utils import get_page_size
+
+
+def cdiv(a, b):
+    return (a + b - 1) // b
+
+
+logger = init_logger(__name__)
+
+
+class Deepseek2PageSizeVariableMemoryManager(PageSizeVariableMemoryManager, 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)
+
+    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(), head_num, head_dim),
+            dtype=dtype,
+            device="cuda",
+        )

lightllm/models/deepseek2/flashattention_infer_struct.py

@@ -4,6 +4,11 @@
 import torch.distributed as dist
 from lightllm.models.deepseek2.infer_struct import Deepseek2InferStateInfo
 from lightllm.utils.dist_utils import get_current_device_id
+from lightllm.utils.envs_utils import get_page_size
+
+
+def cdiv(a, b):
+    return (a + b - 1) // b
 
 
 class Deepseek2FlashAttentionStateInfo(Deepseek2InferStateInfo):
@@ -38,20 +43,24 @@ 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:
-                page_buffer = Deepseek2FlashAttentionStateInfo.get_page_table_buffer(
-                    model.graph_max_batch_size, model.graph_max_len_in_batch
+                length = cdiv(model.graph_max_len_in_batch, 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
                 )
-                self.page_table = page_buffer[self.microbatch_index][
-                    : self.batch_size * model.graph_max_len_in_batch
-                ].reshape(self.batch_size, model.graph_max_len_in_batch)
             else:
-                self.page_table = torch.empty((self.batch_size, self.max_len_in_batch), dtype=torch.int32).to(
-                    input_ids.device
-                )
+                length = cdiv(self.max_len_in_batch, page_size)
+                self.page_table = torch.empty((self.batch_size, length), dtype=torch.int32).to(input_ids.device)
 
-            self.page_table[:, :max_seq_len_k].copy_(
-                model.req_manager.req_to_token_indexs[self.b_req_idx, :max_seq_len_k]
-            )
-            self.page_table[:, max_seq_len_k:].fill_(0)
+            if "page_size_variable" in model.mode:
+                length = cdiv(max_seq_len_k, 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:
+                self.page_table[:, :max_seq_len_k].copy_(
+                    model.req_manager.req_to_token_indexs[self.b_req_idx, :max_seq_len_k]
+                )
+                self.page_table[:, max_seq_len_k:].fill_(0)
         return

lightllm/models/deepseek2/flashinfer_struct.py

@@ -3,16 +3,21 @@
 import numpy as np
 import torch.distributed as dist
 from lightllm.models.deepseek2.infer_struct import Deepseek2InferStateInfo
-from lightllm.utils.envs_utils import get_env_start_args
+from lightllm.utils.envs_utils import get_env_start_args, get_page_size
 from lightllm.models.deepseek2.triton_kernel.repack_kv_index import repack_kv_index
 
 
+def cdiv(a, b):
+    return (a + b - 1) // b
+
+
 class Deepseek2FlashInferStateInfo(Deepseek2InferStateInfo):
     def __init__(self):
         super().__init__()
         self.prefill_wrapper = None
         self.decode_wrapper = None
         self.flashinfer_extra_state = None
+        self.page_size = get_page_size()
 
     def init_some_extra_state(self, model, input_ids: torch.Tensor):
         super().init_some_extra_state(model, input_ids)
@@ -23,24 +28,37 @@ def init_some_extra_state(self, model, input_ids: torch.Tensor):
         if not self.is_prefill:
             if get_env_start_args().enable_flashinfer_decode:
                 self.q_indptr = torch.arange(self.batch_size + 1, dtype=torch.int32).to(input_ids.device)
+                length = cdiv(self.flashinfer_extra_state.max_seq_length, self.page_size)
                 if self.batch_size <= model.graph_max_batch_size:
                     self.kv_indices = self.flashinfer_extra_state.kv_indices_buffer[self.microbatch_index][
-                        : self.batch_size * self.flashinfer_extra_state.max_seq_length
+                        : self.batch_size * length
                     ]
                 else:
                     self.kv_indices = torch.empty(
-                        self.batch_size * self.flashinfer_extra_state.max_seq_length,
+                        self.batch_size * length,
                         dtype=torch.int32,
                         device=input_ids.device,
                     )
-                repack_kv_index(
-                    self.req_manager.req_to_token_indexs,
-                    self.b_req_idx,
-                    self.b_seq_len,
-                    self.b_start_loc,
-                    self.max_len_in_batch,
-                    self.kv_indices,
-                )
+                if "page_size_variable" in model.mode:
+                    b_page_len = cdiv(self.b_seq_len, self.page_size)
+                    self.kv_starts[1:] = b_page_len.cumsum(0)
+                    repack_kv_index(
+                        self.req_manager.req_to_page_indexs,
+                        self.b_req_idx,
+                        b_page_len,
+                        self.kv_starts[:-1],
+                        cdiv(self.max_len_in_batch, self.page_size),
+                        self.kv_indices,
+                    )
+                else:
+                    repack_kv_index(
+                        self.req_manager.req_to_token_indexs,
+                        self.b_req_idx,
+                        self.b_seq_len,
+                        self.b_start_loc,
+                        self.max_len_in_batch,
+                        self.kv_indices,
+                    )
                 if self.decode_wrapper is None:
                     self.decode_wrapper = flashinfer.mla.BatchMLAPagedAttentionWrapper(
                         self.flashinfer_extra_state.workspace_buffer,
@@ -58,7 +76,7 @@ def init_some_extra_state(self, model, input_ids: torch.Tensor):
                         self.flashinfer_extra_state.tp_q_head_num,
                         self.flashinfer_extra_state.kv_lora_rank,
                         self.flashinfer_extra_state.qk_rope_head_dim,
-                        1,
+                        self.page_size,
                         False,  # causal
                         self.flashinfer_extra_state.softmax_scale,
                         self.flashinfer_extra_state.q_data_type,
@@ -97,7 +115,7 @@ def copy_for_cuda_graph(self, new_infer_state):
                 new_infer_state.flashinfer_extra_state.tp_q_head_num,
                 new_infer_state.flashinfer_extra_state.kv_lora_rank,
                 new_infer_state.flashinfer_extra_state.qk_rope_head_dim,
-                1,
+                self.page_size,
                 False,  # causal
                 new_infer_state.flashinfer_extra_state.softmax_scale,
                 new_infer_state.flashinfer_extra_state.q_data_type,

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
+from lightllm.utils.envs_utils import get_env_start_args, get_page_size
 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
@@ -94,6 +94,18 @@ def _bind_attention(self):
             self._token_attention_kernel = partial(
                 Deepseek2TransformerLayerInfer._token_gqa_decode_attention_flashdecoding_fp8, self
             )
+        elif "page_size_variable" in self.mode:
+            self._copy_kv_to_mem_cache = partial(Deepseek2TransformerLayerInfer._copy_kv_to_mem_cache_normal, self)
+            if get_env_start_args().enable_fa3:
+                self._token_attention_kernel = partial(
+                    Deepseek2TransformerLayerInfer._token_gqa_decode_attention_flashattention_paged, self
+                )
+            elif get_env_start_args().enable_flashinfer_decode:
+                self._token_attention_kernel = partial(
+                    Deepseek2TransformerLayerInfer._token_gqa_decode_attention_flashinfer_paged, self
+                )
+            else:
+                raise Exception("Page size variable mode is not supported in other backends.")
         else:
             self._copy_kv_to_mem_cache = partial(Deepseek2TransformerLayerInfer._copy_kv_to_mem_cache_normal, self)
             if get_env_start_args().enable_fa3:
@@ -575,6 +587,36 @@ def _token_gqa_decode_attention_flashattention(
         )
         return o_tensor
 
+    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_descale, v_descale = None, None
+        o_tensor = flash_attn_with_kvcache(
+            q=q_rope,
+            k_cache=k_rope,
+            v_cache=kv_nope,
+            qv=q_nope,
+            page_table=infer_state.page_table,
+            cache_seqlens=infer_state.b_seq_len,
+            cu_seqlens_q=infer_state.cu_seqlens_q,
+            cu_seqlens_k_new=infer_state.cu_seqlens_k,
+            max_seqlen_q=1,
+            softmax_scale=self.softmax_scale,
+            causal=True,
+            window_size=(-1, -1),
+            softcap=0.0,
+            k_descale=k_descale,
+            v_descale=v_descale,
+            return_softmax_lse=False,
+        )
+        return o_tensor
+
     def _token_gqa_decode_attention_flashinfer(
         self, q, infer_state: Deepseek2FlashInferStateInfo, layer_weight: Deepseek2TransformerLayerWeight, out=None
     ):
@@ -594,6 +636,26 @@ def _token_gqa_decode_attention_flashinfer(
         )
         return o_tensor
 
+    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)
+
+        kv = infer_state.mem_manager.kv_buffer[self.layer_num_]
+        o_tensor = self.alloc_tensor(q_nope.shape, dtype=q_nope.dtype)
+
+        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),
+            out=o_tensor,
+            return_lse=False,
+        )
+        return o_tensor
+
     def _token_gqa_decode_attention_flashdecoding(
         self, q, infer_state: Deepseek2InferStateInfo, layer_weight: Deepseek2TransformerLayerWeight, out=None
     ):

lightllm/models/deepseek2/model.py

@@ -10,6 +10,7 @@
 
 from lightllm.models.llama.model import LlamaTpPartModel
 from lightllm.common.deepseek2_mem_manager import Deepseek2MemoryManager
+from lightllm.common.deepseek2_page_size_variable_mem_manager import Deepseek2PageSizeVariableMemoryManager
 from lightllm.common.deepseek2_fp8kv_mem_manager import Deepseek2FP8KVMemoryManager
 from lightllm.utils.log_utils import init_logger
 from lightllm.models.llama.yarn_rotary_utils import get_deepseek_mscale
@@ -97,6 +98,10 @@ def _init_mem_manager(self):
         manager_class = Deepseek2MemoryManager
         if "triton_fp8kv" in self.mode:
             manager_class = Deepseek2FP8KVMemoryManager
+        elif "page_size_variable" in self.mode:
+            manager_class = Deepseek2PageSizeVariableMemoryManager
+        elif self.mode:
+            raise ValueError(f"Unsupported mode for deepseek2: {self.mode}")
 
         # mtp 模式下需要在mem manger上扩展draft model使用的layer
         added_mtp_layer_num = 0

lightllm/models/deepseek2/triton_kernel/destindex_copy_kv.py

@@ -34,7 +34,7 @@ def _fwd_kernel_destindex_copy_kv(
     offs_d_nope = tl.arange(0, BLOCK_DMODEL_NOPE)
     offs_d_rope = tl.arange(0, BLOCK_DMODEL_ROPE)
 
-    dest_index = tl.load(Dest_loc + cur_index)
+    dest_index = tl.load(Dest_loc + cur_index).to(tl.int64)
 
     kv_nope_ptrs = KV_nope + cur_index * stride_kv_nope_bs + stride_kv_nope_d * offs_d_nope[None, :]
     kv_rope_ptrs = KV_rope + cur_index * stride_kv_rope_bs + stride_kv_rope_d * offs_d_rope[None, :]

lightllm/models/deepseek2/triton_kernel/sample_kv.py

@@ -44,7 +44,7 @@ def _sample_kv_kernel(
         Req_to_tokens + stride_req_to_tokens_b * cur_batch_req_idx + offs_m,
         mask=offs_m < block_end_loc,
         other=0,
-    )
+    ).to(tl.int64)
     off_kv_nope = kv_loc[:, None] * stride_input_dim + offs_nope_d[None, :]
     off_kv_rope = kv_loc[:, None] * stride_input_dim + (offs_rope_d + BLOCK_DMODEL)[None, :]
     kv_nope = tl.load(KV_input + off_kv_nope, mask=offs_m[:, None] < block_end_loc, other=0.0)