need fix

Author: sufubao

lightllm/models/deepseek3_2/__init__.py

@@ -9,8 +9,8 @@ def __init__(self):
         self.page_table_size_1 = None
         self.ks = None
         self.ke = None
-
-        self.topk_indices = None
+        self.nsa_cu_seqlens_k = None
+        self.index_topk = 2048
         return
 
     def init_some_extra_state(self, model, input_ids: torch.Tensor):
@@ -24,3 +24,9 @@ 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)
+            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)
+            )   

lightllm/models/deepseek3_2/infer_struct.py

@@ -86,7 +86,7 @@ def _nsa_context_attention_kernel(
         mla_out, _, _ = flash_mla_sparse_fwd(
             q=q_all,
             kv=infer_state.mem_manager.kv_buffer[self.layer_num_],
-            indices=self.topk_indices,
+            indices=self.topk_indices.unsqueeze(1),
             sm_scale=self.softmax_scale,
             d_v=self.kv_lora_rank,
         )
@@ -100,23 +100,17 @@ 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)
-        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=self.topk_indices,
-            cache_seqlens=infer_state.b_att_seq_len,
+            cache_seqlens=infer_state.nsa_cache_seqlens,
             cu_seqlens_q=infer_state.cu_seqlens_q,
-            cu_seqlens_k_new=infer_state.cu_seqlens_k,
+            cu_seqlens_k_new=infer_state.nsa_cu_seqlens_k,
             max_seqlen_q=infer_state.max_q_seq_len,
             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

lightllm/models/deepseek3_2/layer_infer/transformer_layer_infer.py

@@ -16,6 +16,11 @@ class Deepseek3_2TpPartModel(Deepseek2TpPartModel):
     # infer state class
     infer_state_class = Deepseek3_2FlashAttentionStateInfo
 
+    def __init__(self, kvargs):
+        super().__init__(kvargs)
+        self.index_topk = self.config["index_topk"]
+        return
+
     def _init_mem_manager(self):
         manager_class = Deepseek3_2MemoryManager
         if "triton_fp8kv" in self.mode: