Add FlagGems single-op perf benchmark harness

pytest plugin set + driver script for compiler A/B perf comparison:
- reproducible runs: fixed seed, yaml-driven shapes, autotune record/replay
- per-shape ttgir dump of actually-used variants with readable naming
- cudagraph-based timing with documented fallback semantics

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

.gitignore

@@ -0,0 +1,2 @@
+runs/
+__pycache__/

README.md

@@ -1,2 +1,143 @@
-# zl_bench
+# zl_bench — FlagGems 单算子性能测试工具
 
+针对编译器（FlagTree）改动做算子级 A/B 性能对比的 pytest 插件集 + 驱动脚本。在 FlagGems benchmark 体系之上解决三个问题：
+
+1. **可复现**：固定随机种子、指定 shape、固定 autotune config，把 A/B 两次运行之间的差异收敛到"编译器改动"这一个变量；
+2. **可解释**：自动按 shape 收集每次运行实际使用的 ttgir（带可读的变体命名），供 IR 级 diff；
+3. **口径统一**：cudagraph 计时消除 launch 开销，小 kernel 的对比不被 CPU 侧噪声淹没。
+
+## 快速开始
+
+```bash
+# 默认算子（fused_marlin_moe_mxfp4，内置 16 组 MoE shape）
+bash run_pytest.sh
+
+# 换算子：OP=测试函数名去掉 test_ 前缀，OP_FILE=benchmark 文件名去掉 test_ 前缀/.py 后缀
+# （benchmark 文件为 $FLAGGEMS_DIR/benchmark/test_<OP_FILE>.py::test_<OP>）
+OP=softmax OP_FILE=softmax SHAPE_FILE=my_shapes.yaml bash run_pytest.sh
+```
+
+shape yaml 顶层 key 必须是 op 名：
+
+```yaml
+softmax:
+  shapes:
+    - [1024, 1024]
+    - [64, 512, 512]
+```
+
+跑之前先 `nvidia-smi` 确认目标卡空闲——共享机器上别的任务会把 baseline 和被测两边一起等比拖慢，出一份看似自洽实则作废的数据。多卡机器上用 `CUDA_VISIBLE_DEVICES=<空闲卡号>` 明确选卡。
+
+脚本内固定了 `--level core --mode kernel` 和 `USE_FLAGTUNE=1`（FlagTune 扩展调优空间，见"为什么需要 replay"一节）；需要改动这些口径时直接编辑 `run_pytest.sh` 中的 pytest 命令行。
+
+### 环境变量一览
+
+| 变量            | 默认                      | 说明                                             |
+|----------------|---------------------------|--------------------------------------------------|
+| `OP`           | `fused_marlin_moe_mxfp4`  | 测试函数名（`test_` 之后的部分）                    |
+| `OP_FILE`      | `fused_marlin_moe`        | benchmark 文件名（`test_` 与 `.py` 之间的部分）     |
+| `SHAPE_FILE`   | 空（用脚本内置 yaml）       | shape yaml 路径                                   |
+| `FLAGGEMS_DIR` | `/workspace/FlagGems-dev` | FlagGems 仓库路径                                 |
+| `REPLAY_FROM`  | 空（record 模式）          | 指向某次历史 run 目录，replay 其 autotune 选择（见下） |
+
+## 输出目录结构
+
+每次运行产出 `runs/<op>_<时间戳>/`：
+
+```
+runs/softmax_20260716-031752/
+├── run.log                  # 完整日志（含 SUCCESS 行的 latency/speedup 表）
+├── shapes.yaml              # 本次实际使用的 shape（存档）
+├── autotune_records/        # Triton autotune 选中的 config（供 replay）
+│   └── softmax.json
+└── ttgir/                   # 按 shape 分组的 IR 落盘
+    ├── 1024x1024/
+    │   └── softmax_kernel_inner/
+    │       ├── w4s3__bf16.ttgir
+    │       ├── w4s3__fp16.ttgir
+    │       └── w4s3__fp32.ttgir
+    ├── 4096x4096/...
+    ├── index.tsv            # 每个编译变体的完整参数、启动次数、cache hash
+    └── naming.md            # 文件名缩写图例
+```
+
+ttgir 关键设计：
+
+- **只落盘"实际使用"的变体**：autotune sweep 中测过但落选的 config 不拷贝（只在 index.tsv 里留 `sweep loser` 记录）。mm 这类 sweep 上千个变体的算子，最终只留真正被选中的几个文件。
+- **命名 = 组内有区分度的参数**：同一 (shape, kernel) 组内取值相同的 constexpr 不进文件名；多词参数缩写为首字母（`BLOCK_SIZE_M` → `BSM`，图例见 naming.md）；同名冲突依次用 dtype（`__fp16`）、参数对齐特化（`__EMdiv16`）、hash 前缀消歧。
+- **崩溃也能拿到 IR**：dump 挂在进程 atexit 上，CUDA crash 后仍会落盘已编译部分。
+
+## A/B 对比测试标准流程
+
+```bash
+# A 侧（基线编译器）：正常跑，自动 record autotune 选择
+bash run_pytest.sh                          # -> runs/<op>_<ts_A>/
+
+# B 侧（改动后编译器）：replay A 侧的 config，保证两侧同 config
+REPLAY_FROM=$PWD/runs/<op>_<ts_A> bash run_pytest.sh
+```
+
+对比 `run.log` 的 latency 表看性能差异；diff 两侧 `ttgir/<shape>/<kernel>/` 下的同名文件看 IR 差异。
+
+replay 的兜底行为：B 侧遇到记录中没有的 key、或记录的 config 在新编译器下编译失败时，自动回退到现场 autotune 并在 `run.log` 打 `AUTOTUNE_REPLAY_FALLBACK reason=...` 标记——出现该标记的测量点不再满足"同 config"前提，解读时注意。另外 replay 模式的 run 目录不产生 `autotune_records/`，后续 run 的 `REPLAY_FROM` 应始终指向最初 record 的那次 A 侧目录，不要链式指向 replay 产物。
+
+### cudagraph 回退与测量精度
+
+部分算子本身不支持 CUDA graph capture（测量函数内含 host 同步、动态显存分配、不合法的流操作等），这类算子会自动回退到普通 do_bench 计时，`run.log` 中打 `BENCHMARK_DIRECT_NO_CUDAGRAPH` 标记（含具体原因）。
+
+回退本身不影响 A/B 公平性（两侧同一算子回退行为一致），但**回退口径的测量误差更大**：do_bench 每次迭代都走完整的 Python → launch 路径，kernel 越小，launch 开销和 CPU 侧抖动在数字里占比越高——亚毫秒级 kernel 上两种口径可差 2 倍以上，且行间波动更明显。解读这类算子的结果时：
+
+- 小 shape 行的绝对值和小幅（<10%）差异不要过度解读，优先看大 shape 行；
+- 需要更高置信度时，同一配置多跑几次取中位数，或对该算子直接注释掉 `_cudagraph_plugin` 统一用 do_bench 口径（消除同表混两种口径的问题）。
+
+回退是按测量点发生的，同一份 latency 表里可能混有两种口径的行；若不确定，先 `grep BENCHMARK_DIRECT_NO_CUDAGRAPH run.log` 确认哪些行是回退口径再下结论。
+
+### 为什么需要 replay（以及它管不到什么）
+
+本项目涉及两层调优机制，对 A/B 的影响不同：
+
+| 机制 | 结果存储 | 编译器改动后 | A/B 风险 |
+|-----|---------|------------|----------|
+| Triton `@triton.autotune` | 仅进程内存 | 每次进程重新 sweep | 计时噪声可能让 A/B 选中**不同 config** → 用 REPLAY_FROM 固定 |
+| FlagGems `@libtuner`（含 FlagTune 扩展空间） | `~/.flaggems/config_cache/*.db`（sqlite，跨进程持久） | **不失效**（表名只含 kernel 源码与 config 空间的 hash），A/B 自动命中同一 winner | 反向风险：B 侧沿用 A 侧选的旧 winner，测的是"旧 config 下的编译器差异"而非"各自最优" |
+
+libtuner 的持久缓存何时失效：FlagGems kernel 源码改动、tune_configs.yaml / expand yaml / `USE_FLAGTUNE` 开关变化、Triton 大版本或 GPU 型号变化。如果需要"各自最优"口径（让 B 侧重新 sweep），删除 sqlite 中对应表，或设 `FLAGGEMS_DB_URL` 指向一次性文件。**两种口径都合理，报告结论时注明用的哪种。**
+
+## 插件说明
+
+脚本通过 `-p` 加载以下插件（`run_pytest.sh` 的 `PLUGINS` 数组，可按需注释）：
+
+| 插件 | 作用 | 何时关闭 |
+|-----|------|---------|
+| `_seed_plugin` | 固定 random/numpy/torch 种子，数据相关算子（sort/topk 等）输入逐字节一致 | 不关 |
+| `_shape_inject_plugin` | 让 shape yaml 覆盖子类硬编码的 `set_shapes()` | 不关 |
+| `_shape_iter_inject_plugin` | 覆盖在 `get_input_iter` 里硬编码 shape 的类（conv/pool 等） | 不关 |
+| `_bespoke_shape_plugin` | 覆盖特殊输入构造的算子（upsample/flash_mla/cutlass 等，按类名注册） | 测这些算子之外可关 |
+| `_autotune_record_plugin` | record/replay Triton autotune 选择（由 RECORD/REPLAY 环境变量二选一激活） | 不关 |
+| `_cudagraph_plugin` | `do_bench` → `do_bench_cudagraph`（kernel 纯耗时；内部先做跨流同步，修过一个间歇性 illegal instruction）。无法 graph capture 的 kernel 自动回退并打 `BENCHMARK_DIRECT_NO_CUDAGRAPH` 标记 | 需要与他人的普通 do_bench 数据对齐时注释掉 |
+| `_ir_meta_plugin` | 编译期记录每个变体的 constexpr/签名/特化；launch 钩子统计每个 (kernel, shape) 的实际使用；退出时按上述结构 dump ttgir | 不关 |
+| `_mm_cluster_fix_plugin` | Hopper fp16 mm cluster kernel 越界崩溃的运行时规避 | 默认注释；测 fp16 mm 崩溃时打开 |
+
+## 常见问题
+
+**Q: latency 和别人跑的差很多？**
+
+先检查 GPU 是否被共占（`nvidia-smi`）；再确认对方是否开了 cudagraph——M=1 这类小 shape 下 launch 开销占比大，两种口径可差 2 倍以上，大 shape 基本一致。
+
+**Q: 第一次跑某算子特别慢？**
+
+libtuner 算子首跑要做全量 sweep（mm 约 50 分钟），winner 持久化到 `~/.flaggems/config_cache/` 后，之后同 shape 秒级命中。
+
+**Q: `--warmup/--iter` 要设吗？**
+
+不用。cudagraph 计时路径下 warmup 参数本就不生效（内部自带预热），iter 默认 100ms 预算按 kernel 耗时自适应换算次数。
+
+**Q: ttgir 目录里某个 shape 少了文件？**
+
+看 `index.tsv` 的 `launches` 列——没在该 shape 下真正启动过的变体不落盘。`run.log` 里的 `BENCHMARK_DIRECT_NO_CUDAGRAPH` / `AUTOTUNE_REPLAY_FALLBACK` 标记可解释异常回退。
+
+## 依赖假设
+
+- FlagGems benchmark 体系（`benchmark/base.py` 的 `Benchmark` 类、conftest 的 `--shape_file/--level/--mode` 选项）；
+- Triton 需支持 `knobs.compilation.listener`、`kernel_load_end_hook`、`launch_enter_hook`（当前 FlagTree 的 triton 3.6 满足）；
+- 插件通过 monkeypatch 挂钩上游内部结构，FlagGems/Triton 大版本升级后若行为异常，优先检查各插件 pytest_configure 输出的注册日志是否还正常打印。

_autotune_record_plugin.py

@@ -0,0 +1,275 @@
+"""pytest plugin: record or replay autotune configs across benchmark runs.
+
+When comparing two runs (e.g. two compiler builds), replay the first run's
+recorded configs in the second: the autotuner's choice is pinned, so the
+latency delta reflects the change under test rather than the autotuner picking
+a different config under different IR.
+
+Modes (mutually exclusive, selected by env var; op identity from
+FLAGGEMS_PERF_CURRENT_OP):
+- FLAGGEMS_PERF_AUTOTUNE_RECORD_DIR: run full autotune, dump the chosen Config
+  per (kernel, key) to <dir>/<op>.json at session end.
+- FLAGGEMS_PERF_AUTOTUNE_REPLAY_DIR: pre-populate Autotuner.cache from
+  <dir>/<op>.json, skipping the sweep. Missing key / non-compiling config /
+  missing json fall back to full autotune and emit an AUTOTUNE_REPLAY_FALLBACK
+  marker into the log.
+"""
+
+from __future__ import annotations
+
+import atexit
+import json
+import os
+import sys
+import threading
+from pathlib import Path
+from typing import Any, Dict, Optional, Tuple
+
+
+_RECORD_DIR_ENV = "FLAGGEMS_PERF_AUTOTUNE_RECORD_DIR"
+_REPLAY_DIR_ENV = "FLAGGEMS_PERF_AUTOTUNE_REPLAY_DIR"
+_OP_ENV = "FLAGGEMS_PERF_CURRENT_OP"
+
+# Module-global aggregates, accumulated as run() is called and dumped at exit.
+# Lock guards merges in case a kernel uses threads internally.
+_record_map: Dict[str, Dict[str, Dict[str, Any]]] = {}
+_record_lock = threading.Lock()
+_replay_map: Dict[str, Dict[str, Dict[str, Any]]] = {}
+_dump_done = False
+
+
+def _op_name() -> str:
+    return os.environ.get(_OP_ENV, "").strip() or "default"
+
+
+def _record_path(record_dir: str) -> Path:
+    return Path(record_dir) / f"{_op_name()}.json"
+
+
+def _kernel_id(tuner: Any) -> str:
+    # Stable across runs: source module + qualified name. Shapes/keys split
+    # by the inner key dict.
+    name = (
+        getattr(tuner, "__name__", None)
+        or getattr(getattr(tuner, "base_fn", None), "__name__", None)
+        or getattr(getattr(tuner, "fn", None), "__name__", None)
+        or "unknown_kernel"
+    )
+    module = getattr(getattr(tuner, "fn", None), "__module__", "") or ""
+    return f"{module}::{name}"
+
+
+def _serialize_key(key: Tuple[Any, ...]) -> str:
+    # Per-element repr, JSON-encoded into a flat string usable as a JSON key;
+    # stable across runs given seeded RNG + yaml-fixed shapes.
+    return json.dumps([repr(x) for x in key])
+
+
+def _serialize_config(cfg: Any) -> Optional[Dict[str, Any]]:
+    # Drop configs with a pre_hook (un-serializable callable) rather than replay
+    # without it and compute incorrectly.
+    if getattr(cfg, "pre_hook", None) is not None:
+        return None
+    try:
+        return {
+            "kwargs": dict(cfg.kwargs),
+            "num_warps": cfg.num_warps,
+            "num_stages": cfg.num_stages,
+            "num_ctas": cfg.num_ctas,
+            "maxnreg": cfg.maxnreg,
+        }
+    except Exception:
+        return None
+
+
+def _deserialize_config(entry: Dict[str, Any]) -> Any:
+    # Late import (triton loads during collection, after plugin registration).
+    from triton import Config
+    return Config(
+        kwargs=dict(entry["kwargs"]),
+        num_warps=entry.get("num_warps", 4),
+        num_stages=entry.get("num_stages", 3),
+        num_ctas=entry.get("num_ctas", 1),
+        maxnreg=entry.get("maxnreg"),
+    )
+
+
+def _compute_key(tuner: Any, args: Tuple[Any, ...], kwargs: Dict[str, Any]) -> Optional[Tuple[Any, ...]]:
+    # Reproduce Autotuner/LibTuner key computation; return None on any internals
+    # shift (caller then just loses replay for that call).
+    try:
+        arg_names = tuner.arg_names
+        nargs = dict(zip(arg_names, args))
+        all_args = {**nargs, **kwargs}
+        _args = {k: v for k, v in all_args.items() if k in arg_names}
+        if hasattr(tuner, "get_key"):
+            return tuner.get_key(_args)
+        key = [_args[name] for name in tuner.keys if name in _args]
+        for _, arg in _args.items():
+            if hasattr(arg, "dtype"):
+                key.append(str(arg.dtype))
+        return tuple(key)
+    except Exception:
+        return None
+
+
+def _emit_marker(reason: str) -> None:
+    # Marker lands in run.log (grep it to audit replay coverage); flush so it
+    # appears before the kernel runs.
+    print(f"AUTOTUNE_REPLAY_FALLBACK reason={reason}", flush=True)
+
+
+def _record_run(original):
+    # Snapshot self.cache before/after run() to capture the chosen config (works
+    # for both Autotuner and LibTuner). Single-config kernels skip the cache
+    # write, so record configs[0] explicitly for uniform replay.
+    def runner(self, *args, **kwargs):
+        try:
+            keys_before = set(self.cache.keys()) if hasattr(self.cache, "keys") else set()
+        except Exception:
+            keys_before = set()
+        result = original(self, *args, **kwargs)
+        try:
+            keys_after = set(self.cache.keys()) if hasattr(self.cache, "keys") else set()
+        except Exception:
+            keys_after = set()
+        new_keys = keys_after - keys_before
+        if not new_keys:
+            # No new entry: record configs[0] for single-config kernels (cache
+            # write bypassed); otherwise nothing to record (disk-cache hit).
+            if len(getattr(self, "configs", []) or []) == 1:
+                key = _compute_key(self, args, kwargs)
+                if key is not None and key not in self.cache:
+                    cfg = self.configs[0]
+                    entry = _serialize_config(cfg)
+                    if entry is not None:
+                        kid = _kernel_id(self)
+                        with _record_lock:
+                            bucket = _record_map.setdefault(kid, {})
+                            bucket[_serialize_key(key)] = entry
+            return result
+        kid = _kernel_id(self)
+        with _record_lock:
+            bucket = _record_map.setdefault(kid, {})
+            for k in new_keys:
+                try:
+                    cfg = self.cache[k]
+                except Exception:
+                    continue
+                entry = _serialize_config(cfg)
+                if entry is None:
+                    continue
+                bucket[_serialize_key(k)] = entry
+        return result
+    return runner
+
+
+def _replay_run(original):
+    # Inject recorded configs before benchmarking; on key-missing or a recorded
+    # config failing to compile/launch, fall back to full autotune (+ marker).
+    def runner(self, *args, **kwargs):
+        injected_key = None
+        if len(getattr(self, "configs", []) or []) > 1:
+            kid = _kernel_id(self)
+            bucket = _replay_map.get(kid)
+            if bucket:
+                key = _compute_key(self, args, kwargs)
+                if key is not None and key not in self.cache:
+                    rec = bucket.get(_serialize_key(key))
+                    if rec is None:
+                        _emit_marker("key_missing")
+                    else:
+                        try:
+                            cfg = _deserialize_config(rec)
+                            self.cache[key] = cfg
+                            injected_key = key
+                        except Exception as exc:
+                            _emit_marker(f"deserialize_{type(exc).__name__}")
+        try:
+            return original(self, *args, **kwargs)
+        except Exception as exc:
+            if injected_key is not None:
+                # Recorded config failed under the current build: drop it and
+                # let run() autotune.
+                try:
+                    del self.cache[injected_key]
+                except Exception:
+                    pass
+                _emit_marker(f"compile_{type(exc).__name__}")
+                return original(self, *args, **kwargs)
+            raise
+    return runner
+
+
+def _load_replay_dir(replay_dir: str) -> int:
+    path = _record_path(replay_dir)
+    if not path.is_file():
+        _emit_marker("record_missing")
+        return 0
+    try:
+        data = json.loads(path.read_text())
+    except Exception as exc:
+        _emit_marker(f"load_{type(exc).__name__}")
+        return 0
+    if not isinstance(data, dict):
+        return 0
+    count = 0
+    for kid, bucket in data.items():
+        if not isinstance(bucket, dict):
+            continue
+        _replay_map[kid] = dict(bucket)
+        count += len(bucket)
+    return count
+
+
+def _dump_record(record_dir: str) -> None:
+    global _dump_done
+    if _dump_done:
+        return
+    _dump_done = True
+    path = _record_path(record_dir)
+    try:
+        path.parent.mkdir(parents=True, exist_ok=True)
+        with _record_lock:
+            payload = {kid: dict(bucket) for kid, bucket in _record_map.items()}
+        path.write_text(json.dumps(payload, indent=2, ensure_ascii=False) + "\n")
+    except Exception as exc:
+        # Don't break the run if dump fails; the perf row is still produced.
+        print(f"[autotune-record-plugin] dump failed: {exc}", file=sys.stderr, flush=True)
+
+
+def pytest_configure(config):
+    record_dir = os.environ.get(_RECORD_DIR_ENV, "").strip()
+    replay_dir = os.environ.get(_REPLAY_DIR_ENV, "").strip()
+    if record_dir and replay_dir:
+        # Mutually exclusive: recording while replaying is pointless. Surface loudly.
+        print("[autotune-record-plugin] error: both RECORD_DIR and REPLAY_DIR set; "
+              "ignoring both (no-op)", file=sys.stderr, flush=True)
+        return
+    if not record_dir and not replay_dir:
+        return
+
+    import triton.runtime.autotuner as _autotuner
+    wrap = _record_run if record_dir else _replay_run
+    _autotuner.Autotuner.run = wrap(_autotuner.Autotuner.run)
+    patched = ["Autotuner"]
+    try:
+        from flag_gems.utils.libentry import LibTuner
+        if "run" in LibTuner.__dict__:
+            LibTuner.run = wrap(LibTuner.__dict__["run"])
+            patched.append("LibTuner")
+    except Exception:
+        pass
+
+    if record_dir:
+        # atexit (not sessionfinish): persist whatever was recorded even if an op
+        # crash kills the session; a later replay run falls back for missing keys.
+        atexit.register(_dump_record, record_dir)
+        print(f"[autotune-record-plugin] recording autotune configs to "
+              f"{_record_path(record_dir)} ({'+'.join(patched)})",
+              file=sys.stderr, flush=True)
+    else:
+        loaded = _load_replay_dir(replay_dir)
+        print(f"[autotune-record-plugin] replaying {loaded} recorded entries from "
+              f"{_record_path(replay_dir)} ({'+'.join(patched)})",
+              file=sys.stderr, flush=True)

_bespoke_shape_plugin.py

@@ -0,0 +1,150 @@
+"""pytest plugin: make the shape yaml authoritative for ops whose inputs are
+*bespoke* (built inline / from TestParam objects / a sampled kit) and so are
+unreachable by `_shape_inject_plugin` and `_shape_iter_inject_plugin`.
+
+A per-op registry; each handler REUSES the class's own input construction and
+only swaps the shape source for our yaml (so we never measure a different
+kernel). Covered: upsample_bicubic2d_aa_backward, flash_mla_sparse_fwd,
+fused_deepseek_..._quant_insert, cutlass_scaled_mm. Ops with no yaml entry fall
+through unchanged. Couples to upstream internals; a mismatch surfaces as a loud
+benchmark FAIL, not silent corruption.
+"""
+
+from __future__ import annotations
+
+import os
+import sys
+
+import yaml as _yaml
+
+_YAML_CACHE: dict | None = None
+
+
+def _load_yaml(path: str) -> dict:
+    if not path or not os.path.isfile(path):
+        return {}
+    try:
+        with open(path, "r") as f:
+            return _yaml.safe_load(f) or {}
+    except Exception as exc:
+        print(f"[bespoke-shape-plugin] failed to load {path}: {exc}", file=sys.stderr)
+        return {}
+
+
+def _yaml_shapes(op_name):
+    entry = (_YAML_CACHE or {}).get(op_name)
+    if isinstance(entry, dict):
+        shapes = entry.get("shapes")
+        if shapes:
+            return shapes
+    return None
+
+
+# --- per-op builders: reuse the class's own construction, swap the shape source ---
+
+def _upsample_builder(self, dtype, shapes, original):
+    # _cfgs entries are (N,C,Hi,Wi,Ho,Wo,ac,label); yaml carries the 7 numerics.
+    self._cfgs = [tuple(shape) + ("yaml",) for shape in shapes]
+    yield from original(self, dtype)
+
+
+def _fused_deepseek_builder(self, dtype, shapes, original):
+    TestParam = original.__globals__["TestParam"]
+    for shape in shapes:
+        num_tokens, num_heads, num_tokens_insert, block_size, max_pos, eps = shape
+        param = TestParam(
+            num_tokens=num_tokens,
+            num_heads=num_heads,
+            num_tokens_insert=num_tokens_insert,
+            block_size=block_size,
+            max_pos=max_pos,
+            eps=eps,
+        )
+        yield from type(self).make_input(param)
+
+
+def _flash_mla_sparse_builder(self, dtype, shapes, original):
+    TestParam = original.__globals__["TestParam"]
+    for shape in shapes:
+        s_q, s_kv, topk, h_q, d_qk = shape
+        param = TestParam(
+            s_q=s_q, s_kv=s_kv, topk=topk, h_q=h_q, d_qk=d_qk,
+            have_attn_sink=True,  # every hardcoded case uses True
+        )
+        yield from type(self).make_input_flashmla(param)
+
+
+_WRAP_BUILDERS = {
+    "UpsampleBicubic2dAaBackwardBenchmark": _upsample_builder,
+    "FusedDeepseekV4QnormRopeKVRopeQuantInsertBenchmark": _fused_deepseek_builder,
+    "FlashmlaSparseBenchmark": _flash_mla_sparse_builder,
+}
+
+
+def _wrap_get_input_iter(cls, builder) -> None:
+    original = cls.get_input_iter
+
+    def patched(self, dtype):
+        shapes = _yaml_shapes(getattr(self, "op_name", None))
+        if shapes:
+            yield from builder(self, dtype, shapes, original)
+        else:
+            yield from original(self, dtype)
+
+    cls.get_input_iter = patched
+
+
+def _patch_cutlass(cls) -> bool:
+    """Redirect CutlassScaledMMPerfKit's hardcoded mnk to the yaml M,N,K, reusing
+    the kit's own combination/sampling pipeline (the quant-mode sweep stays
+    intrinsic)."""
+    shapes = _yaml_shapes("cutlass_scaled_mm")
+    if not shapes:
+        return False
+    kit = getattr(sys.modules.get(cls.__module__), "CutlassScaledMMPerfKit", None)
+    if kit is None:
+        return False
+    import torch
+    from itertools import product
+
+    mnk = [tuple(shape) for shape in shapes]
+
+    def _get_all_combinations():
+        # Mirror upstream verbatim except mnk (from yaml); keep in sync with
+        # CutlassScaledMMPerfKit._get_all_combinations.
+        scale_shape_types = ["scalar", "vector", "matrix"]
+        if_use_bias = [True, False]
+        dtypes = [(torch.int8, torch.float16), (torch.float8_e4m3fn, torch.bfloat16)]
+        return product(mnk, scale_shape_types, scale_shape_types, if_use_bias, dtypes)
+
+    kit._get_all_combinations = staticmethod(_get_all_combinations)
+    return True
+
+
+def pytest_collection_finish(session):
+    global _YAML_CACHE
+    from benchmark import base as fg_base
+    from benchmark.conftest import Config as fg_Config
+
+    _YAML_CACHE = _load_yaml(getattr(fg_Config, "shape_file", "") or "")
+
+    def _subclasses(cls):
+        for sub in cls.__subclasses__():
+            yield sub
+            yield from _subclasses(sub)
+
+    covered = []
+    seen = set()
+    for cls in _subclasses(fg_base.Benchmark):
+        if cls in seen:
+            continue
+        seen.add(cls)
+        builder = _WRAP_BUILDERS.get(cls.__name__)
+        if builder is not None and "get_input_iter" in cls.__dict__:
+            _wrap_get_input_iter(cls, builder)
+            covered.append(cls.__name__)
+        if cls.__name__ == "CutlassScaledMMBenchmark" and _patch_cutlass(cls):
+            covered.append("CutlassScaledMMBenchmark(mnk)")
+
+    print(f"[bespoke-shape-plugin] yaml-driven inputs for: {', '.join(covered) or 'none'}",
+          file=sys.stderr, flush=True)

_cudagraph_plugin.py

@@ -0,0 +1,82 @@
+"""pytest plugin: swap `triton.testing.do_bench` for cudagraph-based timing.
+
+Loaded via `-p _cudagraph_plugin` in run_pytest.sh's PLUGINS list (comment that
+line out to time with plain do_bench). Monkey-patches do_bench →
+do_bench_cudagraph (kernel-only latency) at session configure, after Triton is
+loaded but before any benchmark calls it. Kernels that can't be graph-captured
+fall back to plain do_bench and print a BENCHMARK_DIRECT_NO_CUDAGRAPH marker
+into the log. The warmup kwarg is dropped (do_bench_cudagraph warms up
+internally).
+"""
+
+from __future__ import annotations
+
+import os
+import sys
+
+import triton.testing as _tt
+
+
+_ORIGINAL_DO_BENCH = _tt.do_bench
+
+
+def _is_fatal_cuda_error(exc: Exception) -> bool:
+    """True for CUDA context errors (illegal address, OOM, device-side assert):
+    the context is poisoned, so re-raise rather than fall back and record a bogus
+    latency as success."""
+    msg = f"{type(exc).__name__}: {exc}".lower()
+    markers = (
+        "illegal memory access",
+        "cudaerrorillegaladdress",
+        "out of memory",
+        "device-side assert",
+        "an illegal instruction",
+        "misaligned address",
+        "uncorrectable ecc",
+    )
+    return any(m in msg for m in markers)
+
+
+def _patched_do_bench(fn, warmup=25, rep=100, grad_to_none=None,
+                      quantiles=None, return_mode="mean"):
+    """Drop-in replacement for triton.testing.do_bench using cudagraph capture."""
+    op = os.environ.get("FLAGGEMS_PERF_CURRENT_OP", "<unknown>")
+    try:
+        # do_bench_cudagraph runs fn on a fresh side stream without syncing
+        # with the default stream first; input tensors produced just before
+        # (quantize/topk in the input iter) may still be in flight, and racing
+        # on them intermittently kills the context (illegal instruction from a
+        # device trap on garbage indices). Sync before switching streams.
+        import torch
+        torch.cuda.synchronize()
+        return _tt.do_bench_cudagraph(
+            fn,
+            rep=rep,
+            grad_to_none=grad_to_none,
+            quantiles=quantiles,
+            return_mode=return_mode,
+        )
+    except Exception as exc:
+        # Re-raise genuine device failures; only graph-capture rejections fall back.
+        if _is_fatal_cuda_error(exc):
+            raise
+        reason = type(exc).__name__
+        detail = " ".join(str(exc).split())[:160]
+        print(
+            f"BENCHMARK_DIRECT_NO_CUDAGRAPH op={op} reason={reason} detail={detail!r}",
+            flush=True,
+        )
+        return _ORIGINAL_DO_BENCH(
+            fn,
+            warmup=warmup,
+            rep=rep,
+            grad_to_none=grad_to_none,
+            quantiles=quantiles,
+            return_mode=return_mode,
+        )
+
+
+def pytest_configure(config):
+    _tt.do_bench = _patched_do_bench
+    print("[cudagraph-plugin] triton.testing.do_bench → do_bench_cudagraph",
+          file=sys.stderr, flush=True)

_ir_meta_plugin.py

@@ -0,0 +1,428 @@
+"""pytest plugin: annotate compiled kernels and dump their ttgir, grouped by
+benchmark shape, at session end.
+
+Pieces, all keyed by the kernel's Triton cache dir:
+
+1. Compilation listener — writes a __constexprs.json sidecar (constexpr values
+   + per-arg specializations like tt.divisibility) next to each compiled
+   kernel. The ttgir alone doesn't say which autotune variant it is.
+2. Shape tracking — wraps every Benchmark subclass's get_input_iter (after the
+   shape-inject plugins have done their patching) to publish the benchmark
+   shape currently being fed, e.g. "16x256x7168x2048x8".
+3. Launch tracking — kernel_load_end_hook maps a loaded function handle to its
+   cache dir (must be the END hook: at load-start the handle is still None);
+   launch_enter_hook counts launches per (kernel cache dir, current shape).
+   Launches made inside Autotuner._bench (the tuning sweep; LibTuner inherits
+   it) are counted separately, so "actually used" = launched at least once
+   OUTSIDE the sweep.
+4. Dump at exit (only when FLAGGEMS_PERF_TTGIR_DUMP_DIR is set) — walks
+   TRITON_CACHE_DIR and copies each ttgir into
+       <dump>/<shape>/<kernel>/<distinguishing constexprs>_w{warps}s{stages}.ttgir
+   for every shape that actually used it. Only constexprs that VARY within a
+   (shape, kernel) group are named, abbreviated per-word (BLOCK_SIZE_M -> BSM;
+   legend in naming.md). Same-name collisions are disambiguated by signature
+   dtype (dtype sweeps, e.g. __fp16), then varying per-arg specializations
+   (e.g. __EMdiv16), then a cache-hash prefix. index.tsv lists every compiled
+   variant per shape with launch counts, including never-used sweep losers
+   (not copied). atexit (not sessionfinish) so a mid-run CUDA crash still
+   dumps whatever was compiled.
+"""
+
+from __future__ import annotations
+
+import atexit
+import json
+import os
+import re
+import sys
+from collections import defaultdict
+
+import pytest
+
+_DUMP_DIR_ENV = "FLAGGEMS_PERF_TTGIR_DUMP_DIR"
+
+_fn_to_cachedir: dict = {}   # GPU function handle -> cache-dir basename
+_launches: dict = {}         # cache-dir basename -> {shape label: [real, sweep]}
+_in_bench = 0                # >0 while inside Autotuner._bench (autotune sweep)
+_current_shape: str | None = None
+_dump_enabled = False
+
+
+# --- compile-time sidecar -------------------------------------------------
+
+def _plain(value):
+    value = getattr(value, "value", value)  # unwrap tl.constexpr
+    if value is None or isinstance(value, (bool, int, float, str)):
+        return value
+    return repr(value)
+
+
+def _compile_listener(*, src, metadata, metadata_group, times, cache_hit):
+    try:
+        fn = getattr(src, "fn", None)
+        arg_names = getattr(fn, "arg_names", None)
+        constants = getattr(src, "constants", None)
+        if not arg_names or constants is None:
+            return  # IRSource or unexpected layout: nothing to record
+        paths = list(metadata_group.values())
+        if not paths:
+            return
+        path = os.path.join(os.path.dirname(paths[0]), "__constexprs.json")
+        if os.path.exists(path):
+            return  # cache hit on a dir we already annotated
+
+        def arg_name(key):
+            # ASTSource keys constants/attrs by arg-index tuples; map back to names.
+            if isinstance(key, tuple):
+                return ".".join(
+                    arg_names[i] if isinstance(i, int) and i < len(arg_names) else str(i)
+                    for i in key
+                )
+            return str(key)
+
+        out = {arg_name(k): _plain(v) for k, v in constants.items()}
+
+        # Argument signature (dtypes): benchmarks sweep dtypes, and dtype is
+        # not a constexpr — without this, dtype variants collide into hash
+        # suffixes. Values look like "*fp16", "i32".
+        sig = {}
+        try:
+            sig = {str(k): str(v) for k, v in (getattr(src, "signature", None) or {}).items()
+                   if str(v) != "constexpr"}
+        except Exception:
+            pass
+
+        # Per-arg specializations, e.g. ("tt.divisibility", 16) -> "div16".
+        specs = {}
+        for key, props in (getattr(src, "attrs", None) or {}).items():
+            encoded = []
+            for p in props or []:
+                try:
+                    pname, pval = p[0], p[1]
+                except (TypeError, IndexError):
+                    encoded.append(str(p))
+                    continue
+                encoded.append(f"div{pval}" if "divisibility" in str(pname)
+                               else f"{pname}={pval}")
+            if encoded:
+                specs[arg_name(key)] = sorted(encoded)
+
+        tmp = f"{path}.tmp.pid{os.getpid()}"
+        with open(tmp, "w") as f:
+            json.dump({"name": getattr(src, "name", "unknown"),
+                       "constexprs": out, "attrs": specs, "signature": sig},
+                      f, indent=1, sort_keys=True)
+        os.replace(tmp, path)
+    except Exception as exc:  # never break compilation over a metadata dump
+        print(f"[ir-meta-plugin] sidecar dump failed: {exc}", file=sys.stderr)
+
+
+# --- shape + launch tracking ----------------------------------------------
+
+def _load_hook(module, function, name, metadata_group, hash):
+    try:
+        paths = list(metadata_group.values())
+        if paths:
+            _fn_to_cachedir[function] = os.path.basename(os.path.dirname(paths[0]))
+    except Exception:
+        pass
+
+
+def _launch_hook(md):
+    try:
+        cachedir = _fn_to_cachedir.get(md.data.get("function"))
+        if cachedir is None:
+            return
+        rec = _launches.setdefault(cachedir, {}).setdefault(
+            _current_shape or "shape_unknown", [0, 0])
+        rec[1 if _in_bench else 0] += 1
+    except Exception:
+        pass
+
+
+def _wrap_bench(original):
+    def wrapped(self, *args, **kwargs):
+        global _in_bench
+        _in_bench += 1
+        try:
+            return original(self, *args, **kwargs)
+        finally:
+            _in_bench -= 1
+    return wrapped
+
+
+def _shape_label(bench, idx):
+    shapes = getattr(bench, "shapes", None) or []
+    if idx < len(shapes):
+        s = shapes[idx]
+        if isinstance(s, (list, tuple)):
+            return "x".join(re.sub(r"\W", "", str(v)) for v in s)
+        return re.sub(r"\W", "", str(s))
+    return f"input{idx}"
+
+
+def _wrap_input_iter(original):
+    # Publish the shape label BEFORE resuming the generator, so kernels
+    # launched while building the inputs attribute to the right shape too.
+    def patched(self, dtype):
+        global _current_shape
+        it = original(self, dtype)
+        idx = 0
+        while True:
+            _current_shape = _shape_label(self, idx)
+            try:
+                item = next(it)
+            except StopIteration:
+                _current_shape = None
+                return
+            yield item
+            idx += 1
+    return patched
+
+
+@pytest.hookimpl(trylast=True)
+def pytest_collection_finish(session):
+    # trylast: run after the shape-inject plugins have re-pointed
+    # get_input_iter, so we wrap the version that will actually execute.
+    if not _dump_enabled:
+        return
+    from benchmark import base as fg_base
+
+    def _subclasses(cls):
+        for sub in cls.__subclasses__():
+            yield sub
+            yield from _subclasses(sub)
+
+    seen, wrapped = set(), 0
+    for cls in (fg_base.Benchmark, *_subclasses(fg_base.Benchmark)):
+        if cls in seen:
+            continue
+        seen.add(cls)
+        own = cls.__dict__.get("get_input_iter")
+        if own is not None:
+            cls.get_input_iter = _wrap_input_iter(own)
+            wrapped += 1
+    print(f"[ir-meta-plugin] shape tracking wrapped on {wrapped} Benchmark classes",
+          file=sys.stderr, flush=True)
+
+
+# --- naming ----------------------------------------------------------------
+
+_SAN = re.compile(r"[^A-Za-z0-9.-]+")
+
+
+def _fmt_value(v):
+    # isinstance check first: 1 == True in Python, a plain dict lookup would
+    # render GROUP_SIZE_M=1 as "T".
+    if isinstance(v, bool):
+        return "T" if v else "F"
+    if v is None:
+        return "-"
+    s = str(v)
+    if "." in s and not re.fullmatch(r"-?\d+(\.\d+)?", s):
+        s = s.split(".")[-1]  # dotted repr like triton.language.bfloat16
+    s = _SAN.sub("", s)
+    return s[:24] or "x"
+
+
+def _abbrev(name):
+    words = [w for w in name.split("_") if w]
+    if len(name) <= 4 or len(words) < 2:
+        return name
+    out = []
+    for w in words:
+        m = re.match(r"^([A-Za-z])[A-Za-z]*?(\d*)$", w)
+        out.append((m.group(1) + m.group(2)) if m else w[0])
+    return "".join(out)
+
+
+def _disambiguate(group):
+    """Suffixes for variants whose constexprs + launch config coincide: name by
+    the signature dtypes that differ (dtype sweeps), else by differing per-arg
+    specializations, else fall back to a hash prefix."""
+    sig_keys = sorted({k for v in group for k in (v.get("signature") or {})})
+    sig_varying = [
+        k for k in sig_keys
+        if len({(v.get("signature") or {}).get(k) for v in group}) > 1
+    ]
+    if sig_varying:
+        # One dtype sweep usually shifts every tensor arg together; the first
+        # varying arg's dtype identifies the variant. Sanitize "*fp16" -> fp16.
+        sufs = []
+        for v in group:
+            toks = [re.sub(r"\W", "", (v.get("signature") or {}).get(k) or "none")
+                    for k in sig_varying]
+            uniq = sorted(set(toks))
+            sufs.append("__" + (uniq[0] if len(uniq) == 1 else "_".join(
+                f"{_abbrev(k)}{t}" for k, t in zip(sig_varying, toks))))
+        if len(set(sufs)) == len(group):
+            return sufs
+    attr_keys = sorted({k for v in group for k in (v.get("attrs") or {})})
+    varying = [
+        k for k in attr_keys
+        if len({tuple((v.get("attrs") or {}).get(k, [])) for v in group}) > 1
+    ]
+    if varying:
+        sufs = []
+        for v in group:
+            toks = [
+                f"{_abbrev(k)}{'.'.join((v.get('attrs') or {}).get(k) or ['none'])}"
+                for k in varying
+            ]
+            sufs.append("__" + "_".join(toks))
+        if len(set(sufs)) == len(group):
+            return sufs
+    return ["__" + v["hash"][:8] for v in group]
+
+
+def _name_group(items):
+    """Filenames for one (shape, kernel) group: only constexprs whose value
+    varies within the group, abbreviated. Returns ([(item, filename)], legend)."""
+    keys = sorted({k for it in items for k in it["constexprs"]})
+    varying = [
+        k for k in keys
+        if len({json.dumps(it["constexprs"].get(k), sort_keys=True) for it in items}) > 1
+    ]
+    by_ab = defaultdict(list)
+    for k in varying:
+        by_ab[_abbrev(k)].append(k)
+    ab = {k: (a if len(ks) == 1 else k) for a, ks in by_ab.items() for k in ks}
+
+    named = defaultdict(list)
+    for it in items:
+        parts = [f"{ab[k]}{_fmt_value(it['constexprs'][k])}"
+                 for k in varying if k in it["constexprs"]]
+        w, s = it.get("warps"), it.get("stages")
+        parts.append(f"w{w}s{s}" if w not in (None, "") else "cfg-unknown")
+        named["_".join(parts)].append(it)
+
+    results = []
+    for base, group in sorted(named.items()):
+        sufs = _disambiguate(group) if len(group) > 1 else [""]
+        for it, suf in zip(group, sufs):
+            results.append((it, f"{base}{suf}.ttgir"))
+    return results, {k: a for k, a in ab.items() if a != k}
+
+
+# --- dump -------------------------------------------------------------------
+
+def _dump_ttgir(cache_dir: str, dump_dir: str) -> None:
+    import pathlib
+    import shutil
+
+    cache, dest = pathlib.Path(cache_dir), pathlib.Path(dump_dir)
+    if not cache.is_dir():
+        print(f"[ir-meta-plugin] no cache dir {cache}; nothing to dump", file=sys.stderr)
+        return
+    dest.mkdir(parents=True, exist_ok=True)
+
+    def load(path):
+        try:
+            return json.loads(path.read_text())
+        except Exception:
+            return {}
+
+    infos = []
+    for ttgir in sorted(cache.rglob("*.ttgir")):
+        meta = load(ttgir.with_suffix(".json"))
+        side = load(ttgir.parent / "__constexprs.json")
+        infos.append({
+            "src": ttgir, "hash": ttgir.parent.name, "kernel": ttgir.stem,
+            "warps": meta.get("num_warps"), "stages": meta.get("num_stages"),
+            "constexprs": side.get("constexprs", {}), "attrs": side.get("attrs", {}),
+            "signature": side.get("signature", {}),
+            "per_shape": _launches.get(ttgir.parent.name, {}),
+        })
+
+    # (shape, kernel) -> variants really used there (launched outside the sweep)
+    groups = defaultdict(list)
+    for info in infos:
+        for shape, (real, _sweep) in info["per_shape"].items():
+            if real > 0:
+                groups[(shape, info["kernel"])].append(info)
+
+    rows, legends, copied = [], defaultdict(dict), 0
+    for (shape, kernel), items in sorted(groups.items()):
+        named, legend = _name_group(items)
+        legends[kernel].update(legend)
+        for it, fname in named:
+            target = dest / shape / kernel / fname
+            target.parent.mkdir(parents=True, exist_ok=True)
+            shutil.copyfile(it["src"], target)
+            copied += 1
+            real, sweep = it["per_shape"][shape]
+            rows.append((shape, kernel, fname, real, sweep,
+                         it["warps"] or "", it["stages"] or "",
+                         json.dumps(it["constexprs"], sort_keys=True), it["hash"]))
+
+    # Compiled but never really used anywhere (autotune losers): index-only.
+    for info in infos:
+        if not any(real > 0 for real, _ in info["per_shape"].values()):
+            sweep = sum(s for _, s in info["per_shape"].values())
+            rows.append(("-", info["kernel"], "(sweep loser, not dumped)", 0, sweep,
+                         info["warps"] or "", info["stages"] or "",
+                         json.dumps(info["constexprs"], sort_keys=True), info["hash"]))
+
+    rows.sort()
+    with open(dest / "index.tsv", "w") as f:
+        f.write("shape\tkernel\tfile\tlaunches\tsweep_launches\t"
+                "num_warps\tnum_stages\tconstexprs\tcache_hash\n")
+        for r in rows:
+            f.write("\t".join(str(x) for x in r) + "\n")
+
+    lines = ["# ttgir naming legend", "",
+             "Layout: `<shape>/<kernel>/<varying constexprs>_w{warps}s{stages}[__spec|__hash8].ttgir`",
+             "Shape dirs mirror the benchmark shape yaml. Only constexprs that vary",
+             "within a (shape, kernel) group appear; `index.tsv` has the full map.",
+             "Values: `T`/`F` = true/false, `-` = none/null."]
+    for kernel in sorted(legends):
+        if not legends[kernel]:
+            continue
+        lines += ["", f"## {kernel}"]
+        width = max(len(a) for a in legends[kernel].values())
+        for full, a in sorted(legends[kernel].items(), key=lambda kv: kv[1]):
+            lines.append(f"- `{a:<{width}}` = {full}")
+    (dest / "naming.md").write_text("\n".join(lines) + "\n")
+
+    shapes = sorted({r[0] for r in rows if r[0] != "-"})
+    losers = sum(1 for r in rows if r[0] == "-")
+    print(f">>> [Dump] ttgir -> {dest} ({copied} files across {len(shapes)} shapes; "
+          f"{losers} unused variants index-only; legend: naming.md)", flush=True)
+    for s in shapes:
+        n = sum(1 for r in rows if r[0] == s)
+        print(f">>> [Dump]   {s}: {n}", flush=True)
+
+
+# --- registration -----------------------------------------------------------
+
+def pytest_configure(config):
+    global _dump_enabled
+    import triton
+    import triton.runtime.autotuner as _autotuner
+
+    prev = triton.knobs.compilation.listener
+    if prev is None:
+        triton.knobs.compilation.listener = _compile_listener
+    else:
+        def chained(**kwargs):
+            prev(**kwargs)
+            _compile_listener(**kwargs)
+        triton.knobs.compilation.listener = chained
+
+    dump_dir = os.environ.get(_DUMP_DIR_ENV, "").strip()
+    if dump_dir:
+        cache_dir = os.environ.get("TRITON_CACHE_DIR", "").strip()
+        if cache_dir:
+            _dump_enabled = True
+            triton.knobs.runtime.kernel_load_end_hook.add(_load_hook)
+            triton.knobs.runtime.launch_enter_hook.add(_launch_hook)
+            _autotuner.Autotuner._bench = _wrap_bench(_autotuner.Autotuner._bench)
+            atexit.register(_dump_ttgir, cache_dir, dump_dir)
+        else:
+            print("[ir-meta-plugin] warning: dump dir set but TRITON_CACHE_DIR "
+                  "is not; ttgir dump disabled", file=sys.stderr, flush=True)
+
+    print("[ir-meta-plugin] compilation listener registered"
+          + (f"; shape/launch tracking on, ttgir dump -> {dump_dir}"
+             if _dump_enabled else ""),
+          file=sys.stderr, flush=True)

_mm_cluster_fix_plugin.py

@@ -0,0 +1,137 @@
+"""pytest plugin: work around the Hopper fp16 cluster-remote GEMM OOB crash via
+runtime monkeypatch (no upstream source change).
+
+Root cause: fp16 mm uses `_cluster_remote_gemm_kernel` (bf16 doesn't). When
+num_pid_n = cdiv(N,BN) is odd (not a multiple of CLUSTER_SIZE=2) AND the GEMM is
+perfectly tiled (USE_MASK off), the tail cluster's extra CTA reads out of range
+-> cudaErrorIllegalAddress (e.g. N=128256 Llama-3 vocab).
+
+Fix: wrap `cluster_remote_mm_scenario` so only that unsafe combination falls
+back to the general/splitk path; all other fp16 shapes keep the cluster kernel.
+The real backend mm module loads lazily under a synthetic name, so we scan
+sys.modules, force a warmup mm to trigger the load, patch, and re-scan per test.
+"""
+from __future__ import annotations
+
+import sys
+import types
+
+
+def _patch_loaded() -> list:
+    import triton
+
+    patched = []
+    for nm, m in list(sys.modules.items()):
+        if not isinstance(m, types.ModuleType):
+            continue
+        fn = getattr(m, "cluster_remote_mm_scenario", None)
+        if not isinstance(fn, types.FunctionType):
+            continue
+        if getattr(fn, "_mm_cluster_guarded", False):
+            continue
+        bn = getattr(m, "TLE_REMOTE_BN", 256)
+        bm = getattr(m, "TLE_REMOTE_BM", 64)
+        bk = getattr(m, "TLE_REMOTE_BK", 64)
+        cs = getattr(m, "TLE_CLUSTER_SIZE", 2)
+
+        def make(orig, bm, bn, bk, cs):
+            def wrapped(a, b, c, M, N, K):
+                if not orig(a, b, c, M, N, K):
+                    return False
+                use_mask = (M % bm != 0) or (N % bn != 0) or (K % bk != 0)
+                # odd N-tile count + no mask -> tail cluster over-runs N
+                if (triton.cdiv(N, bn) % cs != 0) and not use_mask:
+                    return False
+                return True
+            wrapped._mm_cluster_guarded = True
+            return wrapped
+
+        m.cluster_remote_mm_scenario = make(fn, bm, bn, bk, cs)
+        patched.append(nm)
+    return patched
+
+
+def _guarded_modules() -> list:
+    """Loaded modules whose cluster_remote_mm_scenario is already our wrapper
+    (cumulative state, unlike _patch_loaded()'s newly-wrapped-only return)."""
+    out = []
+    for nm, m in list(sys.modules.items()):
+        if not isinstance(m, types.ModuleType):
+            continue
+        fn = getattr(m, "cluster_remote_mm_scenario", None)
+        if isinstance(fn, types.FunctionType) and getattr(fn, "_mm_cluster_guarded", False):
+            out.append(nm)
+    return out
+
+
+def _is_hopper():
+    """True iff a CUDA device with capability >= 9 (the only HW that hits the
+    cluster-remote path) is visible. Returns False on any error / no CUDA."""
+    try:
+        import torch
+        if not torch.cuda.is_available():
+            return False
+        return torch.cuda.get_device_capability()[0] >= 9
+    except Exception:  # noqa: BLE001
+        return False
+
+
+def _warmup_load_backend():
+    """Trigger the lazy backend mm module load via one tiny 64x64 fp16 mm (too
+    small for the cluster path, so it can't trip the bug). A CUDA error here
+    means a pre-broken context -> surface loudly; only import/attr errors are
+    swallowed."""
+    try:
+        import torch
+        import flag_gems
+        if not torch.cuda.is_available():
+            return
+        a = torch.randn(64, 64, dtype=torch.float16, device="cuda")
+        b = torch.randn(64, 64, dtype=torch.float16, device="cuda")
+        with flag_gems.use_gems():
+            torch.mm(a, b)
+            torch.cuda.synchronize()
+    except (ImportError, AttributeError) as e:
+        print(f"[mm-cluster-fix-plugin] warmup skipped (benign): {e!r}",
+              file=sys.stderr, flush=True)
+    except Exception as e:  # noqa: BLE001
+        # e.g. a CUDA error — do NOT hide it; the harness needs to see it.
+        print(f"[mm-cluster-fix-plugin] WARNING: warmup mm failed unexpectedly: {e!r}",
+              file=sys.stderr, flush=True)
+
+
+def pytest_configure(config):
+    _patch_loaded()
+    _warmup_load_backend()
+    _patch_loaded()
+    guarded = _guarded_modules()
+    if guarded:
+        print(
+            "[mm-cluster-fix-plugin] guarded cluster_remote_mm_scenario "
+            f"(odd N-tile unmasked OOB) in: {guarded}",
+            file=sys.stderr,
+            flush=True,
+        )
+    elif _is_hopper():
+        # On Hopper with nothing patched, the guard is a silent no-op and fp16 mm
+        # will crash; fail loudly rather than risk an illegal-access.
+        raise RuntimeError(
+            "[mm-cluster-fix-plugin] FATAL: on a Hopper (cap>=9) device but "
+            "cluster_remote_mm_scenario was not found in any loaded module after "
+            "warmup. The fp16 mm OOB guard is NOT active; aborting rather than "
+            "risk an illegal-memory-access crash. (FlagGems mm module layout may "
+            "have changed.)"
+        )
+    else:
+        # Non-Hopper: the cluster path never runs, so an unpatched state is fine.
+        print(
+            "[mm-cluster-fix-plugin] no cluster_remote_mm_scenario found; "
+            "non-Hopper device, guard not needed (no-op).",
+            file=sys.stderr,
+            flush=True,
+        )
+
+
+def pytest_runtest_setup(item):
+    # backend module may load lazily on the first gems call; re-scan (idempotent).
+    _patch_loaded()

_seed_plugin.py

@@ -0,0 +1,39 @@
+"""pytest plugin: seed random/numpy/torch RNG for reproducible benchmark inputs.
+
+Data-dependent kernels (sort, topk, nonzero, ...) have value-dependent latency;
+a fixed seed makes every run generate byte-identical inputs, so the latency
+delta between two runs reflects the change under test, not the input data.
+"""
+
+from __future__ import annotations
+
+import sys
+
+_SEED = 0
+
+
+def pytest_configure(config):
+    seed = _SEED
+    # random.seed too: some kits (cutlass_scaled_mm) pick cases via random.shuffle.
+    import random
+    random.seed(seed)
+    seeded = ["random"]
+    try:
+        import numpy as _np
+        _np.random.seed(seed)
+        seeded.append("numpy")
+    except Exception:
+        pass
+    try:
+        import torch
+    except Exception as exc:  # torch missing should never happen here, stay safe
+        print(f"[seed-plugin] torch unavailable, seeded {'+'.join(seeded)} only: {exc}",
+              file=sys.stderr)
+        return
+    torch.manual_seed(seed)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed_all(seed)
+    seeded.append("torch")
+    print(f"[seed-plugin] manual_seed({seed}) for {'+'.join(seeded)} "
+          "— reproducible benchmark inputs/cases",
+          file=sys.stderr, flush=True)

_shape_inject_plugin.py

@@ -0,0 +1,99 @@
+"""pytest plugin: make our shape yaml win over subclasses that override
+`set_shapes()` with hardcoded shapes (~9% of upstream Benchmark subclasses).
+
+Wraps `Benchmark.init_user_config` (runs right after `set_shapes()`): if the
+op_name has a yaml entry, overwrite `self.shapes` with the yaml shapes. Ops with
+no yaml entry keep their own shapes. An arity guard skips ops whose subclass
+normalizes shape arity after reading yaml (e.g. BLAS (B,M,N,K)↔(M,N,K)).
+"""
+
+from __future__ import annotations
+
+import os
+import sys
+
+import yaml as _yaml
+
+
+_YAML_CACHE: dict | None = None  # filled lazily from Config.shape_file
+_PATCHED = False
+
+
+def _load_yaml(path: str) -> dict:
+    if not path or not os.path.isfile(path):
+        return {}
+    try:
+        with open(path, "r") as f:
+            return _yaml.safe_load(f) or {}
+    except Exception as exc:
+        print(f"[shape-inject-plugin] failed to load {path}: {exc}", file=sys.stderr)
+        return {}
+
+
+def _normalize_shapes(raw):
+    """Mirror base.py's `[tuple(shape) for shape in self.shapes]`, defensively
+    handling scalar entries (`5` → `(5,)`)."""
+    out = []
+    for shape in raw:
+        if isinstance(shape, (list, tuple)):
+            out.append(tuple(shape))
+        else:
+            out.append((shape,))
+    return out
+
+
+def _patch_init_user_config():
+    global _PATCHED
+    if _PATCHED:
+        return
+    from benchmark import base as fg_base
+    from benchmark.conftest import Config as fg_Config
+
+    original = fg_base.Benchmark.init_user_config
+
+    def _shape_arity(shape):
+        if isinstance(shape, (list, tuple)):
+            return len(shape)
+        return 1
+
+    def patched(self):
+        original(self)
+        global _YAML_CACHE
+        if _YAML_CACHE is None:
+            _YAML_CACHE = _load_yaml(getattr(fg_Config, "shape_file", "") or "")
+        entry = _YAML_CACHE.get(getattr(self, "op_name", None))
+        if not isinstance(entry, dict):
+            return
+        shapes = entry.get("shapes")
+        if not shapes:
+            return
+        # Skip ops whose subclass normalizes shape arity after reading yaml (e.g.
+        # BLAS (B,M,N,K)↔(M,N,K)): detected as both sides being homogeneous in
+        # arity but differing. Heterogeneous subclass shapes are safe to override.
+        if getattr(self, "shapes", None) and len(self.shapes) > 0:
+            cur_arities = {_shape_arity(s) for s in self.shapes}
+            yaml_arities = {_shape_arity(s) for s in shapes}
+            cur_homogeneous = len(cur_arities) == 1
+            yaml_homogeneous = len(yaml_arities) == 1
+            if cur_homogeneous and yaml_homogeneous and cur_arities != yaml_arities:
+                print(
+                    f"[shape-inject-plugin] skip override for op_name={self.op_name!r}: "
+                    f"subclass produced homogeneous arity {next(iter(cur_arities))}, "
+                    f"yaml has homogeneous arity {next(iter(yaml_arities))} "
+                    f"— subclass normalization preserved",
+                    file=sys.stderr,
+                )
+                return
+        self.shapes = _normalize_shapes(shapes)
+        if "shape_desc" in entry:
+            self.shape_desc = entry["shape_desc"]
+
+    fg_base.Benchmark.init_user_config = patched
+    _PATCHED = True
+
+
+def pytest_configure(config):
+    _patch_init_user_config()
+    print("[shape-inject-plugin] Benchmark.init_user_config patched: "
+          "yaml shapes now win over hardcoded subclass set_shapes()",
+          file=sys.stderr, flush=True)

_shape_iter_inject_plugin.py

@@ -0,0 +1,100 @@
+"""pytest plugin: make our yaml shapes win over subclasses that hardcode their
+shape list inside an overridden `get_input_iter` (instead of `self.shapes`),
+which `_shape_inject_plugin` can't reach (conv padding / pool / nll /
+scaled_softmax families, ~17 ops).
+
+At `pytest_collection_finish` (all Benchmark subclasses defined), redirect
+`get_input_iter` ONLY on classes whose own implementation does not reference
+`self.shapes` (the source-level guard that protects well-behaved BLAS/generic
+classes). The redirect yields one input per yaml shape via the op's standard
+`input_fn(shape, dtype, device)`; ops with no yaml entry or a non-standard
+input_fn fall through unchanged. This also bounds compile cost vs the large
+hardcoded lists.
+"""
+
+from __future__ import annotations
+
+import inspect
+import os
+import sys
+
+import yaml as _yaml
+
+
+_YAML_CACHE: dict | None = None
+
+
+def _load_yaml(path: str) -> dict:
+    if not path or not os.path.isfile(path):
+        return {}
+    try:
+        with open(path, "r") as f:
+            return _yaml.safe_load(f) or {}
+    except Exception as exc:
+        print(f"[shape-iter-inject-plugin] failed to load {path}: {exc}", file=sys.stderr)
+        return {}
+
+
+def _reads_self_shapes(func) -> bool:
+    """True if the method references self.shapes (already covered by
+    _shape_inject_plugin, so don't redirect). Defaults True on introspection
+    failure (safe: leave the method alone)."""
+    try:
+        return "self.shapes" in inspect.getsource(func)
+    except (OSError, TypeError):
+        return True
+
+
+def _accepts_shape_dtype_device(fn) -> bool:
+    """True if `fn` can be called as `fn(shape, dtype, device)` (the standard
+    input_fn signature). Guards against blas-style `fn(b,m,n,k,dtype,device,...)`.
+    Defaults False on introspection failure (safe: fall back to original)."""
+    if not callable(fn):
+        return False
+    try:
+        inspect.signature(fn).bind(None, None, None)
+    except (TypeError, ValueError):
+        return False
+    return True
+
+
+def _make_patched(original):
+    def patched(self, dtype):
+        entry = (_YAML_CACHE or {}).get(getattr(self, "op_name", None))
+        shapes = entry.get("shapes") if isinstance(entry, dict) else None
+        input_fn = getattr(self, "input_fn", None)
+        if shapes and _accepts_shape_dtype_device(input_fn):
+            for shape in shapes:
+                yield from input_fn(tuple(shape), dtype, self.device)
+        else:
+            yield from original(self, dtype)
+    return patched
+
+
+def pytest_collection_finish(session):
+    global _YAML_CACHE
+    from benchmark import base as fg_base
+    from benchmark.conftest import Config as fg_Config
+
+    _YAML_CACHE = _load_yaml(getattr(fg_Config, "shape_file", "") or "")
+
+    def _subclasses(cls):
+        for sub in cls.__subclasses__():
+            yield sub
+            yield from _subclasses(sub)
+
+    patched = 0
+    seen = set()
+    for cls in _subclasses(fg_base.Benchmark):
+        if cls in seen:
+            continue
+        seen.add(cls)
+        own = cls.__dict__.get("get_input_iter")
+        # Redirect only classes with their own get_input_iter that ignore self.shapes.
+        if own is not None and not _reads_self_shapes(own):
+            cls.get_input_iter = _make_patched(own)
+            patched += 1
+
+    print(f"[shape-iter-inject-plugin] get_input_iter redirected on {patched} "
+          f"hardcoded-shape Benchmark classes; yaml shapes now win",
+          file=sys.stderr, flush=True)

run_pytest.sh

@@ -0,0 +1,108 @@
+#!/bin/bash
+# Single-operator perf benchmark (manual debug).
+# Edit OP/OP_FILE/INLINE_YAML below (or override via env: OP=... SHAPE_FILE=... ./run_pytest.sh).
+# All artifacts go under runs/<op>_<timestamp>/: run.log, shapes.yaml,
+# autotune_records/, ttgir/.
+set -euo pipefail
+
+OP="${OP:-fused_marlin_moe_mxfp4}"
+OP_FILE="${OP_FILE:-fused_marlin_moe}"   # benchmark file is test_<OP_FILE>.py
+SHAPE_FILE="${SHAPE_FILE:-}"
+
+# Inline shape yaml (used when SHAPE_FILE is empty); top-level key must be the op name.
+read -r -d '' INLINE_YAML <<'YAML' || true
+fused_marlin_moe_mxfp4:
+  # 4 MoE models x 4 token counts (M=1,16,64,256) = 16 shapes
+  shapes:
+    # Mixtral (E=8)
+    - [1,   8, 4096, 14336,  2]
+    - [16,  8, 4096, 14336,  2]
+    - [64,  8, 4096, 14336,  2]
+    - [256, 8, 4096, 14336,  2]
+    # DeepSeek-V3 (E=256, H=7168)
+    - [1,   256, 7168, 2048, 8]
+    - [16,  256, 7168, 2048, 8]
+    - [64,  256, 7168, 2048, 8]
+    - [256, 256, 7168, 2048, 8]
+    # Qwen3 (E=512)
+    - [1,   512, 4096, 1024, 10]
+    - [16,  512, 4096, 1024, 10]
+    - [64,  512, 4096, 1024, 10]
+    - [256, 512, 4096, 1024, 10]
+    # DeepSeek-V4-Flash (E=256, H=4096)
+    - [1,   256, 4096, 2048, 6]
+    - [16,  256, 4096, 2048, 6]
+    - [64,  256, 4096, 2048, 6]
+    - [256, 256, 4096, 2048, 6]
+  shape_desc: "num_tokens, num_experts, hidden_size, intermediate_size, topk"
+YAML
+
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+FLAGGEMS_DIR="${FLAGGEMS_DIR:-/workspace/FlagGems-dev}"
+TEST_FILE="$FLAGGEMS_DIR/benchmark/test_${OP_FILE}.py::test_${OP}"
+
+OUT_DIR="$SCRIPT_DIR/runs/${OP}_$(date +%Y%m%d-%H%M%S)"
+LOG_FILE="$OUT_DIR/run.log"
+mkdir -p "$OUT_DIR"
+
+export PYTHONPATH="$SCRIPT_DIR${PYTHONPATH:+:$PYTHONPATH}"
+PLUGINS=(
+  -p _seed_plugin
+  -p _shape_inject_plugin
+  -p _shape_iter_inject_plugin
+  -p _bespoke_shape_plugin
+  # -p _mm_cluster_fix_plugin
+  -p _autotune_record_plugin
+  -p _cudagraph_plugin
+  -p _ir_meta_plugin
+)
+# Autotune record/replay (see _autotune_record_plugin.py):
+# - default: record chosen configs to $OUT_DIR/autotune_records/<op>.json
+# - REPLAY_FROM=<previous runs/<op>_<ts> dir>: replay that run's recorded
+#   configs instead (for A/B runs where both sides must use the same config).
+REPLAY_FROM="${REPLAY_FROM:-}"
+if [[ -n "$REPLAY_FROM" ]]; then
+  AUTOTUNE_ENV="FLAGGEMS_PERF_AUTOTUNE_REPLAY_DIR=$REPLAY_FROM/autotune_records"
+  [[ -f "$REPLAY_FROM/autotune_records/$OP.json" ]] || \
+    echo ">>> warning: $REPLAY_FROM/autotune_records/$OP.json not found; replay will fall back to autotune" >&2
+else
+  AUTOTUNE_ENV="FLAGGEMS_PERF_AUTOTUNE_RECORD_DIR=$OUT_DIR/autotune_records"
+  mkdir -p "$OUT_DIR/autotune_records"
+fi
+export FLAGGEMS_PERF_CURRENT_OP="$OP"
+
+# Resolve the shape file: use SHAPE_FILE, or write the inline yaml to a temp file.
+if [[ -z "$SHAPE_FILE" ]]; then
+  SHAPE_FILE="$(mktemp --suffix=.yaml)"
+  printf '%s\n' "$INLINE_YAML" > "$SHAPE_FILE"
+  trap 'rm -f "$SHAPE_FILE"' EXIT
+fi
+cp -f "$SHAPE_FILE" "$OUT_DIR/shapes.yaml"  # archive the shape used
+
+export PYTHONUNBUFFERED=1  # unbuffered live output
+
+{
+  # _ir_meta_plugin dumps organized ttgir (only variants actually launched
+  # outside the autotune sweep) into TTGIR_DUMP_DIR at process exit, so the
+  # dump happens even if a CUDA crash kills the run. Layout/legend: see
+  # <dump>/naming.md and index.tsv (which also lists unused sweep losers).
+  CACHE_DIR="$OUT_DIR/.triton_cache"
+  rm -rf "$CACHE_DIR"; mkdir -p "$CACHE_DIR"
+  status=0
+  TRITON_CACHE_DIR="$CACHE_DIR" \
+  FLAGGEMS_PERF_TTGIR_DUMP_DIR="$OUT_DIR/ttgir" \
+  env "$AUTOTUNE_ENV" \
+  USE_FLAGTUNE=1 python -u -m pytest -s "$TEST_FILE" \
+    "${PLUGINS[@]}" \
+    --shape_file "$SHAPE_FILE" \
+    --level core --mode kernel || status=$?
+
+  rm -rf "$CACHE_DIR"
+
+  if (( status == 0 )); then
+    echo ">>> done. outputs in $OUT_DIR"
+  else
+    echo ">>> FAILED (pytest exit $status). partial outputs in $OUT_DIR"
+  fi
+  exit "$status"
+} 2>&1 | tee "$LOG_FILE"