fix

Author: niushengxiao

lightllm/models/deepseek3_2/layer_infer/transformer_layer_infer.py

@@ -227,15 +227,22 @@ def _get_indices(
 
         import deep_gemm
 
-        logits = deep_gemm.fp8_mqa_logits(q_fp8, (k_fp8_, k_scale_), weights.squeeze(-1), ks, ke)
+        logits = deep_gemm.fp8_mqa_logits(
+            q_fp8,
+            (k_fp8_, k_scale_),
+            weights.squeeze(-1),
+            ks,
+            ke,
+            clean_logits=False,
+            max_seqlen_k=infer_state.max_kv_seq_len,
+        )
 
         from sgl_kernel import fast_topk_v2
 
         b_topk_index = fast_topk_v2(
             score=logits,
             lengths=lengths,
             topk=self.index_topk,
-            row_starts=ks,
         )
         b_topk_index = torch.where(b_topk_index != -1, b_topk_index + ks.view(-1, 1), -1)
         # 将 topk index 转化为 mem index

lightllm/models/deepseek3_2/triton_kernel/extract_indexer_ks.py

@@ -112,4 +112,4 @@ def extract_indexer_ks(
         num_stages=1,
     )
 
-    return O_fp8, O_scale
+    return O_fp8, O_scale.squeeze(-1)