"""
Benchmark for Top-K operations including:
- top_k: Basic radix-based top-k selection
- top_k_page_table_transform: Fused top-k + page table gather (for sparse attention)
- top_k_ragged_transform: Fused top-k + offset addition (for sparse attention)
- varlen: Variable-length segment transforms that model production sparse-attention
  workloads, where each row selects top-k over a per-row valid window whose length is
  drawn from a realistic decode/prefill distribution (rather than a fixed seq_len)

Optional comparison with SGLang's sgl_kernel implementation.
"""

import argparse
import math
import os
from contextlib import contextmanager
from dataclasses import dataclass

import numpy as np
import torch

import flashinfer
from flashinfer.topk import TopKTieBreak
from flashinfer.testing.utils import bench_gpu_time
from flashinfer.utils import get_compute_capability


def set_topk_algo(algo: str):
    """Set environment variable to force specific topk algorithm."""
    if algo == "auto":
        os.environ.pop("FLASHINFER_TOPK_ALGO", None)
    else:
        os.environ["FLASHINFER_TOPK_ALGO"] = algo


def classify_benchmark_runtime_error(exc: RuntimeError) -> str | None:
    message = str(exc).lower()
    if "out of memory" in message:
        return "OOM"
    if "invalid" in message or "operation not supported" in message:
        return "UNSUPPORTED"
    return None


# Try to import sgl_kernel for comparison
try:
    import sgl_kernel

    HAS_SGL_KERNEL = True
except ImportError:
    HAS_SGL_KERNEL = False


@contextmanager
def torch_deterministic_algorithms(enabled: bool):
    """Temporarily set PyTorch deterministic algorithm mode."""
    previous = torch.are_deterministic_algorithms_enabled()
    if previous != enabled:
        torch.use_deterministic_algorithms(enabled)
    try:
        yield
    finally:
        if torch.are_deterministic_algorithms_enabled() != previous:
            torch.use_deterministic_algorithms(previous)


def bench_median_ms(fn) -> float:
    measurements = bench_gpu_time(
        fn,
        enable_cupti=True,
        dry_run_iters=10,
        repeat_iters=100,
        use_cuda_graph=False,
    )
    return float(np.median(measurements))


def bench_flashinfer_modes(
    run_flashinfer, deterministic: bool
) -> tuple[float, float | None]:
    selected_ms = bench_median_ms(lambda: run_flashinfer(deterministic))
    nondeterministic_ms = (
        bench_median_ms(lambda: run_flashinfer(False)) if deterministic else None
    )
    return selected_ms, nondeterministic_ms


TIE_BREAK_VARIANTS: tuple[tuple[str, TopKTieBreak], ...] = (
    ("small", TopKTieBreak.SMALL),
    ("large", TopKTieBreak.LARGE),
)


def bench_tie_break_variants(
    run_flashinfer_with_tie_break, baseline_ms: float
) -> dict[str, float]:
    metrics: dict[str, float] = {}
    for suffix, tie_break in TIE_BREAK_VARIANTS:
        try:
            tie_ms = bench_median_ms(lambda: run_flashinfer_with_tie_break(tie_break))
            metrics[f"flashinfer_tie_{suffix}_us"] = tie_ms * 1e3
            metrics[f"tie_{suffix}_slowdown_vs_baseline"] = tie_ms / baseline_ms
        except RuntimeError as exc:
            error_label = classify_benchmark_runtime_error(exc)
            if error_label is None:
                raise
            metrics[f"flashinfer_tie_{suffix}_error"] = error_label
    return metrics


def append_tie_break_header(header: str, enabled: bool) -> str:
    if not enabled:
        return header
    return (
        header
        + f" {'FlashInfer(tie-small)':>21} {'TieSmallSlowdown':>17}"
        + f" {'FlashInfer(tie-large)':>21} {'TieLargeSlowdown':>17}"
    )


def append_tie_break_columns(line: str, result: dict, enabled: bool) -> str:
    if not enabled:
        return line

    def format_variant(suffix: str) -> str:
        us_key = f"flashinfer_tie_{suffix}_us"
        slowdown_key = f"tie_{suffix}_slowdown_vs_baseline"
        error_key = f"flashinfer_tie_{suffix}_error"
        if us_key in result:
            return f" {result[us_key]:>19.2f}us {result[slowdown_key]:>16.2f}x"
        error_label = result.get(error_key, "n/a")
        return f" {error_label:>21} {'n/a':>16}"

    return line + format_variant("small") + format_variant("large")


def bench_top_k_from_scores(
    scores: torch.Tensor,
    k: int,
    deterministic: bool = False,
    compare_tie_break: bool = False,
    compare_torch_deterministic: bool = False,
    compare_sglang: bool = False,
) -> dict:
    """Benchmark top-k on a pre-generated score tensor."""
    batch_size, seq_len = scores.shape

    set_topk_algo("default")
    fi_ms, fi_nondeterministic_ms = bench_flashinfer_modes(
        lambda deterministic_mode: flashinfer.top_k(
            scores,
            k,
            deterministic=deterministic_mode,
            tie_break=TopKTieBreak.NONE,
        ),
        deterministic,
    )

    result = {
        "batch_size": batch_size,
        "seq_len": seq_len,
        "k": k,
        "dtype": str(scores.dtype),
        "flashinfer_us": fi_ms * 1e3,
    }
    if fi_nondeterministic_ms is not None:
        result["flashinfer_nondeterministic_us"] = fi_nondeterministic_ms * 1e3
        result["deterministic_slowdown_vs_nondeterministic"] = (
            fi_ms / fi_nondeterministic_ms
        )

    with torch_deterministic_algorithms(deterministic):
        torch_ms = bench_median_ms(lambda: torch.topk(scores, k, dim=-1))
    result["torch_us"] = torch_ms * 1e3
    result["speedup_vs_torch"] = torch_ms / fi_ms
    if compare_tie_break:
        # Use the same non-deterministic baseline as DetSlowdown when present.
        baseline_ms = (
            fi_nondeterministic_ms if fi_nondeterministic_ms is not None else fi_ms
        )
        result.update(
            bench_tie_break_variants(
                lambda tie_break: flashinfer.top_k(
                    scores,
                    k,
                    deterministic=deterministic,
                    tie_break=tie_break,
                ),
                baseline_ms,
            )
        )

    if compare_torch_deterministic and not deterministic:
        with torch_deterministic_algorithms(True):
            torch_det_ms = bench_median_ms(lambda: torch.topk(scores, k, dim=-1))
        result["torch_deterministic_us"] = torch_det_ms * 1e3
        result["speedup_vs_torch_deterministic"] = torch_det_ms / fi_ms

    set_topk_algo("clusters")
    fast_topk_ms = bench_median_ms(lambda: flashinfer.top_k(scores, k))
    result["fast_topk_us"] = fast_topk_ms * 1e3
    result["speedup_vs_flashinfer"] = fi_ms / fast_topk_ms
    set_topk_algo("auto")

    # SGLang comparison (only supports k=2048 and float32)
    if (
        compare_sglang
        and HAS_SGL_KERNEL
        and k == 2048
        and scores.dtype == torch.float32
    ):
        lengths = torch.full((batch_size,), seq_len, dtype=torch.int32, device="cuda")
        sg_ms = bench_median_ms(
            lambda: sgl_kernel.fast_topk_v2(scores, lengths, k, row_starts=None)
        )
        result["sglang_us"] = sg_ms * 1e3
        result["speedup_vs_sglang"] = sg_ms / fi_ms

    return result


def generate_scores(
    batch_size: int,
    seq_len: int,
    k: int,
    dtype: torch.dtype,
    input_pattern: str,
) -> torch.Tensor:
    """Generate benchmark input scores with controllable tie patterns."""

    if input_pattern == "random":
        return torch.randn(batch_size, seq_len, device="cuda", dtype=dtype)

    if input_pattern in {"quantized_random", "relu_quantized"}:
        base = torch.randn(batch_size, seq_len, device="cuda", dtype=torch.float32)
        if input_pattern == "relu_quantized":
            base = torch.relu(base)
        scores = (torch.round(base * 32.0) / 32.0).to(dtype)
        return scores

    if input_pattern == "tie_heavy":
        pattern = (
            torch.arange(seq_len, device="cuda", dtype=torch.float32) % 64
        ) / 64.0
        return pattern.unsqueeze(0).expand(batch_size, -1).contiguous().to(dtype)

    if input_pattern == "pivot_tie":
        # Severe tie at pivot:
        # - majority entries are identical (1.0)
        # - a small tail region is strictly larger (2.0)
        # This creates truncation in == pivot region when k exceeds tail size.
        scores = torch.ones(batch_size, seq_len, device="cuda", dtype=dtype)
        gt_count = max(1, min(k // 4, seq_len // 8))
        scores[:, seq_len - gt_count :] = 2.0
        return scores

    raise ValueError(f"Unsupported input_pattern: {input_pattern}")


def generate_dsa_scores(
    batch_size: int,
    q_len: int,
    seq_len: int,
    dtype: torch.dtype,
    input_pattern: str,
    causal_chunk: bool,
) -> torch.Tensor:
    """Generate DeepSeek DSA-like indexer score workload.

    Source context:
    - DeepSeek-V3.2-Exp config uses index_topk=2048:
      https://huggingface.co/deepseek-ai/DeepSeek-V3.2-Exp/blob/main/inference/config_671B_v3.2.json
    - Indexer runs topk over index_score last dim:
      https://huggingface.co/deepseek-ai/DeepSeek-V3.2-Exp/blob/main/inference/model.py

    Returns a 2D tensor with shape (batch_size * q_len, seq_len).
    """
    rows = batch_size * q_len

    if input_pattern == "random":
        scores = torch.randn(rows, seq_len, device="cuda", dtype=dtype)
    elif input_pattern in {"quantized_random", "dsa_relu"}:
        base = torch.randn(rows, seq_len, device="cuda", dtype=torch.float32)
        if input_pattern == "dsa_relu":
            base = torch.relu(base)
        scores = (torch.round(base * 32.0) / 32.0).to(dtype)
    else:
        raise ValueError(f"Unsupported dsa input_pattern: {input_pattern}")

    if causal_chunk:
        # Simulate prefill chunk near end of long context:
        # each query row i can only attend [0, start_pos + i].
        start_pos = seq_len - q_len
        lengths = torch.arange(
            start_pos + 1,
            start_pos + q_len + 1,
            device="cuda",
            dtype=torch.int32,
        ).repeat(batch_size)
        col = torch.arange(seq_len, device="cuda", dtype=torch.int32).unsqueeze(0)
        invalid = col >= lengths.unsqueeze(1)
        neg_inf = -torch.inf if dtype == torch.float32 else torch.finfo(dtype).min
        scores = scores.masked_fill(invalid, neg_inf)

    return scores.contiguous()


@dataclass(frozen=True)
class DSATopKCase:
    name: str
    batch_size: int
    q_len: int
    seq_len: int
    causal_chunk: bool


@dataclass(frozen=True)
class TopKCase:
    name: str
    batch_size: int
    seq_len: int
    k: int


def build_top_k_cases(
    batch_sizes: list[int],
    seq_lens: list[int],
    k_values: list[int],
) -> list[TopKCase]:
    cases: list[TopKCase] = []

    for batch_size in batch_sizes:
        for seq_len in seq_lens:
            for k in k_values:
                if k <= seq_len:
                    cases.append(
                        TopKCase(
                            name=f"grid_b{batch_size}_l{seq_len}_k{k}",
                            batch_size=batch_size,
                            seq_len=seq_len,
                            k=k,
                        )
                    )

    # These deterministic large-batch/long-vocab cases are not covered by the
    # original grid but surfaced real correctness/performance differences.
    cases.extend(
        [
            TopKCase("stress_b2048_l131072_k1024", 2048, 131072, 1024),
            TopKCase("stress_b4096_l200000_k1024", 4096, 200000, 1024),
        ]
    )

    return cases


def bench_dsa_top_k(
    batch_size: int,
    q_len: int,
    seq_len: int,
    k: int,
    dtype: torch.dtype = torch.bfloat16,
    input_pattern: str = "dsa_relu",
    deterministic: bool = False,
    compare_tie_break: bool = False,
    compare_torch_deterministic: bool = False,
    compare_sglang: bool = False,
    causal_chunk: bool = False,
) -> dict:
    scores = generate_dsa_scores(
        batch_size=batch_size,
        q_len=q_len,
        seq_len=seq_len,
        dtype=dtype,
        input_pattern=input_pattern,
        causal_chunk=causal_chunk,
    )
    result = bench_top_k_from_scores(
        scores=scores,
        k=k,
        deterministic=deterministic,
        compare_tie_break=compare_tie_break,
        compare_torch_deterministic=compare_torch_deterministic,
        compare_sglang=compare_sglang,
    )
    result["rows"] = batch_size * q_len
    result["q_len"] = q_len
    result["case_type"] = "prefill" if causal_chunk else "decode"
    return result


def bench_top_k(
    batch_size: int,
    seq_len: int,
    k: int,
    dtype: torch.dtype = torch.float32,
    input_pattern: str = "random",
    deterministic: bool = False,
    compare_tie_break: bool = False,
    compare_torch_deterministic: bool = False,
    compare_sglang: bool = False,
) -> dict:
    """Benchmark basic top_k operation."""
    scores = generate_scores(batch_size, seq_len, k, dtype, input_pattern)
    return bench_top_k_from_scores(
        scores=scores,
        k=k,
        deterministic=deterministic,
        compare_tie_break=compare_tie_break,
        compare_torch_deterministic=compare_torch_deterministic,
        compare_sglang=compare_sglang,
    )


def bench_page_table_transform(
    batch_size: int,
    seq_len: int,
    k: int,
    dtype: torch.dtype = torch.float32,
    input_pattern: str = "random",
    deterministic: bool = False,
    compare_tie_break: bool = False,
    compare_sglang: bool = False,
) -> dict:
    """Benchmark fused top_k + page table transform."""
    scores = generate_scores(batch_size, seq_len, k, dtype, input_pattern)
    lengths = torch.full((batch_size,), seq_len, device="cuda", dtype=torch.int32)
    src_page_table = (
        torch.arange(seq_len, device="cuda", dtype=torch.int32)
        .unsqueeze(0)
        .expand(batch_size, -1)
        .contiguous()
    )
    set_topk_algo("default")
    fi_ms, fi_nondeterministic_ms = bench_flashinfer_modes(
        lambda deterministic_mode: flashinfer.top_k_page_table_transform(
            scores,
            src_page_table,
            lengths,
            k,
            deterministic=deterministic_mode,
            tie_break=TopKTieBreak.NONE,
        ),
        deterministic,
    )

    result = {
        "batch_size": batch_size,
        "seq_len": seq_len,
        "k": k,
        "dtype": str(dtype),
        "flashinfer_us": fi_ms * 1e3,
    }
    if fi_nondeterministic_ms is not None:
        result["flashinfer_nondeterministic_us"] = fi_nondeterministic_ms * 1e3
        result["deterministic_slowdown_vs_nondeterministic"] = (
            fi_ms / fi_nondeterministic_ms
        )

    # FlashInfer clusters
    set_topk_algo("clusters")
    fast_topk_ms = bench_median_ms(
        lambda: flashinfer.top_k_page_table_transform(
            scores, src_page_table, lengths, k
        )
    )
    result["fast_topk_us"] = fast_topk_ms * 1e3
    result["speedup_vs_flashinfer"] = fi_ms / fast_topk_ms
    set_topk_algo("auto")

    # SGLang comparison (only supports k=2048 and float32)
    if compare_sglang and HAS_SGL_KERNEL and k == 2048 and dtype == torch.float32:
        cu_seqlens_q = torch.arange(0, batch_size + 1, dtype=torch.int32, device="cuda")
        sg_ms = bench_median_ms(
            lambda: sgl_kernel.fast_topk_transform_fused(
                scores, lengths, src_page_table, cu_seqlens_q, k
            )
        )
        result["sglang_us"] = sg_ms * 1e3
        result["speedup_vs_sglang"] = sg_ms / fi_ms

    if compare_tie_break:
        # Use the same non-deterministic baseline as DetSlowdown when present.
        baseline_ms = (
            fi_nondeterministic_ms if fi_nondeterministic_ms is not None else fi_ms
        )
        result.update(
            bench_tie_break_variants(
                lambda tie_break: flashinfer.top_k_page_table_transform(
                    scores,
                    src_page_table,
                    lengths,
                    k,
                    deterministic=deterministic,
                    tie_break=tie_break,
                ),
                baseline_ms,
            )
        )

    return result


def bench_ragged_transform(
    batch_size: int,
    seq_len: int,
    k: int,
    dtype: torch.dtype = torch.float32,
    input_pattern: str = "random",
    deterministic: bool = False,
    compare_tie_break: bool = False,
    compare_sglang: bool = False,
) -> dict:
    """Benchmark fused top_k + ragged index transform."""
    scores = generate_scores(batch_size, seq_len, k, dtype, input_pattern)
    lengths = torch.full((batch_size,), seq_len, device="cuda", dtype=torch.int32)
    offsets = torch.arange(
        0, batch_size * seq_len, seq_len, device="cuda", dtype=torch.int32
    )
    set_topk_algo("default")
    fi_ms, fi_nondeterministic_ms = bench_flashinfer_modes(
        lambda deterministic_mode: flashinfer.top_k_ragged_transform(
            scores,
            offsets,
            lengths,
            k,
            deterministic=deterministic_mode,
            tie_break=TopKTieBreak.NONE,
        ),
        deterministic,
    )

    result = {
        "batch_size": batch_size,
        "seq_len": seq_len,
        "k": k,
        "dtype": str(dtype),
        "flashinfer_us": fi_ms * 1e3,
    }
    if fi_nondeterministic_ms is not None:
        result["flashinfer_nondeterministic_us"] = fi_nondeterministic_ms * 1e3
        result["deterministic_slowdown_vs_nondeterministic"] = (
            fi_ms / fi_nondeterministic_ms
        )

    # FlashInfer clusters
    set_topk_algo("clusters")
    fast_topk_ms = bench_median_ms(
        lambda: flashinfer.top_k_ragged_transform(scores, offsets, lengths, k),
    )
    result["fast_topk_us"] = fast_topk_ms * 1e3
    result["speedup_vs_flashinfer"] = fi_ms / fast_topk_ms
    set_topk_algo("auto")

    # SGLang comparison (only supports k=2048 and float32)
    if compare_sglang and HAS_SGL_KERNEL and k == 2048 and dtype == torch.float32:
        sg_ms = bench_median_ms(
            lambda: sgl_kernel.fast_topk_transform_ragged_fused(
                scores, lengths, offsets, k
            )
        )
        result["sglang_us"] = sg_ms * 1e3
        result["speedup_vs_sglang"] = sg_ms / fi_ms
    if compare_tie_break:
        baseline_ms = (
            fi_nondeterministic_ms if fi_nondeterministic_ms is not None else fi_ms
        )
        result.update(
            bench_tie_break_variants(
                lambda tie_break: flashinfer.top_k_ragged_transform(
                    scores,
                    offsets,
                    lengths,
                    k,
                    deterministic=deterministic,
                    tie_break=tie_break,
                ),
                baseline_ms,
            )
        )

    return result


# ===================== Variable-Length Segment Benchmark =====================
#
# Production sparse-attention top-k does NOT operate on fixed-size segments: each
# row selects top-k over a per-row valid window described by ``lengths`` (and, for
# the page-table case, a ``row_to_batch`` mapping). The transform benchmarks above
# always pass ``lengths == seq_len`` for every row, which hides load imbalance, the
# trivial ``length <= k`` short-circuit, and multi-CTA scheduling behavior.
#
# This section models two realistic regimes:
#   - decode  : one row per request; ``lengths`` are independent context lengths
#               drawn from a distribution (uniform / lognormal / bimodal).
#   - prefill : ``q_len`` rows per request with causal-monotonic ``lengths`` that
#               grow across query positions (the chunked-prefill pattern), wired via
#               a ``row_to_batch`` mapping.


@dataclass(frozen=True)
class VarLenCase:
    name: str
    regime: str  # "decode" or "prefill"
    length_dist: str  # "uniform" | "lognormal" | "bimodal" | "causal"
    num_requests: int
    q_len: int  # 1 for decode, >1 for prefill
    max_len: int  # padded width of the dense scores matrix
    k: int


def sample_request_lengths(
    num_requests: int,
    max_len: int,
    k: int,
    length_dist: str,
    generator: torch.Generator,
    device: torch.device,
) -> torch.Tensor:
    """Sample per-request valid context lengths for the decode regime."""
    if length_dist == "uniform":
        lengths = torch.randint(
            1,
            max_len + 1,
            (num_requests,),
            generator=generator,
            device=device,
            dtype=torch.int64,
        )
    elif length_dist == "lognormal":
        # Skewed toward shorter contexts with a long tail, typical of a serving
        # mix; median ~ max_len / 8.
        mu = math.log(max(2.0, max_len / 8.0))
        samples = torch.empty(num_requests, device=device, dtype=torch.float32)
        samples.log_normal_(mean=mu, std=1.0, generator=generator)
        lengths = samples.to(torch.int64)
    elif length_dist == "bimodal":
        # Half short (< k, hits the trivial copy path) and half long (near
        # max_len): stresses load imbalance and the trivial short-circuit.
        short_high = max(2, k)
        short = torch.randint(
            1,
            short_high,
            (num_requests,),
            generator=generator,
            device=device,
            dtype=torch.int64,
        )
        long_low = max(short_high, max_len // 2)
        long_seq = torch.randint(
            long_low,
            max_len + 1,
            (num_requests,),
            generator=generator,
            device=device,
            dtype=torch.int64,
        )
        pick_short = torch.rand(num_requests, generator=generator, device=device) < 0.5
        lengths = torch.where(pick_short, short, long_seq)
    else:
        raise ValueError(f"Unsupported decode length_dist: {length_dist}")
    return lengths.clamp_(1, max_len).to(torch.int32)


def build_causal_prefill_lengths(
    num_requests: int,
    q_len: int,
    max_len: int,
    generator: torch.Generator,
    device: torch.device,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Build causal-monotonic lengths for the prefill regime.

    Each request contributes ``q_len`` consecutive rows; query position ``p``
    attends to ``[0, ctx_start + p]`` so its valid length is ``ctx_start + p + 1``.
    Returns ``(lengths, row_to_batch)`` flattened over ``num_requests * q_len`` rows.
    """
    high = max(q_len + 1, max_len - q_len + 1)
    ctx_starts = torch.randint(
        q_len,
        high,
        (num_requests,),
        generator=generator,
        device=device,
        dtype=torch.int64,
    )
    positions = torch.arange(q_len, device=device, dtype=torch.int64)
    lengths = (ctx_starts.unsqueeze(1) + positions.unsqueeze(0) + 1).clamp(1, max_len)
    lengths = lengths.reshape(-1).to(torch.int32)
    row_to_batch = torch.arange(
        num_requests, device=device, dtype=torch.int32
    ).repeat_interleave(q_len)
    return lengths, row_to_batch


def summarize_lengths(lengths: torch.Tensor, k: int) -> tuple[int, float, int, float]:
    """Return ``(min, mean, max, fraction-of-rows-with-length<=k)``."""
    len_min = int(lengths.min().item())
    len_max = int(lengths.max().item())
    len_mean = float(lengths.to(torch.float64).mean().item())
    trivial_frac = float((lengths <= k).to(torch.float64).mean().item())
    return len_min, len_mean, len_max, trivial_frac


def build_masked_scores(scores: torch.Tensor, lengths: torch.Tensor) -> torch.Tensor:
    """Mask positions beyond each row's valid length with -inf.

    Gives ``torch.topk`` (which has no length argument) the same valid window as the
    length-aware kernel. Built once, outside the timed region, so the comparison is
    selection-cost vs selection-cost.
    """
    max_len = scores.size(1)
    col = torch.arange(max_len, device=scores.device).unsqueeze(0)
    invalid = col >= lengths.unsqueeze(1)
    if scores.dtype == torch.float32:
        neg_inf = float("-inf")
    else:
        neg_inf = torch.finfo(scores.dtype).min
    return scores.masked_fill(invalid, neg_inf)


def generate_varlen_inputs(
    case: VarLenCase,
    dtype: torch.dtype,
    generator: torch.Generator,
    device: torch.device,
) -> dict:
    """Build scores + length/offset/page-table tensors for a variable-length case."""
    max_len = case.max_len
    if case.regime == "decode":
        num_requests = case.num_requests
        lengths = sample_request_lengths(
            num_requests, max_len, case.k, case.length_dist, generator, device
        )
        row_to_batch = None
        num_rows = num_requests
    elif case.regime == "prefill":
        num_requests = case.num_requests
        lengths, row_to_batch = build_causal_prefill_lengths(
            num_requests, case.q_len, max_len, generator, device
        )
        num_rows = num_requests * case.q_len
    else:
        raise ValueError(f"Unsupported regime: {case.regime}")

    scores = torch.randn(
        num_rows, max_len, device=device, dtype=dtype, generator=generator
    )

    # Ragged offsets model a tightly packed KV buffer: offsets[i] = sum(lengths[:i]).
    # The absolute offset value is perf-irrelevant (it is just added to the output
    # index); the per-row ``lengths`` drive the actual work.
    offsets = torch.zeros(num_rows, device=device, dtype=torch.int32)
    if num_rows > 1:
        offsets[1:] = torch.cumsum(lengths[:-1].to(torch.int64), dim=0).to(torch.int32)

    # Page table is per request: (num_requests, max_len).
    src_page_table = (
        torch.arange(max_len, device=device, dtype=torch.int32)
        .unsqueeze(0)
        .expand(num_requests, max_len)
        .contiguous()
    )
    return {
        "scores": scores,
        "lengths": lengths,
        "offsets": offsets,
        "src_page_table": src_page_table,
        "row_to_batch": row_to_batch,
        "num_requests": num_requests,
    }


def build_varlen_cases(
    length_dists: list[str],
    max_lens: list[int],
    k_values: list[int],
    decode_batches: list[int],
    prefill_requests: list[int],
    q_len: int,
) -> list[VarLenCase]:
    cases: list[VarLenCase] = []
    decode_dists = [d for d in length_dists if d in ("uniform", "lognormal", "bimodal")]
    for dist in decode_dists:
        for num_requests in decode_batches:
            for max_len in max_lens:
                for k in k_values:
                    if k >= max_len:
                        continue
                    cases.append(
                        VarLenCase(
                            name=f"decode_{dist}_b{num_requests}_l{max_len}_k{k}",
                            regime="decode",
                            length_dist=dist,
                            num_requests=num_requests,
                            q_len=1,
                            max_len=max_len,
                            k=k,
                        )
                    )
    if "causal" in length_dists:
        for num_requests in prefill_requests:
            for max_len in max_lens:
                for k in k_values:
                    if k >= max_len:
                        continue
                    cases.append(
                        VarLenCase(
                            name=f"prefill_causal_r{num_requests}_q{q_len}_l{max_len}_k{k}",
                            regime="prefill",
                            length_dist="causal",
                            num_requests=num_requests,
                            q_len=q_len,
                            max_len=max_len,
                            k=k,
                        )
                    )
    return cases


def bench_varlen_transform(
    case: VarLenCase,
    transform: str,
    dtype: torch.dtype,
    generator: torch.Generator,
    has_clusters: bool,
    deterministic: bool = False,
    compare_tie_break: bool = False,
) -> dict:
    """Benchmark a transform API on realistic variable-length segments."""
    device = torch.device("cuda")
    inputs = generate_varlen_inputs(case, dtype, generator, device)
    scores = inputs["scores"]
    lengths = inputs["lengths"]
    offsets = inputs["offsets"]
    src_page_table = inputs["src_page_table"]
    row_to_batch = inputs["row_to_batch"]
    num_rows = scores.size(0)
    k = case.k

    def run(deterministic_mode, tie_break=TopKTieBreak.NONE):
        if transform == "page_table":
            return flashinfer.top_k_page_table_transform(
                scores,
                src_page_table,
                lengths,
                k,
                row_to_batch=row_to_batch,
                deterministic=deterministic_mode,
                tie_break=tie_break,
            )
        return flashinfer.top_k_ragged_transform(
            scores,
            offsets,
            lengths,
            k,
            deterministic=deterministic_mode,
            tie_break=tie_break,
        )

    set_topk_algo("default")
    fi_ms, fi_nondeterministic_ms = bench_flashinfer_modes(run, deterministic)

    len_min, len_mean, len_max, trivial_frac = summarize_lengths(lengths, k)
    result = {
        "regime": case.regime,
        "length_dist": case.length_dist,
        "transform": transform,
        "num_rows": num_rows,
        "num_requests": inputs["num_requests"],
        "max_len": case.max_len,
        "k": k,
        "len_min": len_min,
        "len_mean": len_mean,
        "len_max": len_max,
        "trivial_frac": trivial_frac,
        "flashinfer_us": fi_ms * 1e3,
    }
    if fi_nondeterministic_ms is not None:
        result["flashinfer_nondeterministic_us"] = fi_nondeterministic_ms * 1e3
        result["deterministic_slowdown_vs_nondeterministic"] = (
            fi_ms / fi_nondeterministic_ms
        )

    # clusters is a non-deterministic SM100 path; it also only dispatches for
    # page_table when row_to_batch is None (the prefill regime sets row_to_batch and
    # falls back to the default path). Only measure/report it when it would actually
    # run, otherwise the "clusters" timing would just be the default path relabeled.
    can_run_clusters = (
        has_clusters
        and not (deterministic or compare_tie_break)
        and not (transform == "page_table" and row_to_batch is not None)
    )
    if can_run_clusters:
        set_topk_algo("clusters")
        clusters_ms = bench_median_ms(lambda: run(False))
        result["clusters_us"] = clusters_ms * 1e3
        result["speedup_clusters_vs_default"] = fi_ms / clusters_ms
    set_topk_algo("auto")

    if compare_tie_break:
        # Use the same non-deterministic baseline as DetSlowdown when present.
        baseline_ms = (
            fi_nondeterministic_ms if fi_nondeterministic_ms is not None else fi_ms
        )
        result.update(
            bench_tie_break_variants(
                lambda tie_break: run(deterministic, tie_break),
                baseline_ms,
            )
        )

    # torch reference operates on a pre-masked tensor (mask built outside timing) so
    # we compare selection cost against the length-aware kernel.
    masked_scores = build_masked_scores(scores, lengths)
    with torch_deterministic_algorithms(deterministic):
        torch_ms = bench_median_ms(lambda: torch.topk(masked_scores, k, dim=-1))
    result["torch_us"] = torch_ms * 1e3
    result["speedup_vs_torch"] = torch_ms / fi_ms

    return result


def parse_dtype(dtype_str: str) -> torch.dtype:
    """Parse dtype string to torch.dtype."""
    dtype_map = {
        "fp32": torch.float32,
        "float32": torch.float32,
        "fp16": torch.float16,
        "float16": torch.float16,
        "half": torch.float16,
        "bf16": torch.bfloat16,
        "bfloat16": torch.bfloat16,
    }
    return dtype_map[dtype_str.lower()]


@torch.inference_mode()
def main():
    parser = argparse.ArgumentParser(description="Benchmark Top-K operations")
    parser.add_argument(
        "--compare-sglang",
        action="store_true",
        help="Compare with SGLang's sgl_kernel (requires sgl_kernel installed)",
    )
    parser.add_argument(
        "--op",
        choices=["all", "top_k", "dsa_topk", "page_table", "ragged", "varlen"],
        default="all",
        help="Which operation to benchmark",
    )
    parser.add_argument(
        "--dtype",
        type=str,
        default="fp32",
        help="Data type: fp32, fp16, or bf16 (default: fp32)",
    )
    parser.add_argument(
        "--compare-algorithms",
        action="store_true",
        help="Compare multi-CTA vs filtered algorithms",
    )
    parser.add_argument(
        "--deterministic",
        action="store_true",
        help="Enable deterministic mode for FlashInfer top-k kernels",
    )
    parser.add_argument(
        "--tie-break",
        action="store_true",
        help=(
            "Also benchmark tie-break variants and report "
            "FlashInfer(tie-small/tie-large) columns with slowdown against "
            "the non-deterministic baseline"
        ),
    )
    parser.add_argument(
        "--compare-torch-deterministic",
        action="store_true",
        help="Also benchmark torch.topk under deterministic algorithm mode",
    )
    parser.add_argument(
        "--input-pattern",
        choices=[
            "random",
            "quantized_random",
            "relu_quantized",
            "tie_heavy",
            "pivot_tie",
        ],
        default="random",
        help=(
            "Input score pattern: random | quantized_random | relu_quantized | "
            "tie_heavy | pivot_tie"
        ),
    )
    parser.add_argument(
        "--dsa-input-pattern",
        choices=["random", "quantized_random", "dsa_relu"],
        default="dsa_relu",
        help="DSA top-k input pattern: random | quantized_random | dsa_relu",
    )
    parser.add_argument(
        "--dsa-case",
        choices=["all", "decode", "prefill"],
        default="all",
        help="DSA case group: all | decode | prefill",
    )
    parser.add_argument(
        "--dsa-topk",
        type=int,
        default=2048,
        help="Top-k for DSA workload (default: 2048, matching DeepSeek DSA config)",
    )
    parser.add_argument(
        "--length-dist",
        choices=["all", "uniform", "lognormal", "bimodal", "causal"],
        default="all",
        help=(
            "Variable-length segment distribution for --op varlen: "
            "uniform | lognormal | bimodal (decode regimes) | causal (prefill regime) | all"
        ),
    )
    parser.add_argument(
        "--varlen-k",
        type=int,
        default=2048,
        help="Top-k for the varlen benchmark (default: 2048, matching DeepSeek DSA config)",
    )
    parser.add_argument(
        "--varlen-q-len",
        type=int,
        default=128,
        help="Query length per request for the varlen prefill (causal) regime (default: 128)",
    )
    args = parser.parse_args()

    if args.varlen_k <= 0:
        parser.error("--varlen-k must be a positive integer")
    if args.varlen_q_len <= 0:
        parser.error("--varlen-q-len must be a positive integer")

    dtype = parse_dtype(args.dtype)

    if args.compare_sglang and not HAS_SGL_KERNEL:
        print("WARNING: sgl_kernel not found, skipping SGLang comparison")
        args.compare_sglang = False

    # Test configurations
    batch_sizes = [1, 16, 64, 256]
    seq_lens = [256, 512, 1024, 2048, 4096, 16384, 65536, 131072, 262144, 524288]
    k_values = [256, 512, 1024, 2048, 4096]
    top_k_cases = build_top_k_cases(
        batch_sizes=batch_sizes,
        seq_lens=seq_lens,
        k_values=k_values,
    )

    dtype_str = args.dtype.upper()

    # Algorithm comparison mode
    if args.compare_algorithms:
        if args.deterministic:
            print(
                "ERROR: --compare-algorithms is only meaningful with non-deterministic mode"
            )
            return
        if args.tie_break:
            print("ERROR: --compare-algorithms does not support --tie-break")
            return
        print("=" * 100)
        print(
            "Algorithm comparison: Multi-CTA vs Filtered "
            f"(dtype={dtype_str}, k=2048, pattern={args.input_pattern})"
        )
        print("=" * 100)
        print(
            f"{'batch':>6} {'seq_len':>10} | {'Multi-CTA':>12} {'Filtered':>12} {'Winner':>12}"
        )
        print("-" * 70)

        for batch_size in batch_sizes:
            for seq_len in seq_lens:
                k = 2048
                if k > seq_len:
                    continue
                try:
                    # Benchmark Multi-CTA
                    set_topk_algo("multi_cta")
                    result_mc = bench_page_table_transform(
                        batch_size, seq_len, k, dtype, args.input_pattern
                    )
                    mc_us = result_mc["flashinfer_us"]

                    # Benchmark Filtered
                    set_topk_algo("filtered")
                    result_f = bench_page_table_transform(
                        batch_size, seq_len, k, dtype, args.input_pattern
                    )
                    f_us = result_f["flashinfer_us"]

                    # Reset to auto
                    set_topk_algo("auto")

                    winner = "Multi-CTA" if mc_us < f_us else "Filtered"
                    speedup = max(mc_us, f_us) / min(mc_us, f_us)
                    print(
                        f"{batch_size:>6} {seq_len:>10} | "
                        f"{mc_us:>10.2f}us {f_us:>10.2f}us "
                        f"{winner:>8} {speedup:.2f}x"
                    )
                except RuntimeError as e:
                    error_label = classify_benchmark_runtime_error(e)
                    if error_label == "OOM":
                        print(f"{batch_size:>6} {seq_len:>10} | OOM")
                        torch.cuda.empty_cache()
                    elif error_label == "UNSUPPORTED":
                        print(f"{batch_size:>6} {seq_len:>10} | UNSUPPORTED")
                    else:
                        raise
        return

    if args.op in ["all", "top_k"]:
        show_det_or_tie = args.deterministic or args.tie_break
        print("=" * 100)
        print(
            "top_k: Basic radix-based top-k selection "
            f"(dtype={dtype_str}, deterministic={args.deterministic}, "
            f"pattern={args.input_pattern}, tie_break={args.tie_break})"
        )
        if args.compare_sglang:
            print("NOTE: SGLang only supports k=2048 and float32")
        if args.deterministic:
            print(
                "NOTE: deterministic mode also benchmarks FlashInfer(non-det) "
                "for direct comparison"
            )
        if args.compare_torch_deterministic:
            print(
                "NOTE: torch.det means torch.topk with torch.use_deterministic_algorithms(True)"
            )
        elif args.deterministic:
            print(
                "NOTE: torch column uses torch.topk with "
                "torch.use_deterministic_algorithms(True)"
            )
        if args.tie_break:
            print(
                f"NOTE: tie-break columns use deterministic={args.deterministic}; "
                "slowdowns use the non-deterministic baseline"
            )
        print(
            "NOTE: default top-k sweep includes two extra large-batch/long-vocab "
            "stress cases beyond the original grid"
        )
        print("=" * 100)

        if show_det_or_tie:
            torch_label = "torch.det" if args.deterministic else "torch.topk"
            header = (
                f"{'batch':>6} {'seq_len':>10} {'k':>6} | "
                f"{'FlashInfer':>12} {'FlashInfer(det)':>14} {'DetSlowdown':>11}"
            )
            header = append_tie_break_header(header, args.tie_break)
            header += f" {torch_label:>12} {'Speedup':>10}"
        else:
            header = (
                f"{'batch':>6} {'seq_len':>10} {'k':>6} | "
                f"{'FlashInfer':>12} {'torch.topk':>12} {'Speedup':>10}"
                f" {'Clusters':>12} {'Speedup Clusters vs. Default':>29}"
            )
        if args.compare_torch_deterministic and not show_det_or_tie:
            header += f" {'torch.det':>12} {'Speedup':>10}"
        if args.compare_sglang:
            header += f" {'SGLang':>12} {'Speedup':>10}"
        print(header)
        print("-" * len(header))

        for case in top_k_cases:
            try:
                result = bench_top_k(
                    case.batch_size,
                    case.seq_len,
                    case.k,
                    dtype,
                    input_pattern=args.input_pattern,
                    deterministic=args.deterministic,
                    compare_tie_break=args.tie_break,
                    compare_torch_deterministic=args.compare_torch_deterministic,
                    compare_sglang=args.compare_sglang,
                )
                if show_det_or_tie:
                    nondet_us = result.get("flashinfer_nondeterministic_us")
                    if nondet_us is None:
                        nondet_us = result["flashinfer_us"]
                    det_us = result["flashinfer_us"] if args.deterministic else None
                    det_us_str = (
                        f"{det_us:>12.2f}us" if det_us is not None else f"{'n/a':>14}"
                    )
                    det_slowdown = result.get(
                        "deterministic_slowdown_vs_nondeterministic"
                    )
                    det_slowdown_str = (
                        f"{det_slowdown:>10.2f}x"
                        if det_slowdown is not None
                        else f"{'n/a':>11}"
                    )
                    line = (
                        f"{result['batch_size']:>6} {result['seq_len']:>10} {result['k']:>6} | "
                        f"{nondet_us:>10.2f}us "
                        f"{det_us_str} "
                        f"{det_slowdown_str}"
                    )
                    line = append_tie_break_columns(line, result, args.tie_break)
                    line += (
                        f" {result['torch_us']:>10.2f}us "
                        f"{result['speedup_vs_torch']:>9.2f}x"
                    )
                else:
                    line = (
                        f"{result['batch_size']:>6} {result['seq_len']:>10} {result['k']:>6} | "
                        f"{result['flashinfer_us']:>12.2f}us {result['torch_us']:>10.2f}us "
                        f"{result['speedup_vs_torch']:>9.2f}x"
                    )
                    if "fast_topk_us" in result:
                        line += f" {result['fast_topk_us']:>10.2f}us {result['speedup_vs_flashinfer']:>28.2f}x"
                if "torch_deterministic_us" in result and not show_det_or_tie:
                    line += (
                        f" {result['torch_deterministic_us']:>10.2f}us "
                        f"{result['speedup_vs_torch_deterministic']:>9.2f}x"
                    )
                if "sglang_us" in result:
                    line += (
                        f" {result['sglang_us']:>10.2f}us "
                        f"{result['speedup_vs_sglang']:>9.2f}x"
                    )
                elif args.compare_sglang and case.k == 2048:
                    line += " (SGLang error)"
                print(line)
            except RuntimeError as e:
                error_label = classify_benchmark_runtime_error(e)
                if error_label is not None:
                    print(
                        f"{case.batch_size:>6} {case.seq_len:>10} {case.k:>6} | {error_label}"
                    )
                    torch.cuda.empty_cache()
                else:
                    raise

    if args.op in ["all", "dsa_topk"]:
        show_det_or_tie = args.deterministic or args.tie_break
        print("\n" + "=" * 100)
        print(
            "dsa_topk: DeepSeek DSA-like indexer top-k workload "
            f"(dtype={dtype_str}, deterministic={args.deterministic}, "
            f"dsa_pattern={args.dsa_input_pattern}, k={args.dsa_topk}, "
            f"tie_break={args.tie_break})"
        )
        if args.deterministic:
            print(
                "NOTE: deterministic mode also benchmarks FlashInfer(non-det) "
                "for direct comparison"
            )
        if args.compare_torch_deterministic:
            print(
                "NOTE: torch.det means torch.topk with torch.use_deterministic_algorithms(True)"
            )
        elif args.deterministic:
            print(
                "NOTE: torch column uses torch.topk with "
                "torch.use_deterministic_algorithms(True)"
            )
        if args.tie_break:
            print(
                f"NOTE: tie-break columns use deterministic={args.deterministic}; "
                "slowdowns use the non-deterministic baseline"
            )
        print("=" * 100)

        if show_det_or_tie:
            torch_label = "torch.det" if args.deterministic else "torch.topk"
            header = (
                f"{'case':>24} {'rows':>8} {'seq_len':>10} {'k':>6} | "
                f"{'FlashInfer':>12} {'FlashInfer(det)':>14} {'DetSlowdown':>11}"
            )
            header = append_tie_break_header(header, args.tie_break)
            header += f" {torch_label:>12} {'Speedup':>10}"
        else:
            header = (
                f"{'case':>24} {'rows':>8} {'seq_len':>10} {'k':>6} | "
                f"{'FlashInfer':>12} {'torch.topk':>12} {'Speedup':>10}"
                f" {'Clusters':>12} {'Speedup Clusters vs. Default':>29}"
            )
        if args.compare_torch_deterministic and not show_det_or_tie:
            header += f" {'torch.det':>12} {'Speedup':>10}"
        print(header)
        print("-" * len(header))

        dsa_cases = [
            # DeepSeek Sparse Attention proxy cases:
            # - decode: q_len=1
            # - prefill chunk: q_len>1 with causal availability growth
            # Ref: https://github.com/deepseek-ai/DeepSeek-V3.2-Exp/blob/main/DeepSeek_V3_2.pdf
            DSATopKCase("decode_b1_q1_l128k", 1, 1, 131072, False),
            DSATopKCase("decode_b8_q1_l64k", 8, 1, 65536, False),
            DSATopKCase("decode_b32_q1_l128k", 32, 1, 131072, False),
            DSATopKCase("prefill_b1_q128_l128k", 1, 128, 131072, True),
        ]

        for case in dsa_cases:
            if args.dsa_case == "decode" and case.causal_chunk:
                continue
            if args.dsa_case == "prefill" and not case.causal_chunk:
                continue
            if args.dsa_topk > case.seq_len:
                continue
            try:
                result = bench_dsa_top_k(
                    batch_size=case.batch_size,
                    q_len=case.q_len,
                    seq_len=case.seq_len,
                    k=args.dsa_topk,
                    dtype=dtype,
                    input_pattern=args.dsa_input_pattern,
                    deterministic=args.deterministic,
                    compare_tie_break=args.tie_break,
                    compare_torch_deterministic=args.compare_torch_deterministic,
                    compare_sglang=False,
                    causal_chunk=case.causal_chunk,
                )
                if show_det_or_tie:
                    nondet_us = result.get("flashinfer_nondeterministic_us")
                    if nondet_us is None:
                        nondet_us = result["flashinfer_us"]
                    det_us = result["flashinfer_us"] if args.deterministic else None
                    det_us_str = (
                        f"{det_us:>12.2f}us" if det_us is not None else f"{'n/a':>14}"
                    )
                    det_slowdown = result.get(
                        "deterministic_slowdown_vs_nondeterministic"
                    )
                    det_slowdown_str = (
                        f"{det_slowdown:>10.2f}x"
                        if det_slowdown is not None
                        else f"{'n/a':>11}"
                    )
                    line = (
                        f"{case.name:>24} {result['rows']:>8} {result['seq_len']:>10} {result['k']:>6} | "
                        f"{nondet_us:>10.2f}us "
                        f"{det_us_str} "
                        f"{det_slowdown_str}"
                    )
                    line = append_tie_break_columns(line, result, args.tie_break)
                    line += (
                        f" {result['torch_us']:>10.2f}us "
                        f"{result['speedup_vs_torch']:>9.2f}x"
                    )
                else:
                    line = (
                        f"{case.name:>24} {result['rows']:>8} {result['seq_len']:>10} {result['k']:>6} | "
                        f"{result['flashinfer_us']:>10.2f}us {result['torch_us']:>10.2f}us "
                        f"{result['speedup_vs_torch']:>9.2f}x"
                    )
                    if "fast_topk_us" in result:
                        line += f" {result['fast_topk_us']:>10.2f}us {result['speedup_vs_flashinfer']:>28.2f}x"
                if "torch_deterministic_us" in result and not show_det_or_tie:
                    line += (
                        f" {result['torch_deterministic_us']:>10.2f}us "
                        f"{result['speedup_vs_torch_deterministic']:>9.2f}x"
                    )
                print(line)
            except RuntimeError as e:
                error_label = classify_benchmark_runtime_error(e)
                if error_label is not None:
                    print(
                        f"{case.name:>24} {case.batch_size * case.q_len:>8} {case.seq_len:>10} "
                        f"{args.dsa_topk:>6} | {error_label}"
                    )
                    torch.cuda.empty_cache()
                else:
                    raise

    if args.op in ["all", "page_table"]:
        show_det_or_tie = args.deterministic or args.tie_break
        print("\n" + "=" * 100)
        print(
            "top_k_page_table_transform: Fused top-k + page table gather "
            f"(dtype={dtype_str}, deterministic={args.deterministic}, "
            f"pattern={args.input_pattern}, tie_break={args.tie_break})"
        )
        if args.compare_sglang:
            print("NOTE: SGLang only supports k=2048 and float32")
        if args.deterministic:
            print(
                "NOTE: deterministic mode also benchmarks FlashInfer(non-det) "
                "for direct comparison"
            )
        if args.tie_break:
            print(
                f"NOTE: tie-break columns use deterministic={args.deterministic}; "
                "slowdowns use the non-deterministic baseline"
            )
        print("=" * 100)

        if show_det_or_tie:
            header = (
                f"{'batch':>6} {'seq_len':>10} {'k':>6} | "
                f"{'FlashInfer':>12} {'FlashInfer(det)':>14} {'DetSlowdown':>11}"
            )
            header = append_tie_break_header(header, args.tie_break)
        else:
            header = f"{'batch':>6} {'seq_len':>10} {'k':>6} | {'FlashInfer':>12} {'Clusters':>12} {'Speedup Clusters vs. Default':>29}"
        if args.compare_sglang:
            header += f" {'SGLang':>12} {'Speedup':>10}"
        print(header)
        print("-" * len(header))

        for batch_size in batch_sizes:
            for seq_len in seq_lens:
                for k in k_values:
                    if k > seq_len:
                        continue
                    try:
                        result = bench_page_table_transform(
                            batch_size,
                            seq_len,
                            k,
                            dtype,
                            input_pattern=args.input_pattern,
                            deterministic=args.deterministic,
                            compare_tie_break=args.tie_break,
                            compare_sglang=args.compare_sglang,
                        )
                        if show_det_or_tie:
                            nondet_us = result.get("flashinfer_nondeterministic_us")
                            if nondet_us is None:
                                nondet_us = result["flashinfer_us"]
                            det_us = (
                                result["flashinfer_us"] if args.deterministic else None
                            )
                            det_us_str = (
                                f"{det_us:>12.2f}us"
                                if det_us is not None
                                else f"{'n/a':>14}"
                            )
                            det_slowdown = result.get(
                                "deterministic_slowdown_vs_nondeterministic"
                            )
                            det_slowdown_str = (
                                f"{det_slowdown:>10.2f}x"
                                if det_slowdown is not None
                                else f"{'n/a':>11}"
                            )
                            line = (
                                f"{result['batch_size']:>6} {result['seq_len']:>10} {result['k']:>6} | "
                                f"{nondet_us:>10.2f}us "
                                f"{det_us_str} "
                                f"{det_slowdown_str}"
                            )
                            line = append_tie_break_columns(
                                line, result, args.tie_break
                            )
                        else:
                            line = (
                                f"{result['batch_size']:>6} {result['seq_len']:>10} {result['k']:>6} | "
                                f"{result['flashinfer_us']:>10.2f}us"
                            )
                            if "fast_topk_us" in result:
                                line += f" {result['fast_topk_us']:>10.2f}us {result['speedup_vs_flashinfer']:>28.2f}x"
                        if "sglang_us" in result:
                            line += (
                                f" {result['sglang_us']:>10.2f}us "
                                f"{result['speedup_vs_sglang']:>9.2f}x"
                            )
                        elif args.compare_sglang and k == 2048:
                            line += " (SGLang error)"
                        print(line)
                    except RuntimeError as e:
                        error_label = classify_benchmark_runtime_error(e)
                        if error_label is not None:
                            print(
                                f"{batch_size:>6} {seq_len:>10} {k:>6} | {error_label}"
                            )
                            torch.cuda.empty_cache()
                        else:
                            raise

    if args.op in ["all", "ragged"]:
        show_det_or_tie = args.deterministic or args.tie_break
        print("\n" + "=" * 100)
        print(
            "top_k_ragged_transform: Fused top-k + ragged index transform "
            f"(dtype={dtype_str}, deterministic={args.deterministic}, "
            f"pattern={args.input_pattern}, tie_break={args.tie_break})"
        )
        if args.compare_sglang:
            print("NOTE: SGLang only supports k=2048 and float32")
        if args.deterministic:
            print(
                "NOTE: deterministic mode also benchmarks FlashInfer(non-det) "
                "for direct comparison"
            )
        if args.tie_break:
            print(
                f"NOTE: tie-break columns use deterministic={args.deterministic}; "
                "slowdowns use the non-deterministic baseline"
            )
        print("=" * 100)

        if show_det_or_tie:
            header = (
                f"{'batch':>6} {'seq_len':>10} {'k':>6} | "
                f"{'FlashInfer':>12} {'FlashInfer(det)':>14} {'DetSlowdown':>11}"
            )
            header = append_tie_break_header(header, args.tie_break)
        else:
            header = f"{'batch':>6} {'seq_len':>10} {'k':>6} | {'FlashInfer':>12} {'Clusters':>12} {'Speedup Clusters vs. Default':>29}"
        if args.compare_sglang:
            header += f" {'SGLang':>12} {'Speedup':>10}"
        print(header)
        print("-" * len(header))

        for batch_size in batch_sizes:
            for seq_len in seq_lens:
                for k in k_values:
                    if k > seq_len:
                        continue
                    try:
                        result = bench_ragged_transform(
                            batch_size,
                            seq_len,
                            k,
                            dtype,
                            input_pattern=args.input_pattern,
                            deterministic=args.deterministic,
                            compare_tie_break=args.tie_break,
                            compare_sglang=args.compare_sglang,
                        )
                        if show_det_or_tie:
                            nondet_us = result.get("flashinfer_nondeterministic_us")
                            if nondet_us is None:
                                nondet_us = result["flashinfer_us"]
                            det_us = (
                                result["flashinfer_us"] if args.deterministic else None
                            )
                            det_us_str = (
                                f"{det_us:>12.2f}us"
                                if det_us is not None
                                else f"{'n/a':>14}"
                            )
                            det_slowdown = result.get(
                                "deterministic_slowdown_vs_nondeterministic"
                            )
                            det_slowdown_str = (
                                f"{det_slowdown:>10.2f}x"
                                if det_slowdown is not None
                                else f"{'n/a':>11}"
                            )
                            line = (
                                f"{result['batch_size']:>6} {result['seq_len']:>10} {result['k']:>6} | "
                                f"{nondet_us:>10.2f}us "
                                f"{det_us_str} "
                                f"{det_slowdown_str}"
                            )
                            line = append_tie_break_columns(
                                line, result, args.tie_break
                            )
                        else:
                            line = (
                                f"{result['batch_size']:>6} {result['seq_len']:>10} {result['k']:>6} | "
                                f"{result['flashinfer_us']:>10.2f}us"
                            )
                            if "fast_topk_us" in result:
                                line += f" {result['fast_topk_us']:>10.2f}us {result['speedup_vs_flashinfer']:>28.2f}x"
                        if "sglang_us" in result:
                            line += (
                                f" {result['sglang_us']:>10.2f}us "
                                f"{result['speedup_vs_sglang']:>9.2f}x"
                            )
                        elif args.compare_sglang and k == 2048:
                            line += " (SGLang error)"
                        print(line)
                    except RuntimeError as e:
                        error_label = classify_benchmark_runtime_error(e)
                        if error_label is not None:
                            print(
                                f"{batch_size:>6} {seq_len:>10} {k:>6} | {error_label}"
                            )
                            torch.cuda.empty_cache()
                        else:
                            raise

    if args.op in ["all", "varlen"]:
        device = torch.device("cuda")
        cap = get_compute_capability(device)
        has_clusters = cap[0] == 10
        length_dists = (
            ["uniform", "lognormal", "bimodal", "causal"]
            if args.length_dist == "all"
            else [args.length_dist]
        )
        varlen_max_lens = [16384, 65536, 131072]
        varlen_k_values = [args.varlen_k]
        varlen_decode_batches = [16, 128]
        varlen_prefill_requests = [4, 16]
        varlen_cases = build_varlen_cases(
            length_dists,
            varlen_max_lens,
            varlen_k_values,
            varlen_decode_batches,
            varlen_prefill_requests,
            args.varlen_q_len,
        )
        # A torch.Generator drives the (otherwise random) length/score draws. It is
        # re-seeded per case in the loop below so the page_table and ragged runs for a
        # given case use identical inputs (and stays reproducible across runs).
        generator = torch.Generator(device=device)

        show_det_or_tie = args.deterministic or args.tie_break
        # clusters is a non-deterministic SM100 path; omit it under deterministic/tie-break.
        show_clusters = has_clusters and not show_det_or_tie

        print("\n" + "=" * 100)
        print(
            "varlen: Variable-length segment top-k transforms (production-realistic) "
            f"(dtype={dtype_str}, length_dist={args.length_dist}, k={args.varlen_k}, "
            f"deterministic={args.deterministic}, tie_break={args.tie_break})"
        )
        print(
            "NOTE: lengths model per-row valid windows; decode = independent context "
            f"lengths, prefill(causal) = monotonic growth within a request (q_len={args.varlen_q_len})"
        )
        print(
            "NOTE: torch(mask) masks invalid positions once (outside timing) then "
            "torch.topk, isolating selection cost vs the length-aware kernel"
        )
        if show_det_or_tie:
            if args.deterministic:
                print(
                    "NOTE: deterministic mode also benchmarks FlashInfer(non-det) "
                    "for direct comparison"
                )
            if args.tie_break:
                print(
                    f"NOTE: tie-break columns use deterministic={args.deterministic}; "
                    "slowdowns use the non-deterministic baseline"
                )
            print(
                "NOTE: Clusters column omitted under deterministic/tie-break "
                "(clusters requires the non-deterministic path)"
            )
        elif not has_clusters:
            print(
                "NOTE: clusters path requires SM100 (Blackwell); omitting Clusters "
                f"column on this device (SM{cap[0]}{cap[1]})"
            )
        print("=" * 100)

        base_header = (
            f"{'regime':>8} {'dist':>10} {'transform':>11} {'rows':>8} {'reqs':>6} "
            f"{'max_len':>9} {'k':>6} | {'len_min':>8} {'len_mean':>9} {'len_max':>8} "
            f"{'triv%':>6} | "
        )
        if show_det_or_tie:
            header = (
                base_header
                + f"{'FlashInfer':>12} {'FlashInfer(det)':>14} {'DetSlowdown':>11}"
            )
            header = append_tie_break_header(header, args.tie_break)
            header += f" {'torch(mask)':>13} {'Speedup':>9}"
        else:
            header = base_header + f"{'FlashInfer':>12}"
            if show_clusters:
                header += f" {'Clusters':>12} {'vsClusters':>10}"
            header += f" {'torch(mask)':>13} {'Speedup':>9}"
        print(header)
        print("-" * len(header))

        for case_idx, case in enumerate(varlen_cases):
            for transform in ["page_table", "ragged"]:
                # Re-seed per case so page_table and ragged see identical inputs,
                # keeping the per-case transform comparison and length stats equivalent.
                generator.manual_seed(1234 + case_idx)
                needs_cache_cleanup = False
                try:
                    result = bench_varlen_transform(
                        case,
                        transform,
                        dtype,
                        generator,
                        has_clusters,
                        deterministic=args.deterministic,
                        compare_tie_break=args.tie_break,
                    )
                    base_line = (
                        f"{result['regime']:>8} {result['length_dist']:>10} "
                        f"{result['transform']:>11} {result['num_rows']:>8} "
                        f"{result['num_requests']:>6} {result['max_len']:>9} "
                        f"{result['k']:>6} | {result['len_min']:>8} "
                        f"{result['len_mean']:>9.1f} {result['len_max']:>8} "
                        f"{100.0 * result['trivial_frac']:>5.1f}% | "
                    )
                    if show_det_or_tie:
                        nondet_us = result.get("flashinfer_nondeterministic_us")
                        if nondet_us is None:
                            nondet_us = result["flashinfer_us"]
                        det_us = result["flashinfer_us"] if args.deterministic else None
                        det_us_str = (
                            f"{det_us:>12.2f}us"
                            if det_us is not None
                            else f"{'n/a':>14}"
                        )
                        det_slowdown = result.get(
                            "deterministic_slowdown_vs_nondeterministic"
                        )
                        det_slowdown_str = (
                            f"{det_slowdown:>10.2f}x"
                            if det_slowdown is not None
                            else f"{'n/a':>11}"
                        )
                        line = (
                            base_line
                            + f"{nondet_us:>10.2f}us {det_us_str} {det_slowdown_str}"
                        )
                        line = append_tie_break_columns(line, result, args.tie_break)
                        line += (
                            f" {result['torch_us']:>11.2f}us "
                            f"{result['speedup_vs_torch']:>8.2f}x"
                        )
                    else:
                        line = base_line + f"{result['flashinfer_us']:>10.2f}us"
                        if show_clusters:
                            # prefill page_table falls back to the default path, so it
                            # has no clusters timing; pad to keep columns aligned.
                            if "clusters_us" in result:
                                line += (
                                    f" {result['clusters_us']:>10.2f}us "
                                    f"{result['speedup_clusters_vs_default']:>9.2f}x"
                                )
                            else:
                                line += f" {'n/a':>12} {'n/a':>10}"
                        line += (
                            f" {result['torch_us']:>11.2f}us "
                            f"{result['speedup_vs_torch']:>8.2f}x"
                        )
                    print(line)
                except RuntimeError as e:
                    error_label = classify_benchmark_runtime_error(e)
                    if error_label is None:
                        raise
                    print(
                        f"{case.regime:>8} {case.length_dist:>10} "
                        f"{transform:>11} {case.max_len:>9} {case.k:>6} | {error_label}"
                    )
                    needs_cache_cleanup = True
                # Reclaim cached memory only after the except block exits. While the
                # handler is active, the in-flight exception keeps bench_varlen_transform's
                # frame (and its large GPU tensors) alive, so empty_cache() inside the
                # handler would be a no-op; releasing here lets the next case start clean.
                if needs_cache_cleanup:
                    torch.cuda.empty_cache()


if __name__ == "__main__":
    main()
