minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

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

Browse Source 
Co-authored-by: chuyeh <chuyeh@users.noreply.github.com>
Co-authored-by: Thomas Wang <1am9trash@gmail.com>
This commit is contained in:
chuyeh
2026-08-15 00:15:52 -07:00
committed by GitHub
co-authored by chuyeh Thomas Wang
parent 5c9ee86d90
commit 6cbfa791d6
4 changed files with 314 additions and 41 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/kernels/ops/attention/verify_mla.py
+41 -27
View File
@@ -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
python/sglang/srt/configs/model_config.py
+9
View File
@@ -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
+29 -14
View File
@@ -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:
test/registered/attention/test_verify_shared_kv.py
+235
View File
@@ -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()