4 Commits
Author SHA1 Message Date
Aharrypotter abf9475ff3 feat: report a per-operator speedup range across two runs
The speedup column in a result table compares gems against torch, not one build
against another, so an A/B needs the two runs paired by shape. batch_summary
gains an --ab mode that indexes both run.logs by operator, dtype and size
detail, computes before/after per pair, and prints the per-shape rows plus the
range. Only rows that succeeded on both sides with a positive latency count, so
an incomplete pairing cannot masquerade as a complete range. When either side
carries a BENCHMARK_DIRECT_NO_CUDAGRAPH marker the range is annotated, because
a fallback row was timed under a different regime than a captured one.

ab_rlc_test.sh drives the pairing for the FlagTree RemoveLayoutConversions
switch the way ab_fold_test.sh does for the FlagGems one: per operator, run with
the enhancement off while recording autotune configs, then replay those configs
with it on, so the two sides differ by that switch alone. Each operator
therefore leaves exactly two run.log files, and the summary of each pair is
printed as it completes.

Checked on an S5000 against gelu_backward over five shapes: both sides captured
with no fallback, five rows paired, range -0.43% to +0.09%.
2026-09-02 12:19:26 +08:00
Aharrypotter 7522143747 feat: drop stale CUDA-only skip guards on MThreads
Several upstream FlagGems tests gate themselves on torch.cuda.is_available()
even though the code under test runs on whatever accelerator FlagGems
dispatches to. On MThreads those guards skip the operator instead of exercising
it, so a working path reports as "not measured".

A guard is dropped only when it is pinned to one file and one exact skip
reason, only when FlagGems has dispatched to the MThreads vendor, and only
after the unskipped test was observed to pass on an S5000. The plugin walks the
node chain rather than the item's own markers, because a module-level pytestmark
does not appear on the item.

Two guards were probed on device and deliberately left in place, which is why
they are named in the docstring rather than silently absent.
test_grid_sample.py builds its inputs through a helper whose device argument
defaults to the literal "cuda": unskipped, all 90 collected cases fail in
aten::empty on a CUDA backend that is not present. test_cutlass_scaled_mm.py
gates on vLLM as well as CUDA, so removing its guard would raise on the missing
dependency instead of measuring anything.
2026-09-02 12:19:26 +08:00
Aharrypotter 3c9529dab2 fix: take graph-capture timing from Triton's device interface
The cudagraph plugin swapped do_bench for do_bench_cudagraph, but upstream's
do_bench_cudagraph names torch.cuda directly: torch.cuda.stream, Event,
synchronize, CUDAGraph, graph. FlagTree's do_bench already resolves the device
through driver.active.get_device_interface(), its cudagraph twin does not. On a
non-CUDA backend every capture therefore raised and every measurement silently
fell back to plain do_bench, so the plugin looked enabled while nothing was
being captured.

The plugin now resolves the interface itself and takes the stream, event,
synchronize and graph calls from it. CUDA still calls upstream's
do_bench_cudagraph verbatim, so its numbers do not move; other backends get an
otherwise step-for-step port. The scratch-RNG guard around capture is skipped
when the interface exposes no default_generators, which is the case on MThreads
and where there is no generator registration to protect against.

The replay timings are handed to _summarize_statistics as a plain list of
floats, matching do_bench and upstream's cudagraph path. 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.

Verified on a Moore Threads S5000 through torch_musa, which supplies Stream,
Event, synchronize, graph and MUSAGraph. softmax and gelu_backward capture and
replay with no fallback marker. RNG operators cannot be captured there yet:
dropout aborts with "Cannot call MUSAGeneratorImpl::current_seed during MUSA
graph capture" and takes the documented BENCHMARK_DIRECT_NO_CUDAGRAPH fallback,
so a corpus mixing RNG and non-RNG operators mixes two timing regimes and the
marker is the thing to grep before comparing them.
2026-09-02 12:19:07 +08:00
Aharrypotter 64baa11a23 chore: pin shipped text formats to LF
The drivers here are shell scripts executed on Linux accelerator hosts, and
consumers record the SHA-256 of the files they ran. With Git's default autocrlf
on Windows a checkout rewrites those files with CRLF, which breaks the shell
drivers outright and changes every digest a receipt refers to.
2026-09-02 12:19:06 +08:00
38 changed files with 263 additions and 6768 deletions

No files matched your search

+20 -310
View File
@@ -6,30 +6,6 @@
2. **可解释**:自动按 shape 收集每次运行实际使用的 ttgir(带可读的变体命名),供 IR 级 diff; 2. **可解释**:自动按 shape 收集每次运行实际使用的 ttgir(带可读的变体命名),供 IR 级 diff;
3. **口径统一**:cudagraph 计时消除 launch 开销,小 kernel 的对比不被 CPU 侧噪声淹没;capture 前按 warmup 预算显式预热吸收 autotune/JIT,首轮即稳态、多轮一致。 3. **口径统一**:cudagraph 计时消除 launch 开销,小 kernel 的对比不被 CPU 侧噪声淹没;capture 前按 warmup 预算显式预热吸收 autotune/JIT,首轮即稳态、多轮一致。
## 先读这个:口径决定数字的含义
同一个算子、同一台机器、同一份代码,换一条测试口径,测出来的数字可以完全不同。
下表每一行都是本仓踩过的,量级是实测的,不是估计:
| 口径轴 | 改一下的量级 | 由谁固定 | 由谁检查 |
|---|---|---|---|
| autotune config(record vs replay) | router_gemm 差 8pp | `REPLAY_FROM` | `caliber_complaints`:两侧都 record 直接拒绝出报告 |
| timer(cuda_graph vs direct) | replication_pad3d 从 +43% 变成 +124.5% | `ZL_TIMER` / backend 能力 | 同上,`timer` 字段不一致即拒绝 |
| libtuner 持久缓存(共享 vs 隔离) | mm 看起来从 945 个变体降到 33 个——实际是 B 侧继承了 A 侧的缓存 | `ZL_ISOLATE_TUNER_CACHE=1` | 共享缓存上 record 会打 `警告 [caliber]` |
| `--level core` vs `comprehensive` | mm 的 shape 数 5 → 29,能不能覆盖到大 GEMM 全看这个 | `BENCH_LEVEL` | `bench_level` 字段不一致即拒绝 |
| 重复轮数 | 收益落在零附近时,单次测量给不出方向:topk 42 个 shape 有 17 个在两次等效运行间翻了符号 | `ROUNDS` / `ZL_DIRECT_REPEAT` | `direct_repeat` 字段不一致即拒绝 |
| 被测开关本身是否真的接上 | 开关名拼错,两侧跑同一份代码,脚本照样报出"无差异" | 开关名正确 | `ab_repeat.sh`:第 1 轮后比两侧 ttgir 指纹,相同即跳过该算子 |
**结论:一次运行的数字没有意义,只有一对同口径运行的差值有意义。** 每次运行把自己的
口径写进 `runs/<run>/run_identity.json`;`batch_summary.py --ab` / `--ab-rounds` 在出
报告前核对这些字段,不一致就拒绝,因为**口径不同的两次运行会产出一份和有效报告长得
一模一样的表**——没有任何一步会报错,只是数字换了含义。确认两次运行确实可比时,用
`--allow-caliber-mismatch` 显式放行。
文档拦不住人:README 里早就写着"B 侧必须走 REPLAY_FROM"和"只有 record 数据支撑的
结论按未验证处理",本仓仍然发布过一个 record 对 record 得到的 router_gemm 数字。
所以上表第四列是代码,不是提醒。
单算子精测流程之上另有批量筛查层(`run_batch.py`,见「批量多算子测试」),可对双仓库上千个算子做例行体检,产物结构与单算子完全同构。 单算子精测流程之上另有批量筛查层(`run_batch.py`,见「批量多算子测试」),可对双仓库上千个算子做例行体检,产物结构与单算子完全同构。
## 快速开始 ## 快速开始
@@ -44,11 +20,6 @@ OP=softmax OP_FILE=softmax SHAPE_FILE=my_shapes.yaml bash run_pytest.sh
# 测 FlagGems-vllm 仓库的算子:换 FLAGGEMS_DIR 即可(插件自动适配其包名 flaggems_vllm) # 测 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_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)。 上游现为两个仓库,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)。
@@ -64,7 +35,7 @@ softmax:
上面这条命令是 **record 模式**:现场 autotune、把选中的 config 存档。它用于单次摸底或给 replay 提供基准,**两次 record 的数字不能互相比**——做对比走下面的「A/B 对比测试标准流程」。 上面这条命令是 **record 模式**:现场 autotune、把选中的 config 存档。它用于单次摸底或给 replay 提供基准,**两次 record 的数字不能互相比**——做对比走下面的「A/B 对比测试标准流程」。
跑之前确认目标卡空闲:共享机器上别的任务会把两侧一起等比拖慢,出一份看似自洽实则作废的数据。用 `ZL_BACKEND` 选运行时(默认 `cuda`),再用该后端的可见性变量选卡(CUDA 为 `CUDA_VISIBLE_DEVICES`,也可设 `ZL_DEVICE`)。**测量始终单进程单设备串行**,这是计时可比的前提;多卡仅用于可选的 sweep 加速(见「加速:多卡并行 sweep」)。非 CUDA 后端默认走上游 `do_bench`,不会尝试 CUDA Graph。 跑之前先 `nvidia-smi` 确认目标卡空闲:共享机器上别的任务会把两侧一起等比拖慢,出一份看似自洽实则作废的数据。多卡机器用 `CUDA_VISIBLE_DEVICES=<空闲卡号>` 选卡。**测量始终单进程单卡串行**,这是计时可比的前提;多卡仅用于可选的 sweep 加速(见「加速:多卡并行 sweep」)。
脚本内固定了 `--level core --mode kernel`,其余可调项都走环境变量(见下);需要改动这两个口径本身时直接编辑 `run_pytest.sh` 中的 pytest 命令行。 脚本内固定了 `--level core --mode kernel`,其余可调项都走环境变量(见下);需要改动这两个口径本身时直接编辑 `run_pytest.sh` 中的 pytest 命令行。
@@ -74,14 +45,10 @@ softmax:
| 变量 | 默认 | 说明 | | 变量 | 默认 | 说明 |
|-----|------|------| |-----|------|------|
| `OP` | `fused_marlin_moe_mxfp4` | 测试函数名(`test_` 之后的部分);设了 `TEST_NODE` 时由节点覆盖 | | `OP` | `fused_marlin_moe_mxfp4` | 测试函数名(`test_` 之后的部分) |
| `OP_FILE` | `fused_marlin_moe` | benchmark 文件名(`test_` 与 `.py` 之间的部分);设了 `TEST_NODE` 时由节点覆盖 | | `OP_FILE` | `fused_marlin_moe` | benchmark 文件名(`test_` 与 `.py` 之间的部分) |
| `TEST_NODE` | 空 | FlagGems pytest 节点(`…/benchmark/test_<op_file>.py::test_<op>`)。设了且未指定 `SHAPE_FILE` 时默认不传 `--shape_file` | | `SHAPE_FILE` | 空(用脚本内置 yaml) | shape yaml 路径 |
| `SHAPE_FILE` | 空 | shape yaml 路径;空且无 `TEST_NODE` 时用脚本内置 yaml;`-`/`upstream`/`none` 表示沿用上游在 `BENCH_LEVEL` 下的原生 shape 集 | | `FLAGGEMS_DIR` | `/workspace/dev/FlagGems` | FlagGems 仓库路径(测 vllm 仓库时指向 `/workspace/dev/FlagGems-vllm`) |
| `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`;批量层绑卡后也会设 |
**怎么测** **怎么测**
@@ -89,16 +56,8 @@ softmax:
|-----|------|------| |-----|------|------|
| `REPLAY_FROM` | 空(record 模式) | 指向某次历史 run 目录,replay 其 autotune 选择。做 A/B 时 B 侧必须设(见「A/B 对比测试标准流程」) | | `REPLAY_FROM` | 空(record 模式) | 指向某次历史 run 目录,replay 其 autotune 选择。做 A/B 时 B 侧必须设(见「A/B 对比测试标准流程」) |
| `USE_FLAGTUNE` | `0` | `0` 走普通 autotune(跳过搜索、快速验证);`1` 用 FlagTune 扩展调优空间(首跑全量搜索、慢)。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 两侧须取同值 |
| `ZL_ISOLATE_TUNER_CACHE` | `0` | FlagGems 的 libtuner winner 存在 `$FLAGGEMS_CACHE_DIR/config_cache` 的 sqlite 里(默认 `~/.flaggems`),跨 run 长存。`1` 给本次 run 一个私有缓存目录,两侧各自冷调优(README「各自最优」口径),跑完随 Triton 缓存一起删——选中的 config 已在 `autotune_records/` 里,DB 本身无留存价值。共享缓存的风险是单向的:先跑的一侧写入 winner,后跑的一侧直接命中、只编译那几个胜出变体。实测 `router_gemm` 共享模式下第二次运行编译变体从 234 掉到 10、耗时 52s→17s,隔离模式下两次都是 234 个变体。走 `REPLAY_FROM` 的标准 A/B 不受影响(config 本就锁死),record 对 record 才会中招。取值无论如何都写进 `run_identity.json` 的 `tuner_cache` |
| `ZL_DIRECT_REPEAT` | `1` | direct 计时的重复次数,取中位。只作用于 direct 路径:整体 `timer=direct` 的运行,以及图捕获运行里逐行回退的那些点(标记会带上 `repeat=N`);捕获成功的行不受影响。`do_bench` 自己已在调用内取中位,压住的是迭代间抖动;跨调用重复补的是它看不见的那层——时钟爬坡、邻居进程进来这类漂移。调用方要 `quantiles` 或 `return_mode="all"` 时返回序列、跨调用中位无定义,自动不重复。在 MThreads 上这不是可选项:cvt 消除量最大的 `sort`/`dropout`/`rand` 恰好都无法捕获,那批结论正是由这条路径产出的。取值写进 `run_identity.json` 的 `direct_repeat` |
| `DTYPES` | 空(上游默认 dtype 扫描) | 空格分隔的 dtype 白名单(如 `"bfloat16 float16"`),逐个转为上游 `--dtypes`。算子不支持时上游报 `can't be supported by this op` | | `DTYPES` | 空(上游默认 dtype 扫描) | 空格分隔的 dtype 白名单(如 `"bfloat16 float16"`),逐个转为上游 `--dtypes`。算子不支持时上游报 `can't be supported by this op` |
| `PARALLEL_WARMUP_GPUS` | 空(串行) | 设为 `N`(≥2)时把 autotune sweep 分片到 N 卡并行,测量仍单卡串行;不设或 `<2` 则完全串行(见「加速:多卡并行 sweep」) | | `PARALLEL_WARMUP_GPUS` | 空(串行) | 设为 `N`(≥2)时把 autotune sweep 分片到 N 卡并行,测量仍单卡串行;不设或 `<2` 则完全串行(见「加速:多卡并行 sweep」) |
| `ZL_ENABLE_CUDAGRAPH` | `auto` | `auto`/`1` 在**声明了 `supports_cuda_graph` 的后端**上走 Graph 计时(目前 cuda 与 musa),其余后端即使设 `1` 也拒绝并保持 `timer=direct`。注意 `timer=cuda_graph` 只表示**尝试过**捕获,不表示有任何一行捕获成功——S5000 上 `dropout`/`rand`/`masked_select` 是逐 shape 全部回退(每行 `BENCHMARK_DIRECT_NO_CUDAGRAPH`,零毒化、不丢行),整张表实际是 direct 计时。这种情况下设 `0` 走的是同一条路径,只是记录更诚实 |
| `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` 列表的插件名 |
**输出** **输出**
@@ -113,13 +72,9 @@ softmax:
``` ```
runs/softmax_20260716-031752/ runs/softmax_20260716-031752/
├── run.log # 完整日志(含 SUCCESS 行的 latency/speedup 表;已去 ANSI 色的纯文本) ├── run.log # 完整日志(含 SUCCESS 行的 latency/speedup 表;已去 ANSI 色的纯文本)
├── run_identity.json # 本次运行的口径:provider/backend/device/timer/vendor/
│ # bench_level/shape_mode/tuner_cache/direct_repeat/autotune_mode
│ # (shell 写 stub,插件合并写入;对比工具据此判可比性)
├── shapes.yaml # 本次实际使用的 shape(存档) ├── shapes.yaml # 本次实际使用的 shape(存档)
├── autotune_records/ # Triton autotune 选中的 config(供 replay) ├── autotune_records/ # Triton autotune 选中的 config(供 replay)
│ ├── softmax.json │ └── softmax.json
│ └── meta.json # provider/backend;跨后端 replay 会被拒绝
└── ttgir/ # 按 shape 分组的 IR 落盘 └── ttgir/ # 按 shape 分组的 IR 落盘
├── 1024x1024/ ├── 1024x1024/
│ └── softmax_kernel_inner/ │ └── softmax_kernel_inner/
@@ -137,137 +92,6 @@ ttgir 关键设计:
- **命名 = 组内有区分度的参数**:同一 (shape, kernel) 组内取值相同的 constexpr 不进文件名;多词参数缩写为首字母(`BLOCK_SIZE_M` → `BSM`,图例见 naming.md);同名冲突依次用 dtype(`__fp16`)、参数对齐特化(`__EMdiv16`)、hash 前缀消歧。 - **命名 = 组内有区分度的参数**:同一 (shape, kernel) 组内取值相同的 constexpr 不进文件名;多词参数缩写为首字母(`BLOCK_SIZE_M` → `BSM`,图例见 naming.md);同名冲突依次用 dtype(`__fp16`)、参数对齐特化(`__EMdiv16`)、hash 前缀消歧。
- **崩溃也能拿到 IR**:dump 挂在进程 atexit 上,CUDA crash 后仍会落盘已编译部分。 - **崩溃也能拿到 IR**:dump 挂在进程 atexit 上,CUDA crash 后仍会落盘已编译部分。
## 多卡跑一批:campaign.sh
`ab_repeat.sh` 是单卡串行的。一批算子要并行时用 `campaign.sh` 编排——它只分活,
测量仍然全部由 `ab_repeat.sh` 做,那个脚本一行不动。
```bash
SWITCH=FLAGTREE_MUSA_RLC_ENHANCE DEVICES=0,1,2,3,4,5,6,7 SCREEN=runs/ab_rlc_XXX/screen.tsv ZL_BACKEND=musa FLAGGEMS_DIR=/path/to/FlagGems DRY_RUN=1 bash campaign.sh
```
`DRY_RUN=1` 先看方案,确认了再去掉它起跑。
- **只取筛选表里 `IR_CHANGED` 的算子**——第一步的判定在这里生效
- **按实测耗时贪心分片**(`partition.py`),不是 hash 取模。S5000 上 16 个算子实测:
hash 分片关键路径 120 分钟、有一张卡空转;贪心 70 分钟。**下界由最重的单个算子
决定**(`topk` 一个就 70.6 分钟),加卡不解决,分片只保证别的卡不空等
- **每卡独立 `AB_ROOT`**——八个 `ab_repeat.sh` 同一秒启动,默认的时间戳目录名会撞
- **强制 `ZL_ISOLATE_TUNER_CACHE=1`**。8 个进程写同一个 sqlite 是要避开的事。代价是
pin 的 config 与共享缓存批次不同,**两批的绝对延迟不可互引**——口径守卫会自己拒绝
- `EXTRA_OPS` 追加筛选表里没有的算子(如 benchmark 刚修好的),放到**负载最小**的卡上
- 跑完合并各卡 `status.tsv`(首列是卡号),并统计非零退出的格与三个标记的总数
设备可见性变量名从 backends 层取,不写死 `MUSA_VISIBLE_DEVICES`。
## 两步:先筛选,再实测
跑一个算子的六轮 A/B 要 13~44 分钟(S5000 实测:`rand` 15 shape 13.0 分钟,
`sort` 25 shape 43.4 分钟)。用 `ab_rlc_test.sh` 筛一对只要约 0.8 分钟。
**差 20~50 倍**,所以先筛后测:
```bash
# 第一步:一对定去留(A 侧 record,B 侧 REPLAY_FROM 它,config 锁死)
OPS="col2im:col2im topk:topk perf_fft:fft" VENDOR=MUSA bash ab_rlc_test.sh
# 第二步:只把 IR_CHANGED 的送进来
SWITCH=FLAGTREE_MUSA_RLC_ENHANCE ROUNDS=6 OPS="topk:topk" bash ab_repeat.sh
```
`ab_rlc_test.sh` 除了原有的 speedup 区间,还给每个算子写一行判定到
`<root>/screen.tsv`:
| 判定 | 含义 | 该怎么办 |
|---|---|---|
| `IR_CHANGED` | 两侧 ttgir 不同 | 排六轮 A/B |
| `NO_IR_DIFF` | 逐字节相同(config 已锁死) | 跳过。跑六轮只会测到噪声,并报成一个"无差异"结论 |
| `NO_TRITON` | 跑出了行,但没编出 `.ttgir` | 本次运行没有任何东西经过 Triton。算子本身可能完全正常——多半跑的是 aten 或厂商库的 kernel——只是被测的编译器 pass 没有作用对象。**既不是"开关无作用",也不是"没有 kernel"**。`run_pytest.sh` 每个 run 用独立 `TRITON_CACHE_DIR`,所以"缓存命中没重新编译"这个解释已被排除 |
| `NO_ROWS` | 没有测量行 | 先修可运行性:多半是 dtype 不被该算子接受,或它在本后端拿不到 config |
两件容易做错的:
- **判据是 ttgir 指纹,不是 `convert_layout` 计数。** cvt 数相同不等于生成代码
相同,而要测的是"生成代码变了没有"。两个方向都会骗人:`addmm` 的 cvt 是 15→15
没变,15 个 ttgir 文件却实质全不同;`tanh`/`softmax` 的 cvt 是 0→0,生成代码也
真的没动。cvt 数照样记,它解释的是变化的性质。
- **指纹先抹掉开关自己盖的注解。** FlagTree 的 `_apply_musa_rlc_policy(mod)` 在开关
打开时把整套 RLC 策略旋钮盖成 module 属性(`ttg.rlc-*`),跟这个 pass 后面有没有
真的改动什么无关。不抹的话判据有个常量偏移——**凡是能编出 kernel 的算子都必然判成
"变了",筛选器谁都不排除**。24 个算子实测:抹之前 23 个 `IR_CHANGED`,抹之后 7 个
落回 `NO_IR_DIFF`。零对照抓不到这个,因为两侧同值时属性也相同。
- **不要全局写死一个 `DTYPES`。** `OPS` 元素支持 `OP:OP_FILE:DTYPES`,留空表示用
该算子自己的全套。`fft` 只吃 complex64,被一刀切的 `DTYPES=float16` 判死过一次
——那是 harness 的问题,被记成了算子的问题。
同一份 `_ir_fingerprint.sh` 也被 `ab_repeat.sh` 用于第 1 轮后的对比轴有效性检查,
两处判据不会漂开。
判定还管住报告:`NO_ROWS` 与 `NO_TRITON` 不出 speedup 区间(前者没有可比的测量,
后者与被测开关无关);`NO_IR_DIFF` 照出,但明说那是**本机对该算子的噪声底**,不是
开关的收益——它离零多远,就是该算子单次测量的可信下限。判定刚说完这一对不可比,
下一行就给个 `+0.3%`,正是这套工具一路在防的"看起来像结论的表"。
`screen.tsv` 末七列是口径(`backend device timer bench_level shape_mode
tuner_cache direct_repeat`),由 `python batch_summary.py --caliber <run_dir>` 取——
`CALIBER_FIELDS` 只有那一处定义,筛选表不重列一遍。两份不同 `--level` 的
`screen.tsv` 放一起比时,这几列是唯一能拦住的东西。
每对有 `CELL_TIMEOUT`(默认 2400 秒)兜底,超时打 `!!! CELL_TIMEOUT` 并按耗时问题
报告,不混进 dtype 那条解释里。`elapsed_before/after` 也进表:两步流程的全部依据
就是"一对 0.8 分钟对六轮 13~44 分钟"这个比例,那这一步该记下自己那一半。
## 算子跑不起来时的排查清单
`NO_ROWS` / `NO_TRITON` 说的是"这一格没测出东西",不是"开关无作用"。按下面的顺序查,
每一条都踩过:
| # | 症状 | 查什么 | 实例 |
|---|---|---|---|
| 1 | `PREFLIGHT_FAILED check=interpreter imports triton + flag_gems` | 用的是哪个 python。triton 通常来自 PYTHONPATH 指向的编译器构建 | 一次噪声标定因此废掉 12 格 + 160 次负载 run |
| 2 | `PREFLIGHT_FAILED check=benchmark file exists` | `OPS` 元素是 `OP:OP_FILE`,写反就找不到文件 | `test_perf_sort` 在 `test_sort.py` 里,要写 `perf_sort:sort` |
| 3 | `ValueError: Given dtype(s) '...' can't be supported by this op` | 该算子只吃特定 dtype,别全局写死 `DTYPES` | `fft` 只吃 complex64,被一刀切的 fp16 判成"跑不出行",而它是 cvt 消除量最大的算子 |
| 4 | `dynamic_func() missing ... BLOCK_M, BLOCK_N ...` 之类 JIT 深处的报错 | **benchmark 是不是导入了别家厂商的实现**(见下) | `router_gemm` 拿 NVIDIA Hopper 的 kernel 跑在 MUSA 上 |
| 5 | 跑出行但 `kernels=0` | 该算子在本后端没走 Triton,走的是 aten 或厂商库 | — |
### 第 4 条:import 成功不等于这个后端支持
FlagGems 的厂商实现放在 `flag_gems/runtime/backend/_<vendor>/ops/`。benchmark 里常见
这种可用性判据:
```python
try:
from flag_gems.runtime.backend._nvidia.hopper.ops.mm import router_gemm
ROUTER_GEMM_AVAILABLE = True # ← import 成功就置 True
except Exception:
ROUTER_GEMM_AVAILABLE = False
```
**`import` 只证明那个 `.py` 在磁盘上,它不碰设备,所以在任何后端上都成功。** gate 于是
形同虚设:MUSA 上照样置 True,benchmark 拿 NVIDIA Hopper 的 kernel 跑在摩尔线程卡上,
最后死在 JIT 深处的 `dynamic_func() missing BLOCK_M, BLOCK_N, BLOCK_K, SPLIT_K`——
**那个报错完全指不向真正的原因**。而 `_mthreads/ops/mm.py` 下明明有一份能用的实现,
benchmark 从头到尾没导入过。
正确写法是先判 vendor(同一个文件里 Cambricon 那支就是对的):
```python
if base.vendor_name == "nvidia":
try:
from flag_gems.runtime.backend._nvidia.hopper.ops.mm import router_gemm
ROUTER_GEMM_AVAILABLE = True
except Exception:
pass
```
`_device_setup_plugin` 会在 session 结束时扫 `sys.modules`,载入过别家厂商实现时打:
```
FOREIGN_VENDOR_MODULE active=mthreads loaded=nvidia
```
**它不拦,只报**——载入别家模块本身不一定有问题(有些是共享工具),但一个算子结果可疑
时,这行是第一个该看的地方。
## A/B 对比测试标准流程 ## A/B 对比测试标准流程
```bash ```bash
@@ -307,104 +131,7 @@ libtuner 的持久缓存何时失效:FlagGems kernel 源码改动、tune_confi
同一层缓存也决定了 `_autotune_record_plugin` 的实现方式:record 不能靠"cache 新增了哪个 key"来判断本次选中的 config(缓存一热就走 cached 分支、不写新 key,键集差集恒为空),改为 `run()` 之后直接读回本次调用的 `self.cache[key]`;replay 同理不能加"key 不在 cache 里才注入"的前置条件,否则注入被跳过、该 run 表面在 replay 实际在自调优。想确认 record 真的覆盖到目标 kernel,查 `runs/<run>/autotune_records/<op>.json` 里有无对应 kernel 条目——漏记时该文件照样生成,只是少了 libtuner 那几个。 同一层缓存也决定了 `_autotune_record_plugin` 的实现方式:record 不能靠"cache 新增了哪个 key"来判断本次选中的 config(缓存一热就走 cached 分支、不写新 key,键集差集恒为空),改为 `run()` 之后直接读回本次调用的 `self.cache[key]`;replay 同理不能加"key 不在 cache 里才注入"的前置条件,否则注入被跳过、该 run 表面在 replay 实际在自调优。想确认 record 真的覆盖到目标 kernel,查 `runs/<run>/autotune_records/<op>.json` 里有无对应 kernel 条目——漏记时该文件照样生成,只是少了 libtuner 那几个。
### 前置检查 ### ab_fold_test.sh:A/B 测试参考示例
`run_pytest.sh` 在启动 pytest 之前校验 FlagGems 目录、benchmark 文件、`SHAPE_FILE`
与 `REPLAY_FROM`,失败时打 `!!! PREFLIGHT_FAILED check=<具名检查>` 并以 3 退出。
这些检查存在的理由是**静默失败会指向错误的方向**:路径不存在时 pytest 收集不到
benchmark 目录,注册 `--level`/`--mode` 的 `conftest.py` 就没被加载,报错变成
`unrecognized arguments: --level core --mode kernel`,看起来像参数或版本问题。
只校验节点的文件那一半;`::test_<fn>` 交给 pytest,函数不存在时它自己会清楚地报
`not found: ...`。
还校验**解释器**——`python -c "import triton, flag_gems"`。triton 通常来自 PYTHONPATH
指向的编译器构建产物、flag_gems 来自源码检出,所以一个能摸到全部正确路径的 shell
仍然可能是错的那个 Python。不查的话,失败以每格一次 `pytest exit 4`(conftest 的
ImportError)出现,多轮 campaign 会在有人读日志之前把每一格都烧掉——这不是假想,
2026-09-10 的噪声标定实验就这么废了 12 格加 160 次负载 run。报错里带上
`$(command -v python)`,因为要看的就是"用的到底是哪个解释器"。
### 口径核对:出报告前的拒绝
`batch_summary.py --ab` / `--ab-rounds` 在解析任何 latency 之前,先读参与对比的每个
run 目录的 `run_identity.json`,逐字段核对 `backend / device / timer / bench_level /
shape_mode / tuner_cache / direct_repeat`。任何一项不一致,或任何一侧缺 identity,
就打 `!!! [caliber] ...` 并以 **4** 退出,不出报告。
```
!!! [caliber] sort: timer 不一致:before=cuda_graph, after=direct
!!! [caliber] sort: 拒绝出报告。确认口径确实可比后加 --allow-caliber-mismatch
```
单独判 `autotune_mode`:一侧 record 定 config、另一侧 replay 它是正常形态;**两侧都
record 直接拒绝**——那正是本仓发布过错误 router_gemm 数字的那次形态。
不拒绝但会警告的一类:在**共享** libtuner 缓存上 record。被 pin 的 config 取决于此前
在这台机器上跑过什么,相对结论仍成立,绝对延迟与变体数不可跨机器引用。
`--allow-caliber-mismatch` 放行,仍会把问题打出来。用它之前先能说清为什么这两次运行
可比;跨口径引用别人的数字,比自己重跑一遍贵得多。
### 多轮重复:ab_repeat.sh 与 --ab-rounds
单对观测无法说明自己稳不稳。S5000 实测 `topk` 消除了 32 个 `convert_layout`,
42 个 shape 里却有 17 个在两次同等有效的运行之间**改变了 delta 的符号**——那不是
计时口径问题(图捕获与 direct 的延迟比是 1.00x),是收益本身落在零附近、单次
测量给不出方向。
```bash
SWITCH=FLAGTREE_MUSA_RLC_ENHANCE OPS="topk masked_select" ROUNDS=6 ZL_BACKEND=musa FLAGGEMS_DIR=/path/to/FlagGems bash ab_repeat.sh
```
口径:第 1 轮 A 侧 record config,此后每一次运行(含 A 侧的后续轮次)都
`REPLAY_FROM` 那一份,config 因此跨轮跨侧全部锁死;逐轮交替先后顺序抵消单向漂移。
有轮次失败的算子不出报告——缺轮的报告会让人误以为是完整的 N 轮。
每跑完一格往 `<root>/status.tsv` 追一行:`round op side switch rc elapsed_s success
failed replay_fallback no_cudagraph capture_poisoned`。README 要求必查的三个标记就在
这张表里,跑的时候就能看见它们从 0 变成非 0,不用等结束再 grep——事后 grep 只在有人
想起来的时候才发生。每格另有 `CELL_TIMEOUT`(默认 2400 秒)兜底,超时记 `rc=124` 并打
`!!! CELL_TIMEOUT`,不会让一格卡死拖住整批。
全程还采一份 `<root>/telemetry.tsv`:
```
timestamp index sm_clock_mhz mem_clock_mhz temp_c power_w util_pct
2026-09-10T07:41:12 7 1888 2500 36 477.68 87
```
列在 cuda 与 musa 上完全一致(`nvidia-smi --query-gpu` 与 `mthreads-gmi -q --json`
各自解析后归一),所以两台机器的 campaign 可以并排读。理由很直接:六轮 A/B 声称的是
0.01% 量级的单侧噪声,**降频和热漂移能移动比这大得多的延迟,而且在 run.log 里不留
任何痕迹**——报告照样干净。`TELEMETRY_INTERVAL`(默认 5 秒)设 0 关闭;采样器缺失或
没权限时 `telemetry.py` 安静退出并把原因写进 `telemetry.err`,不影响测试本身。
第 1 轮 A/B 跑完后,脚本比两侧 `ttgir/` 的指纹(相对路径 + 内容)。**逐字节相同就
跳过这个算子**,并打 `!!! SWITCH_NO_IR_DIFF op=... switch=...`:这个开关没改变该算子的
生成代码,剩下 N-1 轮测的全是噪声,最后还会报成一个"无差异"的结论。先查开关名拼写
与它在被测源码里的读取点;确属只影响运行时、不影响 IR 的开关,用
`SWITCH_ALLOW_NO_IR_DIFF=1` 放行。
报告分开给两种离散度,它们回答不同的问题:
| 指标 | 含义 |
|---|---|
| 逐轮区间 | N 轮各自 delta 的 min~max,说明**判定**稳不稳 |
| 单侧噪声 | 同一侧 N 次延迟的极差,说明**机器**吵不吵 |
| 符号不稳行数 | delta 在轮次间改变符号的 shape 数,这些行的方向不可结论 |
已有产物可以单独重出报告,目录形态是 `<root>/r<k>_<A|B>/run.log`:
```bash
python batch_summary.py --ab-rounds runs/ab_repeat_XXX/topk --label topk
```
NV 上验证过这条路径:`router_gemm` 六轮把逐轮区间收窄到 -16.6% ~ -16.2%、单侧
噪声 0.4%,把单对观测升级成了结论。
### ab_fold_test.sh / ab_rlc_test.sh:A/B 测试参考示例
`ab_fold_test.sh` 是上述 A/B 流程的现成参考:以 FlagGems `FLAGGEMS_MXFP4_FOLDSCALE`(fold_scale 优化)开关为对比对象,对 4 个 DeepSeek-V4-Flash 真实 trace shape 集(`shape/DeepSeek-V4-Flash-p{1024,4096,32768,65536}d1024.yaml`,随仓库提供)各跑一对 FOLD=0(基线,record)→ FOLD=1(replay 同 config),共 8 轮: `ab_fold_test.sh` 是上述 A/B 流程的现成参考:以 FlagGems `FLAGGEMS_MXFP4_FOLDSCALE`(fold_scale 优化)开关为对比对象,对 4 个 DeepSeek-V4-Flash 真实 trace shape 集(`shape/DeepSeek-V4-Flash-p{1024,4096,32768,65536}d1024.yaml`,随仓库提供)各跑一对 FOLD=0(基线,record)→ FOLD=1(replay 同 config),共 8 轮:
@@ -413,8 +140,6 @@ bash ab_fold_test.sh
# 汇总 log:runs/ab_fold_<时间戳>.log(已去色);快速对账:grep -E '^#####|FAILED' <log> # 汇总 log:runs/ab_fold_<时间戳>.log(已去色);快速对账:grep -E '^#####|FAILED' <log>
``` ```
`ab_rlc_test.sh` 是同一口径的另一个实例,对比轴换成 FlagTree 的 `FLAGTREE_<VENDOR>_RLC_ENHANCE`(`VENDOR` 默认 `MUSA`),算子表由 `OPS` 指定,跑完调 `batch_summary.py --ab <before> <after>` 出该算子的 speedup 区间。
可复用的是**口径**——同 config replay、cudagraph 计时、真实 trace shape 集、逐对成组——对比其它开关 / 编译器改动 / 算子,改循环变量与环境变量即可,口径部分不用动。 可复用的是**口径**——同 config replay、cudagraph 计时、真实 trace shape 集、逐对成组——对比其它开关 / 编译器改动 / 算子,改循环变量与环境变量即可,口径部分不用动。
但**它是口径示例,不是可长期直接执行的回归脚本**:对比轴是 FlagGems 侧的算子开关,会随上游演进被改名、改语义或移除,本仓库不跟随同步更新(该脚本引用的开关目前就已被上游移除)。失效后果是**静默的**:脚本照样跑完全部轮次、照样输出完整对比表,只是两侧执行同一份代码,得出"无差异"的假结论。所以套用前先确认对比轴仍有效——`grep -rn '<开关名>' $FLAGGEMS_DIR/src` 要有命中,且两侧 latency 表确有差异。 但**它是口径示例,不是可长期直接执行的回归脚本**:对比轴是 FlagGems 侧的算子开关,会随上游演进被改名、改语义或移除,本仓库不跟随同步更新(该脚本引用的开关目前就已被上游移除)。失效后果是**静默的**:脚本照样跑完全部轮次、照样输出完整对比表,只是两侧执行同一份代码,得出"无差异"的假结论。所以套用前先确认对比轴仍有效——`grep -rn '<开关名>' $FLAGGEMS_DIR/src` 要有命中,且两侧 latency 表确有差异。
@@ -427,7 +152,7 @@ shape 多时墙上时间几乎全在 autotune sweep,而不在测量——被
PARALLEL_WARMUP_GPUS=8 SHAPE_FILE=shape/DeepSeek-V4-Flash-p32768d1024.yaml bash run_pytest.sh PARALLEL_WARMUP_GPUS=8 SHAPE_FILE=shape/DeepSeek-V4-Flash-p32768d1024.yaml bash run_pytest.sh
``` ```
加速比取决于 sweep 占比:shape 多、config 空间大的算子收益最明显,可达数倍。不设该变量(或设 <2)时插件静默 no-op。**产物结构、`REPLAY_FROM` 用法、latency 表格式全不变**——合并后的 config 就写进本次 run 的 `autotune_records/<op>.json`,分片临时目录跑完即删。 加速比取决于 sweep 占比:shape 多、config 空间大的算子收益最明显,可达数倍。不设该变量(或设 <2)时插件静默 no-op,常驻 `PLUGINS` 数组即可。**产物结构、`REPLAY_FROM` 用法、latency 表格式全不变**——合并后的 config 就写进本次 run 的 `autotune_records/<op>.json`,分片临时目录跑完即删。
**测量本身不并行**:N 个进程压满同机多卡会通过功耗墙/散热耦合,单卡 latency 被邻居拖慢且不可复现。并行只用于"决定哪个 config 胜出",测量仍是单进程单卡、与不开插件走同一条路径。 **测量本身不并行**:N 个进程压满同机多卡会通过功耗墙/散热耦合,单卡 latency 被邻居拖慢且不可复现。并行只用于"决定哪个 config 胜出",测量仍是单进程单卡、与不开插件走同一条路径。
@@ -437,17 +162,7 @@ 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 WARNING 1 config key(s) got different winners across shards — parallel interference reached the sweep
``` ```
**冲突数就是这次加速的可信度指标**:0 可放心用;偏多说明 config 选择已被污染,出正式结论前不设该变量重跑一遍。 **冲突数就是这次加速的可信度指标**:0 可放心用;偏多说明 config 选择已被污染,出正式结论前不设该变量重跑一遍。其余行为(`REPLAY_FROM` 已设或可见卡不足 2 张时跳过、分片失败回退串行、轮转分片而非按块切、用子进程而非 xdist 的原因)见插件 docstring。分片绑卡遵守继承来的 `CUDA_VISIBLE_DEVICES`——用它选过卡时,分片只会落在你选的那几张上。
插件会在三种情况下**主动放弃分片、退回串行 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) ## 批量多算子测试(run_batch.py)
@@ -455,8 +170,7 @@ yaml key 的查找与上游对齐:先按 `op_name`,查不到再沿 Benchmark
```bash ```bash
# 全量筛查(双仓库全部可收集算子,排除带无条件 skip 标记的与 ops/dismiss.txt 里的) # 全量筛查(双仓库全部可收集算子,排除带无条件 skip 标记的与 ops/dismiss.txt 里的)
python run_batch.py --all --devices 0,1,2,3,4,5,6,7 --backend cuda --dtypes bfloat16 --op-timeout 1800 python run_batch.py --all --gpus 0,1,2,3,4,5,6,7 --dtypes bfloat16 --op-timeout 1800
# --gpus 仍是 --devices 的别名;两者同时给出会报错
# 指定子集:默认 main 仓,vllm 仓加前缀;同名测试函数用 @文件名 消歧 # 指定子集:默认 main 仓,vllm 仓加前缀;同名测试函数用 @文件名 消歧
python run_batch.py --ops softmax vllm:fused_marlin_moe nextafter_@nextafter_ python run_batch.py --ops softmax vllm:fused_marlin_moe nextafter_@nextafter_
@@ -481,7 +195,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` 据此复用。 **执行语义**:每算子独立子进程(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、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 的行不要当成同口径。 **产物**:`<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` 计数,解读口径时先看这列。
## 计时口径:cudagraph 的预热、回退与精度 ## 计时口径:cudagraph 的预热、回退与精度
@@ -494,29 +208,27 @@ 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 倍以上,且行间波动更明显。解读这类算子的结果时: 回退本身不影响 A/B 公平性(两侧同一算子回退行为一致),但**回退口径的测量误差更大**:do_bench 每次迭代都走完整的 Python → launch 路径,kernel 越小,launch 开销和 CPU 侧抖动在数字里占比越高——亚毫秒级 kernel 上两种口径可差 2 倍以上,且行间波动更明显。解读这类算子的结果时:
- 小 shape 行的绝对值和小幅(<10%)差异不要过度解读,优先看大 shape 行; - 小 shape 行的绝对值和小幅(<10%)差异不要过度解读,优先看大 shape 行;
- 需要更高置信度时,同一配置多跑几次取中位数,或设 `ZL_ENABLE_CUDAGRAPH=0` 统一用 `timer=direct`(消除同表混两种口径的问题)。 - 需要更高置信度时,同一配置多跑几次取中位数,或对该算子直接注释掉 `_cudagraph_plugin` 统一用 do_bench 口径(消除同表混两种口径的问题)。
回退是按测量点发生的,同一份 latency 表里可能混有两种口径的行;若不确定,先 `grep BENCHMARK_DIRECT_NO_CUDAGRAPH run.log` 确认哪些行是回退口径再下结论。 回退是按测量点发生的,同一份 latency 表里可能混有两种口径的行;若不确定,先 `grep BENCHMARK_DIRECT_NO_CUDAGRAPH run.log` 确认哪些行是回退口径再下结论。
## 插件说明 ## 插件说明
`-p` 列表由 `plugin_plan.py` 按 `ZL_BACKEND` 生成,不再在 `run_pytest.sh` 里写死。`_cudagraph_plugin` **不**出现在默认列表上:只有 CUDA 且 `_timing_plugin` 选中 Graph 时才会 import 它。 脚本通过 `-p` 加载以下插件(`run_pytest.sh` 的 `PLUGINS` 数组,可按需注释):
| 插件 | 作用 | 何时关闭 | | 插件 | 作用 | 何时关闭 |
|-----|------|---------| |-----|------|---------|
| `_device_setup_plugin` | 按后端设 `GEMS_VENDOR`(cuda→nvidia),跳过 flag_gems 启动时可能挂死的 `nvidia-smi` 探测,并写入 `run_identity.json`。旧名 `_device_guard_plugin` 仍是兼容 shim | 已设 `GEMS_VENDOR` 时不覆盖 | | `_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 |
| `_parallel_warmup_plugin` | 由 `PARALLEL_WARMUP_GPUS=N` 激活(未设则静默 no-op):autotune sweep 分片到 N 卡并行,测量仍单卡串行、产物结构不变(见「加速:多卡并行 sweep」) | 不设该变量即关闭 | | `_parallel_warmup_plugin` | 由 `PARALLEL_WARMUP_GPUS=N` 激活(未设则静默 no-op):autotune sweep 分片到 N 卡并行,测量仍单卡串行、产物结构不变(见「加速:多卡并行 sweep」) | 不设该变量即关闭 |
| `_seed_plugin` | 固定 random/numpy/torch 种子,数据相关算子(sort/topk 等)输入逐字节一致 | 不关 | | `_seed_plugin` | 固定 random/numpy/torch 种子,数据相关算子(sort/topk 等)输入逐字节一致 | 不关 |
| `_shape_inject_plugin` | 让 shape yaml 覆盖子类硬编码的 `set_shapes()` | 不关 | | `_shape_inject_plugin` | 让 shape yaml 覆盖子类硬编码的 `set_shapes()` | 不关 |
| `_shape_iter_inject_plugin` | 覆盖在 `get_input_iter` 里硬编码 shape 的类(conv/pool 等) | 不关 | | `_shape_iter_inject_plugin` | 覆盖在 `get_input_iter` 里硬编码 shape 的类(conv/pool 等) | 不关 |
| `_bespoke_shape_plugin` | 覆盖特殊输入构造的算子(upsample/flash_mla/cutlass 等,按类名注册) | 测这些算子之外可关 | | `_bespoke_shape_plugin` | 覆盖特殊输入构造的算子(upsample/flash_mla/cutlass 等,按类名注册) | 测这些算子之外可关 |
| `_autotune_record_plugin` | record/replay Triton autotune 选择;`autotune_records/meta.json` 记录 backend,跨后端 replay 拒绝 | `ZL_ENABLE_AUTOTUNE=0` | | `_autotune_record_plugin` | record/replay Triton autotune 选择(由 RECORD/REPLAY 环境变量二选一激活) | 不关 |
| `_timing_plugin` | 按后端的 `supports_cuda_graph` 选计时路径:声明了的走 Graph(内部调用 `_cudagraph_plugin`,逐 shape 回退),没声明的一开始就是上游 `do_bench`(`timer=direct`) | `ZL_ENABLE_CUDAGRAPH=0` 强制 direct | | `_cudagraph_plugin` | `do_bench` → `do_bench_cudagraph`(kernel 纯耗时;内部先做跨流同步,修过一个间歇性 illegal instruction)。**capture 前按 warmup 时间预算显式预热**(先一次丢弃跑吸收 autotune/JIT 编译,再稳态预热),消除首轮抖动。无法 graph capture 的 kernel 自动回退并打 `BENCHMARK_DIRECT_NO_CUDAGRAPH` 标记 | 需要与他人的普通 do_bench 数据对齐时注释掉 |
| `_cudagraph_plugin` | 仅由 `_timing_plugin` 在声明了 `supports_cuda_graph` 的后端上调用。`do_bench` → `do_bench_cudagraph`,capture 前按 warmup 预算显式预热;无法 capture 时回退并打 `BENCHMARK_DIRECT_NO_CUDAGRAPH` | 不要直接 `-p` 加到非 CUDA 运行 | | `_pretty_report_plugin` | 结果表整理 + 着色:全表相同的输入折叠成表头下一行图例,每行 Size Detail 只留随行变化的部分(shape 保持 `torch.Size([...])` 原样,MoE 类算子单行从 ~400 字符缩到一屏内);SUCCESS 绿 / FAILED 红。列名与 `SUCCESS` 字样保持上游原文,`run.log` 的 grep/解析不受影响 | 需要与上游原始表格逐字对齐时注释掉 |
| `_musa_benchmark_compat_plugin` | 剥掉 FlagGems 上游若干 benchmark 上过时的 `torch.cuda.is_available()` 跳过门——被测代码本就经 `flag_gems.device` 分发,在 MThreads 上那些门把可用算子变成「未测量」。每条覆盖都钉死到一个文件与一条确切的 skip 理由,且仅在 `flag_gems.vendor_name == "mthreads"` 时生效 | 由 `backends/musa.py` 的 `compat_plugins` 声明,其他后端不加载 | | `_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` |
## 常见问题 ## 常见问题
@@ -545,7 +257,5 @@ 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 补丁; - FlagGems benchmark 体系(`benchmark/base.py` 的 `Benchmark` 类、conftest 的 `--shape_file/--level/--mode` 选项);主仓与 FlagGems-vllm 仓库同构,`_autotune_record_plugin` 会按实际被 import 的包(`flag_gems` / `flaggems_vllm`)挂 LibTuner 补丁;
- 设备后端由 `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 满足);
- 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 输出的注册日志是否还正常打印。 - 插件通过 monkeypatch 挂钩上游内部结构,FlagGems/Triton 大版本升级后若行为异常,优先检查各插件 pytest_configure 输出的注册日志是否还正常打印。
- 无 GPU 的适配器自检:`python test_backends.py`。
-70
View File
@@ -214,57 +214,6 @@ def _replay_run(original):
return runner 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: def _load_replay_dir(replay_dir: str) -> int:
path = _record_path(replay_dir) path = _record_path(replay_dir)
if not path.is_file(): if not path.is_file():
@@ -307,48 +256,29 @@ def _dump_record(record_dir: str) -> None:
def pytest_configure(config): def pytest_configure(config):
global _wrap global _wrap
from _term_style import tag from _term_style import tag
from _run_identity import update_identity
record_dir = os.environ.get(_RECORD_DIR_ENV, "").strip() record_dir = os.environ.get(_RECORD_DIR_ENV, "").strip()
replay_dir = os.environ.get(_REPLAY_DIR_ENV, "").strip() replay_dir = os.environ.get(_REPLAY_DIR_ENV, "").strip()
if record_dir and replay_dir: if record_dir and replay_dir:
# Mutually exclusive: recording while replaying is pointless. Surface loudly. # Mutually exclusive: recording while replaying is pointless. Surface loudly.
print(f"{tag('[autotune-record-plugin]')} error: both RECORD_DIR and " print(f"{tag('[autotune-record-plugin]')} error: both RECORD_DIR and "
"REPLAY_DIR set; ignoring both (no-op)", file=sys.stderr, flush=True) "REPLAY_DIR set; ignoring both (no-op)", file=sys.stderr, flush=True)
update_identity(autotune_mode="uncontrolled")
return return
if not record_dir and not replay_dir: if not record_dir and not replay_dir:
update_identity(autotune_mode="uncontrolled")
return return
try:
import triton.runtime.autotuner as _autotuner 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 _wrap = _record_run if record_dir else _replay_run
_autotuner.Autotuner.run = _wrap(_autotuner.Autotuner.run) _autotuner.Autotuner.run = _wrap(_autotuner.Autotuner.run)
if record_dir: if record_dir:
# atexit (not sessionfinish): persist whatever was recorded even if an op # atexit (not sessionfinish): persist whatever was recorded even if an op
# crash kills the session; a later replay run falls back for missing keys. # crash kills the session; a later replay run falls back for missing keys.
_write_record_meta(record_dir)
atexit.register(_dump_record, record_dir) atexit.register(_dump_record, record_dir)
update_identity(autotune_mode="record")
print(f"{tag('[autotune-record-plugin]')} recording autotune configs to " print(f"{tag('[autotune-record-plugin]')} recording autotune configs to "
f"{_record_path(record_dir)} (Autotuner)", f"{_record_path(record_dir)} (Autotuner)",
file=sys.stderr, flush=True) file=sys.stderr, flush=True)
else: else:
loaded = _load_replay_dir(replay_dir) loaded = _load_replay_dir(replay_dir)
update_identity(autotune_mode="replay")
print(f"{tag('[autotune-record-plugin]')} replaying {loaded} recorded entries from " print(f"{tag('[autotune-record-plugin]')} replaying {loaded} recorded entries from "
f"{_record_path(replay_dir)} (Autotuner)", f"{_record_path(replay_dir)} (Autotuner)",
file=sys.stderr, flush=True) file=sys.stderr, flush=True)
+74 -120
View File
@@ -1,25 +1,25 @@
"""pytest plugin: swap `triton.testing.do_bench` for graph-capture timing. """pytest plugin: swap `triton.testing.do_bench` for graph-capture timing.
Invoked by `_timing_plugin` when the active backend declares Loaded via `-p _cudagraph_plugin` in run_pytest.sh's PLUGINS list (comment that
`supports_cuda_graph`. Do not put this module on the default `-p` list; the line out to time with plain do_bench). Monkey-patches do_bench →
timing plugin owns the decision. Monkey-patches do_bench →
do_bench_cudagraph (kernel-only latency) at session configure, after Triton is 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 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 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 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 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 settle autotune/JIT for the MoE kernels, which makes the first captured
measurement unstable run-to-run). The capture itself runs under a scratch measurement unstable run-to-run). The capture itself runs under a scratch RNG
RNG state so that an aborted capture cannot poison the process-wide generator state so that an aborted capture cannot poison the process-wide generator
(rationale: _capture_with_scratch_rng). (rationale: _capture_with_scratch_rng).
Stream/event/graph calls come from Triton's active device interface, not from Everything here goes through Triton's active device interface rather than
torch.cuda. FlagTree's do_bench is already written that way; its torch.cuda. FlagTree's do_bench is already written that way, but its
do_bench_cudagraph is not, so on a non-CUDA backend every capture raised and do_bench_cudagraph is not: it hardcodes torch.cuda.stream/Event/CUDAGraph, so
every row silently took the fallback while the log still claimed capture was on. on a non-CUDA backend the graph path raised and every measurement silently fell
CUDA still calls upstream's do_bench_cudagraph verbatim, so its numbers do not back to plain do_bench. On MThreads the interface is torch_musa, which supplies
move. Which failures are recoverable is likewise a backend question, because the Stream, Event, synchronize, graph and MUSAGraph, so the same capture works once
error wording is vendor text; see DeviceBackend.classify_capture_error. the calls are taken from the interface. CUDA still calls upstream's
do_bench_cudagraph, so its numbers are unchanged.
""" """
from __future__ import annotations from __future__ import annotations
@@ -33,47 +33,6 @@ import triton.testing as _tt
_ORIGINAL_DO_BENCH = _tt.do_bench _ORIGINAL_DO_BENCH = _tt.do_bench
def direct_repeat() -> int:
"""ZL_DIRECT_REPEAT, clamped to >= 1. 1 keeps today's single call."""
try:
return max(1, int(os.environ.get("ZL_DIRECT_REPEAT", "1")))
except ValueError:
return 1
def run_direct(fn, warmup, rep, grad_to_none, quantiles, return_mode):
"""Upstream do_bench, optionally repeated with the median taken.
Direct timing walks the full Python launch path once per iteration, so on
small kernels it carries launch overhead and CPU jitter that graph capture
does not. do_bench already returns a median over the iterations inside one
call, which suppresses iteration-level jitter but not drift across a call:
clocks ramping, another process arriving. Repeating the call and taking the
median over calls is the part the inner median cannot see.
It matters most where it is least optional. On MThreads the operators with
the largest RLC convert reductions -- sort, dropout, rand -- are exactly the
ones that cannot be graph-captured, so the headline numbers there are
produced by this path.
Repetition is skipped when the caller asks for quantiles or return_mode
"all", because the result is then a sequence and a median across calls is
not defined for it; the single call is returned unchanged.
"""
n = direct_repeat()
if n == 1 or quantiles is not None or return_mode == "all":
return _ORIGINAL_DO_BENCH(
fn, warmup=warmup, rep=rep, grad_to_none=grad_to_none,
quantiles=quantiles, return_mode=return_mode)
import statistics
samples = [
_ORIGINAL_DO_BENCH(fn, warmup=warmup, rep=rep, grad_to_none=grad_to_none,
quantiles=None, return_mode=return_mode)
for _ in range(n)
]
return statistics.median(samples)
def _device_interface(): def _device_interface():
"""Triton's active device interface: torch.cuda on NVIDIA, torch_musa on MThreads.""" """Triton's active device interface: torch.cuda on NVIDIA, torch_musa on MThreads."""
from triton.runtime import driver from triton.runtime import driver
@@ -99,6 +58,33 @@ def _graph_class(di):
raise RuntimeError("device interface exposes no graph type") 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.
Markers must match the error text itself, not torch's generic footer: newer
torch appends "Compile with `TORCH_USE_CUDA_DSA` to enable device-side
assertions." to EVERY AcceleratorError, so a loose "device-side assert"
marker turned all capture-incompatible ops into fatal re-raises and killed
the documented do_bench fallback path entirely."""
msg = f"{type(exc).__name__}: {exc}".lower()
markers = (
"illegal memory access",
"cudaerrorillegaladdress",
"out of memory",
"device-side assert triggered",
"an illegal instruction",
"misaligned address",
"uncorrectable ecc",
)
return any(m in msg for m in markers)
# Retained under its historical name; the marker list is backend independent.
_is_fatal_cuda_error = _is_fatal_device_error
def _warmup_before_capture(di, fn, warmup_ms, grad_to_none): def _warmup_before_capture(di, fn, warmup_ms, grad_to_none):
"""Warm fn up before graph capture. """Warm fn up before graph capture.
@@ -156,8 +142,9 @@ def _capture_with_scratch_rng(di, call):
can only land on the discarded scratch state, and _seed_plugin determinism can only land on the discarded scratch state, and _seed_plugin determinism
is left untouched. is left untouched.
An interface that exposes no default_generators has no such registration to An interface that does not expose default_generators has no such registration
protect against (torch_musa is one), so the capture simply runs.""" to protect against (torch_musa is one), so the capture simply runs.
"""
import torch import torch
generators = getattr(di, "default_generators", None) generators = getattr(di, "default_generators", None)
@@ -182,14 +169,14 @@ def _do_bench_graph(di, fn, rep, grad_to_none, quantiles, return_mode, n_retries
A step-for-step port of upstream's implementation: estimate one iteration, 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 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 out, then time independent replays. Upstream is used verbatim on CUDA; this
exists because upstream names torch.cuda directly and so cannot run anywhere exists because upstream names torch.cuda directly and so cannot run on any
else. other backend.
Replay timings go to _summarize_statistics as a plain list of floats, which The replay timings are handed to _summarize_statistics as a plain list of
is what both do_bench and upstream's do_bench_cudagraph pass: this fork's floats, which is what both do_bench and upstream's do_bench_cudagraph pass:
helper runs Python's statistics module over them, so wrapping them in a this fork's helper runs Python's statistics module over them, so wrapping
tensor makes "mean" raise and makes a quantile request return tensors that them in a tensor makes "mean" raise and makes a quantile request return
later fail to serialize. tensors that later fail to serialize.
""" """
with di.stream(di.Stream()): with di.stream(di.Stream()):
fn() fn()
@@ -232,33 +219,7 @@ def _patched_do_bench(fn, warmup=25, rep=100, grad_to_none=None,
quantiles=None, return_mode="mean"): 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 graph capture."""
op = os.environ.get("FLAGGEMS_PERF_CURRENT_OP", "<unknown>") op = os.environ.get("FLAGGEMS_PERF_CURRENT_OP", "<unknown>")
phase = "sync" # which step rejected graph capture, for the fallback marker
def _direct(note):
from _term_style import YELLOW, paint
n = direct_repeat()
# The repeat count belongs in the marker: these rows are already a
# different caliber from the captured ones, and now they can also be a
# different number of samples.
rpt = "" if n == 1 else f" repeat={n}"
print(paint(f"BENCHMARK_DIRECT_NO_CUDAGRAPH op={op} {note}{rpt}", YELLOW),
flush=True)
return run_direct(fn, warmup, rep, grad_to_none, quantiles, return_mode)
backend = None
try:
from backends import get_backend
backend = get_backend()
except Exception:
pass
# Known context-poisoning operators never get a capture attempt: the point
# of the deny list is that for these the attempt itself is the damage, so
# discovering it again per shape would cost the run, not just the precision.
if backend is not None and backend.denies_capture(op):
return _direct("phase=denylist reason=KnownPoisoningOp "
f"detail='{backend.name} denies capture for this op'")
phase = "sync" # which step rejected capture, for the marker
try: try:
di = _device_interface() di = _device_interface()
# The graph path runs fn on a fresh side stream without syncing with the # The graph path runs fn on a fresh side stream without syncing with the
@@ -288,42 +249,35 @@ def _patched_do_bench(fn, warmup=25, rep=100, grad_to_none=None,
return_mode=return_mode, return_mode=return_mode,
)) ))
except Exception as exc: except Exception as exc:
verdict = (backend.classify_capture_error(exc) # Re-raise genuine device failures; anything else falls back to plain
if backend is not None else "recoverable") # do_bench, with the failing phase (sync/warmup/capture) in the marker.
if verdict == "fatal": if _is_fatal_device_error(exc):
# The kernel failed on the device. A fallback latency here would
# report a broken kernel as a measurement.
raise raise
if verdict == "poisoned": reason = type(exc).__name__
# The attempt, not the kernel, killed the context. Falling back now
# would measure inside a dead process -- observed as a fallback
# marker followed by a teardown crash and zero rows. Surface it
# under its own marker and stop, so the operator can be re-run with
# ZL_ENABLE_CUDAGRAPH=0 and added to the backend's deny list.
from _term_style import RED, paint
detail = " ".join(str(exc).split())[:160] detail = " ".join(str(exc).split())[:160]
print(paint(f"BENCHMARK_CAPTURE_POISONED op={op} phase={phase} " from _term_style import YELLOW, paint
f"reason={type(exc).__name__} detail={detail!r}", RED), print(
flush=True) paint(f"BENCHMARK_DIRECT_NO_CUDAGRAPH op={op} phase={phase} "
raise f"reason={reason} detail={detail!r}", YELLOW),
detail = " ".join(str(exc).split())[:160] flush=True,
return _direct(f"phase={phase} reason={type(exc).__name__} " )
f"detail={detail!r}") return _ORIGINAL_DO_BENCH(
fn,
warmup=warmup,
rep=rep,
grad_to_none=grad_to_none,
quantiles=quantiles,
return_mode=return_mode,
)
def pytest_configure(config): def pytest_configure(config):
from _term_style import tag from _term_style import tag
try:
from backends import get_backend
if not get_backend().supports_cuda_graph:
print(f"{tag('[cudagraph-plugin]')} skipped: backend "
f"{get_backend().name!r} does not support CUDA Graph",
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 _tt.do_bench = _patched_do_bench
print(f"{tag('[cudagraph-plugin]')} triton.testing.do_bench → do_bench_cudagraph", try:
backend = getattr(_device_interface(), "__name__", "?")
except Exception: # pragma: no cover - reported, never fatal at configure
backend = "?"
print(f"{tag('[cudagraph-plugin]')} triton.testing.do_bench → graph capture "
f"via {backend}",
file=sys.stderr, flush=True) file=sys.stderr, flush=True)
+43 -5
View File
@@ -1,7 +1,45 @@
"""Backward-compatible alias of `_device_setup_plugin`. """Guard against FlagGems' device-detection hang (benchmark-side fix).
Kept so an old `-p _device_guard_plugin` line still sets GEMS_VENDOR. New FlagGems' `runtime.backend.device_finder` falls back to
plugin lists should load `_device_setup_plugin` instead. `subprocess.run(nvidia-smi)` **without a timeout**. Under this conda env's
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.
""" """
from _device_setup_plugin import * # noqa: F401,F403 import os
from _device_setup_plugin import pytest_configure # noqa: F401 import sys
_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()
-112
View File
@@ -1,112 +0,0 @@
"""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()
# FlagGems 的厂商实现按 flag_gems.runtime.backend._<vendor>.ops.* 组织。benchmark
# 常写成 `try: import <某厂商实现>; AVAILABLE = True`——但 import 只证明那个 .py 在
# 磁盘上,它不碰设备,所以在任何后端上都会成功。于是 gate 形同虚设,测出来的是别家
# 的 kernel 跑在本机上。router_gemm 就是这样:test_router_gemm.py 无条件导入
# _nvidia.hopper 版,在 MUSA 上 gate 照过,最后死在 JIT 深处的
# "dynamic_func() missing BLOCK_M, BLOCK_N, BLOCK_K, SPLIT_K"——那个报错完全指不
# 向真正的原因,而 _mthreads 下明明有一份能用的实现。
#
# 这里不拦,只报:跑完看一眼载入过谁家的实现。
_VENDOR_MOD_PREFIX = "flag_gems.runtime.backend._"
def foreign_vendor_modules(active_vendor: str, modules=None) -> list:
"""本次进程里载入过的、不属于当前厂商的 FlagGems 后端模块。"""
if modules is None:
modules = sys.modules
out = set()
for name in list(modules):
if not name.startswith(_VENDOR_MOD_PREFIX):
continue
# 只保留厂商根,_nvidia.hopper.ops.mm 与 _nvidia.ops.mm 算同一条。
# 按段比而不是按前缀:前缀会让 _<vendor>_something 被误认成自家。
root = name[len(_VENDOR_MOD_PREFIX):].split(".")[0]
if root != active_vendor:
out.add(root)
return sorted(out)
def pytest_sessionfinish(session, exitstatus):
from _term_style import tag
try:
vendor = get_backend().flaggems_vendor
except Exception:
return
foreign = foreign_vendor_modules(vendor)
if not foreign:
return
print(f"{tag('[device-setup-plugin]')} FOREIGN_VENDOR_MODULE active={vendor} "
f"loaded={','.join(foreign)}", file=sys.stderr, flush=True)
print(f"{tag('[device-setup-plugin]')} 本次载入了别家厂商的 FlagGems 实现。"
"如果这个算子的结果可疑,先查 benchmark 的可用性判据是不是只做了 import "
"而没判 vendor_name", file=sys.stderr, flush=True)
-31
View File
@@ -1,31 +0,0 @@
#!/bin/bash
# 由 ab_repeat.sh 与 screen.sh 共用;两处必须用同一个判据,否则筛选放行的算子
# 到了 A/B 那步又被拦下,或者反过来。
# 一个 run 生成代码的指纹:只取 .ttgir 的相对路径 + 内容,并先抹掉开关自己盖上去的
# 注解。
#
# 只认 .ttgir:dump 目录里还有 index.tsv 和 naming.md,一个 kernel 都没编出来的 run
# 只剩这两个文件、两侧完全相同,会把"什么都没生成"报成"开关没改变生成代码"。
#
# 抹掉 ttg.rlc-*:FlagTree 的 _apply_musa_rlc_policy(mod) 在开关打开时把整套 RLC 策略
# 旋钮盖成 module 属性,跟这个 pass 后面有没有真的改动什么无关。不抹的话判据有个常量
# 偏移——凡是能编出 kernel 的算子都必然判成"变了",筛选器谁都不排除。实测 24 个算子
# 里有 7 个(tanh gelu mul softmax mm bmm replication_pad3d)除这行属性外逐字节相同。
# 零对照抓不到它:两侧同值时属性也相同。
#
# 零 kernel 返回非零,交给调用方当作"没有可比的 IR",而不是"一致"。
ir_fingerprint() {
local dir="$1/ttgir"
[[ -d "$dir" ]] || return 1
local n
n=$(find "$dir" -name '*.ttgir' -type f 2>/dev/null | wc -l)
(( n > 0 )) || return 1
local out
out=$( { find "$dir" -name '*.ttgir' -type f | LC_ALL=C sort | while IFS= read -r f; do
printf '%s
' "${f#"$dir"/}"
sed -E 's/"ttg[.]rlc-[^"]*" = [0-9-]+ : i32, ?//g' "$f"
done; } | sha256sum | cut -d' ' -f1 )
[[ -n "$out" ]] && printf '%s' "$out"
}
+3 -19
View File
@@ -402,33 +402,17 @@ def _dump_ttgir(cache_dir: str, dump_dir: str) -> None:
def pytest_configure(config): def pytest_configure(config):
global _dump_enabled global _dump_enabled
from _term_style import tag
try:
import triton import triton
import triton.runtime.autotuner as _autotuner 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
try: prev = triton.knobs.compilation.listener
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: if prev is None:
compilation.listener = _compile_listener triton.knobs.compilation.listener = _compile_listener
else: else:
def chained(**kwargs): def chained(**kwargs):
prev(**kwargs) prev(**kwargs)
_compile_listener(**kwargs) _compile_listener(**kwargs)
compilation.listener = chained triton.knobs.compilation.listener = chained
dump_dir = os.environ.get(_DUMP_DIR_ENV, "").strip() dump_dir = os.environ.get(_DUMP_DIR_ENV, "").strip()
if dump_dir: if dump_dir:
-12
View File
@@ -105,18 +105,6 @@ def _warmup_load_backend():
def pytest_configure(config): def pytest_configure(config):
from _term_style import tag 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() _patch_loaded()
_warmup_load_backend() _warmup_load_backend()
_patch_loaded() _patch_loaded()
+27 -168
View File
@@ -33,7 +33,7 @@ import sys
import tempfile import tempfile
import time import time
from pathlib import Path from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple from typing import Any, Dict, List, Tuple
import pytest import pytest
@@ -60,8 +60,8 @@ def _requested_gpus() -> int:
if want < 2: if want < 2:
return 0 return 0
try: try:
from backends import get_backend import torch
have = get_backend().probe_device_count() have = torch.cuda.device_count()
except Exception: except Exception:
return 0 return 0
return max(0, min(want, have)) return max(0, min(want, have))
@@ -70,132 +70,30 @@ def _requested_gpus() -> int:
def _visible_devices() -> List[str]: def _visible_devices() -> List[str]:
"""The device ids a shard may be pinned to, in parent-visible order. """The device ids a shard may be pinned to, in parent-visible order.
Must respect inherited visibility env for the *active* backend. Writing a Must respect an inherited CUDA_VISIBLE_DEVICES: writing a bare shard index
bare shard index into the child would re-interpret it as an absolute id, into the child would re-interpret it as an absolute id, so a caller who
so a caller who picked idle cards (CUDA_VISIBLE_DEVICES=5,6) would silently picked idle cards (CUDA_VISIBLE_DEVICES=5,6) would silently get cards 0,1 —
get cards 0,1 — exactly the busy-GPU case the caller was avoiding. exactly the busy-GPU case the caller was avoiding.
""" """
from backends import get_backend raw = os.environ.get("CUDA_VISIBLE_DEVICES", "").strip()
backend = get_backend()
for key in backend.visibility_env_keys:
raw = os.environ.get(key, "").strip()
if raw: if raw:
return [d.strip() for d in raw.split(",") if d.strip()] return [d.strip() for d in raw.split(",") if d.strip()]
try: try:
n = backend.probe_device_count() import torch
return [str(i) for i in range(n)] if n else [] return [str(i) for i in range(torch.cuda.device_count())]
except Exception: except Exception:
return [] return []
def _attr_name(node: Any) -> str: def _read_shape_yaml(path: str, op: str) -> Tuple[List[Any], Dict[str, Any]]:
"""Last component of a Name/Attribute node: `base.BlasBenchmark` -> `BlasBenchmark`.""" """Shapes for this op, plus the sibling keys to carry into each shard."""
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 import yaml
with open(path) as f: with open(path) as f:
doc = yaml.safe_load(f) or {} doc = yaml.safe_load(f) or {}
keys = [op] entry = doc.get(op) or {}
if not isinstance(doc.get(op), dict) and test_py is not None: shapes = entry.get("shapes") or []
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"} extra = {k: v for k, v in entry.items() if k != "shapes"}
return list(entry["shapes"]), extra, key return list(shapes), extra
return [], {}, ""
def _shard(shapes: List[Any], n: int) -> List[List[Any]]: def _shard(shapes: List[Any], n: int) -> List[List[Any]]:
@@ -213,8 +111,8 @@ def _spawn(shards: List[List[Any]], op: str, extra: Dict[str, Any],
) -> Tuple[List[Path], List[int]]: ) -> Tuple[List[Path], List[int]]:
"""One subprocess per shard, each pinned to its own GPU. """One subprocess per shard, each pinned to its own GPU.
Subprocesses rather than pytest-xdist: `_device_setup_plugin` may initialize Subprocesses rather than pytest-xdist: _device_guard_plugin initializes a
a device context at import time, which does not survive fork-based parallelism. CUDA context at import time, which does not survive fork-based parallelism.
""" """
import yaml import yaml
procs: List[Tuple[int, subprocess.Popen, Any]] = [] procs: List[Tuple[int, subprocess.Popen, Any]] = []
@@ -245,8 +143,7 @@ def _spawn(shards: List[List[Any]], op: str, extra: Dict[str, Any],
child_argv += ["--shape_file", str(w / "shapes.yaml")] child_argv += ["--shape_file", str(w / "shapes.yaml")]
env = dict(os.environ) env = dict(os.environ)
from backends import get_backend env["CUDA_VISIBLE_DEVICES"] = devices[i]
env = get_backend().apply_visibility(env, devices[i])
env[_WORKER_ENV] = "1" # stops recursion env[_WORKER_ENV] = "1" # stops recursion
env[_RECORD_DIR_ENV] = str(rec) # shard records its own picks env[_RECORD_DIR_ENV] = str(rec) # shard records its own picks
env.pop(_REPLAY_DIR_ENV, None) # a shard must sweep env.pop(_REPLAY_DIR_ENV, None) # a shard must sweep
@@ -325,23 +222,6 @@ def pytest_configure(config):
return return
op = _op_name() 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 "" shape_file = getattr(config.option, "shape_file", "") or ""
if not shape_file or not Path(shape_file).is_file(): if not shape_file or not Path(shape_file).is_file():
print(f"{tag('[parallel-warmup-plugin]')} no --shape_file; " print(f"{tag('[parallel-warmup-plugin]')} no --shape_file; "
@@ -352,24 +232,15 @@ def pytest_configure(config):
"already pinned, nothing to sweep, skipping", flush=True) "already pinned, nothing to sweep, skipping", flush=True)
return return
test_py = _test_module_path(list(config.invocation_params.args))
try: try:
shapes, extra, key = _read_shape_yaml(shape_file, op, test_py) shapes, extra = _read_shape_yaml(shape_file, op)
except Exception as exc: except Exception as exc:
print(f"{tag('[parallel-warmup-plugin]')} cannot read shapes " print(f"{tag('[parallel-warmup-plugin]')} cannot read shapes "
f"({type(exc).__name__}: {exc}); skipping", flush=True) f"({type(exc).__name__}: {exc}); skipping", flush=True)
return 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: 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); "
f"under yaml key {key!r}; sharding would not pay off, skipping", "sharding would not pay off, skipping", flush=True)
flush=True)
return return
devices = _visible_devices() devices = _visible_devices()
@@ -380,9 +251,8 @@ def pytest_configure(config):
shards = _shard(shapes, min(gpus, len(shapes), len(devices))) shards = _shard(shapes, min(gpus, len(shapes), len(devices)))
scratch = Path(tempfile.mkdtemp(prefix=f"warmup_{op}_")) scratch = Path(tempfile.mkdtemp(prefix=f"warmup_{op}_"))
print(f"{tag('[parallel-warmup-plugin]')} sweeping {len(shapes)} shapes " print(f"{tag('[parallel-warmup-plugin]')} sweeping {len(shapes)} shapes on "
f"(yaml key {key!r}) on {len(shards)} GPUs " f"{len(shards)} GPUs (shards: {[len(s) for s in shards]}); this phase's "
f"(shards: {[len(s) for s in shards]}); this phase's "
"latency is discarded, only configs are kept", flush=True) "latency is discarded, only configs are kept", flush=True)
keep_scratch = False keep_scratch = False
@@ -390,22 +260,8 @@ def pytest_configure(config):
try: try:
recs, bad = _spawn(shards, op, extra, scratch, recs, bad = _spawn(shards, op, extra, scratch,
list(config.invocation_params.args), devices) list(config.invocation_params.args), devices)
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) merged, total, conflicts = _merge(recs, op)
elapsed = time.time() - t0
if not total: if not total:
keep_scratch = True keep_scratch = True
@@ -431,6 +287,9 @@ def pytest_configure(config):
msg = (f"{tag('[parallel-warmup-plugin]')} sweep done in {elapsed:.0f}s: " msg = (f"{tag('[parallel-warmup-plugin]')} sweep done in {elapsed:.0f}s: "
f"{len(merged)} kernels / {total} configs -> {out / f'{op}.json'}; " f"{len(merged)} kernels / {total} configs -> {out / f'{op}.json'}; "
"measurement continues serially on one GPU") "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) print(msg, flush=True)
if conflicts: if conflicts:
print(f"{tag('[parallel-warmup-plugin]')} WARNING {conflicts} config " print(f"{tag('[parallel-warmup-plugin]')} WARNING {conflicts} config "
-96
View File
@@ -1,96 +0,0 @@
"""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"))
# "isolated" or "shared(<path>)". Recorded because a shared libtuner cache
# lets whichever side runs first hand its winners to the other, and after
# the fact the two runs are otherwise indistinguishable.
p.add_argument("--tuner-cache", default="")
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,
tuner_cache=args.tuner_cache,
)
return 0
if __name__ == "__main__":
raise SystemExit(main())
-70
View File
@@ -1,70 +0,0 @@
"""Select the timing path from backend capabilities, before any shape runs.
Whether capture is attempted comes from the backend's `supports_cuda_graph`,
not from "is this CUDA". CUDA and MUSA both declare it, so both get graph
capture with per-shape fallback to plain do_bench; a backend that declares
nothing leaves `triton.testing.do_bench` in place and records `timer=direct`.
`timer=cuda_graph` therefore means capture was *attempted*, not that any row
was captured. On S5000 every shape of dropout/rand/masked_select falls back
(BENCHMARK_DIRECT_NO_CUDAGRAPH on every row, nothing poisoned, no rows lost),
so the whole table is direct-timed under a cuda_graph label. Setting
ZL_ENABLE_CUDAGRAPH=0 there selects the same timing path and says so plainly.
"""
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)
# Recorded on this branch too: the fallback rows a captured run produces
# are timed by the direct path and take the same repeat count.
update_identity(timer="cuda_graph", graph_capture="attempted",
direct_repeat=_cg.direct_repeat())
print(f"{tag('[timing-plugin]')} timer=cuda_graph backend={backend.name}",
file=sys.stderr, flush=True)
return
# A wholly-direct run gets the same repetition the fallback rows get, so the
# "put everything on one timer" advice does not silently cost the samples.
# Patching only when repeat > 1 keeps the default path untouched: upstream's
# do_bench stays in place exactly as before.
import _cudagraph_plugin as _cg
repeat = _cg.direct_repeat()
if repeat > 1:
def _repeated(fn, warmup=25, rep=100, grad_to_none=None,
quantiles=None, return_mode="mean"):
return _cg.run_direct(fn, warmup, rep, grad_to_none,
quantiles, return_mode)
import triton.testing as _tt
_tt.do_bench = _repeated
update_identity(timer="direct", graph_capture="disabled",
direct_repeat=repeat)
print(f"{tag('[timing-plugin]')} timer=direct backend={backend.name} "
f"(upstream do_bench; graph capture not selected)"
+ (f"; median of {repeat} calls" if repeat > 1 else ""),
file=sys.stderr, flush=True)
-168
View File
@@ -1,168 +0,0 @@
#!/bin/bash
# 多轮 A/B:同一对比在 N 轮独立进程里各跑一次,逐 shape 取轮次中位。
#
# 为什么需要它:单对观测无法说明自己稳不稳。S5000 实测 topk 消除了 32 个
# convert_layout,42 个 shape 里却有 17 个在两次同等有效的运行之间改变了 delta
# 的符号——那不是计时口径的问题(图捕获与 direct 的延迟比是 1.00x),是收益本身
# 落在零附近、单次测量给不出方向。多轮之后才谈得上"这个判定成立"。
#
# 口径:第 1 轮 A 侧 record autotune config,此后每一次运行(含 A 侧自己的后续
# 轮次)都 REPLAY_FROM 那一份。config 因此跨轮跨侧全部锁死,剩下的变量只有被测
# 开关与测量噪声。逐轮交替先后顺序,抵消单向漂移。
#
# 用法:
# SWITCH=FLAGTREE_MUSA_RLC_ENHANCE OPS="topk masked_select" ROUNDS=6 \
# ZL_BACKEND=musa FLAGGEMS_DIR=/path/to/FlagGems bash ab_repeat.sh
#
# OPS 元素为 OP 或 OP:OP_FILE(省略时 OP_FILE=OP)。
# SWITCH 是被测环境变量,A 侧设 0、B 侧设 1。
#
# 对比轴有效性由脚本自己查:第 1 轮 A/B 跑完后比两侧的 ttgir 指纹,完全一致说明
# 这个开关在这个算子上没改变任何生成代码——继续跑满 N 轮只会得到一个"无差异"的假
# 结论(同 ab_fold_test.sh 的告诫)。这类算子默认跳过;开关确实只影响运行时而不影响
# IR 时,用 SWITCH_ALLOW_NO_IR_DIFF=1 放行。
#
# 产物:runs/ab_repeat_<开关>_<时间戳>/<op>/r<k>_<A|B>/,跑完对每个算子调用
# batch_summary.py --ab-rounds 出报告。
set -uo pipefail
cd "$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
SWITCH_ALLOW_NO_IR_DIFF="${SWITCH_ALLOW_NO_IR_DIFF:-0}"
source "$(dirname "${BASH_SOURCE[0]}")/_ir_fingerprint.sh"
: "${SWITCH:?SWITCH is required, e.g. FLAGTREE_MUSA_RLC_ENHANCE}"
OPS="${OPS:-}"
: "${OPS:?OPS is required, e.g. \"topk masked_select\"}"
ROUNDS="${ROUNDS:-6}"
CELL_TIMEOUT="${CELL_TIMEOUT:-2400}"
if [[ -t 1 && -z "${NO_COLOR:-}" ]]; then
export FLAGGEMS_PERF_COLOR=always
fi
# AB_ROOT 供 campaign.sh 指定每卡独立的产物根:八个 ab_repeat.sh 同一秒启动时,
# 默认的时间戳目录名会撞在一起,两张卡往同一个目录里写。
ROOT="${AB_ROOT:-runs/ab_repeat_${SWITCH}_$(date +%Y%m%d-%H%M%S)}"
mkdir -p "$ROOT"
LOG="$ROOT/summary.log"
STATUS="$ROOT/status.tsv"
# 全程采时钟/温度/功耗。六轮 A/B 声称的是 0.01% 量级的单侧噪声,降频和热漂移
# 能移动比这大得多的延迟,而且在 run.log 里不留痕迹——报告照样干净。
# TELEMETRY_INTERVAL=0 关闭。采不到时 telemetry.py 自己安静退出,不影响测试。
TELEMETRY_INTERVAL="${TELEMETRY_INTERVAL:-5}"
TELE_PID=""
if [[ "$TELEMETRY_INTERVAL" != "0" ]]; then
python telemetry.py --backend "${ZL_BACKEND:-cuda}" --device "${ZL_DEVICE:-}" \
--interval "$TELEMETRY_INTERVAL" --run \
> "$ROOT/telemetry.tsv" 2> "$ROOT/telemetry.err" &
TELE_PID=$!
trap '[[ -n "$TELE_PID" ]] && kill "$TELE_PID" 2>/dev/null' EXIT
fi
printf 'round\top\tside\tswitch\trc\telapsed_s\tsuccess\tfailed\treplay_fallback\tno_cudagraph\tcapture_poisoned\n' > "$STATUS"
# 跑一格并逐格登记。事后 grep 也能拿到这些数,但只在有人想起来时才发生;
# 缺格、超时、标记从 0 变成非 0,都要在跑的时候就看得见。
run_cell() { # round side op op_file out_dir [replay_from]
local round="$1" side="$2" op="$3" op_file="$4" dir="$5" replay="${6:-}"
local val=0; [[ "$side" == B ]] && val=1
local start rc elapsed log ok fail fb ng po
start=$(date +%s)
if [[ -n "$replay" ]]; then
OP="$op" OP_FILE="$op_file" OUT_DIR="$dir" REPLAY_FROM="$replay" \
timeout -k 30 "$CELL_TIMEOUT" env "$SWITCH=$val" bash run_pytest.sh
else
OP="$op" OP_FILE="$op_file" OUT_DIR="$dir" \
timeout -k 30 "$CELL_TIMEOUT" env "$SWITCH=$val" bash run_pytest.sh
fi
rc=$?
elapsed=$(( $(date +%s) - start ))
(( rc == 124 || rc == 137 )) && echo "!!! CELL_TIMEOUT ${CELL_TIMEOUT}s op=$op round=$round side=$side"
log="$dir/run.log"
# grep -c 没命中时打印 0 但以 1 退出;不吞掉的话会被当成本格失败
ok=$(grep -c '^SUCCESS' "$log" 2>/dev/null || true)
fail=$(grep -c '^FAILED' "$log" 2>/dev/null || true)
fb=$(grep -c 'AUTOTUNE_REPLAY_FALLBACK' "$log" 2>/dev/null || true)
ng=$(grep -c 'BENCHMARK_DIRECT_NO_CUDAGRAPH' "$log" 2>/dev/null || true)
po=$(grep -c 'BENCHMARK_CAPTURE_POISONED' "$log" 2>/dev/null || true)
printf '%s\t%s\t%s\t%s\t%s\t%s\t%s\t%s\t%s\t%s\t%s\n' \
"$round" "$op" "$side" "$val" "$rc" "$elapsed" \
"${ok:-0}" "${fail:-0}" "${fb:-0}" "${ng:-0}" "${po:-0}" >> "$STATUS"
return $rc
}
{
echo ">>> switch=$SWITCH rounds=$ROUNDS ops=$OPS"
for entry in $OPS; do
op="${entry%%:*}"
op_file="${entry#*:}"
[[ "$op_file" == "$entry" ]] && op_file="$op"
out="$ROOT/$op"
mkdir -p "$out"
# 第 1 轮 A 侧 record,其余全部回放它,config 跨轮跨侧一致
echo "########## $op | round 1 | A ($SWITCH=0, record) ##########"
if ! run_cell 1 A "$op" "$op_file" "$out/r1_A"; then
echo "!!! FAILED $op round 1 A; skipping this operator"
continue
fi
REC="$out/r1_A"
ok=1; reason="有失败轮次"
for ((k = 1; k <= ROUNDS; k++)); do
# 奇数轮先 A 后 B,偶数轮反过来
if (( k % 2 == 1 )); then order="A B"; else order="B A"; fi
for side in $order; do
[[ $k -eq 1 && $side == A ]] && continue # 已经跑过
val=$([[ $side == A ]] && echo 0 || echo 1)
echo "########## $op | round $k | $side ($SWITCH=$val, replay) ##########"
if ! run_cell "$k" "$side" "$op" "$op_file" "$out/r${k}_${side}" "$REC"; then
echo "!!! FAILED $op round $k $side"
ok=0
fi
done
if (( k == 1 && ok )); then
fa=$(ir_fingerprint "$out/r1_A") || fa=""
fb=$(ir_fingerprint "$out/r1_B") || fb=""
if [[ -z "$fa" || -z "$fb" ]]; then
echo "??? $op: 没有编出任何 .ttgir(可能该算子没走 Triton,或本轮失败),"
echo "??? 跳过对比轴有效性检查——这不等于开关无作用"
elif [[ "$fa" == "$fb" ]]; then
echo "!!! SWITCH_NO_IR_DIFF op=$op switch=$SWITCH"
echo "!!! 两侧 ttgir 逐字节相同:这个开关没有改变 $op 的生成代码。"
echo "!!! 继续跑 $ROUNDS 轮只会测出噪声,并报成一个'无差异'结论。"
echo "!!! 先确认开关名拼写与它在被测源码里的读取点;确属只影响运行时的"
echo "!!! 开关,用 SWITCH_ALLOW_NO_IR_DIFF=1 重跑。"
if [[ "$SWITCH_ALLOW_NO_IR_DIFF" != "1" ]]; then
ok=0; reason="开关未改变该算子的生成代码"
break
fi
else
echo ">>> 对比轴有效:$op 两侧 ttgir 不同(A=${fa:0:12} B=${fb:0:12})"
fi
fi
done
echo "########## $op | report ##########"
if (( ok )); then
python batch_summary.py --ab-rounds "$out" --label "$op" \
|| echo "!!! report failed: $op"
else
# 缺轮的算子照样出报告会让人误以为是完整的 N 轮
echo "!!! $op $reason,不出报告;产物在 $out"
fi
done
echo "########## ALL DONE ##########"
} 2>&1 | tee "$LOG"
sed -i -E $'s/\x1b\\[[0-9;]*[A-Za-z]//g' "$LOG"
echo ">>> 总 log: $(pwd)/$LOG"
echo ">>> 逐格状态表: $(pwd)/$STATUS"
[[ -s "$ROOT/telemetry.tsv" ]] && echo ">>> 时钟/温度: $(pwd)/$ROOT/telemetry.tsv"
[[ -s "$ROOT/telemetry.err" ]] && echo ">>> 遥测提示: $(cat "$ROOT/telemetry.err")"
echo ">>> 逐算子重出报告: python batch_summary.py --ab-rounds $ROOT/<op> --label <op>"
echo ">>> 解读要点: 逐轮区间宽说明判定不稳;单侧噪声大说明机器不稳;两者含义不同"
+18 -148
View File
@@ -10,79 +10,15 @@
# #
# 跑完用 batch_summary.py --ab <before> <after> 得到该算子的 speedup 区间。 # 跑完用 batch_summary.py --ab <before> <after> 得到该算子的 speedup 区间。
# #
# 这一轮同时充当 ab_repeat.sh 的筛选前置:一对的代价约 0.8 分钟,六轮 A/B 是
# 13~44 分钟(S5000 实测),差 20~50 倍。所以先用它过一遍,只把 IR_CHANGED 的
# 算子送进 ab_repeat.sh。判定写在 <root>/screen.tsv,四种之一:
#
# IR_CHANGED 两侧 ttgir 不同 -> 排六轮
# NO_IR_DIFF 逐字节相同(config 已锁死)-> 跳过,跑六轮只会测到噪声
# NO_TRITON 跑出行但没编出 .ttgir -> 该算子在本后端没走 Triton,
# 编译器开关无从谈起。不是"开关无作用"
# NO_ROWS 没有测量行 -> 先修可运行性,多半是 dtype
#
# 判据是 ttgir 指纹而不是 convert_layout 计数:cvt 数相同不等于生成代码相同。
# cvt 数照样记,它解释的是变化的性质。
#
# 用法: # 用法:
# OPS="softmax gelu_backward" bash ab_rlc_test.sh # OPS="softmax gelu_backward" bash ab_rlc_test.sh
# OPS="sort:perf_sort topk" bash ab_rlc_test.sh # OP:OP_FILE # OPS="sort:perf_sort topk" bash ab_rlc_test.sh # 冒号分隔 OP:OP_FILE
# OPS="perf_fft:fft:complex64 topk:topk" bash ab_rlc_test.sh # 再加 DTYPES
# #
# DTYPES 留空表示用该算子自己的全套。**不要全局写死一个 dtype**:fft 只吃 # 注意:RLC 开关名带厂商前缀,用 VENDOR 指定(默认 MUSA)。套用前先确认开关有效:
# complex64,被一刀切的 DTYPES=float16 判死过一次——那是 harness 的问题, # grep 该开关名于 FlagTree 源码应有命中,否则两侧跑的是同一份代码。
# 却被记成了算子的问题。
#
# 注意:RLC 开关名带厂商前缀,用 VENDOR 指定(默认 MUSA)。开关名拼错时两侧跑
# 同一份代码,此时判定会是 NO_IR_DIFF——先查拼写,再下"开关无作用"的结论。
set -uo pipefail set -uo pipefail
cd "$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" cd "$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
source "./_ir_fingerprint.sh"
# RESCORE=<已有的 root>:判据改了之后,用当前的 ir_fingerprint 对盘上已有的 ttgir
# 重新打分并改写该 root 的 screen.tsv,不重跑任何测量。
#
# 需要它是因为判据确实改过:_ir_fingerprint.sh 起初把 index.tsv 也算进哈希,后来又
# 发现要抹掉开关自己盖的 ttg.rlc-* 属性。同一批 24 个算子,结论从 23 个 IR_CHANGED
# 变成 16 个。重打分是秒级的,重跑那批筛选要 79 分钟。
#
# 原表存为 screen.tsv.before-rescore,判定有变化的算子逐个打出来。
RESCORE="${RESCORE:-}"
if [[ -n "$RESCORE" ]]; then
[[ -d "$RESCORE" ]] || { echo "!!! 目录不存在: $RESCORE" >&2; exit 2; }
TSV="$RESCORE/screen.tsv"
[[ -f "$TSV" ]] || { echo "!!! 没有 screen.tsv: $TSV" >&2; exit 2; }
cp "$TSV" "$TSV.before-rescore"
hdr=$(head -1 "$TSV.before-rescore")
# 列名 -> 序号,不写死列序
idx() { echo "$hdr" | tr '\t' '\n' | grep -nxF "$1" | cut -d: -f1; }
c_op=$(idx op); c_rb=$(idx rows_before); c_ra=$(idx rows_after); c_v=$(idx verdict)
[[ -n "$c_op" && -n "$c_v" ]] || { echo "!!! screen.tsv 缺 op/verdict 列" >&2; exit 2; }
printf '%s\n' "$hdr" > "$TSV"
changed=0
while IFS= read -r line; do
op=$(printf '%s' "$line" | cut -f"$c_op")
old=$(printf '%s' "$line" | cut -f"$c_v")
rb=$(printf '%s' "$line" | cut -f"${c_rb:-0}"); ra=$(printf '%s' "$line" | cut -f"${c_ra:-0}")
fa=$(ir_fingerprint "$RESCORE/${op}_before") || fa=""
fb=$(ir_fingerprint "$RESCORE/${op}_after") || fb=""
if [[ "${rb:-0}" == 0 || "${ra:-0}" == 0 ]]; then v=NO_ROWS
elif [[ -z "$fa" || -z "$fb" ]]; then v=NO_TRITON
elif [[ "$fa" == "$fb" ]]; then v=NO_IR_DIFF
else v=IR_CHANGED; fi
printf '%s\n' "$line" | awk -v FS='\t' -v OFS='\t' -v i="$c_v" -v v="$v" '{$i=v; print}' >> "$TSV"
if [[ "$old" != "$v" ]]; then
printf ' %-22s %-12s -> %s\n' "$op" "$old" "$v" >&2
changed=$((changed+1))
fi
done < <(tail -n +2 "$TSV.before-rescore")
echo ">>> 重打分完成,$changed 个算子的判定变了;原表存为 $TSV.before-rescore"
awk -v i="$c_v" -F'\t' 'NR>1{c[$i]++} END{for (k in c) printf(" %-12s %d\n", k, c[k])}' "$TSV"
exit 0
fi
VENDOR="${VENDOR:-MUSA}" VENDOR="${VENDOR:-MUSA}"
ENHANCE_VAR="FLAGTREE_${VENDOR}_RLC_ENHANCE" ENHANCE_VAR="FLAGTREE_${VENDOR}_RLC_ENHANCE"
@@ -90,103 +26,38 @@ ENHANCE_VAR="FLAGTREE_${VENDOR}_RLC_ENHANCE"
# 空格分隔的算子表。元素为 OP 或 OP:OP_FILE(省略时 OP_FILE=OP)。 # 空格分隔的算子表。元素为 OP 或 OP:OP_FILE(省略时 OP_FILE=OP)。
OPS="${OPS:-softmax}" OPS="${OPS:-softmax}"
# 筛选是第一次接触一个算子的地方,最容易撞上挂死的那个。没有上限时一个算子
# 能把整批停在那里,而筛选的全部意义就是快。
CELL_TIMEOUT="${CELL_TIMEOUT:-2400}"
if [[ -t 1 && -z "${NO_COLOR:-}" ]]; then if [[ -t 1 && -z "${NO_COLOR:-}" ]]; then
export FLAGGEMS_PERF_COLOR=always export FLAGGEMS_PERF_COLOR=always
fi fi
# 显式 OUT_DIR,不再用 `ls -dt runs/<op>_*` 猜"最新那个就是我刚跑的"。串行下 LOG="runs/ab_rlc_$(date +%Y%m%d-%H%M%S).log"
# 那个假设成立,多卡并行共用 runs/ 就不成立了。 mkdir -p runs
ROOT="runs/ab_rlc_$(date +%Y%m%d-%H%M%S)"
mkdir -p "$ROOT"
LOG="$ROOT/summary.log"
TSV="$ROOT/screen.tsv"
printf 'op op_file dtypes rows_before rows_after kernels_before kernels_after cvt_before cvt_after elapsed_before elapsed_after verdict backend device timer bench_level shape_mode tuner_cache direct_repeat\n' > "$TSV"
count_cvt() { find "$1/ttgir" -name '*.ttgir' -type f -exec cat {} + 2>/dev/null | grep -c convert_layout || true; }
count_kernels() { find "$1/ttgir" -name '*.ttgir' -type f 2>/dev/null | wc -l; }
count_rows() { grep -c '^SUCCESS' "$1/run.log" 2>/dev/null || true; }
{ {
echo ">>> vendor=$VENDOR switch=$ENHANCE_VAR ops=$OPS" echo ">>> vendor=$VENDOR switch=$ENHANCE_VAR ops=$OPS"
for entry in $OPS; do for entry in $OPS; do
op="${entry%%:*}"; rest="${entry#*:}" op="${entry%%:*}"
if [[ "$rest" == "$entry" ]]; then op_file="$op"; dtypes="" op_file="${entry#*:}"
else [[ "$op_file" == "$entry" ]] && op_file="$op"
op_file="${rest%%:*}"; d="${rest#*:}"
[[ "$d" == "$rest" ]] && dtypes="" || dtypes="$d"
fi
BEFORE="$ROOT/${op}_before"; AFTER="$ROOT/${op}_after"
echo "########## $op | $ENHANCE_VAR=0 before (record) ##########" echo "########## $op | $ENHANCE_VAR=0 before (record) ##########"
t0=$(date +%s) if ! OP="$op" OP_FILE="$op_file" env "$ENHANCE_VAR=0" bash run_pytest.sh; then
OP="$op" OP_FILE="$op_file" OUT_DIR="$BEFORE" DTYPES="$dtypes" \ echo "!!! FAILED before: $op"
timeout -k 30 "$CELL_TIMEOUT" env "$ENHANCE_VAR=0" bash run_pytest.sh continue
rc_before=$?; el_before=$(( $(date +%s) - t0 )) fi
BEFORE=$(ls -dt "runs/${op}_"* | head -1)
echo "########## $op | $ENHANCE_VAR=1 after (replay $BEFORE) ##########" echo "########## $op | $ENHANCE_VAR=1 after (replay $BEFORE) ##########"
t0=$(date +%s) if ! OP="$op" OP_FILE="$op_file" REPLAY_FROM="$BEFORE" env "$ENHANCE_VAR=1" bash run_pytest.sh; then
OP="$op" OP_FILE="$op_file" OUT_DIR="$AFTER" DTYPES="$dtypes" \ echo "!!! FAILED after: $op"
REPLAY_FROM="$BEFORE" timeout -k 30 "$CELL_TIMEOUT" env "$ENHANCE_VAR=1" bash run_pytest.sh continue
rc_after=$?; el_after=$(( $(date +%s) - t0 ))
for rc in "$rc_before" "$rc_after"; do
(( rc == 124 || rc == 137 )) && echo "!!! CELL_TIMEOUT ${CELL_TIMEOUT}s op=$op"
done
rb=$(count_rows "$BEFORE"); ra=$(count_rows "$AFTER")
kb=$(count_kernels "$BEFORE"); ka=$(count_kernels "$AFTER")
cb=$(count_cvt "$BEFORE"); ca=$(count_cvt "$AFTER")
fb=$(ir_fingerprint "$BEFORE") || fb=""
fa=$(ir_fingerprint "$AFTER") || fa=""
if (( rc_before != 0 || rc_after != 0 )) || [[ "${rb:-0}" == 0 || "${ra:-0}" == 0 ]]; then
verdict=NO_ROWS
if (( rc_before == 124 || rc_after == 124 || rc_before == 137 || rc_after == 137 )); then
note="超过 CELL_TIMEOUT=${CELL_TIMEOUT}s 被截断(rc=$rc_before/$rc_after)。这是耗时问题,不是对开关的判定——先看它为什么这么慢"
else
note="跑不出测量行(rc=$rc_before/$rc_after)。看 run.log——多半是 dtype 不被该算子接受,或它在本后端拿不到 config"
fi fi
elif [[ -z "$fb" || -z "$fa" ]]; then AFTER=$(ls -dt "runs/${op}_"* | head -1)
verdict=NO_TRITON
note="跑出了行但没编出 .ttgir:本次运行没有任何东西经过 Triton。算子本身可能完全正常(多半跑的是 aten 或厂商库),只是被测的编译器 pass 没有作用对象。不是'开关无作用',也不是'没有 kernel'"
elif [[ "$fb" == "$fa" ]]; then
verdict=NO_IR_DIFF
note="两侧 ttgir 逐字节相同(config 已锁死):开关没改变生成代码。先查开关名拼写,再下结论"
else
verdict=IR_CHANGED
note="生成代码有变(cvt ${cb}→${ca},指纹 ${fb:0:8}/${fa:0:8}):值得排六轮 A/B"
fi
# 口径从 batch_summary.py 取,那里是 CALIBER_FIELDS 的唯一定义。两份不同
# --level 的 screen.tsv 放一起比时,这几列是唯一能拦住的东西。
cal=$(python batch_summary.py --caliber "$BEFORE" 2>/dev/null)
printf '%s\t%s\t%s\t%s\t%s\t%s\t%s\t%s\t%s\t%s\t%s\t%s\t%s\n' \
"$op" "$op_file" "${dtypes:-<all>}" "${rb:-0}" "${ra:-0}" \
"$kb" "$ka" "${cb:-0}" "${ca:-0}" "$el_before" "$el_after" "$verdict" \
"$cal" >> "$TSV"
echo ">>> $op: $verdict — $note"
# 判定说这一对不可比,就别在下一行给出一个看起来像结论的区间
case "$verdict" in
NO_ROWS)
echo "!!! $op 没有可对比的测量,不出 speedup 区间"; continue ;;
NO_TRITON)
echo "!!! $op 没有经过 Triton,speedup 区间与被测开关无关,不出"; continue ;;
NO_IR_DIFF)
echo ">>> $op 两侧生成代码相同:下面这个区间是本机对该算子的噪声底,"
echo ">>> 不是开关的收益。它离零多远,就是该算子单次测量的可信下限。" ;;
esac
echo "########## $op | summary ##########" echo "########## $op | summary ##########"
python batch_summary.py --ab "$BEFORE/run.log" "$AFTER/run.log" --label "$op" \ python batch_summary.py --ab "$BEFORE/run.log" "$AFTER/run.log" --label "$op" \
|| echo "!!! summary failed: $op" || echo "!!! summary failed: $op"
done done
echo "########## 筛选汇总 ##########"
awk -F'\t' 'NR>1{printf " %-12s %-22s cvt %s->%s kernels %s/%s rows %s/%s %ss\n", \
$12,$1,$8,$9,$6,$7,$4,$5,$10+$11}' "$TSV" | sort
echo ">>> 只把 IR_CHANGED 的算子送进 ab_repeat.sh。"
echo "########## ALL DONE ##########" echo "########## ALL DONE ##########"
} 2>&1 | tee "$LOG" } 2>&1 | tee "$LOG"
@@ -194,5 +65,4 @@ count_rows() { grep -c '^SUCCESS' "$1/run.log" 2>/dev/null || true; }
sed -i -E $'s/\x1b\\[[0-9;]*[A-Za-z]//g' "$LOG" sed -i -E $'s/\x1b\\[[0-9;]*[A-Za-z]//g' "$LOG"
echo ">>> 总 log: $(pwd)/$LOG" echo ">>> 总 log: $(pwd)/$LOG"
echo ">>> 筛选表: $(pwd)/$TSV" echo ">>> 每算子两份 run.log 在 runs/<op>_<时间戳>/ 下,按时间先后即 before / after"
echo ">>> 每算子两份 run.log 在 $ROOT/<op>_{before,after}/ 下"
-76
View File
@@ -1,76 +0,0 @@
"""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",
]
-17
View File
@@ -1,17 +0,0 @@
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
-221
View File
@@ -1,221 +0,0 @@
"""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
# --- graph-capture failure classification -----------------------------
# A kernel that refuses capture is normal and falls back. A device error is
# not, and the two are told apart by vendor error text, so the wording lives
# on the backend rather than in the timing plugin.
#
# fatal : the kernel itself failed on the device. Re-raise; recording a
# fallback latency here would report a broken kernel as a result.
# poisoning: the capture ATTEMPT left the context unusable. Falling back runs
# the measurement inside a dead process, which is how an operator
# ends up emitting a fallback marker and then still dying at
# teardown with zero rows. Re-raise too, but under its own marker,
# because the operator is capture-incompatible rather than broken:
# the actionable response is to re-run it with ZL_ENABLE_CUDAGRAPH=0.
fatal_error_markers: Tuple[str, ...] = ()
poisoning_error_markers: Tuple[str, ...] = ()
# Operators whose capture attempt is known to poison the context on this
# backend. Skipped up front so one of them cannot cost a whole run.
capture_deny_ops: Tuple[str, ...] = ()
# One-shot device telemetry sampler. Empty means this backend has none, and
# the campaign runs without a clock/thermal record rather than pretending.
# "{dev}" is dropped from the argv when no device is pinned, so the sampler
# covers every visible device instead of guessing an index.
telemetry_argv: Tuple[str, ...] = ()
def telemetry_command(self, device: str = "") -> List[str]:
if not self.telemetry_argv:
return []
out: List[str] = []
for a in self.telemetry_argv:
if "{dev}" in a:
if not device:
# Drop the flag the placeholder belonged to as well; a bare
# trailing "-i" makes the sampler fail instead of widening.
if out and out[-1].startswith("-"):
out.pop()
continue
a = a.replace("{dev}", device)
out.append(a)
return out
def parse_telemetry(self, raw: str) -> List[Dict[str, str]]:
"""One dict per device: index, sm_clock_mhz, mem_clock_mhz, temp_c,
power_w, util_pct. Normalised so a CUDA and a MUSA campaign produce the
same columns and can be read side by side."""
return []
def classify_capture_error(self, exc: BaseException) -> str:
"""-> "fatal" | "poisoned" | "recoverable".
Matched against the exception text itself, never a generic footer: torch
appends "Compile with `TORCH_USE_CUDA_DSA` to enable device-side
assertions." to every AcceleratorError, so a loose "device-side assert"
marker would classify every capture refusal as fatal and disable the
documented fallback path outright.
"""
msg = f"{type(exc).__name__}: {exc}".lower()
if any(m in msg for m in self.fatal_error_markers):
return "fatal"
if any(m in msg for m in self.poisoning_error_markers):
return "poisoned"
return "recoverable"
def denies_capture(self, op: str) -> bool:
return bool(op) and op in self.capture_deny_ops
# Plugins that patch around upstream assumptions this backend violates, e.g.
# FlagGems benchmarks that gate themselves on torch.cuda.is_available() even
# though the operator under test runs on whatever device FlagGems dispatched
# to. Loaded first so collection-time overrides land before anything reads
# the item list. Each such plugin must scope itself to its own vendor.
compat_plugins: Tuple[str, ...] = ()
# 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
-49
View File
@@ -1,49 +0,0 @@
from __future__ import annotations
from typing import Dict, List
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 ")
# Unchanged from when this list lived in _cudagraph_plugin; CUDA numbers do
# not move because of the classification refactor.
fatal_error_markers = (
"illegal memory access",
"cudaerrorillegaladdress",
"out of memory",
"device-side assert triggered",
"an illegal instruction",
"misaligned address",
"uncorrectable ecc",
)
# No CUDA wording is known to poison the context without also matching one
# of the fatal markers above, so nothing is listed here rather than guessed.
poisoning_error_markers = ()
telemetry_argv = (
"nvidia-smi",
"--query-gpu=index,clocks.sm,clocks.mem,temperature.gpu,power.draw,utilization.gpu",
"--format=csv,noheader,nounits",
"-i", "{dev}",
)
def parse_telemetry(self, raw: str) -> List[Dict[str, str]]:
out: List[Dict[str, str]] = []
for line in raw.splitlines():
parts = [p.strip() for p in line.split(",")]
if len(parts) != 6 or not parts[0].isdigit():
continue
out.append(dict(zip(
("index", "sm_clock_mhz", "mem_clock_mhz",
"temp_c", "power_w", "util_pct"), parts)))
return out
-101
View File
@@ -1,101 +0,0 @@
from __future__ import annotations
from typing import Dict, List
from .base import DeviceBackend
class MusaBackend(DeviceBackend):
name = "musa"
visibility_env_keys = ("MUSA_VISIBLE_DEVICES",)
torch_namespace = "musa"
flaggems_vendor = "musa"
# torch_musa supplies Stream, Event, synchronize, graph and MUSAGraph, so
# capture works once the timing path takes them from Triton's device
# interface instead of torch.cuda. Measured on an S5000 over 15 operators:
# 11 captured, 3 fell back cleanly with a marker (dropout/rand on the RNG
# seed, sort on a capture-illegal stream op), 1 poisoned the context (see
# capture_deny_ops). Capture is worth having because it removes launch
# overhead the direct timer cannot: batch_norm 16x16x8x48 went 0.0758ms ->
# 0.0089ms, tanh/gelu 4096x4096 about 3.3x, while GEMMs stayed at 1.00x.
supports_cuda_graph = True
supports_mm_cluster_fix = False
supports_triton_hooks = True
# MUSA reports its own wording; none of the CUDA markers match it.
fatal_error_markers = (
"musa error: an illegal memory access",
"out of memory",
"device-side assert triggered",
"misaligned address",
"uncorrectable ecc",
)
# "unknown error" during capture is not recoverable on this runtime: the
# fallback runs, and the process then dies in teardown with zero rows. Treat
# it as poisoning so the operator is reported as capture-incompatible rather
# than silently losing its measurement.
poisoning_error_markers = ("musa error: unknown error",)
# Keyed by FLAGGEMS_PERF_CURRENT_OP, which is the pytest FUNCTION name minus
# `test_` -- `unique2`, `perf_sort` -- not the benchmark file name. Getting
# that wrong is silent: the entry simply never matches.
#
# Observed on an S5000: capturing test_unique2 takes two recoverable
# fallbacks and then raises "MUSA error: unknown error", ending the run at
# rc=1 with no measured rows. Skipping the attempt costs that operator's
# timing precision instead of its measurement. Revisit if a newer torch_musa
# changes the behaviour.
capture_deny_ops = ("unique2",)
# Several FlagGems benchmarks skip on `not torch.cuda.is_available()` while
# the code under test dispatches through flag_gems.device. On MThreads that
# turns a working operator into "not measured"; the plugin drops those
# guards, pinned to one file and one exact skip reason each.
compat_plugins = ("_musa_benchmark_compat_plugin",)
# 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
# mthreads-gmi has no CSV query mode; --json is the only machine-readable
# form it offers. Values arrive with their units glued on ("36C", "1888MHz",
# "477.68W", "87%"), so the parser strips trailing non-numerics rather than
# trusting a fixed suffix.
telemetry_argv = (
"mthreads-gmi", "-q",
"-d", "TEMPERATURE,CLOCK,POWER,UTILIZATION",
"--json",
"-i", "{dev}",
)
def parse_telemetry(self, raw: str) -> List[Dict[str, str]]:
import json
import re
def num(v):
m = re.match(r"^\s*(-?[\d.]+)", str(v or ""))
return m.group(1) if m else ""
start = raw.find("{")
if start < 0:
return []
try:
data = json.loads(raw[start:])
except Exception:
return []
gpus = data.get("GPU") or []
if isinstance(gpus, dict):
gpus = [gpus]
out: List[Dict[str, str]] = []
for g in gpus:
clocks = g.get("Clocks") or {}
power = g.get("Power Readings") or {}
# 键名带尾随空格("Power Draw "),别按字面写死
draw = next((v for k, v in power.items()
if k.strip() == "Power Draw"), "")
out.append({
"index": str(g.get("Index", "")).strip(),
"sm_clock_mhz": num(clocks.get("Graphics")),
"mem_clock_mhz": num(clocks.get("Memory")),
"temp_c": num((g.get("Temperature") or {}).get("GPU Current Temp")),
"power_w": num(draw),
"util_pct": num((g.get("Utilization") or {}).get("Gpu")),
})
return out
-16
View File
@@ -1,16 +0,0 @@
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
+7 -277
View File
@@ -16,10 +16,8 @@ from __future__ import annotations
import argparse import argparse
import csv import csv
import itertools
import os import os
import re import re
import statistics
import sys import sys
from pathlib import Path from pathlib import Path
from typing import Dict, List, Optional, Tuple from typing import Dict, List, Optional, Tuple
@@ -41,31 +39,12 @@ DIRECT_NO_CUDAGRAPH_RE = re.compile(r"BENCHMARK_DIRECT_NO_CUDAGRAPH\b")
REPLAY_FALLBACK_RE = re.compile(r"AUTOTUNE_REPLAY_FALLBACK\b") REPLAY_FALLBACK_RE = re.compile(r"AUTOTUNE_REPLAY_FALLBACK\b")
SUMMARY_FIELDS = [ SUMMARY_FIELDS = [
"repo", "op", "backend", "device", "timer", "repo", "op", "fg_op_name", "dtype", "case_id", "row_status",
"fg_op_name", "dtype", "case_id", "row_status",
"torch_lat", "gems_lat", "speedup", "extra_metrics", "size_detail", "torch_lat", "gems_lat", "speedup", "extra_metrics", "size_detail",
"op_status", "notes", "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]]: def parse_run_log(text: str) -> List[Dict[str, str]]:
"""把 run.log 里的全部结果表解析为行列表(跨多个 dtype 表)。""" """把 run.log 里的全部结果表解析为行列表(跨多个 dtype 表)。"""
rows: List[Dict[str, str]] = [] rows: List[Dict[str, str]] = []
@@ -154,11 +133,10 @@ def summarize(batch_dir: Path, out_path: Optional[Path] = None) -> Path:
text = log.read_text(errors="replace") text = log.read_text(errors="replace")
status_row = status_by_dir.get(f"{repo}/{op_dir.name}", {}) status_row = status_by_dir.get(f"{repo}/{op_dir.name}", {})
op = status_row.get("op") or 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) table_rows = parse_run_log(text)
if not table_rows: if not table_rows:
all_rows.append({ all_rows.append({
"repo": repo, "op": op, **ident, "fg_op_name": "", "dtype": "", "repo": repo, "op": op, "fg_op_name": "", "dtype": "",
"case_id": "", "row_status": "", "torch_lat": "", "case_id": "", "row_status": "", "torch_lat": "",
"gems_lat": "", "speedup": "", "extra_metrics": "", "gems_lat": "", "speedup": "", "extra_metrics": "",
"size_detail": "", "size_detail": "",
@@ -172,8 +150,7 @@ def summarize(batch_dir: Path, out_path: Optional[Path] = None) -> Path:
idx = case_counter.get(key, 0) idx = case_counter.get(key, 0)
case_counter[key] = idx + 1 case_counter[key] = idx + 1
all_rows.append({ all_rows.append({
"repo": repo, "op": op, **ident, "repo": repo, "op": op, "case_id": f"case_{idx:03d}",
"case_id": f"case_{idx:03d}",
"op_status": status_row.get("status", ""), "op_status": status_row.get("status", ""),
"notes": status_row.get("notes", ""), "notes": status_row.get("notes", ""),
**row, **row,
@@ -222,16 +199,8 @@ def ab_rows(before_text: str, after_text: str) -> List[Dict[str, str]]:
return paired return paired
def report_ab(before: Path, after: Path, label: str = "", def report_ab(before: Path, after: Path, label: str = "") -> List[Dict[str, str]]:
strict: bool = True) -> List[Dict[str, str]]: """打印一个算子的前后对比:逐行 speedup 与整体区间。"""
"""打印一个算子的前后对比:逐行 speedup 与整体区间。
先核对两侧口径;不一致时默认拒绝出报告,因为不可比的两次运行会产出
一份和有效报告长得一模一样的表。
"""
if not _report_caliber([("before", before.parent), ("after", after.parent)],
label or "ab", strict):
raise CaliberMismatch()
before_text = before.read_text(encoding="utf-8", errors="replace") before_text = before.read_text(encoding="utf-8", errors="replace")
after_text = after.read_text(encoding="utf-8", errors="replace") after_text = after.read_text(encoding="utf-8", errors="replace")
rows = ab_rows(before_text, after_text) rows = ab_rows(before_text, after_text)
@@ -254,225 +223,6 @@ def report_ab(before: Path, after: Path, label: str = "",
return rows return rows
# Fields whose disagreement makes two runs incomparable. Measured magnitudes,
# so the list is not precautionary: an unpinned autotune config moved
# router_gemm by 8pp; the timer moved replication_pad3d from +43% to +124.5%;
# a shared tuner cache made mm look like 945 compiled variants had become 33;
# --level core sees 5 mm shapes where comprehensive sees 29.
CALIBER_FIELDS = (
"backend", "device", "timer", "bench_level", "shape_mode",
"tuner_cache", "direct_repeat",
)
def caliber_of(run_dir: Path) -> Optional[Dict[str, str]]:
"""Caliber fields from a run's run_identity.json, or None if unreadable."""
if not (run_dir / "run_identity.json").is_file():
return None
try:
from _run_identity import load_identity
ident = load_identity(str(run_dir))
except Exception:
return None
return {k: str(ident.get(k, "")) for k in CALIBER_FIELDS} | {
"autotune_mode": str(ident.get("autotune_mode", ""))
}
def caliber_complaints(runs: List[Tuple[str, Path]]) -> List[str]:
"""Everything about this set of runs that makes their comparison unsound.
Reported rather than silently tolerated because the failure is invisible in
the numbers: two runs on different calibers produce a report that looks
exactly like a valid one. This repository has already published three
figures that were an artifact of one such difference.
"""
out: List[str] = []
cals = [(name, caliber_of(d)) for name, d in runs]
missing = [n for n, c in cals if c is None]
if missing:
out.append(f"口径不可核实:{', '.join(missing)} 缺 run_identity.json"
"(早于该字段的产物)")
known = [(n, c) for n, c in cals if c is not None]
if len(known) < 2:
return out
for field in CALIBER_FIELDS:
seen = {c[field] for _, c in known}
if len(seen) > 1:
vals = ", ".join(f"{n}={c[field] or '<空>'}" for n, c in known)
out.append(f"{field} 不一致:{vals}")
# A/B needs one side to pin the configs the other replays. Two record runs
# each sweep independently, and the winner can differ; the README puts that
# at the 10% level and one-directional per shape, which is enough to swamp
# what is being measured.
modes = [(n, c["autotune_mode"]) for n, c in known]
if all(m.startswith("record") for _, m in modes):
out.append("两侧都是 record 模式:config 未锁定,各自 sweep 可能选出不同 "
"winner。按 README,只有 record 数据支撑的结论按未验证处理")
return out
def caliber_hazards(runs: List[Tuple[str, Path]]) -> List[str]:
"""Not mismatches — shared state that makes a matched caliber less solid.
Warned about rather than refused, because refusing would reject the
default configuration every A/B in this repo currently uses.
"""
cals = [(n, c) for n, c in ((n, caliber_of(d)) for n, d in runs) if c]
recording = [n for n, c in cals if c["autotune_mode"].startswith("record")]
shared = [n for n, c in cals if c["tuner_cache"].startswith("shared")]
if recording and shared:
# 一次 record 跑在共享 libtuner cache 上,sweep 的起点是别的运行留下的
# 结果。曾把 mm 的 945 个变体读成 33 个,误报为编译期收益。A/B 两侧都
# replay 同一份 config 时结论仍成立,绝对值不成立。
return [f"{', '.join(recording)} 在共享 tuner cache 上 record:"
"被 pin 的 config 取决于此前跑过什么。相对结论不受影响,"
"绝对延迟/变体数不可跨机器引用(ZL_ISOLATE_TUNER_CACHE=1 可消除)"]
return []
class CaliberMismatch(SystemExit):
"""Strict-mode refusal. Exit 4 so a driver script can tell it from a crash."""
def __init__(self) -> None:
super().__init__(4)
def _report_caliber(runs: List[Tuple[str, Path]], label: str, strict: bool) -> bool:
"""Print caliber problems. Returns False when strict mode should stop."""
for h in caliber_hazards(runs):
print(f"警告 [caliber] {label}: {h}")
problems = caliber_complaints(runs)
if not problems:
return True
head = "!!!" if strict else "警告"
for p in problems:
print(f"{head} [caliber] {label}: {p}")
if strict:
print(f"!!! [caliber] {label}: 拒绝出报告。确认口径确实可比后加 "
"--allow-caliber-mismatch")
return False
return True
def _lat_index(text: str) -> Dict[Tuple[str, str, int], float]:
"""(dtype, size_detail, occurrence) -> gems latency, SUCCESS rows only.
Occurrence disambiguates a shape a benchmark measures twice in one table,
which BLAS does at --level comprehensive.
"""
out: Dict[Tuple[str, str, int], float] = {}
seen: Dict[Tuple[str, str], int] = {}
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:
continue
base = (row.get("dtype", ""), row.get("size_detail", ""))
i = seen.get(base, 0)
seen[base] = i + 1
out[base + (i,)] = lat
return out
def ab_round_rows(root: Path) -> Tuple[List[Dict[str, str]], List[float], int]:
"""Aggregate an N-round A/B laid out as <root>/r<k>_<A|B>/run.log.
Returns the per-shape rows, the per-round deltas, and the round count.
A single pair cannot say whether its own verdict is stable, which is the
question that matters once a delta lands near zero: measured on an S5000,
topk eliminated 32 converts yet 17 of its 42 shapes changed the SIGN of
their delta between two equally valid runs. Two dispersions answer different
halves of that and both are reported:
side spread max/min-1 within one side's samples -> how noisy the machine
was while measuring
round spread min~max of the per-round deltas -> whether the VERDICT holds,
which is what a single pair cannot show
"""
rounds: List[Tuple[Dict, Dict]] = []
for k in itertools.count(1):
a, b = root / f"r{k}_A" / "run.log", root / f"r{k}_B" / "run.log"
if not (a.is_file() and b.is_file()):
break
rounds.append((
_lat_index(a.read_text(encoding="utf-8", errors="replace")),
_lat_index(b.read_text(encoding="utf-8", errors="replace")),
))
if not rounds:
return [], [], 0
keys = set(rounds[0][0]) & set(rounds[0][1])
for a, b in rounds[1:]:
keys &= set(a) & set(b)
per_round = [
statistics.median((a[k] / b[k] - 1.0) * 100.0 for k in keys)
for a, b in rounds
] if keys else []
rows: List[Dict[str, str]] = []
for key in sorted(keys):
av = [a[key] for a, _ in rounds]
bv = [b[key] for _, b in rounds]
am, bm = statistics.median(av), statistics.median(bv)
deltas = [(a / b - 1.0) * 100.0 for a, b in zip(av, bv)]
rows.append({
"dtype": key[0], "size_detail": key[1],
"before_lat": f"{am:.6f}", "after_lat": f"{bm:.6f}",
"speedup_percent": f"{(am / bm - 1.0) * 100.0:.2f}",
"delta_min": f"{min(deltas):.2f}", "delta_max": f"{max(deltas):.2f}",
"sign_stable": "yes" if all(d > 0 for d in deltas)
or all(d <= 0 for d in deltas) else "NO",
"before_spread": f"{(max(av) / min(av) - 1.0) * 100.0:.2f}",
"after_spread": f"{(max(bv) / min(bv) - 1.0) * 100.0:.2f}",
})
return rows, per_round, len(rounds)
def report_ab_rounds(root: Path, label: str = "",
strict: bool = True) -> List[Dict[str, str]]:
"""Print the multi-round A/B: per-shape medians plus both dispersions."""
name = label or root.name
cells = sorted(d for d in root.glob("r*_[AB]") if (d / "run.log").is_file())
if cells and not _report_caliber([(d.name, d) for d in cells], name, strict):
raise CaliberMismatch()
rows, per_round, n = ab_round_rows(root)
if n < 2:
print(f"[ab-rounds] {name}: 需要至少 2 轮 <root>/r<k>_<A|B>/run.log"
f"(找到 {n} 轮)")
return rows
if not rows:
print(f"[ab-rounds] {name}: {n} 轮,但没有全轮次都成功的配对行")
return rows
for r in rows:
flag = "" if r["sign_stable"] == "yes" else " 符号不稳"
print(f" {r['dtype']:<16} {float(r['speedup_percent']):+7.2f}% "
f"[{float(r['delta_min']):+.2f}%~{float(r['delta_max']):+.2f}%] "
f"{r['before_lat']} -> {r['after_lat']} "
f"{r['size_detail']}{flag}")
med = [float(r["speedup_percent"]) for r in rows]
spread = [float(r["before_spread"]) for r in rows] + \
[float(r["after_spread"]) for r in rows]
unstable = sum(1 for r in rows if r["sign_stable"] != "yes")
print(f"[ab-rounds] {name}: {n} 轮 × {len(rows)} shape,"
f"speedup 中位 {statistics.median(med):+.2f}%,"
f"逐轮 {min(per_round):+.2f}% ~ {max(per_round):+.2f}%,"
f"单侧噪声中位 {statistics.median(spread):.2f}%")
if unstable:
# The verdict, not the machine, is what these rows fail to pin down.
print(f"[ab-rounds] {name}: {unstable}/{len(rows)} 个 shape 的 delta "
"在轮次间改变符号——这些行的方向不可结论,与计时口径无关")
return rows
def main() -> None: def main() -> None:
p = argparse.ArgumentParser(description=__doc__.splitlines()[0]) p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("batch_dir", type=Path, nargs="?", p.add_argument("batch_dir", type=Path, nargs="?",
@@ -481,33 +231,13 @@ def main() -> None:
help="输出 CSV(默认 <batch_dir>/summary.csv)") help="输出 CSV(默认 <batch_dir>/summary.csv)")
p.add_argument("--ab", nargs=2, type=Path, metavar=("BEFORE", "AFTER"), p.add_argument("--ab", nargs=2, type=Path, metavar=("BEFORE", "AFTER"),
help="两份 run.log,按 shape 配对输出该算子的 speedup 区间") help="两份 run.log,按 shape 配对输出该算子的 speedup 区间")
p.add_argument("--ab-rounds", type=Path, default=None, metavar="ROOT", p.add_argument("--label", default="", help="--ab 报告里的算子名")
help="多轮 A/B 目录,含 r<k>_A/run.log 与 r<k>_B/run.log;"
"逐 shape 取轮次中位,并报逐轮区间与符号稳定性")
p.add_argument("--label", default="", help="--ab / --ab-rounds 报告里的算子名")
p.add_argument("--caliber", type=Path, default=None, metavar="RUN_DIR",
help="打印该 run 的口径字段(TSV 一行),供筛选表复用")
p.add_argument("--allow-caliber-mismatch", action="store_true",
help="口径不一致时仍出报告(默认拒绝)。只在确认两次运行"
"确实可比时使用")
args = p.parse_args() args = p.parse_args()
if args.caliber is not None:
# 口径的定义只有 CALIBER_FIELDS 一处;筛选表等复用它,不要在别处重列
cal = caliber_of(args.caliber) or {}
print(" ".join(str(cal.get(k, "")) for k in CALIBER_FIELDS))
return
if args.ab_rounds:
if not args.ab_rounds.is_dir():
raise SystemExit(f"目录不存在: {args.ab_rounds}")
report_ab_rounds(args.ab_rounds, args.label,
strict=not args.allow_caliber_mismatch)
return
if args.ab: if args.ab:
for path in args.ab: for path in args.ab:
if not path.is_file(): if not path.is_file():
raise SystemExit(f"run.log 不存在: {path}") raise SystemExit(f"run.log 不存在: {path}")
report_ab(args.ab[0], args.ab[1], args.label, report_ab(args.ab[0], args.ab[1], args.label)
strict=not args.allow_caliber_mismatch)
return return
if args.batch_dir is None: if args.batch_dir is None:
raise SystemExit("需要 batch_dir,或用 --ab BEFORE AFTER") raise SystemExit("需要 batch_dir,或用 --ab BEFORE AFTER")
-128
View File
@@ -1,128 +0,0 @@
#!/bin/bash
# 多卡跑一批算子的六轮 A/B:每卡一个 ab_repeat.sh,卡内串行。
#
# 它只做编排,不做测量。ab_repeat.sh 一行不动——那是已经跑出结论的、验证过的串行
# 路径,不该为了并行去改它。分片交给 partition.py(按实测耗时贪心,比 hash 分片在
# 这批活上快 50 分钟)。
#
# 并行必须隔离 tuner cache:8 个进程写同一个 sqlite 是要避开的事。代价是 pin 下来
# 的 config 与共享缓存下跑的批次不同,两批的绝对延迟不可互引——batch_summary.py
# 的口径守卫会自己拒绝这种对比,不用记着。
#
# 用法:
# SWITCH=FLAGTREE_MUSA_RLC_ENHANCE DEVICES=0,1,2,3,4,5,6,7 \
# SCREEN=runs/ab_rlc_XXX/screen.tsv \
# ZL_BACKEND=musa FLAGGEMS_DIR=/path/to/FlagGems bash campaign.sh
#
# DRY_RUN=1 只打印分片方案,不起跑
# EXTRA_OPS=".." 追加筛选表里没有的算子(如 benchmark 刚修好的),塞给最空的卡
# ROUNDS=6 CELL_TIMEOUT=2400 VERDICT=IR_CHANGED
set -uo pipefail
cd "$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
: "${SWITCH:?SWITCH is required, e.g. FLAGTREE_MUSA_RLC_ENHANCE}"
: "${SCREEN:?SCREEN is required, 指向 ab_rlc_test.sh 的 screen.tsv}"
DEVICES="${DEVICES:-0}"
ROUNDS="${ROUNDS:-6}"
CELL_TIMEOUT="${CELL_TIMEOUT:-2400}"
VERDICT="${VERDICT:-IR_CHANGED}"
EXTRA_OPS="${EXTRA_OPS:-}"
DRY_RUN="${DRY_RUN:-0}"
TAB=$'\t'
[[ -f "$SCREEN" ]] || { echo "!!! 筛选表不存在: $SCREEN" >&2; exit 2; }
# 设备可见性变量名从 backends 层拿,不写死 MUSA_VISIBLE_DEVICES
VIS_KEY=$(python -c "import sys; sys.path.insert(0, '.')
from backends import get_backend
print(get_backend().visibility_env_keys[0])" 2>/dev/null) || VIS_KEY=""
[[ -n "$VIS_KEY" ]] || { echo "!!! 取不到后端的可见性变量名,检查 ZL_BACKEND" >&2; exit 2; }
echo ">>> switch=$SWITCH rounds=$ROUNDS devices=$DEVICES 可见性变量=$VIS_KEY"
python partition.py --screen "$SCREEN" --devices "$DEVICES" \
--verdict "$VERDICT" --rounds "$ROUNDS" || exit 2
# 每卡一行 "<device><TAB><OPS><TAB><权重秒>"
mapfile -t PLAN < <(python partition.py --screen "$SCREEN" --devices "$DEVICES" \
--verdict "$VERDICT" --format shell) || exit 2
if [[ -n "$EXTRA_OPS" ]]; then
# 追加到当前负载最小的卡。按算子个数选是错的——一个重算子能顶三个轻的,
# partition.py 的第三列给的就是权重秒数
lightest=0; least=""
for i in "${!PLAN[@]}"; do
IFS="$TAB" read -r _d _o _w <<< "${PLAN[$i]}"
if [[ -z "$least" ]] || (( ${_w:-0} < least )); then least=${_w:-0}; lightest=$i; fi
done
IFS="$TAB" read -r d o w <<< "${PLAN[$lightest]}"
PLAN[$lightest]="${d}${TAB}${o:+$o }${EXTRA_OPS}${TAB}${w}"
echo ">>> EXTRA_OPS 追加到卡${d}(当前最空,$(( ${w:-0} * ROUNDS / 60 )) 分钟): $EXTRA_OPS"
fi
if [[ "$DRY_RUN" == "1" ]]; then
echo ">>> DRY_RUN=1,不起跑。逐卡将执行:"
for row in "${PLAN[@]}"; do
IFS="$TAB" read -r d o w <<< "$row"
printf ' 卡%s %4s 分钟 OPS="%s"\n' "$d" "$(( ${w:-0} * ROUNDS / 60 ))" "$o"
done
exit 0
fi
ROOT="runs/campaign_${SWITCH}_$(date +%Y%m%d-%H%M%S)"
mkdir -p "$ROOT"
printf '%s\n' "${PLAN[@]}" > "$ROOT/plan.tsv"
echo ">>> 产物: $(pwd)/$ROOT"
declare -a PIDS=() CARDS=()
for row in "${PLAN[@]}"; do
IFS="$TAB" read -r dev ops _w <<< "$row"
[[ -n "$ops" ]] || { echo ">>> 卡$dev 没有分到活,跳过"; continue; }
out="$ROOT/card${dev}"
mkdir -p "$out"
# AB_ROOT 让每卡有独立产物根:八个 ab_repeat.sh 同一秒启动,默认的时间戳目录名会撞
env "$VIS_KEY=$dev" ZL_DEVICE="$dev" ZL_ISOLATE_TUNER_CACHE=1 \
AB_ROOT="$out" SWITCH="$SWITCH" ROUNDS="$ROUNDS" \
CELL_TIMEOUT="$CELL_TIMEOUT" OPS="$ops" NO_COLOR=1 \
bash ab_repeat.sh > "$ROOT/card${dev}.log" 2>&1 &
PIDS+=($!); CARDS+=("$dev")
echo ">>> 卡$dev pid=$! ops=$ops"
done
(( ${#PIDS[@]} )) || { echo "!!! 没有一张卡分到活" >&2; exit 2; }
echo ">>> ${#PIDS[@]} 张卡已启动,等待全部结束……"
fail=0
for i in "${!PIDS[@]}"; do
if wait "${PIDS[$i]}"; then
echo ">>> 卡${CARDS[$i]} 结束 rc=0"
else
rc=$?; fail=$((fail+1))
echo "!!! 卡${CARDS[$i]} 结束 rc=$rc —— 看 $ROOT/card${CARDS[$i]}.log"
fi
done
# 合并逐格记录:卡号加在最前面,便于事后按卡追
{
first=1
for d in "${CARDS[@]}"; do
f="$ROOT/card${d}/status.tsv"
[[ -f "$f" ]] || continue
if (( first )); then printf 'device%s%s\n' "$TAB" "$(head -1 "$f")"; first=0; fi
tail -n +2 "$f" | sed "s/^/${d}${TAB}/"
done
} > "$ROOT/status.tsv"
echo
echo "########## 汇总 ##########"
echo ">>> 失败的卡: $fail / ${#PIDS[@]}"
if [[ -s "$ROOT/status.tsv" ]]; then
echo ">>> 逐格记录: $(pwd)/$ROOT/status.tsv ($(( $(wc -l < "$ROOT/status.tsv") - 1 )) 格)"
awk -F"$TAB" 'NR>1 && ($6+0)!=0 {n++} END{printf ">>> 非零退出的格: %d\n", n+0}' "$ROOT/status.tsv"
awk -F"$TAB" 'NR>1 {fb+=$10; ng+=$11; po+=$12}
END{printf ">>> 三个标记合计: replay_fallback=%d no_cudagraph=%d capture_poisoned=%d\n",
fb+0, ng+0, po+0}' "$ROOT/status.tsv"
fi
echo ">>> 逐算子报告在各卡的 summary.log 里;单独重出:"
echo " python batch_summary.py --ab-rounds $ROOT/card<N>/<op> --label <op>"
echo ">>> 注意:本批 ZL_ISOLATE_TUNER_CACHE=1,绝对延迟不可与共享缓存的批次互引"
-103
View File
@@ -1,103 +0,0 @@
# S5000 MUSA RLC 产品验证 campaign
在 Moore Threads S5000 上,经由 FlagGems public dispatch,测量 FlagTree
`RemoveLayoutConversions`(RLC)增强对 32 个冻结 workload 的影响。
本目录只提供 campaign 层:契约、冻结、执行、判定、重记账。单算子插件、shape、
inventory 仍然复用仓库根目录的既有能力。
## 组成
| 文件 | 作用 |
|---|---|
| `manifest/musa-rlc-32-product-v1.json` | 冻结契约:源码提交、设备、cell、计时口径、32 个 workload |
| `prepare_musa_rlc_product.py` | 校验 manifest 与 `ops/inventory_main.csv`,冻结不可变 task matrix |
| `run_musa_rlc_product.py` | 串行执行 task matrix,产出逐 cell、逐 task、campaign 三级收据 |
| `run_pytest_musa_rlc_product.sh` | 单个算子单个 cell 的 worker,挂载根目录插件 |
| `reaccount_musa_rlc_topology.py` | 对已封存 run 按新口径重记账,不占用设备 |
| `tests/` | 上述四者的单元测试 |
根目录的 `_musa_benchmark_compat_plugin.py` 属于插件族,不在本目录。
## 三个 cell
| cell | `FLAGTREE_MUSA_RLC_ENHANCE` | phase mask | 含义 |
|---|---|---|---|
| `off` | 0 | — | 增强关闭的基线 |
| `mask0` | 1 | 0 | 增强开启但所有 phase 关闭;它的 topology 必须等于 `off`,否则控制面本身有问题 |
| `candidate` | 1 | 13 | 1a、2、3 启用,1b 未启用 |
`off` 记录 autotune config,`mask0` 与 `candidate` 回放同一套 config。回放回退即判定失败,
这样 A/B 之间的差异才收敛到编译器改动这一个变量。
## 跑法
```bash
python campaigns/musa_rlc_product/prepare_musa_rlc_product.py \
campaigns/musa_rlc_product/manifest/musa-rlc-32-product-v1.json \
--inventory-dir ops --selection canary --dry-run
```
`--dry-run` 只做校验与预算估算。去掉它并给出 `--output-root` 才会冻结 task matrix。
`--selection all` 需要 `--canary-receipt` 指向一份通过的 canary 收据。
执行:
```bash
python campaigns/musa_rlc_product/run_musa_rlc_product.py \
--manifest <frozen>/manifest.json \
--task-matrix <frozen>/task-matrix.json \
--build-root <flagtree build root> \
--flaggems <flaggems> --zl-bench <this repo> \
--worker campaigns/musa_rlc_product/run_pytest_musa_rlc_product.sh \
--output-root <run root>
```
runner 在开跑前校验:manifest 与 matrix 的哈希绑定、FlagTree/FlagGems/zl_bench 三处源码身份、
Triton 与 FlagGems 确实从冻结路径导入、设备可见数为 1、所有 benchmark 与 correctness 目标
在真实源码树里存在、设备当前没有其它进程。任何一条不成立就整批拒绝。
单元测试直接跑文件:
```bash
python campaigns/musa_rlc_product/tests/test_musa_rlc_product.py
python campaigns/musa_rlc_product/tests/test_reaccount_musa_rlc_topology.py
```
## 判定
| 类 | 含义 |
|---|---|
| G | 正确,convert topology 下降,且双执行顺序下均稳定提速 ≥5% |
| C | 正确,topology 下降,但收益不足 5% |
| B | 正确,但 candidate 没有减少被测算子的 convert topology |
| R | 正确,但双顺序稳定退化 ≤ -3% |
| K | candidate 特有的正确性失败 |
| X | 不支持、超时、harness 失败,或同配置回放被破坏 |
只有 topology 相对 incumbent 下降时才追加反向顺序计时。convert 数下降不等于变快,
这一点在本 campaign 内已有反例。
## 三条口径边界
**convert 计数只算被测算子。** benchmark 调用会先构造一次输入再在计时区内反复执行算子。
把所有执行过的 kernel 都计入,会把输入构造 kernel 里消掉的 convert 算成收益,而计时器
根本观测不到它。因此每个 dump 的 TTGIR 都与 `index.tsv` 的 launch 数关联,只有 launch
数超过 `measured_launch_floor` 的变体才计入。未收敛总数与被排除文件保留在收据里;
无法与 index 关联的 cell 判为 X,而不是相信一个界定不了的 topology 声明。
**correctness 只跑该 workload 实际测量的变体。** 以整个测试文件为粒度,会让无关变体
决定被测算子的结论。workload 可以钉 node id 或 `-k` 选择器,两者互斥,且任何收窄都必须
写明理由。完整 argv 落在收据里。
**目标在开跑前解析。** inventory 可能落后于固定的 FlagGems 提交,例如上游把一个 benchmark
文件拆成多个。这类任务过去只在 worker 前置检查里静默失败,留下空日志。现在 runner 会
对真实源码树解析全部目标,workload 可以声明 inventory 里没有的目标但必须给出理由,
worker 的每个前置检查失败时自报名称。
## 不属于本 campaign 的结论
一次 terminal 的 campaign 只证明:该冻结语料在该设备该提交下的 dispatch 正确性、
layout topology 响应、以及同契约下的 native-event 延迟。它不证明跨厂商正确性、
性能对等、CI 状态,也不构成打开生产默认开关的依据。manifest 的
`production_default_policy` 固定为 `unchanged-off`。
@@ -1,551 +0,0 @@
{
"schema_version": 1,
"campaign_id": "s5000-musa-rlc-32-product-v1-20260902-v3",
"experiment_kind": "musa-rlc-product-validation",
"permitted_claim": "S5000 MUSA product-dispatch correctness, layout-topology response, and same-contract native-event latency for this frozen corpus only",
"production_default_policy": "unchanged-off",
"source": {
"flagtree_commit": "3059917e2df16f8ea35c5efe1ed1af2d87117396",
"flaggems_commit": "46ef8080329190d8da8738340240de2d5a475950",
"zl_bench_commit": "da2288564584039366e7be07fe4e94539a560c81"
},
"device": {
"ssh_alias": "S5000",
"container": "zixu_flagos_S5000",
"backend": "mthreads",
"model": "Moore Threads S5000",
"physical_device": "7",
"logical_device": "0",
"visible_device_count": 1
},
"protocol": {
"screen_cells": [
"off",
"mask0",
"candidate"
],
"timing_cells": [
"off",
"candidate"
],
"canary_order": [
"dropout",
"softmax",
"instance_norm",
"fft",
"sort",
"batch_norm",
"topk",
"conv1d"
],
"candidate_phase_mask": 13,
"candidate_profitability_policy": true,
"candidate_product_launch_count": 1,
"default_benchmark_dtypes": [
"float16"
],
"timer": "triton.testing.do_bench via MUSA device-interface Event",
"autotune_contract": "off-record then mask0/candidate replay from off",
"fresh_process_per_operator_cell": true,
"same_physical_device_for_all_cells": true,
"serial_measurement": true,
"full_campaign_requires_canary_terminal": true,
"correctness": "FlagGems public dispatch quick tests under candidate, narrowed to the variant each workload measures and run with the MThreads compatibility plugin, plus benchmark execution for every cell",
"classification": {
"G": "correct, convert count decreases, and stable speedup >= 5%",
"C": "correct and convert count decreases, but latency intervals overlap or speedup < 5%",
"B": "correct but candidate does not reduce the executed convert topology",
"R": "correct but stable regression <= -3%",
"K": "candidate correctness failure",
"X": "unsupported, timeout, harness failure, or same-contract replay violation"
},
"timeout_seconds_by_cost": {
"light": 300,
"medium": 900,
"heavy": 1800
},
"measured_launch_floor": 1,
"supported_benchmark_dtypes": [
"float16",
"float32",
"bfloat16",
"int16",
"int32",
"bool",
"complex64"
]
},
"workloads": [
{
"id": "batch_norm",
"repo": "main",
"op": "batch_norm",
"op_file": "batch_norm",
"shape_key": "batch_norm",
"correctness_file": "tests/test_batch_norm.py",
"family": "normalization",
"nv_evidence": "strong-positive",
"nv_gain_percent": 43.1,
"prior_s5000": "B 0->0",
"cost": "medium",
"canary": true
},
{
"id": "replication_pad3d",
"repo": "main",
"op": "replication_pad3d",
"op_file": "replication_pad3d",
"shape_key": "replication_pad3d",
"correctness_file": "tests/test_replication_pad3d.py",
"family": "padding",
"nv_evidence": "strong-positive",
"nv_gain_percent": 38.6,
"prior_s5000": "B 0->0 at author shape",
"cost": "medium",
"canary": false
},
{
"id": "fft",
"repo": "main",
"op": "perf_fft",
"op_file": "fft",
"shape_key": "fft",
"correctness_file": "tests/test_fft.py",
"family": "fft",
"nv_evidence": "strong-positive",
"nv_gain_percent": 36.5,
"prior_s5000": "G 38->2",
"cost": "heavy",
"canary": true,
"benchmark_dtypes": [
"complex64"
]
},
{
"id": "topk",
"repo": "main",
"op": "topk",
"op_file": "topk",
"shape_key": "topk",
"correctness_file": "tests/test_topk.py",
"family": "selection",
"nv_evidence": "strong-positive-shared-network-risk",
"nv_gain_percent": 20.8,
"prior_s5000": "B 0->0 at large author shape",
"cost": "medium",
"canary": true
},
{
"id": "instance_norm",
"repo": "main",
"op": "instance_norm",
"op_file": "instance_norm",
"shape_key": "instance_norm",
"correctness_file": "tests/test_instance_norm.py",
"family": "normalization",
"nv_evidence": "strong-positive",
"nv_gain_percent": 17.8,
"prior_s5000": "G 1->0",
"cost": "medium",
"canary": true
},
{
"id": "dropout",
"repo": "main",
"op": "dropout",
"op_file": "dropout",
"shape_key": "dropout",
"correctness_file": "tests/test_dropout.py",
"family": "rng",
"nv_evidence": "strong-positive-shared-network-risk",
"nv_gain_percent": 15.2,
"prior_s5000": "product G 8->0 +9.15%",
"cost": "light",
"canary": true
},
{
"id": "router_gemm",
"repo": "main",
"op": "perf_router_gemm",
"op_file": "router_gemm",
"shape_key": "router_gemm",
"correctness_file": "tests/test_router_gemm.py",
"family": "mma",
"nv_evidence": "strong-positive-mma",
"nv_gain_percent": 14.5,
"prior_s5000": "public screen available; product v1 pending",
"cost": "heavy",
"canary": false,
"benchmark_dtypes": [
"bfloat16"
]
},
{
"id": "fused_recurrent_gated_delta_rule",
"repo": "main",
"op": "perf_fused_recurrent_gated_delta_rule",
"op_file": "FLA/test_fused_recurrent_gated_delta_rule_perf",
"shape_key": null,
"correctness_file": "tests/test_FLA/test_fused_recurrent_gated_delta_rule.py",
"family": "recurrent",
"nv_evidence": "strong-positive-recurrent",
"nv_gain_percent": 14.2,
"prior_s5000": "not product-qualified",
"cost": "heavy",
"canary": false,
"benchmark_dtypes": [
"bfloat16"
]
},
{
"id": "top_k_per_row_decode",
"repo": "main",
"op": "top_k_per_row_decode",
"op_file": "top_k_per_row_decode",
"shape_key": "top_k_per_row_decode",
"correctness_file": "tests/test_top_k_per_row_decode.py",
"family": "selection",
"nv_evidence": "strong-positive-shared-network-risk",
"nv_gain_percent": 13.0,
"prior_s5000": "not product-qualified",
"cost": "medium",
"canary": false
},
{
"id": "sort",
"repo": "main",
"op": "perf_sort",
"op_file": "sort",
"shape_key": "sort",
"correctness_file": "tests/test_sort.py",
"family": "sort",
"nv_evidence": "strong-positive-shared-network-risk",
"nv_gain_percent": 12.3,
"prior_s5000": "G 2->1",
"cost": "medium",
"canary": true
},
{
"id": "sort_stable",
"repo": "main",
"op": "perf_sort_stable",
"op_file": "sort",
"shape_key": "sort_stable",
"correctness_file": "tests/test_sort.py",
"family": "sort",
"nv_evidence": "strong-positive-shared-network-risk",
"nv_gain_percent": 12.3,
"prior_s5000": "G 2->1",
"cost": "medium",
"canary": false
},
{
"id": "uniform",
"repo": "main",
"op": "uniform_",
"op_file": "uniform_",
"shape_key": "uniform_",
"correctness_file": "tests/test_uniform_.py",
"family": "rng",
"nv_evidence": "strong-positive-rng",
"nv_gain_percent": 10.3,
"prior_s5000": "public boundary risk; product v1 pending",
"cost": "light",
"canary": false
},
{
"id": "rand",
"repo": "main",
"op": "rand",
"op_file": "rand",
"shape_key": "rand",
"correctness_file": "tests/test_rand.py",
"family": "rng",
"nv_evidence": "strong-positive-rng",
"nv_gain_percent": 8.7,
"prior_s5000": "product v1 pending",
"cost": "light",
"canary": false
},
{
"id": "rand_like",
"repo": "main",
"op": "rand_like",
"op_file": "rand_like",
"shape_key": "rand_like",
"correctness_file": "tests/test_rand_like.py",
"family": "rng",
"nv_evidence": "strong-positive-rng",
"nv_gain_percent": 8.6,
"prior_s5000": "product v1 pending",
"cost": "light",
"canary": false
},
{
"id": "max_pool2d_with_indices",
"repo": "main",
"op": "max_pool2d_with_indices",
"op_file": "max_pool2d",
"shape_key": "max_pool2d_with_indices",
"correctness_file": "tests/test_max_pool2d.py",
"family": "reduction-index",
"nv_evidence": "positive-reduction-index",
"nv_gain_percent": 6.1,
"prior_s5000": "B 0->0",
"cost": "medium",
"canary": false
},
{
"id": "randperm",
"repo": "main",
"op": "randperm",
"op_file": "randperm",
"shape_key": "randperm",
"correctness_file": "tests/test_randperm.py",
"family": "rng",
"nv_evidence": "small-positive-rng",
"nv_gain_percent": 4.0,
"prior_s5000": "C 7->1 latency overlap",
"cost": "medium",
"canary": false,
"benchmark_dtypes": [
"int32"
]
},
{
"id": "randn",
"repo": "main",
"op": "randn",
"op_file": "randn",
"shape_key": "randn",
"correctness_file": "tests/test_randn.py",
"family": "rng",
"nv_evidence": "diluted-positive-rng",
"nv_gain_percent": 4.0,
"prior_s5000": "product v1 pending",
"cost": "light",
"canary": false
},
{
"id": "randn_like",
"repo": "main",
"op": "randn_like",
"op_file": "randn_like",
"shape_key": "randn_like",
"correctness_file": "tests/test_randn_like.py",
"family": "rng",
"nv_evidence": "diluted-positive-rng",
"nv_gain_percent": 4.0,
"prior_s5000": "product v1 pending",
"cost": "light",
"canary": false
},
{
"id": "upsample_bicubic2d_aa",
"repo": "main",
"op": "upsample_bicubic2d_aa",
"op_file": "upsample_bicubic2d_aa",
"shape_key": "upsample_bicubic2d_aa",
"correctness_file": "tests/test_upsample_bicubic2d_aa.py",
"family": "interpolation",
"nv_evidence": "diluted-positive-cold-path",
"nv_gain_percent": 1.7,
"prior_s5000": "C 10->0 latency flat",
"cost": "medium",
"canary": false,
"correctness_selector": "not backward",
"correctness_scope_reason": "the benchmarked op is the forward path (benchmark/test_upsample_bicubic2d_aa.py::test_upsample_bicubic2d_aa); the backward baseline defect fails identically under off/mask0/candidate and is tracked as a separate MUSA numerics issue"
},
{
"id": "grid_sample",
"repo": "main",
"op": "grid_sample",
"op_file": "grid_sample",
"shape_key": "grid_sample",
"correctness_file": "tests/test_grid_sample.py",
"family": "sampling",
"nv_evidence": "structure-change-no-win",
"nv_gain_percent": 0.0,
"prior_s5000": "B 0->0",
"cost": "medium",
"canary": false,
"benchmark_selector": "dtype0",
"benchmark_scope_reason": "benchmark/test_grid_sample.py parametrises over consts.FLOAT_DTYPES and declares dtypes=[dtype] per case, so a campaign-wide --dtypes float16 filter makes the float32 and bfloat16 cases raise; dtype0 is the float16 case this workload measures"
},
{
"id": "max_pool2d_backward",
"repo": "main",
"op": "max_pool2d_backward",
"op_file": "max_pool2d",
"shape_key": "max_pool2d_backward",
"correctness_file": "tests/test_max_pool2d.py",
"family": "reduction-index",
"nv_evidence": "structure-change-no-win",
"nv_gain_percent": 0.0,
"prior_s5000": "B 0->0",
"cost": "medium",
"canary": false
},
{
"id": "conv1d",
"repo": "main",
"op": "conv1d",
"op_file": "conv1d",
"shape_key": "conv1d",
"correctness_file": "tests/test_conv1d.py",
"family": "convolution",
"nv_evidence": "light-regression-review",
"nv_gain_percent": -4.2,
"prior_s5000": "C 1->0 via phase3, slight slowdown",
"cost": "heavy",
"canary": true
},
{
"id": "layer_norm_backward",
"repo": "main",
"op": "layer_norm_backward",
"op_file": "layer_norm",
"shape_key": "layer_norm_backward",
"correctness_file": "tests/test_layer_norm.py",
"family": "normalization",
"nv_evidence": "light-regression-review",
"nv_gain_percent": -1.9,
"prior_s5000": "C 4->2 latency flat",
"cost": "medium",
"canary": false
},
{
"id": "cutlass_scaled_mm",
"repo": "main",
"op": "cutlass_scaled_mm_benchmark",
"op_file": "cutlass_scaled_mm",
"shape_key": "cutlass_scaled_mm",
"correctness_file": "tests/test_cutlass_scaled_mm.py",
"family": "mma",
"nv_evidence": "light-regression-mma",
"nv_gain_percent": -1.2,
"prior_s5000": "likely unsupported CUDA-specific control",
"cost": "heavy",
"canary": false
},
{
"id": "gelu_backward",
"repo": "main",
"op": "gelu_backward",
"op_file": "gelu",
"shape_key": "gelu_backward",
"correctness_file": "tests/test_gelu.py",
"family": "elementwise-backward",
"nv_evidence": "direct-no-cudagraph-noise",
"nv_gain_percent": null,
"prior_s5000": "B 0->0 with LLVM optimizations",
"cost": "medium",
"canary": false
},
{
"id": "fused_marlin_moe_mxfp4",
"repo": "main",
"op": "fused_marlin_moe_w4a16_mxfp4",
"op_file": "fused_marlin_moe_w4a16_mxfp4",
"shape_key": null,
"correctness_file": "tests/test_fused_marlin_moe.py",
"family": "quantized-moe",
"nv_evidence": "autotune-config-change",
"nv_gain_percent": null,
"prior_s5000": "likely unsupported CUDA-specific control",
"cost": "heavy",
"canary": false,
"correctness_selector": "mxfp4",
"correctness_scope_reason": "tests/test_fused_marlin_moe.py also holds w4a16_int4 and w8a16_int8 variants; only the mxfp4 variant corresponds to this workload, and the int8 failures are a separate baseline defect",
"inventory_override_reason": "zl_bench da2288564584039366e7be07fe4e94539a560c81 lists benchmark/test_fused_marlin_moe.py, but FlagGems 46ef8080329190d8da8738340240de2d5a475950 split it into per-quantisation files; the mxfp4 target is benchmark/test_fused_marlin_moe_w4a16_mxfp4.py::test_fused_marlin_moe_w4a16_mxfp4"
},
{
"id": "bincount",
"repo": "main",
"op": "bincount",
"op_file": "bincount",
"shape_key": "bincount",
"correctness_file": "tests/test_bincount.py",
"family": "histogram",
"nv_evidence": "autotune-config-change",
"nv_gain_percent": null,
"prior_s5000": "product v1 pending",
"cost": "medium",
"canary": false,
"benchmark_dtypes": [
"float32"
]
},
{
"id": "masked_select",
"repo": "main",
"op": "masked_select",
"op_file": "masked_select",
"shape_key": "masked_select",
"correctness_file": "tests/test_masked_select.py",
"family": "selection",
"nv_evidence": "holdout-shared-network",
"nv_gain_percent": null,
"prior_s5000": "public qualify available",
"cost": "medium",
"canary": false
},
{
"id": "mm",
"repo": "main",
"op": "mm",
"op_file": "mm",
"shape_key": "mm",
"correctness_file": "tests/test_mm.py",
"family": "mma-control",
"nv_evidence": "holdout-mma-control",
"nv_gain_percent": null,
"prior_s5000": "control only",
"cost": "heavy",
"canary": false
},
{
"id": "scaled_dot_product_attention",
"repo": "main",
"op": "scaled_dot_product_attention",
"op_file": "scaled_dot_product_attention",
"shape_key": "scaled_dot_product_attention",
"correctness_file": "tests/test_scaled_dot_product_attention.py",
"family": "attention-control",
"nv_evidence": "holdout-attention-control",
"nv_gain_percent": null,
"prior_s5000": "control only",
"cost": "heavy",
"canary": false
},
{
"id": "softmax",
"repo": "main",
"op": "softmax",
"op_file": "softmax",
"shape_key": "softmax",
"correctness_file": "tests/test_softmax.py",
"family": "reduction-control",
"nv_evidence": "holdout-reduction-control",
"nv_gain_percent": null,
"prior_s5000": "product B 0->0",
"cost": "light",
"canary": true
},
{
"id": "layer_norm",
"repo": "main",
"op": "layer_norm",
"op_file": "layer_norm",
"shape_key": "layer_norm",
"correctness_file": "tests/test_layer_norm.py",
"family": "normalization-control",
"nv_evidence": "holdout-normalization-control",
"nv_gain_percent": null,
"prior_s5000": "control only",
"cost": "medium",
"canary": false
}
],
"predecessor_campaign_id": "s5000-musa-rlc-32-product-v1-20260901-v2"
}
@@ -1,478 +0,0 @@
#!/usr/bin/env python3
"""Validate and freeze the single-device MUSA RLC product campaign."""
from __future__ import annotations
import argparse
import csv
import hashlib
import json
import os
import shutil
from pathlib import Path
from typing import Any
COSTS = {"light", "medium", "heavy"}
CELLS = ["off", "mask0", "candidate"]
TIMING_CELLS = ["off", "candidate"]
ESTIMATED_SECONDS = {"light": 180, "medium": 600, "heavy": 1200}
class CampaignError(RuntimeError):
pass
def _canonical(value: object) -> bytes:
return (json.dumps(value, indent=2, sort_keys=True) + "\n").encode("utf-8")
def _is_git_sha(value: object) -> bool:
return (
isinstance(value, str)
and len(value) == 40
and all(char in "0123456789abcdef" for char in value)
)
def _correctness_scope_errors(row: dict[str, Any], index: int) -> list[str]:
"""Check an optional narrowing of a workload's correctness gate.
Running a whole test file lets an unrelated variant decide the verdict for
the operator actually being measured. A workload may therefore pin node ids
or a ``-k`` selector, but never both, and narrowing always requires a stated
reason so the exclusion stays visible in review.
"""
name = row.get("id", index)
selector = row.get("correctness_selector")
node_ids = row.get("correctness_node_ids")
errors: list[str] = []
if selector is not None and node_ids is not None:
errors.append(
f"workload {name} must not set both correctness_selector and "
"correctness_node_ids"
)
if selector is not None and (not isinstance(selector, str) or not selector.strip()):
errors.append(f"workload {name} correctness_selector must be a non-empty string")
if node_ids is not None and (
not isinstance(node_ids, list)
or not node_ids
or not all(isinstance(value, str) and value.strip() for value in node_ids)
):
errors.append(
f"workload {name} correctness_node_ids must be non-empty strings"
)
if selector is None and node_ids is None:
return errors
reason = row.get("correctness_scope_reason")
if not isinstance(reason, str) or not reason.strip():
errors.append(
f"workload {name} narrows correctness and must state "
"correctness_scope_reason"
)
return errors
def _benchmark_scope_errors(row: dict[str, Any], index: int) -> list[str]:
"""Check an optional narrowing of a workload's benchmark run.
Some benchmarks parametrise one dtype per case, so the campaign-wide dtype
filter makes every non-matching case raise. A workload may pin the case it
measures, but narrowing what gets timed always requires a stated reason.
"""
name = row.get("id", index)
selector = row.get("benchmark_selector")
if selector is None:
return []
errors: list[str] = []
if not isinstance(selector, str) or not selector.strip():
errors.append(f"workload {name} benchmark_selector must be a non-empty string")
reason = row.get("benchmark_scope_reason")
if not isinstance(reason, str) or not reason.strip():
errors.append(
f"workload {name} narrows the benchmark and must state "
"benchmark_scope_reason"
)
return errors
def _inventory_override_errors(row: dict[str, Any], index: int) -> list[str]:
"""Check an explicit departure from the pinned zl_bench inventory.
The inventory can lag the pinned FlagGems commit, for example when upstream
splits one benchmark file into several. A workload may then name the target
that actually exists, but it must say why, so a stale inventory is never
silently routed around.
"""
name = row.get("id", index)
reason = row.get("inventory_override_reason")
if reason is None:
return []
if not isinstance(reason, str) or not reason.strip():
return [
f"workload {name} inventory_override_reason must be a non-empty string"
]
return []
def validate_manifest(manifest: dict[str, Any]) -> list[str]:
errors: list[str] = []
if manifest.get("schema_version") != 1:
errors.append("schema_version must be 1")
if manifest.get("experiment_kind") != "musa-rlc-product-validation":
errors.append("experiment_kind must be musa-rlc-product-validation")
if manifest.get("production_default_policy") != "unchanged-off":
errors.append("production_default_policy must be unchanged-off")
if "S5000 MUSA product-dispatch" not in manifest.get("permitted_claim", ""):
errors.append("permitted_claim must bind S5000 MUSA product dispatch")
source = manifest.get("source")
if not isinstance(source, dict):
errors.append("source must be an object")
else:
for key in ("flagtree_commit", "flaggems_commit", "zl_bench_commit"):
if not _is_git_sha(source.get(key)):
errors.append(f"source.{key} must be a lowercase 40-character Git SHA")
device = manifest.get("device")
expected_device = {
"ssh_alias": "S5000",
"container": "zixu_flagos_S5000",
"backend": "mthreads",
"model": "Moore Threads S5000",
"physical_device": "7",
"logical_device": "0",
"visible_device_count": 1,
}
if not isinstance(device, dict):
errors.append("device must be an object")
else:
for key, expected in expected_device.items():
if device.get(key) != expected:
errors.append(f"device.{key} must equal {expected!r}")
protocol = manifest.get("protocol")
if not isinstance(protocol, dict):
errors.append("protocol must be an object")
else:
if protocol.get("screen_cells") != CELLS:
errors.append("protocol.screen_cells must be off/mask0/candidate")
if protocol.get("timing_cells") != TIMING_CELLS:
errors.append("protocol.timing_cells must be off/candidate")
if protocol.get("candidate_phase_mask") != 13:
errors.append("candidate_phase_mask must be 13")
canary_order = protocol.get("canary_order")
if not isinstance(canary_order, list) or len(canary_order) != 8:
errors.append("protocol.canary_order must contain eight workload ids")
for flag in (
"candidate_profitability_policy",
"fresh_process_per_operator_cell",
"same_physical_device_for_all_cells",
"serial_measurement",
"full_campaign_requires_canary_terminal",
):
if protocol.get(flag) is not True:
errors.append(f"protocol.{flag} must be true")
if protocol.get("timer") != "triton.testing.do_bench via MUSA device-interface Event":
errors.append("protocol.timer must bind the MUSA native-event path")
default_dtypes = protocol.get("default_benchmark_dtypes")
if not isinstance(default_dtypes, list) or not default_dtypes or not all(
isinstance(value, str) and value for value in default_dtypes
):
errors.append("protocol.default_benchmark_dtypes must be non-empty strings")
supported = protocol.get("supported_benchmark_dtypes")
if not isinstance(supported, list) or not supported or not all(
isinstance(value, str) and value for value in supported
):
errors.append(
"protocol.supported_benchmark_dtypes must be non-empty strings"
)
timeouts = protocol.get("timeout_seconds_by_cost")
if not isinstance(timeouts, dict) or set(timeouts) != COSTS:
errors.append("timeout_seconds_by_cost must define light/medium/heavy")
elif not all(isinstance(value, int) and value > 0 for value in timeouts.values()):
errors.append("timeout_seconds_by_cost values must be positive integers")
floor = protocol.get("measured_launch_floor")
if not isinstance(floor, int) or isinstance(floor, bool) or floor < 0:
errors.append(
"protocol.measured_launch_floor must be a non-negative integer"
)
workloads = manifest.get("workloads")
if not isinstance(workloads, list) or len(workloads) != 32:
errors.append("workloads must contain exactly 32 entries")
return errors
ids: list[str] = []
nodes: list[tuple[str, str, str]] = []
canaries = 0
for index, row in enumerate(workloads):
if not isinstance(row, dict):
errors.append(f"workloads[{index}] must be an object")
continue
for key in (
"id", "repo", "op", "op_file", "correctness_file", "family",
"nv_evidence", "prior_s5000", "cost",
):
if not isinstance(row.get(key), str) or not row[key]:
errors.append(f"workloads[{index}] missing {key}")
if row.get("repo") != "main":
errors.append(f"workload {row.get('id', index)} must use the main FlagGems repo")
if row.get("cost") not in COSTS:
errors.append(f"workload {row.get('id', index)} has invalid cost")
if not isinstance(row.get("canary"), bool):
errors.append(f"workload {row.get('id', index)} canary must be boolean")
benchmark_dtypes = row.get(
"benchmark_dtypes", protocol.get("default_benchmark_dtypes", [])
)
if not isinstance(benchmark_dtypes, list) or not benchmark_dtypes or not all(
isinstance(value, str) and value for value in benchmark_dtypes
):
errors.append(
f"workload {row.get('id', index)} benchmark_dtypes must be non-empty strings"
)
else:
# The benchmark CLI rejects a dtype it does not define, and that
# only surfaced as an argparse error inside a device run.
vocabulary = protocol.get("supported_benchmark_dtypes")
if isinstance(vocabulary, list):
unknown = [
value for value in benchmark_dtypes if value not in vocabulary
]
if unknown:
errors.append(
f"workload {row.get('id', index)} declares benchmark dtype(s) "
f"the benchmark CLI does not accept: {', '.join(unknown)}"
)
errors.extend(_correctness_scope_errors(row, index))
errors.extend(_benchmark_scope_errors(row, index))
errors.extend(_inventory_override_errors(row, index))
canaries += row.get("canary") is True
ids.append(row.get("id", ""))
nodes.append((row.get("repo", ""), row.get("op", ""), row.get("op_file", "")))
if len(ids) != len(set(ids)):
errors.append("workload ids must be unique")
if len(nodes) != len(set(nodes)):
errors.append("workload repo/op/op_file triples must be unique")
if canaries != 8:
errors.append("exactly eight workloads must be canaries")
elif set(manifest["protocol"].get("canary_order", [])) != {
row["id"] for row in workloads if row["canary"]
}:
errors.append("protocol.canary_order must contain every canary exactly once")
return errors
def _load_inventory(inventory_dir: Path) -> dict[tuple[str, str, str], dict[str, str]]:
path = inventory_dir / "inventory_main.csv"
if not path.is_file():
raise CampaignError(f"missing inventory: {path}")
rows: dict[tuple[str, str, str], dict[str, str]] = {}
with path.open(newline="", encoding="utf-8") as handle:
for row in csv.DictReader(handle):
key = ("main", row.get("op", ""), row.get("op_file", ""))
if key in rows:
raise CampaignError(f"duplicate inventory row: {key}")
rows[key] = row
return rows
def validate_inventory(workloads: list[dict[str, Any]], inventory_dir: Path) -> list[str]:
inventory = _load_inventory(inventory_dir)
errors: list[str] = []
for row in workloads:
key = (row["repo"], row["op"], row["op_file"])
found = inventory.get(key)
if found is None:
# A stated override admits a target the pinned inventory does not
# know about. The runner still resolves it against the real tree.
if row.get("inventory_override_reason"):
continue
errors.append(f"inventory missing {key}")
continue
if found.get("collected") == "no":
errors.append(f"inventory marks {key} collected=no")
if found.get("marked_skip") == "yes":
errors.append(f"inventory marks {key} skipped")
return errors
def build_plan(manifest: dict[str, Any], selection: str) -> tuple[dict[str, Any], dict[str, Any]]:
workloads = manifest["workloads"]
if selection == "canary":
by_id = {row["id"]: row for row in workloads if row["canary"]}
workloads = [by_id[task_id] for task_id in manifest["protocol"]["canary_order"]]
elif selection != "all":
raise CampaignError(f"unsupported selection: {selection}")
timeouts = manifest["protocol"]["timeout_seconds_by_cost"]
tasks = []
ordinal = 0
for row in workloads:
ordinal += 1
tasks.append(
{
"ordinal": ordinal,
"task_id": row["id"],
"repo": row["repo"],
"op": row["op"],
"op_file": row["op_file"],
"shape_key": row.get("shape_key"),
"correctness_file": row["correctness_file"],
"correctness_selector": row.get("correctness_selector"),
"correctness_node_ids": row.get("correctness_node_ids"),
"correctness_scope_reason": row.get("correctness_scope_reason"),
"benchmark_selector": row.get("benchmark_selector"),
"benchmark_scope_reason": row.get("benchmark_scope_reason"),
"inventory_override_reason": row.get("inventory_override_reason"),
"measured_launch_floor": manifest["protocol"]["measured_launch_floor"],
"family": row["family"],
"nv_evidence": row["nv_evidence"],
"nv_gain_percent": row.get("nv_gain_percent"),
"prior_s5000": row["prior_s5000"],
"cost": row["cost"],
"benchmark_dtypes": row.get(
"benchmark_dtypes",
manifest["protocol"]["default_benchmark_dtypes"],
),
"timeout_seconds": timeouts[row["cost"]],
"estimated_seconds_per_cell": ESTIMATED_SECONDS[row["cost"]],
"cells": CELLS,
"status": "PENDING",
}
)
estimate = sum(task["estimated_seconds_per_cell"] * len(CELLS) for task in tasks)
matrix = {
"schema_version": 1,
"campaign_id": manifest["campaign_id"],
"selection": selection,
"device": manifest["device"],
"screen_cells": CELLS,
"timing_cells": TIMING_CELLS,
"tasks": tasks,
}
budget = {
"schema_version": 1,
"campaign_id": manifest["campaign_id"],
"selection": selection,
"task_count": len(tasks),
"gpu_count": 1,
"serial_measurement": True,
"estimated_wall_seconds": estimate,
"estimated_gpu_hours": round(estimate / 3600.0, 4),
"note": "Scheduling estimate only; it is not performance evidence.",
}
return matrix, budget
def validate_canary_receipt(
manifest: dict[str, Any], selection: str, canary_receipt: Path | None
) -> list[str]:
if selection != "all":
return []
if canary_receipt is None:
return ["full campaign requires --canary-receipt"]
try:
receipt = json.loads(canary_receipt.read_text(encoding="utf-8"))
except (OSError, json.JSONDecodeError) as exc:
return [f"cannot read canary receipt: {exc}"]
errors = []
if receipt.get("campaign_id") != manifest.get("campaign_id"):
errors.append("canary receipt campaign_id mismatch")
if receipt.get("selection") != "canary" or receipt.get("status") != "TERMINAL":
errors.append("canary receipt must be a terminal canary")
if receipt.get("canary_gate_pass") is not True:
errors.append("canary receipt does not authorize the full campaign")
return errors
def prepare(
manifest_path: Path,
inventory_dir: Path,
output_root: Path,
selection: str,
canary_receipt: Path | None = None,
) -> Path:
manifest = json.loads(manifest_path.read_text(encoding="utf-8"))
if not isinstance(manifest, dict):
raise CampaignError("manifest root must be an object")
errors = validate_manifest(manifest)
if not errors:
errors.extend(validate_inventory(manifest["workloads"], inventory_dir))
if not errors:
errors.extend(validate_canary_receipt(manifest, selection, canary_receipt))
if errors:
raise CampaignError("campaign validation failed:\n- " + "\n- ".join(errors))
matrix, budget = build_plan(manifest, selection)
output = output_root.resolve() / f"{manifest['campaign_id']}--{selection}"
if output.exists():
raise CampaignError(f"campaign output already exists and is immutable: {output}")
output.parent.mkdir(parents=True, exist_ok=True)
stage = output.parent / f".{output.name}.staging-{os.getpid()}"
if stage.exists():
raise CampaignError(f"staging directory already exists: {stage}")
try:
stage.mkdir()
shutil.copy2(manifest_path, stage / "manifest.json")
manifest_sha = hashlib.sha256((stage / "manifest.json").read_bytes()).hexdigest()
matrix["manifest_sha256"] = manifest_sha
budget["manifest_sha256"] = manifest_sha
(stage / "manifest.sha256").write_text(
f"{manifest_sha} manifest.json\n", encoding="utf-8", newline="\n"
)
(stage / "task-matrix.json").write_bytes(_canonical(matrix))
(stage / "budget.json").write_bytes(_canonical(budget))
(stage / "ops.txt").write_text(
"".join(f"{task['repo']}:{task['op']}@{task['op_file']}\n" for task in matrix["tasks"]),
encoding="utf-8",
newline="\n",
)
stage.replace(output)
except Exception:
if stage.exists():
shutil.rmtree(stage)
raise
return output
def main() -> int:
parser = argparse.ArgumentParser()
parser.add_argument("manifest", type=Path)
parser.add_argument("--inventory-dir", type=Path, required=True)
parser.add_argument("--selection", choices=("canary", "all"), default="canary")
parser.add_argument("--output-root", type=Path)
parser.add_argument("--canary-receipt", type=Path)
parser.add_argument("--dry-run", action="store_true")
args = parser.parse_args()
try:
manifest = json.loads(args.manifest.read_text(encoding="utf-8"))
errors = validate_manifest(manifest) if isinstance(manifest, dict) else ["manifest root must be an object"]
if not errors:
errors.extend(validate_inventory(manifest["workloads"], args.inventory_dir))
if not errors and not args.dry_run:
errors.extend(validate_canary_receipt(manifest, args.selection, args.canary_receipt))
if errors:
raise CampaignError("campaign validation failed:\n- " + "\n- ".join(errors))
matrix, budget = build_plan(manifest, args.selection)
if args.dry_run:
print(json.dumps({"status": "DRY_RUN", "budget": budget, "tasks": matrix["tasks"]}, indent=2, sort_keys=True))
return 0
if args.output_root is None:
raise CampaignError("--output-root is required unless --dry-run is used")
output = prepare(
args.manifest,
args.inventory_dir,
args.output_root,
args.selection,
args.canary_receipt,
)
except (OSError, json.JSONDecodeError, CampaignError) as exc:
print(str(exc))
return 2
print(json.dumps({"status": "PREPARED", "output": str(output)}, indent=2))
return 0
if __name__ == "__main__":
raise SystemExit(main())
@@ -1,223 +0,0 @@
#!/usr/bin/env python3
"""Re-account a sealed MUSA RLC run under measured-operator convert scoping.
The sealed campaign counted every executed kernel's ``ttg.convert_layout`` ops,
including the kernels a benchmark launches once to build its inputs. Those sit
outside the timed region, so a candidate could be credited with a topology
reduction that no timer could ever observe.
This tool re-reads the sealed TTGIR dumps and their launch index, recomputes the
convert counts scoped to repeatedly launched kernels, and re-runs the campaign
classifier. It never writes inside the sealed run.
A re-accounted verdict can only demote a topology claim. Promoting a task to G
or C needs reversed-order timing, which a sealed run only holds for tasks that
already showed a reduction, so a newly appearing reduction is reported as X for
missing round-b evidence rather than guessed.
"""
from __future__ import annotations
import argparse
import csv
import importlib.util
import json
import sys
from pathlib import Path
from typing import Any
HERE = Path(__file__).resolve().parent
class ReaccountError(RuntimeError):
pass
def _load_runner():
spec = importlib.util.spec_from_file_location(
"run_musa_rlc_product", HERE / "run_musa_rlc_product.py"
)
if spec is None or spec.loader is None:
raise ReaccountError("cannot load run_musa_rlc_product.py")
module = importlib.util.module_from_spec(spec)
spec.loader.exec_module(module)
return module
def _canonical(value: object) -> bytes:
return (json.dumps(value, indent=2, sort_keys=True) + "\n").encode("utf-8")
def reaccount_task(
runner,
task_root: Path,
receipt: dict[str, Any],
measured_launch_floor: int,
) -> dict[str, Any]:
"""Recompute one sealed task receipt under the scoped convert count."""
cells: dict[str, dict[str, Any]] = {}
for key, sealed in receipt.get("cells", {}).items():
cell = dict(sealed)
cell_dir = task_root / str(sealed.get("round")) / str(sealed.get("cell"))
if cell_dir.is_dir():
scope = runner._convert_count(
cell_dir, measured_launch_floor=measured_launch_floor
)
cell.update(
{
"convert_count": scope["measured_convert_count"],
"convert_count_all": scope["total_convert_count"],
"convert_scope_complete": scope["scope_complete"],
"measured_launch_floor": scope["measured_launch_floor"],
"convert_excluded_unmeasured": scope["excluded_unmeasured"],
"convert_unindexed_files": scope["unindexed_files"],
}
)
else:
cell["convert_scope_complete"] = False
cells[key] = cell
classification, detail = runner.classify_task(
receipt["correctness"], cells, receipt.get("correctness_isolation")
)
sealed_class = receipt.get("classification")
excluded = sorted(
{
row["file"]
for cell in cells.values()
for row in cell.get("convert_excluded_unmeasured", [])
}
)
return {
"task_id": receipt["task_id"],
"op": receipt.get("op"),
"family": receipt.get("family"),
"sealed_classification": sealed_class,
"reaccounted_classification": classification,
"changed": classification != sealed_class,
"sealed_reason": receipt.get("classification_detail", {}).get("reason"),
"reaccounted_reason": detail["reason"],
"sealed_convert": {
key: receipt["cells"][key].get("convert_count")
for key in sorted(receipt.get("cells", {}))
},
"reaccounted_convert": {
key: cells[key].get("convert_count") for key in sorted(cells)
},
"excluded_unmeasured_files": excluded,
"classification_detail": detail,
}
def reaccount(
run_root: Path, output_root: Path, measured_launch_floor: int
) -> dict[str, Any]:
run_root = run_root.resolve()
output_root = output_root.resolve()
if not (run_root / "campaign-receipt.json").is_file():
raise ReaccountError(f"not a sealed campaign run root: {run_root}")
if output_root == run_root or run_root in output_root.parents:
raise ReaccountError("refusing to write inside the sealed run root")
runner = _load_runner()
sealed = json.loads((run_root / "campaign-receipt.json").read_text(encoding="utf-8"))
rows: list[dict[str, Any]] = []
for path in sorted((run_root / "tasks").glob("*/task-receipt.json")):
receipt = json.loads(path.read_text(encoding="utf-8"))
rows.append(
reaccount_task(runner, path.parent, receipt, measured_launch_floor)
)
terminal = sorted(runner.TERMINAL)
result = {
"schema_version": 1,
"kind": "musa-rlc-topology-reaccount",
"sealed_campaign_id": sealed.get("campaign_id"),
"sealed_run_root": str(run_root),
"measured_launch_floor": measured_launch_floor,
"task_count": len(rows),
"sealed_counts": {
key: sum(row["sealed_classification"] == key for row in rows)
for key in terminal
},
"reaccounted_counts": {
key: sum(row["reaccounted_classification"] == key for row in rows)
for key in terminal
},
"changed_task_ids": [row["task_id"] for row in rows if row["changed"]],
"note": (
"Derived accounting over sealed artifacts. It supersedes no receipt "
"and produces no new device measurement."
),
"tasks": rows,
}
output_root.mkdir(parents=True, exist_ok=True)
(output_root / "reaccount-receipt.json").write_bytes(_canonical(result))
with (output_root / "reaccount-matrix.csv").open(
"w", newline="", encoding="utf-8"
) as handle:
writer = csv.writer(handle)
writer.writerow(
[
"task_id",
"sealed_class",
"reaccounted_class",
"changed",
"sealed_off",
"sealed_candidate",
"scoped_off",
"scoped_candidate",
"reaccounted_reason",
]
)
for row in rows:
writer.writerow(
[
row["task_id"],
row["sealed_classification"],
row["reaccounted_classification"],
"yes" if row["changed"] else "no",
row["sealed_convert"].get("a_off"),
row["sealed_convert"].get("a_candidate"),
row["reaccounted_convert"].get("a_off"),
row["reaccounted_convert"].get("a_candidate"),
row["reaccounted_reason"],
]
)
return result
def main() -> int:
parser = argparse.ArgumentParser()
parser.add_argument("--run-root", type=Path, required=True)
parser.add_argument("--output-root", type=Path, required=True)
parser.add_argument("--measured-launch-floor", type=int, default=1)
args = parser.parse_args()
try:
if args.measured_launch_floor < 0:
raise ReaccountError("--measured-launch-floor must be non-negative")
result = reaccount(
args.run_root, args.output_root, args.measured_launch_floor
)
except (OSError, json.JSONDecodeError, ReaccountError) as exc:
print(f"MUSA_RLC_REACCOUNT_ERROR: {exc}", file=sys.stderr)
return 2
summary = {
key: result[key]
for key in (
"sealed_campaign_id",
"measured_launch_floor",
"task_count",
"sealed_counts",
"reaccounted_counts",
"changed_task_ids",
)
}
print(json.dumps(summary, indent=2, sort_keys=True))
return 0
if __name__ == "__main__":
raise SystemExit(main())
@@ -1,834 +0,0 @@
#!/usr/bin/env python3
"""Run the frozen S5000 MUSA RLC product campaign serially."""
from __future__ import annotations
import argparse
import csv
import hashlib
import importlib.util
import json
import math
import os
import signal
import statistics
import subprocess
import sys
import time
from pathlib import Path
from typing import Any
TERMINAL = {"G", "C", "B", "R", "K", "X"}
RLC_ENV_KEYS = {
"FLAGTREE_RLC_ENHANCE",
"FLAGTREE_RLC_PHASE_MASK",
"FLAGTREE_MUSA_RLC_ENHANCE",
"FLAGTREE_MUSA_RLC_PHASE_MASK",
"FLAGTREE_MUSA_RLC_PROFITABILITY_POLICY",
"FLAGTREE_MUSA_RLC_PRODUCT_LAUNCH_COUNT",
"FLAGTREE_MUSA_RLC_MIN_ADJUSTED_SAVED_COST_PER_TENSOR_OP",
"FLAGTREE_MUSA_RLC_PHASE3_SAVED_COST_MULTIPLIER",
"FLAGTREE_MUSA_RLC_MAX_EXTERNAL_USE_EDGES",
"FLAGTREE_MUSA_RLC_MIN_REMOVED_CONVERT_DENSITY_PER_1024_PROPOSAL_VALUES",
"FLAGTREE_MUSA_RLC_PRESERVE_INT_TO_FP_CONTIGUITY",
"FLAGTREE_MUSA_RLC_PRESERVE_FP_TO_FP_CONTIGUITY",
}
class CampaignError(RuntimeError):
pass
def _canonical(value: object) -> bytes:
return (json.dumps(value, indent=2, sort_keys=True) + "\n").encode("utf-8")
def _sha256(path: Path) -> str:
digest = hashlib.sha256()
with path.open("rb") as handle:
for chunk in iter(lambda: handle.read(1024 * 1024), b""):
digest.update(chunk)
return digest.hexdigest()
def _atomic_json(path: Path, value: object) -> None:
temporary = path.with_name(path.name + ".tmp")
temporary.write_bytes(_canonical(value))
os.replace(temporary, path)
def _load_module(path: Path, name: str):
spec = importlib.util.spec_from_file_location(name, path)
if spec is None or spec.loader is None:
raise CampaignError(f"cannot load module: {path}")
module = importlib.util.module_from_spec(spec)
spec.loader.exec_module(module)
return module
def _identity(path: Path) -> dict[str, str]:
rows: dict[str, str] = {}
for raw in path.read_text(encoding="utf-8").splitlines():
key, separator, value = raw.partition("=")
if separator:
rows[key] = value
return rows
def _base_env(site: Path, flaggems: Path) -> dict[str, str]:
env = os.environ.copy()
for key in list(env):
if key.startswith("FLAGTREE_") and "RLC" in key:
env.pop(key)
env.pop("DISABLE_LLVM_OPT", None)
env.update(
{
"PYTHONNOUSERSITE": "1",
"PYTHONDONTWRITEBYTECODE": "1",
"PYTHONUNBUFFERED": "1",
"GEMS_VENDOR": "mthreads",
"FLAGGEMS_VENDOR": "mthreads",
"FLAGTREE_BACKEND": "mthreads",
"TRITON_BACKENDS_IN_TREE": "1",
"MUSA_VISIBLE_DEVICES": "7",
"MTHREADS_VISIBLE_DEVICES": "7",
"PYTHONPATH": f"{flaggems / 'src'}:{site}",
"FLAGTREE_RLC_TRACE_REJECTS": "1",
}
)
return env
def _cell_env(base: dict[str, str], cell: str) -> dict[str, str]:
env = base.copy()
configs = {
"off": {"enhance": "0", "mask": "13", "policy": "0", "launches": "0"},
"mask0": {"enhance": "1", "mask": "0", "policy": "1", "launches": "1"},
"candidate": {"enhance": "1", "mask": "13", "policy": "1", "launches": "1"},
}
try:
row = configs[cell]
except KeyError as exc:
raise CampaignError(f"unsupported cell: {cell}") from exc
env.update(
{
"FLAGTREE_MUSA_RLC_ENHANCE": row["enhance"],
"FLAGTREE_MUSA_RLC_PHASE_MASK": row["mask"],
"FLAGTREE_MUSA_RLC_PROFITABILITY_POLICY": row["policy"],
"FLAGTREE_MUSA_RLC_PRODUCT_LAUNCH_COUNT": row["launches"],
"FLAGTREE_MUSA_RLC_MIN_ADJUSTED_SAVED_COST_PER_TENSOR_OP": "1800",
"FLAGTREE_MUSA_RLC_PHASE3_SAVED_COST_MULTIPLIER": "2",
"FLAGTREE_MUSA_RLC_MAX_EXTERNAL_USE_EDGES": "8",
"FLAGTREE_MUSA_RLC_MIN_REMOVED_CONVERT_DENSITY_PER_1024_PROPOSAL_VALUES": "0",
"FLAGTREE_MUSA_RLC_PRESERVE_INT_TO_FP_CONTIGUITY": "1",
"FLAGTREE_MUSA_RLC_PRESERVE_FP_TO_FP_CONTIGUITY": "1",
}
)
return env
def _run_process(
command: list[str], *, cwd: Path, env: dict[str, str], log: Path, timeout: int
) -> tuple[int, bool, float]:
started = time.time()
with log.open("w", encoding="utf-8", newline="\n") as handle:
process = subprocess.Popen(
command,
cwd=str(cwd),
env=env,
stdout=handle,
stderr=subprocess.STDOUT,
text=True,
start_new_session=True,
)
timed_out = False
try:
return_code = process.wait(timeout=timeout)
except subprocess.TimeoutExpired:
timed_out = True
try:
os.killpg(process.pid, signal.SIGTERM)
process.wait(timeout=10)
except subprocess.TimeoutExpired:
os.killpg(process.pid, signal.SIGKILL)
process.wait()
return_code = 124
return return_code, timed_out, time.time() - started
def _indexed_launches(ttgir_root: Path) -> dict[Path, int] | None:
"""Map each dumped TTGIR file to its executed launch count.
Returns None when ``index.tsv`` is missing or unreadable, which means the
measured-operator scope cannot be established for this cell.
"""
index = ttgir_root / "index.tsv"
if not index.is_file():
return None
try:
lines = index.read_text(encoding="utf-8", errors="replace").splitlines()
except OSError:
return None
if not lines:
return None
header = lines[0].split("\t")
try:
shape_at = header.index("shape")
kernel_at = header.index("kernel")
file_at = header.index("file")
launch_at = header.index("launches")
except ValueError:
return None
width = max(shape_at, kernel_at, file_at, launch_at)
launches: dict[Path, int] = {}
for raw in lines[1:]:
fields = raw.split("\t")
if len(fields) <= width:
continue
name = fields[file_at].strip()
# Autotune losers are indexed for audit but never dumped; they carry a
# parenthesised placeholder instead of a file name.
if not name or name.startswith("("):
continue
try:
count = int(fields[launch_at])
except ValueError:
continue
path = ttgir_root / fields[shape_at] / fields[kernel_at] / name
launches[path] = max(launches.get(path, 0), count)
return launches
def _convert_count(path: Path, *, measured_launch_floor: int = 1) -> dict[str, Any]:
"""Count executed convert_layout ops, scoped to the measured operator.
A benchmark call constructs its inputs once and then runs the operator many
times under the timer. Counting every executed kernel therefore credits the
candidate with converts removed from input construction, outside the timed
region. Only variants launched more than ``measured_launch_floor`` times
contribute to the measured count; the unscoped total is preserved too.
"""
ttgir_root = path / "ttgir"
files = sorted(ttgir_root.glob("**/*.ttgir"))
launches = _indexed_launches(ttgir_root)
total = 0
measured = 0
excluded: list[dict[str, Any]] = []
unindexed: list[str] = []
for file in files:
converts = file.read_text(encoding="utf-8", errors="replace").count(
"ttg.convert_layout"
)
total += converts
relative = file.relative_to(ttgir_root).as_posix()
if launches is None:
measured += converts
continue
launch_count = launches.get(file)
if launch_count is None:
# Dumped but absent from the index: keep it and report the gap
# rather than silently dropping executed topology.
measured += converts
unindexed.append(relative)
continue
if launch_count > measured_launch_floor:
measured += converts
elif converts:
excluded.append(
{
"file": relative,
"launches": launch_count,
"convert_count": converts,
}
)
return {
"measured_convert_count": measured,
"total_convert_count": total,
"ttgir_files": len(files),
"measured_launch_floor": measured_launch_floor,
"scope_complete": launches is not None and not unindexed,
"excluded_unmeasured": excluded,
"unindexed_files": unindexed,
}
def _median_gems_latency(rows: list[dict[str, str]]) -> float | None:
values = []
for row in rows:
if row.get("row_status") != "SUCCESS":
continue
try:
value = float(row.get("gems_lat", ""))
except ValueError:
continue
if math.isfinite(value) and value > 0:
values.append(value)
return statistics.median(values) if values else None
def missing_targets(matrix: dict[str, Any], flaggems: Path) -> list[str]:
"""Resolve every task's benchmark and correctness target against the tree.
The inventory a matrix is built from can lag the pinned FlagGems commit, for
example when upstream splits one benchmark file into several. Before this
check, such a task only failed inside the worker's preflight, which exited
silently and left a zero-byte log and a cell receipt with nothing in it.
"""
missing: list[str] = []
for task in matrix["tasks"]:
benchmark = flaggems / "benchmark" / f"test_{task['op_file']}.py"
if not benchmark.is_file():
missing.append(f"{task['task_id']} benchmark -> {benchmark}")
correctness = flaggems / task["correctness_file"]
if not correctness.is_file():
missing.append(f"{task['task_id']} correctness -> {correctness}")
return missing
def validate_contract(
manifest_path: Path,
matrix_path: Path,
build_root: Path,
flaggems: Path,
zl_bench: Path,
python: Path,
) -> tuple[dict[str, Any], dict[str, Any], dict[str, Any]]:
manifest = json.loads(manifest_path.read_text(encoding="utf-8"))
matrix = json.loads(matrix_path.read_text(encoding="utf-8"))
if not isinstance(manifest, dict) or not isinstance(matrix, dict):
raise CampaignError("manifest and task matrix roots must be objects")
manifest_sha = _sha256(manifest_path)
if matrix.get("manifest_sha256") != manifest_sha:
raise CampaignError("task matrix manifest_sha256 does not match manifest")
if matrix.get("campaign_id") != manifest.get("campaign_id"):
raise CampaignError("campaign id mismatch")
if matrix.get("screen_cells") != ["off", "mask0", "candidate"]:
raise CampaignError("task matrix cell contract mismatch")
missing = missing_targets(matrix, flaggems)
if missing:
raise CampaignError(
"targets missing from the frozen FlagGems tree: " + "; ".join(missing)
)
site = build_root / "site"
identity = _identity(build_root / "identity.txt")
expected = manifest["source"]
if identity.get("git_commit") != expected["flagtree_commit"]:
raise CampaignError("FlagTree build identity does not match manifest")
if not (build_root / "build.done").is_file():
raise CampaignError("FlagTree product build is not terminal")
if (flaggems / ".source-commit").read_text(encoding="utf-8").strip() != expected["flaggems_commit"]:
raise CampaignError("FlagGems source identity does not match manifest")
if (zl_bench / ".source-commit").read_text(encoding="utf-8").strip() != expected["zl_bench_commit"]:
raise CampaignError("zl_bench source identity does not match manifest")
probe_env = _base_env(site, flaggems)
probe_code = (
"import json,torch,torch_musa,triton,flag_gems;"
"print(json.dumps({'count':torch.musa.device_count(),"
"'device':str(flag_gems.device),'vendor':str(flag_gems.vendor_name),"
"'triton':triton.__file__,'flaggems':flag_gems.__file__},sort_keys=True))"
)
probe = subprocess.run(
[str(python), "-B", "-c", probe_code],
cwd=str(flaggems),
env=probe_env,
text=True,
capture_output=True,
timeout=60,
)
if probe.returncode:
raise CampaignError(f"product import probe failed: {probe.stderr.strip()}")
observed = json.loads(probe.stdout.strip().splitlines()[-1])
if observed.get("count") != 1 or observed.get("vendor") != "mthreads":
raise CampaignError(f"MUSA device/vendor mismatch: {observed}")
if not Path(observed["triton"]).is_relative_to(site):
raise CampaignError("imported Triton does not come from the frozen product site")
if not Path(observed["flaggems"]).is_relative_to(flaggems / "src"):
raise CampaignError("imported FlagGems does not come from the frozen product source")
observed.update(
{
"flagtree_commit": identity["git_commit"],
"libtriton_sha256": _sha256(site / "triton/_C/libtriton.so"),
"flaggems_commit": expected["flaggems_commit"],
"zl_bench_commit": expected["zl_bench_commit"],
"python": str(python),
}
)
return manifest, matrix, observed
def run_cell(
task: dict[str, Any],
cell: str,
round_name: str,
replay_from: Path | None,
*,
output_root: Path,
worker: Path,
python: Path,
zl_bench: Path,
flaggems: Path,
site: Path,
shape_file: Path,
batch_summary: Any,
base_env: dict[str, str],
) -> dict[str, Any]:
out = output_root / "tasks" / task["task_id"] / round_name / cell
env = _cell_env(base_env, cell)
env.update(
{
"PYTHON_BIN": str(python),
"ZL_BENCH_DIR": str(zl_bench),
"FLAGGEMS_DIR": str(flaggems),
"FLAGTREE_SITE": str(site),
"OP": task["op"],
"OP_FILE": task["op_file"],
"SHAPE_FILE": str(shape_file),
"OUT_DIR": str(out),
"DTYPES": " ".join(task["benchmark_dtypes"]),
}
)
if task.get("benchmark_selector"):
env["BENCH_SELECTOR"] = str(task["benchmark_selector"])
if replay_from is not None:
env["REPLAY_FROM"] = str(replay_from)
command = ["bash", str(worker)]
# The shell worker owns run.log, so use a sibling controller log.
controller_log = out.parent / f"{cell}.controller.log"
controller_log.parent.mkdir(parents=True, exist_ok=True)
return_code, timed_out, elapsed = _run_process(
command,
cwd=worker.parent,
env=env,
log=controller_log,
timeout=int(task["timeout_seconds"]),
)
run_log = out / "run.log"
text = run_log.read_text(encoding="utf-8", errors="replace") if run_log.is_file() else ""
rows = batch_summary.parse_run_log(text)
if timed_out:
status = "TIMEOUT"
else:
status = batch_summary.classify(return_code, rows, text)
scope = _convert_count(
out, measured_launch_floor=int(task.get("measured_launch_floor", 1))
)
# A worker can fail its preflight before creating OUT_DIR (for example,
# when a frozen inventory entry has no backend benchmark file). Preserve
# that failure as a cell receipt instead of aborting the whole campaign.
out.mkdir(parents=True, exist_ok=True)
receipt = {
"task_id": task["task_id"],
"round": round_name,
"cell": cell,
"status": status,
"return_code": return_code,
"timed_out": timed_out,
"elapsed_seconds": round(elapsed, 3),
"rows": rows,
"rows_success": sum(row.get("row_status") == "SUCCESS" for row in rows),
"median_gems_latency_ms": _median_gems_latency(rows),
"convert_count": scope["measured_convert_count"],
"convert_count_all": scope["total_convert_count"],
"convert_scope_complete": scope["scope_complete"],
"measured_launch_floor": scope["measured_launch_floor"],
"convert_excluded_unmeasured": scope["excluded_unmeasured"],
"convert_unindexed_files": scope["unindexed_files"],
"ttgir_files": scope["ttgir_files"],
"replay_fallback_count": text.count("AUTOTUNE_REPLAY_FALLBACK"),
"benchmark_selector": task.get("benchmark_selector"),
"benchmark_dtypes": list(task["benchmark_dtypes"]),
"effective_rlc": {
key: env[key]
for key in sorted(RLC_ENV_KEYS)
if key in env
},
}
_atomic_json(out / "cell-receipt.json", receipt)
return receipt
def correctness_command(task: dict[str, Any], python: Path) -> list[str]:
"""Build the correctness argv, narrowed to the variant this task measures.
A task-level file runs every variant in that file, so an unrelated failing
variant can mask or fabricate a verdict for the operator under test. A
manifest may therefore pin explicit node ids or a ``-k`` selector; the
manifest schema requires a stated reason whenever it does.
"""
command = [str(python), "-B", "-m", "pytest", "-q"]
# The same vendor-scoped compatibility plugin the benchmark worker loads.
# Without it a stale CUDA-only guard skips the operator here too, and a
# SKIP would be read as "not measured" rather than "not gated".
command.extend(["-p", "_musa_benchmark_compat_plugin"])
node_ids = task.get("correctness_node_ids")
if node_ids:
command.extend(str(node) for node in node_ids)
else:
command.append(task["correctness_file"])
selector = task.get("correctness_selector")
if selector:
command.extend(["-k", str(selector)])
# -rs records why a variant was skipped, so a SKIP verdict stays auditable
# without a second run.
command.extend(["-rs", "--quick"])
return command
def run_correctness(
task: dict[str, Any],
cell: str,
*,
output_root: Path,
python: Path,
flaggems: Path,
site: Path,
base_env: dict[str, str],
plugin_dir: Path,
) -> dict[str, Any]:
out = output_root / "correctness" / task["task_id"] / cell
out.mkdir(parents=True, exist_ok=True)
env = _cell_env(base_env, cell)
env["PYTEST_DISABLE_PLUGIN_AUTOLOAD"] = "1"
env["PYTHONPATH"] = f"{plugin_dir}{os.pathsep}{env['PYTHONPATH']}"
env["TRITON_CACHE_DIR"] = str(out / "triton-cache")
env["FLAGGEMS_CACHE_DIR"] = str(out / "flaggems-cache")
(out / "triton-cache").mkdir()
(out / "flaggems-cache").mkdir()
command = correctness_command(task, python)
return_code, timed_out, elapsed = _run_process(
command,
cwd=flaggems,
env=env,
log=out / "correctness.log",
timeout=int(task["timeout_seconds"]),
)
log_text = (out / "correctness.log").read_text(encoding="utf-8", errors="replace")
if timed_out:
status = "TIMEOUT"
elif return_code != 0:
status = "FAIL"
elif " passed" in log_text:
status = "PASS"
elif " skipped" in log_text:
status = "SKIP"
else:
status = "FAIL"
receipt = {
"task_id": task["task_id"],
"correctness_file": task["correctness_file"],
"correctness_node_ids": task.get("correctness_node_ids"),
"correctness_selector": task.get("correctness_selector"),
"correctness_scope_reason": task.get("correctness_scope_reason"),
"command": command,
"cell": cell,
"status": status,
"return_code": return_code,
"timed_out": timed_out,
"elapsed_seconds": round(elapsed, 3),
}
_atomic_json(out / "correctness-receipt.json", receipt)
return receipt
def _speedup(off: dict[str, Any], candidate: dict[str, Any]) -> float | None:
left = off.get("median_gems_latency_ms")
right = candidate.get("median_gems_latency_ms")
if not isinstance(left, (int, float)) or not isinstance(right, (int, float)) or right <= 0:
return None
return (left / right - 1.0) * 100.0
def classify_task(
correctness: dict[str, Any],
cells: dict[str, dict[str, Any]],
correctness_isolation: dict[str, dict[str, Any]] | None = None,
) -> tuple[str, dict[str, Any]]:
detail: dict[str, Any] = {"reason": ""}
if correctness["status"] == "FAIL":
isolation = correctness_isolation or {}
off_status = isolation.get("off", {}).get("status")
mask0_status = isolation.get("mask0", {}).get("status")
detail.update(
{
"off_correctness_status": off_status,
"mask0_correctness_status": mask0_status,
"candidate_correctness_status": correctness["status"],
}
)
if off_status == "FAIL" and mask0_status == "FAIL":
detail["reason"] = "baseline correctness fails under off/mask0/candidate"
return "X", detail
detail["reason"] = "candidate correctness failed"
return "K", detail
if correctness["status"] != "PASS":
detail["reason"] = f"candidate correctness is {correctness['status']}"
return "X", detail
required = ("a_off", "a_mask0", "a_candidate")
if any(cells.get(key, {}).get("status") != "PASS" for key in required):
detail["reason"] = "one or more screen cells did not pass"
return "X", detail
if any(cells[key].get("replay_fallback_count") for key in ("a_mask0", "a_candidate")):
detail["reason"] = "same-config replay fell back"
return "X", detail
# A topology claim is only admissible when every screen cell could be scoped
# to the measured operator. Older receipts predate scoping and are trusted.
if any(cells[key].get("convert_scope_complete", True) is False for key in required):
detail["reason"] = "measured convert topology scope is incomplete"
return "X", detail
off_count = cells["a_off"]["convert_count"]
mask0_count = cells["a_mask0"]["convert_count"]
candidate_count = cells["a_candidate"]["convert_count"]
detail.update(
{
"off_convert_count": off_count,
"mask0_convert_count": mask0_count,
"candidate_convert_count": candidate_count,
"off_convert_count_all": cells["a_off"].get("convert_count_all"),
"candidate_convert_count_all": cells["a_candidate"].get("convert_count_all"),
"round_a_speedup_percent": _speedup(cells["a_off"], cells["a_candidate"]),
}
)
if off_count != mask0_count:
detail["reason"] = "mask0 topology differs from off"
return "X", detail
if candidate_count >= off_count:
speedup = detail["round_a_speedup_percent"]
if speedup is not None and speedup <= -3.0:
detail["reason"] = "no topology reduction and screening regression"
return "R", detail
detail["reason"] = "candidate did not reduce executed convert topology"
return "B", detail
if cells.get("b_off", {}).get("status") != "PASS" or cells.get("b_candidate", {}).get("status") != "PASS":
detail["reason"] = "topology changed but reversed-order timing did not pass"
return "X", detail
if cells["b_off"].get("replay_fallback_count") or cells["b_candidate"].get("replay_fallback_count"):
detail["reason"] = "reversed-order timing replay fell back"
return "X", detail
round_b = _speedup(cells["b_off"], cells["b_candidate"])
round_a = detail["round_a_speedup_percent"]
detail["round_b_speedup_percent"] = round_b
if round_a is None or round_b is None:
detail["reason"] = "missing comparable native-event latency"
return "X", detail
if round_a >= 5.0 and round_b >= 5.0:
detail["reason"] = "topology reduction and stable >=5% speedup"
return "G", detail
if round_a <= -3.0 and round_b <= -3.0:
detail["reason"] = "topology reduction but stable regression"
return "R", detail
detail["reason"] = "topology reduction without stable >=5% speedup"
return "C", detail
def run_campaign(
manifest: dict[str, Any],
matrix: dict[str, Any],
observed: dict[str, Any],
*,
output_root: Path,
worker: Path,
python: Path,
zl_bench: Path,
flaggems: Path,
build_root: Path,
resume: bool,
) -> dict[str, Any]:
if output_root.exists() and not resume:
raise CampaignError(f"output exists; use --resume: {output_root}")
observed = dict(observed)
observed.update(
{
"benchmark_runner_sha256": _sha256(Path(__file__)),
"benchmark_worker_sha256": _sha256(worker),
"musa_benchmark_compat_plugin_sha256": _sha256(
zl_bench / "_musa_benchmark_compat_plugin.py"
),
}
)
output_root.mkdir(parents=True, exist_ok=True)
site = build_root / "site"
shape_file = zl_bench / "ops/shapes_single.yaml"
batch_summary = _load_module(zl_bench / "batch_summary.py", "zl_batch_summary")
base_env = _base_env(site, flaggems)
status_csv = output_root / "status.csv"
rows: list[dict[str, Any]] = []
occupancy = subprocess.run(
["mthreads-gmi"], text=True, capture_output=True, timeout=30
)
if occupancy.returncode:
raise CampaignError(f"mthreads-gmi failed: {occupancy.stderr.strip()}")
(output_root / "occupancy-before.txt").write_text(occupancy.stdout, encoding="utf-8", newline="\n")
if "No running processes found" not in occupancy.stdout:
raise CampaignError("S5000 has an active process; campaign not admitted")
for task in matrix["tasks"]:
task_root = output_root / "tasks" / task["task_id"]
task_receipt_path = task_root / "task-receipt.json"
if resume and task_receipt_path.is_file():
receipt = json.loads(task_receipt_path.read_text(encoding="utf-8"))
if receipt.get("classification") in TERMINAL:
rows.append(receipt)
continue
cells: dict[str, dict[str, Any]] = {}
off_root = task_root / "round-a" / "off"
cells["a_off"] = run_cell(
task, "off", "round-a", None,
output_root=output_root, worker=worker, python=python,
zl_bench=zl_bench, flaggems=flaggems, site=site,
shape_file=shape_file, batch_summary=batch_summary, base_env=base_env,
)
cells["a_mask0"] = run_cell(
task, "mask0", "round-a", off_root,
output_root=output_root, worker=worker, python=python,
zl_bench=zl_bench, flaggems=flaggems, site=site,
shape_file=shape_file, batch_summary=batch_summary, base_env=base_env,
)
cells["a_candidate"] = run_cell(
task, "candidate", "round-a", off_root,
output_root=output_root, worker=worker, python=python,
zl_bench=zl_bench, flaggems=flaggems, site=site,
shape_file=shape_file, batch_summary=batch_summary, base_env=base_env,
)
correctness = run_correctness(
task, "candidate", output_root=output_root, python=python, flaggems=flaggems,
site=site, base_env=base_env, plugin_dir=zl_bench,
)
correctness_isolation: dict[str, dict[str, Any]] = {
"candidate": correctness
}
if correctness["status"] != "PASS":
correctness_isolation["off"] = run_correctness(
task, "off", output_root=output_root, python=python,
flaggems=flaggems, site=site, base_env=base_env,
plugin_dir=zl_bench,
)
correctness_isolation["mask0"] = run_correctness(
task, "mask0", output_root=output_root, python=python,
flaggems=flaggems, site=site, base_env=base_env,
plugin_dir=zl_bench,
)
if (
cells["a_off"]["status"] == "PASS"
and cells["a_candidate"]["status"] == "PASS"
and cells["a_candidate"]["convert_count"] < cells["a_off"]["convert_count"]
):
cells["b_candidate"] = run_cell(
task, "candidate", "round-b", off_root,
output_root=output_root, worker=worker, python=python,
zl_bench=zl_bench, flaggems=flaggems, site=site,
shape_file=shape_file, batch_summary=batch_summary, base_env=base_env,
)
cells["b_off"] = run_cell(
task, "off", "round-b", off_root,
output_root=output_root, worker=worker, python=python,
zl_bench=zl_bench, flaggems=flaggems, site=site,
shape_file=shape_file, batch_summary=batch_summary, base_env=base_env,
)
classification, detail = classify_task(
correctness, cells, correctness_isolation
)
receipt = {
"task_id": task["task_id"],
"op": task["op"],
"op_file": task["op_file"],
"family": task["family"],
"nv_evidence": task["nv_evidence"],
"nv_gain_percent": task.get("nv_gain_percent"),
"prior_s5000": task["prior_s5000"],
"classification": classification,
"classification_detail": detail,
"correctness": correctness,
"correctness_isolation": correctness_isolation,
"cells": cells,
}
task_root.mkdir(parents=True, exist_ok=True)
_atomic_json(task_receipt_path, receipt)
rows.append(receipt)
with status_csv.with_name(status_csv.name + ".tmp").open("w", newline="", encoding="utf-8") as handle:
writer = csv.DictWriter(
handle,
fieldnames=["task_id", "op", "family", "classification", "reason"],
)
writer.writeheader()
for row in rows:
writer.writerow(
{
"task_id": row["task_id"],
"op": row["op"],
"family": row["family"],
"classification": row["classification"],
"reason": row["classification_detail"]["reason"],
}
)
os.replace(status_csv.with_name(status_csv.name + ".tmp"), status_csv)
print(
f"MUSA_RLC_PRODUCT_TASK task={task['task_id']} class={classification} "
f"reason={detail['reason']}",
flush=True,
)
counts = {key: sum(row["classification"] == key for row in rows) for key in sorted(TERMINAL)}
key_rows = {row["task_id"]: row for row in rows}
canary_gate = None
if matrix.get("selection") == "canary":
canary_gate = (
key_rows.get("dropout", {}).get("classification") in {"G", "C"}
and key_rows.get("softmax", {}).get("classification") == "B"
and counts["K"] == 0
)
final = {
"schema_version": 1,
"campaign_id": manifest["campaign_id"],
"selection": matrix["selection"],
"status": "TERMINAL",
"task_count": len(rows),
"counts": counts,
"canary_gate_pass": canary_gate,
"source_observed": observed,
"production_default_policy": "unchanged-off",
"claim_scope": manifest["permitted_claim"],
}
_atomic_json(output_root / "campaign-receipt.json", final)
return final
def main() -> int:
parser = argparse.ArgumentParser()
parser.add_argument("--manifest", type=Path, required=True)
parser.add_argument("--task-matrix", type=Path, required=True)
parser.add_argument("--build-root", type=Path, required=True)
parser.add_argument("--flaggems", type=Path, required=True)
parser.add_argument("--zl-bench", type=Path, required=True)
parser.add_argument("--worker", type=Path, required=True)
parser.add_argument("--python", type=Path, default=Path("/usr/bin/python3"))
parser.add_argument("--output-root", type=Path, required=True)
parser.add_argument("--resume", action="store_true")
args = parser.parse_args()
try:
manifest, matrix, observed = validate_contract(
args.manifest.resolve(), args.task_matrix.resolve(), args.build_root.resolve(),
args.flaggems.resolve(), args.zl_bench.resolve(), args.python.resolve(),
)
final = run_campaign(
manifest, matrix, observed,
output_root=args.output_root.resolve(), worker=args.worker.resolve(),
python=args.python.resolve(), zl_bench=args.zl_bench.resolve(),
flaggems=args.flaggems.resolve(), build_root=args.build_root.resolve(),
resume=args.resume,
)
except (OSError, json.JSONDecodeError, CampaignError, subprocess.SubprocessError) as exc:
print(f"MUSA_RLC_PRODUCT_ERROR: {exc}", file=sys.stderr)
return 2
print(json.dumps(final, indent=2, sort_keys=True))
return 0 if final.get("canary_gate_pass") is not False else 1
if __name__ == "__main__":
raise SystemExit(main())
@@ -1,110 +0,0 @@
#!/usr/bin/env bash
set -Eeuo pipefail
: "${PYTHON_BIN:?PYTHON_BIN is required}"
: "${ZL_BENCH_DIR:?ZL_BENCH_DIR is required}"
: "${FLAGGEMS_DIR:?FLAGGEMS_DIR is required}"
: "${FLAGTREE_SITE:?FLAGTREE_SITE is required}"
: "${OP:?OP is required}"
: "${OP_FILE:?OP_FILE is required}"
: "${SHAPE_FILE:?SHAPE_FILE is required}"
: "${OUT_DIR:?OUT_DIR is required}"
worker_dir="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
# A silent preflight exit leaves a zero-byte controller log and no run.log, so
# the campaign records a failed cell with nothing to diagnose it. Name the check
# that failed instead.
preflight() {
local description="$1"
shift
if ! "$@"; then
printf 'MUSA_RLC_PRODUCT_PREFLIGHT_FAILED op=%s check=%s\n' \
"${OP:-?}" "$description" >&2
exit 3
fi
}
preflight "python interpreter is executable: $PYTHON_BIN" test -x "$PYTHON_BIN"
preflight "zl_bench directory exists: $ZL_BENCH_DIR" test -d "$ZL_BENCH_DIR"
preflight "ir meta plugin exists: $ZL_BENCH_DIR/_ir_meta_plugin.py" \
test -f "$ZL_BENCH_DIR/_ir_meta_plugin.py"
preflight "benchmark file exists: $FLAGGEMS_DIR/benchmark/test_${OP_FILE}.py" \
test -f "$FLAGGEMS_DIR/benchmark/test_${OP_FILE}.py"
preflight "FlagGems source exists: $FLAGGEMS_DIR/src/flag_gems" \
test -d "$FLAGGEMS_DIR/src/flag_gems"
preflight "libtriton is non-empty: $FLAGTREE_SITE/triton/_C/libtriton.so" \
test -s "$FLAGTREE_SITE/triton/_C/libtriton.so"
preflight "shape file exists: $SHAPE_FILE" test -f "$SHAPE_FILE"
preflight "output directory is fresh: $OUT_DIR" test ! -e "$OUT_DIR"
mkdir -p "$OUT_DIR/autotune_records" "$OUT_DIR/triton-cache" "$OUT_DIR/flaggems-cache"
export PYTHONNOUSERSITE=1
export PYTHONDONTWRITEBYTECODE=1
export PYTHONUNBUFFERED=1
export GEMS_VENDOR=mthreads
export FLAGGEMS_VENDOR=mthreads
export FLAGTREE_BACKEND=mthreads
export TRITON_BACKENDS_IN_TREE=1
export MUSA_VISIBLE_DEVICES=7
export MTHREADS_VISIBLE_DEVICES=7
export PYTHONPATH="$worker_dir:$ZL_BENCH_DIR:$FLAGGEMS_DIR/src:$FLAGTREE_SITE"
export FLAGGEMS_PERF_CURRENT_OP="$OP"
export FLAGGEMS_PERF_COLOR=never
export NO_COLOR=1
# The plugin list comes from zl_bench, not from here. Hand-copying it is how
# this campaign drifted: the copy was taken before backends/ existed, so it
# still named _device_guard_plugin (now _device_setup_plugin) and, more
# importantly, carried no timing plugin at all -- which is why every sealed
# measurement is timer=direct while zl_bench's own driver had been capturing.
# Asking the planner means the caliber tracks whatever the backend declares.
export ZL_BACKEND=musa
mapfile -t plugins < <("$PYTHON_BIN" "$ZL_BENCH_DIR/plugin_plan.py" \
--backend "$ZL_BACKEND" --pytest-args | tr ' ' '\n')
preflight "plugin plan is non-empty" test "${#plugins[@]}" -gt 0
dtype_args=()
for dtype in ${DTYPES:-float16}; do
dtype_args+=(--dtypes "$dtype")
done
# Some benchmarks parametrise one dtype per case, so a campaign-wide --dtypes
# filter makes every non-matching case raise. A workload may then pin the case
# it measures; the manifest requires a stated reason for that narrowing.
selector_args=()
if [[ -n "${BENCH_SELECTOR:-}" ]]; then
selector_args+=(-k "$BENCH_SELECTOR")
fi
if [[ -n "${REPLAY_FROM:-}" ]]; then
autotune_env="FLAGGEMS_PERF_AUTOTUNE_REPLAY_DIR=$REPLAY_FROM/autotune_records"
mode_desc="replay($REPLAY_FROM)"
else
autotune_env="FLAGGEMS_PERF_AUTOTUNE_RECORD_DIR=$OUT_DIR/autotune_records"
mode_desc=record
fi
cp "$SHAPE_FILE" "$OUT_DIR/shapes.yaml"
test_file="$FLAGGEMS_DIR/benchmark/test_${OP_FILE}.py::test_${OP}"
status=0
{
printf 'MUSA_RLC_PRODUCT_RUN op=%s mode=%s out=%s\n' "$OP" "$mode_desc" "$OUT_DIR"
env \
TRITON_CACHE_DIR="$OUT_DIR/triton-cache" \
FLAGGEMS_CACHE_DIR="$OUT_DIR/flaggems-cache" \
FLAGGEMS_PERF_TTGIR_DUMP_DIR="$OUT_DIR/ttgir" \
"$autotune_env" \
USE_FLAGTUNE=0 \
"$PYTHON_BIN" -u -m pytest -s "$test_file" \
"${plugins[@]}" \
--color=no \
--shape_file "$SHAPE_FILE" \
"${dtype_args[@]}" \
"${selector_args[@]}" \
--level core --mode kernel || status=$?
printf 'MUSA_RLC_PRODUCT_RUN_RC=%s\n' "$status"
exit "$status"
} 2>&1 | tee "$OUT_DIR/run.log" || status=$?
exit "$status"
@@ -1,573 +0,0 @@
from __future__ import annotations
import copy
import importlib.util
import json
import sys
import tempfile
import unittest
from types import SimpleNamespace
from pathlib import Path
HERE = Path(__file__).resolve().parent
CAMPAIGN = HERE.parent
REPO_ROOT = CAMPAIGN.parents[1]
REFERENCE = CAMPAIGN / "manifest" / "musa-rlc-32-product-v1.json"
SPEC = importlib.util.spec_from_file_location(
"prepare_musa_rlc_product", CAMPAIGN / "prepare_musa_rlc_product.py"
)
assert SPEC is not None and SPEC.loader is not None
MODULE = importlib.util.module_from_spec(SPEC)
SPEC.loader.exec_module(MODULE)
RUN_SPEC = importlib.util.spec_from_file_location(
"run_musa_rlc_product", CAMPAIGN / "run_musa_rlc_product.py"
)
assert RUN_SPEC is not None and RUN_SPEC.loader is not None
RUNNER = importlib.util.module_from_spec(RUN_SPEC)
RUN_SPEC.loader.exec_module(RUNNER)
COMPAT_SPEC = importlib.util.spec_from_file_location(
"_musa_benchmark_compat_plugin", REPO_ROOT / "_musa_benchmark_compat_plugin.py"
)
assert COMPAT_SPEC is not None and COMPAT_SPEC.loader is not None
COMPAT = importlib.util.module_from_spec(COMPAT_SPEC)
COMPAT_SPEC.loader.exec_module(COMPAT)
def _manifest():
return json.loads(REFERENCE.read_text(encoding="utf-8"))
class MusaRlcProductTests(unittest.TestCase):
def test_musa_fft_compatibility_override_is_narrow(self):
stale = SimpleNamespace(
name="skipif", args=(True,), kwargs={"reason": "CUDA is not available"}
)
unrelated = SimpleNamespace(
name="skipif", args=(True,), kwargs={"reason": "another reason"}
)
fft = SimpleNamespace(path=Path("benchmark/test_fft.py"), name="test_perf_fft")
other = SimpleNamespace(path=Path("benchmark/test_fft.py"), name="test_other")
self.assertTrue(COMPAT._is_stale_fft_cuda_guard(fft, stale))
self.assertFalse(COMPAT._is_stale_fft_cuda_guard(fft, unrelated))
self.assertFalse(COMPAT._is_stale_fft_cuda_guard(other, stale))
def test_runner_preserves_preflight_failure_receipts(self):
runner = (CAMPAIGN / "run_musa_rlc_product.py").read_text(encoding="utf-8")
ensure_parent = "out.mkdir(parents=True, exist_ok=True)"
write_receipt = '_atomic_json(out / "cell-receipt.json", receipt)'
self.assertIn(ensure_parent, runner)
self.assertLess(runner.index(ensure_parent), runner.index(write_receipt))
def test_worker_dumps_only_executed_ttgir_variants(self):
worker = (CAMPAIGN / "run_pytest_musa_rlc_product.sh").read_text(
encoding="utf-8"
)
self.assertIn(
'FLAGGEMS_PERF_TTGIR_DUMP_DIR="$OUT_DIR/ttgir"', worker
)
# The plugin list is asked for, not written here: a hand-copied list is
# what let this worker keep _device_guard_plugin and no timing plugin
# long after zl_bench had moved on. What the plan contains for musa --
# the compat plugin included -- is asserted in zl_bench's test_backends.
self.assertIn("plugin_plan.py", worker)
self.assertIn('--backend "$ZL_BACKEND"', worker)
self.assertNotIn("-p _", worker)
self.assertIn(
'PYTHONPATH="$worker_dir:$ZL_BENCH_DIR:$FLAGGEMS_DIR/src:$FLAGTREE_SITE"',
worker,
)
with tempfile.TemporaryDirectory() as temporary:
root = Path(temporary)
dumped = root / "ttgir" / "shape" / "kernel" / "used.ttgir"
cached = root / "triton-cache" / "hash" / "loser.ttgir"
dumped.parent.mkdir(parents=True)
cached.parent.mkdir(parents=True)
dumped.write_text("ttg.convert_layout\nttg.convert_layout\n", encoding="utf-8")
cached.write_text("ttg.convert_layout\n" * 99, encoding="utf-8")
scope = RUNNER._convert_count(root)
self.assertEqual(scope["total_convert_count"], 2)
self.assertEqual(scope["ttgir_files"], 1)
# Without an index the measured scope cannot be established, so the
# count stays unscoped and the cell is marked incomplete.
self.assertEqual(scope["measured_convert_count"], 2)
self.assertFalse(scope["scope_complete"])
@staticmethod
def _write_cell(root, rows):
"""Materialise a cell's ttgir dump plus the index the dumper writes."""
ttgir = root / "ttgir"
ttgir.mkdir(parents=True)
lines = [
"\t".join(
[
"shape", "kernel", "file", "launches", "sweep_launches",
"num_warps", "num_stages", "constexprs", "cache_hash",
]
)
]
for shape, kernel, name, launches, converts in rows:
if name.startswith("("):
lines.append(
"\t".join([shape, kernel, name, str(launches), "0", "4", "3", "{}", "H"])
)
continue
path = ttgir / shape / kernel / name
path.parent.mkdir(parents=True, exist_ok=True)
path.write_text("ttg.convert_layout\n" * converts, encoding="utf-8")
lines.append(
"\t".join([shape, kernel, name, str(launches), "0", "4", "3", "{}", "H"])
)
(ttgir / "index.tsv").write_text("\n".join(lines) + "\n", encoding="utf-8")
def test_convert_count_excludes_input_construction_kernels(self):
# Reproduces the sealed gelu_backward shape: every convert sits in the
# once-launched input builder, while the timed kernel has none.
with tempfile.TemporaryDirectory() as temporary:
root = Path(temporary)
self._write_cell(
root,
[
("1048576", "randn_kernel", "w16s3.ttgir", 1, 4),
("1048576", "gelu_none_kernel_rank_1", "w4s3.ttgir", 1, 0),
("1048576", "gelu_backward_none_kernel_rank_1", "w4s3.ttgir", 3795, 0),
],
)
scope = RUNNER._convert_count(root)
self.assertEqual(scope["total_convert_count"], 4)
self.assertEqual(scope["measured_convert_count"], 0)
self.assertTrue(scope["scope_complete"])
self.assertEqual(
[row["file"] for row in scope["excluded_unmeasured"]],
["1048576/randn_kernel/w16s3.ttgir"],
)
def test_convert_count_keeps_repeatedly_launched_kernels(self):
# Reproduces the sealed sort shape: converts in the dominant kernel and
# in an autotune loser that was indexed but never dumped.
with tempfile.TemporaryDirectory() as temporary:
root = Path(temporary)
self._write_cell(
root,
[
("1024x1024", "sweep", "w4s3.ttgir", 280, 3),
("1024x1024", "broadcast_to_kernel", "w4s3.ttgir", 140, 1),
("-", "arange_kernel", "(sweep loser, not dumped)", 0, 0),
],
)
scope = RUNNER._convert_count(root)
self.assertEqual(scope["measured_convert_count"], 4)
self.assertEqual(scope["total_convert_count"], 4)
self.assertTrue(scope["scope_complete"])
self.assertEqual(scope["excluded_unmeasured"], [])
def test_convert_count_reports_dumped_but_unindexed_files(self):
with tempfile.TemporaryDirectory() as temporary:
root = Path(temporary)
self._write_cell(root, [("s", "measured", "w4s3.ttgir", 140, 2)])
stray = root / "ttgir" / "s" / "stray" / "w4s3.ttgir"
stray.parent.mkdir(parents=True)
stray.write_text("ttg.convert_layout\n", encoding="utf-8")
scope = RUNNER._convert_count(root)
# An unindexed dump is counted rather than dropped, and flagged.
self.assertEqual(scope["measured_convert_count"], 3)
self.assertFalse(scope["scope_complete"])
self.assertEqual(scope["unindexed_files"], ["s/stray/w4s3.ttgir"])
def test_incomplete_convert_scope_fails_closed(self):
def cell(convert_count, latency, complete=True):
return {
"status": "PASS",
"convert_count": convert_count,
"convert_scope_complete": complete,
"median_gems_latency_ms": latency,
"replay_fallback_count": 0,
}
cells = {
"a_off": cell(8, 0.14),
"a_mask0": cell(8, 0.14),
"a_candidate": cell(0, 0.12, complete=False),
"b_candidate": cell(0, 0.121),
"b_off": cell(8, 0.141),
}
classification, detail = RUNNER.classify_task({"status": "PASS"}, cells)
self.assertEqual(classification, "X")
self.assertEqual(
detail["reason"], "measured convert topology scope is incomplete"
)
def test_correctness_command_narrows_to_the_measured_variant(self):
python = Path("/usr/bin/python3")
plain = RUNNER.correctness_command(
{"correctness_file": "tests/test_sort.py"}, python
)
self.assertEqual(plain[-3:], ["tests/test_sort.py", "-rs", "--quick"])
self.assertNotIn("-k", plain)
selected = RUNNER.correctness_command(
{
"correctness_file": "tests/test_upsample_bicubic2d_aa.py",
"correctness_selector": "not backward",
},
python,
)
self.assertIn("tests/test_upsample_bicubic2d_aa.py", selected)
self.assertEqual(selected[selected.index("-k") + 1], "not backward")
pinned = RUNNER.correctness_command(
{
"correctness_file": "tests/test_mm.py",
"correctness_node_ids": ["tests/test_mm.py::test_mm"],
},
python,
)
self.assertIn("tests/test_mm.py::test_mm", pinned)
self.assertNotIn("tests/test_mm.py", pinned)
def test_narrowed_correctness_requires_a_stated_reason(self):
manifest = copy.deepcopy(_manifest())
target = next(
row for row in manifest["workloads"] if row["id"] == "upsample_bicubic2d_aa"
)
self.assertEqual(target["correctness_selector"], "not backward")
target.pop("correctness_scope_reason")
errors = MODULE.validate_manifest(manifest)
self.assertIn(
"workload upsample_bicubic2d_aa narrows correctness and must state "
"correctness_scope_reason",
errors,
)
both = copy.deepcopy(_manifest())
row = next(item for item in both["workloads"] if item["id"] == "sort")
row["correctness_selector"] = "sort"
row["correctness_node_ids"] = ["tests/test_sort.py::test_sort"]
row["correctness_scope_reason"] = "reason"
self.assertIn(
"workload sort must not set both correctness_selector and "
"correctness_node_ids",
MODULE.validate_manifest(both),
)
def test_measured_launch_floor_is_a_frozen_protocol_field(self):
manifest = copy.deepcopy(_manifest())
self.assertEqual(manifest["protocol"]["measured_launch_floor"], 1)
plan, _ = MODULE.build_plan(manifest, "all")
self.assertTrue(
all(task["measured_launch_floor"] == 1 for task in plan["tasks"])
)
scoped = next(
task for task in plan["tasks"] if task["task_id"] == "fused_marlin_moe_mxfp4"
)
self.assertEqual(scoped["correctness_selector"], "mxfp4")
self.assertTrue(scoped["correctness_scope_reason"])
manifest["protocol"]["measured_launch_floor"] = -1
self.assertIn(
"protocol.measured_launch_floor must be a non-negative integer",
MODULE.validate_manifest(manifest),
)
def test_compat_plugin_only_overrides_verified_stale_guards(self):
def mark(reason, name="skipif"):
return SimpleNamespace(name=name, args=(True,), kwargs={"reason": reason})
def item(file_name, test_name):
return SimpleNamespace(path=Path("tests") / file_name, name=test_name)
# Verified on S5000: unskipped, these pass.
self.assertTrue(
COMPAT._is_stale_cuda_guard(
item("test_top_k_per_row_decode.py", "test_accuracy_top_k"),
mark("CUDA device required"),
)
)
self.assertTrue(
COMPAT._is_stale_fft_cuda_guard(
item("test_fft.py", "test_perf_fft"), mark("CUDA is not available")
)
)
# The fft guard stays pinned to the one test it was verified against.
self.assertFalse(
COMPAT._is_stale_cuda_guard(
item("test_fft.py", "test_other"), mark("CUDA is not available")
)
)
# Probed and rejected: grid_sample hardcodes device="cuda" in its own
# helper, and cutlass_scaled_mm gates on vLLM as well as CUDA.
self.assertFalse(
COMPAT._is_stale_cuda_guard(
item("test_grid_sample.py", "test_nearest_zeros_4d_small"),
mark("CUDA is required."),
)
)
self.assertFalse(
COMPAT._is_stale_cuda_guard(
item("test_cutlass_scaled_mm.py", "test_cutlass_scaled_mm"),
mark("Feature requires vLLM and NVIDIA Hopper architecture"),
)
)
# Only skipif marks, and only the exact reason string.
self.assertFalse(
COMPAT._is_stale_cuda_guard(
item("test_top_k_per_row_decode.py", "test_accuracy_top_k"),
mark("CUDA device required", name="xfail"),
)
)
self.assertFalse(
COMPAT._is_stale_cuda_guard(
item("test_top_k_per_row_decode.py", "test_accuracy_top_k"),
mark("some other reason"),
)
)
def test_compat_plugin_strips_module_level_guards(self):
# top_k_per_row_decode carries its guard as a module-level pytestmark,
# which resolves through the node chain rather than the item's own marks.
stale = SimpleNamespace(
name="skipif", args=(True,), kwargs={"reason": "CUDA device required"}
)
keep = SimpleNamespace(
name="skipif", args=(True,), kwargs={"reason": "vLLM not installed"}
)
module = SimpleNamespace(own_markers=[stale, keep])
function = SimpleNamespace(
path=Path("tests/test_top_k_per_row_decode.py"),
name="test_accuracy_top_k",
own_markers=[],
)
function.listchain = lambda: [module, function]
sys.modules.setdefault(
"flag_gems", SimpleNamespace(vendor_name="mthreads")
)
COMPAT.pytest_collection_modifyitems([function])
self.assertEqual(module.own_markers, [keep])
def test_benchmark_selector_requires_a_stated_reason(self):
manifest = copy.deepcopy(_manifest())
target = next(
row for row in manifest["workloads"] if row["id"] == "grid_sample"
)
self.assertEqual(target["benchmark_selector"], "dtype0")
target.pop("benchmark_scope_reason")
self.assertIn(
"workload grid_sample narrows the benchmark and must state "
"benchmark_scope_reason",
MODULE.validate_manifest(manifest),
)
def test_declared_benchmark_dtypes_must_exist_in_the_cli_vocabulary(self):
manifest = copy.deepcopy(_manifest())
self.assertEqual(
next(row for row in manifest["workloads"] if row["id"] == "randperm")[
"benchmark_dtypes"
],
["int32"],
)
# int64 is not a --dtypes choice; the sealed run only found that out on
# the device, as an argparse error inside a screen cell.
next(row for row in manifest["workloads"] if row["id"] == "randperm")[
"benchmark_dtypes"
] = ["int32", "int64"]
self.assertIn(
"workload randperm declares benchmark dtype(s) the benchmark CLI "
"does not accept: int64",
MODULE.validate_manifest(manifest),
)
def test_stale_inventory_needs_an_override_reason(self):
manifest = copy.deepcopy(_manifest())
marlin = next(
row
for row in manifest["workloads"]
if row["id"] == "fused_marlin_moe_mxfp4"
)
self.assertEqual(marlin["op_file"], "fused_marlin_moe_w4a16_mxfp4")
self.assertTrue(marlin["inventory_override_reason"])
with tempfile.TemporaryDirectory() as temporary:
inventory = Path(temporary)
(inventory / "inventory_main.csv").write_text(
"repo,op,op_file,file,marked_skip,collected,op_names\n"
+ "".join(
f"main,{row['op']},{row['op_file']},f.py,,yes,{row['id']}\n"
for row in manifest["workloads"]
if row["id"] != "fused_marlin_moe_mxfp4"
),
encoding="utf-8",
)
# The stated override lets the missing row through.
self.assertEqual(
MODULE.validate_inventory(manifest["workloads"], inventory), []
)
marlin.pop("inventory_override_reason")
self.assertEqual(
MODULE.validate_inventory(manifest["workloads"], inventory),
[
"inventory missing ('main', 'fused_marlin_moe_w4a16_mxfp4', "
"'fused_marlin_moe_w4a16_mxfp4')"
],
)
def test_missing_benchmark_target_fails_before_the_campaign_starts(self):
matrix = {
"tasks": [
{
"task_id": "present",
"op_file": "sort",
"correctness_file": "tests/test_sort.py",
},
{
# The shape of the sealed marlin defect: the inventory named
# a merged benchmark file that upstream had already split.
"task_id": "fused_marlin_moe_mxfp4",
"op_file": "fused_marlin_moe",
"correctness_file": "tests/test_fused_marlin_moe.py",
},
]
}
with tempfile.TemporaryDirectory() as temporary:
flaggems = Path(temporary)
(flaggems / "benchmark").mkdir()
(flaggems / "tests").mkdir()
(flaggems / "benchmark" / "test_sort.py").write_text("", encoding="utf-8")
(flaggems / "tests" / "test_sort.py").write_text("", encoding="utf-8")
(flaggems / "tests" / "test_fused_marlin_moe.py").write_text(
"", encoding="utf-8"
)
missing = RUNNER.missing_targets(matrix, flaggems)
self.assertEqual(len(missing), 1)
self.assertIn("fused_marlin_moe_mxfp4 benchmark", missing[0])
self.assertIn("test_fused_marlin_moe.py", missing[0])
def test_worker_names_the_failed_preflight_check(self):
worker = (CAMPAIGN / "run_pytest_musa_rlc_product.sh").read_text(encoding="utf-8")
self.assertIn("MUSA_RLC_PRODUCT_PREFLIGHT_FAILED", worker)
self.assertIn(
'preflight "benchmark file exists: '
'$FLAGGEMS_DIR/benchmark/test_${OP_FILE}.py"',
worker,
)
# The benchmark selector must reach pytest.
self.assertIn('selector_args+=(-k "$BENCH_SELECTOR")', worker)
self.assertIn('"${selector_args[@]}"', worker)
def test_correctness_runs_with_the_vendor_compat_plugin(self):
command = RUNNER.correctness_command(
{"correctness_file": "tests/test_top_k_per_row_decode.py"},
Path("/usr/bin/python3"),
)
self.assertEqual(
command[command.index("-p") + 1], "_musa_benchmark_compat_plugin"
)
def test_valid_manifest_and_plan_sizes(self):
manifest = _manifest()
self.assertEqual(MODULE.validate_manifest(manifest), [])
canary, canary_budget = MODULE.build_plan(manifest, "canary")
full, full_budget = MODULE.build_plan(manifest, "all")
self.assertEqual(len(canary["tasks"]), 8)
self.assertEqual(len(full["tasks"]), 32)
self.assertEqual(canary["tasks"][0]["task_id"], "dropout")
self.assertEqual(canary["tasks"][1]["task_id"], "softmax")
self.assertEqual(canary["tasks"][0]["benchmark_dtypes"], ["float16"])
self.assertEqual(
next(task for task in canary["tasks"] if task["task_id"] == "fft")[
"benchmark_dtypes"
],
["complex64"],
)
self.assertEqual(canary_budget["gpu_count"], 1)
self.assertEqual(full_budget["gpu_count"], 1)
self.assertTrue(canary_budget["serial_measurement"])
def test_invalid_device_and_canary_count_fail_closed(self):
manifest = copy.deepcopy(_manifest())
manifest["device"]["physical_device"] = "0"
manifest["workloads"][0]["canary"] = False
errors = MODULE.validate_manifest(manifest)
self.assertIn("device.physical_device must equal '7'", errors)
self.assertIn("exactly eight workloads must be canaries", errors)
def test_invalid_timer_and_default_policy_fail_closed(self):
manifest = copy.deepcopy(_manifest())
manifest["protocol"]["timer"] = "host wall clock"
manifest["production_default_policy"] = "enabled"
errors = MODULE.validate_manifest(manifest)
self.assertIn("protocol.timer must bind the MUSA native-event path", errors)
self.assertIn("production_default_policy must be unchanged-off", errors)
def test_runtime_classifier_positive_and_fail_closed_paths(self):
def cell(convert_count, latency, status="PASS", replay=0):
return {
"status": status,
"convert_count": convert_count,
"median_gems_latency_ms": latency,
"replay_fallback_count": replay,
}
correct = {"status": "PASS"}
cells = {
"a_off": cell(8, 0.14),
"a_mask0": cell(8, 0.14),
"a_candidate": cell(0, 0.12),
"b_candidate": cell(0, 0.121),
"b_off": cell(8, 0.141),
}
classification, _ = RUNNER.classify_task(correct, cells)
self.assertEqual(classification, "G")
invalid = copy.deepcopy(cells)
invalid["a_mask0"]["convert_count"] = 7
classification, detail = RUNNER.classify_task(correct, invalid)
self.assertEqual(classification, "X")
self.assertEqual(detail["reason"], "mask0 topology differs from off")
classification, _ = RUNNER.classify_task({"status": "FAIL"}, cells)
self.assertEqual(classification, "K")
classification, detail = RUNNER.classify_task(
{"status": "FAIL"},
cells,
{
"candidate": {"status": "FAIL"},
"mask0": {"status": "FAIL"},
"off": {"status": "FAIL"},
},
)
self.assertEqual(classification, "X")
self.assertEqual(
detail["reason"],
"baseline correctness fails under off/mask0/candidate",
)
def test_full_preparation_requires_passing_canary_receipt(self):
manifest = _manifest()
self.assertEqual(
MODULE.validate_canary_receipt(manifest, "all", None),
["full campaign requires --canary-receipt"],
)
with tempfile.TemporaryDirectory() as temporary:
receipt = Path(temporary) / "campaign-receipt.json"
receipt.write_text(
json.dumps(
{
"campaign_id": manifest["campaign_id"],
"selection": "canary",
"status": "TERMINAL",
"canary_gate_pass": True,
}
),
encoding="utf-8",
)
self.assertEqual(
MODULE.validate_canary_receipt(manifest, "all", receipt), []
)
if __name__ == "__main__":
unittest.main()
@@ -1,157 +0,0 @@
from __future__ import annotations
import importlib.util
import json
import tempfile
import unittest
from pathlib import Path
HERE = Path(__file__).resolve().parent
CAMPAIGN = HERE.parent
SPEC = importlib.util.spec_from_file_location(
"reaccount_musa_rlc_topology", CAMPAIGN / "reaccount_musa_rlc_topology.py"
)
assert SPEC is not None and SPEC.loader is not None
TOOL = importlib.util.module_from_spec(SPEC)
SPEC.loader.exec_module(TOOL)
def _write_cell(cell_dir: Path, rows):
ttgir = cell_dir / "ttgir"
ttgir.mkdir(parents=True)
lines = [
"\t".join(
[
"shape", "kernel", "file", "launches", "sweep_launches",
"num_warps", "num_stages", "constexprs", "cache_hash",
]
)
]
for shape, kernel, name, launches, converts in rows:
path = ttgir / shape / kernel / name
path.parent.mkdir(parents=True, exist_ok=True)
path.write_text("ttg.convert_layout\n" * converts, encoding="utf-8")
lines.append(
"\t".join([shape, kernel, name, str(launches), "0", "4", "3", "{}", "H"])
)
(ttgir / "index.tsv").write_text("\n".join(lines) + "\n", encoding="utf-8")
def _seal_run(root: Path):
"""A minimal sealed run shaped like the gelu_backward task."""
(root / "campaign-receipt.json").write_text(
json.dumps({"campaign_id": "sealed-v2"}), encoding="utf-8"
)
task_root = root / "tasks" / "gelu_backward"
cells = {}
for key, round_name, cell_name, setup, measured in (
("a_off", "round-a", "off", 4, 0),
("a_mask0", "round-a", "mask0", 4, 0),
("a_candidate", "round-a", "candidate", 0, 0),
("b_off", "round-b", "off", 4, 0),
("b_candidate", "round-b", "candidate", 0, 0),
):
cell_dir = task_root / round_name / cell_name
_write_cell(
cell_dir,
[
("s", "randn_kernel", "w16s3.ttgir", 1, setup),
("s", "gelu_backward_kernel", "w4s3.ttgir", 3795, measured),
],
)
cells[key] = {
"round": round_name,
"cell": cell_name,
"status": "PASS",
"convert_count": setup + measured,
"median_gems_latency_ms": 0.1168 if cell_name == "off" else 0.1166,
"replay_fallback_count": 0,
}
task_root.mkdir(parents=True, exist_ok=True)
(task_root / "task-receipt.json").write_text(
json.dumps(
{
"task_id": "gelu_backward",
"op": "gelu_backward",
"family": "elementwise-backward",
"classification": "C",
"classification_detail": {
"reason": "topology reduction without stable >=5% speedup"
},
"correctness": {"status": "PASS"},
"correctness_isolation": {"candidate": {"status": "PASS"}},
"cells": cells,
}
),
encoding="utf-8",
)
class ReaccountTests(unittest.TestCase):
def test_input_construction_reduction_is_demoted_to_b(self):
with tempfile.TemporaryDirectory() as temporary:
base = Path(temporary)
run_root = base / "sealed"
run_root.mkdir()
_seal_run(run_root)
result = TOOL.reaccount(run_root, base / "derived", 1)
self.assertEqual(result["changed_task_ids"], ["gelu_backward"])
task = result["tasks"][0]
self.assertEqual(task["sealed_classification"], "C")
self.assertEqual(task["reaccounted_classification"], "B")
self.assertEqual(task["sealed_convert"]["a_off"], 4)
self.assertEqual(task["reaccounted_convert"]["a_off"], 0)
self.assertEqual(
task["excluded_unmeasured_files"], ["s/randn_kernel/w16s3.ttgir"]
)
def test_a_zero_floor_reproduces_the_sealed_verdict(self):
with tempfile.TemporaryDirectory() as temporary:
base = Path(temporary)
run_root = base / "sealed"
run_root.mkdir()
_seal_run(run_root)
result = TOOL.reaccount(run_root, base / "derived", 0)
# Floor 0 admits the once-launched builder again, so the derived verdict
# must match the sealed one. That equivalence is what makes the scoped
# count auditable rather than an unexplained new number.
self.assertEqual(result["changed_task_ids"], [])
self.assertEqual(result["tasks"][0]["reaccounted_classification"], "C")
def test_derived_output_is_written_and_never_inside_the_sealed_run(self):
with tempfile.TemporaryDirectory() as temporary:
base = Path(temporary)
run_root = base / "sealed"
run_root.mkdir()
_seal_run(run_root)
with self.assertRaises(TOOL.ReaccountError):
TOOL.reaccount(run_root, run_root / "derived", 1)
with self.assertRaises(TOOL.ReaccountError):
TOOL.reaccount(run_root, run_root, 1)
output = base / "derived"
TOOL.reaccount(run_root, output, 1)
self.assertTrue((output / "reaccount-receipt.json").is_file())
self.assertTrue((output / "reaccount-matrix.csv").is_file())
# The sealed receipt must survive untouched.
sealed = json.loads(
(run_root / "tasks" / "gelu_backward" / "task-receipt.json").read_text(
encoding="utf-8"
)
)
self.assertEqual(sealed["classification"], "C")
self.assertEqual(sealed["cells"]["a_off"]["convert_count"], 4)
def test_a_directory_without_a_campaign_receipt_is_rejected(self):
with tempfile.TemporaryDirectory() as temporary:
base = Path(temporary)
with self.assertRaises(TOOL.ReaccountError):
TOOL.reaccount(base, base.parent / "derived", 1)
if __name__ == "__main__":
unittest.main()
-2
View File
@@ -41,8 +41,6 @@ import yaml
OPS_DIR = Path(__file__).resolve().parent 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 = [ DEFAULT_REPOS = [
("main", Path("/workspace/dev/FlagGems")), ("main", Path("/workspace/dev/FlagGems")),
("vllm", Path("/workspace/dev/FlagGems-vllm")), ("vllm", Path("/workspace/dev/FlagGems-vllm")),
-145
View File
@@ -1,145 +0,0 @@
#!/usr/bin/env python3
"""把筛选通过的算子按实测耗时分到各张卡上。
为什么不用 `run_batch.py` 的 hash 分片:它按算子名的 hash 取模,不看重量。S5000 上
17 个算子的六轮 campaign,hash 分片关键路径 120 分钟、有一张卡空转;按耗时贪心是
70 分钟。同一批活,差 50 分钟。
贪心(最长优先塞进当前最空的卡)在这里够用,也不可能更好:`topk` 一个算子就要 70.6
分钟,而理想均分是 45.4——**关键路径的下界由最重的那个算子决定,加卡不解决**。分片
只能保证除它之外的卡不空等。
权重来自 `screen.tsv` 的 `elapsed_before + elapsed_after`(**一对 = 2 格**的秒数)。
N 轮是 2N 格,所以预估耗时是权重的 N 倍——不是 N/6 倍,这里错过一次,把 topk 的 70.6
分钟报成了 11.8。缺列时退化成等权重,并在 stderr 说明——等权重的分片仍然
比不分好,但别让人以为它按耗时排过。
用法:
python partition.py --screen runs/ab_rlc_XXX/screen.tsv --devices 0,1,2,3
python partition.py --screen ... --devices 0,1 --verdict IR_CHANGED --format shell
"""
from __future__ import annotations
import argparse
import csv
import sys
from pathlib import Path
from typing import Dict, List, Optional, Sequence, Tuple
Item = Tuple[str, float]
def plan(items: Sequence[Item], slots: Sequence[str]) -> Dict[str, List[Item]]:
"""最长优先贪心:每次把最重的活放进当前总量最小的卡。
卡少于 1 张是调用方的错,直接抛;活比卡少时后面的卡就是空的,这不是错误。
"""
if not slots:
raise ValueError("需要至少一张卡")
out: Dict[str, List[Item]] = {s: [] for s in slots}
load: Dict[str, float] = {s: 0.0 for s in slots}
# 同重量时按名字排,保证同一份输入每次得到同一个方案
for name, weight in sorted(items, key=lambda kv: (-kv[1], kv[0])):
target = min(slots, key=lambda s: (load[s], s))
out[target].append((name, weight))
load[target] += weight
return out
def imbalance(assignment: Dict[str, List[Item]]) -> float:
"""最慢的卡 / 理想均分。1.0 是完美;下界由最重的单个算子决定。"""
loads = [sum(w for _, w in v) for v in assignment.values()]
total = sum(loads)
if total <= 0:
return 1.0
return max(loads) / (total / len(loads))
def read_screen(path: Path, verdict: str = "IR_CHANGED"
) -> Tuple[List[Item], Dict[str, str], List[str]]:
"""(算子, 一对耗时秒), op -> op_file, 被判定挡掉的算子。"""
with path.open(newline="", encoding="utf-8") as f:
rows = list(csv.DictReader(f, delimiter="\t"))
if not rows:
raise SystemExit(f"筛选表没有数据行: {path}")
items: List[Item] = []
files: Dict[str, str] = {}
skipped: List[str] = []
warned = False
for r in rows:
op = r.get("op", "")
if not op:
continue
if verdict and r.get("verdict") != verdict:
skipped.append(f"{op}({r.get('verdict', '?')})")
continue
try:
w = float(r["elapsed_before"]) + float(r["elapsed_after"])
except (KeyError, TypeError, ValueError):
w = 1.0
if not warned:
print("!!! 筛选表没有 elapsed_* 列,退化成等权重分片——"
"分片仍然比不分好,但它没有按耗时排过", file=sys.stderr)
warned = True
items.append((op, w))
files[op] = r.get("op_file") or op
return items, files, skipped
def main(argv: Optional[List[str]] = None) -> int:
p = argparse.ArgumentParser(description=__doc__.splitlines()[0],
formatter_class=argparse.RawDescriptionHelpFormatter)
p.add_argument("--screen", type=Path, required=True, help="ab_rlc_test.sh 的 screen.tsv")
p.add_argument("--devices", required=True, help="逗号分隔的设备列表,如 0,1,2,3")
p.add_argument("--verdict", default="IR_CHANGED",
help="只取这个判定的算子;空字符串表示全要")
# 权重是一对(2 格)的秒数;N 轮是 2N 格,所以预估耗时是权重的 N 倍
p.add_argument("--rounds", type=int, default=6, help="仅用于打印预估耗时")
p.add_argument("--format", choices=("human", "shell"), default="human",
help="shell 输出每卡一行 `<device>\\t<OPS>`,供驱动脚本读")
args = p.parse_args(argv)
if not args.screen.is_file():
raise SystemExit(f"筛选表不存在: {args.screen}")
devices = [d.strip() for d in args.devices.split(",") if d.strip()]
if not devices:
raise SystemExit("--devices 为空")
items, files, skipped = read_screen(args.screen, args.verdict)
if not items:
raise SystemExit(f"筛选表里没有 verdict={args.verdict} 的算子——"
"先确认筛选跑过,且判据没有把所有算子都挡掉")
assignment = plan(items, devices)
if args.format == "shell":
# 第三列是该卡的权重秒数,供驱动脚本决定往哪张卡追加算子——按负载选,不是
# 按算子个数:一个重算子能顶三个轻的。
# 显式写 "\n" 而不用 print:Windows 上 print 会发 CRLF,读的一侧会把 \r
# 带进最后一个字段,数值比较就废了。
for d in devices:
ops = " ".join(f"{op}:{files[op]}" for op, _ in assignment[d])
load = sum(w for _, w in assignment[d])
sys.stdout.buffer.write(f"{d}\t{ops}\t{load:.0f}\n".encode())
sys.stdout.buffer.flush()
return 0
total = sum(w for _, w in items) * args.rounds
print(f">>> {len(items)} 个算子 / {len(devices)} 张卡 / {args.rounds} 轮")
if skipped:
print(f">>> 判定挡掉 {len(skipped)} 个: {', '.join(skipped)}")
for d in devices:
load = sum(w for _, w in assignment[d]) * args.rounds
names = " ".join(op for op, _ in assignment[d]) or "(空)"
print(f" 卡{d} {load/60:6.1f} 分钟 {names}")
crit = max(sum(w for _, w in v) for v in assignment.values()) * args.rounds
heaviest = max(items, key=lambda kv: kv[1])
print(f">>> 单卡串行 {total/3600:.2f} 小时;{len(devices)} 卡关键路径 "
f"{crit/60:.1f} 分钟,失衡 {imbalance(assignment):.2f}x")
print(f">>> 下界由最重的算子决定:{heaviest[0]} 单独就要 "
f"{heaviest[1]*args.rounds/60:.1f} 分钟,加卡不解决")
return 0
if __name__ == "__main__":
raise SystemExit(main())
-109
View File
@@ -1,109 +0,0 @@
"""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)
# Compat first: these override collection-time markers, so they have to be
# in place before anything else inspects the collected items.
plugins: List[str] = list(backend.compat_plugins) + 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())
-1
View File
@@ -1 +0,0 @@
"""Benchmark providers. Phase 1 only implements FlagGems."""
-219
View File
@@ -1,219 +0,0 @@
"""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())
+47 -84
View File
@@ -1,16 +1,13 @@
#!/usr/bin/env python3 #!/usr/bin/env python3
"""多算子批量测试驱动:按算子清单逐个调 run_pytest.sh,多卡分片、串行测量。 """多算子批量测试驱动:按算子清单逐个调 run_pytest.sh,多卡分片、串行测量。
每个算子一个独立子进程(设备 crash 只废单个算子)、独立产物目录 每个算子一个独立子进程(CUDA crash 只废单个算子)、独立产物目录
<batch>/<repo>/<op>/(与单算子跑 run_pytest.sh 的产物结构完全一致,含 <batch>/<repo>/<op>/(与单算子跑 run_pytest.sh 的产物结构完全一致,含
run.log / shapes.yaml / autotune_records / ttgir / run_identity.json)。 run.log / shapes.yaml / autotune_records / ttgir)。批量层只做编排:
批量层只做编排:
- 分片:按算子名稳定哈希到设备(每卡内部串行,测量口径与单算子一致; - 分片:按算子名稳定哈希到 GPU(每卡内部串行,测量口径与单算子一致;
跨卡并行测量存在功耗耦合噪声,本模式定位为筛查口径,可疑算子回 跨卡并行测量存在功耗耦合噪声,本模式定位为筛查口径,可疑算子回
单算子串行流程确认); 单算子串行流程确认);
- 后端:--backend / ZL_BACKEND 选择 cuda/musa/rocm/ascend,设备可见性
由 backends/ 写入,不再写死 CUDA_VISIBLE_DEVICES;
- 超时:--op-timeout 杀死挂死算子(记 TIMEOUT,不重试,继续跑); - 超时:--op-timeout 杀死挂死算子(记 TIMEOUT,不重试,继续跑);
- 重试:pytest rc=1(用例失败/flaky)与信号杀(rc>=128)最多重试 2 次, - 重试:pytest rc=1(用例失败/flaky)与信号杀(rc>=128)最多重试 2 次,
连续两次同 rc 视为确定性失败提前止损;--dtypes 下两级降级——遇上游 连续两次同 rc 视为确定性失败提前止损;--dtypes 下两级降级——遇上游
@@ -26,7 +23,7 @@ run.log / shapes.yaml / autotune_records / ttgir / run_identity.json)。
示例: 示例:
# 全量筛查(两仓库全部可收集算子,排除 marked_skip 与 dismiss) # 全量筛查(两仓库全部可收集算子,排除 marked_skip 与 dismiss)
python run_batch.py --all --devices 0,1,2,3,4,5,6,7 --backend cuda --op-timeout 1800 python run_batch.py --all --gpus 0,1,2,3,4,5,6,7 --op-timeout 1800
# 指定算子子集(repo 前缀默认 main;重名用 @文件名 消歧) # 指定算子子集(repo 前缀默认 main;重名用 @文件名 消歧)
python run_batch.py --ops softmax vllm:fused_marlin_moe nextafter_@nextafter_ python run_batch.py --ops softmax vllm:fused_marlin_moe nextafter_@nextafter_
@@ -53,8 +50,6 @@ from pathlib import Path
from typing import Dict, List, Optional, Set, Tuple from typing import Dict, List, Optional, Set, Tuple
import batch_summary import batch_summary
from backends import UnknownBackendError, get_backend
from _run_identity import load_identity
ZL_BENCH = Path(__file__).resolve().parent ZL_BENCH = Path(__file__).resolve().parent
OPS_DIR = ZL_BENCH / "ops" OPS_DIR = ZL_BENCH / "ops"
@@ -67,9 +62,7 @@ REPO_DIRS = {
} }
STATUS_FIELDS = [ STATUS_FIELDS = [
"repo", "op", "op_file", "run_name", "repo", "op", "op_file", "run_name", "gpu", "status", "return_code",
"backend", "device", "gpu", "timer",
"status", "return_code",
"elapsed_ms", "rows_total", "rows_success", "notes", "out_dir", "elapsed_ms", "rows_total", "rows_success", "notes", "out_dir",
] ]
@@ -201,44 +194,48 @@ def resolve_selection(args: argparse.Namespace,
return cases return cases
# ---------------- 设备分片 ---------------- # ---------------- GPU 分片 ----------------
def resolve_device_arg(devices: Optional[str], gpus: Optional[str]) -> Optional[str]: def visible_gpus(arg: Optional[str]) -> List[str]:
"""Prefer --devices; --gpus is an alias. Passing both is an error.""" if arg:
if devices is not None and gpus is not None: return [g.strip() for g in arg.split(",") if g.strip()]
raise ValueError("do not pass both --devices and --gpus") env = os.environ.get("CUDA_VISIBLE_DEVICES", "").strip()
return devices if devices is not None else gpus if env:
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 visible_devices(backend, explicit: Optional[str] = None) -> List[str]: def partition(cases: List[Case], gpus: List[str]) -> Dict[str, List[Case]]:
return list(backend.enumerate_devices(explicit))
def partition(cases: List[Case], devices: List[str]) -> Dict[str, List[Case]]:
"""按 (repo, run_name) 稳定哈希分片:与选择顺序无关,resume 时同卡复用。""" """按 (repo, run_name) 稳定哈希分片:与选择顺序无关,resume 时同卡复用。"""
parts: Dict[str, List[Case]] = {d: [] for d in devices} parts: Dict[str, List[Case]] = {g: [] for g in gpus}
for c in cases: for c in cases:
digest = hashlib.sha256(f"{c.repo}/{c.run_name}".encode()).digest() digest = hashlib.sha256(f"{c.repo}/{c.run_name}".encode()).digest()
parts[devices[int.from_bytes(digest[:8], "big") % len(devices)]].append(c) parts[gpus[int.from_bytes(digest[:8], "big") % len(gpus)]].append(c)
return parts return parts
# ---------------- 单算子执行 ---------------- # ---------------- 单算子执行 ----------------
def _spawn(case: Case, out_dir: Path, device: str, args: argparse.Namespace, def _spawn(case: Case, out_dir: Path, gpu: str, args: argparse.Namespace,
dtypes: str) -> subprocess.Popen: dtypes: str) -> subprocess.Popen:
env = os.environ.copy() env = os.environ.copy()
env.pop("PARALLEL_WARMUP_GPUS", None) # 批量层已分片,算子内不再并行 env.pop("PARALLEL_WARMUP_GPUS", None) # 批量层已分片,算子内不再并行
env.pop("REPLAY_FROM", None) env.pop("REPLAY_FROM", None)
env.pop("OUT_DIR", None) env.pop("OUT_DIR", None)
env = args.backend_obj.apply_visibility(env, device)
env.update({ env.update({
"OP": case.op, "OP": case.op,
"OP_FILE": case.op_file, "OP_FILE": case.op_file,
"FLAGGEMS_DIR": REPO_DIRS[case.repo], "FLAGGEMS_DIR": REPO_DIRS[case.repo],
"OUT_DIR": str(out_dir), "OUT_DIR": str(out_dir),
"ZL_BACKEND": args.backend, "CUDA_VISIBLE_DEVICES": gpu,
"ZL_DEVICE": device,
"FLAGGEMS_PERF_COLOR": "never", "FLAGGEMS_PERF_COLOR": "never",
"NO_COLOR": "1", "NO_COLOR": "1",
"USE_FLAGTUNE": str(args.use_flagtune), "USE_FLAGTUNE": str(args.use_flagtune),
@@ -274,14 +271,13 @@ def _wait(proc: subprocess.Popen, timeout: int) -> Tuple[int, bool]:
return 124, True return 124, True
def run_case(case: Case, batch_dir: Path, device: str, def run_case(case: Case, batch_dir: Path, gpu: str,
args: argparse.Namespace) -> Dict[str, str]: args: argparse.Namespace) -> Dict[str, str]:
out_dir = batch_dir / case.repo / case.run_name out_dir = batch_dir / case.repo / case.run_name
marker = out_dir / ".complete" marker = out_dir / ".complete"
if args.resume and marker.is_file() and (out_dir / "run.log").is_file(): if args.resume and marker.is_file() and (out_dir / "run.log").is_file():
row = _classify(case, out_dir, device, 0, 0, "", timed_out=False, row = _classify(case, out_dir, gpu, 0, 0, "", timed_out=False)
args=args)
row["notes"] = _join("resumed", row["notes"]) row["notes"] = _join("resumed", row["notes"])
return row return row
@@ -299,7 +295,7 @@ def run_case(case: Case, batch_dir: Path, device: str,
log = out_dir / "run.log" log = out_dir / "run.log"
if log.is_file(): if log.is_file():
log.rename(out_dir / f"run.attempt{attempt - 1}.log") log.rename(out_dir / f"run.attempt{attempt - 1}.log")
rc, timed_out = _wait(_spawn(case, out_dir, device, args, dtypes), rc, timed_out = _wait(_spawn(case, out_dir, gpu, args, dtypes),
args.op_timeout) args.op_timeout)
rcs.append(rc) rcs.append(rc)
if timed_out: if timed_out:
@@ -319,22 +315,20 @@ def run_case(case: Case, batch_dir: Path, device: str,
# 覆盖 baseline 对受限 dtype 编译失败(torch jiterator 的 special 函数族等) # 覆盖 baseline 对受限 dtype 编译失败(torch jiterator 的 special 函数族等)
# 但报错文案不是上游 "can't be supported by this op" 的场景。 # 但报错文案不是上游 "can't be supported by this op" 的场景。
if dtypes and not timed_out: if dtypes and not timed_out:
probe = _classify(case, out_dir, device, rcs[-1], 0, "", False, probe = _classify(case, out_dir, gpu, rcs[-1], 0, "", False)
args=args)
if probe["status"] == "FAIL" and probe["rows_success"] == "0": if probe["status"] == "FAIL" and probe["rows_success"] == "0":
log = out_dir / "run.log" log = out_dir / "run.log"
if log.is_file(): if log.is_file():
log.rename(out_dir / f"run.attempt{len(rcs) - 1}.log") log.rename(out_dir / f"run.attempt{len(rcs) - 1}.log")
notes = _join(notes, "dtype_rescue") notes = _join(notes, "dtype_rescue")
rc, timed_out = _wait(_spawn(case, out_dir, device, args, ""), rc, timed_out = _wait(_spawn(case, out_dir, gpu, args, ""),
args.op_timeout) args.op_timeout)
rcs.append(rc) rcs.append(rc)
elapsed_ms = int((time.time() - t0) * 1000) elapsed_ms = int((time.time() - t0) * 1000)
if len(rcs) > 1: if len(rcs) > 1:
notes = _join(notes, "retry_rc=" + "->".join(map(str, rcs))) notes = _join(notes, "retry_rc=" + "->".join(map(str, rcs)))
row = _classify(case, out_dir, device, rcs[-1], elapsed_ms, notes, timed_out, row = _classify(case, out_dir, gpu, rcs[-1], elapsed_ms, notes, timed_out)
args=args)
if row["status"] in ("PASS", "SKIP"): if row["status"] in ("PASS", "SKIP"):
tmp = marker.with_name(".complete.tmp") tmp = marker.with_name(".complete.tmp")
tmp.write_text(row["status"] + "\n") tmp.write_text(row["status"] + "\n")
@@ -351,21 +345,8 @@ def _join(*notes: str) -> str:
return ";".join(n for n in notes if n) return ";".join(n for n in notes if n)
def _identity_fields(out_dir: Path, device: str, backend: str) -> Dict[str, str]: def _classify(case: Case, out_dir: Path, gpu: str, rc: int, elapsed_ms: int,
ident = load_identity(str(out_dir)) notes: str, timed_out: bool) -> Dict[str, str]:
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) text = _read_log(out_dir)
rows = batch_summary.parse_run_log(text) rows = batch_summary.parse_run_log(text)
if timed_out: if timed_out:
@@ -375,13 +356,12 @@ def _classify(case: Case, out_dir: Path, device: str, rc: int, elapsed_ms: int,
n_skipped = batch_summary.pytest_skipped_count(text) n_skipped = batch_summary.pytest_skipped_count(text)
if status == "PASS" and n_skipped: if status == "PASS" and n_skipped:
notes = _join(notes, f"partial_skip={n_skipped}") notes = _join(notes, f"partial_skip={n_skipped}")
ident = _identity_fields(out_dir, device, args.backend)
return { return {
"repo": case.repo, "repo": case.repo,
"op": case.op, "op": case.op,
"op_file": case.op_file, "op_file": case.op_file,
"run_name": case.run_name, "run_name": case.run_name,
**ident, "gpu": gpu,
"status": status, "status": status,
"return_code": str(rc), "return_code": str(rc),
"elapsed_ms": str(elapsed_ms), "elapsed_ms": str(elapsed_ms),
@@ -413,12 +393,8 @@ def parse_args() -> argparse.Namespace:
help="--all 时不剔除带无条件 skip 标记的算子") help="--all 时不剔除带无条件 skip 标记的算子")
run = p.add_argument_group("执行") 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, run.add_argument("--gpus", default=None,
help="--devices 的别名(两者同时给出则报错)") help="逗号分隔 GPU 列表(默认 CUDA_VISIBLE_DEVICES 或全部)")
run.add_argument("--batch-dir", type=Path, default=None, run.add_argument("--batch-dir", type=Path, default=None,
help="产物目录(默认 runs/batch_<时间戳>;resume 时必须指定)") help="产物目录(默认 runs/batch_<时间戳>;resume 时必须指定)")
run.add_argument("--op-timeout", type=int, default=1800, run.add_argument("--op-timeout", type=int, default=1800,
@@ -453,30 +429,19 @@ def parse_args() -> argparse.Namespace:
def main() -> None: def main() -> None:
args = parse_args() 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() inv = load_inventory()
cases = resolve_selection(args, inv) cases = resolve_selection(args, inv)
if not cases: if not cases:
raise SystemExit(f"{LOG} 选择结果为空") raise SystemExit(f"{LOG} 选择结果为空")
gpus = visible_gpus(args.gpus)
batch_dir = (args.batch_dir or batch_dir = (args.batch_dir or
ZL_BENCH / "runs" / f"batch_{datetime.now():%Y%m%d-%H%M%S}") ZL_BENCH / "runs" / f"batch_{datetime.now():%Y%m%d-%H%M%S}")
batch_dir = batch_dir.resolve() batch_dir = batch_dir.resolve()
batch_dir.mkdir(parents=True, exist_ok=True) batch_dir.mkdir(parents=True, exist_ok=True)
parts = {d: cs for d, cs in partition(cases, devices).items() if cs} parts = {g: cs for g, cs in partition(cases, gpus).items() if cs}
print(f"{LOG} backend={args.backend} {len(cases)} 个算子 -> {len(parts)} 台设备 " print(f"{LOG} {len(cases)} 个算子 -> {len(parts)} 张卡 "
f"(shards: {[len(cs) for cs in parts.values()]}) out={batch_dir}") f"(shards: {[len(cs) for cs in parts.values()]}) out={batch_dir}")
if args.replay_root: if args.replay_root:
print(f"{LOG} replay 模式: REPLAY_FROM={args.replay_root}/<repo>/<op>") print(f"{LOG} replay 模式: REPLAY_FROM={args.replay_root}/<repo>/<op>")
@@ -496,19 +461,17 @@ def main() -> None:
writer.writerows(rows) writer.writerows(rows)
os.replace(tmp, status_csv) os.replace(tmp, status_csv)
def run_partition(device: str, part: List[Case]) -> None: def run_partition(gpu: str, part: List[Case]) -> None:
for case in part: for case in part:
try: try:
row = run_case(case, batch_dir, device, args) row = run_case(case, batch_dir, gpu, args)
except Exception as exc: # 编排层异常不拖垮整卡分片 except Exception as exc: # 编排层异常不拖垮整卡分片
out_dir = batch_dir / case.repo / case.run_name
ident = _identity_fields(out_dir, device, args.backend)
row = { row = {
"repo": case.repo, "op": case.op, "op_file": case.op_file, "repo": case.repo, "op": case.op, "op_file": case.op_file,
"run_name": case.run_name, **ident, "status": "FAIL", "run_name": case.run_name, "gpu": gpu, "status": "FAIL",
"return_code": "EXC", "elapsed_ms": "0", "rows_total": "0", "return_code": "EXC", "elapsed_ms": "0", "rows_total": "0",
"rows_success": "0", "notes": f"harness_exc={exc}", "rows_success": "0", "notes": f"harness_exc={exc}",
"out_dir": str(out_dir), "out_dir": str(batch_dir / case.repo / case.run_name),
} }
with lock: with lock:
done[0] += 1 done[0] += 1
@@ -517,16 +480,16 @@ def main() -> None:
print(f" ({done[0]:>4}/{len(cases)}) {row['status']:<7} " print(f" ({done[0]:>4}/{len(cases)}) {row['status']:<7} "
f"{row['elapsed_ms']:>8}ms rows={row['rows_success']}/" f"{row['elapsed_ms']:>8}ms rows={row['rows_success']}/"
f"{row['rows_total']:<4} {case.repo}:{case.run_name:<{wname}} " f"{row['rows_total']:<4} {case.repo}:{case.run_name:<{wname}} "
f"[device{device}]" f"[gpu{gpu}]"
+ (f" {row['notes']}" if row["notes"] else ""), + (f" {row['notes']}" if row["notes"] else ""),
flush=True) flush=True)
if len(parts) == 1: if len(parts) == 1:
device, part = next(iter(parts.items())) gpu, part = next(iter(parts.items()))
run_partition(device, part) run_partition(gpu, part)
else: else:
with ThreadPoolExecutor(max_workers=len(parts)) as pool: with ThreadPoolExecutor(max_workers=len(parts)) as pool:
futures = [pool.submit(run_partition, d, p) for d, p in parts.items()] futures = [pool.submit(run_partition, g, p) for g, p in parts.items()]
for f in futures: for f in futures:
f.result() f.result()
+24 -117
View File
@@ -8,21 +8,15 @@ set -euo pipefail
# 每项均可用同名环境变量覆盖,详见 README「环境变量一览」。 # 每项均可用同名环境变量覆盖,详见 README「环境变量一览」。
# 被测算子,对应 $FLAGGEMS_DIR/benchmark/test_<OP_FILE>.py::test_<OP> # 被测算子,对应 $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="${OP:-fused_marlin_moe_mxfp4}"
OP_FILE="${OP_FILE:-fused_marlin_moe}" OP_FILE="${OP_FILE:-fused_marlin_moe}"
# 上游已拆分为两个仓库:主仓 /workspace/dev/FlagGems(包名 flag_gems)与 # 上游已拆分为两个仓库:主仓 /workspace/dev/FlagGems(包名 flag_gems)与
# /workspace/dev/FlagGems-vllm(包名 flaggems_vllm)。测 vllm 仓库算子时覆盖本变量即可。 # /workspace/dev/FlagGems-vllm(包名 flaggems_vllm)。测 vllm 仓库算子时覆盖本变量即可。
# NV 产品路径也常见 /workspace/FlagGems;TEST_NODE 或缺失测试文件时会回退到该路径。
FLAGGEMS_DIR="${FLAGGEMS_DIR:-/workspace/dev/FlagGems}" FLAGGEMS_DIR="${FLAGGEMS_DIR:-/workspace/dev/FlagGems}"
TEST_NODE="${TEST_NODE:-${ZL_TEST_NODE:-}}"
# shape 来源:非空则用该 yaml;SHAPE_MODE=upstream 则不传 --shape_file,沿用 FlagGems # shape 来源:非空则用该 yaml,否则用 INLINE_YAML。顶层 key 必须是 op 名。
# 在 $BENCH_LEVEL 下的原生 shape 集。设了 TEST_NODE 且未指定 SHAPE_FILE 时默认 upstream。
SHAPE_FILE="${SHAPE_FILE:-}" SHAPE_FILE="${SHAPE_FILE:-}"
SHAPE_MODE="${SHAPE_MODE:-}"
read -r -d '' INLINE_YAML <<'YAML' || true read -r -d '' INLINE_YAML <<'YAML' || true
fused_marlin_moe_mxfp4: fused_marlin_moe_mxfp4:
shapes: shapes:
@@ -59,73 +53,25 @@ REPLAY_FROM="${REPLAY_FROM:-}"
# always/never 强制开关终端颜色,空则按 tty 判断(run.log 始终去色) # always/never 强制开关终端颜色,空则按 tty 判断(run.log 始终去色)
FLAGGEMS_PERF_COLOR="${FLAGGEMS_PERF_COLOR:-}" FLAGGEMS_PERF_COLOR="${FLAGGEMS_PERF_COLOR:-}"
# 设备后端:cuda / musa / rocm / ascend(别名见 backends/__init__.py)。 # 各插件作用见 README「插件说明」;注释掉某行即停用该插件
# 插件列表由 plugin_plan.py 按后端能力生成,不再在此硬编码 -p。 PLUGINS=(
# 非 CUDA 默认走上游 do_bench,不会加载 _cudagraph_plugin。 -p _device_guard_plugin
ZL_BACKEND="${ZL_BACKEND:-cuda}" -p _parallel_warmup_plugin
-p _seed_plugin
# 上游 --level:core(默认)只用各算子的 core_shapes;comprehensive 额外并入 -p _shape_inject_plugin
# 子类的 set_more_shapes(),并对同一 shape 再产出一次变体(BLAS 是转置侧)。 -p _shape_iter_inject_plugin
# 差距可以很大:mm 在 core 下是 5 个 shape,comprehensive 下并入 4 个大 K -p _bespoke_shape_plugin
# (N/K 到 152064)与 models_shapes/*.yaml 的模型 shape,共 29 个,且用例翻倍。 # -p _mm_cluster_fix_plugin
# 代价同步放大——sweep 是墙上时间的大头,shape 多时配合 PARALLEL_WARMUP_GPUS。 -p _autotune_record_plugin
BENCH_LEVEL="${BENCH_LEVEL:-core}" -p _cudagraph_plugin
case "$BENCH_LEVEL" in -p _pretty_report_plugin
core|comprehensive) ;; -p _ir_meta_plugin
*) echo "!!! BENCH_LEVEL 只能是 core 或 comprehensive(当前:$BENCH_LEVEL)" >&2; exit 2 ;; )
esac
# FlagGems 的 libtuner winner 存在 $FLAGGEMS_CACHE_DIR/config_cache 的 sqlite 里,
# 默认 ~/.flaggems,跨 run 长存。1=本次 run 用独立缓存目录,两侧各自冷调优
# (README「各自最优」口径);0=沿用共享缓存(默认,快)。
#
# 共享缓存的风险是单向的:先跑的一侧把 winner 写进去,后跑的一侧直接命中,
# 于是它只编译那几个胜出变体。实测 mm 上后跑侧的 TTGIR 变体从 945 塌到 33、
# cvt 从 136 变 10,看起来像被测改动的效果,其实只是继承了对侧的 config。
# 走 REPLAY_FROM 的标准 A/B 不受此影响(config 本来就锁死),record 对 record
# 才会中招。无论取哪个值都会写进 run_identity.json,口径不再隐形。
ZL_ISOLATE_TUNER_CACHE="${ZL_ISOLATE_TUNER_CACHE:-0}"
# ============================== 执行逻辑 ============================== # ============================== 执行逻辑 ==============================
# 前置检查。静默失败最贵:路径不存在时 pytest 收集不到 benchmark 目录,注册
# --level/--mode 的 conftest 就没加载,报错变成 "unrecognized arguments",指向
# 完全错误的方向。每项检查自报名字,失败即停。
preflight() {
local what="$1"; shift
if ! "$@"; then
echo "!!! PREFLIGHT_FAILED check=$what" >&2
exit 3
fi
}
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
export PYTHONPATH="$SCRIPT_DIR${PYTHONPATH:+:$PYTHONPATH}" TEST_FILE="$FLAGGEMS_DIR/benchmark/test_${OP_FILE}.py::test_${OP}"
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}}"
preflight "FlagGems directory exists: $FLAGGEMS_DIR" test -d "$FLAGGEMS_DIR"
# The node's file half; the ::test_<fn> half is left to pytest, which reports a
# missing function clearly ("not found: ...") while a missing FILE does not.
preflight "benchmark file exists: ${TEST_FILE%%::*}" test -f "${TEST_FILE%%::*}"
# The interpreter, not just the paths. triton and flag_gems normally come from a
# PYTHONPATH pointing at a compiler build and a source checkout, so a shell that
# merely reaches the right files can still be the wrong Python. Without this the
# failure surfaces per cell as `pytest exit 4` from a conftest ImportError, and a
# multi-round campaign burns every cell before anyone reads a log.
preflight "interpreter imports triton + flag_gems ($(command -v python))" sh -c 'python -c "import triton, flag_gems" 2>/dev/null'
if [[ -n "$SHAPE_FILE" ]]; then
preflight "shape file exists: $SHAPE_FILE" test -f "$SHAPE_FILE"
fi
if [[ -n "$REPLAY_FROM" ]]; then
preflight "replay run directory exists: $REPLAY_FROM" test -d "$REPLAY_FROM"
fi
# 产物目录,可用 OUT_DIR 指定。时间戳只到秒,并发启动会撞名,故用不带 -p 的 # 产物目录,可用 OUT_DIR 指定。时间戳只到秒,并发启动会撞名,故用不带 -p 的
# mkdir 抢占(已存在即失败),撞上就退让到 -2、-3…… # mkdir 抢占(已存在即失败),撞上就退让到 -2、-3……
@@ -158,12 +104,9 @@ else
PYTEST_COLOR=() PYTEST_COLOR=()
fi fi
export PYTHONPATH="$SCRIPT_DIR${PYTHONPATH:+:$PYTHONPATH}"
export FLAGGEMS_PERF_CURRENT_OP="$OP" export FLAGGEMS_PERF_CURRENT_OP="$OP"
export PYTHONUNBUFFERED=1 # 实时输出不缓冲 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 一次) # DTYPES 非空时逐个转为上游 --dtypes 选项(action=append,每个 dtype 一次)
DTYPE_ARGS=() DTYPE_ARGS=()
@@ -183,52 +126,17 @@ else
MODE_DESC=record MODE_DESC=record
fi fi
# 未指定 shape 文件时:inline 用 INLINE_YAML;upstream 不传 --shape_file。 # 未指定 shape 文件时,把 INLINE_YAML 落到临时文件供 pytest 读取
SHAPE_ARGS=() if [[ -z "$SHAPE_FILE" ]]; then
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)" SHAPE_FILE="$(mktemp --suffix=.yaml)"
printf '%s\n' "$INLINE_YAML" > "$SHAPE_FILE" printf '%s\n' "$INLINE_YAML" > "$SHAPE_FILE"
trap 'rm -f "$SHAPE_FILE"' EXIT trap 'rm -f "$SHAPE_FILE"' EXIT
fi
cp -f "$SHAPE_FILE" "$OUT_DIR/shapes.yaml" # 存档本次实际使用的 shape
SHAPE_ARGS=(--shape_file "$SHAPE_FILE")
fi fi
cp -f "$SHAPE_FILE" "$OUT_DIR/shapes.yaml" # 存档本次实际使用的 shape
# libtuner 缓存:隔离时用本 run 的私有目录,跑完随 Triton 缓存一起删(选中的
# config 已经在 autotune_records/ 里,DB 本身没有留存价值)。
TUNER_CACHE_ARGS=()
if [[ "$ZL_ISOLATE_TUNER_CACHE" == "1" ]]; then
TUNER_CACHE_DIR="$OUT_DIR/.flaggems_cache"
rm -rf "$TUNER_CACHE_DIR"; mkdir -p "$TUNER_CACHE_DIR"
TUNER_CACHE_ARGS=(FLAGGEMS_CACHE_DIR="$TUNER_CACHE_DIR")
TUNER_DESC=isolated
else
TUNER_DESC="shared(${FLAGGEMS_CACHE_DIR:-~/.flaggems})"
fi
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" \
--tuner-cache "$TUNER_DESC"
status=0 status=0
{ {
echo "${C_BOLD}>>> op=$OP backend=$ZL_BACKEND mode=$MODE_DESC USE_FLAGTUNE=$USE_FLAGTUNE shape=$SHAPE_MODE level=$BENCH_LEVEL tuner=$TUNER_DESC${C_RESET}" echo "${C_BOLD}>>> op=$OP mode=$MODE_DESC USE_FLAGTUNE=$USE_FLAGTUNE${C_RESET}"
echo "${C_DIM}>>> node=$TEST_FILE${C_RESET}"
echo "${C_DIM}>>> out=$OUT_DIR${C_RESET}" echo "${C_DIM}>>> out=$OUT_DIR${C_RESET}"
# 每次用独立的 Triton 缓存目录,跑完即删:保证编译过程可复现,且 ttgir 落盘 # 每次用独立的 Triton 缓存目录,跑完即删:保证编译过程可复现,且 ttgir 落盘
@@ -238,14 +146,13 @@ status=0
status=0 status=0
TRITON_CACHE_DIR="$CACHE_DIR" \ TRITON_CACHE_DIR="$CACHE_DIR" \
FLAGGEMS_PERF_TTGIR_DUMP_DIR="$OUT_DIR/ttgir" \ FLAGGEMS_PERF_TTGIR_DUMP_DIR="$OUT_DIR/ttgir" \
env "$AUTOTUNE_ENV" ${TUNER_CACHE_ARGS[@]+"${TUNER_CACHE_ARGS[@]}"} \ env "$AUTOTUNE_ENV" \
USE_FLAGTUNE=$USE_FLAGTUNE python -u -m pytest -s "$TEST_FILE" \ USE_FLAGTUNE=$USE_FLAGTUNE python -u -m pytest -s "$TEST_FILE" \
"${PLUGINS[@]}" "${PYTEST_COLOR[@]}" \ "${PLUGINS[@]}" "${PYTEST_COLOR[@]}" \
${SHAPE_ARGS[@]+"${SHAPE_ARGS[@]}"} ${DTYPE_ARGS[@]+"${DTYPE_ARGS[@]}"} \ --shape_file "$SHAPE_FILE" ${DTYPE_ARGS[@]+"${DTYPE_ARGS[@]}"} \
--level "$BENCH_LEVEL" --mode kernel || status=$? --level core --mode kernel || status=$?
rm -rf "$CACHE_DIR" rm -rf "$CACHE_DIR"
[[ -n "${TUNER_CACHE_DIR:-}" ]] && rm -rf "$TUNER_CACHE_DIR"
if (( status == 0 )); then if (( status == 0 )); then
echo "${C_GREEN}>>> done. outputs in $OUT_DIR${C_RESET}" echo "${C_GREEN}>>> done. outputs in $OUT_DIR${C_RESET}"
-110
View File
@@ -1,110 +0,0 @@
#!/usr/bin/env python3
"""采样设备时钟/温度/功耗,跨后端归一化成同一组列。
为什么要有它:六轮 A/B 会声称单侧噪声在 0.01% 量级。降频和热漂移足以移动比这大
得多的延迟,而且在 run.log 里不留任何痕迹——报告照样干净。NV 上复现 PR763 的那次
全程 `nvidia-smi -l 5` 写 clocks.csv,S5000 这边一直没有对应的东西。
输出每设备每次采样一行 TSV:
timestamp index sm_clock_mhz mem_clock_mhz temp_c power_w util_pct
列在 cuda 与 musa 上完全一致,所以两边的 campaign 可以并排读。采样器不存在时安静
退出:拿不到遥测是缺一份证据,不是让整批测试失败的理由。
用法:
python telemetry.py --backend musa --device 7 --interval 5 --run > telemetry.tsv
python telemetry.py --backend cuda # 只打印将要执行的采样命令
"""
from __future__ import annotations
import argparse
import os
import subprocess
import sys
import time
from typing import List, Optional
FIELDS = ("index", "sm_clock_mhz", "mem_clock_mhz", "temp_c", "power_w", "util_pct")
def resolve_device(backend, explicit: Optional[str]) -> str:
"""显式设备优先,否则从后端自己的可见性变量取。
单值才用得上:`MUSA_VISIBLE_DEVICES=7` 指一张卡,`0,1,2` 指一组,后者不传
`-i`、采全部,比挑第一张更诚实。
"""
if explicit:
return explicit.strip()
for key in getattr(backend, "visibility_env_keys", ()):
raw = os.environ.get(key, "").strip()
if raw and "," not in raw:
return raw
return ""
def sample_once(backend, argv: List[str]) -> List[dict]:
try:
out = subprocess.run(argv, capture_output=True, text=True, timeout=30)
except (OSError, subprocess.SubprocessError):
return []
return backend.parse_telemetry(out.stdout or "")
def run(backend, device: str, interval: float, out=sys.stdout) -> int:
argv = backend.telemetry_command(device)
if not argv:
print(f"[telemetry] backend={backend.name} 没有采样器;本批不记时钟/温度",
file=sys.stderr, flush=True)
return 0
if not sample_once(backend, argv):
# 第一次就取不到,说明工具缺失或没权限。不要留一个空文件让人以为采过了
print(f"[telemetry] {argv[0]} 取不到读数;本批不记时钟/温度",
file=sys.stderr, flush=True)
return 0
print("timestamp\t" + "\t".join(FIELDS), file=out, flush=True)
while True:
ts = time.strftime("%Y-%m-%dT%H:%M:%S")
for row in sample_once(backend, argv):
print(ts + "\t" + "\t".join(row.get(f, "") for f in FIELDS),
file=out, flush=True)
time.sleep(interval)
def main(argv: Optional[List[str]] = None) -> int:
p = argparse.ArgumentParser(description=__doc__,
formatter_class=argparse.RawDescriptionHelpFormatter)
p.add_argument("--backend", default=os.environ.get("ZL_BACKEND", "cuda"))
p.add_argument("--device", default=os.environ.get("ZL_DEVICE", ""))
p.add_argument("--interval", type=float, default=5.0)
p.add_argument("--run", action="store_true",
help="持续采样直到被杀;不给则只打印采样命令")
args = p.parse_args(argv)
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from backends import get_backend
backend = get_backend(args.backend)
device = resolve_device(backend, args.device)
if not args.run:
print(" ".join(backend.telemetry_command(device)))
return 0
try:
return run(backend, device, args.interval)
except KeyboardInterrupt:
return 0
except BrokenPipeError:
# 读端先走(管给 head,或 campaign 结束时收掉)。这是正常终止,不该在
# telemetry.err 里留一段 traceback 让人以为遥测本身坏了。按 CPython 的
# 惯用做法把 stdout 指到 devnull——只吞异常不够,解释器退出时还会 flush
# 一次,再抛一条 "Exception ignored" 到 stderr。
try:
devnull = os.open(os.devnull, os.O_WRONLY)
os.dup2(devnull, sys.stdout.fileno())
except Exception:
pass
return 0
if __name__ == "__main__":
raise SystemExit(main())
-738
View File
@@ -1,738 +0,0 @@
#!/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 types
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_loads_its_compat_plugin_first(self):
names = self._plan("musa")
self.assertIn("_musa_benchmark_compat_plugin", names)
# It rewrites collection-time markers, so nothing may inspect the item
# list before it runs.
self.assertEqual(names[0], "_musa_benchmark_compat_plugin")
def test_compat_plugins_do_not_leak_to_other_backends(self):
for name in ("cuda", "rocm", "ascend"):
self.assertNotIn("_musa_benchmark_compat_plugin", self._plan(name), name)
def test_musa_same_core_without_cluster_fix(self):
# _cudagraph_plugin is absent from the plan on every backend: the timing
# plugin loads it, so its absence here says nothing about capture support.
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 CaptureClassificationTests(unittest.TestCase):
"""Graph-capture failures split three ways; the wording is vendor text."""
def test_cuda_and_musa_both_capture(self):
self.assertTrue(get_backend("cuda").supports_cuda_graph)
# torch_musa supplies Stream/Event/graph/MUSAGraph; measured on S5000.
self.assertTrue(get_backend("musa").supports_cuda_graph)
def test_recoverable_is_the_default(self):
# A kernel refusing capture is ordinary and must reach the fallback.
for name in ("cuda", "musa"):
v = get_backend(name).classify_capture_error(
RuntimeError("operation not permitted when stream is capturing"))
self.assertEqual(v, "recoverable", name)
def test_cuda_fatal_markers(self):
b = get_backend("cuda")
self.assertEqual(
b.classify_capture_error(RuntimeError(
"CUDA error: an illegal memory access was encountered")), "fatal")
self.assertEqual(
b.classify_capture_error(RuntimeError("CUDA out of memory")), "fatal")
def test_torch_dsa_footer_does_not_make_everything_fatal(self):
# torch appends this to every AcceleratorError; matching it loosely
# would classify each capture refusal as fatal and kill the fallback.
exc = RuntimeError("Cannot call current_seed during capture. Compile with "
"`TORCH_USE_CUDA_DSA` to enable device-side assertions.")
self.assertEqual(get_backend("cuda").classify_capture_error(exc),
"recoverable")
def test_musa_unknown_error_is_poisoning_not_recoverable(self):
b = get_backend("musa")
exc = RuntimeError("MUSA error: unknown error")
self.assertEqual(b.classify_capture_error(exc), "poisoned")
# The same text carries no meaning for CUDA, so it must not leak across.
self.assertEqual(get_backend("cuda").classify_capture_error(exc),
"recoverable")
def test_musa_rng_seed_failure_stays_recoverable(self):
exc = RuntimeError("Cannot call MUSAGeneratorImpl::current_seed during "
"MUSA graph capture")
self.assertEqual(get_backend("musa").classify_capture_error(exc),
"recoverable")
def test_deny_list_is_per_backend(self):
self.assertTrue(get_backend("musa").denies_capture("unique2"))
self.assertFalse(get_backend("musa").denies_capture("softmax"))
self.assertFalse(get_backend("cuda").denies_capture("unique2"))
self.assertFalse(get_backend("musa").denies_capture(""))
def test_deny_list_is_keyed_by_pytest_function_name(self):
# FLAGGEMS_PERF_CURRENT_OP holds the pytest function name minus `test_`,
# so a deny entry written as the benchmark FILE name never matches and
# fails open. Observed once: "unique" listed, op=unique2 at runtime.
musa = get_backend("musa")
self.assertTrue(musa.denies_capture("unique2"))
self.assertFalse(musa.denies_capture("unique"))
class DirectRepeatTests(unittest.TestCase):
"""Repetition on the direct timer.
_cudagraph_plugin imports triton at module scope, so it is loaded here
against a stub: the repeat logic is plain arithmetic and should stay
testable on a machine with no GPU and no Triton.
"""
@classmethod
def setUpClass(cls):
cls._saved = {k: sys.modules.get(k) for k in ("triton", "triton.testing")}
triton = types.ModuleType("triton")
testing = types.ModuleType("triton.testing")
testing.do_bench = lambda *a, **k: 0.0
testing.do_bench_cudagraph = lambda *a, **k: 0.0
testing._summarize_statistics = lambda *a, **k: 0.0
triton.testing = testing
sys.modules["triton"] = triton
sys.modules["triton.testing"] = testing
import _cudagraph_plugin # noqa: E402
cls.cg = _cudagraph_plugin
@classmethod
def tearDownClass(cls):
for name, mod in cls._saved.items():
if mod is None:
sys.modules.pop(name, None)
else:
sys.modules[name] = mod
sys.modules.pop("_cudagraph_plugin", None)
def _run(self, value, **kwargs):
"""-> (result, how many times do_bench was called)"""
seq, calls = [3.0, 1.0, 2.0, 9.0, 5.0], []
self.cg._ORIGINAL_DO_BENCH = (
lambda fn, **kw: (calls.append(1), seq[(len(calls) - 1) % len(seq)])[1])
old = os.environ.get("ZL_DIRECT_REPEAT")
os.environ["ZL_DIRECT_REPEAT"] = value
try:
out = self.cg.run_direct(lambda: None, 25, 100, None,
kwargs.get("quantiles"),
kwargs.get("return_mode", "mean"))
finally:
if old is None:
os.environ.pop("ZL_DIRECT_REPEAT", None)
else:
os.environ["ZL_DIRECT_REPEAT"] = old
return out, len(calls)
def test_default_is_one_call(self):
self.assertEqual(self._run("1"), (3.0, 1))
def test_unset_is_one_call(self):
old = os.environ.pop("ZL_DIRECT_REPEAT", None)
try:
self.assertEqual(self.cg.direct_repeat(), 1)
finally:
if old is not None:
os.environ["ZL_DIRECT_REPEAT"] = old
def test_median_over_calls(self):
# 3, 1, 2, 9, 5 -> 3
self.assertEqual(self._run("5"), (3.0, 5))
def test_garbage_and_zero_fall_back_to_one(self):
self.assertEqual(self._run("abc"), (3.0, 1))
self.assertEqual(self._run("0"), (3.0, 1))
self.assertEqual(self._run("-4"), (3.0, 1))
def test_sequence_results_are_not_repeated(self):
# quantiles / return_mode="all" make do_bench return a sequence, and a
# median across calls is undefined for it; take the single call.
self.assertEqual(self._run("5", quantiles=[0.5])[1], 1)
self.assertEqual(self._run("5", return_mode="all")[1], 1)
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))
class CaliberGuardTests(unittest.TestCase):
"""两次运行口径不同产出的报告,和有效报告长得一模一样。
这些用例钉住的是本仓已经踩过的三次:record 对 record、共享 tuner cache、
timer 不同。它们都不会让任何一步报错,只会让数字变成另一个含义。
"""
BASE = {
"backend": "musa", "device": "S5000", "timer": "cuda_graph",
"bench_level": "core", "shape_mode": "default",
"tuner_cache": "isolated", "direct_repeat": 5,
"autotune_mode": "replay",
}
def _run(self, root, name, **over):
d = Path(root, name)
d.mkdir(parents=True)
if over.get("_no_identity"):
return d
ident = dict(self.BASE)
ident.update({k: v for k, v in over.items() if not k.startswith("_")})
Path(d, "run_identity.json").write_text(
json.dumps(ident), encoding="utf-8")
return d
def test_same_caliber_has_no_complaints(self):
from batch_summary import caliber_complaints
with tempfile.TemporaryDirectory() as tmp:
runs = [("A", self._run(tmp, "a")), ("B", self._run(tmp, "b"))]
self.assertEqual(caliber_complaints(runs), [])
def test_record_side_pinning_the_other_is_normal(self):
"""A/B 的正常形态:一侧 record 定 config,另一侧 replay 它。"""
from batch_summary import caliber_complaints
with tempfile.TemporaryDirectory() as tmp:
runs = [("r1_A", self._run(tmp, "a", autotune_mode="record")),
("r1_B", self._run(tmp, "b"))]
self.assertEqual(caliber_complaints(runs), [])
def test_record_against_record_is_flagged(self):
from batch_summary import caliber_complaints
with tempfile.TemporaryDirectory() as tmp:
runs = [("A", self._run(tmp, "a", autotune_mode="record")),
("B", self._run(tmp, "b", autotune_mode="record"))]
self.assertTrue(any("record" in c for c in caliber_complaints(runs)))
def test_each_caliber_field_is_checked(self):
from batch_summary import CALIBER_FIELDS, caliber_complaints
others = {"backend": "cuda", "device": "H20", "timer": "direct",
"bench_level": "comprehensive", "shape_mode": "yaml",
"tuner_cache": "shared", "direct_repeat": 1}
self.assertEqual(set(others), set(CALIBER_FIELDS))
for field, other in others.items():
with tempfile.TemporaryDirectory() as tmp:
runs = [("A", self._run(tmp, "a")),
("B", self._run(tmp, "b", **{field: other}))]
got = caliber_complaints(runs)
self.assertTrue(any(c.startswith(field) for c in got),
f"{field} 不一致时没被拦下: {got}")
def test_missing_identity_is_flagged_not_assumed_equal(self):
from batch_summary import caliber_complaints
with tempfile.TemporaryDirectory() as tmp:
runs = [("A", self._run(tmp, "a")),
("B", self._run(tmp, "b", _no_identity=True))]
self.assertTrue(any("run_identity" in c
for c in caliber_complaints(runs)))
def test_shared_cache_record_warns_but_does_not_block(self):
from batch_summary import caliber_complaints, caliber_hazards
with tempfile.TemporaryDirectory() as tmp:
runs = [("r1_A", self._run(tmp, "a", autotune_mode="record",
tuner_cache="shared(~/.flaggems)")),
("r1_B", self._run(tmp, "b",
tuner_cache="shared(~/.flaggems)"))]
self.assertEqual(caliber_complaints(runs), [])
self.assertTrue(caliber_hazards(runs))
def test_isolated_cache_has_no_hazard(self):
from batch_summary import caliber_hazards
with tempfile.TemporaryDirectory() as tmp:
runs = [("A", self._run(tmp, "a", autotune_mode="record")),
("B", self._run(tmp, "b"))]
self.assertEqual(caliber_hazards(runs), [])
def test_strict_refuses_and_permissive_continues(self):
from batch_summary import _report_caliber
with tempfile.TemporaryDirectory() as tmp:
runs = [("A", self._run(tmp, "a")),
("B", self._run(tmp, "b", timer="direct"))]
self.assertFalse(_report_caliber(runs, "op", strict=True))
self.assertTrue(_report_caliber(runs, "op", strict=False))
LOG = ("Operator: op (dtype=torch.float16, mode=kernel, level=core)\n"
"SUCCESS 1.000 2.000 0.500 shape\n")
def test_report_ab_refuses_mismatched_pair(self):
from batch_summary import CaliberMismatch, report_ab
with tempfile.TemporaryDirectory() as tmp:
a = self._run(tmp, "a")
b = self._run(tmp, "b", timer="direct")
for d in (a, b):
Path(d, "run.log").write_text(self.LOG, encoding="utf-8")
with self.assertRaises(CaliberMismatch):
report_ab(a / "run.log", b / "run.log", "op")
# 放行开关在,但必须是显式的
self.assertEqual(
len(report_ab(a / "run.log", b / "run.log", "op", strict=False)), 1)
class TelemetryTests(unittest.TestCase):
"""采样器命令与解析。两边归一到同一组列,campaign 才能并排读。"""
# 取自 S5000 实机:值带单位,"Power Draw " 的键名有尾随空格
MUSA_JSON = """
{
"Attached GPUs": "8",
"GPU": [
{
"Index": "7",
"Temperature": {"GPU Current Temp": "36C"},
"Clocks": {"Graphics": "1888MHz", "Memory": "2500MHz", "Video": "0MHz"},
"Power Readings": {"Power Draw ": "477.68W", "Current Power Limit": "950.00W"},
"Utilization": {"Gpu": "87%", "Memory": "14%"}
}
]
}
"""
CUDA_CSV = "0, 1980, 2619, 41, 122.35, 63\n"
def test_musa_parses_units_off_values(self):
rows = get_backend("musa").parse_telemetry(self.MUSA_JSON)
self.assertEqual(rows, [{
"index": "7", "sm_clock_mhz": "1888", "mem_clock_mhz": "2500",
"temp_c": "36", "power_w": "477.68", "util_pct": "87"}])
def test_cuda_parses_csv(self):
rows = get_backend("cuda").parse_telemetry(self.CUDA_CSV)
self.assertEqual(rows, [{
"index": "0", "sm_clock_mhz": "1980", "mem_clock_mhz": "2619",
"temp_c": "41", "power_w": "122.35", "util_pct": "63"}])
def test_both_backends_yield_the_same_columns(self):
import telemetry
a = get_backend("musa").parse_telemetry(self.MUSA_JSON)[0]
b = get_backend("cuda").parse_telemetry(self.CUDA_CSV)[0]
self.assertEqual(set(a), set(b))
self.assertEqual(set(a), set(telemetry.FIELDS))
def test_garbage_parses_to_nothing_rather_than_raising(self):
for backend in ("musa", "cuda"):
for raw in ("", "no such tool", "{broken json"):
self.assertEqual(get_backend(backend).parse_telemetry(raw), [])
def test_no_device_drops_the_flag_not_just_the_value(self):
# 留下裸 "-i" 会让采样器直接失败,比采全部更糟
argv = get_backend("musa").telemetry_command("")
self.assertNotIn("-i", argv)
self.assertIn("--json", argv)
self.assertEqual(get_backend("musa").telemetry_command("7")[-2:], ["-i", "7"])
def test_backend_without_sampler_returns_empty(self):
from backends.base import DeviceBackend
self.assertEqual(DeviceBackend().telemetry_command("0"), [])
def test_device_comes_from_the_visibility_var_when_unpinned(self):
import telemetry
backend = get_backend("musa")
old = os.environ.get("MUSA_VISIBLE_DEVICES")
try:
os.environ["MUSA_VISIBLE_DEVICES"] = "7"
self.assertEqual(telemetry.resolve_device(backend, ""), "7")
# 一组卡不挑第一张,交给采样器采全部
os.environ["MUSA_VISIBLE_DEVICES"] = "0,1,2"
self.assertEqual(telemetry.resolve_device(backend, ""), "")
self.assertEqual(telemetry.resolve_device(backend, "3"), "3")
finally:
if old is None:
os.environ.pop("MUSA_VISIBLE_DEVICES", None)
else:
os.environ["MUSA_VISIBLE_DEVICES"] = old
class ForeignVendorModuleTests(unittest.TestCase):
"""benchmark 的可用性判据常写成 `try: import <厂商实现>; AVAILABLE = True`。
import 只证明那个 .py 在磁盘上——它不碰设备,所以在任何后端上都成功,gate 于是
形同虚设。router_gemm 就这样让 NVIDIA Hopper 的 kernel 跑在了 MUSA 卡上,最后
死在 JIT 深处的 "dynamic_func() missing BLOCK_M ...",报错完全指不向真正的原因。
"""
P = "flag_gems.runtime.backend._"
def _mods(self, *names):
return {self.P + n: object() for n in names} | {"torch": object()}
def test_only_own_vendor_is_clean(self):
from _device_setup_plugin import foreign_vendor_modules
mods = self._mods("mthreads", "mthreads.ops.mm")
self.assertEqual(foreign_vendor_modules("mthreads", mods), [])
def test_foreign_vendor_is_reported_once_per_root(self):
from _device_setup_plugin import foreign_vendor_modules
# 同一厂商的多个子模块只报一条
mods = self._mods("mthreads.ops.mm",
"nvidia.hopper.ops.mm", "nvidia.ops.mm", "nvidia")
self.assertEqual(foreign_vendor_modules("mthreads", mods), ["nvidia"])
def test_several_foreign_vendors(self):
from _device_setup_plugin import foreign_vendor_modules
mods = self._mods("mthreads", "nvidia.ops.mm", "cambricon.ops.mm")
self.assertEqual(foreign_vendor_modules("mthreads", mods),
["cambricon", "nvidia"])
def test_a_vendor_whose_name_extends_ours_is_still_foreign(self):
# 按前缀比会把 _mthreads_extra 认成自家
from _device_setup_plugin import foreign_vendor_modules
mods = self._mods("mthreads", "mthreads_extra.ops")
self.assertEqual(foreign_vendor_modules("mthreads", mods),
["mthreads_extra"])
def test_unrelated_modules_are_ignored(self):
from _device_setup_plugin import foreign_vendor_modules
mods = {"torch": object(), "triton": object(),
"flag_gems.ops.mm": object()}
self.assertEqual(foreign_vendor_modules("mthreads", mods), [])
class PartitionTests(unittest.TestCase):
"""按耗时把算子分到各卡。hash 分片在这批活上比它慢 50 分钟。"""
HEADER = "op\top_file\telapsed_before\telapsed_after\tverdict"
def _load(self, assignment):
return sorted(sum(w for _, w in v) for v in assignment.values())
def _screen(self, tmp, rows, header=None):
p = Path(tmp, "screen.tsv")
p.write_text((header or self.HEADER) + "\n" + "\n".join(rows) + "\n",
encoding="utf-8")
return p
def test_heavy_first_keeps_cards_even(self):
from partition import plan
items = [("a", 10), ("b", 9), ("c", 8), ("d", 7), ("e", 6), ("f", 5)]
self.assertEqual(self._load(plan(items, ["0", "1", "2"])), [15, 15, 15])
def test_one_item_heavier_than_the_ideal_share_sets_the_floor(self):
# topk 单个 70 分钟、理想均分 45——加卡不解决,分片只保证别人不空等
from partition import imbalance, plan
items = [("topk", 70)] + [(f"op{i}", 5) for i in range(10)]
a = plan(items, [str(i) for i in range(8)])
self.assertEqual(max(sum(w for _, w in v) for v in a.values()), 70)
self.assertGreater(imbalance(a), 1.0)
def test_more_cards_than_operators_leaves_some_empty(self):
from partition import plan
a = plan([("a", 1), ("b", 1)], ["0", "1", "2", "3"])
self.assertEqual(sum(1 for v in a.values() if not v), 2)
def test_same_input_gives_the_same_plan(self):
from partition import plan
items = [("a", 5), ("b", 5), ("c", 5), ("d", 5)]
first = plan(items, ["0", "1"])
for _ in range(5):
self.assertEqual(plan(list(reversed(items)), ["0", "1"]), first)
def test_no_cards_is_the_callers_error(self):
from partition import plan
with self.assertRaises(ValueError):
plan([("a", 1)], [])
def test_perfect_balance_is_one(self):
from partition import imbalance, plan
self.assertEqual(imbalance(plan([("a", 5), ("b", 5)], ["0", "1"])), 1.0)
def test_verdict_filters_and_reports_what_it_dropped(self):
from partition import read_screen
with tempfile.TemporaryDirectory() as tmp:
f = self._screen(tmp, ["a\ta\t10\t10\tIR_CHANGED",
"b\tbb\t5\t5\tNO_IR_DIFF",
"c\tc\t1\t1\tNO_TRITON"])
items, files, skipped = read_screen(f)
self.assertEqual(items, [("a", 20.0)])
self.assertEqual(files["a"], "a")
self.assertEqual(len(skipped), 2)
def test_op_file_defaults_to_op_name(self):
from partition import read_screen
with tempfile.TemporaryDirectory() as tmp:
f = self._screen(tmp, ["a\t\t1\t1\tIR_CHANGED"])
_, files, _ = read_screen(f)
self.assertEqual(files["a"], "a")
def test_missing_elapsed_falls_back_to_equal_weights(self):
from partition import read_screen
with tempfile.TemporaryDirectory() as tmp:
f = self._screen(tmp, ["a\ta\tIR_CHANGED", "b\tb\tIR_CHANGED"],
header="op\top_file\tverdict")
items, _, _ = read_screen(f)
self.assertEqual([w for _, w in items], [1.0, 1.0])
if __name__ == "__main__":
unittest.main()