Aharrypotter/zl_bench
1
0 Fork 0
forked from zhoulin/zl_bench
Code Pull Requests Activity

feat: make the harness backend-agnostic

Browse Source 
Everything about which accelerator a run targets was previously spelled CUDA:
device discovery went through nvidia-smi, pinning through CUDA_VISIBLE_DEVICES,
the vendor guard only recognised NVIDIA, and the -p plugin list was hardcoded in
run_pytest.sh with no way to drop the CUDA-only members. Running the same
operator corpus on another vendor's device therefore meant editing the drivers.

The selection now hangs off one ZL_BACKEND variable (cuda/musa/rocm/ascend, with
the usual aliases):

- backends/ owns per-backend device enumeration, visibility env keys and
  capability flags, so nothing above it names a vendor;
- plugin_plan.py derives the -p list from those capabilities instead of a
  hardcoded array, which is what lets a non-CUDA run omit the graph-capture and
  cluster-fix plugins rather than watch them fail;
- _timing_plugin selects the timer from the same capabilities: CUDA keeps
  graph capture with its per-shape fallback, other backends stay on upstream
  do_bench and record timer=direct, because capture-then-fallback on a device
  without graph support only produces noise;
- _device_setup_plugin replaces the NVIDIA-only vendor guard;
- providers/flaggems.py resolves the pytest node, so a caller may pass TEST_NODE
  directly instead of the OP/OP_FILE pair, and SHAPE_MODE=upstream can defer to
  FlagGems' own core_shapes;
- _run_identity.py writes run_identity.json next to each run, recording backend,
  device, timer, autotune mode and the resolved test node. Two runs whose
  calibers differ are otherwise indistinguishable after the fact.

run_batch gains --backend and --devices (--gpus kept as an alias) and carries
backend/device/timer into ops_status.csv. test_backends.py covers the registry,
alias resolution, visibility env and plugin planning.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
This commit is contained in:
AharrypotterandClaude Opus 5 committed 2026-09-09 16:39:14 +08:00
1 parent 35a928a516
commit c8ea586d0f
24 files changed
+1448 -159
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

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