diff --git a/benchmark/kernels/verify_splitkv_triton/bench_verify_splitkv.py b/benchmark/kernels/verify_splitkv_triton/bench_verify_splitkv.py new file mode 100644 index 000000000..eb35e8ae2 --- /dev/null +++ b/benchmark/kernels/verify_splitkv_triton/bench_verify_splitkv.py @@ -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() diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index 6b85f2f98..cea221453 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -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. diff --git a/python/sglang/srt/layers/attention/triton_backend.py b/python/sglang/srt/layers/attention/triton_backend.py index a44dbd2fc..ad1151ed4 100644 --- a/python/sglang/srt/layers/attention/triton_backend.py +++ b/python/sglang/srt/layers/attention/triton_backend.py @@ -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(), diff --git a/python/sglang/srt/layers/attention/triton_ops/verify_splitkv.py b/python/sglang/srt/layers/attention/triton_ops/verify_splitkv.py new file mode 100644 index 000000000..dd6d3fa35 --- /dev/null +++ b/python/sglang/srt/layers/attention/triton_ops/verify_splitkv.py @@ -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 diff --git a/test/registered/attention/test_verify_splitkv.py b/test/registered/attention/test_verify_splitkv.py new file mode 100644 index 000000000..9895e2fae --- /dev/null +++ b/test/registered/attention/test_verify_splitkv.py @@ -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()