[AMD][Spec] Accelerate Qwen3.5 verification with grouped-head shared KV (#34517)

Co-authored-by: chuyeh <chuyeh@users.noreply.github.com>
Co-authored-by: Thomas Wang <1am9trash@gmail.com>

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,6 +131,7 @@ def _verify_mla_prefix_stage1(
             mask=row_mask[:, None] & (offs_dn[None, :] < NOPE_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),
@@ -150,17 +155,19 @@ def _verify_mla_prefix_stage1(
                 mask=(offs_dn[:, None] < NOPE_DIM) & n_mask[None, :],
                 other=0.0,
             )
+            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_nope, k_nope) + tl.dot(q_pe, k_pe)
+                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
+            # MLA exposes its latent V through V_Buffer; ordinary shared-KV
-            # V instead of reusing k_nope.
+            # 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
+    """Return True iff the grouped-head split-KV verify path can serve this
-    with the same result as extend_attention_fwd. Conservative: anything not
+    exact problem with the same result as extend_attention_fwd. Conservative:
-    explicitly handled -> False -> caller falls back to the baseline.
+    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
+    Grouped-head drop-in for extend_attention_fwd on a topk==1 target-verify
-    (topk==1) shape. Returns True if it ran (o_extend written), False if unsupported
+    shape. Returns True if it ran (o_extend written), False if unsupported
-    (caller falls back). Requires h_kv == 1 (MLA single latent).
+    (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

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",

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).
+        # Grouped-head split-KV verify kernel for MLA or one shared local KV head.
-        self.verify_mla_fwd = torch.compiler.disable(verify_mla_fwd)
+        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
+        # The grouped-head verify kernel is tuned for Kimi-K3 MLA and Qwen3.5
-        # Kimi-K3 absorbed-MLA shape; gate it on K3.
+        # GQA with exactly one TP-local KV head.
-        self.use_verify_mla = (
+        self.use_verify_shared_kv = _should_use_verify_shared_kv(
-            is_gfx95_supported()
+            model_runner.model_config,
-            and self.topk == 1
+            self.topk,
-            and self.use_mla
+            self.use_mla,
-            and is_kimi_k3(model_runner.model_config.hf_config)
+            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:
+        if self.use_verify_shared_kv:
-            verify_fwd = self.verify_mla_fwd
+            verify_fwd = self.verify_shared_kv_fwd
         elif self.use_verify_splitkv:
             verify_fwd = self.verify_splitkv_fwd
         else:

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()