/test/registered/attention/test_verify_splitkv.py

"""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.kernels.ops.attention.extend_attention import (
    extend_attention_fwd,
)
from sglang.kernels.ops.attention.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()