run like deepseek v3

Author: sufubao

lightllm/models/deepseek2/model.py

@@ -50,7 +50,7 @@ def __init__(self, model):
                 self.softmax_scale = self.softmax_scale * mscale * mscale
 
 
-@ModelRegistry(["deepseek_v2", "deepseek_v3"])
+@ModelRegistry(["deepseek_v2", "deepseek_v3", "deepseek_v32"])
 class Deepseek2TpPartModel(LlamaTpPartModel):
     # weight class
     transformer_weight_class = Deepseek2TransformerLayerWeight

lightllm/models/deepseek3_2/infer_struct.py

@@ -1,5 +1,6 @@
 import torch
 from lightllm.models.deepseek2.flashattention_infer_struct import Deepseek2FlashAttentionStateInfo
+from lightllm.models.deepseek3_2.mem_manager import Deepseek3_2MemoryManager
 
 class Deepseek3_2FlashAttentionStateInfo(Deepseek2FlashAttentionStateInfo):
 
@@ -15,6 +16,9 @@ def __init__(self):
 
     def init_some_extra_state(self, model, input_ids: torch.Tensor):
         super().init_some_extra_state(model, input_ids)
+        assert isinstance(self.mem_manager, Deepseek3_2MemoryManager)
+        self.indexer_ks_mem_manager = self.mem_manager.indexer_ks_mem_manager
+
         # Ensure b_ready_cache_len is set for both prefill and decode modes
         if self.is_prefill:
             # b_ready_cache_len is already set in basemodel.py for prefill
@@ -24,9 +28,42 @@ def init_some_extra_state(self, model, input_ids: torch.Tensor):
             # since b_q_seq_len represents the new tokens being processed
             if self.b_ready_cache_len is None:
                 self.b_ready_cache_len = self.b_seq_len - self.b_q_seq_len
-        
-            self.nsa_cache_seqlens = self.b_att_seq_len.clamp(max=model.index_topk)
+
+            self.nsa_cache_seqlens = self.b_att_seq_len.clamp(max=self.index_topk)
             assert self.nsa_cache_seqlens.dtype == torch.int32
             self.nsa_cu_seqlens_k =  torch.nn.functional.pad(
                 torch.cumsum(self.nsa_cache_seqlens, dim=0, dtype=torch.int32), (1, 0)
-            )   
+            )
+
+        # Pre-compute NSA indexer indexing structures
+        self._init_nsa_indexing_structures()
+
+    def _init_nsa_indexing_structures(self):
+        """Pre-compute ks, ke, lengths, and page_table_size_1 for NSA indexer"""
+        mem_index_list = []
+        ks_list = []
+        ke_list = []
+        lengths_list = []
+        offset = 0
+        num_seq_len = self.b_req_idx.shape[0]
+        self.page_table_size_1 = torch.zeros((num_seq_len, self.b_seq_len.max()), dtype=torch.int, device='cuda')
+
+        for i in range(num_seq_len):
+            seq_len = self.b_seq_len[i]
+            q_seq_len = self.b_q_seq_len[i]
+            mem_index = self.req_manager.req_to_token_indexs[i, :seq_len]
+            mem_index_list.append(mem_index)
+            self.page_table_size_1[i, :seq_len] = mem_index
+            ks = torch.zeros(q_seq_len, dtype=torch.int, device='cuda') + offset
+            ke = torch.arange(q_seq_len, dtype=torch.int, device='cuda') + offset + 1
+            ks_list.append(ks)
+            ke_list.append(ke)
+            lengths_list.append(torch.arange(seq_len - q_seq_len + 1, seq_len + 1, dtype=torch.int, device='cuda'))
+            offset += seq_len
+
+        self.mem_index = torch.cat(mem_index_list, dim=0)
+        # ks : [seq_len_q] 标志kv的起始位置
+        # ke : [seq_len_q] 标志kv的结束位置
+        self.ks = torch.cat(ks_list, dim=0)
+        self.ke = torch.cat(ke_list, dim=0)
+        self.lengths = torch.cat(lengths_list, dim=0)

lightllm/models/deepseek3_2/layer_infer/nsa_indexer_layer_inder.py

@@ -10,7 +10,9 @@
 from lightllm.models.deepseek3_2.triton_kernel.act_quant import act_quant
 from lightllm.models.deepseek3_2.mem_manager import Deepseek3_2MemoryManager
 from lightllm.models.deepseek3_2.triton_kernel.destindex_copy_indexer_ks import destindex_copy_indexer_ks
-# from lightllm.models.deepseek3_2.triton_kernel.fp8_mqa_logits import fp8_mqa_logits
+from lightllm.utils.log_utils import init_logger
+
+logger = init_logger(__name__)
 
 class NSAIndexerInfer(BaseLayerInfer):
     def __init__(self, layer_idx, network_config, mode=[]):
@@ -66,70 +68,37 @@ def get_indices(self, hidden_states: torch.Tensor, q_lora: torch.Tensor,
         q_fp8, q_scale = act_quant(q, self.block_size, self.scale_fmt)
         k_fp8, k_scale = act_quant(k, self.block_size, self.scale_fmt)
 
-        self._copy_ks_to_mem_cache(k_fp8, k_scale, infer_state.mem_index, infer_state.mem_manager)
+        destindex_copy_indexer_ks(
+            k_fp8.unsqueeze(1), 
+            k_scale.unsqueeze(1), 
+            infer_state.mem_index,
+            infer_state.indexer_ks_mem_manager.kv_buffer[self.layer_idx_]
+        )
 
         weights = layer_weight.weights_proj_.mm(hidden_states) * self.index_n_heads_scale
         weights = weights.unsqueeze(-1) * q_scale
 
-        ks_buffer = infer_state.mem_manager.indexer_ks_mem_manager.kv_buffer[self.layer_idx_]
-
-        k_fp8_list = []
-        k_scale_list = []
-        ks_list = []
-        ke_list = []
-        offset = 0
-        for i in range(infer_state.batch_size):
-            q_len = infer_state.b_q_seq_len[i]
-            cache_len = infer_state.b_ready_cache_len[i]
-            mem_indexes = infer_state.req_manager.req_to_token_indexs[infer_state.b_req_idx[i], :cache_len+q_len]
-            k_fp8 = ks_buffer[mem_indexes, 0, :128].view(torch.float8_e4m3fn).contiguous()
-            k_scale = ks_buffer[mem_indexes, 0, 128:].view(torch.float32).contiguous()
-            ks = torch.full((q_len,), offset, dtype=torch.int32, device="cuda")
-            ke = ks + torch.arange(q_len, dtype=torch.int32, device="cuda") + 1
-            k_fp8_list.append(k_fp8)
-            k_scale_list.append(k_scale)
-            ks_list.append(ks)
-            ke_list.append(ke)
-            offset += q_len 
-
-        k_fp8 = torch.cat(k_fp8_list, dim=0).view(torch.float8_e4m3fn)
-        k_scale = torch.cat(k_scale_list, dim=0).view(torch.float32).squeeze(-1)
-        kv_fp8 = (k_fp8, k_scale)
-        ks = torch.cat(ks_list, dim=0)
-        ke = torch.cat(ke_list, dim=0)
-
-        logits = deep_gemm.fp8_mqa_logits(
-            q_fp8,
-            kv_fp8,
-            weights.squeeze(-1),
-            ks,
-            ke,
-            clean_logits=False,
-        )
-
-        return self.get_topk(logits, infer_state)
-
-    def get_topk(self, logits, infer_state: Deepseek3_2FlashAttentionStateInfo):
-        topk_indices_list = []
-        offset = 0
-
-        for i in range(infer_state.batch_size):
-            q_len = infer_state.b_q_seq_len[i]
-            cache_len = infer_state.b_ready_cache_len[i]
-            end_pos = q_len + cache_len
-            # Slice logits for this batch (both query and sequence dimensions)
-            batch_logits = logits[offset:offset + q_len, :end_pos]
-            topk_indices = batch_logits.topk(min(self.index_topk, end_pos), dim=-1)[1]
-            mem_indexes = infer_state.req_manager.req_to_token_indexs[infer_state.b_req_idx[i], :cache_len+q_len]
-            indices = torch.full((q_len, self.index_topk), -1, dtype=torch.int32, device="cuda")
-            for j in range(q_len):
-                indices[j, :topk_indices[j].shape[0]] = mem_indexes[topk_indices[j]]
-            topk_indices_list.append(indices)
-            offset += q_len
+        # Use pre-computed indexing structures from infer_state
+        mem_index = infer_state.mem_index
+        ks = infer_state.ks
+        ke = infer_state.ke
+        lengths = infer_state.lengths
+        page_table_1 = infer_state.page_table_size_1
 
-        topk_indices_ = torch.cat(topk_indices_list, dim=0)
+        # TODO
+        k_fp8_ = infer_state.indexer_ks_mem_manager.kv_buffer[self.layer_idx_][mem_index, :, :128].view(torch.float8_e4m3fn).squeeze(1).contiguous()
+        k_scale_ = infer_state.indexer_ks_mem_manager.kv_buffer[self.layer_idx_][mem_index, :, 128:].view(torch.float32)[:, 0, 0].contiguous()
 
-        return topk_indices_
+        logits = deep_gemm.fp8_mqa_logits(q_fp8, (k_fp8_, k_scale_), weights.squeeze(-1), ks, ke) 
+        
+        # 返回 ： [seq_q_len, topk] 无效的位置使用-1填充
+        return fast_topk_transform_fused(
+            score=logits, # [seq_len_q, seq_len_kv]
+            lengths=lengths, # [seq_len_q]
+            page_table_size_1=page_table_1, # [seq_len_q, max(lengths)] 无效的使用0填充
+            cu_seqlens_q=infer_state.cu_seqlens_q, # [seq_len_q + 1]
+            topk=self.index_topk,
+        )
 
 
     def get_k_float32_from_buffer(self, buffer: torch.Tensor):
@@ -152,8 +121,9 @@ def _rotate_activation(x: torch.Tensor) -> torch.Tensor:
     def _get_q_k_bf16(self, hidden_states: torch.Tensor, q_lora: torch.Tensor,
                      infer_state: Deepseek3_2FlashAttentionStateInfo, layer_weight: NSAIndexerWeight):
         q = layer_weight.wq_b_proj_.mm(q_lora).view(-1, self.index_n_heads, self.index_head_dim)
-
         k = layer_weight.wk_proj_.mm(hidden_states)
+
+        # TODO
         k = F.layer_norm(
             k.float(), (self.index_head_dim,), layer_weight.k_norm_.weight, layer_weight.k_norm_.bias, self.eps
         ).type_as(k)
@@ -168,17 +138,3 @@ def _get_q_k_bf16(self, hidden_states: torch.Tensor, q_lora: torch.Tensor,
         q = self._rotate_activation(q)
         k = self._rotate_activation(k)
         return q, k
-
-    def _copy_ks_to_mem_cache(self, k_fp8, k_scale, mem_index, mem_manager: Deepseek3_2MemoryManager):
-        # k_fp8 : [seq_len, 128] torch.fp8_e4m3
-        # k_scale : [seq_len, 1] torch.float32
-        # mem_index : [seq_len] torch.int32
-        # buffer : [10000000, 1, 132] torch.uint8
-        buffer = mem_manager.indexer_ks_mem_manager.kv_buffer[self.layer_idx_]
-        destindex_copy_indexer_ks(
-            k_fp8.unsqueeze(1),  # Add head dimension: [seq_len, 1, 128]
-            k_scale.unsqueeze(1),  # Add head dimension: [seq_len, 1, 1]
-            mem_index,
-            buffer
-        )
-        return

lightllm/models/deepseek3_2/layer_infer/transformer_layer_infer.py

@@ -82,10 +82,9 @@ def _nsa_context_attention_kernel(
         q_nope, q_rope = torch.split(q, [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1)
         q_nope = layer_weight.k_b_proj_.bmm(q_nope.transpose(0, 1)).transpose(0, 1)
         q_all = torch.cat([q_nope, q_rope], dim=-1)
-
         mla_out, _, _ = flash_mla_sparse_fwd(
-            q=q_all,
-            kv=infer_state.mem_manager.kv_buffer[self.layer_num_],
+            q=q_all, # [seq_len_q, q_num_head, qk_dim]
+            kv=infer_state.mem_manager.kv_buffer[self.layer_num_], # [size, 1, qk_dim]
             indices=self.topk_indices.unsqueeze(1),
             sm_scale=self.softmax_scale,
             d_v=self.kv_lora_rank,
@@ -100,15 +99,16 @@ def _nsa_token_attention_kernel(
         kv = infer_state.mem_manager.kv_buffer[self.layer_num_]
         k_rope = kv[:, :, -self.qk_rope_head_dim :].reshape(-1, 1, 1, self.qk_rope_head_dim)
         kv_nope = kv[:, :, : -self.qk_rope_head_dim].reshape(-1, 1, 1, self.kv_lora_rank)
+
         o_tensor = flash_attn_with_kvcache(
-            q=q_rope,
-            k_cache=k_rope,
-            v_cache=kv_nope,
-            qv=q_nope,
-            page_table=self.topk_indices,
-            cache_seqlens=infer_state.nsa_cache_seqlens,
-            cu_seqlens_q=infer_state.cu_seqlens_q,
-            cu_seqlens_k_new=infer_state.nsa_cu_seqlens_k,
+            q=q_rope, # (q_seqlen, nheads, qk_headdim)
+            k_cache=k_rope, # (kv_size, 1, 1, qk_head_dim)
+            v_cache=kv_nope, # (kv_size, 1, 1, kv_lora_rank)
+            qv=q_nope, # (q_seqlen, nheads, kv_lora_rank)
+            page_table=self.topk_indices, # (q_seqlen, max_seq_len) 
+            cache_seqlens=infer_state.nsa_cache_seqlens, # (q_seqlen) # 表示当前kv长度，用于读取page_table.
+            cu_seqlens_q=infer_state.cu_seqlens_q, # (batch_size+1) [0,1]
+            cu_seqlens_k_new=infer_state.nsa_cu_seqlens_k, #(batch_size+1) [0,9]
             max_seqlen_q=infer_state.max_q_seq_len,
             softmax_scale=self.softmax_scale,
             causal=True,

lightllm/models/deepseek3_2/model.py

@@ -5,7 +5,7 @@
 from lightllm.utils.envs_utils import get_env_start_args
 from lightllm.models.deepseek3_2.infer_struct import Deepseek3_2FlashAttentionStateInfo
 from lightllm.models.deepseek3_2.mem_manager import Deepseek3_2MemoryManager, Deepseek3_2FP8KVMemoryManager
-@ModelRegistry(["deepseek_v32"])
+# @ModelRegistry(["deepseek_v32"])
 class Deepseek3_2TpPartModel(Deepseek2TpPartModel):
     # weight class
     transformer_weight_class = Deepseek3_2TransformerLayerWeight
@@ -21,6 +21,9 @@ def __init__(self, kvargs):
         self.index_topk = self.config["index_topk"]
         return
 
+    def _init_inferstate_cls(self):
+        self.infer_state_class = Deepseek3_2FlashAttentionStateInfo
+
     def _init_mem_manager(self):
         manager_class = Deepseek3_2MemoryManager
         if "triton_fp8kv" in self.mode: