Stabilize first-run cudagraph timing and colorize terminal output
- _cudagraph_plugin: warm up autotune/JIT explicitly before graph capture (internal 5-iter warmup is too short), tag fallback markers with the failing phase; first-run latency no longer jitters run-to-run. - _device_guard_plugin (new): set GEMS_VENDOR via torch probe before importing flag_gems, avoiding its timeout-less nvidia-smi subprocess probe that can hang import in fork-broken environments. - _pretty_report_plugin (new) + _term_style (new): fold inputs identical across all result rows into a legend line, color status/plugin tags/markers on the live terminal; run.log is ANSI-stripped and keeps upstream SUCCESS/column wording for grep compatibility. - run_pytest.sh: make USE_FLAGTUNE overridable, group all knobs into a config section with Chinese comments, add start/end banners. - README: document the warmup semantics, new plugins/env vars, and the A/B rule that both sides must use the same USE_FLAGTUNE.
This commit is contained in:
@@ -4,7 +4,7 @@
|
||||
|
||||
1. **可复现**：固定随机种子、指定 shape、固定 autotune config，把 A/B 两次运行之间的差异收敛到"编译器改动"这一个变量；
|
||||
2. **可解释**：自动按 shape 收集每次运行实际使用的 ttgir（带可读的变体命名），供 IR 级 diff；
|
||||
3. **口径统一**：cudagraph 计时消除 launch 开销，小 kernel 的对比不被 CPU 侧噪声淹没。
|
||||
3. **口径统一**：cudagraph 计时消除 launch 开销，小 kernel 的对比不被 CPU 侧噪声淹没；capture 前按 warmup 预算显式预热吸收 autotune/JIT，首轮即稳态、多轮一致。
|
||||
|
||||
## 快速开始
|
||||
|
||||
@@ -28,7 +28,7 @@ softmax:
|
||||
|
||||
跑之前先 `nvidia-smi` 确认目标卡空闲——共享机器上别的任务会把 baseline 和被测两边一起等比拖慢，出一份看似自洽实则作废的数据。多卡机器上用 `CUDA_VISIBLE_DEVICES=<空闲卡号>` 明确选卡。
|
||||
|
||||
脚本内固定了 `--level core --mode kernel` 和 `USE_FLAGTUNE=1`（FlagTune 扩展调优空间，见"为什么需要 replay"一节）；需要改动这些口径时直接编辑 `run_pytest.sh` 中的 pytest 命令行。
|
||||
脚本内固定了 `--level core --mode kernel`；`USE_FLAGTUNE` 默认 `1`（FlagTune 扩展调优空间，见"为什么需要 replay"一节），可用环境变量覆盖（`USE_FLAGTUNE=0 bash run_pytest.sh` 走普通 autotune，跳过全量搜索，快速验证时用）；需要改动其它口径时直接编辑 `run_pytest.sh` 中的 pytest 命令行。
|
||||
|
||||
### 环境变量一览
|
||||
|
||||
@@ -39,6 +39,8 @@ softmax:
|
||||
| `SHAPE_FILE` | 空（用脚本内置 yaml） | shape yaml 路径 |
|
||||
| `FLAGGEMS_DIR` | `/workspace/FlagGems-dev` | FlagGems 仓库路径 |
|
||||
| `REPLAY_FROM` | 空（record 模式） | 指向某次历史 run 目录，replay 其 autotune 选择（见下） |
|
||||
| `USE_FLAGTUNE` | `1` | `1` 用 FlagTune 扩展调优空间（首跑全量搜索、慢）；`0` 走普通 autotune，跳过搜索、快速验证 |
|
||||
| `FLAGGEMS_PERF_COLOR` | 空（按 tty 自动判断） | `always`/`never` 强制开/关终端颜色；`run.log` 始终为去色纯文本 |
|
||||
|
||||
## 输出目录结构
|
||||
|
||||
@@ -46,7 +48,7 @@ softmax:
|
||||
|
||||
```
|
||||
runs/softmax_20260716-031752/
|
||||
├── run.log # 完整日志（含 SUCCESS 行的 latency/speedup 表）
|
||||
├── run.log # 完整日志（含 SUCCESS 行的 latency/speedup 表；已去 ANSI 色的纯文本）
|
||||
├── shapes.yaml # 本次实际使用的 shape（存档）
|
||||
├── autotune_records/ # Triton autotune 选中的 config（供 replay）
|
||||
│ └── softmax.json
|
||||
@@ -77,13 +79,17 @@ bash run_pytest.sh # -> runs/<op>_<ts_A>/
|
||||
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 产物。
|
||||
|
||||
### cudagraph 回退与测量精度
|
||||
### cudagraph 计时的 warmup、回退与精度
|
||||
|
||||
部分算子本身不支持 CUDA graph capture（测量函数内含 host 同步、动态显存分配、不合法的流操作等），这类算子会自动回退到普通 do_bench 计时，`run.log` 中打 `BENCHMARK_DIRECT_NO_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 倍以上，且行间波动更明显。解读这类算子的结果时：
|
||||
|
||||
@@ -109,12 +115,14 @@ libtuner 的持久缓存何时失效：FlagGems kernel 源码改动、tune_confi
|
||||
|
||||
| 插件 | 作用 | 何时关闭 |
|
||||
|-----|------|---------|
|
||||
| `_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）。无法 graph capture 的 kernel 自动回退并打 `BENCHMARK_DIRECT_NO_CUDAGRAPH` 标记 | 需要与他人的普通 do_bench 数据对齐时注释掉 |
|
||||
| `_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 崩溃时打开 |
|
||||
|
||||
@@ -126,11 +134,11 @@ libtuner 的持久缓存何时失效：FlagGems kernel 源码改动、tune_confi
|
||||
|
||||
**Q: 第一次跑某算子特别慢？**
|
||||
|
||||
libtuner 算子首跑要做全量 sweep（mm 约 50 分钟），winner 持久化到 `~/.flaggems/config_cache/` 后，之后同 shape 秒级命中。
|
||||
两个来源：（1）默认 `USE_FLAGTUNE=1`，首跑要在 FlagTune 扩展空间做全量搜索（fused_marlin_moe_mxfp4 单进程可达数分钟甚至十几分钟，期间 GPU 满载、`run.log` 停在测试名不动属正常，不是卡死）；（2）libtuner 算子首跑的全量 sweep（mm 约 50 分钟）。winner 持久化到 `~/.flaggems/config_cache/` 后同 shape 秒级命中（但该缓存会因源码/开关/GPU 变化失效，失效后又需重搜）。快速验证（不追求"各自最优 config"、只看功能/稳定性）时直接 `USE_FLAGTUNE=0 bash run_pytest.sh` 走普通 autotune，单 shape 通常 10 秒级出结果。
|
||||
|
||||
**Q: `--warmup/--iter` 要设吗？**
|
||||
|
||||
不用。cudagraph 计时路径下 warmup 参数本就不生效（内部自带预热），iter 默认 100ms 预算按 kernel 耗时自适应换算次数。
|
||||
不用。cudagraph 计时路径下，传入的 warmup 时间预算会被用来在 capture 前显式预热（先吸收 autotune/JIT 编译再稳态预热，稳定首轮，见"cudagraph 计时的 warmup、回退与精度"），iter 默认 100ms 预算按 kernel 耗时自适应换算次数。
|
||||
|
||||
**Q: ttgir 目录里某个 shape 少了文件？**
|
||||
|
||||
|
||||
Reference in New Issue
Block a user