diff --git a/H20-TRITON-PLAN.md b/H20-TRITON-PLAN.md
new file mode 100644
index 0000000..42b7b79
--- /dev/null
+++ b/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 表）。
+
+## 七、进度
+
+见文末「执行日志」。
diff --git a/deepinfer/model_executor/layers/tilelang_ops/ops.py b/deepinfer/model_executor/layers/tilelang_ops/ops.py
index 88546e6..5835467 100644
--- a/deepinfer/model_executor/layers/tilelang_ops/ops.py
+++ b/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):
diff --git a/deepinfer/model_executor/layers/tilelang_ops/triton/__init__.py b/deepinfer/model_executor/layers/tilelang_ops/triton/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/deepinfer/model_executor/layers/tilelang_ops/triton/common.py b/deepinfer/model_executor/layers/tilelang_ops/triton/common.py
new file mode 100644
index 0000000..b439dff
--- /dev/null
+++ b/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)
diff --git a/deepinfer/model_executor/layers/tilelang_ops/triton/gather.py b/deepinfer/model_executor/layers/tilelang_ops/triton/gather.py
new file mode 100644
index 0000000..f0cf58b
--- /dev/null
+++ b/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
diff --git a/pyproject.toml b/pyproject.toml
index f448e48..909f4f6 100644
--- a/pyproject.toml
+++ b/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"]
diff --git a/scripts/ab_cases.py b/scripts/ab_cases.py
new file mode 100644
index 0000000..8a13fcd
--- /dev/null
+++ b/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,
+}
diff --git a/scripts/h20_ops_ab.py b/scripts/h20_ops_ab.py
new file mode 100644
index 0000000..8a36b6c
--- /dev/null
+++ b/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()