chore: import upstream snapshot with attribution

python/sglang/kernels/ops/attention/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