zl_bench/README.md

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 侧噪声淹没；capture 前按 warmup 预算显式预热吸收 autotune/JIT，首轮即稳态、多轮一致。

快速开始

# 默认算子（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 名：

softmax:
  shapes:
    - [1024, 1024]
    - [64, 512, 512]

跑之前先 nvidia-smi 确认目标卡空闲——共享机器上别的任务会把 baseline 和被测两边一起等比拖慢，出一份看似自洽实则作废的数据。多卡机器上用 CUDA_VISIBLE_DEVICES=<空闲卡号> 明确选卡。

脚本内固定了 --level core --mode kernel；USE_FLAGTUNE 默认 0（普通 autotune，快速出结果），需要 FlagTune 扩展调优空间时用 USE_FLAGTUNE=1 bash run_pytest.sh（首跑全量搜索、慢，见"为什么需要 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 选择（见下）
USE_FLAGTUNE	0	0 走普通 autotune（跳过搜索、快速验证）；1 用 FlagTune 扩展调优空间（首跑全量搜索、慢）
FLAGGEMS_PERF_COLOR	空（按 tty 自动判断）	always/never 强制开/关终端颜色；run.log 始终为去色纯文本

输出目录结构

每次运行产出 runs/<op>_<时间戳>/：

runs/softmax_20260716-031752/
├── run.log                  # 完整日志（含 SUCCESS 行的 latency/speedup 表；已去 ANSI 色的纯文本）
├── 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 对比测试标准流程

# 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

A/B 两侧的 USE_FLAGTUNE 必须取同值：该开关决定调优空间，两侧不一致时 replay 会大量 fallback，libtuner 持久缓存也各自独立命中，"同 config"前提不再成立。

对比 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 产物。

ab_fold_test.sh：FOLD 开关 A/B（现成示例）

ab_fold_test.sh 把上述流程封装成 FlagGems FLAGGEMS_MXFP4_FOLDSCALE（fold_scale 优化）开关的 A/B 验证：对 4 个 DeepSeek-V4-Flash 真实 trace shape 集（runs/shape/DeepSeek-V4-Flash-p{1024,4096,32768,65536}d1024.yaml，随仓库提供）各跑一对 FOLD=0（基线，record）→ FOLD=1（replay 同 config），共 8 轮：

bash ab_fold_test.sh
# 汇总 log：runs/ab_fold_<时间戳>.log（已去色）；快速对账：grep -E '^#####|FAILED' <log>

对比其它开关 / 算子时可直接改脚本内的循环变量与环境变量，A/B 口径（同 config replay、cudagraph 计时）无需改动。

cudagraph 计时的 warmup、回退与精度

capture 前的显式 warmup：do_bench_cudagraph 自带的内部预热只有 5 次迭代，对首跑要触发 autotune 编译（尤其含 FlagTune 扩展空间）、libtuner 选择、lazy JIT 的算子远远不够——这些一次性开销若漏进被捕获的 graph 或第一个计时迭代，测出的 latency 会 run-to-run 抖动（M=1 多 kernel 路径最明显，实测首轮可低到稳态的 ~1/3）。_cudagraph_plugin 因此在 capture 前显式预热：先跑一次并丢弃（吸收 autotune/JIT 编译），再按调用方传入的 warmup 时间预算（Config.warm_up）循环稳态预热，然后才 capture。预热次数按稳态单次耗时换算，并夹在 5–200 次之间——亚毫秒 kernel 的实际预热时长因此低于名义预算（1000ms），实测足够；若换新算子仍见首轮抖动，优先调大 _warmup_before_capture 里的次数上限。这样首轮即稳态、多轮一致（实测同一 M=1 shape 两轮 speedup 差 <0.1%）。

部分算子本身不支持 CUDA graph capture（测量函数内含 host 同步、动态显存分配、不合法的流操作等），这类算子会自动回退到普通 do_bench 计时，run.log 中打 BENCHMARK_DIRECT_NO_CUDAGRAPH 标记（含失败阶段与具体原因；phase=warmup 表示预热阶段就失败了，并非 graph capture 被拒）。

回退本身不影响 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 数组，可按需注释）：

插件	作用	何时关闭
_device_guard_plugin	在 import flag_gems 之前直接用 torch 探测到 NVIDIA 卡就设 GEMS_VENDOR=nvidia，跳过 flag_gems 启动时 nvidia-smi 子进程探测（该探测在部分 fork 环境下会挂在 wait4 上导致 import 卡死）	一般无需关：已设 GEMS_VENDOR/FLAGGEMS_VENDOR 等 env 时自动跳过，非 NVIDIA 卡上自动 no-op
_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 崩溃时打开

常见问题

Q: latency 和别人跑的差很多？

先检查 GPU 是否被共占（nvidia-smi）；再确认对方是否开了 cudagraph——M=1 这类小 shape 下 launch 开销占比大，两种口径可差 2 倍以上，大 shape 基本一致。

Q: 第一次跑某算子特别慢？

两个来源：（1）开了 USE_FLAGTUNE=1 时，首跑要在 FlagTune 扩展空间做全量搜索（fused_marlin_moe_mxfp4 单进程可达数分钟甚至十几分钟，期间 GPU 满载、run.log 停在测试名不动属正常，不是卡死）；（2）libtuner 算子首跑的全量 sweep（mm 约 50 分钟）。winner 持久化到 ~/.flaggems/config_cache/ 后同 shape 秒级命中（但该缓存会因源码/开关/GPU 变化失效，失效后又需重搜）。默认 USE_FLAGTUNE=0 走普通 autotune，单 shape 通常 10 秒级出结果。

Q: --warmup/--iter 要设吗？

不用。cudagraph 计时路径下，传入的 warmup 时间预算会被用来在 capture 前显式预热（先吸收 autotune/JIT 编译再稳态预热，稳定首轮，见"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 输出的注册日志是否还正常打印。