feat(ops): add a selectable Triton backend, and measure the real kernel surface

Groundwork for replacing TileLang with Triton on the H20 path, plus the first ported kernel.

The scope is measured rather than guessed. `DEEPINFER_TRACE_KERNELS=<path>` proxies every `*_tl`
factory in the dispatcher and records its first use; one full text-generation run on 8xH20 says
exactly 24 kernels execute. `H20-TRITON-PLAN.md` lists them, in dependency order.

`DEEPINFER_OPS_BACKEND` selects the backend package (unset keeps the platform default, so nothing
changes by omission) and `DEEPINFER_OPS_FALLBACK` names individual kernels that must stay on the
default. Kernels the selected backend does not export at all are taken from the default with a
one-line notice, which lets a partial port run end to end the whole way through instead of only
at the end.

That coexistence is the reason for a parallel backend rather than an in-place rewrite: the
TileLang kernel is a validated baseline, so each Triton kernel can be diffed against it on real
shapes through the same public interface function, which is a much stronger check than comparing
against a fresh Python reference. `scripts/h20_ops_ab.py` does that swap and reports per case.

First kernel: `gather_tl`, all 9 cases bitwise identical (zeros filled for out-of-range indices,
int8/bfloat16 payloads, strided rows, int32/int64 indices, empty and single-row inputs).

H20-TRITON-PLAN.md

@@ -0,0 +1,107 @@
+# H20 Triton 改造计划
+
+把 DeepInfer 在 H20（SM90）上推理所需的 TileLang kernel 全部换成 Triton。
+
+分支：`h20-sm90-triton`，基于 `h20-sm90-support`（PR #180）。
+
+## 一、范围是怎么定下来的
+
+不靠读代码猜，而是实测。`ops.py` 里加了一个 `DEEPINFER_TRACE_KERNELS=<path>` 探针
+（`_KernelTrace` 代理每个 `*_tl` 工厂，第一次被调用时记名），然后跑一次完整文本生成，
+得到真实执行面。
+
+**24 个 kernel**，就是本计划要换的全部：
+
+| # | kernel | 文件 | 阶段 |
+|---|---|---|---|
+| 1 | `gather_tl` | `gather.py` | P1 |
+| 2 | `gather_input_ids_zos_tl` | `gather_input_ids_zos.py` | P1 |
+| 3 | `gather_engram_tokens_zos_tl` | `gather_engram_tokens_zos.py` | P1 |
+| 4 | `expand_swa_index_tl` | `swa_index.py` | P1 |
+| 5 | `kv_index_cat_tl` | `kv_index_cat.py` | P1 |
+| 6 | `convert_kv_index_tl` | `flash_mla_utils.py` | P1 |
+| 7 | `load_kv_index_tl` | `get_kv_index.py` | P1 |
+| 8 | `map_lookup_tl` | `sampler.py` | P1 |
+| 9 | `norm_tl` | `norm.py` | P1 |
+| 10 | `copy_block_tables_tl` | `prepare_decode_metadata.py` | P1 |
+| 11 | `prepare_v4_decode_metadata_tl` | `prepare_decode_metadata.py` | P1 |
+| 12 | `expand_sampling_params_tl` | `sampler.py` | P1 |
+| 13 | `store_mtp_info_tl` | `store_mtp_info.py` | P1 |
+| 14 | `store_kv_cache_tl` | `store_kv_cache.py` | P2 |
+| 15 | `load_paged_kv_cache_tl` | `load_paged_kv_cache.py` | P2 |
+| 16 | `fused_norm_rope_hadamard_cast_store_kvcache_tl` | `fused_transform.py` | P2 |
+| 17 | `prepare_verification_metadata_tl` | `sampler.py` | P3 |
+| 18 | `get_accept_result_tl` | `sampler.py` | P3 |
+| 19 | `finalize_verification_tl` | `sampler.py` | P3 |
+| 20 | `remap_draft_topk_tl` | `sampler.py` | P3 |
+| 21 | `confidence_head_fused_tl` | `dflash.py` | P3 |
+| 22 | `dflash_input_store_tl` | `dflash.py` | P3 |
+| 23 | `get_fused_dflash_get_input_tl` | `get_kv_index.py` | P3 |
+| 24 | `fused_sample_tl` | `fused_sample.py` | P4 |
+
+不在此列的（vision、structured output、engram 主表、compress、candidate_blocks 等）
+在这条路径上没被执行，不在本次范围。
+
+## 二、关键架构决策：并行后端，而不是原地替换
+
+**不做原地改写。** 新建 `tilelang_ops/triton/` 包，与 `cuda/` 平级，暴露同名 `*_tl` 工厂；
+`ops.py` 按 `DEEPINFER_OPS_BACKEND` 环境变量选后端（默认仍是 `cuda`，保持现状不回归）。
+
+理由：
+
+1. **每个 kernel 都能 A/B 对照**。同一输入跑 TileLang 版和 Triton 版，比输出。这比只和
+   Python ref 比强得多 —— TileLang 版是已经在整机验证过的基线。
+2. **失败可回退**。某个 kernel 没达标时，可以单独把那个名字指回 cuda 后端（见下），
+   其余 23 个照样走 Triton，端到端仍然能跑。这样交付是渐进的，不是全有或全无。
+3. **review 友好**。PR 里能直接看到「同一契约、两套实现」，而不是一大片删除。
+
+`ops.py` 的解析顺序：`DEEPINFER_OPS_BACKEND`（整包） →
+`DEEPINFER_OPS_FALLBACK`（逗号分隔的 kernel 名，回退到 cuda）。
+
+## 三、验证协议（每个 kernel 都必须过三关）
+
+1. **契约一致**：Triton 版与 TileLang 版在真实 shape 上，输出逐元素一致（浮点按 dtype
+   容差；索引/计数类要求精确相等）。
+2. **边界覆盖**：空输入、`-1` padding、跨页、多段、奇偶长度、超出范围的值 ——
+   按各 kernel 自己的语义，从 TileLang 版的测试用例里继承（`cuda/*.py` 里已有
+   `test_*` 参数化矩阵，直接复用）。
+3. **不退化**：不能因为它们跑通就悄悄降级语义（例如把排序改成不排序、把原子累加
+   改成串行），必须保持与 TileLang 版相同的契约。
+
+## 四、阶段与顺序
+
+依赖顺序决定阶段顺序：底层 gather/metadata 在前，cache 在后，采样最后。
+
+- **P0 基础设施**：`triton/` 包骨架、后端选择与 fallback、A/B 对照脚本。
+- **P1 轻量 kernel（13 个）**：纯搬运/索引/逐元素，无矩阵乘。先易后难建立信心。
+- **P2 KV cache（3 个）**：量化写、分页读、fused norm+rope+hadamard+cast+store。
+  `store_kv_cache` / `load_paged_kv_cache` 已有闭环 round-trip 脚本可复用
+  （`scripts/h20_kv_roundtrip.py`），直接拿来验 Triton 版。
+- **P3 MTP 链路（7 个）**：验证元数据、接受判定、draft top-k 重映射、confidence head。
+  这条链路只在 MTP 开启时执行，需要额外用默认配置跑。
+- **P4 采样（1 个）**：`fused_sample_tl`，最复杂的一个（top-k + 温度 + 拒绝采样）。
+- **P5 端到端**：`h20_ask.py` 输出与 TileLang 版逐 token 一致，两种配置都过
+  （`DEEPINFER_USE_NEXTN_DECODE=0` 和默认开启）。
+
+## 五、已知风险与对策
+
+| 风险 | 对策 |
+|---|---|
+| `T.pdl_sync()`（Hopper PDL）在 Triton 没有直接对应 | 该同步保证后续 kernel 读到已写数据；Triton 侧去掉依赖 PDL 的假设，改用 stream 顺序语义。若某 kernel 依赖 PDL 才能正确，就在 Python 侧显式 `wait_event`，不靠隐式 |
+| TileLang `T.dynamic(...)` 动态 shape | Triton 用 constexpr 分块 + 运行期循环边界，shape 从 tensor 推导；避免为每个 shape 编译 |
+| `T.StridedTensor` 非连续视图 | Triton 传入 stride 作为参数，或先 `contiguous()`（会记录在注释里并计入性能） |
+| 原子操作 | `atomic_add` 在 Triton 可用，但需注意非确定性；若原 kernel 有确定性保证，需复现（例如按固定顺序 reduce） |
+| fp8 / fp4 位级操作 | Triton 支持 `tl.float8e4nv` 等；fp4 打包需手工位运算，参考 `quant.py` 的现有实现 |
+| 编译期开销 | Triton 首次编译慢；用 `tl.constexpr` 收敛变体数量，必要时预编译缓存 |
+
+## 六、完成判据
+
+- 24 个 kernel 全部有 `triton/` 实现，且 A/B 对照通过；
+- `DEEPINFER_OPS_BACKEND=triton` 下 `h20_ask.py` 端到端输出与 TileLang 版**逐 token 相同**
+  （MTP 开/关两种配置）；
+- `ruff check --no-cache` 与 `ruff format --check --no-cache` 通过；
+- 性能不要求超过 TileLang，但需测量并记录差距（附 benchmark 表）。
+
+## 七、进度
+
+见文末「执行日志」。

deepinfer/model_executor/layers/tilelang_ops/ops.py

@@ -3,23 +3,115 @@
 ``ops.xxx_tl``; backend selection happens only in this file.
 
 Re-exports all ``xxx_tl`` factories of the selected backend.
+
+Backend resolution:
+
+* ``DEEPINFER_OPS_BACKEND`` picks the backend package for every kernel. Unset means the platform
+  default, which is what the engine has always used, so nothing changes unless asked for.
+* ``DEEPINFER_OPS_FALLBACK`` is a comma-separated list of kernel names that must come from the
+  platform default even when another backend is selected. It exists so a partially ported backend
+  can still run end to end: the ported kernels exercise the new code, the rest keep working, and
+  the comparison stays meaningful throughout the port.
+
+When the selected backend does not export a kernel at all, the default backend's implementation
+is used and a line naming it is printed once, so a partial port is never mistaken for a complete
+one.
 """
 
 import importlib
+import os
 import pkgutil
+import sys
+from pathlib import Path
 from types import ModuleType
 from typing import Any
 
 from deepinfer import PLATFORM
 
-_pkg = importlib.import_module(f"{__package__}.{PLATFORM.tilelang_ops_backend}")
-for _info in pkgutil.iter_modules(_pkg.__path__):
-    _mod = importlib.import_module(f"{_pkg.__name__}.{_info.name}")
-    for _name, _obj in vars(_mod).items():
-        if _name.startswith("_") or isinstance(_obj, ModuleType) or not _name.endswith("_tl"):
-            continue
-        assert globals().get(_name, _obj) is _obj, f"backend symbol collision: {_name}"
-        globals()[_name] = _obj
+_TRACE = bool(os.getenv("DEEPINFER_TRACE_KERNELS"))
+
+
+def _record(name: str) -> None:
+    """Log one use of a backend kernel factory, so the real surface can be measured.
+
+    Called once per kernel *construction* rather than per launch: TileLang caches compiled
+    kernels, so construction happens at warmup and the count stays meaningful.
+    """
+    path = Path(os.environ["DEEPINFER_TRACE_KERNELS"])
+    with path.open("a") as f:
+        f.write(f"{name}\n")
+
+
+class _KernelTrace:
+    """Proxy a `*_tl` factory, recording its first use.
+
+    A plain function wrapper is not enough: callers read attributes off the factory (the TileLang
+    `@tilelang.jit` decorator attaches `pass_configs`, for one), so attribute access has to keep
+    working. Everything unknown is forwarded to the real object.
+    """
+
+    def __init__(self, name: str, factory: Any) -> None:
+        self._name = name
+        self._factory = factory
+        self._seen = False
+
+    def __call__(self, *args: Any, **kwargs: Any) -> Any:
+        if not self._seen:
+            self._seen = True
+            _record(self._name)
+        return self._factory(*args, **kwargs)
+
+    def __getattr__(self, item: str) -> Any:
+        return getattr(self._factory, item)
+
+    def __repr__(self) -> str:
+        return f"<traced {self._name}>"
+
+
+def _collect(backend: str) -> dict[str, Any]:
+    """Collect every ``*_tl`` factory exported by a backend package, by name."""
+    pkg = importlib.import_module(f"{__package__}.{backend}")
+    found: dict[str, Any] = {}
+    for info in pkgutil.iter_modules(pkg.__path__):
+        mod = importlib.import_module(f"{pkg.__name__}.{info.name}")
+        for name, obj in vars(mod).items():
+            if name.startswith("_") or isinstance(obj, ModuleType) or not name.endswith("_tl"):
+                continue
+            if name in found:
+                raise RuntimeError(f"backend {backend!r} exports {name!r} twice")
+            found[name] = obj
+    return found
+
+
+_DEFAULT_BACKEND = PLATFORM.tilelang_ops_backend
+_BACKEND = os.getenv("DEEPINFER_OPS_BACKEND") or _DEFAULT_BACKEND
+_FALLBACK = (
+    {name.strip() for name in os.getenv("DEEPINFER_OPS_FALLBACK", "").split(",") if name.strip()} if _BACKEND != _DEFAULT_BACKEND else set()
+)
+
+_SELECTED = _collect(_BACKEND)
+_DEFAULT: dict[str, Any] = _SELECTED if _BACKEND == _DEFAULT_BACKEND else _collect(_DEFAULT_BACKEND)
+
+_unknown = sorted(_FALLBACK - set(_SELECTED))
+if _unknown:
+    raise RuntimeError(f"DEEPINFER_OPS_FALLBACK names kernel(s) backend {_BACKEND!r} does not export: {_unknown}")
+
+for _name, _obj in _SELECTED.items():
+    globals()[_name] = _KernelTrace(_name, _obj) if _TRACE else _obj
+
+for _name in _FALLBACK:
+    globals()[_name] = _DEFAULT[_name]
+
+_INHERITED = sorted(set(_DEFAULT) - set(_SELECTED) - _FALLBACK)
+for _name in _INHERITED:
+    globals()[_name] = _DEFAULT[_name]
+
+if _INHERITED:
+    print(
+        f"[ops] backend {_BACKEND!r} lacks {len(_INHERITED)} kernel(s), using {_DEFAULT_BACKEND!r} for: " + ", ".join(_INHERITED),
+        file=sys.stderr,
+        flush=True,
+    )
 
 
 class KernelUnavailableError(AttributeError):

deepinfer/model_executor/layers/tilelang_ops/triton/common.py

@@ -0,0 +1,60 @@
+"""Shared helpers for the Triton kernel backend.
+
+The backend mirrors ``cuda/``: each module exports the same ``*_tl`` factory names with the same
+signatures, and each factory returns a callable that takes tensors by position and writes its
+result into an output argument. Call sites are untouched, which is what makes per-kernel A/B
+comparison possible -- the interface layer is shared, only the kernel differs.
+
+Triton-specific conventions used throughout:
+
+* Everything the kernel needs to know at compile time (shapes, strides, tile sizes) arrives as a
+  factory argument and is passed to the kernel as ``tl.constexpr``. Nothing is recompiled per
+  launch for shape reasons.
+* Runtime lengths are passed as ordinary arguments so a single compiled kernel covers many cases.
+* These kernels make no assumptions about program execution order beyond what CUDA provides on a
+  single stream, so the Hopper programmatic-dependent-launch (``T.pdl_sync``) hints in the
+  TileLang versions have no counterpart and are not needed: consecutive launches on one stream are
+  already ordered.
+"""
+
+from typing import Any
+
+import torch
+import triton.language as tl
+
+_TRITON_DTYPES: dict[torch.dtype, Any] = {
+    torch.int8: tl.int8,
+    torch.uint8: tl.uint8,
+    torch.int16: tl.int16,
+    torch.int32: tl.int32,
+    torch.int64: tl.int64,
+    torch.float16: tl.float16,
+    torch.bfloat16: tl.bfloat16,
+    torch.float32: tl.float32,
+    torch.float64: tl.float64,
+    torch.float8_e4m3fn: tl.float8e4nv,
+    torch.float8_e5m2: tl.float8e5,
+    # Packed FP4 is carried as raw bytes; kernels that inspect nibbles do it themselves.
+    torch.bool: tl.int1,
+}
+
+
+def triton_dtype(dtype: torch.dtype) -> Any:
+    """Map a torch dtype to its Triton counterpart."""
+    try:
+        return _TRITON_DTYPES[dtype]
+    except KeyError as exc:  # pragma: no cover - only on a new dtype
+        raise NotImplementedError(f"triton backend has no mapping for {dtype}") from exc
+
+
+def next_power_of_2(n: int) -> int:
+    return 1 if n <= 1 else 1 << (n - 1).bit_length()
+
+
+def cdiv(a: int, b: int) -> int:
+    return -(-a // b)
+
+
+def grid2(n: int, m: int) -> tuple[int, int]:
+    """2D launch grid; Triton rejects a zero-sized dimension, so clamp to one block."""
+    return max(1, n), max(1, m)

deepinfer/model_executor/layers/tilelang_ops/triton/gather.py

@@ -0,0 +1,85 @@
+"""Row gather ``y = x[index]`` with out-of-range indices zero-filled, in Triton."""
+
+from typing import Any
+
+import torch
+import triton
+import triton.language as tl
+
+from deepinfer.model_executor.layers.tilelang_ops.triton.common import cdiv, next_power_of_2, triton_dtype
+
+
+@triton.jit
+def _gather_kernel(
+    X,
+    INDEX,
+    Y,
+    n,
+    m,
+    D,
+    x_stride,
+    BLOCK_M: tl.constexpr,
+    BLOCK_D: tl.constexpr,
+):
+    pid_m = tl.program_id(0)
+    pid_d = tl.program_id(1)
+
+    offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
+    offs_d = pid_d * BLOCK_D + tl.arange(0, BLOCK_D)
+    row_in_range = offs_m < m
+
+    # `other=-1` marks masked rows as out of range, which the validity test already rejects, so
+    # no separate check is needed for the tail of the last block.
+    idx = tl.load(INDEX + offs_m, mask=row_in_range, other=-1)
+    row_valid = row_in_range & (idx >= 0) & (idx < n)
+
+    col_in_range = offs_d < D
+    src = X + idx[:, None] * x_stride + offs_d[None, :]
+    values = tl.load(src, mask=row_valid[:, None] & col_in_range[None, :], other=0)
+
+    dst = Y + offs_m[:, None] * D + offs_d[None, :]
+    # Rows past the end of `y` are not written; rows whose index was invalid are written as 0,
+    # which is what makes the zero-fill observable rather than left as uninitialised memory.
+    tl.store(dst, values, mask=row_in_range[:, None] & col_in_range[None, :])
+
+
+def gather_tl(
+    D: int,
+    dtype: torch.dtype,
+    index_dtype: torch.dtype,
+    x_stride: int,
+    threads: int = 256,
+) -> Any:
+    """Build the gather kernel for a fixed row width, dtype and source stride."""
+    out_dtype = triton_dtype(dtype)
+    idx_dtype = triton_dtype(index_dtype)
+
+    block_d = min(next_power_of_2(D), 1024)
+    # One warp-group's worth of rows per block keeps coalescing good for narrow rows; wide rows
+    # want fewer rows so the register tile stays bounded.
+    block_m = 64 if block_d <= 128 else 16
+
+    def launch(x: torch.Tensor, index: torch.Tensor, y: torch.Tensor) -> None:
+        m = index.shape[0]
+        n = x.shape[0]
+        if m == 0:
+            return
+        grid = (cdiv(m, block_m), cdiv(D, block_d))
+        _gather_kernel[grid](
+            x,
+            index,
+            y,
+            n,
+            m,
+            D,
+            x_stride,
+            BLOCK_M=block_m,
+            BLOCK_D=block_d,
+            num_warps=max(1, min(8, threads // 32)),
+        )
+
+    # Kept for parity with the TileLang factory, which also closes over these; the interface layer
+    # asserts the dtypes match before calling, so nothing here needs to re-check.
+    launch.out_dtype = out_dtype  # type: ignore[attr-defined]
+    launch.idx_dtype = idx_dtype  # type: ignore[attr-defined]
+    return launch

pyproject.toml

@@ -76,6 +76,12 @@ ignore = [
     "B023",
 ]
 
+[tool.ruff.lint.per-file-ignores]
+# Triton kernel bodies: parameters are pointers and scalars that Triton itself types, and the
+# kernels are launched through `kernel[grid](...)` rather than called directly, so the usual
+# signature annotations would be either wrong or rejected by the JIT.
+"deepinfer/model_executor/layers/tilelang_ops/triton/*.py" = ["ANN001", "ANN201", "ANN202"]
+
 [tool.ruff]
 line-length = 140
 include = ["./**/*.py"]

scripts/ab_cases.py

@@ -0,0 +1,53 @@
+"""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.gather import gather
+
+
+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)),
+    ]
+
+
+CASES: dict[str, Callable[[], list[tuple[str, bool, Callable[[], Any]]]]] = {
+    "gather_tl": _cases_gather,
+}

scripts/h20_ops_ab.py

@@ -0,0 +1,132 @@
+"""Compare a Triton kernel against the TileLang one it replaces, through the shared interface.
+
+Every interface module (``gather.py``, ``norm.py``, ...) calls ``ops.<name>_tl`` and then uses the
+result. That means the strongest available check is not "does Triton match a Python reference"
+but "does the same public call give the same answer when the kernel behind it changes". This
+script does exactly that: it swaps ``ops.<name>_tl`` between the two backends and compares the
+output of the public function.
+
+    python scripts/h20_ops_ab.py            # every case
+    python scripts/h20_ops_ab.py gather norm
+"""
+
+from __future__ import annotations
+
+import argparse
+import sys
+from collections.abc import Callable
+from typing import Any
+
+import torch
+
+from deepinfer.model_executor.layers.tilelang_ops import ops
+
+# Each backend is a package whose submodules export `*_tl`; the packages' own `__init__` does not
+# re-export, so collect the factories the same way the dispatcher does.
+CUDA = ops._collect("cuda")  # noqa: SLF001 - the collector is the one place that knows the layout
+TRITON = ops._collect("triton")  # noqa: SLF001
+
+
+def swap(name: str, backend: str) -> Callable[[], None]:
+    """Point ``ops.<name>`` at one backend, returning a function that restores it."""
+    previous = getattr(ops, name)
+    source = CUDA if backend == "cuda" else TRITON
+    if name not in source:
+        raise AttributeError(f"backend {backend!r} does not export {name!r}")
+    setattr(ops, name, source[name])
+
+    def restore() -> None:
+        setattr(ops, name, previous)
+
+    return restore
+
+
+def run_both(name: str, fn: Callable[[], Any]) -> tuple[Any, Any]:
+    """Run ``fn`` once against each backend, returning (cuda result, triton result)."""
+    restore = swap(name, "cuda")
+    try:
+        expected = fn()
+    finally:
+        restore()
+
+    restore = swap(name, "triton")
+    try:
+        actual = fn()
+    finally:
+        restore()
+    return expected, actual
+
+
+def compare(name: str, expected: Any, actual: Any, *, exact: bool, rtol: float = 1e-5, atol: float = 1e-6) -> bool:
+    """Report whether one case agrees; returns True on success."""
+    if isinstance(expected, torch.Tensor):
+        if expected.dtype != actual.dtype:
+            print(f"  FAIL {name}: dtype {expected.dtype} vs {actual.dtype}")
+            return False
+        if expected.shape != actual.shape:
+            print(f"  FAIL {name}: shape {tuple(expected.shape)} vs {tuple(actual.shape)}")
+            return False
+        if exact or not expected.is_floating_point():
+            same = torch.equal(expected, actual)
+            if not same:
+                bad = expected != actual
+                print(f"  FAIL {name}: {int(bad.sum())}/{expected.numel()} elements differ")
+                idx = bad.nonzero()[0].tolist()
+                print(f"        first at {idx}: {expected[tuple(idx)].item()} vs {actual[tuple(idx)].item()}")
+            return bool(same)
+        close = torch.allclose(expected.float(), actual.float(), rtol=rtol, atol=atol)
+        if not close:
+            diff = (expected.float() - actual.float()).abs()
+            scale = expected.float().abs().clamp_min(1e-12)
+            print(f"  FAIL {name}: max abs {diff.max().item():.3e}, max rel {(diff / scale).max().item():.3e}")
+        return bool(close)
+
+    same = expected == actual
+    if not same:
+        print(f"  FAIL {name}: {expected!r} vs {actual!r}")
+    return bool(same)
+
+
+def cases() -> dict[str, Callable[[], list[tuple[str, bool, Callable[[], Any]]]]]:
+    """Registry of per-kernel case builders, filled in as kernels are ported."""
+    from scripts.ab_cases import CASES
+
+    return CASES
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("kernels", nargs="*", help="subset of kernel names; default is all ported")
+    parser.add_argument("--verbose", action="store_true")
+    args = parser.parse_args()
+
+    registry = cases()
+    selected = args.kernels or sorted(registry)
+    unknown = [name for name in selected if name not in registry]
+    if unknown:
+        raise SystemExit(f"no A/B cases for: {', '.join(unknown)}")
+
+    failures = 0
+    for kernel in selected:
+        print(f"[{kernel}]")
+        for case_name, exact, fn in registry[kernel]():
+            try:
+                expected, actual = run_both(kernel, fn)
+            except Exception as exc:  # noqa: BLE001 - a raise on one backend is itself a result
+                print(f"  FAIL {case_name}: {type(exc).__name__}: {str(exc).splitlines()[0][:160]}")
+                failures += 1
+                continue
+            if not compare(case_name, expected, actual, exact=exact):
+                failures += 1
+            elif args.verbose:
+                print(f"  ok   {case_name}")
+
+    print()
+    if failures:
+        print(f"AB_FAILED {failures}")
+        sys.exit(1)
+    print("AB_OK")
+
+
+if __name__ == "__main__":
+    main()