[SM120] Use exact query-head widths for DeepSeek-V4 sparse MLA decode (#36655)

python/sglang/kernels/ops/attention/flash_mla_sm120.py

@@ -13,7 +13,8 @@ separate region at the end of each page.
 
 import logging
 import math
-from typing import Optional
+from functools import lru_cache
+from typing import FrozenSet, Optional, Tuple
 
 import torch
 import triton
@@ -266,6 +267,34 @@ def _flash_mla_sm120_prefill(
     return (output.unsqueeze(1), None)
 
 
+@lru_cache(maxsize=1)
+def _flashinfer_dsv4_decode_capabilities() -> Tuple[int, FrozenSet[int]]:
+    """Read the installed FlashInfer DSV4 decode capabilities once."""
+    try:
+        from flashinfer.mla._sparse_mla_sm120 import (
+            _DECODE_DSV4_DISPATCH,
+            _DECODE_MAX_TOKENS,
+        )
+    except (AttributeError, ImportError):
+        return 0, frozenset()
+
+    return int(_DECODE_MAX_TOKENS), frozenset(
+        heads for heads, _ in _DECODE_DSV4_DISPATCH
+    )
+
+
+def flashinfer_dsv4_decode_supports_num_heads(num_heads: int, num_tokens: int) -> bool:
+    """Return whether FlashInfer supports this DSV4 decode head count.
+
+    Keep this capability check fail-closed because SGLang can be used with a
+    locally installed FlashInfer even though the release dependency is pinned.
+    The padded 64-head decode path remains the safe fallback for older builds.
+    Prefill head selection is handled separately by the caller.
+    """
+    decode_max_tokens, supported_heads = _flashinfer_dsv4_decode_capabilities()
+    return num_tokens <= decode_max_tokens and num_heads in supported_heads
+
+
 def flash_mla_with_kvcache_sm120(**kwargs):
     """SM120 FlashMLA sparse decode entry point.
 

python/sglang/srt/models/deepseek_v4.py

@@ -164,7 +164,6 @@ from sglang.srt.utils import (
     is_gfx95_supported,
     is_gfx942_supported,
     is_gfx1250_supported,
-    is_sm120_supported,
     log_info_on_rank0,
     make_layers,
 )
@@ -821,17 +820,52 @@ class MqaAttentionBase(nn.Module):
         self.register_buffer("freqs_cis", freqs_cis, persistent=False)
         self.freqs_cis: torch.Tensor
 
-    def _local_attn_sink(self) -> torch.Tensor:
+    def _kernel_num_heads(self, num_tokens: int) -> int:
+        if self.attn_tp_size == 1:
+            return self.n_local_heads
+
+        if get_platform().is_sm120:
+            # Prefill already accepts the native per-rank query width.
+            if num_tokens > SM120_DECODE_MAX_TOKENS:
+                return self.n_local_heads
+
+            if envs.SGLANG_SM120_FLASHMLA_BACKEND.get() == "flashinfer":
+                from sglang.kernels.ops.attention.flash_mla_sm120 import (
+                    flashinfer_dsv4_decode_supports_num_heads,
+                )
+
+                if flashinfer_dsv4_decode_supports_num_heads(
+                    self.n_local_heads, num_tokens
+                ):
+                    return self.n_local_heads
+
+        # Other FlashMLA implementations retain their existing padded shape.
+        return 64 if self.n_local_heads <= 64 else self.n_heads
+
+    def _local_attn_sink(self, kernel_num_heads: Optional[int] = None) -> torch.Tensor:
         if self.attn_tp_size == 1:
             return self.attn_sink
+
+        rank = self.attn_tp_rank
+        num_heads = self.n_local_heads
+        padded_num_heads = 64 if num_heads <= 64 else self.n_heads
+        if kernel_num_heads is None:
+            # Preserve the legacy contract for subclasses such as DSpark that
+            # always pad their attention query independently of this helper.
+            kernel_num_heads = padded_num_heads
+        assert kernel_num_heads >= num_heads
+
+        # Keep one fallback-width allocation and return a view matching Q.
+        # Prefill and decode can alternate, and CUDA graphs can retain the
+        # view, so replacing this tensor when the path changes would
+        # both reallocate every transition and risk invalidating a captured
+        # pointer.
+        sink_num_heads = max(kernel_num_heads, padded_num_heads)
         if self._attn_sink_local is None:
-            rank = self.attn_tp_rank
-            num_heads = self.n_local_heads
-            padded_num_heads = 64 if num_heads <= 64 else self.n_heads
-            sink = self.attn_sink.new_zeros(padded_num_heads)
+            sink = self.attn_sink.new_zeros(sink_num_heads)
             sink[:num_heads] = self.attn_sink[rank * num_heads : (rank + 1) * num_heads]
             self._attn_sink_local = sink
-        return self._attn_sink_local
+        return self._attn_sink_local[:kernel_num_heads]
 
     @contextmanager
     def maybe_use_decode_attn_tp(self, forward_batch: ForwardBatch):
@@ -1624,30 +1658,21 @@ class MQALayer(MqaAttentionBase):
         )
 
         tp_slice, q_padded, q_out = slice(None), None, None
-        # Above this the SM120 route is the prefill kernel, which takes
-        # arbitrary h_q, so the decode pad below would just be sliced back off.
-        skip_decode_pad = is_sm120_supported() and x.shape[0] > SM120_DECODE_MAX_TOKENS
-        if self.attn_tp_size > 1:
-            # FlashMLA's fp8 sparse decode kernel only specializes h_q for {64, 128}.
-            # Pad the per-rank heads to 64 (not the full n_heads) when they fit, to
-            # dispatch the cheaper decode::head64 variant; attn_sink is sliced to
-            # this rank and padded to match.
-            padded_num_heads = (
-                self.n_local_heads
-                if skip_decode_pad
-                else (64 if self.n_local_heads <= 64 else self.n_heads)
-            )
+        kernel_num_heads = self._kernel_num_heads(x.shape[0])
+        if kernel_num_heads != self.n_local_heads:
+            # Backends without an exact-head specialization retain the existing
+            # padded shape. attn_sink is sliced to this rank and padded to match.
             # Only [0:n_local_heads] is written below. Uninitialized padded TP
             # heads inject NaN into attention on gfx942 (fnuz), so zero-init
             # there; other archs tolerate new_empty and skip the per-forward
             # memset.
             if _is_gfx942_supported:
-                q_padded = x.new_zeros(x.shape[0], padded_num_heads, self.head_dim)
+                q_padded = x.new_zeros(x.shape[0], kernel_num_heads, self.head_dim)
             else:
-                q_padded = x.new_empty(x.shape[0], padded_num_heads, self.head_dim)
+                q_padded = x.new_empty(x.shape[0], kernel_num_heads, self.head_dim)
             tp_slice = slice(0, self.n_local_heads)
             q_out = q_padded[:, tp_slice, :]
-        attn_sink = self._local_attn_sink()
+        attn_sink = self._local_attn_sink(kernel_num_heads)
 
         if enable_multi_stream:
             # Multi-stream path always fuses cache write into the K kernel,