diff --git a/python/sglang/kernels/ops/attention/verify_mla.py b/python/sglang/kernels/ops/attention/verify_mla.py index 6c5b1e1a0..443829885 100644 --- a/python/sglang/kernels/ops/attention/verify_mla.py +++ b/python/sglang/kernels/ops/attention/verify_mla.py @@ -1,9 +1,10 @@ """ -MLA split-KV attention for EAGLE/DSpark speculative *verify* (topk==1). +Grouped-head split-KV attention for speculative *verify* (topk==1). Following the pattern of ``python/sglang/kernels/ops/attention/verify_splitkv.py``. Grid is ``(bs, n_head_blocks, split)``; each program handles ``BLOCK_H`` query -heads x ALL ``L_EXT`` draft queries. +heads x ALL ``L_EXT`` draft queries. It supports absorbed MLA and ordinary +MHA/GQA when exactly one TP-local KV head is shared by all local query heads. Correctness matches ``extend_attention_fwd`` for the topk==1 causal verify case. @@ -27,6 +28,7 @@ DEFAULT_BLOCK_N = 64 DEFAULT_NUM_WARPS = 8 _BLOCK_CONFIG = { # head_dim: (BLOCK_H, BLOCK_N, num_warps) + 256: (4, 64, 8), # Qwen3.5 TP2 / TP4 / TP8 576: (4, 64, 8), # K3 MLA (kv_lora_rank 512 + qk_rope 64) } @@ -117,7 +119,9 @@ def _verify_mla_prefix_stage1( split_kv_id == active - 1, cur_batch_seq_len, split_start + kv_len_per_split ) - # load q_nope and q_pe + # For absorbed MLA, NOPE_DIM is the latent width and PE_DIM is the + # appended RoPE width. For ordinary shared-KV attention (Qwen3.5), + # NOPE_DIM is the complete already-rotated Q/K head and PE_DIM is zero. q_row = tl.reshape( (cur_q_start + offs_l)[None, :] * stride_qbs + offs_h[:, None] * stride_qh, (R,), @@ -127,11 +131,12 @@ def _verify_mla_prefix_stage1( mask=row_mask[:, None] & (offs_dn[None, :] < NOPE_DIM), other=0.0, ).to(K_Buffer.dtype.element_ty) - q_pe = tl.load( - Q + q_row[:, None] + (NOPE_DIM + offs_dp)[None, :], - mask=row_mask[:, None] & (offs_dp[None, :] < PE_DIM), - other=0.0, - ).to(K_Buffer.dtype.element_ty) + if PE_DIM > 0: + q_pe = tl.load( + Q + q_row[:, None] + (NOPE_DIM + offs_dp)[None, :], + mask=row_mask[:, None] & (offs_dp[None, :] < PE_DIM), + other=0.0, + ).to(K_Buffer.dtype.element_ty) e_max = tl.zeros([R], dtype=tl.float32) - float("inf") e_sum = tl.zeros([R], dtype=tl.float32) @@ -150,17 +155,19 @@ def _verify_mla_prefix_stage1( mask=(offs_dn[:, None] < NOPE_DIM) & n_mask[None, :], other=0.0, ) - k_pe = tl.load( - K_Buffer + base + (NOPE_DIM + offs_dp)[:, None], - mask=(offs_dp[:, None] < PE_DIM) & n_mask[None, :], - other=0.0, - ) - qk = tl.dot(q_nope, k_nope) + tl.dot(q_pe, k_pe) + qk = tl.dot(q_nope, k_nope) + if PE_DIM > 0: + k_pe = tl.load( + K_Buffer + base + (NOPE_DIM + offs_dp)[:, None], + mask=(offs_dp[:, None] < PE_DIM) & n_mask[None, :], + other=0.0, + ) + qk += tl.dot(q_pe, k_pe) qk *= sm_scale * k_scale qk = tl.where(n_mask[None, :], qk, float("-inf")) - # V is the same as k_nope but transposed; tl.trans is slow, so re-load - # V instead of reusing k_nope. + # MLA exposes its latent V through V_Buffer; ordinary shared-KV + # attention has an independent V cache. Both use this same load. v = tl.load( V_Buffer + kv_loc[:, None] * stride_buf_vbs + offs_dv[None, :], mask=n_mask[:, None] & (offs_dv[None, :] < V_HEAD_DIM), @@ -348,7 +355,9 @@ class VerifyMLA: self.nope_dim = v_head_dim self.pe_dim = head_dim - v_head_dim self.l_ext = l_ext - self.l_pad = triton.next_power_of_2(l_ext) + # tl.dot requires BLOCK_H * L_EXT to cover at least 16 rows. + min_l_pad = triton.next_power_of_2(triton.cdiv(16, block_h)) + self.l_pad = max(min_l_pad, triton.next_power_of_2(l_ext)) self.device = device self.block_h = block_h self.block_n = block_n @@ -424,7 +433,7 @@ class VerifyMLA: PE_DIM=self.pe_dim, V_HEAD_DIM=self.v_head_dim, BLOCK_DNOPE=triton.next_power_of_2(self.nope_dim), - BLOCK_DPE=triton.next_power_of_2(self.pe_dim), + BLOCK_DPE=max(1, triton.next_power_of_2(self.pe_dim)), BLOCK_DV=triton.next_power_of_2(self.v_head_dim), BLOCK_N=self.block_n, num_warps=self.num_warps, @@ -573,15 +582,16 @@ def can_handle( logit_cap=0.0, xai_temperature_len=-1, ): - """Return True iff the MLA split-KV verify path can serve this exact problem - with the same result as extend_attention_fwd. Conservative: anything not - explicitly handled -> False -> caller falls back to the baseline. + """Return True iff the grouped-head split-KV verify path can serve this + exact problem with the same result as extend_attention_fwd. Conservative: + anything not explicitly handled -> False -> caller falls back to the + baseline. IMPORTANT: ``custom_mask`` is intentionally NOT inspected (its values can't be read inside a captured HIP graph without a host sync). The kernel always computes pure-causal attention, which equals the tree mask ONLY at speculative topk == 1. The caller therefore MUST gate enablement on topk == 1 - (TritonAttnBackend does: ``use_verify_mla = ... and self.topk == 1``). + (TritonAttnBackend does: ``use_verify_shared_kv = ... and self.topk == 1``). At topk > 1 the tree is not causal and this path must stay disabled.""" # No exotic features. if sinks is not None: @@ -638,7 +648,7 @@ def can_handle( return True -def verify_mla_fwd( +def verify_shared_kv_fwd( q_extend, k_extend, v_extend, @@ -664,9 +674,9 @@ def verify_mla_fwd( max_bs=None, ): """ - MLA-native drop-in for extend_attention_fwd on the EAGLE target-verify - (topk==1) shape. Returns True if it ran (o_extend written), False if unsupported - (caller falls back). Requires h_kv == 1 (MLA single latent). + Grouped-head drop-in for extend_attention_fwd on a topk==1 target-verify + shape. Returns True if it ran (o_extend written), False if unsupported + (caller falls back). Requires exactly one TP-local KV head. """ if not can_handle( q_extend, @@ -687,7 +697,11 @@ def verify_mla_fwd( xai_temperature_len=xai_temperature_len, ): return False - if k_extend.shape[1] != 1: # MLA: single shared latent head + if k_extend.shape[1] != 1: + return False + if q_extend.shape[2] < v_extend.shape[2]: + return False + if kv_indices.numel() == 0: return False bs = qo_indptr.shape[0] - 1 diff --git a/python/sglang/srt/configs/model_config.py b/python/sglang/srt/configs/model_config.py index 0b24cd87c..9811e663c 100644 --- a/python/sglang/srt/configs/model_config.py +++ b/python/sglang/srt/configs/model_config.py @@ -125,6 +125,15 @@ def is_kimi_k3(config) -> bool: return _hf_arch(config) == "KimiK3ForConditionalGeneration" +def is_qwen3_5(config) -> bool: + return _hf_arch(config) in ( + "Qwen3_5ForConditionalGeneration", + "Qwen3_5MoeForConditionalGeneration", + "Qwen3_5ForCausalLM", + "Qwen3_5MoeForCausalLM", + ) + + def is_deepseek_v4(config) -> bool: return _hf_arch(config) in ( "DeepseekV4ForCausalLM", diff --git a/python/sglang/srt/layers/attention/triton_backend.py b/python/sglang/srt/layers/attention/triton_backend.py index c6b9cb92b..ab8add40b 100644 --- a/python/sglang/srt/layers/attention/triton_backend.py +++ b/python/sglang/srt/layers/attention/triton_backend.py @@ -11,7 +11,7 @@ from sglang.kernels.ops.kvcache.kv_indices import ( create_flashinfer_kv_indices_triton, ) from sglang.srt.configs.hybrid_arch import mambaish_config -from sglang.srt.configs.model_config import AttentionArch, is_kimi_k3 +from sglang.srt.configs.model_config import AttentionArch, is_kimi_k3, is_qwen3_5 from sglang.srt.distributed.device_communicators.pynccl_allocator import ( use_symmetric_memory, ) @@ -76,6 +76,21 @@ def _mla_decode_kv_splits_cap( return max(base_max_kv_splits, min(sm_cap, ctx_cap)) +def _should_use_verify_shared_kv(model_config, topk, use_mla, use_verify_splitkv): + if not is_gfx95_supported() or topk != 1: + return False + if use_mla: + return is_kimi_k3(model_config.hf_config) + return ( + use_verify_splitkv + and is_qwen3_5(model_config.hf_config) + and model_config.get_num_kv_heads( + get_parallel().attn_tp_size, get_parallel().attn_dcp_size + ) + == 1 + ) + + def logit_capping_mod(logit_capping_method, logit_cap): # positive logit_cap -> tanh cap if logit_capping_method == "tanh": @@ -134,7 +149,7 @@ class TritonAttnBackend(AttentionBackend): extend_attention_fwd_unified, ) from sglang.kernels.ops.attention.verify_mla import ( - verify_mla_fwd, + verify_shared_kv_fwd, ) from sglang.kernels.ops.attention.verify_splitkv import ( verify_splitkv_fwd, @@ -150,8 +165,8 @@ class TritonAttnBackend(AttentionBackend): self.build_unified_kv_indices = torch.compiler.disable(build_unified_kv_indices) # Split-KV EAGLE-verify kernel; enabled below once topk is known (valid only at topk == 1). self.verify_splitkv_fwd = torch.compiler.disable(verify_splitkv_fwd) - # MLA split-KV EAGLE-verify kernel; enabled below once topk is known (valid only at topk == 1). - self.verify_mla_fwd = torch.compiler.disable(verify_mla_fwd) + # Grouped-head split-KV verify kernel for MLA or one shared local KV head. + self.verify_shared_kv_fwd = torch.compiler.disable(verify_shared_kv_fwd) # Parse args self.skip_prefill = skip_prefill @@ -184,13 +199,13 @@ class TritonAttnBackend(AttentionBackend): and self.topk == 1 ) self.use_mla = model_runner.model_config.attention_arch == AttentionArch.MLA - # The MLA verify kernel (verify_mla_fwd) is tuned and validated for the - # Kimi-K3 absorbed-MLA shape; gate it on K3. - self.use_verify_mla = ( - is_gfx95_supported() - and self.topk == 1 - and self.use_mla - and is_kimi_k3(model_runner.model_config.hf_config) + # The grouped-head verify kernel is tuned for Kimi-K3 MLA and Qwen3.5 + # GQA with exactly one TP-local KV head. + self.use_verify_shared_kv = _should_use_verify_shared_kv( + model_runner.model_config, + self.topk, + self.use_mla, + self.use_verify_splitkv, ) self.dcp_size = get_parallel().attn_dcp_size self.dcp_rank = get_parallel().attn_dcp_rank @@ -1405,10 +1420,10 @@ class TritonAttnBackend(AttentionBackend): # serve bit-equivalently (its can_handle() gates on non-causal / sinks / # sliding-window / ragged / topk>1), so we fall through to # extend_attention_fwd below. Correctness is never at risk. - # Route target-verify to the K3-tuned MLA kernel when eligible, else the + # Route target-verify to the grouped-head kernel when eligible, else the # per-head split-KV kernel. - if self.use_verify_mla: - verify_fwd = self.verify_mla_fwd + if self.use_verify_shared_kv: + verify_fwd = self.verify_shared_kv_fwd elif self.use_verify_splitkv: verify_fwd = self.verify_splitkv_fwd else: diff --git a/test/registered/attention/test_verify_shared_kv.py b/test/registered/attention/test_verify_shared_kv.py new file mode 100644 index 000000000..330d4e37d --- /dev/null +++ b/test/registered/attention/test_verify_shared_kv.py @@ -0,0 +1,235 @@ +"""Correctness tests for grouped-head target-verify attention.""" + +import unittest +from types import SimpleNamespace +from unittest.mock import patch + +import torch + +from sglang.kernels.ops.attention.extend_attention import extend_attention_fwd +from sglang.kernels.ops.attention.verify_mla import verify_shared_kv_fwd +from sglang.srt.layers.attention.triton_backend import ( + _should_use_verify_shared_kv, +) +from sglang.test.ci.ci_register import register_amd_ci +from sglang.test.test_utils import CustomTestCase + +register_amd_ci(est_time=30, suite="stage-b-test-1-gpu-small-amd-mi35x") + +BF16_ATOL = 2e-2 +BF16_RTOL = 1e-2 +FP8_ATOL = 8e-2 +FP8_RTOL = 2e-2 + + +def _build_inputs( + prefix_lens, + l_ext, + h_q, + head_dim, + v_head_dim, + cache_dtype=torch.bfloat16, +): + device = "cuda" + dtype = torch.bfloat16 + generator = torch.Generator(device=device).manual_seed(0) + prefix_lens_t = torch.tensor(prefix_lens, dtype=torch.int32, device=device) + total_prefix = sum(prefix_lens) + batch_size = len(prefix_lens) + num_extend_tokens = batch_size * l_ext + + def randn(*shape): + return torch.randn(*shape, dtype=dtype, device=device, generator=generator) + + q = randn(num_extend_tokens, h_q, head_dim) + k = randn(num_extend_tokens, 1, head_dim) + v = randn(num_extend_tokens, 1, v_head_dim) + k_buffer = randn(total_prefix, 1, head_dim).to(cache_dtype) + v_buffer = randn(total_prefix, 1, v_head_dim).to(cache_dtype) + qo_indptr = torch.arange( + 0, num_extend_tokens + 1, l_ext, dtype=torch.int32, device=device + ) + kv_indptr = torch.zeros(batch_size + 1, dtype=torch.int32, device=device) + kv_indptr[1:] = torch.cumsum(prefix_lens_t, dim=0) + kv_indices = torch.arange(total_prefix, dtype=torch.int64, device=device) + return q, k, v, k_buffer, v_buffer, qo_indptr, kv_indptr, kv_indices + + +@unittest.skipIf(not torch.cuda.is_available(), "GPU required") +class TestVerifySharedKV(CustomTestCase): + def _run_parity( + self, + head_dim, + v_head_dim, + h_q=4, + cache_dtype=torch.bfloat16, + k_scale=1.0, + v_scale=1.0, + l_ext=4, + atol=BF16_ATOL, + rtol=BF16_RTOL, + ): + inputs = _build_inputs( + prefix_lens=[512, 2048], + l_ext=l_ext, + h_q=h_q, + head_dim=head_dim, + v_head_dim=v_head_dim, + cache_dtype=cache_dtype, + ) + q, k, v, k_buffer, v_buffer, qo_indptr, kv_indptr, kv_indices = inputs + output_shape = (q.shape[0], q.shape[1], v_head_dim) + reference = torch.empty(output_shape, dtype=q.dtype, device=q.device) + actual = torch.empty_like(reference) + scale = head_dim**-0.5 + + extend_attention_fwd( + q, + k, + v, + reference, + k_buffer, + v_buffer, + qo_indptr, + kv_indptr, + kv_indices, + None, + True, + None, + l_ext, + k_scale, + v_scale, + sm_scale=scale, + ) + ran = verify_shared_kv_fwd( + q, + k, + v, + actual, + k_buffer, + v_buffer, + qo_indptr, + kv_indptr, + kv_indices, + None, + True, + None, + l_ext, + k_scale, + v_scale, + sm_scale=scale, + max_bs=len(kv_indptr) - 1, + ) + + self.assertTrue(ran) + torch.testing.assert_close(actual, reference, atol=atol, rtol=rtol) + + def test_qwen3_5_tp_shapes(self): + # Qwen3.5 has 32 global query heads. TP8, TP4, and InferenceX's TP2 + # configurations expose 4, 8, and 16 local query heads respectively, + # all sharing one TP-local KV head. + for h_q in (4, 8, 16): + with self.subTest(h_q=h_q): + self._run_parity(head_dim=256, v_head_dim=256, h_q=h_q) + + def test_qwen3_5_short_verify_widths(self): + for l_ext in (1, 2, 3): + with self.subTest(l_ext=l_ext): + self._run_parity(head_dim=256, v_head_dim=256, l_ext=l_ext) + + def test_qwen3_5_fp8_kv_cache(self): + self._run_parity( + head_dim=256, + v_head_dim=256, + h_q=8, + cache_dtype=torch.float8_e4m3fn, + k_scale=0.5, + v_scale=0.25, + atol=FP8_ATOL, + rtol=FP8_RTOL, + ) + + def test_kimi_k3_absorbed_mla_shape(self): + self._run_parity(head_dim=576, v_head_dim=512) + + def test_rejects_multiple_local_kv_heads(self): + inputs = list( + _build_inputs( + prefix_lens=[512], + l_ext=4, + h_q=4, + head_dim=256, + v_head_dim=256, + ) + ) + for index in (1, 2, 3, 4): + inputs[index] = inputs[index].expand(-1, 2, -1).contiguous() + q, k, v, k_buffer, v_buffer, qo_indptr, kv_indptr, kv_indices = inputs + output = torch.empty_like(q) + self.assertFalse( + verify_shared_kv_fwd( + q, + k, + v, + output, + k_buffer, + v_buffer, + qo_indptr, + kv_indptr, + kv_indices, + None, + True, + None, + 4, + 1.0, + 1.0, + ) + ) + + @patch( + "sglang.srt.layers.attention.triton_backend.is_gfx95_supported", + return_value=True, + ) + @patch("sglang.srt.layers.attention.triton_backend.get_parallel") + def test_backend_dispatch_gate(self, get_parallel_mock, _is_gfx95_mock): + get_parallel_mock.return_value = SimpleNamespace( + attn_tp_size=8, attn_dcp_size=1 + ) + + def model_config(architecture, local_kv_heads=1): + return SimpleNamespace( + hf_config=SimpleNamespace(architectures=[architecture]), + get_num_kv_heads=lambda _tp, _dcp: local_kv_heads, + ) + + qwen = model_config("Qwen3_5MoeForCausalLM") + self.assertTrue(_should_use_verify_shared_kv(qwen, 1, False, True)) + self.assertFalse(_should_use_verify_shared_kv(qwen, 2, False, True)) + self.assertFalse(_should_use_verify_shared_kv(qwen, 1, False, False)) + self.assertFalse( + _should_use_verify_shared_kv( + model_config("Qwen3_5MoeForCausalLM", local_kv_heads=2), + 1, + False, + True, + ) + ) + self.assertFalse( + _should_use_verify_shared_kv( + model_config("LlamaForCausalLM"), 1, False, True + ) + ) + self.assertTrue( + _should_use_verify_shared_kv( + model_config("KimiK3ForConditionalGeneration"), 1, True, False + ) + ) + with patch( + "sglang.srt.layers.attention.triton_backend.is_gfx95_supported", + return_value=False, + ): + self.assertFalse(_should_use_verify_shared_kv(qwen, 1, False, True)) + + +if __name__ == "__main__": + unittest.main()