diff --git a/python/sglang/kernels/ops/attention/flash_mla_sm120.py b/python/sglang/kernels/ops/attention/flash_mla_sm120.py index 906d4aa22..5267908ee 100644 --- a/python/sglang/kernels/ops/attention/flash_mla_sm120.py +++ b/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. diff --git a/python/sglang/srt/models/deepseek_v4.py b/python/sglang/srt/models/deepseek_v4.py index e05ab612e..9af352e22 100644 --- a/python/sglang/srt/models/deepseek_v4.py +++ b/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, diff --git a/test/registered/kernels/ops/attention/test_flash_mla_backends.py b/test/registered/kernels/ops/attention/test_flash_mla_backends.py index a41568432..d105bdfc7 100644 --- a/test/registered/kernels/ops/attention/test_flash_mla_backends.py +++ b/test/registered/kernels/ops/attention/test_flash_mla_backends.py @@ -43,6 +43,7 @@ from sglang.kernels.ops.attention.flash_mla_sm120 import ( _sm120_sparse_decode_fwd, _split_kv_pages_to_64, flash_mla_with_kvcache_sm120, + flashinfer_dsv4_decode_supports_num_heads, ) from sglang.kernels.ops.attention.flash_mla_sm120_triton import ( _apply_attn_sink, @@ -502,6 +503,87 @@ class TestEntryPointDispatch(CustomTestCase): rtol=5e-2, ) + def test_flashinfer_exact_heads_match_padded_64_heads(self): + """Native TP4/TP8 heads agree with padding across the prefill boundary.""" + num_pages, page_size, topk = 2, 64, 128 + k_cache, _ = _build_kvcache(num_pages, page_size, device=self.device, seed=17) + extra_cache, _ = _build_kvcache( + num_pages, page_size, device=self.device, seed=23 + ) + for num_heads in (8, 16): + with self.subTest(heads=num_heads): + if not flashinfer_dsv4_decode_supports_num_heads(num_heads, 1): + self.skipTest( + f"FlashInfer has no {num_heads}-head DSV4 decode specialization" + ) + self.assertTrue( + flashinfer_dsv4_decode_supports_num_heads(num_heads, 64) + ) + self.assertFalse( + flashinfer_dsv4_decode_supports_num_heads(num_heads, 65) + ) + for num_tokens in (1, 64, 65): + for dual_cache in (False, True): + with self.subTest(tokens=num_tokens, dual_cache=dual_cache): + q, indices = _build_q_indices( + num_tokens, + num_heads, + topk, + num_pages, + page_size, + device=self.device, + seed=29, + ) + topk_length = torch.full( + (num_tokens,), + topk, + dtype=torch.int32, + device=self.device, + ) + sink = torch.linspace( + -1.0, + 1.0, + num_heads, + dtype=torch.float32, + device=self.device, + ) + q_padded = q.new_zeros(num_tokens, 1, 64, _D) + q_padded[:, :, :num_heads].copy_(q) + sink_padded = sink.new_zeros(64) + sink_padded[:num_heads].copy_(sink) + + common = dict( + k_cache=k_cache, + indices=indices, + topk_length=topk_length, + head_dim_v=_D, + softmax_scale=_D**-0.5, + extra_k_cache=extra_cache if dual_cache else None, + extra_indices_in_kvcache=indices + if dual_cache + else None, + extra_topk_length=topk_length if dual_cache else None, + ) + with mock.patch.object( + fmod, "_sm120_default_backend", "flashinfer" + ): + out_exact, _ = flash_mla_with_kvcache_sm120( + q=q, attn_sink=sink, **common + ) + out_padded, _ = flash_mla_with_kvcache_sm120( + q=q_padded, attn_sink=sink_padded, **common + ) + + self.assertEqual( + out_exact.shape, (num_tokens, 1, num_heads, _D) + ) + torch.testing.assert_close( + out_exact.float(), + out_padded[:, :, :num_heads].float(), + atol=5e-2, + rtol=5e-2, + ) + @unittest.skipUnless(_IS_SM120, "SM120 (compute capability 12.0) required") class TestTouchedPageSplit(CustomTestCase):