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

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

python/sglang/srt/environ.py

@@ -643,6 +643,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

@@ -7,6 +7,7 @@
 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 @@
     get_device_core_count,
     get_int_env_var,
     is_cuda,
+    is_gfx95_supported,
     is_gfx942_supported,
     next_power_of_2,
 )
@@ -107,6 +109,9 @@ def __init__(
             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 @@ def __init__(
             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 @@ def __init__(
         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_attention_tp_size()
@@ -1094,6 +1117,43 @@ def forward_extend(
             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(),