minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 15 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: