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

[AMD][Perf] Split-KV flash-decode attention for EAGLE target-verify (Triton backend) (#27382)

Browse Source
This commit is contained in:
giang_ng_tr
2026-06-18 19:11:10 -07:00
committed by GitHub
parent 1c8551169d
commit b36360dc5b
5 changed files with 1259 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
benchmark/kernels/verify_splitkv_triton/bench_verify_splitkv.py
+152
View File
@@ -0,0 +1,152 @@
"""Micro-benchmark: split-KV EAGLE-verify kernel vs extend_attention_fwd.
Times ``verify_splitkv_fwd`` against the baseline ``extend_attention_fwd`` on
the verify shape (a few draft-token queries over a long prefix KV) across
context lengths and head dims, and reports the per-kernel latency, the speedup,
and the achieved KV-read bandwidth. Model-independent (head_dim is just a shape
parameter).
NOTE: this benchmark targets AMD MI35x (gfx950). The verify kernel's block config
and its CDNA-only Triton launch hints (waves_per_eu, matrix_instr_nonkdim) are
tuned and validated only on gfx950, and the kernel is gated to gfx95 in production
-- so these numbers are meaningful only on MI35x. GPU + Triton required.
python3 benchmark/kernels/verify_splitkv_triton/bench_verify_splitkv.py
"""
import argparse
import torch
import triton
from sglang.srt.layers.attention.triton_ops.extend_attention import (
extend_attention_fwd,
)
from sglang.srt.layers.attention.triton_ops.verify_splitkv import verify_splitkv_fwd
from sglang.srt.utils import is_gfx95_supported
def build_inputs(prefix_len, l_ext, h_q, h_kv, head_dim, v_head_dim, dtype, device):
"""One verify-shaped sequence repeated to batch size 1 per call here; the
kernels are timed at bs=1 to isolate the per-(seq,head) bandwidth story."""
total_prefix = prefix_len
k_buffer = torch.randn(total_prefix, h_kv, head_dim, dtype=dtype, device=device)
v_buffer = torch.randn(total_prefix, h_kv, v_head_dim, dtype=dtype, device=device)
kv_indptr = torch.tensor([0, total_prefix], dtype=torch.int32, device=device)
# kv_indices is int64 in production (TritonAttnBackend allocates int64).
kv_indices = torch.arange(total_prefix, dtype=torch.int64, device=device)
q = torch.randn(l_ext, h_q, head_dim, dtype=dtype, device=device)
k = torch.randn(l_ext, h_kv, head_dim, dtype=dtype, device=device)
v = torch.randn(l_ext, h_kv, v_head_dim, dtype=dtype, device=device)
qo_indptr = torch.tensor([0, l_ext], dtype=torch.int32, device=device)
return q, k, v, k_buffer, v_buffer, qo_indptr, kv_indptr, kv_indices, l_ext
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--l-ext", type=int, default=4, help="draft tokens per seq")
ap.add_argument("--h-q", type=int, default=16)
ap.add_argument("--h-kv", type=int, default=2)
ap.add_argument("--head-dim", type=int, default=256)
args = ap.parse_args()
if not is_gfx95_supported():
raise SystemExit(
"This benchmark is for AMD MI35x (gfx950) only: the verify kernel's "
"block config and CDNA launch hints are tuned/validated there, and the "
"kernel is gated to gfx95 in production, so results on other hardware "
"are not representative."
)
if not torch.cuda.is_available():
raise SystemExit("GPU required")
device, dtype = "cuda", torch.bfloat16
hd, vhd = args.head_dim, args.head_dim
sm_scale = 1.0 / (hd**0.5)
kv_bytes_per_tok = 2 * args.h_kv * hd * torch.tensor([], dtype=dtype).element_size()
print(
f"verify-split-KV vs extend_attention_fwd "
f"(l_ext={args.l_ext}, H_Q={args.h_q}, H_KV={args.h_kv}, head_dim={hd}, bf16)\n"
)
print(
f"{'ctx':>8} {'extend(ms)':>12} {'splitkv(ms)':>12} {'speedup':>9} {'splitkv GB/s':>13}"
)
for ctx in (1024, 2048, 4096, 8192, 16384):
q, k, v, kb, vb, qo, kvp, kvi, mle = build_inputs(
ctx, args.l_ext, args.h_q, args.h_kv, hd, vhd, dtype, device
)
o = torch.empty(q.shape[0], args.h_q, vhd, dtype=dtype, device=device)
def run_extend():
extend_attention_fwd(
q,
k,
v,
o,
kb,
vb,
qo,
kvp,
kvi,
None,
True,
None,
mle,
1.0,
1.0,
sm_scale=sm_scale,
)
def run_split():
verify_splitkv_fwd(
q,
k,
v,
o,
kb,
vb,
qo,
kvp,
kvi,
None,
True,
None,
mle,
1.0,
1.0,
sm_scale=sm_scale,
)
# Ensure the split-KV path actually handled this shape before timing it
# (verify_splitkv_fwd returns False + no-ops on unsupported cases).
assert verify_splitkv_fwd(
q,
k,
v,
o,
kb,
vb,
qo,
kvp,
kvi,
None,
True,
None,
mle,
1.0,
1.0,
sm_scale=sm_scale,
), "verify_splitkv_fwd did not handle the verify shape"
t_ext = triton.testing.do_bench(run_extend)
t_spl = triton.testing.do_bench(run_split)
kv_bytes = int(kv_bytes_per_tok) * ctx
gbs = kv_bytes / (t_spl * 1e-3) / 1e9
print(
f"{ctx:>8} {t_ext:>12.3f} {t_spl:>12.3f} {t_ext / t_spl:>8.2f}x {gbs:>12.0f}"
)
if __name__ == "__main__":
main()
python/sglang/srt/environ.py
+5
View File
@@ -644,6 +644,11 @@ class Envs:
# Saves the per-step draft forward, but the draft KV goes stale: an upshift
# back to steps>0 starts from a cold draft state (low accept until it recovers).
SGLANG_SPEC_SKIP_ZERO_STEP_DRAFT_EXTEND = EnvBool(False)
# Use the split-KV (flash-decode) kernel for EAGLE target-verify on the
# Triton backend (ROCm). Only active at speculative topk == 1; falls back to
# extend_attention_fwd for unsupported cases or when set false (e.g. for
# debugging). Correctness is unaffected; this only changes performance.
SGLANG_ENABLE_SPLITKV_VERIFY = EnvBool(True)
# Master switch for all async-asserted invariant probes (NaN, Inf, OOB,
# page alignment). Off in prod; tests turn it on to fail-fast on
# numerical / index violations instead of getting silent NaN cascades.
python/sglang/srt/layers/attention/triton_backend.py
+60
View File
@@ -7,6 +7,7 @@ import torch
import triton
from sglang.srt.configs.model_config import AttentionArch
from sglang.srt.environ import envs
from sglang.srt.layers.attention.base_attn_backend import AttentionBackend
from sglang.srt.layers.attention.triton_ops.kv_indices import (
create_flashinfer_kv_indices_triton,
@@ -28,6 +29,7 @@ from sglang.srt.utils import (
get_device_core_count,
get_int_env_var,
is_cuda,
is_gfx95_supported,
is_gfx942_supported,
next_power_of_2,
)
@@ -107,6 +109,9 @@ class TritonAttnBackend(AttentionBackend):
extend_attention_fwd,
extend_attention_fwd_unified,
)
from sglang.srt.layers.attention.triton_ops.verify_splitkv import (
verify_splitkv_fwd,
)
super().__init__()
@@ -116,6 +121,9 @@ class TritonAttnBackend(AttentionBackend):
extend_attention_fwd_unified
)
self.build_unified_kv_indices = torch.compiler.disable(build_unified_kv_indices)
# Split-KV EAGLE-verify kernel (ROCm/Triton). Registered here; enabled
# below once topk is known (the path is only valid at topk == 1).
self.verify_splitkv_fwd = torch.compiler.disable(verify_splitkv_fwd)
# Parse args
self.skip_prefill = skip_prefill
@@ -130,6 +138,21 @@ class TritonAttnBackend(AttentionBackend):
self.use_sliding_window_kv_pool = isinstance(self.token_to_kv_pool, SWAKVPool)
self.num_draft_tokens = model_runner.server_args.speculative_num_draft_tokens
self.speculative_num_steps = model_runner.server_args.speculative_num_steps
self.topk = model_runner.server_args.speculative_eagle_topk or 0
# Split-KV verify matches extend_attention_fwd only when the EAGLE tree
# reduces to a pure-causal chain, i.e. topk == 1 (the same condition the
# aiter backend's unified-verify uses). For topk > 1 the tree custom_mask
# is not causal, so leave the path off and fall back to the baseline.
# gfx95-only (MI350X/CDNA4): the kernel uses ROCm/CDNA Triton launch hints
# (waves_per_eu, matrix_instr_nonkdim) and its block config is tuned and
# validated only on gfx950. NVIDIA's Triton rejects those kwargs, and the
# path is unvalidated on NV and on other AMD archs, so restrict it to gfx95
# and fall back to extend_attention_fwd everywhere else.
self.use_verify_splitkv = (
is_gfx95_supported()
and envs.SGLANG_ENABLE_SPLITKV_VERIFY.get()
and self.topk == 1
)
self.use_mla = model_runner.model_config.attention_arch == AttentionArch.MLA
self.num_head = (
model_runner.model_config.num_attention_heads // get_parallel().attn_tp_size
@@ -1094,6 +1117,43 @@ class TritonAttnBackend(AttentionBackend):
k_descale = 1.0
v_descale = 1.0
# Split-KV EAGLE-verify fast path (ROCm/Triton). On target-verify
# (topk=1 causal chain), run the bandwidth-efficient split-KV kernel
# instead of the serial-prefix extend kernel. verify_splitkv_fwd()
# returns True if it ran (o written), or False for any case it cannot
# 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.
if (
self.use_verify_splitkv
and forward_batch.forward_mode.is_target_verify()
and self.verify_splitkv_fwd(
q.view(-1, layer.tp_q_head_num, layer.qk_head_dim),
k.contiguous(),
v.contiguous(),
o.view(-1, layer.tp_q_head_num, layer.v_head_dim),
self.token_to_kv_pool.get_key_buffer(layer.layer_id),
self.token_to_kv_pool.get_value_buffer(layer.layer_id),
self.forward_metadata.qo_indptr,
kv_indptr,
kv_indices,
self.forward_metadata.custom_mask,
causal,
self.forward_metadata.mask_indptr,
self.forward_metadata.max_extend_len,
k_descale,
v_descale,
layer.scaling,
logit_cap=logits_soft_cap,
sliding_window_size=sliding_window_size,
sinks=sinks,
window_kv_offsets=window_kv_offsets,
xai_temperature_len=layer.xai_temperature_len,
max_bs=self.req_to_token_pool.size,
)
):
return o
self.extend_attention_fwd(
q.view(-1, layer.tp_q_head_num, layer.qk_head_dim),
k.contiguous(),
python/sglang/srt/layers/attention/triton_ops/verify_splitkv.py
+803
View File
@@ -0,0 +1,803 @@
"""Split-KV (flash-decode) attention for EAGLE speculative *verify*.
Only valid when speculative ``topk == 1`` (the EAGLE tree reduces to a pure
causal chain); the caller gates on that. ``topk > 1`` trees fall back to
``extend_attention_fwd``.
On the Triton backend, EAGLE target-verify runs through the prefill
``extend_attention_fwd``, which loops the (long) prefix KV serially per
(sequence, head). With only a few draft-token queries, that leaves the GPU
memory system far under-utilized at long context. This kernel instead splits
the prefix KV across parallel programs (flash-decode style) and combines the
partials with a log-sum-exp merge, then handles the small causal draft-draft
block -- recovering memory bandwidth on the verify path.
Two Triton kernels:
* ``_verify_prefix_stage1``: split-KV over the shared prefix. Applies the fp8
dequant multipliers ``k_scale`` (on the QK score) and ``v_scale`` (on the
prefix output), matching ``extend_attention_fwd``'s ``_fwd_kernel``
(qk *= sm_scale * k_scale; acc += dot(p, v) * v_scale on the prefix loop;
NO scaling on the draft-draft loop, whose K/V are the fresh bf16 draft
tensors, not the fp8 pool). fp8 K/V buffers are handled by casting q to the
buffer dtype before the dot (mirrors ``q.to(k.dtype)`` in the baseline).
* ``_verify_combine_stage2``: combines the prefix splits (LSE merge) with the
small causal draft-draft block and writes the output.
``verify_splitkv_fwd(...)`` takes the SAME positional args as
``extend_attention_fwd``; it runs the split-KV path when it can serve the case
bit-equivalently and returns True, otherwise returns False (doing nothing) so
the caller falls back to ``extend_attention_fwd``. Supported case: causal
(topk=1) verify with a constant per-sequence extend length, no sinks /
sliding-window / logit-cap / xai-temperature. Correctness is never violated.
"""
import torch
import triton
import triton.language as tl
from sglang.srt.utils import is_hip
_MIN_BLOCK_KV = 32
# AMD/CDNA-only Triton launch hints (waves_per_eu, matrix_instr_nonkdim); NVIDIA's
# Triton rejects these kwargs, so only pass them on ROCm. In production this kernel
# is dispatched only on AMD (see TritonAttnBackend); keeping it NV-safe lets the
# numerics test run on the CUDA CI lane.
_IS_HIP = is_hip()
_AMD_LAUNCH_KWARGS = {"waves_per_eu": 4, "matrix_instr_nonkdim": 16} if _IS_HIP else {}
# Block-size config keyed on head_dim. The (BLOCK_N, num_warps) tile that best
# hides latency depends on head_dim: at head_dim=256 (Qwen3 family) a narrower
# BLOCK_N with more warps wins, since the 256-wide QK/PV tiles are register
# heavy. head_dim=256 is the value validated on MI350X; other head dims use a
# conservative default. Block size affects PERFORMANCE only, never correctness
# (any valid block size produces the same result).
DEFAULT_N_SPLITS = 8
DEFAULT_BLOCK_N = 32
DEFAULT_NUM_WARPS = 4
_BLOCK_CONFIG = {
# head_dim: (BLOCK_N, num_warps)
256: (32, 4),
}
def block_config(head_dim):
"""Return (BLOCK_N, num_warps) for a head_dim; default for untuned dims."""
return _BLOCK_CONFIG.get(head_dim, (DEFAULT_BLOCK_N, DEFAULT_NUM_WARPS))
# ---------------------------------------------------------------------------
# Adaptive N_SPLITS.
# ---------------------------------------------------------------------------
# The prefix split-KV stage launches a (bs, h_q, N_SPLITS) grid; each (b,h,s)
# program handles kv_len_per_split = cdiv(cdiv(seqlen, N_SPLITS), MIN)*MIN keys.
# A fixed N_SPLITS=16 over-splits short/mid contexts (each split does too little
# work -> launch + reduction overhead dominates) and under-splits very long ones
# (too few parallel waves to saturate the device, raising tail latency on the
# slow split). Mirror the decode kernel's intent (decode_attention.py
# get_num_kv_splits): pick the split count per-dispatch from the representative
# sequence length, growing gradually with seqlen and capped at MAX.
#
# CRITICAL: this must be computed from STATIC shapes only (no .item()/.cpu()
# sync), because the verify/draft-extend step runs inside a captured HIP graph
# where a device->host copy raises hipErrorStreamCaptureUnsupported. We use the
# average prefix length = kv_indices.shape[0] / bs, which is a pure python int
# from tensor shapes -- no device read. N_SPLITS is then a power of two so the
# stage2 reduction tile (tl.arange(0, N_SPLITS)) stays cheap.
#
# Split-count bounds (internal constants). MAX=16 is the MI350X cap: 32
# oversubscribes the device and regresses, per tuning.
ADAPTIVE_SPLITS = True
MAX_N_SPLITS = 16
MIN_N_SPLITS = 4
def choose_n_splits(avg_seqlen):
"""Pick N_SPLITS (power of two, in [MIN_N_SPLITS, MAX_N_SPLITS]) from the
average prefix length. Tuned by the real-shape sweep (head_dim=256, BS*H_Q
=128 base programs on ~132 CUs):
ctx < 4k -> 4 (short: extra splits add launch/reduction overhead)
4k <= ctx < 8k -> 8 (sweet spot: best across 1k-16k in the sweep)
ctx >= 8k -> 16 (long: a few more splits help latency-bound tail)
Never 32 (4096 grid blocks oversubscribes the device and regresses, per the
sweep). Computed from a static shape (avg prefix = kv_indices.shape[0]/bs),
so it is HIP-graph-capture safe (no device->host sync)."""
if not ADAPTIVE_SPLITS:
return DEFAULT_N_SPLITS
s = int(avg_seqlen)
if s < 4096:
n = 4
elif s < 8192:
n = 8
else:
n = 16
if n < MIN_N_SPLITS:
n = MIN_N_SPLITS
if n > MAX_N_SPLITS:
n = MAX_N_SPLITS
return n
@triton.jit
def _verify_prefix_stage1(
Q, # [extend_tokens, H_Q, D]
K_Buffer, # [pool_tokens, H_KV, D]
V_Buffer, # [pool_tokens, H_KV, Dv]
sm_scale,
k_scale, # fp8 dequant multiplier for prefix K (1.0 if bf16)
v_scale, # fp8 dequant multiplier for prefix V (1.0 if bf16)
qo_indptr, # [BS+1] int32 -> rows of Q (draft queries)
kv_indptr, # [BS+1] int32 -> rows of kv_indices (prefix)
kv_indices, # [sum prefix] int64
Att_Out, # [BS, H_Q, N_SPLITS, L_EXT, Dv] fp32
Att_Lse, # [BS, H_Q, N_SPLITS, L_EXT] fp32
stride_qbs,
stride_qh,
stride_buf_kbs,
stride_buf_kh,
stride_buf_vbs,
stride_buf_vh,
stride_ob,
stride_oh,
stride_os,
stride_ol,
stride_lb,
stride_lh,
stride_ls,
kv_group_num: tl.constexpr,
N_SPLITS: tl.constexpr,
L_EXT: tl.constexpr, # padded power-of-2 row tile (>= real l_ext)
BLOCK_DMODEL: tl.constexpr,
BLOCK_DV: tl.constexpr,
BLOCK_N: tl.constexpr,
MIN_BLOCK_KV: tl.constexpr,
):
cur_batch = tl.program_id(0)
cur_head = tl.program_id(1)
split_kv_id = tl.program_id(2)
cur_kv_head = cur_head // kv_group_num
offs_d = tl.arange(0, BLOCK_DMODEL)
offs_dv = tl.arange(0, BLOCK_DV)
offs_l = tl.arange(0, L_EXT)
# real number of draft query tokens for this seq
cur_q_start = tl.load(qo_indptr + cur_batch)
l_ext = tl.load(qo_indptr + cur_batch + 1) - cur_q_start
mask_l = offs_l < l_ext
cur_batch_kv_start_idx = tl.load(kv_indptr + cur_batch)
cur_batch_seq_len = tl.load(kv_indptr + cur_batch + 1) - cur_batch_kv_start_idx
# split sizing identical to the decode kernel
kv_len_per_split = (
tl.cdiv(tl.cdiv(cur_batch_seq_len, N_SPLITS), MIN_BLOCK_KV) * MIN_BLOCK_KV
)
split_kv_start = kv_len_per_split * split_kv_id
split_kv_end = tl.minimum(split_kv_start + kv_len_per_split, cur_batch_seq_len)
e_max = tl.zeros([L_EXT], dtype=tl.float32) - float("inf")
e_sum = tl.zeros([L_EXT], dtype=tl.float32)
acc = tl.zeros([L_EXT, BLOCK_DV], dtype=tl.float32)
if split_kv_end > split_kv_start:
# q tile: [L_EXT, D]
offs_q = (
(cur_q_start + offs_l)[:, None] * stride_qbs
+ cur_head * stride_qh
+ offs_d[None, :]
)
q = tl.load(Q + offs_q, mask=mask_l[:, None], other=0.0)
q_k = q.to(K_Buffer.dtype.element_ty)
base_offs_k = cur_kv_head * stride_buf_kh + offs_d[:, None]
base_offs_v = cur_kv_head * stride_buf_vh + offs_dv[None, :]
for start_n in tl.range(split_kv_start, split_kv_end, BLOCK_N):
offs_n = start_n + tl.arange(0, BLOCK_N)
n_mask = offs_n < split_kv_end
kv_loc = tl.load(
kv_indices + cur_batch_kv_start_idx + offs_n,
mask=n_mask,
other=0,
)
# K block: [D, BLOCK_N]
offs_buf_k = kv_loc[None, :] * stride_buf_kbs + base_offs_k
k = tl.load(K_Buffer + offs_buf_k, mask=n_mask[None, :], other=0.0)
qk = tl.dot(q_k, k) # [L_EXT, BLOCK_N]
qk *= sm_scale * k_scale # fp8 dequant of prefix K (k_scale==1 if bf16)
# NO causal mask: full prefix is visible to all draft tokens.
qk = tl.where(n_mask[None, :], qk, float("-inf"))
# V block: [BLOCK_N, Dv]
offs_buf_v = kv_loc[:, None] * stride_buf_vbs + base_offs_v
v = tl.load(V_Buffer + offs_buf_v, mask=n_mask[:, None], other=0.0)
n_e_max = tl.maximum(tl.max(qk, 1), e_max)
re_scale = tl.exp(e_max - n_e_max)
p = tl.exp(qk - n_e_max[:, None])
acc *= re_scale[:, None]
acc += tl.dot(p.to(v.dtype), v)
e_sum = e_sum * re_scale + tl.sum(p, 1)
e_max = n_e_max
# fp8 dequant of prefix V: scale the accumulated (pre-normalised) output.
acc *= v_scale
offs_o = (
cur_batch * stride_ob
+ cur_head * stride_oh
+ split_kv_id * stride_os
+ offs_l[:, None] * stride_ol
+ offs_dv[None, :]
)
tl.store(Att_Out + offs_o, acc / e_sum[:, None], mask=mask_l[:, None])
offs_lse = (
cur_batch * stride_lb
+ cur_head * stride_lh
+ split_kv_id * stride_ls
+ offs_l
)
tl.store(Att_Lse + offs_lse, e_max + tl.log(e_sum), mask=mask_l)
else:
# split did not run: write a sentinel lse so stage2 can ignore it.
offs_lse = (
cur_batch * stride_lb
+ cur_head * stride_lh
+ split_kv_id * stride_ls
+ offs_l
)
tl.store(
Att_Lse + offs_lse,
tl.zeros([L_EXT], tl.float32) - float("inf"),
mask=mask_l,
)
@triton.jit
def _verify_combine_stage2(
Att_Out, # [BS, H_Q, N_SPLITS, L_EXT, Dv] fp32
Att_Lse, # [BS, H_Q, N_SPLITS, L_EXT] fp32
Q, # [extend_tokens, H_Q, D] (draft queries)
K_Extend, # [extend_tokens, H_KV, D]
V_Extend, # [extend_tokens, H_KV, Dv]
O_Out, # [extend_tokens, H_Q, Dv] (final, written)
sm_scale,
qo_indptr, # [BS+1] int32
stride_ob,
stride_oh,
stride_os,
stride_ol,
stride_lb,
stride_lh,
stride_ls,
stride_qbs,
stride_qh,
stride_kebs,
stride_keh,
stride_vebs,
stride_veh,
stride_oobs,
stride_ooh,
kv_group_num: tl.constexpr,
N_SPLITS: tl.constexpr,
L_EXT: tl.constexpr,
BLOCK_DMODEL: tl.constexpr,
BLOCK_DV: tl.constexpr,
):
cur_batch = tl.program_id(0)
cur_head = tl.program_id(1)
cur_kv_head = cur_head // kv_group_num
offs_d = tl.arange(0, BLOCK_DMODEL)
offs_dv = tl.arange(0, BLOCK_DV)
offs_l = tl.arange(0, L_EXT)
offs_s = tl.arange(0, N_SPLITS)
cur_q_start = tl.load(qo_indptr + cur_batch)
l_ext = tl.load(qo_indptr + cur_batch + 1) - cur_q_start
mask_l = offs_l < l_ext
# ---- (a) combine prefix splits (logsumexp) ----------------------------
# lse: [N_SPLITS, L_EXT]
offs_lse = (
cur_batch * stride_lb
+ cur_head * stride_lh
+ offs_s[:, None] * stride_ls
+ offs_l[None, :]
)
lse = tl.load(offs_lse + Att_Lse) # [N_SPLITS, L_EXT]
m_p = tl.max(lse, 0) # [L_EXT]
w = tl.exp(lse - m_p[None, :]) # [N_SPLITS, L_EXT]; -inf->0
denom_p = tl.sum(w, 0) # [L_EXT]
# weighted-sum of partial outputs: o_prefix[L_EXT, Dv]
# Att_Out[b,h,s,l,dv]
offs_ao = (
cur_batch * stride_ob
+ cur_head * stride_oh
+ offs_s[:, None, None] * stride_os
+ offs_l[None, :, None] * stride_ol
+ offs_dv[None, None, :]
)
ao = tl.load(offs_ao + Att_Out) # [N_SPLITS, L_EXT, Dv]
o_prefix = tl.sum(ao * w[:, :, None], 0) # [L_EXT, Dv]
o_prefix = o_prefix / denom_p[:, None]
lse_prefix = m_p + tl.log(denom_p) # [L_EXT]
# ---- (b) draft-draft causal attention (L_EXT x L_EXT) -----------------
# load draft queries [L_EXT, D], draft K/V [L_EXT, D]/[L_EXT, Dv]
offs_q = (
(cur_q_start + offs_l)[:, None] * stride_qbs
+ cur_head * stride_qh
+ offs_d[None, :]
)
q = tl.load(Q + offs_q, mask=mask_l[:, None], other=0.0).to(tl.float32)
offs_ke = (
(cur_q_start + offs_l)[:, None] * stride_kebs
+ cur_kv_head * stride_keh
+ offs_d[None, :]
)
ke = tl.load(K_Extend + offs_ke, mask=mask_l[:, None], other=0.0).to(tl.float32)
offs_ve = (
(cur_q_start + offs_l)[:, None] * stride_vebs
+ cur_kv_head * stride_veh
+ offs_dv[None, :]
)
ve = tl.load(V_Extend + offs_ve, mask=mask_l[:, None], other=0.0).to(tl.float32)
# scores[i,j] = q_i . k_j (i query, j key) -> [L_EXT, L_EXT]
qk = tl.sum(q[:, None, :] * ke[None, :, :], 2) * sm_scale
# causal among drafts: query i sees key j iff j <= i, and both valid
causal = (offs_l[None, :] <= offs_l[:, None]) & mask_l[None, :] & mask_l[:, None]
qk = tl.where(causal, qk, float("-inf"))
m_d = tl.max(qk, 1) # [L_EXT]
pd = tl.exp(qk - m_d[:, None]) # [L_EXT, L_EXT]
denom_d = tl.sum(pd, 1) # [L_EXT]
o_draft = tl.sum(pd[:, :, None] * ve[None, :, :], 1) # [L_EXT, Dv]
o_draft = o_draft / denom_d[:, None]
lse_draft = m_d + tl.log(denom_d) # [L_EXT]
# ---- (c) final LSE merge (prefix vs draft) ----------------------------
m = tl.maximum(lse_prefix, lse_draft)
wp = tl.exp(lse_prefix - m)
wd = tl.exp(lse_draft - m)
o = (o_prefix * wp[:, None] + o_draft * wd[:, None]) / (wp + wd)[:, None]
offs_oo = (
(cur_q_start + offs_l)[:, None] * stride_oobs
+ cur_head * stride_ooh
+ offs_dv[None, :]
)
tl.store(O_Out + offs_oo, o.to(O_Out.dtype.element_ty), mask=mask_l[:, None])
class VerifySplitKV:
"""Pre-allocates scratch buffers for a problem shape and runs the split-KV
verify attention end to end (two Triton launches: prefix split-KV + fused
combine/draft/merge). Buffers are sized by ``max_bs`` (constant for the
server lifetime) and reused for every batch size <= max_bs, so their
addresses stay fixed (CUDA/HIP-graph safe) and GPU memory does not grow per
batch size. The kernel grid uses the actual per-call bs (<= max_bs)."""
def __init__(
self,
max_bs,
h_q,
h_kv,
head_dim,
v_head_dim,
l_ext,
device="cuda",
n_splits=DEFAULT_N_SPLITS,
block_n=DEFAULT_BLOCK_N,
num_warps=DEFAULT_NUM_WARPS,
):
self.h_q = h_q
self.h_kv = h_kv
self.group = h_q // h_kv
self.head_dim = head_dim
self.v_head_dim = v_head_dim
self.l_ext = l_ext # real draft tokens per seq (fixed == 4)
self.l_pad = triton.next_power_of_2(l_ext)
self.device = device
self.n_splits = n_splits
self.block_n = block_n
self.num_warps = num_warps
self._alloc(max_bs)
def _alloc(self, max_bs):
# prefix split partials (fp32), sized for the maximum batch size.
self.max_bs = max_bs
self.att_out = torch.empty(
(max_bs, self.h_q, self.n_splits, self.l_pad, self.v_head_dim),
dtype=torch.float32,
device=self.device,
)
self.att_lse = torch.empty(
(max_bs, self.h_q, self.n_splits, self.l_pad),
dtype=torch.float32,
device=self.device,
)
def grow_buffers(self, max_bs):
if max_bs > self.max_bs:
self._alloc(max_bs)
def _run_prefix_kernel(
self,
bs,
q_extend,
k_buffer,
v_buffer,
qo_indptr,
kv_indptr,
kv_indices,
sm_scale,
k_scale,
v_scale,
):
grid = (bs, self.h_q, self.n_splits)
_verify_prefix_stage1[grid](
q_extend,
k_buffer,
v_buffer,
sm_scale,
k_scale,
v_scale,
qo_indptr,
kv_indptr,
kv_indices,
self.att_out,
self.att_lse,
q_extend.stride(0),
q_extend.stride(1),
k_buffer.stride(0),
k_buffer.stride(1),
v_buffer.stride(0),
v_buffer.stride(1),
self.att_out.stride(0),
self.att_out.stride(1),
self.att_out.stride(2),
self.att_out.stride(3),
self.att_lse.stride(0),
self.att_lse.stride(1),
self.att_lse.stride(2),
kv_group_num=self.group,
N_SPLITS=self.n_splits,
L_EXT=self.l_pad,
BLOCK_DMODEL=triton.next_power_of_2(self.head_dim),
BLOCK_DV=triton.next_power_of_2(self.v_head_dim),
BLOCK_N=self.block_n,
MIN_BLOCK_KV=_MIN_BLOCK_KV,
num_warps=self.num_warps,
num_stages=1,
**_AMD_LAUNCH_KWARGS,
)
def _run_combine_kernel(
self, bs, q_extend, k_extend, v_extend, o_out, qo_indptr, sm_scale
):
grid = (bs, self.h_q)
_verify_combine_stage2[grid](
self.att_out,
self.att_lse,
q_extend,
k_extend,
v_extend,
o_out,
sm_scale,
qo_indptr,
self.att_out.stride(0),
self.att_out.stride(1),
self.att_out.stride(2),
self.att_out.stride(3),
self.att_lse.stride(0),
self.att_lse.stride(1),
self.att_lse.stride(2),
q_extend.stride(0),
q_extend.stride(1),
k_extend.stride(0),
k_extend.stride(1),
v_extend.stride(0),
v_extend.stride(1),
o_out.stride(0),
o_out.stride(1),
kv_group_num=self.group,
N_SPLITS=self.n_splits,
L_EXT=self.l_pad,
BLOCK_DMODEL=triton.next_power_of_2(self.head_dim),
BLOCK_DV=triton.next_power_of_2(self.v_head_dim),
num_warps=1,
num_stages=1,
)
def __call__(
self,
q_extend,
k_extend,
v_extend,
k_buffer,
v_buffer,
qo_indptr,
kv_indptr,
kv_indices,
sm_scale,
o_out=None,
k_scale=1.0,
v_scale=1.0,
):
if o_out is None:
o_out = torch.empty(
(q_extend.shape[0], self.h_q, self.v_head_dim),
dtype=q_extend.dtype,
device=q_extend.device,
)
# actual batch size for this call (<= max_bs); the grid uses it while the
# scratch buffers stay max_bs-sized (only the first bs slices are touched).
bs = qo_indptr.shape[0] - 1
# 1. prefix split-KV
self._run_prefix_kernel(
bs,
q_extend,
k_buffer,
v_buffer,
qo_indptr,
kv_indptr,
kv_indices,
sm_scale,
k_scale,
v_scale,
)
# 2+3+4. fused combine + draft-draft + merge
self._run_combine_kernel(
bs,
q_extend,
k_extend,
v_extend,
o_out,
qo_indptr,
sm_scale,
)
return o_out
# ---------------------------------------------------------------------------
# Live-server dispatch entry.
# ---------------------------------------------------------------------------
# Cache one VerifySplitKV instance per (h_q, h_kv, head_dim, v_head_dim, l_ext,
# device, n_splits) shape -- NOT keyed on the dynamic batch size. Buffers are
# sized by the stable max_bs (grown only if a larger one is ever requested), so
# a single instance serves every batch size: addresses stay fixed (graph-safe)
# and GPU memory does not grow per batch size.
_VK_CACHE = {}
def _get_vk(
max_bs, h_q, h_kv, head_dim, v_head_dim, l_ext, device, n_splits=DEFAULT_N_SPLITS
):
key = (h_q, h_kv, head_dim, v_head_dim, l_ext, str(device), n_splits)
vk = _VK_CACHE.get(key)
if vk is None:
block_n, num_warps = block_config(head_dim)
vk = VerifySplitKV(
max_bs,
h_q,
h_kv,
head_dim,
v_head_dim,
l_ext,
device=device,
n_splits=n_splits,
block_n=block_n,
num_warps=num_warps,
)
_VK_CACHE[key] = vk
else:
vk.grow_buffers(max_bs)
return vk
def can_handle(
q_extend,
k_extend,
v_extend,
k_buffer,
v_buffer,
qo_indptr,
kv_indptr,
kv_indices,
custom_mask,
is_causal,
mask_indptr,
max_len_extend,
sliding_window_size=-1,
sinks=None,
logit_cap=0.0,
xai_temperature_len=-1,
):
"""Return True iff the 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_splitkv = ... 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:
return False
if sliding_window_size is not None and sliding_window_size > 0:
return False
if logit_cap and logit_cap > 0:
return False
if xai_temperature_len is not None and xai_temperature_len > 0:
return False
if not is_causal:
return False
# q layout must be [tokens, H_Q, D]; head dims handled by power-of-2 pad.
if q_extend.dim() != 3 or k_extend.dim() != 3 or v_extend.dim() != 3:
return False
# GQA group must divide evenly.
h_q = q_extend.shape[1]
h_kv = k_extend.shape[1]
if h_kv == 0 or h_q % h_kv != 0:
return False
# head dims must match buffers.
if k_buffer.shape[1] != h_kv or v_buffer.shape[1] != h_kv:
return False
if q_extend.shape[2] != k_extend.shape[2]:
return False
if q_extend.shape[2] != k_buffer.shape[2]:
return False
if v_extend.shape[2] != v_buffer.shape[2]:
return False
# NOTE: must NOT read any tensor *values* here (no .item()/.cpu()): the
# target-verify step runs inside a captured CUDA/HIP graph, where a
# device->host sync raises hipErrorStreamCaptureUnsupported. We therefore
# gate purely on static shapes/dtypes/python scalars.
bs = qo_indptr.shape[0] - 1
if bs < 1:
return False
# max_len_extend must be a known positive python int (it is the static
# server_args.speculative_num_draft_tokens for the verify path). For
# topk=1 the per-seq extend len is constant == num_draft_tokens ==
# max_len_extend by construction of qo_indptr (arange with that step), so
# the L_EXT row-tile mask is exactly right and the tree custom_mask equals
# causal -- no value inspection required.
try:
mle = int(max_len_extend)
except (TypeError, ValueError):
return False
if mle < 1:
return False
# The packed extend tensor must hold exactly bs * max_len_extend rows
# (constant extend len). This is a pure shape check (no sync) and rejects
# any ragged/variable-extend batch -> falls back to the baseline.
if q_extend.shape[0] != bs * mle:
return False
return True
def verify_splitkv_fwd(
q_extend,
k_extend,
v_extend,
o_extend,
k_buffer,
v_buffer,
qo_indptr,
kv_indptr,
kv_indices,
custom_mask,
is_causal,
mask_indptr,
max_len_extend,
k_scale,
v_scale,
sm_scale=None,
logit_cap=0.0,
skip_prefix_custom_mask=True,
sliding_window_size=-1,
sinks=None,
window_kv_offsets=None,
xai_temperature_len=-1,
max_bs=None,
):
"""Drop-in for extend_attention_fwd on the EAGLE target-verify (topk=1)
shape. Returns True if it ran (o_extend written), False if the case is
unsupported and the caller must fall back to extend_attention_fwd.
``max_bs`` (optional) is the stable maximum batch size used to size the
cached scratch buffers; the backend passes its req_to_token_pool size. If
omitted it defaults to this call's bs.
Arg order mirrors extend_attention_fwd exactly so the call site is a
one-line swap.
"""
if not can_handle(
q_extend,
k_extend,
v_extend,
k_buffer,
v_buffer,
qo_indptr,
kv_indptr,
kv_indices,
custom_mask,
is_causal,
mask_indptr,
max_len_extend,
sliding_window_size=sliding_window_size,
sinks=sinks,
logit_cap=logit_cap,
xai_temperature_len=xai_temperature_len,
):
return False
bs = qo_indptr.shape[0] - 1
h_q = q_extend.shape[1]
h_kv = k_extend.shape[1]
head_dim = q_extend.shape[2]
v_head_dim = v_extend.shape[2]
l_ext = int(max_len_extend)
if sm_scale is None:
sm_scale = 1.0 / (head_dim**0.5)
# k_scale/v_scale may be float or 0-d tensor; coerce to python float.
try:
k_scale = float(k_scale)
except (TypeError, ValueError):
k_scale = 1.0
try:
v_scale = float(v_scale)
except (TypeError, ValueError):
v_scale = 1.0
# Adaptive split count from the average prefix length. This is a
# pure-shape derivation (kv_indices.shape[0] / bs) -- no device->host sync,
# so it is safe inside a captured HIP graph. The whole batch shares one
# N_SPLITS (the grid dim must be a launch constexpr); the per-split kernel
# logic still clamps each split's [start,end) to that seq's real length, so
# mixed-length batches stay correct -- shorter seqs simply write fewer
# active splits (the rest emit the -inf lse sentinel, ignored in stage2).
avg_seqlen = kv_indices.shape[0] / max(1, bs)
n_splits = choose_n_splits(avg_seqlen)
# Size scratch by the stable max_bs (backend passes req_to_token_pool size);
# fall back to this call's bs if not provided / smaller.
if max_bs is None or max_bs < bs:
max_bs = bs
vk = _get_vk(
max_bs,
h_q,
h_kv,
head_dim,
v_head_dim,
l_ext,
q_extend.device,
n_splits=n_splits,
)
vk(
q_extend,
k_extend.contiguous(),
v_extend.contiguous(),
k_buffer,
v_buffer,
qo_indptr,
kv_indptr,
kv_indices,
sm_scale,
o_out=o_extend,
k_scale=k_scale,
v_scale=v_scale,
)
return True
test/registered/attention/test_verify_splitkv.py
+239
View File
@@ -0,0 +1,239 @@
"""Correctness tests for the split-KV EAGLE-verify attention kernel.
``verify_splitkv_fwd`` is a drop-in for ``extend_attention_fwd`` on the topk=1
causal (target-verify) path. These tests check:
(a) numerical parity with ``extend_attention_fwd`` on the pure-causal path
(the operative case at topk=1), across head dims / GQA ratios / prefix
lengths / extend lengths / KV scales;
(b) ``can_handle()`` rejects cases the kernel cannot serve bit-equivalently
(non-causal, sinks, sliding-window, logit-cap, ragged extend), so the
backend falls back to ``extend_attention_fwd``.
The topk>1 case is gated off in the backend (TritonAttnBackend enables this path
only when ``self.topk == 1``), since the kernel ignores the tree custom_mask;
that gate is exercised end-to-end by the nightly ROCm spec accuracy test.
GPU + Triton required. Runs on the CUDA PR lane and the AMD MI35x lane.
"""
import unittest
import torch
from sglang.srt.layers.attention.triton_ops.extend_attention import (
extend_attention_fwd,
)
from sglang.srt.layers.attention.triton_ops.verify_splitkv import (
can_handle,
verify_splitkv_fwd,
)
from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
from sglang.test.test_utils import CustomTestCase
register_cuda_ci(est_time=30, stage="base-b", runner_config="1-gpu-small")
register_amd_ci(est_time=30, suite="stage-b-test-1-gpu-small-amd-mi35x")
# Split-KV accumulates the prefix in parallel splits and merges via log-sum-exp;
# it differs from the single-pass extend kernel only by reduction order. On-GPU
# (gfx950) the max abs diff across these shapes was ~2e-3 (rising to ~6e-2 only
# at much longer ctx); 2e-2 keeps a ~10x margin over the observed noise.
ATOL = 2e-2
RTOL = 1e-2
def _build_verify_inputs(
prefix_lens, l_ext, h_q, h_kv, head_dim, v_head_dim, dtype, device
):
"""Build a verify-shaped problem: a constant extend length ``l_ext`` per
sequence, with the prefix cache addressed contiguously by ``kv_indices``.
Returns the positional args shared by extend_attention_fwd / verify_splitkv_fwd.
"""
B = len(prefix_lens)
prefix_lens_t = torch.tensor(prefix_lens, dtype=torch.int32, device=device)
total_prefix = int(prefix_lens_t.sum())
# Prefix KV cache laid out contiguously; kv_indices is just arange over it.
# kv_indices is int64 in production (TritonAttnBackend allocates int64) -- match it.
k_buffer = torch.randn(total_prefix, h_kv, head_dim, dtype=dtype, device=device)
v_buffer = torch.randn(total_prefix, h_kv, v_head_dim, dtype=dtype, device=device)
kv_indptr = torch.zeros(B + 1, dtype=torch.int32, device=device)
kv_indptr[1:] = torch.cumsum(prefix_lens_t, 0)
kv_indices = torch.arange(total_prefix, dtype=torch.int64, device=device)
# Draft (extend) tensors: constant l_ext rows per sequence.
n_ext = B * l_ext
q_extend = torch.randn(n_ext, h_q, head_dim, dtype=dtype, device=device)
k_extend = torch.randn(n_ext, h_kv, head_dim, dtype=dtype, device=device)
v_extend = torch.randn(n_ext, h_kv, v_head_dim, dtype=dtype, device=device)
qo_indptr = torch.arange(0, n_ext + 1, l_ext, dtype=torch.int32, device=device)
return (
q_extend,
k_extend,
v_extend,
k_buffer,
v_buffer,
qo_indptr,
kv_indptr,
kv_indices,
l_ext,
)
@unittest.skipIf(not torch.cuda.is_available(), "GPU required")
class TestVerifySplitKV(CustomTestCase):
def _run_parity(
self,
prefix_lens,
l_ext=4,
h_q=16,
h_kv=2,
head_dim=256,
v_head_dim=256,
k_scale=1.0,
v_scale=1.0,
dtype=torch.bfloat16,
):
device = "cuda"
q, k, v, kb, vb, qo, kvp, kvi, mle = _build_verify_inputs(
prefix_lens, l_ext, h_q, h_kv, head_dim, v_head_dim, dtype, device
)
sm_scale = 1.0 / (head_dim**0.5)
# Reference: extend_attention_fwd, pure causal (custom_mask=None) -- the
# topk=1 operative case, with the same KV scales.
o_ref = torch.empty(q.shape[0], h_q, v_head_dim, dtype=dtype, device=device)
extend_attention_fwd(
q,
k,
v,
o_ref,
kb,
vb,
qo,
kvp,
kvi,
None,
True,
None,
mle,
k_scale,
v_scale,
sm_scale=sm_scale,
)
o_split = torch.empty_like(o_ref)
ran = verify_splitkv_fwd(
q,
k,
v,
o_split,
kb,
vb,
qo,
kvp,
kvi,
None,
True,
None,
mle,
k_scale,
v_scale,
sm_scale=sm_scale,
)
self.assertTrue(ran, "verify_splitkv_fwd must handle the topk=1 causal case")
torch.testing.assert_close(o_split, o_ref, atol=ATOL, rtol=RTOL)
def test_numerics_head_dim_256(self):
# head_dim=256 is the validated Qwen3 value (the tuned block config).
for prefix_lens in ([512, 512, 512], [768, 1536, 3072], [4096, 8192]):
with self.subTest(prefix_lens=prefix_lens):
self._run_parity(prefix_lens)
def test_numerics_head_dim_128(self):
# A head_dim without a tuned block entry must still be correct (default).
self._run_parity([1024, 2048], head_dim=128, v_head_dim=128)
def test_numerics_gqa_ratios(self):
# Sweep GQA group sizes incl. MQA (h_kv=1); the kv_group_num arithmetic
# in the kernel must be correct across ratios.
for h_q, h_kv in ((16, 1), (8, 1), (8, 2), (8, 4), (8, 8)):
with self.subTest(h_q=h_q, h_kv=h_kv):
self._run_parity([1024, 2048], h_q=h_q, h_kv=h_kv)
def test_numerics_extend_len_variants(self):
for l_ext in (1, 2, 4, 8):
with self.subTest(l_ext=l_ext):
self._run_parity([1024, 1024], l_ext=l_ext)
def test_numerics_with_kv_scales(self):
# Exercise the k_scale/v_scale dequant-multiplier path (same multipliers
# the fp8 KV-cache path applies); both kernels must apply them identically.
self._run_parity([1024, 2048], k_scale=0.5, v_scale=0.25)
# --- fallback: can_handle() must reject what the kernel can't serve --------
# (topk>1 is gated off in the backend, not here -- can_handle never inspects
# the tree custom_mask; see verify_splitkv.can_handle docstring.)
def _inputs(self):
return _build_verify_inputs(
[512, 512], 4, 16, 2, 256, 256, torch.bfloat16, "cuda"
)
def test_fallback_non_causal(self):
q, k, v, kb, vb, qo, kvp, kvi, mle = self._inputs()
self.assertFalse(
can_handle(q, k, v, kb, vb, qo, kvp, kvi, None, False, None, mle)
)
def test_fallback_exotic_features(self):
q, k, v, kb, vb, qo, kvp, kvi, mle = self._inputs()
self.assertFalse(
can_handle(
q,
k,
v,
kb,
vb,
qo,
kvp,
kvi,
None,
True,
None,
mle,
sinks=torch.zeros(16, device="cuda"),
)
)
self.assertFalse(
can_handle(
q,
k,
v,
kb,
vb,
qo,
kvp,
kvi,
None,
True,
None,
mle,
sliding_window_size=128,
)
)
self.assertFalse(
can_handle(
q, k, v, kb, vb, qo, kvp, kvi, None, True, None, mle, logit_cap=30.0
)
)
def test_fallback_ragged_extend(self):
# q rows (bs*l_ext) inconsistent with the claimed max_len_extend -> reject.
q, k, v, kb, vb, qo, kvp, kvi, mle = self._inputs()
self.assertFalse(
can_handle(q, k, v, kb, vb, qo, kvp, kvi, None, True, None, mle + 1)
)
if __name__ == "__main__":
unittest.main()