zl_bench/README.md

# zl_bench — FlagGems 单算子性能测试工具

针对编译器（FlagTree）改动做算子级 A/B 性能对比的 pytest 插件集 + 驱动脚本。在 FlagGems benchmark 体系之上解决三个问题：

1. **可复现**：固定随机种子、指定 shape，并通过 record/replay 让 A/B 两侧锁定同一套 autotune config，把两次运行之间的差异收敛到"编译器改动"这一个变量（**config 必须靠 replay 锁定，不会自己稳定**，见「为什么需要 replay」）；
2. **可解释**：自动按 shape 收集每次运行实际使用的 ttgir（带可读的变体命名），供 IR 级 diff；
3. **口径统一**：cudagraph 计时消除 launch 开销，小 kernel 的对比不被 CPU 侧噪声淹没；capture 前按 warmup 预算显式预热吸收 autotune/JIT，首轮即稳态、多轮一致。

## 快速开始

```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

# 测 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
```

上游现为两个仓库，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）。

shape yaml 顶层 key 必须是 op 名：

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

上面这条命令是 **record 模式**：现场 autotune、把选中的 config 存档。它用于单次摸底或给 replay 提供基准，**两次 record 的数字不能互相比**——做对比走下面的「A/B 对比测试标准流程」。

跑之前先 `nvidia-smi` 确认目标卡空闲：共享机器上别的任务会把两侧一起等比拖慢，出一份看似自洽实则作废的数据。多卡机器用 `CUDA_VISIBLE_DEVICES=<空闲卡号>` 选卡。**测量始终单进程单卡串行**，这是计时可比的前提；多卡仅用于可选的 sweep 加速（见「加速：多卡并行 sweep」）。

脚本内固定了 `--level core --mode kernel`，其余可调项都走环境变量（见下）；需要改动这两个口径本身时直接编辑 `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/dev/FlagGems` | FlagGems 仓库路径（测 vllm 仓库时指向 `/workspace/dev/FlagGems-vllm`） |

**怎么测**

| 变量 | 默认 | 说明 |
|-----|------|------|
| `REPLAY_FROM` | 空（record 模式） | 指向某次历史 run 目录，replay 其 autotune 选择。做 A/B 时 B 侧必须设（见「A/B 对比测试标准流程」） |
| `USE_FLAGTUNE` | `0` | `0` 走普通 autotune（跳过搜索、快速验证）；`1` 用 FlagTune 扩展调优空间（首跑全量搜索、慢）。A/B 两侧须取同值 |
| `PARALLEL_WARMUP_GPUS` | 空（串行） | 设为 `N`（≥2）时把 autotune sweep 分片到 N 卡并行，测量仍单卡串行；不设或 `<2` 则完全串行（见「加速：多卡并行 sweep」） |

**输出**

| 变量 | 默认 | 说明 |
|-----|------|------|
| `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 对比测试标准流程

```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
```

B 侧必须走 `REPLAY_FROM`。不能靠「改动前后各跑一次」来比——两次 record 会各自重新 sweep、可能选出不同 config，这个差异足以盖过被测改动本身（成因见「为什么需要 replay」）。

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

对比 `run.log` 的 latency 表看性能差异；diff 两侧 `ttgir/<shape>/<kernel>/` 下的同名文件看 IR 差异。

带 `pre_hook` 的 config（如 hopper mm 的 TMA configs）：libtuner kernel 正常 record/replay——注入的 config 会由上游按 kwargs 匹配自动接回 pre_hook（与其自身 ConfigCache 的 DB 回读同一条路径）；普通 `@triton.autotune` kernel 无此恢复机制，这类 config 不落 record，replay 时表现为下述 `key_missing` fallback。

replay 的兜底行为：B 侧遇到记录中没有的 key、或记录的 config 在新编译器下编译失败时，自动回退到现场 autotune 并在 `run.log` 打 `AUTOTUNE_REPLAY_FALLBACK reason=...` 标记——出现该标记的测量点不再满足"同 config"前提，解读时注意。标记以 `_no_evict` 结尾时更弱一层：坏 config 存在 libtuner 的 sqlite 缓存里、没有 `__delitem__` 可摘除，重试会再读到同一个 config——这类测量点按未验证处理。另外 replay 模式的 run 目录不产生 `autotune_records/`，后续 run 的 `REPLAY_FROM` 应始终指向最初 record 的那次 A 侧目录，不要链式指向 replay 产物。

### 为什么需要 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 在同一 db 里还持久化了 **BenchmarkCache**（sweep 中每个 config 的实测 latency）：db 已热时 record 模式的 sweep 也不重测，直接按历史延迟选 winner——"record 每次现场 sweep"仅在冷 db 下严格成立。下文"删 db 换各自最优口径"的操作会同时清掉 winner 与延迟两层，仍然有效。

**这不是理论风险，量级足以吞掉被测优化本身。** 同一份 `shapes.yaml`、同一口径、相隔十几分钟的两次 record，平均 speedup 可以差出 10% 量级，且逐 shape 单向偏移（不是随机噪声）。diff 两侧 `autotune_records/*.json` 能看到差异往往不是微调而是换挡——`num_stages`、`BLOCK_SIZE_*`、`num_warps` 整档跳变。怀疑遇到这种情况时，先 diff 两侧的 config 再看 latency。

所以 record 模式的数字只用来给 replay 提供 config 基准。**若某次结论只有 record 数据支撑，按未验证处理、重跑补 replay。**

libtuner 的持久缓存何时失效：FlagGems kernel 源码改动、tune_configs.yaml / expand yaml / `USE_FLAGTUNE` 开关变化、Triton 大版本或 GPU 型号变化。该 db 跨 run 长存，**本脚本只清理自己的 `TRITON_CACHE_DIR`（Triton 编译产物），不碰它**——清编译缓存不等于重新调优，config 命中后照样沿用旧 winner。如果需要"各自最优"口径（让两侧各自重新 sweep），**删掉 `~/.flaggems/config_cache/TunedConfig_*.db` 或设 `FLAGGEMS_DB_URL` 指向一次性文件**——注意"各自 fresh tune"不等于"不设 `REPLAY_FROM`"：缓存已热时两侧都会直接命中同一 winner，看似独立调优实则同 config。**两种口径都合理，报告结论时注明用的哪种。**

同一层缓存也决定了 `_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 测试参考示例

`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 轮：

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

可复用的是**口径**——同 config replay、cudagraph 计时、真实 trace shape 集、逐对成组——对比其它开关 / 编译器改动 / 算子，改循环变量与环境变量即可，口径部分不用动。

但**它是口径示例，不是可长期直接执行的回归脚本**：对比轴是 FlagGems 侧的算子开关，会随上游演进被改名、改语义或移除，本仓库不跟随同步更新（该脚本引用的开关目前就已被上游移除）。失效后果是**静默的**：脚本照样跑完全部轮次、照样输出完整对比表，只是两侧执行同一份代码，得出"无差异"的假结论。所以套用前先确认对比轴仍有效——`grep -rn '<开关名>' $FLAGGEMS_DIR/src` 要有命中，且两侧 latency 表确有差异。

## 加速：多卡并行 sweep（可选）

shape 多时墙上时间几乎全在 autotune sweep，而不在测量——被测 kernel 本身往往只有毫秒量级，绝大部分编译变体是 sweep 里测完就丢的 loser（`ttgir/index.tsv` 的 `sweep loser` 行）。`_parallel_warmup_plugin` 把 sweep 分片到 N 卡并行，再交回单卡串行测量：

```bash
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`，分片临时目录跑完即删。

**测量本身不并行**：N 个进程压满同机多卡会通过功耗墙/散热耦合，单卡 latency 被邻居拖慢且不可复现。并行只用于"决定哪个 config 胜出"，测量仍是单进程单卡、与不开插件走同一条路径。

因此要留意日志里的冲突警告——同一 config key 在不同分片选出了不同 winner，说明互扰已经影响到 sweep 结果：

```
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`——用它选过卡时，分片只会落在你选的那几张上。

## 计时口径：cudagraph 的预热、回退与精度

**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` 确认哪些行是回退口径再下结论。

## 插件说明

脚本通过 `-p` 加载以下插件（`run_pytest.sh` 的 `PLUGINS` 数组，可按需注释）：

| 插件 | 作用 | 何时关闭 |
|-----|------|---------|
| `_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」） | 不设该变量即关闭 |
| `_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: 同一份 shape、什么都没改，两次跑出来的 speedup 不一样？**

正常，且幅度可能不小——record 模式每次进程都重新 sweep Triton autotune，计时噪声会让不同 run 选中不同 config。想让两次可比，B 侧必须 `REPLAY_FROM` A 侧的 record 目录；diff 两侧 `autotune_records/*.json` 可确认 config 是否真的一致。实测幅度与成因见「为什么需要 replay」。

**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 秒级出结果。

shape 多时这部分会主导墙上时间，可设 `PARALLEL_WARMUP_GPUS=N` 并行做 sweep，见「加速：多卡并行 sweep」。

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

不用。cudagraph 计时路径下，传入的 warmup 时间预算会被用来在 capture 前显式预热（先吸收 autotune/JIT 编译再稳态预热，稳定首轮，见「计时口径：cudagraph 的预热、回退与精度」），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` 选项）；主仓与 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 满足）；
- 插件通过 monkeypatch 挂钩上游内部结构，FlagGems/Triton 大版本升级后若行为异常，优先检查各插件 pytest_configure 输出的注册日志是否还正常打印。