fix: never let the parallel sweep cover only part of the shape set

The sharded autotune sweep could publish a config set with holes, and a run
using it reports itself as replaying while every uncovered key silently
self-tunes. That is the exact failure the record/replay caliber exists to
prevent, and it was reachable two ways.

Shape lookup was op_name-only, but BLAS operators keep their shapes under the
class name: core_shapes.yaml has BlasBenchmark, not mm. Every BLAS operator
therefore resolved to zero shapes and the plugin disabled itself with a message
that read like a normal skip. It now walks the benchmark class chain the way
FlagGems' set_shapes does. The chain is recovered statically -- the class comes
from the op_name= keyword of the constructor call in the test module, its bases
from that module plus benchmark/base.py -- because the plugin runs at
pytest_configure, before any Benchmark instance exists. Resolved keys are named
in the log, and a total miss lists the candidates it tried.

Two cases can still produce partial coverage, and both now refuse rather than
shard a subset:

- --level comprehensive, where set_shapes() merges set_more_shapes() on top of
  the yaml. That addition is arbitrary per-subclass Python and invisible here.
  Measured on mm: 5 of 29 shapes swept, 90 AUTOTUNE_REPLAY_FALLBACK
  reason=key_missing in the run that consumed the result.
- a shard exiting non-zero, whose shapes were never swept. This used to append a
  WARNING to a success message and publish the merge anyway; observed with 3 of
  5 shards dead. The whole sweep is now discarded, worker logs kept.

Both fall back to the normal serial sweep, which costs wall time rather than the
same-config premise.

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

README.md

@@ -20,6 +20,11 @@ OP=softmax OP_FILE=softmax SHAPE_FILE=my_shapes.yaml bash run_pytest.sh
 
 # 测 FlagGems-vllm 仓库的算子：换 FLAGGEMS_DIR 即可（插件自动适配其包名 flaggems_vllm）
 FLAGGEMS_DIR=/workspace/dev/FlagGems-vllm OP=fused_marlin_moe OP_FILE=fused_marlin_moe bash run_pytest.sh
+
+# 对齐 FlagGems 原生节点（不传 --shape_file，沿用 --level core 的 core_shapes）
+ZL_BACKEND=cuda \
+TEST_NODE=/workspace/FlagGems/benchmark/test_fused_marlin_moe.py::test_fused_marlin_moe_mxfp4 \
+bash run_pytest.sh
 ```
 
 上游现为两个仓库，benchmark 体系同构、均受支持：主仓 `/workspace/dev/FlagGems`（Python 包名 `flag_gems`，默认；`fused_marlin_moe_mxfp4` 仅此仓有）与 `/workspace/dev/FlagGems-vllm`（包名 `flaggems_vllm`）。两包需各自 `pip install -e <仓库> --no-deps` 安装（`--no-deps` 避免动 FlagTree 的 triton）。
@@ -35,7 +40,7 @@ softmax:
 
 上面这条命令是 **record 模式**：现场 autotune、把选中的 config 存档。它用于单次摸底或给 replay 提供基准，**两次 record 的数字不能互相比**——做对比走下面的「A/B 对比测试标准流程」。
 
-跑之前先 `nvidia-smi` 确认目标卡空闲：共享机器上别的任务会把两侧一起等比拖慢，出一份看似自洽实则作废的数据。多卡机器用 `CUDA_VISIBLE_DEVICES=<空闲卡号>` 选卡。**测量始终单进程单卡串行**，这是计时可比的前提；多卡仅用于可选的 sweep 加速（见「加速：多卡并行 sweep」）。
+跑之前确认目标卡空闲：共享机器上别的任务会把两侧一起等比拖慢，出一份看似自洽实则作废的数据。用 `ZL_BACKEND` 选运行时（默认 `cuda`），再用该后端的可见性变量选卡（CUDA 为 `CUDA_VISIBLE_DEVICES`，也可设 `ZL_DEVICE`）。**测量始终单进程单设备串行**，这是计时可比的前提；多卡仅用于可选的 sweep 加速（见「加速：多卡并行 sweep」）。非 CUDA 后端默认走上游 `do_bench`，不会尝试 CUDA Graph。
 
 脚本内固定了 `--level core --mode kernel`，其余可调项都走环境变量（见下）；需要改动这两个口径本身时直接编辑 `run_pytest.sh` 中的 pytest 命令行。
 
@@ -45,10 +50,14 @@ softmax:
 
 | 变量 | 默认 | 说明 |
 |-----|------|------|
-| `OP` | `fused_marlin_moe_mxfp4` | 测试函数名（`test_` 之后的部分） |
+| `OP` | `fused_marlin_moe_mxfp4` | 测试函数名（`test_` 之后的部分）；设了 `TEST_NODE` 时由节点覆盖 |
-| `OP_FILE` | `fused_marlin_moe` | benchmark 文件名（`test_` 与 `.py` 之间的部分） |
+| `OP_FILE` | `fused_marlin_moe` | benchmark 文件名（`test_` 与 `.py` 之间的部分）；设了 `TEST_NODE` 时由节点覆盖 |
-| `SHAPE_FILE` | 空（用脚本内置 yaml） | shape yaml 路径 |
+| `TEST_NODE` | 空 | FlagGems pytest 节点（`…/benchmark/test_<op_file>.py::test_<op>`）。设了且未指定 `SHAPE_FILE` 时默认不传 `--shape_file` |
-| `FLAGGEMS_DIR` | `/workspace/dev/FlagGems` | FlagGems 仓库路径（测 vllm 仓库时指向 `/workspace/dev/FlagGems-vllm`） |
+| `SHAPE_FILE` | 空 | shape yaml 路径；空且无 `TEST_NODE` 时用脚本内置 yaml；`-`/`upstream`/`none` 表示沿用上游在 `BENCH_LEVEL` 下的原生 shape 集 |
+| `SHAPE_MODE` | 空 | `inline` / `file` / `upstream`；通常不用手设，`TEST_NODE` 会把空 `SHAPE_FILE` 推成 `upstream` |
+| `FLAGGEMS_DIR` | `/workspace/dev/FlagGems` | FlagGems 仓库路径。测试文件缺失时回退 `/workspace/FlagGems`；`TEST_NODE` 也可直接带出仓库根 |
+| `ZL_BACKEND` | `cuda` | 设备后端：`cuda` / `musa` / `rocm` / `ascend`（别名 `nvidia`/`amd`/`npu` 等；未知名直接失败） |
+| `ZL_DEVICE` | 空 | 父进程看到的设备 id，写入 `run_identity.json`；批量层绑卡后也会设 |
 
 **怎么测**
 
@@ -56,8 +65,14 @@ softmax:
 |-----|------|------|
 | `REPLAY_FROM` | 空（record 模式） | 指向某次历史 run 目录，replay 其 autotune 选择。做 A/B 时 B 侧必须设（见「A/B 对比测试标准流程」） |
 | `USE_FLAGTUNE` | `0` | `0` 走普通 autotune（跳过搜索、快速验证）；`1` 用 FlagTune 扩展调优空间（首跑全量搜索、慢）。A/B 两侧须取同值 |
+| `BENCH_LEVEL` | `core` | 透传上游 `--level`。`comprehensive` 相对 `core` 做两件事：① 在 `set_shapes()` 里并入子类 `set_more_shapes()`（`mm` 因此从 5 个 shape 变 29 个，含 N/K 到 152064 的大 K 与 `models_shapes/*.yaml` 的模型 shape）；② 在 `get_input_iter()` 里对每个 shape 再产出一次变体（BLAS 是转置侧），用例数翻倍。注意 ① 只在 `SHAPE_MODE=upstream` 下可见——传了 `SHAPE_FILE` 时 `_shape_inject_plugin` 会在上游合并之后再覆盖一次、把 shape 集拉回 yaml 内容（BLAS 那类 arity 不一致的情形除外，插件会退让）；② 不经插件，任何 shape 来源下都生效。sweep 成本同步放大，shape 多时配合 `PARALLEL_WARMUP_GPUS`。A/B 两侧须取同值 |
 | `DTYPES` | 空（上游默认 dtype 扫描） | 空格分隔的 dtype 白名单（如 `"bfloat16 float16"`），逐个转为上游 `--dtypes`。算子不支持时上游报 `can't be supported by this op` |
 | `PARALLEL_WARMUP_GPUS` | 空（串行） | 设为 `N`（≥2）时把 autotune sweep 分片到 N 卡并行，测量仍单卡串行；不设或 `<2` 则完全串行（见「加速：多卡并行 sweep」） |
+| `ZL_ENABLE_CUDAGRAPH` | `auto` | `auto`/`1` 仅在 CUDA 上走 Graph 计时；非 CUDA 即使设 `1` 也拒绝并保持 `timer=direct` |
+| `ZL_ENABLE_AUTOTUNE` | `1` | `0` 关闭 autotune record/replay 插件 |
+| `ZL_ENABLE_IR` | `1` | `0` 关闭 TTGIR / 编译元数据插件 |
+| `ZL_ENABLE_MM_CLUSTER` | `0` | `1` 且后端支持时加载 Hopper mm cluster 规避 |
+| `ZL_EXTRA_PLUGINS` | 空 | 空格分隔、追加到 `-p` 列表的插件名 |
 
 **输出**
 
@@ -72,9 +87,11 @@ softmax:
 ```
 runs/softmax_20260716-031752/
 ├── run.log                  # 完整日志（含 SUCCESS 行的 latency/speedup 表；已去 ANSI 色的纯文本）
+├── run_identity.json        # provider/backend/device/timer/vendor（插件合并写入）
 ├── shapes.yaml              # 本次实际使用的 shape（存档）
 ├── autotune_records/        # Triton autotune 选中的 config（供 replay）
-│   └── softmax.json
+│   ├── softmax.json
+│   └── meta.json            # provider/backend；跨后端 replay 会被拒绝
 └── ttgir/                   # 按 shape 分组的 IR 落盘
     ├── 1024x1024/
     │   └── softmax_kernel_inner/
@@ -152,7 +169,7 @@ shape 多时墙上时间几乎全在 autotune sweep，而不在测量——被
 PARALLEL_WARMUP_GPUS=8 SHAPE_FILE=shape/DeepSeek-V4-Flash-p32768d1024.yaml bash run_pytest.sh
 ```
 
-加速比取决于 sweep 占比：shape 多、config 空间大的算子收益最明显，可达数倍。不设该变量（或设 <2）时插件静默 no-op，常驻 `PLUGINS` 数组即可。**产物结构、`REPLAY_FROM` 用法、latency 表格式全不变**——合并后的 config 就写进本次 run 的 `autotune_records/<op>.json`，分片临时目录跑完即删。
+加速比取决于 sweep 占比：shape 多、config 空间大的算子收益最明显，可达数倍。不设该变量（或设 <2）时插件静默 no-op。**产物结构、`REPLAY_FROM` 用法、latency 表格式全不变**——合并后的 config 就写进本次 run 的 `autotune_records/<op>.json`，分片临时目录跑完即删。
 
 **测量本身不并行**：N 个进程压满同机多卡会通过功耗墙/散热耦合，单卡 latency 被邻居拖慢且不可复现。并行只用于"决定哪个 config 胜出"，测量仍是单进程单卡、与不开插件走同一条路径。
 
@@ -162,7 +179,17 @@ PARALLEL_WARMUP_GPUS=8 SHAPE_FILE=shape/DeepSeek-V4-Flash-p32768d1024.yaml bash
 WARNING 1 config key(s) got different winners across shards — parallel interference reached the sweep
 ```
 
-**冲突数就是这次加速的可信度指标**：0 可放心用；偏多说明 config 选择已被污染，出正式结论前不设该变量重跑一遍。其余行为（`REPLAY_FROM` 已设或可见卡不足 2 张时跳过、分片失败回退串行、轮转分片而非按块切、用子进程而非 xdist 的原因）见插件 docstring。分片绑卡遵守继承来的 `CUDA_VISIBLE_DEVICES`——用它选过卡时，分片只会落在你选的那几张上。
+**冲突数就是这次加速的可信度指标**：0 可放心用；偏多说明 config 选择已被污染，出正式结论前不设该变量重跑一遍。
+
+插件会在三种情况下**主动放弃分片、退回串行 sweep**，日志各有独立措辞，都不是静默的：
+
+- **`BENCH_LEVEL` 不是 `core`**。`comprehensive` 会让上游在 `set_shapes()` 里并入 `set_more_shapes()`，那部分是各子类自己的 Python（硬编码元组加 `models_shapes/*.yaml`），插件只看得到 yaml 里的 shape 列表。按子集分片会让未覆盖的 key 全部现场自调优，而这一轮却自称在 replay——`mm` 上实测是 29 个 shape 只扫到 5 个、90 次 `AUTOTUNE_REPLAY_FALLBACK reason=key_missing`。
+- **任一分片非零退出**。死掉的分片的 shape 从未被扫过，合并出来的 config 集是带洞的，publish 之后每个洞都会静默自调优。因此整轮 sweep 直接作废，worker 日志保留在临时目录里备查。
+- yaml 里查不到该算子的 shape（见下）、`REPLAY_FROM` 已设、或可见卡不足 2 张。
+
+yaml key 的查找与上游对齐：先按 `op_name`，查不到再沿 Benchmark 类链逐级尝试类名（`mm` 因此命中 `BlasBenchmark`）。类链是静态解析测试模块里 `op_name=` 那次构造调用得到的，不 import FlagGems。日志会打出命中的 key；一个都没命中时会列出尝试过的候选。
+
+其余行为（轮转分片而非按块切、用子进程而非 xdist 的原因）见插件 docstring。分片绑卡遵守当前后端的可见性变量（CUDA 为 `CUDA_VISIBLE_DEVICES`）——用它选过卡时，分片只会落在你选的那几张上。
 
 ## 批量多算子测试（run_batch.py）
 
@@ -170,7 +197,8 @@ WARNING 1 config key(s) got different winners across shards — parallel interfe
 
 ```bash
 # 全量筛查（双仓库全部可收集算子，排除带无条件 skip 标记的与 ops/dismiss.txt 里的）
-python run_batch.py --all --gpus 0,1,2,3,4,5,6,7 --dtypes bfloat16 --op-timeout 1800
+python run_batch.py --all --devices 0,1,2,3,4,5,6,7 --backend cuda --dtypes bfloat16 --op-timeout 1800
+# --gpus 仍是 --devices 的别名；两者同时给出会报错
 
 # 指定子集：默认 main 仓，vllm 仓加前缀；同名测试函数用 @文件名 消歧
 python run_batch.py --ops softmax vllm:fused_marlin_moe nextafter_@nextafter_
@@ -195,7 +223,7 @@ python run_batch.py --ops-file my_ops.txt --replay-root runs/batch_xxx
 
 **执行语义**：每算子独立子进程（CUDA crash 只废单个算子）；`--op-timeout` 对进程组 SIGTERM→SIGKILL，记 `TIMEOUT` 不重试、批次继续；pytest rc=1 与信号杀最多重试 2 次，连续同 rc 视为确定性失败提前止损；`--dtypes` 有两级降级——遇上游 `can't be supported by this op` 即去掉限制重跑（备注 `dtype_fallback`），其余失败且无成功行时也去掉限制最后救一次（备注 `dtype_rescue`，覆盖 torch baseline 对受限 dtype 编译失败的场景）。runtime skipif 与上游 `xfail` 都判 SKIP（上游声明的无信号状态）；`.complete` 标记只给 PASS/SKIP，`--resume` 据此复用。
 
-**产物**：`<batch>/<repo>/<op>/` 与单算子 run 目录结构完全一致（run.log、autotune_records/、ttgir/……），所以任何一个算子都可以事后单独 `REPLAY_FROM` 复测。批量层额外产出 `ops_status.csv`（每算子一行：状态/耗时/行数/备注，边跑边原子更新）与 `summary.csv`（每测量行一行：dtype、latency、speedup、Size Detail，来源是 run.log 的结果表，`batch_summary.py` 也可单独对旧批次重跑）。备注列聚合 `no_cudagraph_rows`/`replay_fallback` 计数，解读口径时先看这列。
+**产物**：`<batch>/<repo>/<op>/` 与单算子 run 目录结构完全一致（run.log、run_identity.json、autotune_records/、ttgir/……），所以任何一个算子都可以事后单独 `REPLAY_FROM` 复测。批量层额外产出 `ops_status.csv`（每算子一行：backend/device/timer、状态/耗时/行数/备注，边跑边原子更新；`gpu` 列与 `device` 同值）与 `summary.csv`（每测量行一行：backend/device/timer、dtype、latency、speedup、Size Detail，来源是 run.log 的结果表，`batch_summary.py` 也可单独对旧批次重跑）。备注列聚合 `no_cudagraph_rows`/`replay_fallback` 计数，解读口径时先看这列。不同后端、不同 timer 的行不要当成同口径。
 
 ## 计时口径：cudagraph 的预热、回退与精度
 
@@ -208,27 +236,28 @@ python run_batch.py --ops-file my_ops.txt --replay-root runs/batch_xxx
 回退本身不影响 A/B 公平性（两侧同一算子回退行为一致），但**回退口径的测量误差更大**：do_bench 每次迭代都走完整的 Python → launch 路径，kernel 越小，launch 开销和 CPU 侧抖动在数字里占比越高——亚毫秒级 kernel 上两种口径可差 2 倍以上，且行间波动更明显。解读这类算子的结果时：
 
 - 小 shape 行的绝对值和小幅（<10%）差异不要过度解读，优先看大 shape 行；
-- 需要更高置信度时，同一配置多跑几次取中位数，或对该算子直接注释掉 `_cudagraph_plugin` 统一用 do_bench 口径（消除同表混两种口径的问题）。
+- 需要更高置信度时，同一配置多跑几次取中位数，或设 `ZL_ENABLE_CUDAGRAPH=0` 统一用 `timer=direct`（消除同表混两种口径的问题）。
 
 回退是按测量点发生的，同一份 latency 表里可能混有两种口径的行；若不确定，先 `grep BENCHMARK_DIRECT_NO_CUDAGRAPH run.log` 确认哪些行是回退口径再下结论。
 
 ## 插件说明
 
-脚本通过 `-p` 加载以下插件（`run_pytest.sh` 的 `PLUGINS` 数组，可按需注释）：
+`-p` 列表由 `plugin_plan.py` 按 `ZL_BACKEND` 生成，不再在 `run_pytest.sh` 里写死。`_cudagraph_plugin` **不**出现在默认列表上：只有 CUDA 且 `_timing_plugin` 选中 Graph 时才会 import 它。
 
 | 插件 | 作用 | 何时关闭 |
 |-----|------|---------|
-| `_device_guard_plugin` | 在 import flag_gems 之前直接用 torch 探测到 NVIDIA 卡就设 `GEMS_VENDOR=nvidia`，跳过 flag_gems 启动时 `nvidia-smi` 子进程探测（该探测在部分 fork 环境下会挂在 `wait4` 上导致 import 卡死）。副作用：import 期即初始化 CUDA context，故与 pytest-xdist 不兼容（本框架不用 xdist） | 一般无需关：已设 `GEMS_VENDOR`/`FLAGGEMS_VENDOR` 等 env 时自动跳过，非 NVIDIA 卡上自动 no-op |
+| `_device_setup_plugin` | 按后端设 `GEMS_VENDOR`（cuda→nvidia），跳过 flag_gems 启动时可能挂死的 `nvidia-smi` 探测，并写入 `run_identity.json`。旧名 `_device_guard_plugin` 仍是兼容 shim | 已设 `GEMS_VENDOR` 时不覆盖 |
 | `_parallel_warmup_plugin` | 由 `PARALLEL_WARMUP_GPUS=N` 激活（未设则静默 no-op）：autotune sweep 分片到 N 卡并行，测量仍单卡串行、产物结构不变（见「加速：多卡并行 sweep」） | 不设该变量即关闭 |
 | `_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 环境变量二选一激活） | 不关 |
+| `_autotune_record_plugin` | record/replay Triton autotune 选择；`autotune_records/meta.json` 记录 backend，跨后端 replay 拒绝 | `ZL_ENABLE_AUTOTUNE=0` |
-| `_cudagraph_plugin` | `do_bench` → `do_bench_cudagraph`（kernel 纯耗时；内部先做跨流同步，修过一个间歇性 illegal instruction）。**capture 前按 warmup 时间预算显式预热**（先一次丢弃跑吸收 autotune/JIT 编译，再稳态预热），消除首轮抖动。无法 graph capture 的 kernel 自动回退并打 `BENCHMARK_DIRECT_NO_CUDAGRAPH` 标记 | 需要与他人的普通 do_bench 数据对齐时注释掉 |
+| `_timing_plugin` | 按后端选择计时路径：CUDA 可走 Graph（内部调用 `_cudagraph_plugin`），其余后端一开始就是上游 `do_bench`（`timer=direct`） | `ZL_ENABLE_CUDAGRAPH=0` 强制 direct |
-| `_pretty_report_plugin` | 结果表整理 + 着色：全表相同的输入折叠成表头下一行图例，每行 Size Detail 只留随行变化的部分（shape 保持 `torch.Size([...])` 原样，MoE 类算子单行从 ~400 字符缩到一屏内）；SUCCESS 绿 / FAILED 红。列名与 `SUCCESS` 字样保持上游原文，`run.log` 的 grep/解析不受影响 | 需要与上游原始表格逐字对齐时注释掉 |
+| `_cudagraph_plugin` | 仅由 `_timing_plugin` 在 CUDA 上调用。`do_bench` → `do_bench_cudagraph`，capture 前按 warmup 预算显式预热；无法 capture 时回退并打 `BENCHMARK_DIRECT_NO_CUDAGRAPH` | 不要直接 `-p` 加到非 CUDA 运行 |
-| `_ir_meta_plugin` | 编译期记录每个变体的 constexpr/签名/特化；launch 钩子统计每个 (kernel, shape) 的实际使用；退出时按上述结构 dump ttgir | 不关 |
+| `_pretty_report_plugin` | 结果表整理 + 着色：全表相同的输入折叠成表头下一行图例，每行 Size Detail 只留随行变化的部分（shape 保持 `torch.Size([...])` 原样，MoE 类算子单行从 ~400 字符缩到一屏内）；SUCCESS 绿 / FAILED 红。列名与 `SUCCESS` 字样保持上游原文，`run.log` 的 grep/解析不受影响 | 需要与上游原始表格逐字对齐时去掉 |
-| `_mm_cluster_fix_plugin` | Hopper fp16 mm cluster kernel 越界崩溃的运行时规避 | 默认注释；测 fp16 mm 崩溃时打开 |
+| `_ir_meta_plugin` | 编译期记录每个变体的 constexpr/签名/特化；launch 钩子统计每个 (kernel, shape) 的实际使用；退出时按上述结构 dump ttgir。Triton knobs 缺失时 fail-soft | `ZL_ENABLE_IR=0` |
+| `_mm_cluster_fix_plugin` | Hopper fp16 mm cluster kernel 越界崩溃的运行时规避 | 默认关；`ZL_ENABLE_MM_CLUSTER=1` 或 `ZL_EXTRA_PLUGINS` |
 
 ## 常见问题
 
@@ -257,5 +286,7 @@ shape 多时这部分会主导墙上时间，可设 `PARALLEL_WARMUP_GPUS=N` 并
 ## 依赖假设
 
 - FlagGems benchmark 体系（`benchmark/base.py` 的 `Benchmark` 类、conftest 的 `--shape_file/--level/--mode` 选项）；主仓与 FlagGems-vllm 仓库同构，`_autotune_record_plugin` 会按实际被 import 的包（`flag_gems` / `flaggems_vllm`）挂 LibTuner 补丁；
-- Triton 需支持 `knobs.compilation.listener`、`kernel_load_end_hook`、`launch_enter_hook`（当前 FlagTree 的 triton 3.6 满足）；
+- 设备后端由 `ZL_BACKEND` / `--backend` 选择（`cuda` / `musa` / `rocm` / `ascend`）。inventory 仍是 FlagGems 逻辑算子清单，同一 case 可在多个后端上跑，不要按硬件复制 CSV；
+- Triton 需支持 `knobs.compilation.listener`、`kernel_load_end_hook`、`launch_enter_hook`（当前 FlagTree 的 triton 3.6 满足）。缺 hook 时 IR / autotune 插件 fail-soft，不会把进程打挂；
 - 插件通过 monkeypatch 挂钩上游内部结构，FlagGems/Triton 大版本升级后若行为异常，优先检查各插件 pytest_configure 输出的注册日志是否还正常打印。
+- 无 GPU 的适配器自检：`python test_backends.py`。

_autotune_record_plugin.py

@@ -214,6 +214,57 @@ def _replay_run(original):
     return runner
 
 
+def _write_record_meta(record_dir: str) -> None:
+    """Sidecar next to <op>.json. Replay refuses a different backend.
+
+    The config JSON itself stays `{kernel: {key: cfg}}` so existing REPLAY_FROM
+    directories keep working.
+    """
+    try:
+        from backends import get_backend
+        backend_name = get_backend().name
+    except Exception:
+        backend_name = os.environ.get("ZL_BACKEND", "")
+    meta = {
+        "schema_version": 1,
+        "provider": "flaggems",
+        "backend": backend_name,
+        "op": _op_name(),
+    }
+    path = Path(record_dir) / "meta.json"
+    try:
+        path.parent.mkdir(parents=True, exist_ok=True)
+        path.write_text(json.dumps(meta, indent=2, ensure_ascii=False) + "\n")
+    except Exception as exc:
+        from _term_style import tag
+        print(f"{tag('[autotune-record-plugin]')} meta.json write failed: {exc}",
+              file=sys.stderr, flush=True)
+
+
+def _replay_backend_ok(replay_dir: str) -> bool:
+    path = Path(replay_dir) / "meta.json"
+    if not path.is_file():
+        return True
+    try:
+        meta = json.loads(path.read_text())
+    except Exception:
+        return True
+    if not isinstance(meta, dict):
+        return True
+    recorded = str(meta.get("backend") or "").strip()
+    if not recorded:
+        return True
+    try:
+        from backends import get_backend
+        current = get_backend().name
+    except Exception:
+        current = os.environ.get("ZL_BACKEND", "").strip()
+    if recorded != current:
+        _emit_marker(f"backend_mismatch recorded={recorded} current={current}")
+        return False
+    return True
+
+
 def _load_replay_dir(replay_dir: str) -> int:
     path = _record_path(replay_dir)
     if not path.is_file():
@@ -256,29 +307,48 @@ def _dump_record(record_dir: str) -> None:
 def pytest_configure(config):
     global _wrap
     from _term_style import tag
+    from _run_identity import update_identity
     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(f"{tag('[autotune-record-plugin]')} error: both RECORD_DIR and "
               "REPLAY_DIR set; ignoring both (no-op)", file=sys.stderr, flush=True)
+        update_identity(autotune_mode="uncontrolled")
         return
     if not record_dir and not replay_dir:
+        update_identity(autotune_mode="uncontrolled")
         return
 
+    try:
         import triton.runtime.autotuner as _autotuner
+    except Exception as exc:
+        print(f"{tag('[autotune-record-plugin]')} Triton Autotuner unavailable "
+              f"({exc}); autotune_mode=uncontrolled", file=sys.stderr, flush=True)
+        update_identity(autotune_mode="uncontrolled")
+        return
+
+    if replay_dir and not _replay_backend_ok(replay_dir):
+        print(f"{tag('[autotune-record-plugin]')} refusing replay across backends",
+              file=sys.stderr, flush=True)
+        update_identity(autotune_mode="replay_fallback")
+        return
+
     _wrap = _record_run if record_dir else _replay_run
     _autotuner.Autotuner.run = _wrap(_autotuner.Autotuner.run)
 
     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.
+        _write_record_meta(record_dir)
         atexit.register(_dump_record, record_dir)
+        update_identity(autotune_mode="record")
         print(f"{tag('[autotune-record-plugin]')} recording autotune configs to "
               f"{_record_path(record_dir)} (Autotuner)",
               file=sys.stderr, flush=True)
     else:
         loaded = _load_replay_dir(replay_dir)
+        update_identity(autotune_mode="replay")
         print(f"{tag('[autotune-record-plugin]')} replaying {loaded} recorded entries from "
               f"{_record_path(replay_dir)} (Autotuner)",
               file=sys.stderr, flush=True)

_cudagraph_plugin.py

@@ -1,25 +1,18 @@
-"""pytest plugin: swap `triton.testing.do_bench` for graph-capture timing.
+"""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
+Invoked by `_timing_plugin` when the active backend declares
-line out to time with plain do_bench). Monkey-patches do_bench →
+`supports_cuda_graph` (CUDA only). Do not put this module on the default `-p`
+list for MUSA/ROCm/Ascend: those backends must start on plain do_bench, not
+attempt capture-then-fallback. 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. Before capturing, fn is explicitly warmed up for the caller's
 warmup time budget (do_bench_cudagraph's own 5-iter warmup is too short to
 settle autotune/JIT for the MoE kernels, which makes the first captured
-measurement unstable run-to-run). The capture itself runs under a scratch RNG
+measurement unstable run-to-run). The capture itself runs under a scratch CUDA
-state so that an aborted capture cannot poison the process-wide generator
+RNG state so that an aborted capture cannot poison the process-wide generator
 (rationale: _capture_with_scratch_rng).
-
-Everything here goes through Triton's active device interface rather than
-torch.cuda. FlagTree's do_bench is already written that way, but its
-do_bench_cudagraph is not: it hardcodes torch.cuda.stream/Event/CUDAGraph, so
-on a non-CUDA backend the graph path raised and every measurement silently fell
-back to plain do_bench. On MThreads the interface is torch_musa, which supplies
-Stream, Event, synchronize, graph and MUSAGraph, so the same capture works once
-the calls are taken from the interface. CUDA still calls upstream's
-do_bench_cudagraph, so its numbers are unchanged.
 """
 
 from __future__ import annotations
@@ -33,33 +26,8 @@ import triton.testing as _tt
 _ORIGINAL_DO_BENCH = _tt.do_bench
 
 
-def _device_interface():
+def _is_fatal_cuda_error(exc: Exception) -> bool:
-    """Triton's active device interface: torch.cuda on NVIDIA, torch_musa on MThreads."""
+    """True for CUDA context errors (illegal address, OOM, device-side assert):
-    from triton.runtime import driver
-
-    return driver.active.get_device_interface()
-
-
-def _is_cuda(di) -> bool:
-    import torch
-
-    return di is torch.cuda
-
-
-def _graph_class(di):
-    """The interface's graph type: CUDAGraph, MUSAGraph, or whatever it names."""
-    for name in ("CUDAGraph", "MUSAGraph"):
-        graph = getattr(di, name, None)
-        if graph is not None:
-            return graph
-    for name in dir(di):
-        if name.endswith("Graph"):
-            return getattr(di, name)
-    raise RuntimeError("device interface exposes no graph type")
-
-
-def _is_fatal_device_error(exc: Exception) -> bool:
-    """True for device 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.
 
@@ -81,12 +49,8 @@ def _is_fatal_device_error(exc: Exception) -> bool:
     return any(m in msg for m in markers)
 
 
-# Retained under its historical name; the marker list is backend independent.
+def _warmup_before_capture(fn, warmup_ms, grad_to_none):
-_is_fatal_cuda_error = _is_fatal_device_error
+    """Warm fn up before cudagraph capture.
-
-
-def _warmup_before_capture(di, fn, warmup_ms, grad_to_none):
-    """Warm fn up before graph capture.
 
     do_bench_cudagraph only warms up 5 iterations internally — too few for the
     MoE kernels, where the first calls trigger Triton autotune config
@@ -104,6 +68,8 @@ def _warmup_before_capture(di, fn, warmup_ms, grad_to_none):
     """
     import time
 
+    import torch
+
     def _run():
         if grad_to_none is not None:
             for x in grad_to_none:
@@ -112,23 +78,23 @@ def _warmup_before_capture(di, fn, warmup_ms, grad_to_none):
 
     # First call absorbs autotune/JIT compilation (can take seconds); discard it.
     _run()
-    di.synchronize()
+    torch.cuda.synchronize()
     # Size the warmup loop from one timed steady-state run. dt_ms includes
     # launch overhead, so n errs low for tiny kernels; warmup needn't be exact.
     t0 = time.perf_counter()
     _run()
-    di.synchronize()
+    torch.cuda.synchronize()
     dt_ms = (time.perf_counter() - t0) * 1000.0
     budget = warmup_ms if (warmup_ms and warmup_ms > 0) else 25.0
     n = int(budget / dt_ms) if dt_ms > 0 else 25
     n = max(5, min(n, 200))
     for _ in range(n):
         _run()
-    di.synchronize()
+    torch.cuda.synchronize()
 
 
-def _capture_with_scratch_rng(di, call):
+def _capture_with_scratch_rng(call):
-    """Run a graph capture under a throwaway device RNG state.
+    """Run a graph capture under a throwaway CUDA RNG state.
 
     torch.cuda.graph's capture_begin unconditionally registers the default CUDA
     RNG generator and sets its capturing flag; when a capture aborts midway the
@@ -140,21 +106,12 @@ def _capture_with_scratch_rng(di, call):
     before capturing and restore the original state object afterwards, success
     or failure. The real generator never takes part in a capture, so poisoning
     can only land on the discarded scratch state, and _seed_plugin determinism
-    is left untouched.
+    is left untouched."""
-
-    An interface that does not expose default_generators has no such registration
-    to protect against (torch_musa is one), so the capture simply runs.
-    """
     import torch
-
+    dev = torch.cuda.current_device()
-    generators = getattr(di, "default_generators", None)
+    gen = torch.cuda.default_generators[dev]
-    if not generators:
-        return call()
-    gen = generators[di.current_device()]
-    if not hasattr(gen, "graphsafe_get_state"):
-        return call()
     saved = gen.graphsafe_get_state()
-    scratch = torch.Generator(device=torch.device(gen.device))
+    scratch = torch.Generator(device="cuda")
     scratch.manual_seed(gen.initial_seed())
     gen.graphsafe_set_state(scratch)
     try:
@@ -163,85 +120,24 @@ def _capture_with_scratch_rng(di, call):
         gen.graphsafe_set_state(saved)
 
 
-def _do_bench_graph(di, fn, rep, grad_to_none, quantiles, return_mode, n_retries=10):
-    """do_bench_cudagraph, taken from the device interface instead of torch.cuda.
-
-    A step-for-step port of upstream's implementation: estimate one iteration,
-    unroll that many calls into a captured graph so host launch overhead drops
-    out, then time independent replays. Upstream is used verbatim on CUDA; this
-    exists because upstream names torch.cuda directly and so cannot run on any
-    other backend.
-
-    The replay timings are handed to _summarize_statistics as a plain list of
-    floats, which is what both do_bench and upstream's do_bench_cudagraph pass:
-    this fork's helper runs Python's statistics module over them, so wrapping
-    them in a tensor makes "mean" raise and makes a quantile request return
-    tensors that later fail to serialize.
-    """
-    with di.stream(di.Stream()):
-        fn()
-        if grad_to_none is not None:
-            for x in grad_to_none:
-                x.detach_()
-                x.requires_grad_(True)
-                x.grad = None
-        start_event = di.Event(enable_timing=True)
-        end_event = di.Event(enable_timing=True)
-        start_event.record()
-        for _ in range(5):
-            fn()
-        end_event.record()
-        di.synchronize()
-        estimate_ms = start_event.elapsed_time(end_event) / 5
-        n_repeat = 1000 if estimate_ms == 0 else max(1, int(rep / estimate_ms))
-        graph = _graph_class(di)()
-        with di.graph(graph):
-            for _ in range(n_repeat):
-                if grad_to_none is not None:
-                    for x in grad_to_none:
-                        x.grad = None
-                fn()
-        di.synchronize()
-        ret = []
-        for _ in range(n_retries):
-            start_event = di.Event(enable_timing=True)
-            end_event = di.Event(enable_timing=True)
-            start_event.record()
-            graph.replay()
-            end_event.record()
-            di.synchronize()
-            ret.append(start_event.elapsed_time(end_event) / n_repeat)
-        del graph
-        return _tt._summarize_statistics(ret, quantiles, return_mode)
-
-
 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 graph capture."""
+    """Drop-in replacement for triton.testing.do_bench using cudagraph capture."""
     op = os.environ.get("FLAGGEMS_PERF_CURRENT_OP", "<unknown>")
-    phase = "sync"  # which step rejected graph capture, for the fallback marker
+    phase = "sync"  # which step rejected cudagraph, for the fallback marker
     try:
-        di = _device_interface()
+        # do_bench_cudagraph runs fn on a fresh side stream without syncing
-        # The graph path runs fn on a fresh side stream without syncing with the
+        # with the default stream first; input tensors produced just before
-        # default stream first; input tensors produced just before (quantize/topk
+        # (quantize/topk in the input iter) may still be in flight, and racing
-        # in the input iter) may still be in flight, and racing on them
+        # on them intermittently kills the context (illegal instruction from a
-        # intermittently kills the context (illegal instruction from a device
+        # device trap on garbage indices). Sync before switching streams.
-        # trap on garbage indices). Sync before switching streams.
+        import torch
-        di.synchronize()
+        torch.cuda.synchronize()
         phase = "warmup"
         # Settle autotune/JIT before capture (rationale: _warmup_before_capture).
-        _warmup_before_capture(di, fn, warmup, grad_to_none)
+        _warmup_before_capture(fn, warmup, grad_to_none)
         phase = "capture"
-        if _is_cuda(di):
+        return _capture_with_scratch_rng(lambda: _tt.do_bench_cudagraph(
-            return _capture_with_scratch_rng(di, lambda: _tt.do_bench_cudagraph(
-                fn,
-                rep=rep,
-                grad_to_none=grad_to_none,
-                quantiles=quantiles,
-                return_mode=return_mode,
-            ))
-        return _capture_with_scratch_rng(di, lambda: _do_bench_graph(
-            di,
             fn,
             rep=rep,
             grad_to_none=grad_to_none,
@@ -251,7 +147,7 @@ def _patched_do_bench(fn, warmup=25, rep=100, grad_to_none=None,
     except Exception as exc:
         # Re-raise genuine device failures; anything else falls back to plain
         # do_bench, with the failing phase (sync/warmup/capture) in the marker.
-        if _is_fatal_device_error(exc):
+        if _is_fatal_cuda_error(exc):
             raise
         reason = type(exc).__name__
         detail = " ".join(str(exc).split())[:160]
@@ -273,11 +169,17 @@ def _patched_do_bench(fn, warmup=25, rep=100, grad_to_none=None,
 
 def pytest_configure(config):
     from _term_style import tag
-    _tt.do_bench = _patched_do_bench
     try:
-        backend = getattr(_device_interface(), "__name__", "?")
+        from backends import get_backend
-    except Exception:  # pragma: no cover - reported, never fatal at configure
+        if not get_backend().supports_cuda_graph:
-        backend = "?"
+            print(f"{tag('[cudagraph-plugin]')} skipped: backend "
-    print(f"{tag('[cudagraph-plugin]')} triton.testing.do_bench → graph capture "
+                  f"{get_backend().name!r} does not support CUDA Graph",
-          f"via {backend}",
+                  file=sys.stderr, flush=True)
+            return
+    except Exception as exc:
+        print(f"{tag('[cudagraph-plugin]')} backend probe failed ({exc}); "
+              "not patching do_bench", file=sys.stderr, flush=True)
+        return
+    _tt.do_bench = _patched_do_bench
+    print(f"{tag('[cudagraph-plugin]')} triton.testing.do_bench → do_bench_cudagraph",
           file=sys.stderr, flush=True)

_device_guard_plugin.py

@@ -1,45 +1,7 @@
-"""Guard against FlagGems' device-detection hang (benchmark-side fix).
+"""Backward-compatible alias of `_device_setup_plugin`.
 
-FlagGems' `runtime.backend.device_finder` falls back to
+Kept so an old `-p _device_guard_plugin` line still sets GEMS_VENDOR. New
-`subprocess.run(nvidia-smi)` **without a timeout**. Under this conda env's
+plugin lists should load `_device_setup_plugin` instead.
-forked/inconsistent subprocess (`_posixsubprocess` symbol mismatch), that
-child can hang indefinitely, leaving `import flag_gems` stuck in `wait4`
-(seen as run_pytest hanging with no output, GPU at 0%).
-
-This plugin runs at import time — before any test module imports flag_gems —
-detects the vendor via torch (no subprocess), and sets GEMS_VENDOR so
-device_finder takes its env fast-path (`_get_vendor_from_env`) and never
-reaches the hang-prone subprocess probe. No change to the FlagGems repo.
-
-Side effect: get_device_properties initializes the CUDA context very early in
-the pytest process. Fine for the current single-process runs; revisit if
-fork-based parallelism (e.g. pytest-xdist) is ever introduced.
 """
-import os
+from _device_setup_plugin import *  # noqa: F401,F403
-import sys
+from _device_setup_plugin import pytest_configure  # noqa: F401
-
-_VENDOR_ENV_KEYS = ("GEMS_VENDOR", "FLAGGEMS_VENDOR", "GEMS_BACKEND", "FLAGGEMS_BACKEND")
-
-
-def _guard_device_vendor():
-    # Respect an explicit choice if the user already set one.
-    if any(k in os.environ for k in _VENDOR_ENV_KEYS):
-        return
-    try:
-        import torch
-
-        if torch.cuda.is_available():
-            name = torch.cuda.get_device_properties(0).name.upper()
-            if "NVIDIA" in name:
-                os.environ["GEMS_VENDOR"] = "nvidia"
-                from _term_style import tag
-                print(f"{tag('[device-guard-plugin]')} set GEMS_VENDOR=nvidia "
-                      "(skip flag_gems nvidia-smi subprocess probe, avoids import hang)",
-                      file=sys.stderr, flush=True)
-    except Exception as e:  # pragma: no cover
-        print(f"[device-guard-plugin] torch vendor probe failed ({e}); "
-              "leaving detection to flag_gems",
-              file=sys.stderr, flush=True)
-
-
-_guard_device_vendor()

_device_setup_plugin.py

@@ -0,0 +1,69 @@
+"""Set FlagGems vendor + record backend identity before flag_gems is imported.
+
+Replaces the NVIDIA-only `_device_guard_plugin`. Setting GEMS_VENDOR from the
+zl_bench backend name skips FlagGems' hang-prone nvidia-smi subprocess probe
+on CUDA, and maps musa/rocm/ascend onto the matching FlagGems vendor.
+
+Import-time side effect is intentional: FlagGems reads GEMS_VENDOR while the
+benchmark module is being collected. Device name / arch are filled later in
+pytest_configure so import does not initialize a device context.
+"""
+
+from __future__ import annotations
+
+import os
+import sys
+
+from backends import get_backend
+from providers.flaggems import apply_vendor_env
+from _run_identity import update_identity
+
+
+def _setup_at_import() -> None:
+    try:
+        backend = get_backend()
+    except Exception as exc:
+        print(f"[device-setup-plugin] cannot resolve backend ({exc}); "
+              "leaving GEMS_VENDOR to flag_gems",
+              file=sys.stderr, flush=True)
+        return
+    vendor = apply_vendor_env(backend)
+    os.environ["ZL_BACKEND"] = backend.name
+    device = os.environ.get("ZL_DEVICE", "").strip()
+    if not device:
+        for key in backend.visibility_env_keys:
+            raw = os.environ.get(key, "").strip()
+            if raw:
+                device = raw.split(",")[0].strip()
+                break
+    update_identity(
+        backend=backend.name,
+        device=device,
+        flaggems_vendor=vendor or "",
+        visibility_env=backend.visibility_env_keys[0] if backend.visibility_env_keys else "",
+    )
+    if vendor:
+        from _term_style import tag
+        print(f"{tag('[device-setup-plugin]')} backend={backend.name} "
+              f"GEMS_VENDOR={vendor}",
+              file=sys.stderr, flush=True)
+
+
+def pytest_configure(config):
+    try:
+        backend = get_backend()
+        ident = backend.device_identity()
+        update_identity(
+            backend=backend.name,
+            device_name=ident.get("device_name", ""),
+            device_arch=ident.get("device_arch", ""),
+            timer_api=ident.get("timer_api", ""),
+            flaggems_vendor=ident.get("flaggems_vendor", ""),
+        )
+    except Exception as exc:
+        from _term_style import tag
+        print(f"{tag('[device-setup-plugin]')} identity probe failed ({exc})",
+              file=sys.stderr, flush=True)
+
+
+_setup_at_import()

_ir_meta_plugin.py

@@ -402,17 +402,33 @@ def _dump_ttgir(cache_dir: str, dump_dir: str) -> None:
 
 def pytest_configure(config):
     global _dump_enabled
+    from _term_style import tag
+    try:
         import triton
         import triton.runtime.autotuner as _autotuner
+    except Exception as exc:
+        print(f"{tag('[ir-meta-plugin]')} Triton knobs unavailable ({exc}); "
+              "IR collection disabled", file=sys.stderr, flush=True)
+        return
+    if not hasattr(triton, "knobs"):
+        print(f"{tag('[ir-meta-plugin]')} triton.knobs missing; "
+              "IR collection disabled", file=sys.stderr, flush=True)
+        return
 
-    prev = triton.knobs.compilation.listener
+    try:
+        compilation = triton.knobs.compilation
+        prev = compilation.listener
+    except Exception as exc:
+        print(f"{tag('[ir-meta-plugin]')} compilation listener unavailable "
+              f"({exc}); IR collection disabled", file=sys.stderr, flush=True)
+        return
     if prev is None:
-        triton.knobs.compilation.listener = _compile_listener
+        compilation.listener = _compile_listener
     else:
         def chained(**kwargs):
             prev(**kwargs)
             _compile_listener(**kwargs)
-        triton.knobs.compilation.listener = chained
+        compilation.listener = chained
 
     dump_dir = os.environ.get(_DUMP_DIR_ENV, "").strip()
     if dump_dir:

_mm_cluster_fix_plugin.py

@@ -105,6 +105,18 @@ def _warmup_load_backend():
 
 def pytest_configure(config):
     from _term_style import tag
+    try:
+        from backends import get_backend
+        if not get_backend().supports_mm_cluster_fix:
+            print(f"{tag('[mm-cluster-fix-plugin]')} skipped: backend "
+                  f"{get_backend().name!r} is not the CUDA Hopper path",
+                  file=sys.stderr, flush=True)
+            return
+    except Exception as exc:
+        print(f"{tag('[mm-cluster-fix-plugin]')} backend probe failed ({exc}); "
+              "leaving cluster_remote_mm_scenario unpatched",
+              file=sys.stderr, flush=True)
+        return
     _patch_loaded()
     _warmup_load_backend()
     _patch_loaded()

_musa_benchmark_compat_plugin.py

@@ -1,89 +0,0 @@
-"""Narrow compatibility fixes for the frozen MUSA product corpus.
-
-Several upstream FlagGems tests gate themselves on ``torch.cuda.is_available()``
-even though the code under test runs on any accelerator FlagGems dispatches to.
-On MThreads those guards skip the test instead of exercising the operator, which
-turns a working path into a verdict of "not measured".
-
-Every override here is pinned to one file and one exact skip reason, applies
-only when FlagGems has dispatched to the MThreads vendor, and is only added
-after the unskipped test was observed to pass on the device.
-
-Two files were probed on S5000 and deliberately left out, because their guards
-are not stale:
-
-- ``test_cutlass_scaled_mm.py`` gates on vLLM as well as CUDA, so removing the
-  guard would raise on the missing dependency instead of measuring anything.
-- ``test_grid_sample.py`` builds its inputs through a helper whose device
-  argument defaults to the literal ``"cuda"``. Its guard is consistent with its
-  own body: unskipped, all 90 collected cases fail in ``aten::empty`` on a CUDA
-  backend that does not exist here. Making that file portable means editing
-  upstream test source, which is outside what a compatibility plugin may do.
-"""
-
-from __future__ import annotations
-
-import sys
-from pathlib import Path
-
-
-# file name -> exact skipif reasons that are stale on a dispatching backend.
-_STALE_CUDA_GUARDS: dict[str, frozenset[str]] = {
-    "test_fft.py": frozenset({"CUDA is not available"}),
-    "test_top_k_per_row_decode.py": frozenset({"CUDA device required"}),
-}
-
-# Guards that are only stale for one test in the file stay pinned to it.
-_PINNED_TESTS: dict[str, str] = {
-    "test_fft.py": "test_perf_fft",
-}
-
-
-def _is_stale_cuda_guard(item, mark) -> bool:
-    if mark.name != "skipif":
-        return False
-    name = Path(str(item.path)).name
-    reasons = _STALE_CUDA_GUARDS.get(name)
-    if reasons is None or mark.kwargs.get("reason") not in reasons:
-        return False
-    pinned = _PINNED_TESTS.get(name)
-    return pinned is None or item.name == pinned
-
-
-# Retained under its historical name: the sealed campaign's tests reference it.
-def _is_stale_fft_cuda_guard(item, mark) -> bool:
-    return Path(str(item.path)).name == "test_fft.py" and _is_stale_cuda_guard(
-        item, mark
-    )
-
-
-def pytest_collection_modifyitems(items):
-    import flag_gems
-
-    if flag_gems.vendor_name != "mthreads":
-        return
-    removed: dict[str, int] = {}
-    # A guard can sit on the test, its class, or the module's pytestmark, and
-    # skipif is resolved by walking that whole chain. Strip every level.
-    seen: set[int] = set()
-    for item in items:
-        for node in item.listchain():
-            if id(node) in seen:
-                continue
-            seen.add(id(node))
-            markers = getattr(node, "own_markers", None)
-            if not markers:
-                continue
-            kept = [mark for mark in markers if not _is_stale_cuda_guard(item, mark)]
-            dropped = len(markers) - len(kept)
-            if dropped:
-                markers[:] = kept
-                name = Path(str(item.path)).name
-                removed[name] = removed.get(name, 0) + dropped
-    for name, count in sorted(removed.items()):
-        print(
-            f"[musa-benchmark-compat] removed {count} stale CUDA-only skip "
-            f"guard(s) from {name}",
-            file=sys.stderr,
-            flush=True,
-        )

_parallel_warmup_plugin.py

@@ -33,7 +33,7 @@ import sys
 import tempfile
 import time
 from pathlib import Path
-from typing import Any, Dict, List, Tuple
+from typing import Any, Dict, List, Optional, Tuple
 
 import pytest
 
@@ -60,8 +60,8 @@ def _requested_gpus() -> int:
     if want < 2:
         return 0
     try:
-        import torch
+        from backends import get_backend
-        have = torch.cuda.device_count()
+        have = get_backend().probe_device_count()
     except Exception:
         return 0
     return max(0, min(want, have))
@@ -70,30 +70,132 @@ def _requested_gpus() -> int:
 def _visible_devices() -> List[str]:
     """The device ids a shard may be pinned to, in parent-visible order.
 
-    Must respect an inherited CUDA_VISIBLE_DEVICES: writing a bare shard index
+    Must respect inherited visibility env for the *active* backend. Writing a
-    into the child would re-interpret it as an absolute id, so a caller who
+    bare shard index into the child would re-interpret it as an absolute id,
-    picked idle cards (CUDA_VISIBLE_DEVICES=5,6) would silently get cards 0,1 —
+    so a caller who picked idle cards (CUDA_VISIBLE_DEVICES=5,6) would silently
-    exactly the busy-GPU case the caller was avoiding.
+    get cards 0,1 — exactly the busy-GPU case the caller was avoiding.
     """
-    raw = os.environ.get("CUDA_VISIBLE_DEVICES", "").strip()
+    from backends import get_backend
+    backend = get_backend()
+    for key in backend.visibility_env_keys:
+        raw = os.environ.get(key, "").strip()
         if raw:
             return [d.strip() for d in raw.split(",") if d.strip()]
     try:
-        import torch
+        n = backend.probe_device_count()
-        return [str(i) for i in range(torch.cuda.device_count())]
+        return [str(i) for i in range(n)] if n else []
     except Exception:
         return []
 
 
-def _read_shape_yaml(path: str, op: str) -> Tuple[List[Any], Dict[str, Any]]:
+def _attr_name(node: Any) -> str:
-    """Shapes for this op, plus the sibling keys to carry into each shard."""
+    """Last component of a Name/Attribute node: `base.BlasBenchmark` -> `BlasBenchmark`."""
+    import ast
+    if isinstance(node, ast.Name):
+        return node.id
+    if isinstance(node, ast.Attribute):
+        return node.attr
+    return ""
+
+
+def _class_bases(py_path: Path) -> Dict[str, List[str]]:
+    """{class name: [base class names]} parsed statically from one module."""
+    import ast
+    try:
+        tree = ast.parse(py_path.read_text(errors="replace"))
+    except Exception:
+        return {}
+    return {
+        node.name: [n for n in (_attr_name(b) for b in node.bases) if n]
+        for node in ast.walk(tree)
+        if isinstance(node, ast.ClassDef)
+    }
+
+
+def _class_chain(test_py: Path, op: str) -> List[str]:
+    """Benchmark class names to try for `op`, nearest ancestor first.
+
+    FlagGems' `Benchmark.set_shapes` looks up `op_name` and, failing that, walks
+    `type(self).__mro__` trying each class name as a yaml key -- which is how
+    BLAS ops reach the shared `BlasBenchmark` entry in core_shapes.yaml. This
+    plugin runs at pytest_configure, before any Benchmark exists, so the real
+    MRO is unavailable; the class is recovered statically instead, from the
+    `op_name=` keyword of the constructor call inside the test module, then its
+    bases from that module plus benchmark/base.py.
+
+    Breadth-first over the inheritance graph, not a true C3 linearisation: the
+    benchmark classes are single-inheritance in practice, and a miss only costs
+    the sharding (see the caller's skip path), never correctness.
+    """
+    import ast
+    try:
+        tree = ast.parse(test_py.read_text(errors="replace"))
+    except Exception:
+        return []
+    start = ""
+    for node in ast.walk(tree):
+        if not isinstance(node, ast.Call):
+            continue
+        for kw in node.keywords:
+            if (kw.arg == "op_name" and isinstance(kw.value, ast.Constant)
+                    and kw.value.value == op):
+                start = _attr_name(node.func)
+                break
+        if start:
+            break
+    if not start:
+        return []
+
+    graph = _class_bases(test_py)
+    for name, bases in _class_bases(test_py.parent / "base.py").items():
+        graph.setdefault(name, bases)
+
+    chain: List[str] = []
+    seen: set = set()
+    queue = [start]
+    while queue:
+        name = queue.pop(0)
+        if not name or name in seen:
+            continue
+        seen.add(name)
+        chain.append(name)
+        queue.extend(graph.get(name, []))
+    return chain
+
+
+def _test_module_path(argv: List[str]) -> Optional[Path]:
+    """The benchmark module under test, from pytest's own argv."""
+    for a in argv:
+        if a.startswith("-"):
+            continue
+        p = Path(a.split("::", 1)[0])
+        if p.suffix == ".py" and p.is_file():
+            return p
+    return None
+
+
+def _read_shape_yaml(path: str, op: str, test_py: Optional[Path] = None
+                     ) -> Tuple[List[Any], Dict[str, Any], str]:
+    """Shapes for this op, the sibling keys to carry into each shard, and the
+    yaml key they came from ("" when nothing matched).
+
+    Tries `op` first, then the benchmark class chain, mirroring FlagGems. Without
+    the class fallback every BLAS op (mm/addmm/bmm/baddbmm, whose shapes live
+    under `BlasBenchmark`) resolved to zero shapes and sharding silently turned
+    itself off.
+    """
     import yaml
     with open(path) as f:
         doc = yaml.safe_load(f) or {}
-    entry = doc.get(op) or {}
+    keys = [op]
-    shapes = entry.get("shapes") or []
+    if not isinstance(doc.get(op), dict) and test_py is not None:
+        keys += _class_chain(test_py, op)
+    for key in keys:
+        entry = doc.get(key)
+        if isinstance(entry, dict) and entry.get("shapes"):
             extra = {k: v for k, v in entry.items() if k != "shapes"}
-    return list(shapes), extra
+            return list(entry["shapes"]), extra, key
+    return [], {}, ""
 
 
 def _shard(shapes: List[Any], n: int) -> List[List[Any]]:
@@ -111,8 +213,8 @@ def _spawn(shards: List[List[Any]], op: str, extra: Dict[str, Any],
            ) -> Tuple[List[Path], List[int]]:
     """One subprocess per shard, each pinned to its own GPU.
 
-    Subprocesses rather than pytest-xdist: _device_guard_plugin initializes a
+    Subprocesses rather than pytest-xdist: `_device_setup_plugin` may initialize
-    CUDA context at import time, which does not survive fork-based parallelism.
+    a device context at import time, which does not survive fork-based parallelism.
     """
     import yaml
     procs: List[Tuple[int, subprocess.Popen, Any]] = []
@@ -143,7 +245,8 @@ def _spawn(shards: List[List[Any]], op: str, extra: Dict[str, Any],
         child_argv += ["--shape_file", str(w / "shapes.yaml")]
 
         env = dict(os.environ)
-        env["CUDA_VISIBLE_DEVICES"] = devices[i]
+        from backends import get_backend
+        env = get_backend().apply_visibility(env, devices[i])
         env[_WORKER_ENV] = "1"                    # stops recursion
         env[_RECORD_DIR_ENV] = str(rec)           # shard records its own picks
         env.pop(_REPLAY_DIR_ENV, None)            # a shard must sweep
@@ -222,6 +325,23 @@ def pytest_configure(config):
         return
 
     op = _op_name()
+
+    # Only --level core is shardable. At comprehensive, FlagGems' set_shapes()
+    # merges set_more_shapes() on top of whatever the yaml holds, and that
+    # addition is arbitrary per-subclass Python (hardcoded tuples plus
+    # models_shapes/*.yaml) which this plugin cannot see. Sharding the yaml list
+    # alone would sweep a strict subset -- measured on mm: 5 of 29 shapes, 90
+    # AUTOTUNE_REPLAY_FALLBACK reason=key_missing -- and every uncovered key
+    # then self-tunes while the run reports itself as replaying. Refusing costs
+    # wall time; sharding a subset costs the same-config premise.
+    level = str(getattr(config.option, "level", "") or "").strip().lower()
+    if level and level != "core":
+        print(f"{tag('[parallel-warmup-plugin]')} --level {level} merges shapes "
+              "this plugin cannot enumerate (set_more_shapes); sharding a subset "
+              "would break the same-config premise, so sweeping serially instead",
+              flush=True)
+        return
+
     shape_file = getattr(config.option, "shape_file", "") or ""
     if not shape_file or not Path(shape_file).is_file():
         print(f"{tag('[parallel-warmup-plugin]')} no --shape_file; "
@@ -232,15 +352,24 @@ def pytest_configure(config):
               "already pinned, nothing to sweep, skipping", flush=True)
         return
 
+    test_py = _test_module_path(list(config.invocation_params.args))
     try:
-        shapes, extra = _read_shape_yaml(shape_file, op)
+        shapes, extra, key = _read_shape_yaml(shape_file, op, test_py)
     except Exception as exc:
         print(f"{tag('[parallel-warmup-plugin]')} cannot read shapes "
               f"({type(exc).__name__}: {exc}); skipping", flush=True)
         return
+    if not key:
+        # Distinct from "found a short list": nothing in the yaml describes this
+        # op at all, so the shard files would be built from an empty set.
+        print(f"{tag('[parallel-warmup-plugin]')} no yaml entry for {op!r} "
+              f"(tried: {', '.join([op] + _class_chain(test_py, op)) if test_py else op}); "
+              "skipping", flush=True)
+        return
     if len(shapes) < 2:
-        print(f"{tag('[parallel-warmup-plugin]')} only {len(shapes)} shape(s); "
+        print(f"{tag('[parallel-warmup-plugin]')} only {len(shapes)} shape(s) "
-              "sharding would not pay off, skipping", flush=True)
+              f"under yaml key {key!r}; sharding would not pay off, skipping",
+              flush=True)
         return
 
     devices = _visible_devices()
@@ -251,8 +380,9 @@ def pytest_configure(config):
 
     shards = _shard(shapes, min(gpus, len(shapes), len(devices)))
     scratch = Path(tempfile.mkdtemp(prefix=f"warmup_{op}_"))
-    print(f"{tag('[parallel-warmup-plugin]')} sweeping {len(shapes)} shapes on "
+    print(f"{tag('[parallel-warmup-plugin]')} sweeping {len(shapes)} shapes "
-          f"{len(shards)} GPUs (shards: {[len(s) for s in shards]}); this phase's "
+          f"(yaml key {key!r}) on {len(shards)} GPUs "
+          f"(shards: {[len(s) for s in shards]}); this phase's "
           "latency is discarded, only configs are kept", flush=True)
 
     keep_scratch = False
@@ -260,9 +390,23 @@ def pytest_configure(config):
     try:
         recs, bad = _spawn(shards, op, extra, scratch,
                            list(config.invocation_params.args), devices)
-        merged, total, conflicts = _merge(recs, op)
         elapsed = time.time() - t0
 
+        if bad:
+            # A dead shard means its shapes were never swept, so publishing the
+            # merge would hand the measurement a config set with holes; every
+            # hole then silently self-tunes under a run that claims to replay.
+            # Discard the whole sweep rather than publish a partial one.
+            keep_scratch = True
+            print(f"{tag('[parallel-warmup-plugin]')} shard(s) {bad} exited "
+                  f"non-zero after {elapsed:.0f}s; their shapes were never swept, "
+                  "so the merged configs would have holes. Discarding this sweep "
+                  f"and falling back to normal serial autotune (worker logs kept "
+                  f"in {scratch})", flush=True)
+            return
+
+        merged, total, conflicts = _merge(recs, op)
+
         if not total:
             keep_scratch = True
             print(f"{tag('[parallel-warmup-plugin]')} no configs recovered in "
@@ -287,9 +431,6 @@ def pytest_configure(config):
         msg = (f"{tag('[parallel-warmup-plugin]')} sweep done in {elapsed:.0f}s: "
                f"{len(merged)} kernels / {total} configs -> {out / f'{op}.json'}; "
                "measurement continues serially on one GPU")
-        if bad:
-            msg += f" (WARNING shard(s) {bad} exited non-zero; missing keys will "
-            msg += "fall back to live autotune)"
         print(msg, flush=True)
         if conflicts:
             print(f"{tag('[parallel-warmup-plugin]')} WARNING {conflicts} config "

_run_identity.py

@@ -0,0 +1,91 @@
+"""Merge-update `$OUT_DIR/run_identity.json` so plugins can record backend facts.
+
+Shell writes a stub before pytest; plugins add device_name / timer / vendor
+without clobbering earlier keys. Missing OUT_DIR is a silent no-op so unit
+tests and ad-hoc pytest invocations do not fail.
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import os
+from pathlib import Path
+from typing import Any, Dict, Optional
+
+
+IDENTITY_NAME = "run_identity.json"
+
+
+def identity_path(out_dir: Optional[str] = None) -> Optional[Path]:
+    root = out_dir or os.environ.get("OUT_DIR", "").strip()
+    if not root:
+        return None
+    return Path(root) / IDENTITY_NAME
+
+
+def load_identity(out_dir: Optional[str] = None) -> Dict[str, Any]:
+    path = identity_path(out_dir)
+    if path is None or not path.is_file():
+        return {}
+    try:
+        data = json.loads(path.read_text(encoding="utf-8"))
+    except Exception:
+        return {}
+    return data if isinstance(data, dict) else {}
+
+
+def update_identity(out_dir: Optional[str] = None, **fields: Any) -> Dict[str, Any]:
+    path = identity_path(out_dir)
+    if path is None:
+        return dict(fields)
+    path.parent.mkdir(parents=True, exist_ok=True)
+    current = load_identity(out_dir)
+    for key, value in fields.items():
+        if value is None:
+            continue
+        current[key] = value
+    tmp = path.with_name(path.name + ".tmp")
+    tmp.write_text(json.dumps(current, indent=2, ensure_ascii=False) + "\n",
+                   encoding="utf-8")
+    os.replace(tmp, path)
+    return current
+
+
+def init_identity(out_dir: str, **fields: Any) -> Dict[str, Any]:
+    payload = {"schema_version": 1, "provider": "flaggems"}
+    payload.update(fields)
+    os.environ["OUT_DIR"] = str(out_dir)
+    return update_identity(out_dir, **payload)
+
+
+def main(argv=None) -> int:
+    p = argparse.ArgumentParser()
+    p.add_argument("out_dir")
+    p.add_argument("--backend", default=os.environ.get("ZL_BACKEND", "cuda"))
+    p.add_argument("--op", default=os.environ.get("OP", ""))
+    p.add_argument("--op-file", default=os.environ.get("OP_FILE", ""))
+    p.add_argument("--flaggems-dir", default=os.environ.get("FLAGGEMS_DIR", ""))
+    p.add_argument("--autotune-mode", default="")
+    p.add_argument("--device", default=os.environ.get("ZL_DEVICE", ""))
+    p.add_argument("--test-file", default=os.environ.get("TEST_FILE", ""))
+    p.add_argument("--shape-mode", default=os.environ.get("SHAPE_MODE", ""))
+    p.add_argument("--bench-level", default=os.environ.get("BENCH_LEVEL", "core"))
+    args = p.parse_args(argv)
+    init_identity(
+        args.out_dir,
+        backend=args.backend,
+        op=args.op,
+        op_file=args.op_file,
+        flaggems_dir=args.flaggems_dir,
+        autotune_mode=args.autotune_mode,
+        device=args.device,
+        test_file=args.test_file,
+        shape_mode=args.shape_mode,
+        bench_level=args.bench_level,
+    )
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

_timing_plugin.py

@@ -0,0 +1,47 @@
+"""Select the timing path from backend capabilities, before any shape runs.
+
+CUDA keeps the existing `_cudagraph_plugin` behaviour (graph capture, with
+per-shape fallback to plain do_bench). Other backends do not attempt CUDA
+Graph at all: they leave `triton.testing.do_bench` in place and record
+`timer=direct`. Trying capture-then-fallback on a non-CUDA device would
+produce FATAL_DEVICE noise and mixed-caliber rows.
+"""
+
+from __future__ import annotations
+
+import os
+import sys
+
+from backends import get_backend
+from _run_identity import update_identity
+
+
+def _graph_requested(backend) -> bool:
+    raw = os.environ.get("ZL_ENABLE_CUDAGRAPH", "auto").strip().lower()
+    if raw in ("0", "false", "no", "off", "never", "direct"):
+        return False
+    if raw in ("1", "true", "yes", "on", "always"):
+        return bool(backend.supports_cuda_graph)
+    return bool(backend.supports_cuda_graph)
+
+
+def pytest_configure(config):
+    from _term_style import tag
+    backend = get_backend()
+    raw = os.environ.get("ZL_ENABLE_CUDAGRAPH", "auto").strip().lower()
+    force_on = raw in ("1", "true", "yes", "on", "always")
+    if force_on and not backend.supports_cuda_graph:
+        print(f"{tag('[timing-plugin]')} ZL_ENABLE_CUDAGRAPH=1 ignored on "
+              f"backend={backend.name} (no CUDA Graph); using timer=direct",
+              file=sys.stderr, flush=True)
+    if _graph_requested(backend):
+        import _cudagraph_plugin as _cg
+        _cg.pytest_configure(config)
+        update_identity(timer="cuda_graph", graph_capture="attempted")
+        print(f"{tag('[timing-plugin]')} timer=cuda_graph backend={backend.name}",
+              file=sys.stderr, flush=True)
+        return
+    update_identity(timer="direct", graph_capture="disabled")
+    print(f"{tag('[timing-plugin]')} timer=direct backend={backend.name} "
+          "(upstream do_bench; graph capture not selected)",
+          file=sys.stderr, flush=True)

ab_rlc_test.sh

@@ -1,68 +0,0 @@
-#!/bin/bash
-# A/B 对比：FlagTree RemoveLayoutConversions 增强关 vs 开，逐算子两轮。
-#
-# 口径：每个算子先跑关（record autotune config），再用 REPLAY_FROM 回放同一套
-# config 跑开，两侧只差 FLAGTREE_<VENDOR>_RLC_ENHANCE 这一个开关。产物沿用
-# run_pytest.sh 的目录结构，所以每个算子恰好留下两份 run.log：
-#
-#   runs/<op>_<时间戳>/run.log   关（before）
-#   runs/<op>_<时间戳>/run.log   开（after）
-#
-# 跑完用 batch_summary.py --ab <before> <after> 得到该算子的 speedup 区间。
-#
-# 用法：
-#   OPS="softmax gelu_backward" bash ab_rlc_test.sh
-#   OPS="sort:perf_sort topk" bash ab_rlc_test.sh     # 冒号分隔 OP:OP_FILE
-#
-# 注意：RLC 开关名带厂商前缀，用 VENDOR 指定（默认 MUSA）。套用前先确认开关有效：
-# grep 该开关名于 FlagTree 源码应有命中，否则两侧跑的是同一份代码。
-set -uo pipefail
-
-cd "$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
-
-VENDOR="${VENDOR:-MUSA}"
-ENHANCE_VAR="FLAGTREE_${VENDOR}_RLC_ENHANCE"
-
-# 空格分隔的算子表。元素为 OP 或 OP:OP_FILE（省略时 OP_FILE=OP）。
-OPS="${OPS:-softmax}"
-
-if [[ -t 1 && -z "${NO_COLOR:-}" ]]; then
-  export FLAGGEMS_PERF_COLOR=always
-fi
-
-LOG="runs/ab_rlc_$(date +%Y%m%d-%H%M%S).log"
-mkdir -p runs
-
-{
-  echo ">>> vendor=$VENDOR  switch=$ENHANCE_VAR  ops=$OPS"
-  for entry in $OPS; do
-    op="${entry%%:*}"
-    op_file="${entry#*:}"
-    [[ "$op_file" == "$entry" ]] && op_file="$op"
-
-    echo "########## $op | $ENHANCE_VAR=0 before (record) ##########"
-    if ! OP="$op" OP_FILE="$op_file" env "$ENHANCE_VAR=0" bash run_pytest.sh; then
-      echo "!!! FAILED before: $op"
-      continue
-    fi
-    BEFORE=$(ls -dt "runs/${op}_"* | head -1)
-
-    echo "########## $op | $ENHANCE_VAR=1 after (replay $BEFORE) ##########"
-    if ! OP="$op" OP_FILE="$op_file" REPLAY_FROM="$BEFORE" env "$ENHANCE_VAR=1" bash run_pytest.sh; then
-      echo "!!! FAILED after: $op"
-      continue
-    fi
-    AFTER=$(ls -dt "runs/${op}_"* | head -1)
-
-    echo "########## $op | summary ##########"
-    python batch_summary.py --ab "$BEFORE/run.log" "$AFTER/run.log" --label "$op" \
-      || echo "!!! summary failed: $op"
-  done
-  echo "########## ALL DONE ##########"
-} 2>&1 | tee "$LOG"
-
-# 去掉 ANSI 转义以便 grep/diff（终端输出保留颜色）
-sed -i -E $'s/\x1b\\[[0-9;]*[A-Za-z]//g' "$LOG"
-
-echo ">>> 总 log: $(pwd)/$LOG"
-echo ">>> 每算子两份 run.log 在 runs/<op>_<时间戳>/ 下，按时间先后即 before / after"

backends/__init__.py

@@ -0,0 +1,76 @@
+"""Registry of zl_bench device backends.
+
+Aliases map common vendor spellings onto the canonical backend name. Unknown
+names fail closed rather than defaulting to CUDA: a silent CUDA fallback would
+load the cudagraph plugin on the wrong device.
+"""
+
+from __future__ import annotations
+
+import os
+from typing import Dict, Optional
+
+from .ascend import AscendBackend
+from .base import DeviceBackend, parse_device_list
+from .cuda import CudaBackend
+from .musa import MusaBackend
+from .rocm import RocmBackend
+
+ALIASES: Dict[str, str] = {
+    "nvidia": "cuda",
+    "nv": "cuda",
+    "amd": "rocm",
+    "hip": "rocm",
+    "npu": "ascend",
+    "huawei": "ascend",
+}
+
+_INSTANCES: Dict[str, DeviceBackend] = {
+    "cuda": CudaBackend(),
+    "musa": MusaBackend(),
+    "rocm": RocmBackend(),
+    "ascend": AscendBackend(),
+}
+
+
+class UnknownBackendError(ValueError):
+    pass
+
+
+def known_backends() -> tuple[str, ...]:
+    return tuple(sorted(_INSTANCES))
+
+
+def normalize_backend_name(name: str) -> str:
+    key = (name or "").strip().lower()
+    if not key:
+        return "cuda"
+    return ALIASES.get(key, key)
+
+
+def get_backend(name: Optional[str] = None) -> DeviceBackend:
+    raw = name if name is not None else os.environ.get("ZL_BACKEND", "cuda")
+    key = normalize_backend_name(raw)
+    backend = _INSTANCES.get(key)
+    if backend is None:
+        raise UnknownBackendError(
+            f"unknown backend {raw!r}; known: {', '.join(known_backends())} "
+            f"(aliases: {', '.join(sorted(ALIASES))})"
+        )
+    return backend
+
+
+def enumerate_devices(backend: Optional[DeviceBackend] = None,
+                      explicit: Optional[str] = None) -> list[str]:
+    return (backend or get_backend()).enumerate_devices(explicit)
+
+
+__all__ = [
+    "DeviceBackend",
+    "UnknownBackendError",
+    "enumerate_devices",
+    "get_backend",
+    "known_backends",
+    "normalize_backend_name",
+    "parse_device_list",
+]

backends/ascend.py

@@ -0,0 +1,17 @@
+from __future__ import annotations
+
+from .base import DeviceBackend
+
+
+class AscendBackend(DeviceBackend):
+    name = "ascend"
+    visibility_env_keys = (
+        "ASCEND_RT_VISIBLE_DEVICES",
+        "ASCEND_VISIBLE_DEVICES",
+    )
+    torch_namespace = "npu"
+    flaggems_vendor = "ascend"
+    supports_cuda_graph = False
+    supports_mm_cluster_fix = False
+    supports_triton_hooks = True
+    list_command = None

backends/base.py

@@ -0,0 +1,148 @@
+"""Device backend contract for zl_bench screening runs.
+
+This is the orchestration-layer adapter: device enumeration, process visibility
+env, and capability flags that decide which pytest plugins are safe to load.
+It does not implement kernel timing receipts or compiler-artifact identity;
+those stay with the caller (FlagGems pytest for screening, a separate harness
+for qualification).
+"""
+
+from __future__ import annotations
+
+import os
+import subprocess
+from typing import Dict, Iterable, List, Optional, Sequence, Tuple
+
+# Every vendor visibility variable zl_bench knows about. apply_visibility()
+# clears the whole set before writing the active backend's keys so a parent
+# CUDA_VISIBLE_DEVICES cannot leak onto a MUSA/Ascend child.
+ALL_VISIBILITY_KEYS: Tuple[str, ...] = (
+    "CUDA_VISIBLE_DEVICES",
+    "MUSA_VISIBLE_DEVICES",
+    "HIP_VISIBLE_DEVICES",
+    "ROCR_VISIBLE_DEVICES",
+    "ASCEND_RT_VISIBLE_DEVICES",
+    "ASCEND_VISIBLE_DEVICES",
+)
+
+
+def parse_device_list(raw: Optional[str]) -> List[str]:
+    if not raw:
+        return []
+    return [tok.strip() for tok in raw.split(",") if tok.strip()]
+
+
+class DeviceBackend:
+    """One hardware runtime (CUDA, MUSA, ROCm, Ascend, ...).
+
+    Subclasses only fill class attributes and optional probe hooks. Callers
+    should not branch on ``backend.name`` for visibility or plugin policy;
+    they should read these attributes / methods.
+    """
+
+    name: str = ""
+    # First key is the canonical one written into run_identity.json.
+    visibility_env_keys: Tuple[str, ...] = ()
+    # torch.<namespace> used for optional device_count / identity probes.
+    torch_namespace: str = ""
+    # FlagGems GEMS_VENDOR value. Lives on the backend object because zl_bench
+    # still only drives FlagGems, but the mapping is owned by providers.flaggems
+    # and copied here so DeviceBackend.prepare_environment() can stay generic.
+    flaggems_vendor: str = ""
+    supports_cuda_graph: bool = False
+    supports_mm_cluster_fix: bool = False
+    supports_triton_hooks: bool = True
+    # Optional listing command; (argv, line_prefix). prefix="" means count
+    # non-empty lines. Must be a real vendor CLI; do not invent flags.
+    list_command: Optional[Tuple[Sequence[str], str]] = None
+
+    def visibility_env(self, device: str) -> Dict[str, str]:
+        return {key: device for key in self.visibility_env_keys}
+
+    def apply_visibility(self, env: Dict[str, str], device: str) -> Dict[str, str]:
+        out = dict(env)
+        for key in ALL_VISIBILITY_KEYS:
+            out.pop(key, None)
+        out.update(self.visibility_env(device))
+        return out
+
+    def enumerate_devices(self, explicit: Optional[str] = None) -> List[str]:
+        listed = parse_device_list(explicit)
+        if listed:
+            return listed
+        for key in self.visibility_env_keys:
+            listed = parse_device_list(os.environ.get(key, ""))
+            if listed:
+                return listed
+        n = self.probe_device_count()
+        if n > 0:
+            return [str(i) for i in range(n)]
+        return ["0"]
+
+    def probe_device_count(self) -> int:
+        n = self._probe_list_command()
+        if n > 0:
+            return n
+        return self._probe_torch_count()
+
+    def _probe_list_command(self) -> int:
+        if not self.list_command:
+            return 0
+        argv, prefix = self.list_command
+        try:
+            proc = subprocess.run(
+                list(argv), text=True, capture_output=True, timeout=30,
+            )
+        except Exception:
+            return 0
+        if proc.returncode != 0:
+            return 0
+        lines = [ln for ln in proc.stdout.splitlines() if ln.strip()]
+        if prefix:
+            lines = [ln for ln in lines if ln.startswith(prefix)]
+        return len(lines)
+
+    def _probe_torch_count(self) -> int:
+        if not self.torch_namespace:
+            return 0
+        try:
+            import torch
+            api = getattr(torch, self.torch_namespace, None)
+            if api is None or not hasattr(api, "device_count"):
+                return 0
+            return int(api.device_count())
+        except Exception:
+            return 0
+
+    def device_identity(self) -> Dict[str, str]:
+        """Best-effort identity for screening logs. Empty strings are allowed."""
+        info = {
+            "backend": self.name,
+            "device_name": "",
+            "device_arch": "",
+            "timer_api": f"torch.{self.torch_namespace}" if self.torch_namespace else "",
+            "flaggems_vendor": self.flaggems_vendor,
+        }
+        if not self.torch_namespace:
+            return info
+        try:
+            import torch
+            api = getattr(torch, self.torch_namespace, None)
+            if api is None:
+                return info
+            count = int(api.device_count()) if hasattr(api, "device_count") else 0
+            if count <= 0:
+                return info
+            index = int(api.current_device()) if hasattr(api, "current_device") else 0
+            if hasattr(api, "get_device_name"):
+                info["device_name"] = str(api.get_device_name(index))
+            if hasattr(api, "get_device_capability"):
+                cap = api.get_device_capability(index)
+                info["device_arch"] = f"{cap[0]}.{cap[1]}"
+            elif hasattr(api, "get_device_properties"):
+                props = api.get_device_properties(index)
+                arch = getattr(props, "name", "") or getattr(props, "arch", "")
+                info["device_arch"] = str(arch)
+        except Exception:
+            return info
+        return info

backends/cuda.py

@@ -0,0 +1,14 @@
+from __future__ import annotations
+
+from .base import DeviceBackend
+
+
+class CudaBackend(DeviceBackend):
+    name = "cuda"
+    visibility_env_keys = ("CUDA_VISIBLE_DEVICES",)
+    torch_namespace = "cuda"
+    flaggems_vendor = "nvidia"
+    supports_cuda_graph = True
+    supports_mm_cluster_fix = True
+    supports_triton_hooks = True
+    list_command = (("nvidia-smi", "--list-gpus"), "GPU ")

backends/musa.py

@@ -0,0 +1,16 @@
+from __future__ import annotations
+
+from .base import DeviceBackend
+
+
+class MusaBackend(DeviceBackend):
+    name = "musa"
+    visibility_env_keys = ("MUSA_VISIBLE_DEVICES",)
+    torch_namespace = "musa"
+    flaggems_vendor = "musa"
+    supports_cuda_graph = False
+    supports_mm_cluster_fix = False
+    supports_triton_hooks = True
+    # Device count comes from torch.musa or MUSA_VISIBLE_DEVICES. Do not guess
+    # a vendor CLI here; confirm the real listing tool before adding it.
+    list_command = None

backends/rocm.py

@@ -0,0 +1,16 @@
+from __future__ import annotations
+
+from .base import DeviceBackend
+
+
+class RocmBackend(DeviceBackend):
+    name = "rocm"
+    # HIP runtimes honor both; write both so a parent CUDA/HIP mix cannot
+    # leave the child looking at the wrong cards.
+    visibility_env_keys = ("HIP_VISIBLE_DEVICES", "ROCR_VISIBLE_DEVICES")
+    torch_namespace = "cuda"  # ROCm PyTorch still exposes torch.cuda
+    flaggems_vendor = "amd"
+    supports_cuda_graph = False
+    supports_mm_cluster_fix = False
+    supports_triton_hooks = True
+    list_command = None

batch_summary.py

@@ -39,12 +39,31 @@ DIRECT_NO_CUDAGRAPH_RE = re.compile(r"BENCHMARK_DIRECT_NO_CUDAGRAPH\b")
 REPLAY_FALLBACK_RE = re.compile(r"AUTOTUNE_REPLAY_FALLBACK\b")
 
 SUMMARY_FIELDS = [
-    "repo", "op", "fg_op_name", "dtype", "case_id", "row_status",
+    "repo", "op", "backend", "device", "timer",
+    "fg_op_name", "dtype", "case_id", "row_status",
     "torch_lat", "gems_lat", "speedup", "extra_metrics", "size_detail",
     "op_status", "notes",
 ]
 
 
+def _identity_from_status(status_row: Dict[str, str], op_dir: Path) -> Dict[str, str]:
+    """Backend/device/timer from ops_status.csv, falling back to run_identity.json."""
+    backend = (status_row.get("backend") or "").strip()
+    device = (status_row.get("device") or status_row.get("gpu") or "").strip()
+    timer = (status_row.get("timer") or "").strip()
+    ident_path = op_dir / "run_identity.json"
+    if ident_path.is_file() and not (backend and device and timer):
+        try:
+            from _run_identity import load_identity
+            ident = load_identity(str(op_dir))
+            backend = backend or str(ident.get("backend") or "")
+            device = device or str(ident.get("device") or ident.get("gpu") or "")
+            timer = timer or str(ident.get("timer") or "")
+        except Exception:
+            pass
+    return {"backend": backend, "device": device, "timer": timer}
+
+
 def parse_run_log(text: str) -> List[Dict[str, str]]:
     """把 run.log 里的全部结果表解析为行列表（跨多个 dtype 表）。"""
     rows: List[Dict[str, str]] = []
@@ -133,10 +152,11 @@ def summarize(batch_dir: Path, out_path: Optional[Path] = None) -> Path:
             text = log.read_text(errors="replace")
             status_row = status_by_dir.get(f"{repo}/{op_dir.name}", {})
             op = status_row.get("op") or op_dir.name
+            ident = _identity_from_status(status_row, op_dir)
             table_rows = parse_run_log(text)
             if not table_rows:
                 all_rows.append({
-                    "repo": repo, "op": op, "fg_op_name": "", "dtype": "",
+                    "repo": repo, "op": op, **ident, "fg_op_name": "", "dtype": "",
                     "case_id": "", "row_status": "", "torch_lat": "",
                     "gems_lat": "", "speedup": "", "extra_metrics": "",
                     "size_detail": "",
@@ -150,7 +170,8 @@ def summarize(batch_dir: Path, out_path: Optional[Path] = None) -> Path:
                 idx = case_counter.get(key, 0)
                 case_counter[key] = idx + 1
                 all_rows.append({
-                    "repo": repo, "op": op, "case_id": f"case_{idx:03d}",
+                    "repo": repo, "op": op, **ident,
+                    "case_id": f"case_{idx:03d}",
                     "op_status": status_row.get("status", ""),
                     "notes": status_row.get("notes", ""),
                     **row,
@@ -165,82 +186,12 @@ def summarize(batch_dir: Path, out_path: Optional[Path] = None) -> Path:
     return out_path
 
 
-def ab_rows(before_text: str, after_text: str) -> List[Dict[str, str]]:
-    """把两次 run.log 的测量行按 (op, dtype, size_detail) 配对。
-
-    speedup 列是 gems 对 torch，不是前后对比，所以这里用 gems_lat 自己算
-    before/after - 1。只有两侧都 SUCCESS 且延迟为正的行才计入；对不上的行
-    单独返回，避免用不完整的配对报一个看起来完整的区间。
-    """
-    def index(text: str) -> Dict[Tuple[str, str, str], float]:
-        out: Dict[Tuple[str, str, str], float] = {}
-        for row in parse_run_log(text):
-            if row.get("row_status") != "SUCCESS":
-                continue
-            try:
-                lat = float(row.get("gems_lat", ""))
-            except ValueError:
-                continue
-            if lat > 0:
-                out[(row.get("fg_op_name", ""), row.get("dtype", ""),
-                     row.get("size_detail", ""))] = lat
-        return out
-
-    before, after = index(before_text), index(after_text)
-    paired: List[Dict[str, str]] = []
-    for key in sorted(before.keys() & after.keys()):
-        op, dtype, size = key
-        b, a = before[key], after[key]
-        paired.append({
-            "fg_op_name": op, "dtype": dtype, "size_detail": size,
-            "before_lat": f"{b:.6f}", "after_lat": f"{a:.6f}",
-            "speedup_percent": f"{(b / a - 1.0) * 100.0:.2f}",
-        })
-    return paired
-
-
-def report_ab(before: Path, after: Path, label: str = "") -> List[Dict[str, str]]:
-    """打印一个算子的前后对比：逐行 speedup 与整体区间。"""
-    before_text = before.read_text(encoding="utf-8", errors="replace")
-    after_text = after.read_text(encoding="utf-8", errors="replace")
-    rows = ab_rows(before_text, after_text)
-    name = label or (rows[0]["fg_op_name"] if rows else "?")
-    if not rows:
-        print(f"[ab] {name}: 没有可配对的测量行"
-              f"（before={before} after={after}）")
-        return rows
-    values = [float(r["speedup_percent"]) for r in rows]
-    for r in rows:
-        print(f"  {r['dtype']:<16} {float(r['speedup_percent']):+7.2f}%  "
-              f"{r['before_lat']} -> {r['after_lat']}  {r['size_detail']}")
-    fallback = ""
-    if DIRECT_NO_CUDAGRAPH_RE.search(before_text) or \
-            DIRECT_NO_CUDAGRAPH_RE.search(after_text):
-        # 回退行用的是 do_bench 口径，与 cudagraph 行不同表可比性
-        fallback = "  [含 BENCHMARK_DIRECT_NO_CUDAGRAPH 回退行]"
-    print(f"[ab] {name}: {len(rows)} 行配对，speedup "
-          f"{min(values):+.2f}% ~ {max(values):+.2f}%{fallback}")
-    return rows
-
-
 def main() -> None:
     p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
-    p.add_argument("batch_dir", type=Path, nargs="?",
+    p.add_argument("batch_dir", type=Path, help="run_batch.py 的产物目录")
-                   help="run_batch.py 的产物目录")
     p.add_argument("-o", "--output", type=Path, default=None,
                    help="输出 CSV（默认 <batch_dir>/summary.csv）")
-    p.add_argument("--ab", nargs=2, type=Path, metavar=("BEFORE", "AFTER"),
-                   help="两份 run.log，按 shape 配对输出该算子的 speedup 区间")
-    p.add_argument("--label", default="", help="--ab 报告里的算子名")
     args = p.parse_args()
-    if args.ab:
-        for path in args.ab:
-            if not path.is_file():
-                raise SystemExit(f"run.log 不存在: {path}")
-        report_ab(args.ab[0], args.ab[1], args.label)
-        return
-    if args.batch_dir is None:
-        raise SystemExit("需要 batch_dir，或用 --ab BEFORE AFTER")
     if not args.batch_dir.is_dir():
         raise SystemExit(f"目录不存在: {args.batch_dir}")
     out = summarize(args.batch_dir, args.output)

ops/gen_inventory.py

@@ -41,6 +41,8 @@ import yaml
 
 OPS_DIR = Path(__file__).resolve().parent
 
+# Inventory is FlagGems-logical (repo/op/op_file), not per-hardware.
+# One case is reused across backends; do not duplicate CSVs per vendor.
 DEFAULT_REPOS = [
     ("main", Path("/workspace/dev/FlagGems")),
     ("vllm", Path("/workspace/dev/FlagGems-vllm")),

plugin_plan.py

@@ -0,0 +1,107 @@
+"""Decide which pytest plugins a zl_bench run should load.
+
+CUDA-specific plugins (cudagraph patch, Hopper mm cluster workaround) are not
+placed on the default list; `_timing_plugin` enables graph capture only when
+the backend declares `supports_cuda_graph`. IR / autotune plugins stay on for
+Triton-compatible compilers and can be turned off with env flags.
+"""
+
+from __future__ import annotations
+
+import argparse
+import os
+from typing import List, Optional, Sequence
+
+from backends import get_backend
+
+CORE_PLUGINS: tuple[str, ...] = (
+    "_device_setup_plugin",
+    "_parallel_warmup_plugin",
+    "_seed_plugin",
+    "_shape_inject_plugin",
+    "_shape_iter_inject_plugin",
+    "_bespoke_shape_plugin",
+)
+
+ALWAYS_AFTER: tuple[str, ...] = (
+    "_pretty_report_plugin",
+)
+
+
+def _env_flag(name: str, default: str = "1") -> bool:
+    raw = os.environ.get(name, default).strip().lower()
+    return raw not in ("0", "false", "no", "off", "")
+
+
+def plan_plugins(
+    backend_name: Optional[str] = None,
+    extra: Optional[Sequence[str]] = None,
+    enable_autotune: Optional[bool] = None,
+    enable_ir: Optional[bool] = None,
+    enable_mm_cluster: Optional[bool] = None,
+) -> List[str]:
+    backend = get_backend(backend_name)
+    plugins: List[str] = list(CORE_PLUGINS)
+
+    autotune = (
+        backend.supports_triton_hooks
+        if enable_autotune is None
+        else enable_autotune
+    )
+    if enable_autotune is None:
+        autotune = autotune and _env_flag("ZL_ENABLE_AUTOTUNE", "1")
+    if autotune:
+        plugins.append("_autotune_record_plugin")
+
+    plugins.append("_timing_plugin")
+    plugins.extend(ALWAYS_AFTER)
+
+    ir = backend.supports_triton_hooks if enable_ir is None else enable_ir
+    if enable_ir is None:
+        ir = ir and _env_flag("ZL_ENABLE_IR", "1")
+    if ir:
+        plugins.append("_ir_meta_plugin")
+
+    mm = enable_mm_cluster
+    if mm is None:
+        mm = backend.supports_mm_cluster_fix and _env_flag("ZL_ENABLE_MM_CLUSTER", "0")
+    elif mm:
+        mm = backend.supports_mm_cluster_fix
+    if mm:
+        plugins.append("_mm_cluster_fix_plugin")
+
+    if extra is None:
+        extras = [p for p in os.environ.get("ZL_EXTRA_PLUGINS", "").split() if p]
+    else:
+        extras = [p for p in extra if p]
+    for name in extras:
+        if name not in plugins:
+            plugins.append(name)
+    return plugins
+
+
+def pytest_args(plugins: Optional[Sequence[str]] = None, **kwargs) -> List[str]:
+    names = list(plugins) if plugins is not None else plan_plugins(**kwargs)
+    args: List[str] = []
+    for name in names:
+        args.extend(("-p", name))
+    return args
+
+
+def main(argv: Optional[Sequence[str]] = None) -> int:
+    p = argparse.ArgumentParser(description="Print zl_bench pytest -p args")
+    p.add_argument("--backend", default=None)
+    p.add_argument("--pytest-args", action="store_true",
+                   help="print `-p name` tokens on one line")
+    p.add_argument("--one-per-line", action="store_true")
+    args = p.parse_args(argv)
+    names = plan_plugins(backend_name=args.backend)
+    if args.one_per_line:
+        print("\n".join(names))
+    else:
+        print(" ".join(pytest_args(names)))
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

providers/__init__.py

@@ -0,0 +1 @@
+"""Benchmark providers. Phase 1 only implements FlagGems."""

providers/flaggems.py

@@ -0,0 +1,219 @@
+"""FlagGems-specific mapping kept out of the generic scheduler.
+
+GEMS_VENDOR values are FlagGems' vendor names, not zl_bench backend names.
+cuda → nvidia is the important one: FlagGems' device_finder otherwise shells
+out to nvidia-smi without a timeout.
+"""
+
+from __future__ import annotations
+
+import argparse
+import os
+import shlex
+import sys
+from pathlib import Path
+from typing import Dict, Optional
+
+_ROOT = Path(__file__).resolve().parent.parent
+if str(_ROOT) not in sys.path:
+    sys.path.insert(0, str(_ROOT))
+
+from backends.base import DeviceBackend
+
+# Live checkouts used by the NV FlagGems product path. Order is the fallback
+# sequence when the configured root does not contain the requested test file.
+KNOWN_FLAGGEMS_ROOTS = (
+    "/workspace/FlagGems",
+    "/workspace/dev/FlagGems",
+)
+
+VENDOR_ENV_KEYS = (
+    "GEMS_VENDOR", "FLAGGEMS_VENDOR", "GEMS_BACKEND", "FLAGGEMS_BACKEND",
+)
+
+VENDOR_BY_BACKEND: Dict[str, str] = {
+    "cuda": "nvidia",
+    "musa": "musa",
+    "rocm": "amd",
+    "ascend": "ascend",
+}
+
+
+def flaggems_vendor(backend: DeviceBackend) -> str:
+    return backend.flaggems_vendor or VENDOR_BY_BACKEND[backend.name]
+
+
+def apply_vendor_env(backend: DeviceBackend) -> Optional[str]:
+    """Set GEMS_VENDOR unless the user already chose one. Returns the value."""
+    for key in VENDOR_ENV_KEYS:
+        if key in os.environ and os.environ[key].strip():
+            return os.environ[key].strip()
+    vendor = flaggems_vendor(backend)
+    if vendor:
+        os.environ["GEMS_VENDOR"] = vendor
+    return vendor or None
+
+
+def make_test_node(flaggems_dir: str, op_file: str, op: str) -> str:
+    root = flaggems_dir.replace("\\", "/").rstrip("/")
+    return f"{root}/benchmark/test_{op_file}.py::test_{op}"
+
+
+# Back-compat alias. Do not call this from emit_run_env: the parameter of the
+# same name would shadow it and raise TypeError.
+test_node = make_test_node
+
+
+def parse_test_node(node: str) -> Dict[str, Optional[str]]:
+    """Parse a FlagGems pytest node into op / op_file / optional repo root.
+
+    Example:
+      /workspace/FlagGems/benchmark/test_fused_marlin_moe.py::test_fused_marlin_moe_mxfp4
+    """
+    raw = (node or "").strip()
+    if "::" not in raw:
+        raise ValueError(f"pytest node must contain '::': {node!r}")
+    file_part, func = raw.rsplit("::", 1)
+    if not func.startswith("test_"):
+        raise ValueError(f"test function must start with 'test_': {func!r}")
+    stem = Path(file_part.replace("\\", "/").rsplit("/", 1)[-1]).stem
+    if not stem.startswith("test_"):
+        raise ValueError(f"benchmark filename must start with 'test_': {file_part!r}")
+    posix_file = file_part.replace("\\", "/")
+    parent = posix_file.rsplit("/", 1)[0] if "/" in posix_file else ""
+    parent_name = parent.rsplit("/", 1)[-1] if parent else ""
+    flaggems_dir = (
+        parent.rsplit("/", 1)[0] if parent_name == "benchmark" and "/" in parent else None
+    )
+    return {
+        "test_file": f"{file_part}::{func}",
+        "file": file_part,
+        "op": func[len("test_"):],
+        "op_file": stem[len("test_"):],
+        "flaggems_dir": flaggems_dir,
+    }
+
+
+def _test_py(flaggems_dir: str, op_file: str) -> Path:
+    return Path(flaggems_dir) / "benchmark" / f"test_{op_file}.py"
+
+
+def resolve_flaggems_dir(configured: str, op_file: str) -> str:
+    """Return the first root that actually contains benchmark/test_<op_file>.py."""
+    candidates = []
+    for item in (configured, *KNOWN_FLAGGEMS_ROOTS):
+        if item and item not in candidates:
+            candidates.append(item)
+    for root in candidates:
+        if _test_py(root, op_file).is_file():
+            return root
+    return configured
+
+
+def resolve_existing_test_file(test_file: str, op_file: str) -> str:
+    """If the node path is missing, retry the same relative path under known roots."""
+    if "::" not in test_file:
+        return test_file
+    file_part, func = test_file.rsplit("::", 1)
+    path = Path(file_part)
+    if path.is_file():
+        return test_file
+    for root in KNOWN_FLAGGEMS_ROOTS:
+        candidate = _test_py(root, op_file)
+        if candidate.is_file():
+            return f"{candidate.as_posix()}::{func}"
+    return test_file
+
+
+def emit_run_env(
+    test_node: str = "",
+    op: str = "",
+    op_file: str = "",
+    flaggems_dir: str = "",
+    shape_file: str = "",
+    shape_mode: str = "",
+) -> Dict[str, str]:
+    """Resolve OP / OP_FILE / TEST_FILE / SHAPE_MODE for run_pytest.sh."""
+    parsed = parse_test_node(test_node) if test_node.strip() else None
+    if parsed:
+        op = parsed["op"] or op
+        op_file = parsed["op_file"] or op_file
+        if parsed["flaggems_dir"]:
+            flaggems_dir = parsed["flaggems_dir"]
+        test_file = parsed["test_file"] or ""
+    else:
+        flaggems_dir = resolve_flaggems_dir(flaggems_dir, op_file)
+        test_file = make_test_node(flaggems_dir, op_file, op)
+
+    flaggems_dir = resolve_flaggems_dir(flaggems_dir, op_file)
+    test_file = resolve_existing_test_file(test_file, op_file)
+
+    shape_file = shape_file.strip()
+    shape_mode = shape_mode.strip().lower()
+    if shape_file in ("-", "upstream", "none"):
+        shape_file = ""
+        shape_mode = "upstream"
+    if not shape_mode:
+        shape_mode = "upstream" if (test_node.strip() and not shape_file) else "inline"
+    if shape_file:
+        shape_mode = "file"
+        pass_shape_file = "1"
+    elif shape_mode == "upstream":
+        pass_shape_file = "0"
+    else:
+        pass_shape_file = "1"
+        shape_mode = "inline"
+
+    return {
+        "OP": op,
+        "OP_FILE": op_file,
+        "FLAGGEMS_DIR": flaggems_dir,
+        "TEST_FILE": test_file,
+        "SHAPE_FILE": shape_file,
+        "SHAPE_MODE": shape_mode,
+        "PASS_SHAPE_FILE": pass_shape_file,
+    }
+
+
+def _print_shell(env: Dict[str, str]) -> None:
+    for key, value in env.items():
+        print(f"{key}={shlex.quote(str(value))}")
+
+
+def main(argv: Optional[list] = None) -> int:
+    p = argparse.ArgumentParser(description="FlagGems pytest node helpers")
+    p.add_argument("--parse-node", default="")
+    p.add_argument("--emit-shell", action="store_true")
+    p.add_argument("--op", default="")
+    p.add_argument("--op-file", default="")
+    p.add_argument("--flaggems-dir", default="")
+    p.add_argument("--shape-file", default="")
+    p.add_argument("--shape-mode", default="")
+    p.add_argument("--shell", action="store_true",
+                   help="with --parse-node, print OP/OP_FILE/FLAGGEMS_DIR/TEST_FILE")
+    args = p.parse_args(argv)
+    if args.emit_shell:
+        _print_shell(emit_run_env(
+            test_node=args.parse_node,
+            op=args.op,
+            op_file=args.op_file,
+            flaggems_dir=args.flaggems_dir,
+            shape_file=args.shape_file,
+            shape_mode=args.shape_mode,
+        ))
+        return 0
+    if args.parse_node:
+        parsed = parse_test_node(args.parse_node)
+        if args.shell:
+            _print_shell({k: v or "" for k, v in parsed.items()
+                          if k in ("op", "op_file", "flaggems_dir", "test_file")})
+            return 0
+        for key, value in parsed.items():
+            print(f"{key}={value or ''}")
+        return 0
+    p.error("pass --emit-shell or --parse-node")
+    return 2
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

run_batch.py

@@ -1,13 +1,16 @@
 #!/usr/bin/env python3
 """多算子批量测试驱动：按算子清单逐个调 run_pytest.sh，多卡分片、串行测量。
 
-每个算子一个独立子进程（CUDA crash 只废单个算子）、独立产物目录
+每个算子一个独立子进程（设备 crash 只废单个算子）、独立产物目录
 <batch>/<repo>/<op>/（与单算子跑 run_pytest.sh 的产物结构完全一致，含
-run.log / shapes.yaml / autotune_records / ttgir）。批量层只做编排：
+run.log / shapes.yaml / autotune_records / ttgir / run_identity.json）。
+批量层只做编排：
 
-- 分片：按算子名稳定哈希到 GPU（每卡内部串行，测量口径与单算子一致；
+- 分片：按算子名稳定哈希到设备（每卡内部串行，测量口径与单算子一致；
   跨卡并行测量存在功耗耦合噪声，本模式定位为筛查口径，可疑算子回
   单算子串行流程确认）；
+- 后端：--backend / ZL_BACKEND 选择 cuda/musa/rocm/ascend，设备可见性
+  由 backends/ 写入，不再写死 CUDA_VISIBLE_DEVICES；
 - 超时：--op-timeout 杀死挂死算子（记 TIMEOUT，不重试，继续跑）；
 - 重试：pytest rc=1（用例失败/flaky）与信号杀（rc>=128）最多重试 2 次，
   连续两次同 rc 视为确定性失败提前止损；--dtypes 下两级降级——遇上游
@@ -23,7 +26,7 @@ run.log / shapes.yaml / autotune_records / ttgir）。批量层只做编排：
 
 示例：
   # 全量筛查（两仓库全部可收集算子，排除 marked_skip 与 dismiss）
-  python run_batch.py --all --gpus 0,1,2,3,4,5,6,7 --op-timeout 1800
+  python run_batch.py --all --devices 0,1,2,3,4,5,6,7 --backend cuda --op-timeout 1800
 
   # 指定算子子集（repo 前缀默认 main；重名用 @文件名 消歧）
   python run_batch.py --ops softmax vllm:fused_marlin_moe nextafter_@nextafter_
@@ -50,6 +53,8 @@ from pathlib import Path
 from typing import Dict, List, Optional, Set, Tuple
 
 import batch_summary
+from backends import UnknownBackendError, get_backend
+from _run_identity import load_identity
 
 ZL_BENCH = Path(__file__).resolve().parent
 OPS_DIR = ZL_BENCH / "ops"
@@ -62,7 +67,9 @@ REPO_DIRS = {
 }
 
 STATUS_FIELDS = [
-    "repo", "op", "op_file", "run_name", "gpu", "status", "return_code",
+    "repo", "op", "op_file", "run_name",
+    "backend", "device", "gpu", "timer",
+    "status", "return_code",
     "elapsed_ms", "rows_total", "rows_success", "notes", "out_dir",
 ]
 
@@ -194,48 +201,44 @@ def resolve_selection(args: argparse.Namespace,
     return cases
 
 
-# ---------------- GPU 分片 ----------------
+# ---------------- 设备分片 ----------------
 
-def visible_gpus(arg: Optional[str]) -> List[str]:
+def resolve_device_arg(devices: Optional[str], gpus: Optional[str]) -> Optional[str]:
-    if arg:
+    """Prefer --devices; --gpus is an alias. Passing both is an error."""
-        return [g.strip() for g in arg.split(",") if g.strip()]
+    if devices is not None and gpus is not None:
-    env = os.environ.get("CUDA_VISIBLE_DEVICES", "").strip()
+        raise ValueError("do not pass both --devices and --gpus")
-    if env:
+    return devices if devices is not None else gpus
-        return [g.strip() for g in env.split(",") if g.strip()]
-    try:
-        out = subprocess.run(["nvidia-smi", "--list-gpus"], text=True,
-                             capture_output=True, timeout=30).stdout
-        n = sum(1 for line in out.splitlines() if line.startswith("GPU "))
-        if n:
-            return [str(i) for i in range(n)]
-    except Exception:
-        pass
-    return ["0"]
 
 
-def partition(cases: List[Case], gpus: List[str]) -> Dict[str, List[Case]]:
+def visible_devices(backend, explicit: Optional[str] = None) -> List[str]:
+    return list(backend.enumerate_devices(explicit))
+
+
+def partition(cases: List[Case], devices: List[str]) -> Dict[str, List[Case]]:
     """按 (repo, run_name) 稳定哈希分片：与选择顺序无关，resume 时同卡复用。"""
-    parts: Dict[str, List[Case]] = {g: [] for g in gpus}
+    parts: Dict[str, List[Case]] = {d: [] for d in devices}
     for c in cases:
         digest = hashlib.sha256(f"{c.repo}/{c.run_name}".encode()).digest()
-        parts[gpus[int.from_bytes(digest[:8], "big") % len(gpus)]].append(c)
+        parts[devices[int.from_bytes(digest[:8], "big") % len(devices)]].append(c)
     return parts
 
 
 # ---------------- 单算子执行 ----------------
 
-def _spawn(case: Case, out_dir: Path, gpu: str, args: argparse.Namespace,
+def _spawn(case: Case, out_dir: Path, device: str, args: argparse.Namespace,
            dtypes: str) -> subprocess.Popen:
     env = os.environ.copy()
     env.pop("PARALLEL_WARMUP_GPUS", None)   # 批量层已分片，算子内不再并行
     env.pop("REPLAY_FROM", None)
     env.pop("OUT_DIR", None)
+    env = args.backend_obj.apply_visibility(env, device)
     env.update({
         "OP": case.op,
         "OP_FILE": case.op_file,
         "FLAGGEMS_DIR": REPO_DIRS[case.repo],
         "OUT_DIR": str(out_dir),
-        "CUDA_VISIBLE_DEVICES": gpu,
+        "ZL_BACKEND": args.backend,
+        "ZL_DEVICE": device,
         "FLAGGEMS_PERF_COLOR": "never",
         "NO_COLOR": "1",
         "USE_FLAGTUNE": str(args.use_flagtune),
@@ -271,13 +274,14 @@ def _wait(proc: subprocess.Popen, timeout: int) -> Tuple[int, bool]:
         return 124, True
 
 
-def run_case(case: Case, batch_dir: Path, gpu: str,
+def run_case(case: Case, batch_dir: Path, device: str,
              args: argparse.Namespace) -> Dict[str, str]:
     out_dir = batch_dir / case.repo / case.run_name
     marker = out_dir / ".complete"
 
     if args.resume and marker.is_file() and (out_dir / "run.log").is_file():
-        row = _classify(case, out_dir, gpu, 0, 0, "", timed_out=False)
+        row = _classify(case, out_dir, device, 0, 0, "", timed_out=False,
+                        args=args)
         row["notes"] = _join("resumed", row["notes"])
         return row
 
@@ -295,7 +299,7 @@ def run_case(case: Case, batch_dir: Path, gpu: str,
             log = out_dir / "run.log"
             if log.is_file():
                 log.rename(out_dir / f"run.attempt{attempt - 1}.log")
-        rc, timed_out = _wait(_spawn(case, out_dir, gpu, args, dtypes),
+        rc, timed_out = _wait(_spawn(case, out_dir, device, args, dtypes),
                               args.op_timeout)
         rcs.append(rc)
         if timed_out:
@@ -315,20 +319,22 @@ def run_case(case: Case, batch_dir: Path, gpu: str,
     # 覆盖 baseline 对受限 dtype 编译失败（torch jiterator 的 special 函数族等）
     # 但报错文案不是上游 "can't be supported by this op" 的场景。
     if dtypes and not timed_out:
-        probe = _classify(case, out_dir, gpu, rcs[-1], 0, "", False)
+        probe = _classify(case, out_dir, device, rcs[-1], 0, "", False,
+                          args=args)
         if probe["status"] == "FAIL" and probe["rows_success"] == "0":
             log = out_dir / "run.log"
             if log.is_file():
                 log.rename(out_dir / f"run.attempt{len(rcs) - 1}.log")
             notes = _join(notes, "dtype_rescue")
-            rc, timed_out = _wait(_spawn(case, out_dir, gpu, args, ""),
+            rc, timed_out = _wait(_spawn(case, out_dir, device, args, ""),
                                   args.op_timeout)
             rcs.append(rc)
     elapsed_ms = int((time.time() - t0) * 1000)
 
     if len(rcs) > 1:
         notes = _join(notes, "retry_rc=" + "->".join(map(str, rcs)))
-    row = _classify(case, out_dir, gpu, rcs[-1], elapsed_ms, notes, timed_out)
+    row = _classify(case, out_dir, device, rcs[-1], elapsed_ms, notes, timed_out,
+                    args=args)
     if row["status"] in ("PASS", "SKIP"):
         tmp = marker.with_name(".complete.tmp")
         tmp.write_text(row["status"] + "\n")
@@ -345,8 +351,21 @@ def _join(*notes: str) -> str:
     return ";".join(n for n in notes if n)
 
 
-def _classify(case: Case, out_dir: Path, gpu: str, rc: int, elapsed_ms: int,
+def _identity_fields(out_dir: Path, device: str, backend: str) -> Dict[str, str]:
-              notes: str, timed_out: bool) -> Dict[str, str]:
+    ident = load_identity(str(out_dir))
+    backend_v = str(ident.get("backend") or backend or "")
+    device_v = str(ident.get("device") or device or "")
+    timer_v = str(ident.get("timer") or "")
+    return {
+        "backend": backend_v,
+        "device": device_v,
+        "gpu": device_v,
+        "timer": timer_v,
+    }
+
+
+def _classify(case: Case, out_dir: Path, device: str, rc: int, elapsed_ms: int,
+              notes: str, timed_out: bool, args: argparse.Namespace) -> Dict[str, str]:
     text = _read_log(out_dir)
     rows = batch_summary.parse_run_log(text)
     if timed_out:
@@ -356,12 +375,13 @@ def _classify(case: Case, out_dir: Path, gpu: str, rc: int, elapsed_ms: int,
     n_skipped = batch_summary.pytest_skipped_count(text)
     if status == "PASS" and n_skipped:
         notes = _join(notes, f"partial_skip={n_skipped}")
+    ident = _identity_fields(out_dir, device, args.backend)
     return {
         "repo": case.repo,
         "op": case.op,
         "op_file": case.op_file,
         "run_name": case.run_name,
-        "gpu": gpu,
+        **ident,
         "status": status,
         "return_code": str(rc),
         "elapsed_ms": str(elapsed_ms),
@@ -393,8 +413,12 @@ def parse_args() -> argparse.Namespace:
                      help="--all 时不剔除带无条件 skip 标记的算子")
 
     run = p.add_argument_group("执行")
+    run.add_argument("--backend", default=os.environ.get("ZL_BACKEND", "cuda"),
+                     help="设备后端：cuda/musa/rocm/ascend（默认 ZL_BACKEND 或 cuda）")
+    run.add_argument("--devices", default=None,
+                     help="逗号分隔设备列表（默认该后端可见设备；优先于 --gpus）")
     run.add_argument("--gpus", default=None,
-                     help="逗号分隔 GPU 列表（默认 CUDA_VISIBLE_DEVICES 或全部）")
+                     help="--devices 的别名（两者同时给出则报错）")
     run.add_argument("--batch-dir", type=Path, default=None,
                      help="产物目录（默认 runs/batch_<时间戳>；resume 时必须指定）")
     run.add_argument("--op-timeout", type=int, default=1800,
@@ -429,19 +453,30 @@ def parse_args() -> argparse.Namespace:
 
 def main() -> None:
     args = parse_args()
+    try:
+        backend = get_backend(args.backend)
+    except UnknownBackendError as exc:
+        raise SystemExit(f"{LOG} {exc}") from exc
+    args.backend_obj = backend
+    args.backend = backend.name
+    try:
+        explicit = resolve_device_arg(args.devices, args.gpus)
+    except ValueError as exc:
+        raise SystemExit(f"{LOG} {exc}") from exc
+    devices = visible_devices(backend, explicit)
+
     inv = load_inventory()
     cases = resolve_selection(args, inv)
     if not cases:
         raise SystemExit(f"{LOG} 选择结果为空")
 
-    gpus = visible_gpus(args.gpus)
     batch_dir = (args.batch_dir or
                  ZL_BENCH / "runs" / f"batch_{datetime.now():%Y%m%d-%H%M%S}")
     batch_dir = batch_dir.resolve()
     batch_dir.mkdir(parents=True, exist_ok=True)
 
-    parts = {g: cs for g, cs in partition(cases, gpus).items() if cs}
+    parts = {d: cs for d, cs in partition(cases, devices).items() if cs}
-    print(f"{LOG} {len(cases)} 个算子 -> {len(parts)} 张卡 "
+    print(f"{LOG} backend={args.backend}  {len(cases)} 个算子 -> {len(parts)} 台设备 "
           f"(shards: {[len(cs) for cs in parts.values()]})  out={batch_dir}")
     if args.replay_root:
         print(f"{LOG} replay 模式: REPLAY_FROM={args.replay_root}/<repo>/<op>")
@@ -461,17 +496,19 @@ def main() -> None:
             writer.writerows(rows)
         os.replace(tmp, status_csv)
 
-    def run_partition(gpu: str, part: List[Case]) -> None:
+    def run_partition(device: str, part: List[Case]) -> None:
         for case in part:
             try:
-                row = run_case(case, batch_dir, gpu, args)
+                row = run_case(case, batch_dir, device, args)
             except Exception as exc:  # 编排层异常不拖垮整卡分片
+                out_dir = batch_dir / case.repo / case.run_name
+                ident = _identity_fields(out_dir, device, args.backend)
                 row = {
                     "repo": case.repo, "op": case.op, "op_file": case.op_file,
-                    "run_name": case.run_name, "gpu": gpu, "status": "FAIL",
+                    "run_name": case.run_name, **ident, "status": "FAIL",
                     "return_code": "EXC", "elapsed_ms": "0", "rows_total": "0",
                     "rows_success": "0", "notes": f"harness_exc={exc}",
-                    "out_dir": str(batch_dir / case.repo / case.run_name),
+                    "out_dir": str(out_dir),
                 }
             with lock:
                 done[0] += 1
@@ -480,16 +517,16 @@ def main() -> None:
                 print(f"  ({done[0]:>4}/{len(cases)})  {row['status']:<7} "
                       f"{row['elapsed_ms']:>8}ms  rows={row['rows_success']}/"
                       f"{row['rows_total']:<4} {case.repo}:{case.run_name:<{wname}} "
-                      f"[gpu{gpu}]"
+                      f"[device{device}]"
                       + (f"  {row['notes']}" if row["notes"] else ""),
                       flush=True)
 
     if len(parts) == 1:
-        gpu, part = next(iter(parts.items()))
+        device, part = next(iter(parts.items()))
-        run_partition(gpu, part)
+        run_partition(device, part)
     else:
         with ThreadPoolExecutor(max_workers=len(parts)) as pool:
-            futures = [pool.submit(run_partition, g, p) for g, p in parts.items()]
+            futures = [pool.submit(run_partition, d, p) for d, p in parts.items()]
             for f in futures:
                 f.result()
 

run_pytest.sh

@@ -8,15 +8,21 @@ set -euo pipefail
 # 每项均可用同名环境变量覆盖，详见 README「环境变量一览」。
 
 # 被测算子，对应 $FLAGGEMS_DIR/benchmark/test_<OP_FILE>.py::test_<OP>
+# 也可直接传 FlagGems pytest 节点，例如：
+#   TEST_NODE=/workspace/FlagGems/benchmark/test_fused_marlin_moe.py::test_fused_marlin_moe_mxfp4
 OP="${OP:-fused_marlin_moe_mxfp4}"
 OP_FILE="${OP_FILE:-fused_marlin_moe}"
 
 # 上游已拆分为两个仓库：主仓 /workspace/dev/FlagGems（包名 flag_gems）与
 # /workspace/dev/FlagGems-vllm（包名 flaggems_vllm）。测 vllm 仓库算子时覆盖本变量即可。
+# NV 产品路径也常见 /workspace/FlagGems；TEST_NODE 或缺失测试文件时会回退到该路径。
 FLAGGEMS_DIR="${FLAGGEMS_DIR:-/workspace/dev/FlagGems}"
+TEST_NODE="${TEST_NODE:-${ZL_TEST_NODE:-}}"
 
-# shape 来源：非空则用该 yaml，否则用 INLINE_YAML。顶层 key 必须是 op 名。
+# shape 来源：非空则用该 yaml；SHAPE_MODE=upstream 则不传 --shape_file，沿用 FlagGems
+# 在 $BENCH_LEVEL 下的原生 shape 集。设了 TEST_NODE 且未指定 SHAPE_FILE 时默认 upstream。
 SHAPE_FILE="${SHAPE_FILE:-}"
+SHAPE_MODE="${SHAPE_MODE:-}"
 read -r -d '' INLINE_YAML <<'YAML' || true
 fused_marlin_moe_mxfp4:
   shapes:
@@ -53,25 +59,34 @@ REPLAY_FROM="${REPLAY_FROM:-}"
 # always/never 强制开关终端颜色，空则按 tty 判断（run.log 始终去色）
 FLAGGEMS_PERF_COLOR="${FLAGGEMS_PERF_COLOR:-}"
 
-# 各插件作用见 README「插件说明」；注释掉某行即停用该插件
+# 设备后端：cuda / musa / rocm / ascend（别名见 backends/__init__.py）。
-PLUGINS=(
+# 插件列表由 plugin_plan.py 按后端能力生成，不再在此硬编码 -p。
-  -p _device_guard_plugin
+# 非 CUDA 默认走上游 do_bench，不会加载 _cudagraph_plugin。
-  -p _parallel_warmup_plugin
+ZL_BACKEND="${ZL_BACKEND:-cuda}"
-  -p _seed_plugin
+
-  -p _shape_inject_plugin
+# 上游 --level：core（默认）只用各算子的 core_shapes；comprehensive 额外并入
-  -p _shape_iter_inject_plugin
+# 子类的 set_more_shapes()，并对同一 shape 再产出一次变体（BLAS 是转置侧）。
-  -p _bespoke_shape_plugin
+# 差距可以很大：mm 在 core 下是 5 个 shape，comprehensive 下并入 4 个大 K
-  # -p _mm_cluster_fix_plugin
+# （N/K 到 152064）与 models_shapes/*.yaml 的模型 shape，共 29 个，且用例翻倍。
-  -p _autotune_record_plugin
+# 代价同步放大——sweep 是墙上时间的大头，shape 多时配合 PARALLEL_WARMUP_GPUS。
-  -p _cudagraph_plugin
+BENCH_LEVEL="${BENCH_LEVEL:-core}"
-  -p _pretty_report_plugin
+case "$BENCH_LEVEL" in
-  -p _ir_meta_plugin
+  core|comprehensive) ;;
-)
+  *) echo "!!! BENCH_LEVEL 只能是 core 或 comprehensive（当前：$BENCH_LEVEL）" >&2; exit 2 ;;
+esac
 
 # ============================== 执行逻辑 ==============================
 
 SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
-TEST_FILE="$FLAGGEMS_DIR/benchmark/test_${OP_FILE}.py::test_${OP}"
+export PYTHONPATH="$SCRIPT_DIR${PYTHONPATH:+:$PYTHONPATH}"
+eval "$(python "$SCRIPT_DIR/providers/flaggems.py" --emit-shell \
+  --parse-node "$TEST_NODE" \
+  --op "$OP" \
+  --op-file "$OP_FILE" \
+  --flaggems-dir "$FLAGGEMS_DIR" \
+  --shape-file "$SHAPE_FILE" \
+  --shape-mode "$SHAPE_MODE")"
+: "${TEST_FILE:=$FLAGGEMS_DIR/benchmark/test_${OP_FILE}.py::test_${OP}}"
 
 # 产物目录，可用 OUT_DIR 指定。时间戳只到秒，并发启动会撞名，故用不带 -p 的
 # mkdir 抢占（已存在即失败），撞上就退让到 -2、-3……
@@ -104,9 +119,12 @@ else
   PYTEST_COLOR=()
 fi
 
-export PYTHONPATH="$SCRIPT_DIR${PYTHONPATH:+:$PYTHONPATH}"
 export FLAGGEMS_PERF_CURRENT_OP="$OP"
 export PYTHONUNBUFFERED=1  # 实时输出不缓冲
+export ZL_BACKEND
+
+# 按 ZL_BACKEND / ZL_ENABLE_* 生成 -p 列表（见 plugin_plan.py）
+read -r -a PLUGINS <<< "$(python "$SCRIPT_DIR/plugin_plan.py" --backend "$ZL_BACKEND" --pytest-args)"
 
 # DTYPES 非空时逐个转为上游 --dtypes 选项（action=append，每个 dtype 一次）
 DTYPE_ARGS=()
@@ -126,17 +144,39 @@ else
   MODE_DESC=record
 fi
 
-# 未指定 shape 文件时，把 INLINE_YAML 落到临时文件供 pytest 读取
+# 未指定 shape 文件时：inline 用 INLINE_YAML；upstream 不传 --shape_file。
-if [[ -z "$SHAPE_FILE" ]]; then
+SHAPE_ARGS=()
+if [[ "${PASS_SHAPE_FILE:-1}" == "0" ]]; then
+  if [[ -f "$FLAGGEMS_DIR/benchmark/core_shapes.yaml" ]]; then
+    cp -f "$FLAGGEMS_DIR/benchmark/core_shapes.yaml" "$OUT_DIR/shapes.yaml"
+  else
+    printf '%s\n' "source: upstream" "level: $BENCH_LEVEL" > "$OUT_DIR/shapes.yaml"
+  fi
+else
+  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"  # 存档本次实际使用的 shape
+  SHAPE_ARGS=(--shape_file "$SHAPE_FILE")
 fi
-cp -f "$SHAPE_FILE" "$OUT_DIR/shapes.yaml"  # 存档本次实际使用的 shape
+
+python "$SCRIPT_DIR/_run_identity.py" "$OUT_DIR" \
+  --backend "$ZL_BACKEND" \
+  --op "$OP" \
+  --op-file "$OP_FILE" \
+  --flaggems-dir "$FLAGGEMS_DIR" \
+  --autotune-mode "$MODE_DESC" \
+  --device "${ZL_DEVICE:-}" \
+  --test-file "$TEST_FILE" \
+  --shape-mode "$SHAPE_MODE" \
+  --bench-level "$BENCH_LEVEL"
 
 status=0
 {
-  echo "${C_BOLD}>>> op=$OP  mode=$MODE_DESC  USE_FLAGTUNE=$USE_FLAGTUNE${C_RESET}"
+  echo "${C_BOLD}>>> op=$OP  backend=$ZL_BACKEND  mode=$MODE_DESC  USE_FLAGTUNE=$USE_FLAGTUNE  shape=$SHAPE_MODE  level=$BENCH_LEVEL${C_RESET}"
+  echo "${C_DIM}>>> node=$TEST_FILE${C_RESET}"
   echo "${C_DIM}>>> out=$OUT_DIR${C_RESET}"
 
   # 每次用独立的 Triton 缓存目录，跑完即删：保证编译过程可复现，且 ttgir 落盘
@@ -149,8 +189,8 @@ status=0
   env "$AUTOTUNE_ENV" \
   USE_FLAGTUNE=$USE_FLAGTUNE python -u -m pytest -s "$TEST_FILE" \
     "${PLUGINS[@]}" "${PYTEST_COLOR[@]}" \
-    --shape_file "$SHAPE_FILE" ${DTYPE_ARGS[@]+"${DTYPE_ARGS[@]}"} \
+    ${SHAPE_ARGS[@]+"${SHAPE_ARGS[@]}"} ${DTYPE_ARGS[@]+"${DTYPE_ARGS[@]}"} \
-    --level core --mode kernel || status=$?
+    --level "$BENCH_LEVEL" --mode kernel || status=$?
 
   rm -rf "$CACHE_DIR"
 

test_backends.py

@@ -0,0 +1,285 @@
+#!/usr/bin/env python3
+"""No-GPU unit tests for zl_bench backend adapters and plugin policy.
+
+FlagGems pytest invocations pass an explicit node, so this file is not
+collected during operator runs. Run it with `python test_backends.py`.
+"""
+from __future__ import annotations
+
+import json
+import os
+import sys
+import tempfile
+import unittest
+from pathlib import Path
+
+ROOT = Path(__file__).resolve().parent
+if str(ROOT) not in sys.path:
+    sys.path.insert(0, str(ROOT))
+
+from backends import (  # noqa: E402
+    UnknownBackendError,
+    get_backend,
+    normalize_backend_name,
+)
+from backends.base import ALL_VISIBILITY_KEYS  # noqa: E402
+from batch_summary import SUMMARY_FIELDS  # noqa: E402
+from plugin_plan import plan_plugins  # noqa: E402
+from providers.flaggems import (  # noqa: E402
+    VENDOR_BY_BACKEND,
+    emit_run_env,
+    flaggems_vendor,
+    parse_test_node,
+)
+from run_batch import STATUS_FIELDS, resolve_device_arg  # noqa: E402
+from _run_identity import init_identity, load_identity, update_identity  # noqa: E402
+
+
+class BackendResolveTests(unittest.TestCase):
+    def test_empty_and_alias_resolve_to_cuda(self):
+        self.assertEqual(normalize_backend_name(""), "cuda")
+        self.assertEqual(get_backend("nvidia").name, "cuda")
+        self.assertEqual(get_backend("nv").name, "cuda")
+        self.assertEqual(get_backend("").name, "cuda")
+
+    def test_other_aliases(self):
+        self.assertEqual(get_backend("amd").name, "rocm")
+        self.assertEqual(get_backend("hip").name, "rocm")
+        self.assertEqual(get_backend("npu").name, "ascend")
+        self.assertEqual(get_backend("huawei").name, "ascend")
+
+    def test_unknown_backend_fails_closed(self):
+        with self.assertRaises(UnknownBackendError):
+            get_backend("tpu")
+        with self.assertRaises(UnknownBackendError):
+            get_backend("not-a-backend")
+
+
+class VisibilityTests(unittest.TestCase):
+    def test_musa_drops_leftover_cuda(self):
+        env = {"CUDA_VISIBLE_DEVICES": "7", "PATH": "/bin"}
+        out = get_backend("musa").apply_visibility(env, "1")
+        self.assertNotIn("CUDA_VISIBLE_DEVICES", out)
+        self.assertEqual(out["MUSA_VISIBLE_DEVICES"], "1")
+        self.assertEqual(out["PATH"], "/bin")
+        for key in ALL_VISIBILITY_KEYS:
+            if key != "MUSA_VISIBLE_DEVICES":
+                self.assertNotIn(key, out)
+
+    def test_rocm_writes_hip_and_rocr(self):
+        out = get_backend("rocm").apply_visibility(
+            {"CUDA_VISIBLE_DEVICES": "0"}, "3")
+        self.assertEqual(out["HIP_VISIBLE_DEVICES"], "3")
+        self.assertEqual(out["ROCR_VISIBLE_DEVICES"], "3")
+        self.assertNotIn("CUDA_VISIBLE_DEVICES", out)
+
+    def test_ascend_writes_both_keys(self):
+        out = get_backend("ascend").apply_visibility({}, "2")
+        self.assertEqual(out["ASCEND_RT_VISIBLE_DEVICES"], "2")
+        self.assertEqual(out["ASCEND_VISIBLE_DEVICES"], "2")
+        self.assertNotIn("CUDA_VISIBLE_DEVICES", out)
+
+    def test_cuda_only_sets_cuda(self):
+        out = get_backend("cuda").apply_visibility(
+            {"MUSA_VISIBLE_DEVICES": "9"}, "0")
+        self.assertEqual(out["CUDA_VISIBLE_DEVICES"], "0")
+        self.assertNotIn("MUSA_VISIBLE_DEVICES", out)
+
+
+class FlagGemsVendorTests(unittest.TestCase):
+    def test_vendor_map(self):
+        self.assertEqual(VENDOR_BY_BACKEND["cuda"], "nvidia")
+        self.assertEqual(VENDOR_BY_BACKEND["rocm"], "amd")
+        self.assertEqual(flaggems_vendor(get_backend("cuda")), "nvidia")
+        self.assertEqual(flaggems_vendor(get_backend("rocm")), "amd")
+        self.assertEqual(flaggems_vendor(get_backend("musa")), "musa")
+        self.assertEqual(flaggems_vendor(get_backend("ascend")), "ascend")
+
+
+class FlagGemsNodeTests(unittest.TestCase):
+    _NODE = (
+        "/workspace/FlagGems/benchmark/test_fused_marlin_moe.py"
+        "::test_fused_marlin_moe_mxfp4"
+    )
+
+    def test_parse_user_command_node(self):
+        parsed = parse_test_node(self._NODE)
+        self.assertEqual(parsed["op"], "fused_marlin_moe_mxfp4")
+        self.assertEqual(parsed["op_file"], "fused_marlin_moe")
+        self.assertEqual(parsed["flaggems_dir"], "/workspace/FlagGems")
+        self.assertTrue(parsed["test_file"].endswith(
+            "test_fused_marlin_moe.py::test_fused_marlin_moe_mxfp4"))
+
+    def test_test_node_defaults_to_upstream_shapes(self):
+        env = emit_run_env(
+            test_node=self._NODE,
+            op="softmax",
+            op_file="softmax",
+            flaggems_dir="/workspace/dev/FlagGems",
+        )
+        self.assertEqual(env["OP"], "fused_marlin_moe_mxfp4")
+        self.assertEqual(env["OP_FILE"], "fused_marlin_moe")
+        self.assertEqual(env["FLAGGEMS_DIR"], "/workspace/FlagGems")
+        self.assertEqual(env["SHAPE_MODE"], "upstream")
+        self.assertEqual(env["PASS_SHAPE_FILE"], "0")
+        self.assertEqual(env["TEST_FILE"], self._NODE)
+
+    def test_plain_run_keeps_inline_shapes(self):
+        env = emit_run_env(
+            op="fused_marlin_moe_mxfp4",
+            op_file="fused_marlin_moe",
+            flaggems_dir="/workspace/dev/FlagGems",
+        )
+        self.assertEqual(env["SHAPE_MODE"], "inline")
+        self.assertEqual(env["PASS_SHAPE_FILE"], "1")
+        self.assertIn("test_fused_marlin_moe.py::test_fused_marlin_moe_mxfp4",
+                      env["TEST_FILE"])
+
+    def test_explicit_shape_file_wins(self):
+        env = emit_run_env(
+            test_node=self._NODE,
+            shape_file="/tmp/my_shapes.yaml",
+        )
+        self.assertEqual(env["SHAPE_MODE"], "file")
+        self.assertEqual(env["PASS_SHAPE_FILE"], "1")
+
+    def test_resolve_uses_existing_local_root(self):
+        with tempfile.TemporaryDirectory() as tmp:
+            root = Path(tmp) / "FlagGems"
+            py = root / "benchmark" / "test_fused_marlin_moe.py"
+            py.parent.mkdir(parents=True)
+            py.write_text("# stub\n", encoding="utf-8")
+            env = emit_run_env(
+                op="fused_marlin_moe_mxfp4",
+                op_file="fused_marlin_moe",
+                flaggems_dir=str(root),
+            )
+            self.assertEqual(env["FLAGGEMS_DIR"], str(root))
+            self.assertTrue(Path(env["TEST_FILE"].split("::", 1)[0]).is_file())
+
+
+class PluginPlanTests(unittest.TestCase):
+    def _plan(self, backend, **kwargs):
+        kwargs.setdefault("extra", [])
+        kwargs.setdefault("enable_autotune", True)
+        kwargs.setdefault("enable_ir", True)
+        kwargs.setdefault("enable_mm_cluster", False)
+        return plan_plugins(backend, **kwargs)
+
+    def test_cuda_does_not_put_cudagraph_on_the_list(self):
+        names = self._plan("cuda")
+        self.assertIn("_timing_plugin", names)
+        self.assertIn("_device_setup_plugin", names)
+        self.assertNotIn("_cudagraph_plugin", names)
+        self.assertNotIn("_mm_cluster_fix_plugin", names)
+
+    def test_musa_same_core_without_graph_or_cluster(self):
+        names = self._plan("musa")
+        self.assertIn("_timing_plugin", names)
+        self.assertNotIn("_cudagraph_plugin", names)
+        self.assertNotIn("_mm_cluster_fix_plugin", names)
+
+    def test_mm_cluster_ignored_on_non_cuda(self):
+        names = self._plan("musa", enable_mm_cluster=True)
+        self.assertNotIn("_mm_cluster_fix_plugin", names)
+        names = self._plan("cuda", enable_mm_cluster=True)
+        self.assertIn("_mm_cluster_fix_plugin", names)
+
+    def test_extra_plugins_isolated(self):
+        names = self._plan("cuda", extra=["_cudagraph_plugin"])
+        self.assertIn("_cudagraph_plugin", names)
+        names = self._plan("cuda", extra=[])
+        self.assertNotIn("_cudagraph_plugin", names)
+
+    def test_extra_none_reads_env(self):
+        old = os.environ.get("ZL_EXTRA_PLUGINS")
+        os.environ["ZL_EXTRA_PLUGINS"] = "_cudagraph_plugin"
+        try:
+            names = plan_plugins(
+                "cuda", extra=None, enable_autotune=True,
+                enable_ir=True, enable_mm_cluster=False)
+            self.assertIn("_cudagraph_plugin", names)
+        finally:
+            if old is None:
+                os.environ.pop("ZL_EXTRA_PLUGINS", None)
+            else:
+                os.environ["ZL_EXTRA_PLUGINS"] = old
+
+
+class DeviceArgTests(unittest.TestCase):
+    def test_devices_preferred(self):
+        self.assertEqual(resolve_device_arg("0,1", None), "0,1")
+        self.assertEqual(resolve_device_arg(None, "2"), "2")
+        self.assertIsNone(resolve_device_arg(None, None))
+
+    def test_both_set_is_error(self):
+        with self.assertRaises(ValueError):
+            resolve_device_arg("0", "1")
+
+
+class CsvFieldTests(unittest.TestCase):
+    def test_status_fields(self):
+        for name in ("backend", "device", "gpu", "timer"):
+            self.assertIn(name, STATUS_FIELDS)
+
+    def test_summary_fields(self):
+        for name in ("backend", "device", "timer"):
+            self.assertIn(name, SUMMARY_FIELDS)
+        self.assertLess(SUMMARY_FIELDS.index("backend"),
+                        SUMMARY_FIELDS.index("fg_op_name"))
+
+
+class IdentityMergeTests(unittest.TestCase):
+    def test_init_then_update(self):
+        old_out = os.environ.get("OUT_DIR")
+        try:
+            with tempfile.TemporaryDirectory() as tmp:
+                init_identity(tmp, backend="cuda", device="0", op="softmax")
+                update_identity(tmp, timer="direct")
+                ident = load_identity(tmp)
+                self.assertEqual(ident["schema_version"], 1)
+                self.assertEqual(ident["provider"], "flaggems")
+                self.assertEqual(ident["backend"], "cuda")
+                self.assertEqual(ident["device"], "0")
+                self.assertEqual(ident["timer"], "direct")
+                self.assertEqual(ident["op"], "softmax")
+        finally:
+            if old_out is None:
+                os.environ.pop("OUT_DIR", None)
+            else:
+                os.environ["OUT_DIR"] = old_out
+
+
+class AutotuneReplayBackendTests(unittest.TestCase):
+    def setUp(self):
+        self._old = os.environ.get("ZL_BACKEND")
+        os.environ["ZL_BACKEND"] = "cuda"
+
+    def tearDown(self):
+        if self._old is None:
+            os.environ.pop("ZL_BACKEND", None)
+        else:
+            os.environ["ZL_BACKEND"] = self._old
+
+    def test_missing_meta_still_replays(self):
+        from _autotune_record_plugin import _replay_backend_ok
+        with tempfile.TemporaryDirectory() as tmp:
+            self.assertTrue(_replay_backend_ok(tmp))
+
+    def test_matching_backend_ok(self):
+        from _autotune_record_plugin import _replay_backend_ok
+        with tempfile.TemporaryDirectory() as tmp:
+            Path(tmp, "meta.json").write_text(
+                json.dumps({"backend": "cuda"}), encoding="utf-8")
+            self.assertTrue(_replay_backend_ok(tmp))
+
+    def test_mismatch_refuses_replay(self):
+        from _autotune_record_plugin import _replay_backend_ok
+        with tempfile.TemporaryDirectory() as tmp:
+            Path(tmp, "meta.json").write_text(
+                json.dumps({"backend": "musa"}), encoding="utf-8")
+            self.assertFalse(_replay_backend_ok(tmp))
+
+
+if __name__ == "__main__":
+    unittest.main()