"""Per-kernel A/B cases: run the public interface function once per backend and compare.

Each entry maps a kernel name to a builder returning ``(case_name, exact, callable)`` triples.
``exact=True`` demands bitwise equality, which is right for index and metadata kernels; floating
point kernels are compared with a tolerance and are labelled ``exact=False``.

Shapes and dtypes reuse the ones the engine actually passes, taken from the TileLang test matrices
in ``cuda/*.py`` and from the traced inference run.
"""

from __future__ import annotations

from collections.abc import Callable
from typing import Any

import torch

from deepinfer.model_executor.layers.tilelang_ops import ops
from deepinfer.model_executor.layers.tilelang_ops.cache_common import (
    FP8_FLOAT32_SF_INDEX_KV_FORMAT,
    V41_FORMAT,
)
from deepinfer.model_executor.layers.tilelang_ops.compress import (
    compress_and_update_state_decode,
    compress_and_update_state_prefill,
)
from deepinfer.model_executor.layers.tilelang_ops.flash_mla_utils import convert_kv_index
from deepinfer.model_executor.layers.tilelang_ops.fused_transform import (
    CastArgs,
    KVCacheOutArgs,
    NormArgs,
    RopeArgs,
    fused_norm_rope_hadamard_cast_store_kvcache,
)
from deepinfer.model_executor.layers.tilelang_ops.gather import gather
from deepinfer.model_executor.layers.tilelang_ops.gather_engram_tokens_zos import gather_engram_tokens_zos
from deepinfer.model_executor.layers.tilelang_ops.gather_input_ids_zos import gather_input_ids_zos
from deepinfer.model_executor.layers.tilelang_ops.get_kv_index import get_kv_ca_index_decode, get_kv_ca_index_prefill
from deepinfer.model_executor.layers.tilelang_ops.kv_index_cat import kv_index_cat
from deepinfer.model_executor.layers.tilelang_ops.load_paged_kv_cache import load_paged_kv_cache
from deepinfer.model_executor.layers.tilelang_ops.norm import norm
from deepinfer.model_executor.layers.tilelang_ops.prepare_v4_decode_metadata import prepare_v4_decode_metadata
from deepinfer.model_executor.layers.tilelang_ops.sampler import map_lookup
from deepinfer.model_executor.layers.tilelang_ops.store_kv_cache import store_kv_cache
from deepinfer.model_executor.layers.tilelang_ops.store_mtp_info import store_mtp_info
from deepinfer.model_executor.layers.tilelang_ops.swa_index import expand_swa_index


def _cases_gather() -> list[tuple[str, bool, Callable[[], Any]]]:
    """Row gather across the widths the engine uses: metadata rows, heads, and hidden states."""

    def make(n: int, rows: int, d: int, dtype: torch.dtype, index_dtype: torch.dtype, strided: bool):
        def run() -> Any:
            torch.manual_seed(0)
            base = (
                torch.randn(n, d, device="cuda").to(dtype)
                if dtype.is_floating_point
                else torch.randint(0, 100, (n, d), device="cuda", dtype=dtype)
            )
            x = base[:, ::1] if strided else base
            index = torch.randint(-3, n + 3, (rows,), device="cuda", dtype=index_dtype)
            return gather(x, index)

        return run

    return [
        ("n1k_m233_d128_int32_idx", True, make(1024, 233, 128, torch.int8, torch.int32, False)),
        ("n1k_m1_d4_int8", True, make(1024, 1, 4, torch.int8, torch.int32, False)),
        ("n64_m256_d528_bytes", True, make(64, 256, 528, torch.int8, torch.int32, False)),
        ("n512_m300_d4096_bytes", True, make(512, 300, 4096, torch.int8, torch.int32, False)),
        ("n512_m300_d7168_int64_idx", True, make(512, 300, 7168, torch.int8, torch.int64, False)),
        ("n256_m256_d4096_bf16", True, make(256, 256, 4096, torch.bfloat16, torch.int32, False)),
        ("n256_m0_d128_empty", True, make(256, 0, 128, torch.int8, torch.int32, False)),
        ("n1_m17_d128_single_row", True, make(1, 17, 128, torch.int8, torch.int32, False)),
        ("n300_m233_d1025_odd", True, make(300, 233, 1025, torch.int8, torch.int32, False)),
    ]


def _cases_gather_input_ids() -> list[tuple[str, bool, Callable[[], Any]]]:
    """Draft-token replay: remapped requests replay previous ids, fresh ones use their own row."""

    def make(n_seq: int, max_seqlen_q: int, prev_n_seq: int, mix: bool):
        def run() -> Any:
            torch.manual_seed(0)
            nextn = max_seqlen_q - 1
            fresh = torch.randint(0, 1000, (n_seq, max_seqlen_q), device="cuda", dtype=torch.int32)
            prev_sampled = torch.randint(0, 1000, (prev_n_seq,), device="cuda", dtype=torch.int32)
            prev_predict = torch.randint(0, 1000, (prev_n_seq, nextn), device="cuda", dtype=torch.int32)
            if mix:
                remap = torch.tensor([-1 if i % 2 == 0 else i % prev_n_seq for i in range(n_seq)], device="cuda", dtype=torch.int32)
            else:
                remap = torch.full((n_seq,), -1, device="cuda", dtype=torch.int32)
            seqlens = torch.randint(1, max_seqlen_q + 1, (n_seq,), device="cuda", dtype=torch.int32)
            cu = torch.cumsum(seqlens, 0, dtype=torch.int32)
            out = torch.full((int(cu[-1].item()),), -7, device="cuda", dtype=torch.int32)
            gather_input_ids_zos(fresh, prev_sampled, prev_predict, remap, seqlens, cu, out, max_seqlen_q=max_seqlen_q)
            return out

        return run

    return [
        ("n8_nextn5_mixed", True, make(8, 6, 8, True)),
        ("n8_nextn5_all_fresh", True, make(8, 6, 8, False)),
        ("n1_nextn1_mixed", True, make(1, 2, 3, True)),
        ("n64_nextn5_mixed", True, make(64, 6, 64, True)),
        ("n5_nextn2_full_len", True, make(5, 3, 5, True)),
    ]


def _cases_gather_engram() -> list[tuple[str, bool, Callable[[], Any]]]:
    """Engram n-gram windows, including the image-sentinel padding rule."""

    def make(n_seq: int, max_seqlen_q: int, max_ngram: int, vocab: int, with_image: bool):
        def run() -> Any:
            torch.manual_seed(0)
            nextn = max_seqlen_q - 1
            window = nextn + max_ngram - 1
            pad_token = 7
            sentinel = -1
            win = torch.randint(0, 50, (n_seq, window), device="cuda", dtype=torch.int32)
            if with_image:
                win[:, window // 2] = sentinel
            fresh = torch.randint(0, 50, (n_seq, max_seqlen_q), device="cuda", dtype=torch.int32)
            if with_image:
                fresh[:, -1] = sentinel
            prev_sampled = torch.randint(0, vocab, (n_seq,), device="cuda", dtype=torch.int32)
            prev_predict = torch.randint(0, vocab, (n_seq, nextn), device="cuda", dtype=torch.int32)
            prev_accepted = torch.randint(0, nextn + 1, (n_seq,), device="cuda", dtype=torch.int32)
            remap = torch.tensor([-1 if i % 2 == 0 else i for i in range(n_seq)], device="cuda", dtype=torch.int32)
            lookup = torch.arange(vocab, device="cuda", dtype=torch.int32) % 13
            seqlens = torch.randint(1, max_seqlen_q + 1, (n_seq,), device="cuda", dtype=torch.int32)
            cu = torch.cumsum(seqlens, 0, dtype=torch.int32)
            out = torch.full((int(cu[-1].item()), max_ngram), -7, device="cuda", dtype=torch.int32)
            gather_engram_tokens_zos(
                win,
                fresh,
                prev_sampled,
                prev_predict,
                prev_accepted,
                remap,
                lookup,
                seqlens,
                cu,
                out,
                max_seqlen_q=max_seqlen_q,
                max_ngram_size=max_ngram,
                engram_compressed_pad_token=pad_token,
                image_sentinel=sentinel,
            )
            return out

        return run

    return [
        ("n8_nextn5_ngram3", True, make(8, 6, 3, 64, False)),
        ("n8_nextn5_ngram1", True, make(8, 6, 1, 64, False)),
        ("n8_nextn5_ngram3_image", True, make(8, 6, 3, 64, True)),
        ("n1_nextn1_ngram2", True, make(1, 2, 2, 64, False)),
    ]


def _cases_swa_index() -> list[tuple[str, bool, Callable[[], Any]]]:
    """Sliding-window index expansion for the widths the engine uses."""

    def make(n: int, width: int, overflow: bool):
        def run() -> Any:
            torch.manual_seed(0)
            pivots = torch.zeros((n, 4), device="cuda", dtype=torch.int32)
            pivots[:, 0] = torch.randint(0, 4096, (n,), device="cuda", dtype=torch.int32)
            if overflow:
                # len1 + len2 can exceed width; the third branch must still win past len1 + len2.
                pivots[:, 1] = torch.randint(0, width + 1, (n,), device="cuda", dtype=torch.int32)
                pivots[:, 2] = torch.randint(0, 4096, (n,), device="cuda", dtype=torch.int32)
                pivots[:, 3] = torch.randint(0, width + 1, (n,), device="cuda", dtype=torch.int32)
            else:
                total = torch.randint(0, width + 1, (n,), device="cuda", dtype=torch.int32)
                pivots[:, 1] = total // 2
                pivots[:, 2] = torch.randint(0, 4096, (n,), device="cuda", dtype=torch.int32)
                pivots[:, 3] = total - total // 2
            return expand_swa_index(pivots, width)

        return run

    return [
        ("n32_w512", True, make(32, 512, False)),
        ("n32_w513_overflow", True, make(32, 513, True)),
        ("n1_w1", True, make(1, 1, False)),
        ("n37_w2048_overflow", True, make(37, 2048, True)),
        ("n8_w64", True, make(8, 64, False)),
    ]


def _cases_kv_index_cat() -> list[tuple[str, bool, Callable[[], Any]]]:
    """Window + compressed index concatenation with the -1 padding preserved."""

    def make(total_q: int, window: int, topk: int):
        def run() -> Any:
            torch.manual_seed(0)
            kv_index = torch.randint(-1, 4000, (total_q, window), device="cuda", dtype=torch.int32)
            kv_ca = torch.randint(-1, 300, (total_q, topk), device="cuda", dtype=torch.int32)
            return kv_index_cat(kv_index, kv_ca, 512, 2000)

        return run

    return [
        ("q8_w512_k128", True, make(8, 512, 128)),
        ("q1_w64_k32", True, make(1, 64, 32)),
        ("q5_w2048_k1024", True, make(5, 2048, 1024)),
    ]


def _cases_convert_kv_index() -> list[tuple[str, bool, Callable[[], Any]]]:
    """Logical to physical index translation through a page table."""

    def make(b: int, s: int, k: int, block_size: int, max_blocks: int):
        def run() -> Any:
            torch.manual_seed(0)
            topk = torch.randint(0, max_blocks * block_size, (b, s, k), device="cuda", dtype=torch.int32)
            ctx = torch.randint(0, k + 1, (b, s), device="cuda", dtype=torch.int32)
            table = torch.randint(0, 50, (b, max_blocks), device="cuda", dtype=torch.int32)
            convert_kv_index(topk, block_size, ctx, table)
            return topk

        return run

    return [
        ("b2_s4_k64_bs64", True, make(2, 4, 64, 64, 32)),
        ("b1_s1_k128_bs128", True, make(1, 1, 128, 128, 16)),
        ("b3_s2_k16_bs8", True, make(3, 2, 16, 8, 64)),
    ]


def _cases_load_kv_index() -> list[tuple[str, bool, Callable[[], Any]]]:
    """Prefill (contiguous) and decode (paged) index expansion."""

    def make_prefill(n: int, length: int):
        def run() -> Any:
            torch.manual_seed(0)
            off_len = torch.zeros((n, 2), device="cuda", dtype=torch.int32)
            off_len[:, 0] = torch.randint(0, 500, (n,), device="cuda", dtype=torch.int32)
            off_len[:, 1] = torch.randint(0, length + 1, (n,), device="cuda", dtype=torch.int32)
            return get_kv_ca_index_prefill(off_len, length)

        return run

    def make_decode(n: int, length: int, tpb: int, num_blocks: int):
        def run() -> Any:
            torch.manual_seed(0)
            off_len = torch.zeros((n, 2), device="cuda", dtype=torch.int32)
            off_len[:, 0] = torch.randint(0, 500, (n,), device="cuda", dtype=torch.int32)
            off_len[:, 1] = torch.randint(0, length + 1, (n,), device="cuda", dtype=torch.int32)
            table = torch.randint(0, num_blocks, (num_blocks,), device="cuda", dtype=torch.int32)
            return get_kv_ca_index_decode(off_len, length, tpb, table)

        return run

    return [
        ("prefill_n8_len512", True, make_prefill(8, 512)),
        ("prefill_n1_len1024", True, make_prefill(1, 1024)),
        ("prefill_n33_len100", True, make_prefill(33, 100)),
        ("decode_n8_len512_tpb64", True, make_decode(8, 512, 64, 128)),
        ("decode_n5_len128_tpb128", True, make_decode(5, 128, 128, 64)),
    ]


def _cases_store_mtp_info() -> list[tuple[str, bool, Callable[[], Any]]]:
    """Draft logprob pooling, including the slot-0 no-write path."""

    def make(req_cnt: int, nextn: int, vocab: int, has_index: bool, has_confidence: bool, zero_slots: bool):
        def run() -> Any:
            torch.manual_seed(0)
            # The pool is addressed as slot * nextn + step and slots run 0..req_cnt, so it needs
            # room for every slot including the reserved one.
            pool_rows = (req_cnt + 1) * nextn
            write_slots = torch.arange(1, req_cnt + 1, device="cuda", dtype=torch.int32)
            if zero_slots:
                write_slots[0] = 0
            token_ids = torch.randint(0, vocab, (req_cnt, nextn), device="cuda", dtype=torch.int32)
            logprobs = torch.randn((req_cnt, nextn, vocab), device="cuda", dtype=torch.float32)
            logprobs_index = torch.randint(0, vocab, (req_cnt, nextn, vocab), device="cuda", dtype=torch.int32)
            confidence = torch.rand((req_cnt, nextn), device="cuda", dtype=torch.float32)
            token_ids_pool = torch.full((pool_rows,), -7, device="cuda", dtype=torch.int32)
            logprobs_pool = torch.full((pool_rows, vocab), -7.0, device="cuda", dtype=torch.float32)
            logprobs_index_pool = torch.full((pool_rows, vocab), -7, device="cuda", dtype=torch.int32)
            confidence_pool = torch.full((pool_rows,), -7.0, device="cuda", dtype=torch.float32)
            store_mtp_info(
                write_slots,
                token_ids,
                logprobs,
                logprobs_index,
                confidence,
                token_ids_pool,
                logprobs_pool,
                logprobs_index_pool,
                confidence_pool,
            )
            return (
                token_ids_pool.clone(),
                logprobs_pool.clone(),
                logprobs_index_pool.clone(),
                confidence_pool.clone(),
            )

        return run

    return [
        ("n4_nextn5_v1k_full", True, make(4, 5, 1024, True, True, False)),
        ("n4_nextn5_v1k_no_index", True, make(4, 5, 1024, False, True, False)),
        ("n4_nextn5_v1k_no_conf", True, make(4, 5, 1024, True, False, False)),
        ("n4_nextn5_v1k_slot0", True, make(4, 5, 1024, True, True, True)),
        ("n2_nextn1_v129280", True, make(2, 1, 129280, True, True, False)),
    ]


def _cases_norm() -> list[tuple[str, bool, Callable[[], Any]]]:
    """RMSNorm in the forms the engine uses: bf16 out, fp8 with fp32 scales, and both at once."""

    def make(rows: int, d: int, with_residual: bool, as_fp8: bool, return_bf16: bool, round_sf: bool):
        def run() -> Any:
            torch.manual_seed(0)
            x = torch.randn((rows, d), device="cuda", dtype=torch.bfloat16)
            w = torch.randn((d,), device="cuda", dtype=torch.bfloat16)
            residual = torch.randn((rows, d), device="cuda", dtype=torch.bfloat16) if with_residual else None
            kwargs: dict[str, Any] = {"eps": 1e-20, "out_scale": 1.5}
            if as_fp8:
                kwargs.update(
                    fmt="e4m3",
                    num_per_channels=128,
                    round_sf=round_sf,
                    use_tma_aligned_col_major_sf=True,
                )
            return norm(x, w, residual=residual, return_bf16=return_bf16, **kwargs)

        return run

    return [
        ("r128_d5120_bf16", False, make(128, 5120, False, False, False, False)),
        ("r128_d5120_bf16_residual", False, make(128, 5120, True, False, False, False)),
        ("r64_d5120_fp8_round", False, make(64, 5120, False, True, False, True)),
        ("r64_d5120_fp8_no_round", False, make(64, 5120, False, True, False, False)),
        ("r64_d5120_fp8_round_bf16", False, make(64, 5120, False, True, True, True)),
        ("r33_d5120_fp8_round_bf16", False, make(33, 5120, False, True, True, True)),
        ("r16_d7168_fp8_round", False, make(16, 7168, False, True, False, True)),
        ("r8_d512_fp8_round", False, make(8, 512, False, True, False, True)),
    ]


def _cases_map_lookup() -> list[tuple[str, bool, Callable[[], Any]]]:
    """Both lookup modes, including out-of-range and missing keys."""

    def make_index(rows: int, width: int, dtype: torch.dtype):
        def run() -> Any:
            torch.manual_seed(0)
            value = (
                torch.randn((rows, width), device="cuda", dtype=dtype)
                if dtype.is_floating_point
                else torch.randint(0, 1000, (rows, width), device="cuda", dtype=dtype)
            )
            query = torch.randint(-width, 2 * width, (rows, 1), device="cuda", dtype=torch.int32)
            return map_lookup(None, value, query, mode="by_index")

        return run

    def make_key(rows: int, width: int):
        def run() -> Any:
            torch.manual_seed(0)
            # Distinct keys per row: that is what the indexer produces, and it is also the only
            # case where the kernel (last match wins) and the Python reference (max) agree.
            key = torch.randperm(width, device="cuda").to(torch.int32).unsqueeze(0).repeat(rows, 1)
            value = torch.randn((rows, width), device="cuda", dtype=torch.float32)
            # Half the queries exist in the row, half do not.
            present = key.gather(-1, torch.randint(0, width, (rows, 1), device="cuda").long())
            missing = torch.full((rows, 1), 10_000, device="cuda", dtype=torch.int32)
            keep = torch.randint(0, 2, (rows, 1), device="cuda") == 1
            query = torch.where(keep, present, missing)
            return map_lookup(key, value, query, mode="by_key")

        return run

    return [
        ("index_f32", True, make_index(256, 4096, torch.float32)),
        ("index_i32", True, make_index(256, 4096, torch.int32)),
        ("key_f32", False, make_key(128, 4096)),
    ]


def _cases_expand_sampling_params() -> list[tuple[str, bool, Callable[[], Any]]]:
    """Per-row parameter expansion for the draft rows."""

    def make(req_cnt: int, total_q: int):
        def run() -> Any:
            torch.manual_seed(0)
            temps = torch.rand((req_cnt,), device="cuda", dtype=torch.float32)
            top_ps = torch.rand((req_cnt,), device="cuda", dtype=torch.float32)
            top_ks = torch.randint(0, 4096, (req_cnt,), device="cuda", dtype=torch.int32)
            positions = torch.randint(0, 1000, (req_cnt,), device="cuda", dtype=torch.int32)
            seeds = torch.randint(0, 2**62, (req_cnt,), device="cuda", dtype=torch.uint64)
            row_to_req = torch.randint(0, req_cnt, (total_q,), device="cuda", dtype=torch.int32)
            local_step = torch.randint(0, 6, (total_q,), device="cuda", dtype=torch.int32)
            outs = [
                torch.empty((total_q,), device="cuda", dtype=torch.float32),
                torch.empty((total_q,), device="cuda", dtype=torch.float32),
                torch.empty((total_q,), device="cuda", dtype=torch.int32),
                torch.empty((total_q,), device="cuda", dtype=torch.int32),
                torch.empty((total_q,), device="cuda", dtype=torch.uint64),
            ]
            kernel = ops.expand_sampling_params_tl()
            kernel(temps, top_ps, top_ks, positions, seeds, row_to_req, local_step, *outs)
            return tuple(outs)

        return run

    return [("n4_q32", True, make(4, 32)), ("n1_q5", True, make(1, 5)), ("n8_q100", True, make(8, 100))]


def _cases_prepare_verification() -> list[tuple[str, bool, Callable[[], Any]]]:
    """Draft-row metadata: request index, local step and the verify flag."""

    def make(req_cnt: int, max_len: int, build_steps: int):
        def run() -> Any:
            torch.manual_seed(0)
            seqlen_q = torch.randint(1, max_len + 1, (req_cnt,), device="cuda", dtype=torch.int32)
            cu = torch.cumsum(seqlen_q, 0, dtype=torch.int32)
            total_q = int(cu[-1].item())
            ftp = torch.zeros((req_cnt,), device="cuda", dtype=torch.int32)
            ftp[1:] = cu[:-1]
            # The sentinel makes "wrote nothing" visible. `row_to_req` is a `torch.empty` buffer in
            # the engine, so a kernel that fills fewer than `total_q` rows leaves garbage behind and
            # every consumer then indexes with it -- and both backends would still agree with each
            # other here, which is why this check has to be explicit rather than a comparison.
            row_to_req = torch.full((total_q,), -7, device="cuda", dtype=torch.int32)
            local_step = torch.full((total_q,), -7, device="cuda", dtype=torch.int32)
            is_verify = torch.zeros((total_q,), device="cuda", dtype=torch.bool)
            # Driven through the factory: the interface wrapper allocates its own outputs, which
            # would leave nothing to compare.
            kernel = ops.prepare_verification_metadata_tl(build_steps)
            kernel(ftp, seqlen_q, row_to_req, local_step, is_verify)
            assert int((row_to_req == -7).sum()) == 0, "kernel left rows unwritten"
            assert int((local_step == -7).sum()) == 0, "kernel left rows unwritten"
            return row_to_req, local_step, is_verify

        return run

    return [
        ("n4_maxlen6", True, make(4, 6, 6)),
        ("n1_maxlen1", True, make(1, 1, 1)),
        ("n8_maxlen6", True, make(8, 6, 6)),
        # `num_steps=0` is what the engine builds when speculative decoding is off, so it must
        # still write every row.
        ("n8_zero_steps", True, make(8, 6, 0)),
        # A request longer than one tile (32 lanes) exercises the runtime tile loop.
        ("n2_maxlen130", True, make(2, 130, 130)),
        ("n64_maxlen6", True, make(64, 6, 6)),
    ]


def _cases_copy_block_tables() -> list[tuple[str, bool, Callable[[], Any]]]:
    """Compact 1D block table expanded to one row per token."""

    def make(n_seq: int, max_seqlen_q: int, max_blocks: int):
        def run() -> Any:
            torch.manual_seed(0)
            seqlens = torch.randint(1, max_seqlen_q + 1, (n_seq,), device="cuda", dtype=torch.int32)
            cu = torch.cumsum(seqlens, 0, dtype=torch.int32)
            total_q = int(cu[-1].item())
            counts = torch.randint(1, max_blocks + 1, (n_seq,), device="cuda", dtype=torch.int32)
            offsets = torch.zeros((n_seq + 1,), device="cuda", dtype=torch.int32)
            offsets[1:] = torch.cumsum(counts, 0, dtype=torch.int32)
            flat = torch.randint(0, 512, (int(offsets[-1].item()),), device="cuda", dtype=torch.int32)
            out = torch.full((total_q, max_blocks), -7, device="cuda", dtype=torch.int32)
            # No standalone interface function: this kernel is only reached from
            # `prepare_v4_decode_metadata`, so the case drives the factory directly. The harness
            # swaps `ops.copy_block_tables_tl` between backends, which is all that is needed.
            kernel = ops.copy_block_tables_tl(max_seqlen_q, max_blocks)
            kernel(seqlens, cu, flat, offsets, out)
            return out

        return run

    return [
        ("n4_q6_b64", True, make(4, 6, 64)),
        ("n1_q1_b32", True, make(1, 1, 32)),
        ("n8_q6_b1400_multi_chunk", True, make(8, 6, 1400)),
    ]


def _cases_store_kv_cache() -> list[tuple[str, bool, Callable[[], Any]]]:
    """Quantised cache writes: the whole cache buffer must match byte for byte."""

    def make(layout: Any, block_size: int, n: int, with_invalid: bool):
        def run() -> Any:
            torch.manual_seed(0)
            num_blocks = (n + block_size - 1) // block_size + 2
            kv = torch.randn((n, layout.num_elements), device="cuda", dtype=torch.bfloat16)
            slots = torch.arange(n, dtype=torch.int32, device="cuda") + block_size
            if with_invalid:
                slots[::7] = -1
            cache = torch.zeros((num_blocks, block_size, 1, layout.bytes_per_token), dtype=torch.uint8, device="cuda")
            store_kv_cache(kv, cache, slots, layout=layout)
            return cache.clone()

        return run

    return [
        ("main_fp8_b32", True, make(V41_FORMAT, 32, 100, False)),
        ("main_fp8_b64", True, make(V41_FORMAT, 64, 233, True)),
        ("main_fp8_b134", True, make(V41_FORMAT, 134, 500, False)),
        ("index_fp32sf_b32", True, make(FP8_FLOAT32_SF_INDEX_KV_FORMAT, 32, 64, False)),
        ("index_fp32sf_b64", True, make(FP8_FLOAT32_SF_INDEX_KV_FORMAT, 64, 200, True)),
    ]


def _cases_load_paged_kv_cache() -> list[tuple[str, bool, Callable[[], Any]]]:
    """Paged reads, including holes (src_block_id == -1) and partial segments.

    The cache is filled by `store_kv_cache` rather than with random bytes: random bytes produce
    fp8 NaN patterns and NaN scales, and those compare unequal to themselves, which would report a
    mismatch that says nothing about the load kernel.
    """

    def make(layout: Any, block_size: int, num_blocks: int, num_segments: int, with_hole: bool):
        torch.manual_seed(0)
        total = num_blocks * block_size
        kv = torch.randn((total, layout.num_elements), device="cuda", dtype=torch.bfloat16)
        cache = torch.zeros((num_blocks, block_size, 1, layout.bytes_per_token), dtype=torch.uint8, device="cuda")
        store_kv_cache(kv, cache, torch.arange(total, dtype=torch.int32, device="cuda"), layout=layout)

        segments = []
        dst = 0
        for i in range(num_segments):
            begin = i % block_size
            end = min(block_size, begin + 1 + (i * 7) % block_size)
            src = -1 if (with_hole and i == 1) else i % num_blocks
            segments.append([src, begin, end, dst])
            dst += end - begin
        seg = torch.tensor(segments, dtype=torch.int32, device="cuda")

        def run() -> Any:
            result = load_paged_kv_cache(cache, seg, dst, layout=layout)
            if result.bf16_data is not None:
                return result.bf16_data
            # Byte views for the index layout: both payload and scales are passthrough copies, so
            # exact bytes are the right comparison and NaN bit patterns are harmless.
            return result.quantized_data.view(torch.uint8), result.sf_data.view(torch.uint8)

        return run

    return [
        ("main_fp8_b32", True, make(V41_FORMAT, 32, 8, 6, False)),
        ("main_fp8_b64_holes", True, make(V41_FORMAT, 64, 8, 6, True)),
        ("index_fp32sf_b32", True, make(FP8_FLOAT32_SF_INDEX_KV_FORMAT, 32, 8, 6, False)),
        ("index_fp32sf_b64_holes", True, make(FP8_FLOAT32_SF_INDEX_KV_FORMAT, 64, 8, 6, True)),
    ]


def _cases_fused_transform() -> list[tuple[str, bool, Callable[[], Any]]]:
    """The norm -> RoPE -> fp8 pipeline, in every form the engine can ask for."""

    def make(
        n: int,
        num_heads: int,
        head_dim: int,
        rotary: int,
        *,
        norm: bool,
        rope: bool,
        cast: bool,
        cast_back: bool,
        is_neox: bool = False,
        conjugate: bool = False,
        npc: int = 32,
        sf_format: str = "ue8m0",
    ):
        def run() -> Any:
            torch.manual_seed(0)
            x = torch.randn((n, num_heads, head_dim), device="cuda", dtype=torch.bfloat16)
            norm_args = NormArgs(weight=torch.randn((head_dim,), device="cuda", dtype=torch.bfloat16), eps=1e-20) if norm else None
            rope_args = None
            if rope:
                rotary_dim = rotary
                half = rotary_dim // 2
                angles = torch.randn((4096, half), device="cuda", dtype=torch.float32).cumsum(-1) * 0.01
                cache = torch.cat([angles.cos(), angles.sin()], dim=-1)
                positions = torch.randint(0, 4096, (n,), device="cuda", dtype=torch.int32)
                rope_args = RopeArgs(
                    positions=positions,
                    cos_sin_cache=cache,
                    rotary_dim=rotary_dim,
                    is_neox=is_neox,
                    offset=head_dim - rotary_dim,
                    conjugate=conjugate,
                )
            cast_args = CastArgs("e4m3", npc, round_sf=True, sf_format=sf_format, cast_back=cast_back) if cast else None
            return fused_norm_rope_hadamard_cast_store_kvcache(x, norm_args, rope_args, None, cast_args, None)

        return run

    def make_index_cache(n: int, num_heads: int, head_dim: int, rotary: int, num_blocks: int):
        """Index-K path: fp32 scales, per-128 groups, into the index KV cache."""

        def run() -> Any:
            torch.manual_seed(0)
            layout = FP8_FLOAT32_SF_INDEX_KV_FORMAT
            block_size = 64
            x = torch.randn((n, num_heads, head_dim), device="cuda", dtype=torch.bfloat16)
            half = rotary // 2
            angles = torch.randn((4096, half), device="cuda", dtype=torch.float32).cumsum(-1) * 0.01
            cos_sin = torch.cat([angles.cos(), angles.sin()], dim=-1)
            positions = torch.randint(0, 4096, (n,), device="cuda", dtype=torch.int32)
            kv_cache = torch.zeros((num_blocks, block_size, 1, layout.bytes_per_token), dtype=torch.uint8, device="cuda")
            # Distinct slots: two tokens landing on the same slot would race, and the winner
            # would depend on CTA scheduling rather than on the kernel under test.
            slots = torch.randperm(num_blocks * block_size, device="cuda")[:n].to(torch.int32)
            fused_norm_rope_hadamard_cast_store_kvcache(
                x,
                None,
                RopeArgs(positions, cos_sin, rotary, False, offset=head_dim - rotary),
                None,
                CastArgs.from_cache_format(layout),
                KVCacheOutArgs(kv_cache, slots, layout=layout),
            )
            return kv_cache

        return run

    def make_cache(n: int, block_size: int, num_blocks: int, with_bf16: bool, with_rope: bool = True):
        """The main-KV path: quantise and write straight into the paged cache."""

        def run() -> Any:
            torch.manual_seed(0)
            head_dim, rotary = 512, 64
            layout = V41_FORMAT
            x = torch.randn((n, 1, head_dim), device="cuda", dtype=torch.bfloat16)
            weight = torch.randn((head_dim,), device="cuda", dtype=torch.bfloat16)
            half = rotary // 2
            angles = torch.randn((4096, half), device="cuda", dtype=torch.float32).cumsum(-1) * 0.01
            cache_cs = torch.cat([angles.cos(), angles.sin()], dim=-1)
            positions = torch.randint(0, 4096, (n,), device="cuda", dtype=torch.int32)
            rope_args = RopeArgs(positions, cache_cs, rotary, False, offset=head_dim - rotary) if with_rope else None
            kv_cache = torch.zeros((num_blocks, block_size, 1, layout.bytes_per_token), dtype=torch.uint8, device="cuda")
            # Distinct slots: two tokens landing on the same slot would race, and the winner
            # would depend on CTA scheduling rather than on the kernel under test.
            slots = torch.randperm(num_blocks * block_size, device="cuda")[:n].to(torch.int32)
            out_bf16 = torch.empty((n, 1, head_dim), device="cuda", dtype=torch.bfloat16) if with_bf16 else None
            fused_norm_rope_hadamard_cast_store_kvcache(
                x,
                NormArgs(weight, 1e-20),
                rope_args,
                None,
                CastArgs.from_cache_format(layout),
                KVCacheOutArgs(kv_cache, slots, out=out_bf16, out_stage="norm", layout=layout),
            )
            return (kv_cache, out_bf16) if with_bf16 else kv_cache

        return run

    return [
        # The configuration the engine builds, from DEEPINFER_TRACE_KERNEL_ARGS on a full run.
        ("engine_cfg_n1_h512", False, make(16, 1, 512, 64, norm=True, rope=True, cast=True, cast_back=True)),
        # The engine's other call site: norm + rope + quantise into the paged KV cache, with the
        # post-norm row also returned for the indexer.
        ("engine_cache", False, make_cache(16, 134, 4, True)),
        ("engine_cache_no_bf16", False, make_cache(16, 134, 4, False)),
        ("engine_cache_b32", False, make_cache(9, 32, 8, True)),
        ("engine_cache_no_rope", False, make_cache(8, 134, 2, True, with_rope=False)),
        ("engine_cfg_n128", False, make(128, 1, 512, 64, norm=True, rope=True, cast=True, cast_back=True)),
        ("engine_cfg_n1_token", False, make(1, 1, 512, 64, norm=True, rope=True, cast=True, cast_back=True)),
        ("neox", False, make(16, 1, 512, 64, norm=True, rope=True, cast=True, cast_back=True, is_neox=True)),
        ("conjugate", False, make(16, 1, 512, 64, norm=True, rope=True, cast=True, cast_back=True, conjugate=True)),
        ("no_rope", False, make(64, 1, 512, 64, norm=True, rope=False, cast=True, cast_back=True)),
        ("no_norm", False, make(64, 1, 512, 64, norm=False, rope=True, cast=True, cast_back=True)),
        ("bf16_only", False, make(64, 1, 512, 64, norm=True, rope=True, cast=False, cast_back=False)),
        ("return_quant", False, make(64, 1, 512, 64, norm=True, rope=True, cast=True, cast_back=False)),
        ("multi_head_h128", False, make(32, 4, 128, 128, norm=True, rope=True, cast=True, cast_back=True)),
        ("h128_rope64", False, make(32, 2, 128, 64, norm=True, rope=True, cast=True, cast_back=True)),
        # The indexer's Q quantisation: no norm, per-128 groups, fp32 scales, quantised output.
        ("index_q_fp32sf", False, make(64, 4, 128, 64, norm=False, rope=True, cast=True, cast_back=False, npc=128, sf_format="fp32")),
        (
            "index_q_fp32sf_many",
            False,
            make(256, 32, 128, 64, norm=False, rope=True, cast=True, cast_back=False, npc=128, sf_format="fp32"),
        ),
        # Same, but written straight into the index KV cache.
        # The index KV cache holds a single KV head (the indexer is MQA), which the baseline
        # asserts too.
        ("index_k_cache", False, make_index_cache(48, 1, 128, 64, 16)),
    ]


def _cases_prepare_v4_decode_metadata() -> list[tuple[str, bool, Callable[[], Any]]]:
    """Every output of the decode-metadata kernel, including the ratio-bucketed ones."""

    def make(
        n_seq: int,
        max_seqlen_q: int,
        max_num_blocks: int,
        window_size: int,
        swa_block_size: int,
        block_size: int,
        ca_ratios: tuple[int, ...],
        dense_ratios: tuple[int, ...],
        sparse_ratios: tuple[int, ...],
        placeholders: bool,
    ):
        def run() -> Any:
            torch.manual_seed(0)
            seqlens_q = torch.randint(1, max_seqlen_q + 1, (n_seq,), device="cuda", dtype=torch.int32)
            seqlens_k = seqlens_q + torch.randint(0, 40, (n_seq,), device="cuda", dtype=torch.int32)
            cu = torch.cumsum(seqlens_q, 0, dtype=torch.int32)
            total_q = int(cu[-1].item())
            per_seq_blocks = torch.randint(1, max_num_blocks + 1, (n_seq,), device="cuda", dtype=torch.int32)
            offsets = torch.zeros((n_seq + 1,), device="cuda", dtype=torch.int32)
            offsets[1:] = torch.cumsum(per_seq_blocks, 0, dtype=torch.int32)
            total_blocks = int(offsets[-1].item())
            block_tables_in = torch.randint(0, 40, (total_blocks,), device="cuda", dtype=torch.int32)
            state_block_idx = torch.randint(0, 8, (n_seq,), device="cuda", dtype=torch.int32)
            is_placeholder = torch.zeros((n_seq,), device="cuda", dtype=torch.bool)
            if placeholders:
                is_placeholder[::2] = True

            n_ca = max(len(ca_ratios), 1)
            n_dense = max(len(dense_ratios), 1)
            n_sparse = max(len(sparse_ratios), 1)
            i32 = {"device": "cuda", "dtype": torch.int32}
            outs = [
                torch.full((total_q,), -7, **i32),  # positions
                torch.zeros((total_q,), device="cuda", dtype=torch.bool),  # valid_token_mask
                torch.zeros((total_q,), device="cuda", dtype=torch.bool),  # gate_force_random
                torch.ones((total_q,), device="cuda", dtype=torch.bool),  # fix_routing_mask
                torch.full((total_q,), -7, **i32),  # context_lens
                torch.full((n_sparse, total_q), -7, **i32),  # klen
                torch.full((n_dense, total_q, 2), -7, **i32),  # kv_ca_index_offset_len
                torch.full((total_q, 2), -7, **i32),  # kv_swa_index_offset_len
                torch.full((total_q,), -7, **i32),  # slot_mapping_store_swa
                torch.full((total_q,), -7, **i32),  # query_indices
                torch.full((total_q,), -7, **i32),  # state_block_idx_out
                torch.full((total_q, max_num_blocks), -7, **i32),  # block_tables_out
                torch.full((n_ca, total_q), -7, **i32),  # positions_c
                torch.full((n_ca, total_q), -7, **i32),  # slot_mapping_store_ca
            ]
            prepare_v4_decode_metadata(
                seqlens_q,
                seqlens_k,
                cu,
                block_tables_in,
                offsets,
                state_block_idx,
                is_placeholder,
                *outs,
                max_seqlen_q=max_seqlen_q,
                max_num_blocks=max_num_blocks,
                window_size=window_size,
                swa_block_size=swa_block_size,
                block_size=block_size,
                ca_ratios=ca_ratios,
                dense_ratios=dense_ratios,
                sparse_ratios=sparse_ratios,
                force_load_balance=False,
            )
            # Nothing should be left at the sentinel: a partially written output is a bug even when
            # both backends agree on it.
            for name, tensor in (
                ("positions", outs[0]),
                ("context_lens", outs[4]),
                ("kv_swa_index_offset_len", outs[7]),
                ("slot_mapping_store_swa", outs[8]),
                ("query_indices", outs[9]),
                ("positions_c", outs[12]),
                ("slot_mapping_store_ca", outs[13]),
            ):
                assert int((tensor == -7).sum()) == 0, f"{name} left unwritten rows"
            return tuple(outs)

        return run

    return [
        # The engine's own configuration: window 128, SWA block 134, ratios {1, 2} all compressed.
        ("engine_cfg", False, make(8, 6, 64, 128, 134, 128, (1, 2), (), (1, 2), False)),
        ("engine_cfg_placeholders", False, make(8, 6, 64, 128, 134, 128, (1, 2), (), (1, 2), True)),
        ("with_dense", False, make(6, 6, 32, 64, 128, 128, (1, 2, 4), (1, 2), (4,), False)),
        ("single_seq", False, make(1, 1, 8, 32, 32, 32, (1,), (), (1,), False)),
        ("prefill_len8", False, make(5, 8, 16, 512, 512, 64, (1, 2, 4), (), (1, 2, 4), False)),
    ]


def _cases_compress_prefill() -> list[tuple[str, bool, Callable[[], Any]]]:
    """Block compression plus the ring-cache tail write, including a window that spans chunks."""

    def make(dim: int, ratio: int, n_seq: int, cache_size: int, slots: int, spans_chunks: bool):
        def run() -> Any:
            torch.manual_seed(0)
            qlens = torch.randint(3, 7, (n_seq,), device="cuda", dtype=torch.int32)
            # qstart > 0 means the first compression window reaches into the previous chunk, which
            # is the only way the cache-read branch is exercised.
            qstart = (
                torch.randint(1, 4, (n_seq,), device="cuda", dtype=torch.int32)
                if spans_chunks
                else torch.zeros((n_seq,), device="cuda", dtype=torch.int32)
            )
            klen = qstart + qlens
            # All four prefix sums are exclusive with a leading zero (`n_seq + 1` entries); the
            # kernel's ABI rejects the inclusive form.
            zero = torch.zeros((1,), device="cuda", dtype=torch.int32)
            cu_q = torch.cat([zero, torch.cumsum(qlens, 0, dtype=torch.int32)])
            cu_k = torch.cat([zero, torch.cumsum(klen, 0, dtype=torch.int32)])
            total_q = int(cu_q[-1].item())

            new_blocks = (klen // ratio - qstart // ratio).tolist()
            reserved = [max(1, b) for b in new_blocks]
            mapping = torch.cat([torch.full((r,), i, device="cuda", dtype=torch.int32) for i, r in enumerate(reserved)])
            # Exclusive prefix sum with a leading zero, i.e. `n_seq + 1` entries: the kernel binds
            # its `num_seqs` symbol from this tensor, so a shorter one fails the ABI check.
            cu_c = torch.zeros((n_seq + 1,), device="cuda", dtype=torch.int32)
            acc = 0
            for i, r in enumerate(reserved):
                cu_c[i] = acc
                acc += r
            cu_c[n_seq] = acc
            # Distinct slots per sequence, as the engine assigns them. Sharing a slot would let one
            # sequence's tail copy overwrite the other's cache, which makes the result depend on
            # kernel interleaving rather than on the code under test.
            assert n_seq <= slots
            state_block_idx = torch.randperm(slots, device="cuda")[:n_seq].to(torch.int32)

            score = torch.randn((total_q, dim), device="cuda", dtype=torch.float32)
            latent = torch.randn((total_q, dim), device="cuda", dtype=torch.float32)
            state_cache = torch.randn((slots, cache_size, 2, dim), device="cuda", dtype=torch.float32)

            kv = compress_and_update_state_prefill(
                score,
                latent,
                state_cache,
                state_block_idx,
                mapping,
                cu_q,
                cu_k,
                cu_c,
                ratio,
                None,
            )
            return kv, state_cache

        return run

    return [
        ("engine_r2", False, make(512, 2, 3, 16, 4, True)),
        ("engine_r2_no_prev", False, make(512, 2, 3, 16, 4, False)),
        ("r1", False, make(512, 1, 3, 16, 4, True)),
        ("single_seq", False, make(512, 2, 1, 8, 1, True)),
        # Tightest legal ring: at least one chunk's worth of tokens, so the tail copy cannot wrap
        # onto a slot the compression step is reading. The interface asserts this relation too.
        ("tight_cache", False, make(512, 2, 2, 6, 2, True)),
    ]


def _cases_compress_decode() -> list[tuple[str, bool, Callable[[], Any]]]:
    """Per-token cache update and the block-closing compression."""

    def make(dim: int, ratio: int, total_q: int, max_seqlen_q: int, slots: int, cache_size: int):
        def run() -> Any:
            torch.manual_seed(0)
            score = torch.randn((total_q, dim), device="cuda", dtype=torch.float32)
            latent = torch.randn((total_q, dim), device="cuda", dtype=torch.float32)
            state_cache = torch.randn((slots, cache_size, 2, dim), device="cuda", dtype=torch.float32)
            # One distinct state slot per token: the engine gives every sequence its own slot, and
            # sharing one would make the cache writes overlap and the result order-dependent.
            assert total_q <= slots
            state_block_idx = torch.randperm(slots, device="cuda")[:total_q].to(torch.int32)
            context_lens = torch.randint(1, 40, (total_q,), device="cuda", dtype=torch.int32)
            kv = compress_and_update_state_decode(
                score,
                latent,
                state_cache,
                state_block_idx,
                context_lens,
                max_seqlen_q,
                ratio,
            )
            return kv, state_cache

        return run

    return [
        ("engine_r2", False, make(512, 2, 8, 6, 8, 8)),
        ("r1", False, make(512, 1, 8, 6, 8, 8)),
        ("single_token", False, make(512, 2, 1, 1, 1, 4)),
        ("many_slots", False, make(512, 2, 16, 6, 16, 16)),
    ]


CASES: dict[str, Callable[[], list[tuple[str, bool, Callable[[], Any]]]]] = {
    "store_kv_cache_tl": _cases_store_kv_cache,
    "load_paged_kv_cache_tl": _cases_load_paged_kv_cache,
    "norm_tl": _cases_norm,
    "map_lookup_tl": _cases_map_lookup,
    "expand_sampling_params_tl": _cases_expand_sampling_params,
    "prepare_verification_metadata_tl": _cases_prepare_verification,
    "fused_norm_rope_hadamard_cast_store_kvcache_tl": _cases_fused_transform,
    "prepare_v4_decode_metadata_tl": _cases_prepare_v4_decode_metadata,
    "compress_and_update_state_prefill_tl": _cases_compress_prefill,
    "compress_and_update_state_decode_tl": _cases_compress_decode,
    "copy_block_tables_tl": _cases_copy_block_tables,
    "gather_tl": _cases_gather,
    "gather_input_ids_zos_tl": _cases_gather_input_ids,
    "gather_engram_tokens_zos_tl": _cases_gather_engram,
    "expand_swa_index_tl": _cases_swa_index,
    "kv_index_cat_tl": _cases_kv_index_cat,
    "convert_kv_index_tl": _cases_convert_kv_index,
    "load_kv_index_tl": _cases_load_kv_index,
    "store_mtp_info_tl": _cases_store_mtp_info,
}