feat(h20): enable FlagGems end to end, from an allow-list of ops proven equal to PyTorch

`DEEPINFER_ENABLE_FLAGGEMS=1` turns it on. Verified on all 8 DP ranks: exit=0, 8/8 ranks agreeing,
no kernel fallbacks, `hard failures: none`, correct answers. 19 ops actually dispatched, 152 call
sites, 43 registrations from 37 requested names.

The default is `flag_gems.only_enable(include=[...])`, not `enable()`, and that is the whole design:

* allow-list -- every intermediate state is safe (unproven ops stay on PyTorch, which is correct by
  construction); an incomplete list costs throughput, not correctness.
* deny-list -- I tried it first. Five full engine runs, one new defect each: `cat` (e8m0 dtype),
  `mul` (broadcast shape), `topk` (compile blowup), `rms_norm` (signature), `mm` (compile assert).
  The method requires being bitten before it can exclude, and every state along the way is unsafe.

That matches the instruction this was built to: use what works, leave the rest.

Where it is called is not incidental. `import flag_gems` initialises a CUDA context -- measured,
`torch.cuda.is_initialized()` flips on the import alone, before `enable()`. A `sitecustomize`
injector therefore put all 8 ranks on device 0 and NCCL refused with "Multiple Ranks are using the
same GPU". It now runs inside `create_llm`, after `setup_cuda_visible_devices`, which is the same
hazard already documented there for the profiler.

Defects found and documented rather than worked around, because each one is worth knowing:

* `mul` returns the *first operand's shape* when broadcasting expands it: (1,64,8)*(4,1,8) gives
  (1,64,8), and (64,1)*(1,128) gives (64,1). Silent, no error, fewer elements computed. It made
  `engram_hash_ref` disagree with PyTorch during a real run. add/sub/div/maximum/minimum all
  broadcast correctly, so it is mul's own bug.
* `cumsum` drifts by up to 1.0 over a 512-element bf16 row (tolerance ~0.17) -- bf16 accumulation.
* `topk`'s compile time explodes with k: 4.8 s at k=8, 19 s at k=64, >180 s at k=512, and the
  sampler asks for 4096. py-spy put all 8 ranks at 100% CPU inside `make_llir` with every GPU idle.
* `cat` raises `KeyError: 'float8_e8m0fnu'` on the model's UE8M0 scale tensors during weight load.
* `mm` passes every isolated probe yet fails inside the engine at `general_mm`; trigger unknown, so
  GEMMs stay on cuBLAS, which is faster anyway.

Two methodological errors in my own measurements, both fixed and both recorded:

1. After `enable()`, `torch.mul` *is* `flag_gems.mul`, so the first version of the mul A/B compared
   FlagGems with itself and reported a perfect score. All references are now computed before enable.
2. "ok" did not mean FlagGems ran. `torch.bitwise_xor` has no FlagGems registration at all, so that
   "match" was PyTorch against PyTorch. `screen_ops.py` now checks registration first and reports
   such cases as `skip` rather than pass.

Name mismatch is its own trap: FlagGems registers `add` bare but `eq.Scalar`, `where.self`,
`repeat_interleave.self_int`, `_softmax` (not `softmax`), `div.true_divide`. Asking for `softmax`
silently registers nothing. The bootstrap now expands names against `flag_gems._FULL_CONFIG` and
warns on anything unresolved -- which is how `softmax` was found to need `_softmax`.

Output is not token-identical to the no-FlagGems build (diverges at the first token), but two runs of
the same FlagGems config are bit-identical, so the difference is Triton-versus-eager accumulation
order, not nondeterminism. Answers stay correct and coherent.

H20-BRANCHES.md

@@ -23,7 +23,7 @@ main (8b16149)
 | `h20-sm90-triton` | `73b8090` | 引擎侧 29 个 kernel 加 Triton 实现（17 个可用），`DEEPINFER_OPS_BACKEND=triton` 切换 | 融合算子库 `tile_kernels` 仍要 TileLang |
 | `h20-sm90-triton-flagtree` | `c6a65fe` | 在 flagtree 镜像（torch 2.11 + cu13）里重编 `deepinfer._C` / `deep_gemm` / `deep_ep` / `deep_select`；`tile_kernels` 加 compat 层走 torch 参考，从而不 import TileLang；修掉参考实现在 sync 守卫下抛错的问题 | attention 仍要 TileLang |
 | `h20-sm90-sparse-mla-triton` | `6d337a8` | **attention 也换成 FlagGems 的 Triton kernel，于是「屏蔽 tilelang 也能端到端」成立** | 采样器的 rejection-sampling 变体仍要 TileLang（关 MTP 可绕）；不与 TileLang 版逐 token 一致 |
-| **`h20-sm90-flaggems`** | **进行中** | **把 `tile_kernels` 的 torch 参考换成 FlagGems 的 Triton 实现**（先做已验证的 per-token quant） | 刚起步；`mhc_*` / `engram_*` 多数没有 1:1 对应 |
+| **`h20-sm90-flaggems`** | **进行中** | FlagGems 白名单启用已端到端跑通（19 算子、8/8 DP、0 回退）；per-token quant 已验证逐位相同 | 只覆盖点状/reduction 算子；GEMM 留 cuBLAS；token 级与不用时不同但确定性 |
 
 ## `h20-sm90-sparse-mla-triton` 到底做了什么
 

H20-FLAGGEMS-PLAN.md

@@ -78,6 +78,66 @@ use_tma_aligned_col_major_sf=True)`，5 个形状：
 
 `tilelang_ops/` 已经有 Triton 后端，`DEEPINFER_OPS_BACKEND=triton` 切换，与 flag_gems 无关。
 
+## 端到端已跑通（2026-10-26）：白名单模式
+
+`DEEPINFER_ENABLE_FLAGGEMS=1` → `deepinfer/flaggems_bootstrap.py` 在 `create_llm` 里、
+`setup_cuda_visible_devices()` **之后**调 `flag_gems.only_enable(include=[...])`。
+
+结果：
+
+| | 值 |
+|---|---|
+| exit code | **0** |
+| 后端 | sparse MLA = triton，FlagGems mode=allow |
+| DP 一致 | **8/8** |
+| kernel 回退 | **无** |
+| hard failures | **none** |
+| 回答 | 「中国的首都→北京」「日本的首都→东京」，带分段格式 |
+| 实际派发算子 | **19 个**、152 处调用点 |
+| 注册 | 37 个请求名 → 96 个变体 → 43 个实际注册 |
+| 警告 | 0 |
+
+实际在跑的是 `abs/amax/bincount/clamp(3变体)/cos/exp/gather/index_select/neg/nonzero/
+reciprocal/relu/rsqrt/sigmoid/sin/sqrt/stack` 这些——点状与 reduction 算子，正是 eager
+PyTorch 一次一个 launch 吃亏、Triton 能合并的地方。
+
+### 为什么改成白名单，而不是「全开 + 排除坏的」
+
+我先按「全开+排除」做了 5 轮，每轮都在完整引擎里撞出一个新问题，每轮约 5 分钟：
+`cat`(e8m0 dtype) → `mul`(广播形状) → `topk`(编译爆炸) → `rms_norm`(签名) → `mm`(编译断言)。
+这条路的性质是「必须先被坑到才能修」，且中途任何状态都是不安全的。
+
+白名单相反：**任何中间状态都是安全的**（没验证的算子留在 PyTorch），不完整只是性能问题，
+不是正确性问题。这正好对应「能用就用、不能用就拉到」。
+
+### 踩到的两个方法论错误（都已修）
+
+1. **`torch.mul` 被替换后不能再当参考**。第一版 `ab_mul_dtype.py` 在 `enable()` 之后同时用了
+   `torch.mul` 和 `flag_gems.mul`，于是「bf16 mul 0 差异」是自比自。现在所有参考值都在
+   `enable()` **之前**算好并 clone。
+2. **「ok」可能根本没走 FlagGems**。`torch.bitwise_xor` 在 FlagGems 里**没有实现**，所以
+   「bitwise_xor 通过」是 PyTorch 比 PyTorch。`screen_ops.py` 现在先判断该算子是否真的被注册，
+   否则标 `skip` 并说明「PyTorch answered」。`square` 同理（FlagGems 用别的名字）。
+
+### 名字对不上是坑
+
+FlagGems 的注册名是混合风格：`add`/`mul` 是裸名，但 `eq.Scalar`、`lt.Tensor`、`where.self`、
+`repeat_interleave.self_int`、`_softmax`（不是 `softmax`）、`div.true_divide`、`pow.pow_scalar`
+是 aten 风格的。只写 `softmax` 会**静默什么都不注册**。现在 `_expand()` 对着
+`flag_gems._FULL_CONFIG` 自己做前缀展开，并在匹配不到时打印警告（这次就是靠它发现 `softmax`
+应写 `_softmax`）。
+
+### 仍未解决
+
+- **`mm` 家族**：隔离探测（M/N/K 扫描、NT 布局、batch、addmm）全部正确，但引擎里仍在
+  `general_mm` 撞 `accumulator shape is incompatible`。触发条件没定位，因此 GEMM 继续走 cuBLAS
+  （本来也更快）。
+- **token 级不一致，但是确定性的**：与不用 FlagGems 的版本对比，第 1 个 token 就分叉（上一版
+  sparse-MLA-only 是前 15 个一致）。但**同配置连跑两次逐 token 完全相同**，所以差异来自
+  Triton kernel 与 eager PyTorch 的累加顺序，不是不确定性。回答仍然正确连贯，与 SM90 已知
+  2.7% logits 漂移同源，被贪心解码放大。
+- **性能完全没测**。白名单里只有 19 个算子真被用到，收益未知。
+
 ## 建议的推进顺序
 
 每一步都是「先 A/B 单测、通过才接线」，沿用 sparse MLA 那次的规矩。

deepinfer/flaggems_bootstrap.py

@@ -0,0 +1,227 @@
+"""Optionally route PyTorch operators through FlagGems, via its documented `enable()` / `only_enable()`.
+
+FlagGems' usage docs (`docs/content/en/usage/basic.md`, `.../selective.md`) prescribe:
+
+    import flag_gems
+    flag_gems.enable()                      # all supported ops
+    flag_gems.only_enable(include=[...])    # only the named ops
+
+This module defaults to the **allow-list** form, because flag_gems 5.0.2 on SM90 has operators that are
+outright wrong, and a silently wrong operator is worse than a slow correct one. Only ops proven equal
+to PyTorch on this build are enabled; everything else stays on PyTorch, which is correct by
+construction. Measured by `scripts/flaggems_port/screen_ops.py`; extend `_VERIFIED_OPS` only after that
+screen (or an equivalent A/B) passes.
+
+Why an allow-list rather than "enable everything then exclude what breaks": a deny-list can only grow
+after each failure is hit inside a full engine run, and each cycle surfaces one operator. The
+allow-list is safe at every intermediate state, so an incomplete list is a performance question rather
+than a correctness one.
+
+Env:
+
+`DEEPINFER_ENABLE_FLAGGEMS=1`          turn it on
+`DEEPINFER_FLAGGEMS_MODE=allow|deny`   default `allow` (curated) vs `deny` (all but the broken ones)
+`DEEPINFER_FLAGGEMS_INCLUDE=a,b`       extra ops for allow mode
+`DEEPINFER_FLAGGEMS_EXCLUDE=a,b`       extra ops to keep on PyTorch, in either mode
+`DEEPINFER_FLAGGEMS_RECORD_DIR=<dir>`  write `<dir>/oplist-<pid>.log`, one line per call site
+
+**Where it is called matters more than what it does.** `import flag_gems` initialises a CUDA context
+(measured: `torch.cuda.is_initialized()` flips False -> True on the import alone, before `enable()`).
+If that happens before this process is pinned to its GPU, every rank initialises on device 0 and the
+engine fails at the first collective with
+
+    NCCL error: invalid usage
+    Last error: Multiple Ranks are using the same GPU/Partition.
+
+That is what a `sitecustomize` injector produced, because `sitecustomize` runs at interpreter startup
+while the per-rank `CUDA_VISIBLE_DEVICES` is set later, inside `create_llm`. The same hazard is already
+documented there for the profiler ("CUPTI 启动即初始化 CUDA，会使此后设置的 CUDA_VISIBLE_DEVICES 失效").
+Hence this is called from `create_llm` immediately after `setup_cuda_visible_devices`, and there is no
+`sitecustomize` for it.
+
+**Measured on flag_gems 5.0.2 / H20 (SM90), which is why the lists below look the way they do:**
+
+* `mul` returns the **first operand's shape** when broadcasting would expand it -- `(1,64,8)*(4,1,8)`
+  gives `(1,64,8)` instead of `(4,64,8)`; `(64,1)*(1,128)` gives `(64,1)`. Silent, no error, fewer
+  elements computed. Same-shape and scalar `mul` are exact. `add`/`sub`/`div`/`maximum`/`minimum` all
+  broadcast correctly, so it is `mul`'s own bug. It produced a wrong result inside `engram_hash_ref`
+  during a real engine run.
+* `cumsum` disagrees with PyTorch by up to 1.0 over a 512-element bf16 row where the tolerance is
+  ~0.17 -- bf16 accumulation where PyTorch uses fp32.
+* `rms_norm` requires `weight` positionally where `torch.nn.functional.rms_norm` allows omitting it,
+  and with `weight` supplied still differs by ~3e-2 on O(1) output.
+* `topk`'s compile time is superlinear in `k` and the sampler uses `k=4096`: k=8 -> 4.8 s, k=64 -> 19 s,
+  k=512 -> >180 s. The startup profile run never returned; `py-spy` put all 8 ranks at 100% CPU inside
+  `triton/backends/nvidia/compiler.py:make_llir` with every GPU idle.
+* `cat` raises `KeyError: 'float8_e8m0fnu'` from Triton's dtype table on the UE8M0 scale tensors this
+  model stores, during weight loading. Triton has no e8m0 element type.
+* the `mm` family passes every isolated probe (M/N/K swept, NT layout, batch, addmm) yet the engine
+  still hit `accumulator shape is incompatible` inside `general_mm`; trigger not identified, so GEMMs
+  stay on cuBLAS -- which is also the faster choice.
+"""
+
+from __future__ import annotations
+
+import os
+
+_ENV_ENABLE = "DEEPINFER_ENABLE_FLAGGEMS"
+_ENV_MODE = "DEEPINFER_FLAGGEMS_MODE"
+_ENV_INCLUDE = "DEEPINFER_FLAGGEMS_INCLUDE"
+_ENV_EXCLUDE = "DEEPINFER_FLAGGEMS_EXCLUDE"
+_ENV_RECORD_DIR = "DEEPINFER_FLAGGEMS_RECORD_DIR"
+
+# Ops measured equal to PyTorch by `scripts/flaggems_port/screen_ops.py`, across the shape variants
+# that screen covers (same-shape, broadcast, scalar, unary, reduction). Pointwise and reduction ops
+# are where FlagGems can actually win: one kernel launch instead of several eager calls. GEMMs are
+# deliberately absent -- cuBLAS wins there regardless.
+_VERIFIED_OPS: tuple[str, ...] = (
+    "abs",
+    "add",
+    "amax",
+    "bincount",
+    "bitwise_and",
+    "clamp",
+    "cos",
+    "div",
+    "eq",
+    "exp",
+    "gather",
+    "ge",
+    "gelu",
+    "index_select",
+    "lt",
+    "masked_fill",
+    "maximum",
+    "mean",
+    "minimum",
+    "neg",
+    "nonzero",
+    "pow",
+    "reciprocal",
+    "relu",
+    "remainder",
+    "repeat_interleave",
+    "rsqrt",
+    "sigmoid",
+    "silu",
+    "sin",
+    "_softmax",
+    "sort",
+    "sqrt",
+    "stack",
+    "sub",
+    "sum",
+    "where",
+)
+
+# Everything known broken on this build, documented and used by deny mode. Nothing here may move into
+# `_VERIFIED_OPS` without a passing A/B.
+_BROKEN_OPS: tuple[str, ...] = ("cat", "cumsum", "mul", "rms_norm", "topk")
+
+_enabled = False
+
+
+def enabled() -> bool:
+    """Whether FlagGems has been enabled in this process."""
+    return _enabled
+
+
+def _names(raw: str) -> list[str]:
+    return [op.strip() for op in raw.split(",") if op.strip()]
+
+
+def _expand(names: list[str], available: list[str]) -> tuple[list[str], list[str]]:
+    """Resolve plain op names to the names FlagGems actually registers.
+
+    FlagGems registers a mix of plain names (`add`, `mul`) and aten-style dotted or suffixed names:
+    `eq.Scalar`, `lt.Tensor`, `where.self`, `repeat_interleave.self_int`, `_softmax`, `div.true_divide`,
+    `pow.pow_scalar`. Asking for `eq` by name therefore registers nothing, silently -- caught here as a
+    warning, but the fix is to expand to every variant rather than spell out each one in
+    `_VERIFIED_OPS`.
+
+    Returns (expanded, unresolved). `available` is FlagGems' own registration table, so this cannot
+    drift from what the installed version actually offers.
+    """
+    expanded: list[str] = []
+    unresolved: list[str] = []
+    for name in names:
+        hits = [
+            n
+            for n in available
+            if n == name or n.startswith(f"{name}_") or n.startswith(f"{name}.")
+        ]
+        if hits:
+            expanded.extend(hits)
+        else:
+            unresolved.append(name)
+    # dict.fromkeys: preserve order, drop duplicates (a name may match several wanted entries).
+    return list(dict.fromkeys(expanded)), unresolved
+
+
+def maybe_enable() -> str:
+    """Enable FlagGems if asked to. Returns a one-line status; no-op when the env var is unset."""
+    global _enabled
+
+    if _enabled:
+        return "already enabled"
+    if not os.environ.get(_ENV_ENABLE):
+        return "disabled"
+
+    try:
+        import flag_gems
+    except Exception as exc:  # noqa: BLE001 - a requested-but-broken import must be loud
+        raise RuntimeError(
+            f"{_ENV_ENABLE} is set but flag_gems is not importable: {type(exc).__name__}: {exc}. "
+            f"Unset it, or fix the install (it also needs libcudart on LD_LIBRARY_PATH for Triton TLE)."
+        ) from exc
+
+    mode = (os.environ.get(_ENV_MODE, "allow").strip().lower() or "allow")
+    extra_in = _names(os.environ.get(_ENV_INCLUDE, ""))
+    extra_out = _names(os.environ.get(_ENV_EXCLUDE, ""))
+
+    log_kwargs: dict[str, object] = {}
+    record_dir = os.environ.get(_ENV_RECORD_DIR, "").strip()
+    if record_dir:
+        from pathlib import Path
+
+        Path(record_dir).mkdir(parents=True, exist_ok=True)
+        # Per-PID: `setup_flaggems_logging` opens the path with mode="w", so ranks sharing one path
+        # would leave only whichever wrote last.
+        log_kwargs = {
+            "record": True,
+            "path": str(Path(record_dir) / f"oplist-{os.getpid()}.log"),
+            "once": True,
+        }
+
+    include: list[str] | None = None
+    if mode == "allow":
+        excluded = set(extra_out)
+        wanted = [op for op in (*_VERIFIED_OPS, *extra_in) if op not in excluded]
+        # Expand against FlagGems' own table, not a hand-maintained list of alias names. Entries are
+        # variable-length tuples whose first field is the op name.
+        available = [entry[0] for entry in flag_gems._FULL_CONFIG]  # noqa: SLF001
+        include, unresolved = _expand(wanted, available)
+        flag_gems.only_enable(include=include, **log_kwargs)
+        how = f"only_enable({len(wanted)} wanted -> {len(include)} registered)"
+    elif mode == "deny":
+        unused = [op for op in (*_BROKEN_OPS, *extra_out)]
+        unresolved = []
+        flag_gems.enable(unused=unused, **log_kwargs)
+        how = f"enable(unused={unused})"
+    else:
+        raise RuntimeError(f"{_ENV_MODE}={mode!r} is not recognised; use 'allow' or 'deny'")
+
+    registered = set(flag_gems.all_registered_ops())
+    if unresolved:
+        # Nothing would be enabled for these, which reads as an unexplained slowdown rather than a
+        # failure. Say it once, loudly.
+        print(f"[flaggems] WARNING: no registered op matches: {unresolved}", flush=True)
+
+    _enabled = True
+    status = (
+        f"enabled pid={os.getpid()} flag_gems={flag_gems.__version__} "
+        f"device={flag_gems.device}/{flag_gems.vendor_name} mode={mode} {how} "
+        f"registered={len(registered)}"
+    )
+    print(f"[flaggems] {status}", flush=True)
+    return status

examples/_config.py

@@ -107,6 +107,13 @@ def create_llm(
 
     import torch
 
+    # After the device pin, not before: importing flag_gems initialises a CUDA context, so doing it
+    # earlier puts every rank on device 0 and NCCL then refuses with "Multiple Ranks are using the
+    # same GPU". No-op unless DEEPINFER_ENABLE_FLAGGEMS is set.
+    from deepinfer.flaggems_bootstrap import maybe_enable as maybe_enable_flaggems
+
+    maybe_enable_flaggems()
+
     from deepinfer.entrypoints.llm import LLM
 
     with torch.profiler.profile() if profile_init else nullcontext():

scripts/flaggems_port/ab_mul_dtype.py

@@ -0,0 +1,111 @@
+"""Is FlagGems' `mul` correct, and for which dtypes?
+
+Method note, because the first version of this script was wrong: once `flag_gems.enable()` runs it also
+captures `torch.mul`, so comparing `torch.mul(...)` with `flag_gems.mul(...)` compares FlagGems with
+itself. Every reference here is therefore computed **before** `enable()`, in the same process.
+
+The defect it measures: `engram_hash_ref` does `int64 (1,T,N) * int64 (L,1,N)`, and FlagGems returns a
+tensor of the *first operand's* shape rather than the broadcast shape, so most of the result is never
+written. Excluding `mul` from `enable()` restores the reference result exactly, which is what pointed
+here in the first place.
+
+Run: PYTHONPATH=. python scripts/flaggems_port/ab_mul_dtype.py
+"""
+
+from __future__ import annotations
+
+import sys
+from collections.abc import Callable
+
+import torch
+
+FAILED: list[str] = []
+
+
+def check(tag: str, cond: bool, detail: str = "") -> None:
+    print(f"  {'ok  ' if cond else 'BAD '} {tag}{'  ' + detail if detail else ''}")
+    if not cond:
+        FAILED.append(tag)
+
+
+def compare(tag: str, ref: torch.Tensor, got: torch.Tensor) -> None:
+    if ref.shape != got.shape:
+        check(tag, False, f"shape ref {tuple(ref.shape)} vs got {tuple(got.shape)}")
+        return
+    if ref.dtype != got.dtype:
+        check(tag, False, f"dtype ref {ref.dtype} vs got {got.dtype}")
+        return
+    if ref.dtype.is_floating_point:
+        diff = (ref.float() - got.float()).abs().max().item()
+        scale = ref.float().abs().max().item()
+        tol = 2.0**-8 * max(scale, 1e-9) * 2
+        check(tag, diff <= tol, f"max diff {diff:.3e} tol {tol:.3e}")
+        return
+    same = torch.equal(ref, got)
+    detail = ""
+    if not same:
+        bad = (ref != got).sum().item()
+        idx = (ref != got).nonzero()[:3]
+        sample = [(ref[tuple(i)].item(), got[tuple(i)].item()) for i in idx]
+        detail = f"{bad}/{ref.numel()} wrong, (ref, got) {sample}"
+    check(tag, same, detail)
+
+
+def main() -> int:
+    torch.manual_seed(0)
+    dev = "cuda"
+
+    # ---- references, computed while torch.mul is still PyTorch's ----
+    a = torch.randint(0, 1000, (1, 64, 8), device=dev, dtype=torch.int64)
+    b = torch.randint(1, 100, (4, 1, 8), device=dev, dtype=torch.int64)
+    ref_bcast = torch.mul(a, b).clone()
+
+    a_l = torch.randint(2**40, 2**41, (256,), device=dev, dtype=torch.int64)
+    ref_large = torch.mul(a_l, a_l).clone()
+
+    a_i32 = torch.randint(0, 4096, (1024,), device=dev, dtype=torch.int32)
+    ref_i32 = torch.mul(a_i32, a_i32).clone()
+
+    a_bf = torch.randn((2048, 4096), device=dev, dtype=torch.bfloat16) * 0.1
+    b_bf = torch.randn((2048, 4096), device=dev, dtype=torch.bfloat16) * 0.1
+    ref_bf16 = torch.mul(a_bf, b_bf).clone()
+
+    a_bf_s = torch.randn((4096,), device=dev, dtype=torch.bfloat16)
+    ref_bf16_scalar = torch.mul(a_bf_s, 3).clone()
+    ref_i64_scalar = torch.mul(a, 3).clone()
+
+    import flag_gems
+
+    flag_gems.enable()
+    print(f"flag_gems {flag_gems.__version__}: enable() done")
+    print(f"reference broadcast shape: {tuple(ref_bcast.shape)}\n")
+
+    cases: list[tuple[str, torch.Tensor, Callable[[], torch.Tensor]]] = [
+        ("int64 mul, broadcasting (1,64,8)*(4,1,8)", ref_bcast, lambda: flag_gems.mul(a, b)),
+        ("int64 mul, large values", ref_large, lambda: flag_gems.mul(a_l, a_l)),
+        ("int32 mul", ref_i32, lambda: flag_gems.mul(a_i32, a_i32)),
+        ("bf16 mul, 8M elements", ref_bf16, lambda: flag_gems.mul(a_bf, b_bf)),
+        ("bf16 mul by scalar", ref_bf16_scalar, lambda: flag_gems.mul(a_bf_s, 3)),
+        ("int64 mul by scalar", ref_i64_scalar, lambda: flag_gems.mul(a, 3)),
+    ]
+    for tag, ref, fn in cases:
+        got = fn()
+        torch.cuda.synchronize()
+        compare(tag, ref, got)
+
+    # torch.mul must show the same defect. If only the direct call did, the dispatcher would be
+    # fine and the fix would live elsewhere.
+    got_dispatch = torch.mul(a, b)
+    torch.cuda.synchronize()
+    compare("torch.mul (dispatched) shows the same defect", ref_bcast, got_dispatch)
+
+    print()
+    if FAILED:
+        print(f"DEFECTS: {FAILED}")
+        return 1
+    print("no defect found in the dtypes tested")
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())

scripts/flaggems_port/repro_engram_hash.py

@@ -0,0 +1,73 @@
+"""Reproduce the `engram_hash_ref` failure seen with FlagGems enabled, and bisect it.
+
+The end-to-end run failed inside `tile_kernels.torch.engram.engram_hash_ref` with
+
+    IndexError: index 1 is out of bounds for dimension 0 with size 1
+
+That function is a plain PyTorch reference, so either FlagGems changed an operator's semantics, or
+the engine's own call path is at fault and enabling FlagGems merely shifted when it surfaced. Those
+have different fixes, so this script answers "is it FlagGems" first, in-process, without a full
+engine run.
+
+Run with FLAGGEMS_REPRO_MODE=torch|full|unused:
+    torch  -- do not enable (control)
+    full   -- flag_gems.enable()
+    unused -- flag_gems.enable(unused=[...]) with FLAGGEMS_REPRO_EXCLUDE=op1,op2
+"""
+
+from __future__ import annotations
+
+import os
+import sys
+import traceback
+
+import torch
+from tile_kernels.torch.engram import engram_hash_ref
+
+MODE = os.environ.get("FLAGGEMS_REPRO_MODE", "torch")
+
+
+def enable() -> str:
+    if MODE == "torch":
+        return "flag_gems: off (control)"
+    import flag_gems
+
+    if MODE == "full":
+        flag_gems.enable()
+        return f"flag_gems: full enable, {len(flag_gems.all_registered_ops())} ops"
+    exclude = [op.strip() for op in os.environ.get("FLAGGEMS_REPRO_EXCLUDE", "").split(",") if op.strip()]
+    flag_gems.enable(unused=exclude)
+    return f"flag_gems: enable(unused={exclude})"
+
+
+def build(device: str = "cuda") -> tuple[torch.Tensor, ...]:
+    """Shapes from the engine's config rather than invented ones."""
+    num_ngram_layers = int(os.environ.get("REPRO_LAYERS", "1"))
+    max_ngram_size = int(os.environ.get("REPRO_NGRAM", "2"))
+    tables = 2
+    num_tokens = int(os.environ.get("REPRO_TOKENS", "16"))
+    vocab = 4096
+
+    torch.manual_seed(0)
+    # (num_tokens, max_ngram_size) per the docstring: unsqueeze(0) then broadcast against
+    # multipliers.unsqueeze(1) is what builds (layers, tokens, ngram).
+    ngram_token_ids = torch.randint(0, 1000, (num_tokens, max_ngram_size), device=device, dtype=torch.int32)
+    multipliers = torch.randint(1, 100, (num_ngram_layers, max_ngram_size), device=device, dtype=torch.int64)
+    vocab_sizes = torch.full((num_ngram_layers, max_ngram_size - 1, tables), vocab, device=device, dtype=torch.int32)
+    offsets = torch.zeros((num_ngram_layers, (max_ngram_size - 1) * tables), device=device, dtype=torch.int32)
+    return ngram_token_ids, multipliers, vocab_sizes, offsets
+
+
+if __name__ == "__main__":
+    print(enable(), flush=True)
+    args = build()
+    try:
+        out = engram_hash_ref(*args)
+        torch.cuda.synchronize()
+        print(f"OK   out {tuple(out.shape)} {out.dtype} finite={bool(torch.isfinite(out.float()).all())}")
+        # Print a checksum so runs with different exclude lists can be compared.
+        print(f"     checksum {int(out.to(torch.int64).abs().sum())}")
+    except Exception:
+        print("FAIL")
+        traceback.print_exc()
+        sys.exit(1)

scripts/flaggems_port/screen_ops.py

@@ -0,0 +1,228 @@
+"""Screen FlagGems ops for the defect class that broke `mul`, so the exclusion list is evidence-based.
+
+`mul` returns the *first operand's shape* whenever broadcasting would expand it:
+
+    (1,64,8) * (4,1,8)  ->  FlagGems (1,64,8),   PyTorch (4,64,8)
+    (64,1)   * (1,128)  ->  FlagGems (64,1),     PyTorch (64,128)
+
+Silent, no error. It broke `engram_hash_ref` inside the engine, and it is not dtype-specific.
+
+The traceback went through `flag_gems/utils/pointwise_dynamic.py` and `shape_utils.broadcast_shapes`,
+which is shared by every pointwise op -- so the question is which *other* ops inherit it. This script
+answers that: each op is run twice, same-shape and broadcasting, and the reference is always taken
+before `enable()` so the comparison cannot be FlagGems-against-itself.
+
+Output is a table, which is what the exclusion list in `deepinfer/flaggems_bootstrap.py` is built
+from.
+
+Run: PYTHONPATH=. python scripts/flaggems_port/screen_ops.py
+"""
+
+from __future__ import annotations
+
+import sys
+from collections.abc import Callable
+from dataclasses import dataclass
+from typing import Any
+
+import torch
+
+DEV = "cuda"
+
+
+@dataclass
+class Case:
+    op: str
+    label: str
+    run: Callable[[], Any]
+    ref: Any
+
+
+def ref_of(fn: Callable[[], Any]) -> Any:
+    out = fn()
+    if isinstance(out, torch.Tensor):
+        return out.clone()
+    if isinstance(out, (tuple, list)):
+        return type(out)(t.clone() if isinstance(t, torch.Tensor) else t for t in out)
+    return out
+
+
+def same(ref: Any, got: Any) -> tuple[bool, str]:
+    if isinstance(ref, torch.Tensor):
+        if not isinstance(got, torch.Tensor):
+            return False, f"type {type(got).__name__}"
+        if ref.shape != got.shape:
+            return False, f"shape {tuple(ref.shape)} -> {tuple(got.shape)}"
+        if ref.dtype != got.dtype:
+            return False, f"dtype {ref.dtype} -> {got.dtype}"
+        if ref.dtype.is_floating_point:
+            d = (ref.float() - got.float()).abs().max().item()
+            s = ref.float().abs().max().item()
+            tol = 2.0**-8 * max(s, 1e-9) * 2
+            return (d <= tol, f"max diff {d:.2e}")
+        eq = torch.equal(ref, got)
+        if eq:
+            return True, "exact"
+        return False, f"{(ref != got).sum().item()}/{ref.numel()} wrong"
+    if isinstance(ref, (tuple, list)):
+        for r, g in zip(ref, got, strict=False):
+            ok, msg = same(r, g)
+            if not ok:
+                return False, msg
+        return True, "exact"
+    return ref == got, ""
+
+
+def build_cases() -> list[Case]:
+    torch.manual_seed(0)
+    b = 64
+    n = 8
+
+    # (1, b, n) vs (4, 1, n): the exact shape pattern that broke engram_hash.
+    x3 = torch.randn((1, b, n), device=DEV, dtype=torch.bfloat16)
+    y3 = torch.randn((4, 1, n), device=DEV, dtype=torch.bfloat16)
+    i3 = torch.randint(0, 1000, (1, b, n), device=DEV, dtype=torch.int64)
+    j3 = torch.randint(1, 100, (4, 1, n), device=DEV, dtype=torch.int64)
+
+    # (b, 1) vs (1, m): the most common broadcast in transformer code (masks, scales).
+    x2 = torch.randn((b, 1), device=DEV, dtype=torch.bfloat16)
+    y2 = torch.randn((1, 128), device=DEV, dtype=torch.bfloat16)
+
+    xs = torch.randn((256, 512), device=DEV, dtype=torch.bfloat16)
+    ys = torch.randn((256, 512), device=DEV, dtype=torch.bfloat16)
+
+    cases: list[Case] = []
+
+    def add(name: str, label: str, fn: Callable[[], Any]) -> None:
+        cases.append(Case(name, label, fn, ref_of(fn)))
+
+    # binary pointwise, same shape (control: isolates broadcast as the cause)
+    add("add", "same-shape", lambda: torch.add(xs, ys))
+    add("mul", "same-shape", lambda: torch.mul(xs, ys))
+    add("sub", "same-shape", lambda: torch.sub(xs, ys))
+    add("div", "same-shape", lambda: torch.div(xs, ys))
+    add("maximum", "same-shape", lambda: torch.maximum(xs, ys))
+    add("minimum", "same-shape", lambda: torch.minimum(xs, ys))
+
+    # binary pointwise, broadcasting (the suspect class)
+    add("add", "broadcast 3d", lambda: torch.add(x3, y3))
+    add("mul", "broadcast 3d", lambda: torch.mul(x3, y3))
+    add("sub", "broadcast 3d", lambda: torch.sub(x3, y3))
+    add("div", "broadcast 3d", lambda: torch.div(x3, y3))
+    add("maximum", "broadcast 3d", lambda: torch.maximum(x3, y3))
+    add("minimum", "broadcast 3d", lambda: torch.minimum(x3, y3))
+    add("add", "broadcast 2d", lambda: torch.add(x2, y2))
+    add("mul", "broadcast 2d", lambda: torch.mul(x2, y2))
+    add("sub", "broadcast 2d", lambda: torch.sub(x2, y2))
+
+    # integer bitwise / remainder: what engram_hash also used
+    add("bitwise_xor_", "broadcast int64", lambda: i3.bitwise_xor(j3))
+    add("remainder", "broadcast int64", lambda: torch.remainder(i3, j3))
+
+    # scalar variants: same code path, different signature
+    add("mul", "scalar", lambda: torch.mul(xs, 3.0))
+    add("add", "scalar", lambda: torch.add(xs, 3.0))
+
+    # unary: should not involve broadcasting at all
+    add("silu", "unary", lambda: torch.nn.functional.silu(xs))
+    add("gelu", "unary", lambda: torch.nn.functional.gelu(xs))
+    add("abs", "unary", lambda: torch.abs(xs))
+
+    # reductions
+    add("softmax", "last dim", lambda: torch.softmax(xs, dim=-1))
+    add("amax", "last dim", lambda: torch.amax(xs, dim=-1))
+    add("sum", "last dim", lambda: torch.sum(xs, dim=-1))
+
+    # normalisation; guarded because rms_norm is only in newer torch
+    if hasattr(torch.nn.functional, "rms_norm"):
+        add("rms_norm", "1x512", lambda: torch.nn.functional.rms_norm(xs, (512,)))
+
+    # more pointwise / reduction candidates for an allow-list
+    add("exp", "unary", lambda: torch.exp(xs))
+    add("sigmoid", "unary", lambda: torch.sigmoid(xs))
+    add("rsqrt", "unary", lambda: torch.rsqrt(xs.abs() + 1))
+    add("neg", "unary", lambda: torch.neg(xs))
+    add("relu", "unary", lambda: torch.relu(xs))
+    add("sqr", "unary", lambda: torch.square(xs))
+    add("clamp", "min/max", lambda: torch.clamp(xs, -0.5, 0.5))
+    add("pow", "scalar", lambda: torch.pow(xs, 2))
+    add("reciprocal", "unary", lambda: torch.reciprocal(xs.abs() + 1))
+    add("mean", "last dim", lambda: torch.mean(xs, dim=-1))
+    add("cumsum", "last dim", lambda: torch.cumsum(xs, dim=-1))
+    add("masked_fill", "scalar", lambda: xs.masked_fill(xs > 0, -1.0))
+    add("where", "self", lambda: torch.where(xs > 0, xs, -xs))
+    add("eq", "scalar", lambda: torch.eq(xs, 0.0))
+    add("lt", "scalar", lambda: torch.lt(xs, 0.0))
+    add("ge", "scalar", lambda: torch.ge(xs, 0.0))
+    add("nonzero", "unary", lambda: torch.nonzero(xs > 0.5))
+    add("gather", "dim0", lambda: torch.gather(xs, 0, torch.zeros((8, 512), device=DEV, dtype=torch.int64)))
+    add("index_select", "dim0", lambda: torch.index_select(xs, 0, torch.arange(8, device=DEV)))
+    add("sort", "last dim", lambda: torch.sort(xs, dim=-1).values)
+    add("sin", "unary", lambda: torch.sin(xs))
+    add("cos", "unary", lambda: torch.cos(xs))
+    add("sqrt", "unary", lambda: torch.sqrt(xs.abs()))
+    add("bitwise_and", "scalar", lambda: torch.bitwise_and(i3, 255))
+    add("bincount", "int", lambda: torch.bincount(torch.zeros(64, device=DEV, dtype=torch.int64), minlength=8))
+
+    # shape ops
+    add("cat", "2 tensors", lambda: torch.cat([xs, xs], dim=0))
+    add("stack", "2 tensors", lambda: torch.stack([xs, xs], dim=0))
+    add("repeat_interleave", "self_int 2", lambda: torch.repeat_interleave(xs, 2, dim=0))
+
+    return cases
+
+
+def main() -> int:
+    cases = build_cases()  # references captured here, before enable()
+    import flag_gems
+
+    flag_gems.enable()
+    registered = sorted(flag_gems.all_registered_ops())
+    print(f"flag_gems {flag_gems.__version__}, {len(registered)} ops registered\n")
+
+    def served(op: str) -> bool:
+        """Did FlagGems actually take this op?
+
+        Without this the screen lies by omission: `torch.bitwise_xor` has no FlagGems registration at
+        all, so a "match" there is PyTorch compared with PyTorch and says nothing about FlagGems. Any
+        case that is not served is reported as such and cannot enter an allow-list.
+        """
+        return any(n == op or n.startswith(f"{op}.") or n.startswith(f"{op}_") for n in registered)
+
+    bad: dict[str, list[str]] = {}
+    unserved: set[str] = set()
+    rows = []
+    for c in cases:
+        if not served(c.op):
+            unserved.add(c.op)
+            rows.append((c.op, c.label, None, "not registered in flag_gems (PyTorch answered)"))
+            continue
+        try:
+            got = c.run()
+            torch.cuda.synchronize()
+            ok, msg = same(c.ref, got)
+        except Exception as exc:  # noqa: BLE001 - an exception is itself a finding
+            ok, msg = False, f"raised {type(exc).__name__}: {str(exc)[:70]}"
+        rows.append((c.op, c.label, ok, msg))
+        if not ok:
+            bad.setdefault(c.op, []).append(c.label)
+
+    width = max(len(r[0]) for r in rows) + 2
+    for op, label, ok, msg in rows:
+        mark = "skip" if ok is None else ("ok  " if ok else "BAD ")
+        print(f"  {mark} {op:<{width}} {label:<16} {msg}")
+
+    if unserved:
+        print(f"\nnot served by flag_gems (cannot be judged here): {sorted(unserved)}")
+    print()
+    if bad:
+        print("UNSAFE OPS (exclude these):")
+        for op, labels in sorted(bad.items()):
+            print(f"  {op}: {', '.join(labels)}")
+        return 1
+    print("all served ops matched")
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())