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