diff --git a/README.md b/README.md index c3f1b0c..0dd1ad2 100644 --- a/README.md +++ b/README.md @@ -2,7 +2,7 @@ 针对编译器（FlagTree）改动做算子级 A/B 性能对比的 pytest 插件集 + 驱动脚本。在 FlagGems benchmark 体系之上解决三个问题： -1. **可复现**：固定随机种子、指定 shape、固定 autotune config，把 A/B 两次运行之间的差异收敛到"编译器改动"这一个变量； +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，首轮即稳态、多轮一致。 @@ -26,20 +26,35 @@ softmax: - [64, 512, 512] ``` -跑之前先 `nvidia-smi` 确认目标卡空闲——共享机器上别的任务会把 baseline 和被测两边一起等比拖慢，出一份看似自洽实则作废的数据。多卡机器上用 `CUDA_VISIBLE_DEVICES=<空闲卡号>` 明确选卡。 +上面这条命令是 **record 模式**：现场 autotune、把选中的 config 存档。它用于单次摸底或给 replay 提供基准，**两次 record 的数字不能互相比**——做对比走下面的「A/B 对比测试标准流程」。 -脚本内固定了 `--level core --mode kernel`；`USE_FLAGTUNE` 默认 `0`（普通 autotune，快速出结果），需要 FlagTune 扩展调优空间时用 `USE_FLAGTUNE=1 bash run_pytest.sh`（首跑全量搜索、慢，见"为什么需要 replay"一节）；需要改动其它口径时直接编辑 `run_pytest.sh` 中的 pytest 命令行。 +跑之前先 `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/FlagGems-dev` | FlagGems 仓库路径 | -| `REPLAY_FROM` | 空（record 模式） | 指向某次历史 run 目录，replay 其 autotune 选择（见下） | -| `USE_FLAGTUNE` | `0` | `0` 走普通 autotune（跳过搜索、快速验证）；`1` 用 FlagTune 扩展调优空间（首跑全量搜索、慢） | +**测什么** + +| 变量 | 默认 | 说明 | +|-----|------|------| +| `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 选择。做 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` 始终为去色纯文本 | ## 输出目录结构 @@ -79,11 +94,30 @@ bash run_pytest.sh # -> runs/_/ REPLAY_FROM=$PWD/runs/_ bash run_pytest.sh ``` +B 侧必须走 `REPLAY_FROM`。不能靠「改动前后各跑一次」来比——两次 record 会各自重新 sweep、可能选出不同 config，这个差异足以盖过被测改动本身（成因见「为什么需要 replay」）。 + A/B 两侧的 `USE_FLAGTUNE` 必须取同值：该开关决定调优空间，两侧不一致时 replay 会大量 fallback，libtuner 持久缓存也各自独立命中，"同 config"前提不再成立。 对比 `run.log` 的 latency 表看性能差异；diff 两侧 `ttgir///` 下的同名文件看 IR 差异。 -replay 的兜底行为：B 侧遇到记录中没有的 key、或记录的 config 在新编译器下编译失败时，自动回退到现场 autotune 并在 `run.log` 打 `AUTOTUNE_REPLAY_FALLBACK reason=...` 标记——出现该标记的测量点不再满足"同 config"前提，解读时注意。另外 replay 模式的 run 目录不产生 `autotune_records/`，后续 run 的 `REPLAY_FROM` 应始终指向最初 record 的那次 A 侧目录，不要链式指向 replay 产物。 +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 下的编译器差异"而非"各自最优" | + +**这不是理论风险，量级足以吞掉被测优化本身。** 同一份 `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//autotune_records/.json` 里有无对应 kernel 条目——漏记时该文件照样生成，只是少了 libtuner 那几个。 ### ab_fold_test.sh：A/B 测试参考示例 @@ -94,9 +128,31 @@ bash ab_fold_test.sh # 汇总 log：runs/ab_fold_<时间戳>.log（已去色）；快速对账：grep -E '^#####|FAILED' ``` -要对比其它开关 / 编译器改动 / 算子，套用该脚本改循环变量与环境变量即可，A/B 口径（同 config replay、cudagraph 计时）无需改动。 +可复用的是**口径**——同 config replay、cudagraph 计时、真实 trace shape 集、逐对成组——对比其它开关 / 编译器改动 / 算子，改循环变量与环境变量即可，口径部分不用动。 -### cudagraph 计时的 warmup、回退与精度 +但**它是口径示例，不是可长期直接执行的回归脚本**：对比轴是 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/.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%）。 @@ -109,26 +165,14 @@ bash ab_fold_test.sh 回退是按测量点发生的，同一份 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 型号变化。如果需要"各自最优"口径（让两侧各自重新 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//autotune_records/.json` 里有无对应 kernel 条目——漏记时该文件照样生成，只是少了 libtuner 那几个。 - ## 插件说明 脚本通过 `-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 | +| `_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 等） | 不关 | @@ -145,13 +189,19 @@ libtuner 的持久缓存何时失效：FlagGems kernel 源码改动、tune_confi 先检查 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 计时的 warmup、回退与精度"），iter 默认 100ms 预算按 kernel 耗时自适应换算次数。 +不用。cudagraph 计时路径下，传入的 warmup 时间预算会被用来在 capture 前显式预热（先吸收 autotune/JIT 编译再稳态预热，稳定首轮，见「计时口径：cudagraph 的预热、回退与精度」），iter 默认 100ms 预算按 kernel 耗时自适应换算次数。 **Q: ttgir 目录里某个 shape 少了文件？** diff --git a/_autotune_record_plugin.py b/_autotune_record_plugin.py index 6084167..4b016a9 100644 --- a/_autotune_record_plugin.py +++ b/_autotune_record_plugin.py @@ -180,12 +180,21 @@ def _replay_run(original): except Exception as exc: if injected_key is not None: # Recorded config failed under the current build: drop it and - # let run() autotune. + # let run() autotune. Only Autotuner.cache is a plain dict; + # LibTuner.cache is a sqlite-backed ConfigCache with no + # __delitem__, so eviction is impossible there -- say so in the + # marker instead of retrying with the same bad config and + # reporting a clean fallback. + evicted = True try: del self.cache[injected_key] except Exception: - pass - _emit_marker(f"compile_{type(exc).__name__}") + evicted = False + # _no_evict means the retry below re-reads the same recorded + # config, so it is not a clean "fell back to live autotune": + # treat those measurement points as unverified. + suffix = "" if evicted else "_no_evict" + _emit_marker(f"compile_{type(exc).__name__}{suffix}") return original(self, *args, **kwargs) raise return runner diff --git a/_parallel_warmup_plugin.py b/_parallel_warmup_plugin.py new file mode 100644 index 0000000..8bafa2e --- /dev/null +++ b/_parallel_warmup_plugin.py @@ -0,0 +1,306 @@ +"""pytest plugin: cut wall time by doing the autotune sweep on N GPUs in parallel. + +Why: for shape-heavy ops the sweep dominates, not the measurement. Measured on +fused_marlin_moe_mxfp4 / 53 shapes: 3755s total, of which the summed kernel time +is 49ms; ttgir/index.tsv holds 2858 compiled variants, 2645 of them `sweep +loser`. The same 53 shapes replayed (configs pinned, no sweep) take 196s. The +sweep is ~19x the measurement, so that is the part worth parallelizing. + +What this does NOT parallelize: the measurement. N processes hammering N GPUs of +one box couple through the power/thermal budget, so per-card latency gets dragged +by whatever the neighbours are doing, by an amount that does not reproduce. This +plugin only parallelizes "decide which config wins", then hands the configs to +the normal single-process single-GPU path, which times things exactly as it would +without the plugin. + +Enable with PARALLEL_WARMUP_GPUS=; unset (or <2) makes the plugin a no-op, so +`-p _parallel_warmup_plugin` can stay in the plugin list permanently. + +No new output structure: the merged configs land in the run's own +autotune_records/.json — the same path record mode writes and REPLAY_FROM +reads — and the sweep's scratch dirs are deleted. The run just finishes sooner, +plus a few log lines. (It is the config actually used for the measurement, so the +artifact stays faithful to what was run.) +""" + +from __future__ import annotations + +import json +import os +import shutil +import subprocess +import sys +import tempfile +import time +from pathlib import Path +from typing import Any, Dict, List, Tuple + +import pytest + +_GPUS_ENV = "PARALLEL_WARMUP_GPUS" +_WORKER_ENV = "FLAGGEMS_PERF_WARMUP_WORKER" +_RECORD_DIR_ENV = "FLAGGEMS_PERF_AUTOTUNE_RECORD_DIR" +_REPLAY_DIR_ENV = "FLAGGEMS_PERF_AUTOTUNE_REPLAY_DIR" +_OP_ENV = "FLAGGEMS_PERF_CURRENT_OP" + + +def _op_name() -> str: + return os.environ.get(_OP_ENV, "").strip() or "default" + + +def _requested_gpus() -> int: + """N from the env, clamped to what the box has. 0 disables.""" + raw = os.environ.get(_GPUS_ENV, "").strip() + if not raw or os.environ.get(_WORKER_ENV): + return 0 + try: + want = int(raw) + except ValueError: + return 0 + if want < 2: + return 0 + try: + import torch + have = torch.cuda.device_count() + except Exception: + return 0 + return max(0, min(want, have)) + + +def _visible_devices() -> List[str]: + """The device ids a shard may be pinned to, in parent-visible order. + + Must respect an inherited CUDA_VISIBLE_DEVICES: writing a bare shard index + into the child would re-interpret it as an absolute id, so a caller who + picked idle cards (CUDA_VISIBLE_DEVICES=5,6) would silently get cards 0,1 — + exactly the busy-GPU case the caller was avoiding. + """ + raw = os.environ.get("CUDA_VISIBLE_DEVICES", "").strip() + if raw: + return [d.strip() for d in raw.split(",") if d.strip()] + try: + import torch + return [str(i) for i in range(torch.cuda.device_count())] + except Exception: + return [] + + +def _read_shape_yaml(path: str, op: str) -> Tuple[List[Any], Dict[str, Any]]: + """Shapes for this op, plus the sibling keys to carry into each shard.""" + import yaml + with open(path) as f: + doc = yaml.safe_load(f) or {} + entry = doc.get(op) or {} + shapes = entry.get("shapes") or [] + extra = {k: v for k, v in entry.items() if k != "shapes"} + return list(shapes), extra + + +def _shard(shapes: List[Any], n: int) -> List[List[Any]]: + """Round-robin, not contiguous blocks. + + Shape lists are usually sorted ascending, so a block split hands one worker + every big (slowest-to-compile) shape and the wall time collapses to that + worker. Round-robin spreads compile cost evenly. + """ + return [s for s in (shapes[i::n] for i in range(n)) if s] + + +def _spawn(shards: List[List[Any]], op: str, extra: Dict[str, Any], + scratch: Path, argv: List[str], devices: List[str] + ) -> Tuple[List[Path], List[int]]: + """One subprocess per shard, each pinned to its own GPU. + + Subprocesses rather than pytest-xdist: _device_guard_plugin initializes a + CUDA context at import time, which does not survive fork-based parallelism. + """ + import yaml + procs: List[Tuple[int, subprocess.Popen, Any]] = [] + recs: List[Path] = [] + for i, shard in enumerate(shards): + w = scratch / f"w{i}" + w.mkdir(parents=True, exist_ok=True) + body = dict(extra) + body["shapes"] = shard + (w / "shapes.yaml").write_text( + yaml.safe_dump({op: body}, sort_keys=False, + default_flow_style=None, allow_unicode=True)) + rec = w / "rec" + rec.mkdir(exist_ok=True) + + child_argv: List[str] = [] + skip_next = False + for a in argv: + if skip_next: + skip_next = False + continue + if a == "--shape_file": + skip_next = True + continue + if a.startswith("--shape_file="): + continue + child_argv.append(a) + child_argv += ["--shape_file", str(w / "shapes.yaml")] + + env = dict(os.environ) + env["CUDA_VISIBLE_DEVICES"] = devices[i] + env[_WORKER_ENV] = "1" # stops recursion + env[_RECORD_DIR_ENV] = str(rec) # shard records its own picks + env.pop(_REPLAY_DIR_ENV, None) # a shard must sweep + env.pop(_GPUS_ENV, None) + env["TRITON_CACHE_DIR"] = str(w / ".triton_cache") + env.pop("FLAGGEMS_PERF_TTGIR_DUMP_DIR", None) # IR comes from the real run + env["FLAGGEMS_PERF_COLOR"] = "never" + + log = open(w / "worker.log", "w") + procs.append((i, subprocess.Popen( + [sys.executable, "-u", "-m", "pytest", *child_argv], + stdout=log, stderr=subprocess.STDOUT, env=env), log)) + recs.append(rec) + + try: + for _, p, _log in procs: + p.wait() + except BaseException: + # Ctrl-C (or anything else) must not leave N children holding GPUs: + # terminate, then reap, then let the exception continue. + for _, p, _log in procs: + if p.poll() is None: + p.terminate() + for _, p, _log in procs: + try: + p.wait(timeout=10) + except subprocess.TimeoutExpired: + p.kill() + raise + finally: + for _, _p, log in procs: + log.close() + return recs, [i for i, p, _ in procs if p.returncode != 0] + + +def _merge(recs: List[Path], op: str) -> Tuple[Dict[str, Any], int, int]: + """Union the shards' {kernel: {key: config}} maps. + + A key present in several shards with *different* values means that key's + winner moved under parallel interference; count those — the count is how much + to trust this warmup. First value wins. + """ + merged: Dict[str, Dict[str, Any]] = {} + conflicts = 0 + for rec in recs: + p = rec / f"{op}.json" + if not p.is_file(): + continue + try: + data = json.loads(p.read_text()) + except Exception: + continue + for kernel, bucket in data.items(): + if not isinstance(bucket, dict): + continue + tgt = merged.setdefault(kernel, {}) + for key, cfg in bucket.items(): + if key in tgt: + conflicts += tgt[key] != cfg + continue + tgt[key] = cfg + return merged, sum(len(b) for b in merged.values()), conflicts + + +@pytest.hookimpl(tryfirst=True) +def pytest_configure(config): + """Do the parallel sweep here, before _autotune_record_plugin configures. + + That plugin decides record-vs-replay and patches Autotuner.run inside its own + pytest_configure, so the swap to replay has to be in place before it runs — + hence tryfirst. Ordering within the -p list is not relied upon. + """ + from _term_style import tag + gpus = _requested_gpus() + if not gpus: + return + + op = _op_name() + shape_file = getattr(config.option, "shape_file", "") or "" + if not shape_file or not Path(shape_file).is_file(): + print(f"{tag('[parallel-warmup-plugin]')} no --shape_file; " + "parallel warmup needs an explicit shape set, skipping", flush=True) + return + if os.environ.get(_REPLAY_DIR_ENV, "").strip(): + print(f"{tag('[parallel-warmup-plugin]')} REPLAY_FROM is set; configs are " + "already pinned, nothing to sweep, skipping", flush=True) + return + + try: + shapes, extra = _read_shape_yaml(shape_file, op) + except Exception as exc: + print(f"{tag('[parallel-warmup-plugin]')} cannot read shapes " + f"({type(exc).__name__}: {exc}); skipping", flush=True) + return + if len(shapes) < 2: + print(f"{tag('[parallel-warmup-plugin]')} only {len(shapes)} shape(s); " + "sharding would not pay off, skipping", flush=True) + return + + devices = _visible_devices() + if len(devices) < 2: + print(f"{tag('[parallel-warmup-plugin]')} only {len(devices)} visible " + "GPU(s); nothing to parallelize, skipping", flush=True) + return + + shards = _shard(shapes, min(gpus, len(shapes), len(devices))) + scratch = Path(tempfile.mkdtemp(prefix=f"warmup_{op}_")) + print(f"{tag('[parallel-warmup-plugin]')} sweeping {len(shapes)} shapes on " + f"{len(shards)} GPUs (shards: {[len(s) for s in shards]}); this phase's " + "latency is discarded, only configs are kept", flush=True) + + keep_scratch = False + t0 = time.time() + try: + recs, bad = _spawn(shards, op, extra, scratch, + list(config.invocation_params.args), devices) + merged, total, conflicts = _merge(recs, op) + elapsed = time.time() - t0 + + if not total: + keep_scratch = True + print(f"{tag('[parallel-warmup-plugin]')} no configs recovered in " + f"{elapsed:.0f}s; falling back to normal serial autotune " + f"(worker logs kept in {scratch})", flush=True) + return + + # Publish the merged configs where record mode would have written them, + # so the artifact layout is unchanged and REPLAY_FROM still works. Then + # point the record plugin at that file in replay mode: the serial + # measurement below reuses these configs instead of sweeping again. + rec_dir = os.environ.get(_RECORD_DIR_ENV, "").strip() + out = Path(rec_dir) if rec_dir else (scratch / "merged") + out.mkdir(parents=True, exist_ok=True) + (out / f"{op}.json").write_text( + json.dumps(merged, indent=2, ensure_ascii=False) + "\n") + if not rec_dir: + keep_scratch = True # nowhere else to keep the configs + os.environ.pop(_RECORD_DIR_ENV, None) # record+replay is rejected + os.environ[_REPLAY_DIR_ENV] = str(out) + + msg = (f"{tag('[parallel-warmup-plugin]')} sweep done in {elapsed:.0f}s: " + f"{len(merged)} kernels / {total} configs -> {out / f'{op}.json'}; " + "measurement continues serially on one GPU") + if bad: + msg += f" (WARNING shard(s) {bad} exited non-zero; missing keys will " + msg += "fall back to live autotune)" + print(msg, flush=True) + if conflicts: + print(f"{tag('[parallel-warmup-plugin]')} WARNING {conflicts} config " + "key(s) got different winners across shards — parallel " + "interference reached the sweep. Treat this run as a rough " + f"pass; unset {_GPUS_ENV} for a clean baseline", flush=True) + except Exception as exc: + keep_scratch = True + print(f"{tag('[parallel-warmup-plugin]')} warmup failed " + f"({type(exc).__name__}: {exc}); falling back to normal serial " + f"autotune (scratch kept in {scratch})", flush=True) + finally: + if not keep_scratch: + shutil.rmtree(scratch, ignore_errors=True) diff --git a/ab_fold_test.sh b/ab_fold_test.sh index 04708e4..7386edf 100755 --- a/ab_fold_test.sh +++ b/ab_fold_test.sh @@ -1,17 +1,19 @@ #!/bin/bash -# A/B 测试参考示例：FLAGGEMS_MXFP4_FOLDSCALE（fold_scale）开关验证 -# （4 个 DeepSeek-V4-Flash trace yaml × FOLD=0/1 共 8 轮）。 -# 口径：USE_FLAGTUNE=0；FOLD=0（基线）先跑并 record autotune config，FOLD=1 用 -# REPLAY_FROM 回放同一套 config（两侧同配置比 wall-time）。 -# 对比其它开关/算子时，改循环变量与环境变量即可复用同一口径。 -# 全部输出汇总到 runs/ab_fold_<时间戳>.log；各轮 runs/_<时间戳>/ 产物照常存档。 -# 用法：bash ab_fold_test.sh +# A/B 对比口径示例：4 个 trace shape 集 × 开关 0/1，共 8 轮。 +# +# 口径：每个 shape 集先跑基线（record autotune config），再用 REPLAY_FROM 回放 +# 同一套 config 跑对照侧，两侧只差被测开关。对比其它开关/算子时改循环变量即可。 +# +# 注意：对比轴 FLAGGEMS_MXFP4_FOLDSCALE 属于 FlagGems，已被上游移除。开关不存在 +# 时本脚本仍会跑完并输出完整表格，但两侧执行同一份代码——套用前先确认对比轴有效 +# （grep 该开关名于 $FLAGGEMS_DIR/src 应有命中）。详见 README。 +# +# 输出：runs/ab_fold_<时间戳>.log（汇总）+ 各轮 runs/_<时间戳>/ set -uo pipefail cd "$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" -# 终端颜色：下方 tee 管道会让 run_pytest.sh 侧检测不到 tty，在这里先判断并下传 -# （与 run_pytest.sh 自身的做法一致）；总 log 最后统一去色。 +# tee 管道会让子脚本检测不到 tty，故在此判断后下传（同 run_pytest.sh 的做法） if [[ -t 1 && -z "${NO_COLOR:-}" ]]; then export FLAGGEMS_PERF_COLOR=always fi @@ -31,7 +33,7 @@ LOG="runs/ab_fold_$(date +%Y%m%d-%H%M%S).log" echo "########## ALL DONE ##########" } 2>&1 | tee "$LOG" -# 总 log 去掉 ANSI 转义，保证 grep/diff 面对纯文本（终端输出保留颜色） +# 去掉 ANSI 转义以便 grep/diff（终端输出保留颜色） sed -i -E $'s/\x1b\\[[0-9;]*[A-Za-z]//g' "$LOG" echo ">>> 总 log: $(pwd)/$LOG" diff --git a/run_pytest.sh b/run_pytest.sh index 1fbf568..9d537e8 100644 --- a/run_pytest.sh +++ b/run_pytest.sh @@ -1,23 +1,19 @@ #!/bin/bash -# 单算子性能测试入口（手动调试用）。 -# 用法：改下方【配置区】，或用环境变量覆盖，例如： +# 单算子性能测试入口。改配置区或用同名环境变量覆盖： # OP=softmax OP_FILE=softmax SHAPE_FILE=my_shapes.yaml bash run_pytest.sh -# 产物统一落在 runs/_<时间戳>/：run.log、shapes.yaml、autotune_records/、ttgir/。 +# 产物落在 runs/_<时间戳>/：run.log、shapes.yaml、autotune_records/、ttgir/ set -euo pipefail # ============================== 配置区 ============================== # 每项均可用同名环境变量覆盖，详见 README「环境变量一览」。 -# 被测算子：OP=测试函数名（去掉 test_ 前缀）；OP_FILE=benchmark 文件名 -# （对应 $FLAGGEMS_DIR/benchmark/test_.py::test_） +# 被测算子，对应 $FLAGGEMS_DIR/benchmark/test_.py::test_ OP="${OP:-fused_marlin_moe_mxfp4}" OP_FILE="${OP_FILE:-fused_marlin_moe}" -# FlagGems 仓库路径 FLAGGEMS_DIR="${FLAGGEMS_DIR:-/workspace/FlagGems-dev}" -# shape 来源：SHAPE_FILE 非空则用该 yaml；为空则用下方内置 yaml。 -# yaml 顶层 key 必须是 op 名。 +# shape 来源：非空则用该 yaml，否则用 INLINE_YAML。顶层 key 必须是 op 名。 SHAPE_FILE="${SHAPE_FILE:-}" read -r -d '' INLINE_YAML <<'YAML' || true fused_marlin_moe_mxfp4: @@ -44,17 +40,17 @@ YAML # 调优空间：0=普通 autotune（默认，快速验证）；1=FlagTune 扩展空间（首跑全量搜索、慢） USE_FLAGTUNE="${USE_FLAGTUNE:-0}" -# autotune record/replay（见 _autotune_record_plugin.py）： -# - 空：record 模式，本次选中的 config 记录到 $OUT_DIR/autotune_records/.json； -# - 指向某次历史 runs/_<时间戳> 目录：replay 该次记录（A/B 两侧同 config）。 +# 空=record 模式，把本次选中的 config 记入 autotune_records/.json； +# 指向某次历史 run 目录则 replay 其记录，用于 A/B 两侧锁同一套 config。 REPLAY_FROM="${REPLAY_FROM:-}" -# 终端颜色：always/never 强制开/关；为空则按 tty 自动判断（run.log 始终去色） +# always/never 强制开关终端颜色，空则按 tty 判断（run.log 始终去色） FLAGGEMS_PERF_COLOR="${FLAGGEMS_PERF_COLOR:-}" -# pytest 插件列表（可按需注释停用；各插件作用见 README「插件说明」） +# 各插件作用见 README「插件说明」；注释掉某行即停用该插件 PLUGINS=( -p _device_guard_plugin + -p _parallel_warmup_plugin -p _seed_plugin -p _shape_inject_plugin -p _shape_iter_inject_plugin @@ -71,12 +67,24 @@ PLUGINS=( SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" TEST_FILE="$FLAGGEMS_DIR/benchmark/test_${OP_FILE}.py::test_${OP}" -OUT_DIR="$SCRIPT_DIR/runs/${OP}_$(date +%Y%m%d-%H%M%S)" +# 产物目录，可用 OUT_DIR 指定。时间戳只到秒，并发启动会撞名，故用不带 -p 的 +# mkdir 抢占（已存在即失败），撞上就退让到 -2、-3…… +if [[ -n "${OUT_DIR:-}" ]]; then + mkdir -p "$OUT_DIR" +else + mkdir -p "$SCRIPT_DIR/runs" + BASE="$SCRIPT_DIR/runs/${OP}_$(date +%Y%m%d-%H%M%S)" + OUT_DIR="$BASE" + n=1 + until mkdir "$OUT_DIR" 2>/dev/null; do + n=$((n + 1)) + OUT_DIR="${BASE}-${n}" + (( n > 99 )) && { echo "!!! 无法创建产物目录（$BASE 及后缀均被占用）" >&2; exit 1; } + done +fi LOG_FILE="$OUT_DIR/run.log" -mkdir -p "$OUT_DIR" -# 颜色决策：下方 tee 管道会让 python 侧看不到 tty，所以在这里判断， -# 并通过 FLAGGEMS_PERF_COLOR 下传给各插件（_term_style.py 消费）。 +# tee 管道会让 python 侧检测不到 tty，故在 shell 层判断后经环境变量下传 if [[ -z "$FLAGGEMS_PERF_COLOR" && -t 1 && -z "${NO_COLOR:-}" ]]; then FLAGGEMS_PERF_COLOR=always fi @@ -94,7 +102,7 @@ export PYTHONPATH="$SCRIPT_DIR${PYTHONPATH:+:$PYTHONPATH}" export FLAGGEMS_PERF_CURRENT_OP="$OP" export PYTHONUNBUFFERED=1 # 实时输出不缓冲 -# record/replay 二选一，通过环境变量激活对应模式 +# record/replay 互斥，各由自己的环境变量激活 if [[ -n "$REPLAY_FROM" ]]; then AUTOTUNE_ENV="FLAGGEMS_PERF_AUTOTUNE_REPLAY_DIR=$REPLAY_FROM/autotune_records" [[ -f "$REPLAY_FROM/autotune_records/$OP.json" ]] || \ @@ -106,7 +114,7 @@ else MODE_DESC=record fi -# 解析 shape 文件：SHAPE_FILE 为空时把内置 yaml 写到临时文件 +# 未指定 shape 文件时，把 INLINE_YAML 落到临时文件供 pytest 读取 if [[ -z "$SHAPE_FILE" ]]; then SHAPE_FILE="$(mktemp --suffix=.yaml)" printf '%s\n' "$INLINE_YAML" > "$SHAPE_FILE" @@ -119,9 +127,8 @@ status=0 echo "${C_BOLD}>>> op=$OP mode=$MODE_DESC USE_FLAGTUNE=$USE_FLAGTUNE${C_RESET}" echo "${C_DIM}>>> out=$OUT_DIR${C_RESET}" - # _ir_meta_plugin 在进程退出时把"实际被使用"的变体的 ttgir 按 shape 整理落盘 - # （挂在 atexit 上，CUDA crash 后已编译部分仍可拿到）。 - # 目录结构与命名图例见 /naming.md 和 index.tsv（后者也记录落选的 sweep 变体）。 + # 每次用独立的 Triton 缓存目录，跑完即删：保证编译过程可复现，且 ttgir 落盘 + # 只包含本次的变体（_ir_meta_plugin 在 atexit 里按 shape 整理）。 CACHE_DIR="$OUT_DIR/.triton_cache" rm -rf "$CACHE_DIR"; mkdir -p "$CACHE_DIR" status=0 @@ -143,7 +150,6 @@ status=0 exit "$status" } 2>&1 | tee "$LOG_FILE" || status=$? -# 终端保留颜色；落盘的 run.log 去掉 ANSI 转义，保证 grep/diff 面对纯文本 -# （Ctrl-C 中断时会跳过去色，仅影响观感）。 +# run.log 去掉 ANSI 转义以便 grep/diff（终端输出保留颜色；Ctrl-C 时会跳过这步） sed -i -E $'s/\x1b\\[[0-9;]*[A-Za-z]//g' "$LOG_FILE" exit "$status"