DeepSeek-V3.2: Add Adaptive MHA Attention Pathway for Short-Sequence Prefill (#11892)

python/sglang/srt/models/deepseek_v2.py

@@ -398,6 +398,34 @@ def handle_attention_aiter(attn, forward_batch):
 
 
 def handle_attention_nsa(attn, forward_batch):
+    """
+    Select MHA or MLA based on sequence length for optimal performance.
+
+    - Decode: MLA (avoids per-token decompression)
+    - Prefill <= 2048: MHA (topk ineffective, MHA has lower FLOPs)
+    - Prefill > 2048: MLA (topk filtering reduces computation significantly)
+
+    TODO: B200 (SM100) MHA path is temporarily disabled due to FA4 gpqa accuracy issues.
+    """
+    if forward_batch.forward_mode.is_decode_or_idle():
+        return AttnForwardMethod.MLA
+
+    if _is_extend_without_speculative(forward_batch):
+        assert forward_batch.seq_lens_cpu is not None
+        max_kv_len = forward_batch.seq_lens_cpu.max().item()
+
+        # B200 (SM100) is temporarily disabled for MHA due to FA4 accuracy issues
+        # Currently only H200 (SM90) with FA3 is allowed to use MHA path
+        is_hopper = _device_sm == 90
+
+        if max_kv_len <= attn.indexer.index_topk and is_hopper:
+            # NSA backend uses varlen kernel which supports MHA_ONE_SHOT
+            # Check if total sequence length fits in chunk capacity
+            sum_seq_lens = sum(forward_batch.seq_lens_cpu)
+            # Use MHA_ONE_SHOT for best performance
+            if sum_seq_lens <= forward_batch.get_max_chunk_capacity():
+                return AttnForwardMethod.MHA_ONE_SHOT
+
     return AttnForwardMethod.MLA
 
 
@@ -1466,8 +1494,21 @@ class DeepseekV2AttentionMLA(nn.Module):
             q, latent_cache = self.fused_qkv_a_proj_with_mqa(hidden_states)[0].split(
                 [self.q_lora_rank, self.kv_lora_rank + self.qk_rope_head_dim], dim=-1
             )
-            q = self.q_a_layernorm(q)
-            q = self.q_b_proj(q)[0].view(-1, self.num_local_heads, self.qk_head_dim)
+            q_lora = self.q_a_layernorm(q)
+            q = self.q_b_proj(q_lora)[0].view(
+                -1, self.num_local_heads, self.qk_head_dim
+            )
+
+            # NSA Indexer: cache quantized keys, auto-skip topk for sequences <= nsa_index_topk
+            if self.use_nsa and _is_extend_without_speculative(forward_batch):
+                _ = self.indexer(
+                    x=hidden_states,
+                    q_lora=q_lora,
+                    positions=positions,
+                    forward_batch=forward_batch,
+                    layer_id=self.layer_id,
+                    return_indices=False,
+                )
         else:
             q = self.q_proj(hidden_states)[0].view(
                 -1, self.num_local_heads, self.qk_head_dim