"""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.flash_mla_utils import convert_kv_index
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.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, total_q: int, max_len: int):
        def run() -> Any:
            torch.manual_seed(0)
            seqlen_q = torch.randint(1, max_len + 1, (req_cnt,), device="cuda", dtype=torch.int32)
            ftp = torch.zeros((req_cnt,), device="cuda", dtype=torch.int32)
            ftp[1:] = torch.cumsum(seqlen_q, 0, dtype=torch.int32)[:-1]
            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(max_len)
            kernel(ftp, seqlen_q, row_to_req, local_step, is_verify)
            return row_to_req, local_step, is_verify

        return run

    return [("n4_q24", True, make(4, 24, 6)), ("n1_q1", True, make(1, 1, 1)), ("n8_q48", True, make(8, 48, 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)),
    ]


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,
    "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,
}