Compare commits

...

5 Commits

Author SHA1 Message Date
zhoulin da22885645 Add multi-op batch screening layer; fix cudagraph fallback regressions
run_batch.py drives run_pytest.sh per operator across GPUs (stable-hash
sharding, per-op subprocess isolation, process-group timeouts, retry with
deterministic-failure cutoff, two-level dtype fallback, .complete resume,
per-op REPLAY_FROM). batch_summary.py aggregates run.log tables into
summary.csv. ops/ holds the curated assets: dual-repo inventories rebuilt
via AST scan + pytest collect verification, shape sets migrated from the
old regression harness and merged with upstream core_shapes class-name
keys (upstream's set_shapes falls back op_name -> MRO class name ->
1-D DEFAULT_SHAPES, so replacing the shape file without class keys
crashes the BLAS family), and a dismiss list where all 76 entries carry
verified reasons. Validated end to end: 1036-op full screen with zero
failures.

Also fix two cudagraph plugin regressions: newer torch appends "enable
device-side assertions" to every CUDA error, so the loose fatal-error
marker disabled the documented do_bench fallback entirely; and an aborted
graph capture can leave the default CUDA RNG generator stuck in capturing
state, poisoning every later torch.randn - captures now run under a
throwaway RNG state. run_pytest.sh gains an optional DTYPES passthrough.
2026-08-12 19:04:09 +00:00
zhoulin 95895cea0e Adapt to upstream FlagGems repo split and record pre_hook configs for libtuner
Upstream moved from single /workspace/FlagGems-dev to /workspace/dev/FlagGems
(flag_gems) plus /workspace/dev/FlagGems-vllm (flaggems_vllm). Patch LibTuner
per actually-imported package at collection finish instead of hardcoding
flag_gems, and update the default FLAGGEMS_DIR.

Also stop dropping pre_hook configs for libtuner kernels: hopper mm's TMA
configs carry a pre_hook even with USE_FLAGTUNE=0, so record silently skipped
them and replay fell back. Upstream LibTuner.run now re-attaches the hook by
kwargs match on cache read (same path as its own ConfigCache round-trip), so
pre_hook-less injection is safe there; plain triton Autotuner keeps the drop.
2026-08-12 10:55:32 +00:00
zhoulin 1e1d612031 Add opt-in multi-GPU sweep parallelism; fix GPU pinning and run dir races
For shape-heavy ops the autotune sweep dominates wall time, not the
measurement. _parallel_warmup_plugin shards the sweep across N GPUs and
hands the merged configs to the normal single-GPU serial measurement, so
timing stays comparable. Activated by PARALLEL_WARMUP_GPUS=N (silent
no-op when unset, so it can stay in the plugin list); output layout,
REPLAY_FROM usage and the latency table are unchanged.

Measurement is deliberately not parallelized: N processes saturating one
box couple through the power/thermal budget, so per-card latency gets
dragged by its neighbours by an amount that does not reproduce. Config
keys whose winner differs across shards are counted and reported as a
WARNING -- that count is how much to trust the run.

Fixes found while auditing:

- Shards no longer overwrite CUDA_VISIBLE_DEVICES with a bare shard
  index, which a caller who had selected idle cards (e.g. 6,7) would see
  re-interpreted as absolute ids 0,1 -- silently benchmarking on the busy
  cards they were avoiding.
- Interrupting the sweep no longer leaves N subprocesses holding GPUs;
  children are terminated and reaped before the exception propagates
  (BaseException, since KeyboardInterrupt is not an Exception).
- run_pytest.sh claims its output dir with a bare mkdir and retreats to
  a -2/-3 suffix on collision. The second-resolution timestamp meant two
  concurrent runs shared one directory and overwrote each other.
- Replay eviction failures now emit a distinct compile_*_no_evict marker.
  LibTuner.cache is a sqlite-backed ConfigCache with no __delitem__, so
  the bad config could not be dropped and the retry re-read it, while the
  log still claimed a clean fallback to live autotune.

Also trims comment density in both shell scripts and reworks the README:
promotes the parallel and cudagraph sections out from under the A/B flow,
groups the env table by purpose, documents that two record-mode runs are
not comparable, and marks ab_fold_test.sh as a caliber example whose
switch upstream has already removed.
2026-07-29 12:14:12 +00:00
zhoulin 7e0e8648f1 Fix autotune record/replay silently skipping libtuner kernels
LibTuner.cache is a view onto flag_gems' persistent sqlite config DB
(~/.flaggems/config_cache/TunedConfig_*.db), which survives across runs.
Both record and replay assumed a cold, in-process cache:

- record captured the chosen config by diffing cache keys before/after
  run(). On a warm DB the key is already present, LibTuner.run takes the
  cached branch without writing it again, so the diff was always empty
  and nothing was recorded for any libtuner kernel. Only @triton.autotune
  kernels (in-process cache) made it into the json -- e.g. a
  fused_marlin_moe_mxfp4 run recorded moe_sum_kernel alone, missing both
  MXFP4 GEMMs.
- replay only injected when the key was absent from the cache, so a warm
  DB skipped injection entirely: the run reported "replaying N entries"
  while actually self-tuning.

Record now reads back this call's own self.cache[key] after run();
replay overwrites unconditionally. Verified on fused_marlin_moe_mxfp4:
recorded entries 1 -> 6 (both GEMMs present), replay injects all 6 with
zero AUTOTUNE_REPLAY_FALLBACK and reproduces latency.

README: note that "fresh tune per side" requires dropping the sqlite DB
(not merely omitting REPLAY_FROM), and how to verify record coverage.
2026-07-29 00:38:14 +00:00
zhoulin f98ec10fd0 Renew scripts 2026-07-24 10:17:33 +00:00
16 changed files with 15242 additions and 151 deletions
+125 -26
View File
@@ -1,11 +1,13 @@
# zl_bench — FlagGems 算子性能测试工具
# zl_bench — FlagGems 算子性能测试工具
针对编译器(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,首轮即稳态、多轮一致。
单算子精测流程之上另有批量筛查层(`run_batch.py`,见「批量多算子测试」),可对双仓库上千个算子做例行体检,产物结构与单算子完全同构。
## 快速开始
```bash
@@ -15,8 +17,13 @@ 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
@@ -26,20 +33,36 @@ 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 扩展调优空间(首跑全量搜索、慢) |
| `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 两侧须取同值 |
| `DTYPES` | 空(上游默认 dtype 扫描) | 空格分隔的 dtype 白名单(如 `"bfloat16 float16"`),逐个转为上游 `--dtypes`。算子不支持时上游报 `can't be supported by this op` |
| `PARALLEL_WARMUP_GPUS` | 空(串行) | 设为 `N`(≥2)时把 autotune sweep 分片到 N 卡并行,测量仍单卡串行;不设或 `<2` 则完全串行(见「加速:多卡并行 sweep」) |
**输出**
| 变量 | 默认 | 说明 |
|-----|------|------|
| `FLAGGEMS_PERF_COLOR` | 空(按 tty 自动判断) | `always`/`never` 强制开/关终端颜色;`run.log` 始终为去色纯文本 |
## 输出目录结构
@@ -79,11 +102,34 @@ bash run_pytest.sh # -> runs/<op>_<ts_A>/
REPLAY_FROM=$PWD/runs/<op>_<ts_A> bash run_pytest.sh
```
B 侧必须走 `REPLAY_FROM`。不能靠「改动前后各跑一次」来比——两次 record 会各自重新 sweep、可能选出不同 config,这个差异足以盖过被测改动本身(成因见「为什么需要 replay」)。
A/B 两侧的 `USE_FLAGTUNE` 必须取同值:该开关决定调优空间,两侧不一致时 replay 会大量 fallbacklibtuner 持久缓存也各自独立命中,"同 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 产物
`pre_hook` 的 config(如 hopper mm 的 TMA configs):libtuner kernel 正常 record/replay——注入的 config 会由上游按 kwargs 匹配自动接回 pre_hook(与其自身 ConfigCache 的 DB 回读同一条路径);普通 `@triton.autotune` kernel 无此恢复机制,这类 config 不落 recordreplay 时表现为下述 `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.shA/B 测试参考示例
@@ -94,13 +140,70 @@ bash ab_fold_test.sh
# 汇总 logruns/ab_fold_<时间戳>.log(已去色);快速对账:grep -E '^#####|FAILED' <log>
```
对比其它开关 / 编译器改动 / 算子,套用该脚本改循环变量与环境变量即可,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/<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`——用它选过卡时,分片只会落在你选的那几张上。
## 批量多算子测试(run_batch.py
单算子流程之上的批量筛查层:按清单逐算子调 `run_pytest.sh`,多卡分片、每卡内部串行,产出逐算子状态表和逐测量行总表。**定位是筛查口径**——跨卡并行测量存在功耗/散热耦合噪声(幅度可到百分之几),发现可疑算子后回单算子串行流程(前几节)确认,不要直接拿批量数字下精细结论。
```bash
# 全量筛查(双仓库全部可收集算子,排除带无条件 skip 标记的与 ops/dismiss.txt 里的)
python run_batch.py --all --gpus 0,1,2,3,4,5,6,7 --dtypes bfloat16 --op-timeout 1800
# 指定子集:默认 main 仓,vllm 仓加前缀;同名测试函数用 @文件名 消歧
python run_batch.py --ops softmax vllm:fused_marlin_moe nextafter_@nextafter_
# 断点续跑(PASS/SKIP 的复用,FAIL/TIMEOUT 的重跑)
python run_batch.py --all --batch-dir runs/batch_xxx --resume
# 锁 config 复测:逐算子 REPLAY_FROM <root>/<repo>/<op>,语义同单算子 REPLAY_FROM
python run_batch.py --ops-file my_ops.txt --replay-root runs/batch_xxx
```
**算子资产(`ops/`**
| 文件 | 内容 | 维护方式 |
|-----|------|---------|
| `inventory_main.csv` / `inventory_vllm.csv` | 双仓库全部 benchmark 测试函数(op、op_file、无条件 skip 标记、pytest collect 核实结果、op_name 提取) | 上游更新后重跑 `python ops/gen_inventory.py --verify-collect` |
| `shapes_single.yaml` | 每算子单 shape 的筛查集(537 个精选 op_name 键 + 上游 core_shapes 的类名键底座——上游 shape 回退链是 op_name → MRO 类名 → 基类默认,缺类名键会让 BLAS 族跌到 1 维默认值崩溃) | 手工增补;上游变动后原地刷新:`python ops/gen_inventory.py --migrate-shapes ops/shapes_single.yaml --shapes-out ops/shapes_single.yaml`(未匹配键落 `.unmatched.yaml` 供复核) |
| `shapes_multi.yaml` | 每算子多 shape 的深查集(结构同上) | 同上 |
| `dismiss.txt` | 批量排除清单(`[repo:]op` 每行一个),只放本仓复核确认的失败项并注明原因/日期 | 批量跑出 FAIL 并确认原因后手工添加 |
无 shape 条目的算子用上游默认 shape,照常可测。`--shape-file ops/shapes_multi.yaml` 切换深查集。
**执行语义**:每算子独立子进程(CUDA crash 只废单个算子);`--op-timeout` 对进程组 SIGTERM→SIGKILL,记 `TIMEOUT` 不重试、批次继续;pytest rc=1 与信号杀最多重试 2 次,连续同 rc 视为确定性失败提前止损;`--dtypes` 有两级降级——遇上游 `can't be supported by this op` 即去掉限制重跑(备注 `dtype_fallback`),其余失败且无成功行时也去掉限制最后救一次(备注 `dtype_rescue`,覆盖 torch baseline 对受限 dtype 编译失败的场景)。runtime skipif 与上游 `xfail` 都判 SKIP(上游声明的无信号状态);`.complete` 标记只给 PASS/SKIP`--resume` 据此复用。
**产物**`<batch>/<repo>/<op>/` 与单算子 run 目录结构完全一致(run.log、autotune_records/、ttgir/……),所以任何一个算子都可以事后单独 `REPLAY_FROM` 复测。批量层额外产出 `ops_status.csv`(每算子一行:状态/耗时/行数/备注,边跑边原子更新)与 `summary.csv`(每测量行一行:dtype、latency、speedup、Size Detail,来源是 run.log 的结果表,`batch_summary.py` 也可单独对旧批次重跑)。备注列聚合 `no_cudagraph_rows`/`replay_fallback` 计数,解读口径时先看这列。
## 计时口径: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 被拒)。
部分算子本身不支持 CUDA graph capture(测量函数内含 host 同步、动态显存分配、autograd backward、不合法的流操作等),这类算子会自动回退到普通 do_bench 计时,`run.log` 中打 `BENCHMARK_DIRECT_NO_CUDAGRAPH` 标记(含失败阶段与具体原因;`phase=warmup` 表示预热阶段就失败了,并非 graph capture 被拒)。
回退路径的两个防御措施(都吃过亏):(1)致命错误判定按错误原文匹配而非宽泛子串——新版 torch 给每个 CUDA 错误都追加 "enable device-side assertions" 提示语,宽松匹配会把所有 capture 失败误判为致命错误、令回退路径整体失效;(2)capture 在一次性 RNG state 替身下执行——`torch.cuda.graph` 会无条件注册默认 CUDA RNG 生成器,capture 中途失败可能让生成器卡在 capturing 态,此后进程内任何 `torch.randn` 都抛 "Offset increment outside graph capture"(毒化显形时机不定,事后修复不可靠),替身隔离让真身永不参与 capture,成败都原样归位。
回退本身不影响 A/B 公平性(两侧同一算子回退行为一致),但**回退口径的测量误差更大**:do_bench 每次迭代都走完整的 Python → launch 路径,kernel 越小,launch 开销和 CPU 侧抖动在数字里占比越高——亚毫秒级 kernel 上两种口径可差 2 倍以上,且行间波动更明显。解读这类算子的结果时:
@@ -109,24 +212,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 型号变化。如果需要"各自最优"口径(让 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 |
| `_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 等) | 不关 |
@@ -143,13 +236,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 算子首跑的全量 sweepmm 约 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 少了文件?**
@@ -157,6 +256,6 @@ libtuner 的持久缓存何时失效:FlagGems kernel 源码改动、tune_confi
## 依赖假设
- FlagGems benchmark 体系(`benchmark/base.py``Benchmark` 类、conftest 的 `--shape_file/--level/--mode` 选项);
- 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 输出的注册日志是否还正常打印。
+83 -53
View File
@@ -18,6 +18,7 @@ FLAGGEMS_PERF_CURRENT_OP):
from __future__ import annotations
import atexit
import importlib
import json
import os
import sys
@@ -30,12 +31,17 @@ _RECORD_DIR_ENV = "FLAGGEMS_PERF_AUTOTUNE_RECORD_DIR"
_REPLAY_DIR_ENV = "FLAGGEMS_PERF_AUTOTUNE_REPLAY_DIR"
_OP_ENV = "FLAGGEMS_PERF_CURRENT_OP"
# Upstream split into two repos with different package names; either may host
# LibTuner-decorated kernels depending on which repo's benchmark is running.
_LIBTUNER_PKGS = ("flag_gems", "flaggems_vllm")
# Module-global aggregates, accumulated as run() is called and dumped at exit.
# Lock guards merges in case a kernel uses threads internally.
_record_map: Dict[str, Dict[str, Dict[str, Any]]] = {}
_record_lock = threading.Lock()
_replay_map: Dict[str, Dict[str, Dict[str, Any]]] = {}
_dump_done = False
_wrap = None # set in pytest_configure when record/replay is active
def _op_name() -> str:
@@ -65,10 +71,16 @@ def _serialize_key(key: Tuple[Any, ...]) -> str:
return json.dumps([repr(x) for x in key])
def _serialize_config(cfg: Any) -> Optional[Dict[str, Any]]:
# Drop configs with a pre_hook (un-serializable callable) rather than replay
# without it and compute incorrectly.
if getattr(cfg, "pre_hook", None) is not None:
def _serialize_config(cfg: Any, allow_pre_hook: bool = False) -> Optional[Dict[str, Any]]:
# pre_hook is an un-serializable callable. For plain triton Autotuner drop
# such configs rather than replay without the hook and compute incorrectly.
# For LibTuner callers pass allow_pre_hook=True: upstream's own ConfigCache
# round-trips configs pre_hook-less and LibTuner.run re-attaches the hook by
# matching all_kwargs against self.configs, and replay-injected configs go
# through that same path (hopper mm's TMA configs carry a pre_hook even
# with USE_FLAGTUNE=0 — dropping them would record nothing for those
# kernels and silently defeat replay).
if getattr(cfg, "pre_hook", None) is not None and not allow_pre_hook:
return None
try:
return {
@@ -122,46 +134,32 @@ def _emit_marker(reason: str) -> None:
def _record_run(original):
# Snapshot self.cache before/after run() to capture the chosen config (works
# for both Autotuner and LibTuner). Single-config kernels skip the cache
# write, so record configs[0] explicitly for uniform replay.
# Read back this call's own key after run(), rather than diffing cache keys.
# LibTuner.cache is backed by a persistent sqlite DB (flag_gems libcache), so
# on a warm DB the key is already present and a before/after diff comes up
# empty -- which silently recorded nothing for every libtuner kernel.
# Single-config kernels bypass the cache write, so fall back to configs[0].
def runner(self, *args, **kwargs):
try:
keys_before = set(self.cache.keys()) if hasattr(self.cache, "keys") else set()
except Exception:
keys_before = set()
result = original(self, *args, **kwargs)
try:
keys_after = set(self.cache.keys()) if hasattr(self.cache, "keys") else set()
except Exception:
keys_after = set()
new_keys = keys_after - keys_before
if not new_keys:
# No new entry: record configs[0] for single-config kernels (cache
# write bypassed); otherwise nothing to record (disk-cache hit).
if len(getattr(self, "configs", []) or []) == 1:
key = _compute_key(self, args, kwargs)
if key is not None and key not in self.cache:
if key is None:
return result
cfg = None
try:
cfg = self.cache[key]
except Exception:
cfg = None
if cfg is None and len(getattr(self, "configs", []) or []) == 1:
cfg = self.configs[0]
entry = _serialize_config(cfg)
if cfg is None:
return result
# get_key marks LibTuner (both flag_gems and flaggems_vllm); plain
# triton Autotuner has no pre_hook re-attach on cache read, LibTuner does.
entry = _serialize_config(cfg, allow_pre_hook=hasattr(self, "get_key"))
if entry is not None:
kid = _kernel_id(self)
with _record_lock:
bucket = _record_map.setdefault(kid, {})
bucket[_serialize_key(key)] = entry
return result
kid = _kernel_id(self)
with _record_lock:
bucket = _record_map.setdefault(kid, {})
for k in new_keys:
try:
cfg = self.cache[k]
except Exception:
continue
entry = _serialize_config(cfg)
if entry is None:
continue
bucket[_serialize_key(k)] = entry
_record_map.setdefault(kid, {})[_serialize_key(key)] = entry
return result
return runner
@@ -176,7 +174,11 @@ def _replay_run(original):
bucket = _replay_map.get(kid)
if bucket:
key = _compute_key(self, args, kwargs)
if key is not None and key not in self.cache:
# Overwrite unconditionally: LibTuner.cache is backed by a
# persistent sqlite DB, so gating on `key not in self.cache`
# would skip injection whenever that DB is warm -- leaving the
# run silently self-tuned instead of replaying.
if key is not None:
rec = bucket.get(_serialize_key(key))
if rec is None:
_emit_marker("key_missing")
@@ -192,12 +194,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
@@ -243,6 +254,7 @@ def _dump_record(record_dir: str) -> None:
def pytest_configure(config):
global _wrap
from _term_style import tag
record_dir = os.environ.get(_RECORD_DIR_ENV, "").strip()
replay_dir = os.environ.get(_REPLAY_DIR_ENV, "").strip()
@@ -255,26 +267,44 @@ def pytest_configure(config):
return
import triton.runtime.autotuner as _autotuner
wrap = _record_run if record_dir else _replay_run
_autotuner.Autotuner.run = wrap(_autotuner.Autotuner.run)
patched = ["Autotuner"]
try:
from flag_gems.utils.libentry import LibTuner
if "run" in LibTuner.__dict__:
LibTuner.run = wrap(LibTuner.__dict__["run"])
patched.append("LibTuner")
except Exception:
pass
_wrap = _record_run if record_dir else _replay_run
_autotuner.Autotuner.run = _wrap(_autotuner.Autotuner.run)
if record_dir:
# atexit (not sessionfinish): persist whatever was recorded even if an op
# crash kills the session; a later replay run falls back for missing keys.
atexit.register(_dump_record, record_dir)
print(f"{tag('[autotune-record-plugin]')} recording autotune configs to "
f"{_record_path(record_dir)} ({'+'.join(patched)})",
f"{_record_path(record_dir)} (Autotuner)",
file=sys.stderr, flush=True)
else:
loaded = _load_replay_dir(replay_dir)
print(f"{tag('[autotune-record-plugin]')} replaying {loaded} recorded entries from "
f"{_record_path(replay_dir)} ({'+'.join(patched)})",
f"{_record_path(replay_dir)} (Autotuner)",
file=sys.stderr, flush=True)
def pytest_collection_finish(session):
# LibTuner is patched here rather than in pytest_configure: only after
# collection (which imports the benchmark module) do we know which FlagGems
# package is actually in use, and importing the unused one just to patch it
# would initialize a second runtime in this process for nothing.
if _wrap is None:
return
from _term_style import tag
patched = []
for pkg in _LIBTUNER_PKGS:
if pkg not in sys.modules:
continue
try:
libentry = importlib.import_module(f"{pkg}.utils.libentry")
tuner_cls = libentry.LibTuner
if "run" in tuner_cls.__dict__:
tuner_cls.run = _wrap(tuner_cls.__dict__["run"])
patched.append(f"LibTuner[{pkg}]")
except Exception as exc:
print(f"{tag('[autotune-record-plugin]')} warning: LibTuner patch "
f"failed for {pkg}: {exc}", file=sys.stderr, flush=True)
print(f"{tag('[autotune-record-plugin]')} libtuner coverage: "
f"{', '.join(patched) or 'none (no FlagGems package imported?)'}",
file=sys.stderr, flush=True)
+40 -5
View File
@@ -8,7 +8,9 @@ fall back to plain do_bench and print a BENCHMARK_DIRECT_NO_CUDAGRAPH marker
into the log. Before capturing, fn is explicitly warmed up for the caller's
warmup time budget (do_bench_cudagraph's own 5-iter warmup is too short to
settle autotune/JIT for the MoE kernels, which makes the first captured
measurement unstable run-to-run).
measurement unstable run-to-run). The capture itself runs under a scratch CUDA
RNG state so that an aborted capture cannot poison the process-wide generator
(rationale: _capture_with_scratch_rng).
"""
from __future__ import annotations
@@ -25,13 +27,19 @@ _ORIGINAL_DO_BENCH = _tt.do_bench
def _is_fatal_cuda_error(exc: Exception) -> bool:
"""True for CUDA context errors (illegal address, OOM, device-side assert):
the context is poisoned, so re-raise rather than fall back and record a bogus
latency as success."""
latency as success.
Markers must match the error text itself, not torch's generic footer: newer
torch appends "Compile with `TORCH_USE_CUDA_DSA` to enable device-side
assertions." to EVERY AcceleratorError, so a loose "device-side assert"
marker turned all capture-incompatible ops into fatal re-raises and killed
the documented do_bench fallback path entirely."""
msg = f"{type(exc).__name__}: {exc}".lower()
markers = (
"illegal memory access",
"cudaerrorillegaladdress",
"out of memory",
"device-side assert",
"device-side assert triggered",
"an illegal instruction",
"misaligned address",
"uncorrectable ecc",
@@ -83,6 +91,33 @@ def _warmup_before_capture(fn, warmup_ms, grad_to_none):
torch.cuda.synchronize()
def _capture_with_scratch_rng(call):
"""Run a graph capture under a throwaway CUDA RNG state.
torch.cuda.graph's capture_begin unconditionally registers the default CUDA
RNG generator and sets its capturing flag; when a capture aborts midway the
epilogue may never run, the flag sticks, and every later CUDA RNG call in
the process (the next shape's torch.randn included) raises "Offset
increment outside graph capture". Repairing after the fact is unreliable —
when the poisoning surfaces depends on which stage the capture died in —
so prevent it instead: swap in a scratch state via graphsafe_set_state
before capturing and restore the original state object afterwards, success
or failure. The real generator never takes part in a capture, so poisoning
can only land on the discarded scratch state, and _seed_plugin determinism
is left untouched."""
import torch
dev = torch.cuda.current_device()
gen = torch.cuda.default_generators[dev]
saved = gen.graphsafe_get_state()
scratch = torch.Generator(device="cuda")
scratch.manual_seed(gen.initial_seed())
gen.graphsafe_set_state(scratch)
try:
return call()
finally:
gen.graphsafe_set_state(saved)
def _patched_do_bench(fn, warmup=25, rep=100, grad_to_none=None,
quantiles=None, return_mode="mean"):
"""Drop-in replacement for triton.testing.do_bench using cudagraph capture."""
@@ -100,13 +135,13 @@ def _patched_do_bench(fn, warmup=25, rep=100, grad_to_none=None,
# Settle autotune/JIT before capture (rationale: _warmup_before_capture).
_warmup_before_capture(fn, warmup, grad_to_none)
phase = "capture"
return _tt.do_bench_cudagraph(
return _capture_with_scratch_rng(lambda: _tt.do_bench_cudagraph(
fn,
rep=rep,
grad_to_none=grad_to_none,
quantiles=quantiles,
return_mode=return_mode,
)
))
except Exception as exc:
# Re-raise genuine device failures; anything else falls back to plain
# do_bench, with the failing phase (sync/warmup/capture) in the marker.
+1 -1
View File
@@ -4,7 +4,7 @@ FlagGems' `runtime.backend.device_finder` falls back to
`subprocess.run(nvidia-smi)` **without a timeout**. Under this conda env's
forked/inconsistent subprocess (`_posixsubprocess` symbol mismatch), that
child can hang indefinitely, leaving `import flag_gems` stuck in `wait4`
(seen as run_pytest "卡住 with no result", GPU 0%).
(seen as run_pytest hanging with no output, GPU at 0%).
This plugin runs at import time — before any test module imports flag_gems —
detects the vendor via torch (no subprocess), and sets GEMS_VENDOR so
+306
View File
@@ -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=<N>; 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/<op>.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)
+12 -10
View File
@@ -1,17 +1,19 @@
#!/bin/bash
# A/B 测试参考示例:FLAGGEMS_MXFP4_FOLDSCALEfold_scale)开关验证
# 4 个 DeepSeek-V4-Flash trace yaml × FOLD=0/1 共 8 轮)。
# 口径:USE_FLAGTUNE=0FOLD=0(基线)先跑并 record autotune configFOLD=1 用
# REPLAY_FROM 回放同一套 config(两侧同配置比 wall-time
# 对比其它开关/算子时,改循环变量与环境变量即可复用同一口径。
# 全部输出汇总到 runs/ab_fold_<时间戳>.log;各轮 runs/<op>_<时间戳>/ 产物照常存档。
# 用法: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/<op>_<时间戳>/
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"
+186
View File
@@ -0,0 +1,186 @@
#!/usr/bin/env python3
"""聚合一次批量跑(run_batch.py 产物目录)的结果为 summary.csv。
数据源是各算子 run 目录里的 run.logrun_pytest.sh 已去 ANSI 色):
- 逐表解析 `Operator: <op_name> (dtype=..., mode=..., level=...)` 标题
(兼容 _pretty_report_plugin 与上游原生两种格式);
- 逐行解析 SUCCESS/FAILED 表行:状态 + 前三个数值列固定为
torch_lat / gems_lat / speeduptflops/gbps 等追加列收进 extra_metrics),
余下为 size_detail
- 汇入 run_batch.py 写的 ops_status.csv(每算子总状态 / 备注)。
也可单独使用:python batch_summary.py runs/batch_xxx [-o summary.csv]
"""
from __future__ import annotations
import argparse
import csv
import os
import re
import sys
from pathlib import Path
from typing import Dict, List, Optional, Tuple
# 兼容两种标题:
# pretty: Operator: softmax (dtype=torch.float16, mode=kernel, level=core)
# upstream: Operator: softmax Performance Test (dtype=torch.float16, mode=kernel,level=core)
OPERATOR_HEADER_RE = re.compile(
r"^Operator:\s+(?P<name>\S+)\s+(?:Performance Test\s*)?"
r"\(dtype=(?P<dtype>[^,]+),\s*mode=(?P<mode>[^,]+),\s*level=(?P<level>[^)]+)\)"
)
ROW_RE = re.compile(r"^(?P<status>SUCCESS|FAILED)\s+(?P<rest>\S.*)$")
NUM_TOKEN_RE = re.compile(
r"^(?:N/A|nan|-?inf"
r"|-?(?:\d+(?:\.\d*)?|\.\d+)(?:[eE][+-]?\d+)?)$"
)
PYTEST_SKIPPED_RE = re.compile(r"(?P<n>\d+)\s+(?:skipped|xfailed)\b")
DIRECT_NO_CUDAGRAPH_RE = re.compile(r"BENCHMARK_DIRECT_NO_CUDAGRAPH\b")
REPLAY_FALLBACK_RE = re.compile(r"AUTOTUNE_REPLAY_FALLBACK\b")
SUMMARY_FIELDS = [
"repo", "op", "fg_op_name", "dtype", "case_id", "row_status",
"torch_lat", "gems_lat", "speedup", "extra_metrics", "size_detail",
"op_status", "notes",
]
def parse_run_log(text: str) -> List[Dict[str, str]]:
"""把 run.log 里的全部结果表解析为行列表(跨多个 dtype 表)。"""
rows: List[Dict[str, str]] = []
cur_op, cur_dtype = "", ""
for line in text.splitlines():
stripped = line.strip()
header = OPERATOR_HEADER_RE.match(stripped)
if header:
cur_op = header.group("name")
cur_dtype = header.group("dtype").strip()
continue
m = ROW_RE.match(stripped)
if not m:
continue
tokens = m.group("rest").split()
metrics: List[str] = []
while tokens and len(metrics) < 7 and NUM_TOKEN_RE.match(tokens[0]):
metrics.append(tokens.pop(0))
# 前三个数值列固定:torch_lat, gems_lat, speedup(上游列序承诺不变)
while len(metrics) < 3:
metrics.append("")
rows.append({
"fg_op_name": cur_op,
"dtype": cur_dtype,
"row_status": m.group("status"),
"torch_lat": metrics[0],
"gems_lat": metrics[1],
"speedup": metrics[2],
"extra_metrics": " ".join(metrics[3:]),
"size_detail": " ".join(tokens),
})
return rows
def pytest_skipped_count(text: str) -> int:
"""runtime skipif 与上游 xfail 标记都算"无测量信号",供 SKIP 判定。"""
count = 0
for m in PYTEST_SKIPPED_RE.finditer(text[-4000:]):
count += int(m.group("n"))
return count
def marker_notes(text: str) -> str:
"""cudagraph 回退 / replay 兜底标记计数,供备注列(口径解读用)。"""
notes = []
n = len(DIRECT_NO_CUDAGRAPH_RE.findall(text))
if n:
notes.append(f"no_cudagraph_rows={n}")
n = len(REPLAY_FALLBACK_RE.findall(text))
if n:
notes.append(f"replay_fallback={n}")
return ";".join(notes)
def classify(rc: int, rows: List[Dict[str, str]], log_text: str) -> str:
"""算子级状态:FAIL / SKIP / PASSTIMEOUT 由 run_batch 在外层判)。"""
if rc != 0:
return "FAIL"
if not rows:
return "SKIP" if pytest_skipped_count(log_text) else "FAIL"
if any(r["row_status"] != "SUCCESS" for r in rows):
return "FAIL"
return "PASS"
def _load_ops_status(batch_dir: Path) -> Dict[str, Dict[str, str]]:
path = batch_dir / "ops_status.csv"
if not path.is_file():
return {}
with path.open(newline="") as f:
return {f"{r['repo']}/{r['run_name']}": r for r in csv.DictReader(f)}
def summarize(batch_dir: Path, out_path: Optional[Path] = None) -> Path:
batch_dir = batch_dir.resolve()
out_path = out_path or (batch_dir / "summary.csv")
status_by_dir = _load_ops_status(batch_dir)
all_rows: List[Dict[str, str]] = []
for repo_dir in sorted(p for p in batch_dir.iterdir() if p.is_dir()):
repo = repo_dir.name
for op_dir in sorted(p for p in repo_dir.iterdir() if p.is_dir()):
log = op_dir / "run.log"
if not log.is_file():
continue
text = log.read_text(errors="replace")
status_row = status_by_dir.get(f"{repo}/{op_dir.name}", {})
op = status_row.get("op") or op_dir.name
table_rows = parse_run_log(text)
if not table_rows:
all_rows.append({
"repo": repo, "op": op, "fg_op_name": "", "dtype": "",
"case_id": "", "row_status": "", "torch_lat": "",
"gems_lat": "", "speedup": "", "extra_metrics": "",
"size_detail": "",
"op_status": status_row.get("status", ""),
"notes": status_row.get("notes", ""),
})
continue
case_counter: Dict[Tuple[str, str], int] = {}
for row in table_rows:
key = (row["fg_op_name"], row["dtype"])
idx = case_counter.get(key, 0)
case_counter[key] = idx + 1
all_rows.append({
"repo": repo, "op": op, "case_id": f"case_{idx:03d}",
"op_status": status_row.get("status", ""),
"notes": status_row.get("notes", ""),
**row,
})
tmp = out_path.with_name(out_path.name + ".tmp")
with tmp.open("w", newline="") as f:
writer = csv.DictWriter(f, fieldnames=SUMMARY_FIELDS)
writer.writeheader()
writer.writerows(all_rows)
os.replace(tmp, out_path)
return out_path
def main() -> None:
p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("batch_dir", type=Path, help="run_batch.py 的产物目录")
p.add_argument("-o", "--output", type=Path, default=None,
help="输出 CSV(默认 <batch_dir>/summary.csv")
args = p.parse_args()
if not args.batch_dir.is_dir():
raise SystemExit(f"目录不存在: {args.batch_dir}")
out = summarize(args.batch_dir, args.output)
with out.open(newline="") as f:
rows = list(csv.DictReader(f))
n_ok = sum(1 for r in rows if r["row_status"] == "SUCCESS")
n_bad = sum(1 for r in rows if r["row_status"] == "FAILED")
print(f"[summary] {out}: {len(rows)} 行(SUCCESS={n_ok} FAILED={n_bad}",
file=sys.stderr)
if __name__ == "__main__":
main()
+120
View File
@@ -0,0 +1,120 @@
# 批量测试排除清单。每行一个算子:`op`(默认 main 仓)或 `repo:op`。
# run_batch.py --all 会跳过这里列出的算子;带无条件 pytest.mark.skip 的算子由
# inventory 的 marked_skip 列自动剔除,不必写在这里。
# 只放"经本仓批量跑复核确认失败"的条目,注明原因与复核日期。
#
# 2026-08-12 复核记录:
# - 旧回归清单 24 个候选全部重验:14 个新上游已修复(移出),3 个系本仓
# _cudagraph_plugin 回退失效 bug(已修复,移出),7 个确认为上游/环境问题;
# - 全量筛查(1103 算子)新增确认 67 个:43 个 torch jiterator 环境问题 +
# 24 个上游 bug/环境限制。shape 语义类失败(28 个)已通过 shape yaml 合并
# 上游 core_shapes 类名键修复,不在此列;
# - test_FLA/ 补齐后新增确认 2 个(硬件限制 + 专用 harness),共 76 条。
# ==== 上游 kernel 编译错误 [2026-08-12]
# CompilationError(关 cudagraph 后仍复现)
hstack
scaled_mm_benchmark
scaled_mm_out_benchmark
# tt.load result/ptr type 校验失败
flash_mla_with_kvcache
# PassManager::run failedFlagTree triton 编译失败)
flash_mla_sparse_fwd
vllm:flash_mla_sparse_fwd
# capture 失败回退后仍报 kernel 错误
vllm:flash_mla
vllm:flash_mla_with_kvcache
linalg_lu_factor_ex
# ==== 上游 benchmark 代码 bug [2026-08-12]
# baseline torch.trace 收到非矩阵输入(expected a matrix
trace
# set_shapes() 签名与基类调用不符(TypeError)
get_paged_mqa_logits_metadata
# 输入解包 ValueErrorexpected 3, got 1
grouped_mm
# 上游 benchmark assert 失败
get_scheduler_metadata
# is_backward 对未参与计算图的张量求梯度(allow_unused
embedding_backward
margin_ranking_loss_backward
# baseline 函数签名不符(unexpected keyword 'q' / 缺 constexpr 实参)
fp8_fp4_paged_mqa_logits
vllm:fp8_fp4_mqa_logits
# baseline 返回 NoneNoneType has no attribute 'device'
special_gammainc
# normalized_shape=[](上游默认 shape 语义错误)
native_layer_norm
# 上游断言 K>=16 与其默认 attention shape 冲突
perf_scaled_dot_product_flash_attention_backward
perf_scaled_dot_product_cudnn_attention_backward
# Triton Error [CUDA]: invalid argument(上游默认 shape 下 kernel 启动参数非法)
index_select_backward
pairwise_distance
# test_blas_perf_parallel 专用 harnessshape kind 语义与通用注入冲突;常规 mul 已覆盖)
perf_mul
# 需要 pytest-benchmark fixture(非 FlagGems Benchmark 体系)
vllm:triton_unified_attention_perf
# 自定义输出格式(非 FlagGems Benchmark 表)且预编译超 40min
vllm:perf_chunk_gla
# ==== 硬件限制 [2026-08-12]
# kernel 共享内存需求 245-335KB 超 H20 上限 227KBautotune 全候选 OutOfResources
vllm:chunk_gdn2
# ==== 环境限制(换环境后应重验)[2026-08-12]
# vllm 0.20.2 无 vllm.utils.deep_gemm.fp8_mqa_logits
fp8_mqa_logits
# FA2 不支持 num_splits > 1flash-attn 版本)
flash_attn_varlen_opt_func
# magma 显存分配失败(cannot allocate memory on GPU, info=-113
cholesky_solve
cholesky_solve_out
# cusolver INTERNAL_ERRORXgeev
linalg_eigvals
# ---- torch jiterator/NVRTC 在本环境编译失败(fp32 也复现),
# ---- sinc/bessel/special 族的 torch baseline 整体不可用:
erfc
erfc_
erfinv
erfinv_
lcm
lcm_
lgamma
lgamma_
mvlgamma
mvlgamma_
polygamma
polygamma_inplace
polygamma_out
sinc
sinc_
special_airy_ai
special_airy_ai_out
special_bessel_j0
special_bessel_j1
special_chebyshev_polynomial_u
special_chebyshev_polynomial_v
special_chebyshev_polynomial_w
special_chebyshev_polynomial_w_out
special_erfc
special_erfcx
special_erfinv
special_erfinv_out
special_gammaln
special_gammaln_out
special_hermite_polynomial_h
special_i1e
special_legendre_polynomial_p
special_log1p_non_tensor
special_modified_bessel_k0
special_modified_bessel_k0_out
special_modified_bessel_k1
special_modified_bessel_k1_out
special_scaled_modified_bessel_k1
special_scaled_modified_bessel_k1_out
special_shifted_chebyshev_polynomial_u
special_shifted_chebyshev_polynomial_u_
special_shifted_chebyshev_polynomial_v
special_shifted_chebyshev_polynomial_w
special_sinc
+294
View File
@@ -0,0 +1,294 @@
#!/usr/bin/env python3
"""重建批量测试的算子清单(inventory),并可校验迁移 shape yaml。
对每个 FlagGems 仓库静态扫描 benchmark/ 下的 test_*.py(含 test_ 前缀的
子目录,如 test_FLA/——op_file 记相对路径去前后缀,run_pytest.sh 的
`benchmark/test_${OP_FILE}.py` 拼接对其天然成立;models_benchmark/ 等
非 test_ 前缀路径不在性能批量范围):
- 每个模块级 test_* 函数出一行 inventoryop(函数名去 test_ 前缀,即
run_pytest.sh 的 OP)、op_file(文件名去 test_ 前缀/.py 后缀,即 OP_FILE);
- 标注 marked_skip(函数带无条件 pytest.mark.skip,跑了必 SKIP,批量驱动默认剔除);
- 提取 op_name 字符串常量(op_name= 关键字实参,以及 *Benchmark 类
__init__ 里 super().__init__ 的首个字符串位置实参),用于 shape yaml 键校验;
- --verify-collect 时额外跑一次 pytest --collect-only 核实函数确实可被收集
AST 见到 ≠ pytest 收得到,import 失败/条件定义都会导致差异)。
产物:
ops/inventory_<label>.csv # repo,op,op_file,file,marked_skip,collected,op_names
--migrate-shapes 时按 op_name 并集过滤输入 yaml、合并上游 core_shapes 底座:
<out>.yaml + <out>.unmatched.yaml(未匹配任何 op_name 的键,供人工复核后删除)
用法示例(默认扫描两个新仓库):
python ops/gen_inventory.py --verify-collect
# 上游 op_name/core_shapes 变动后,对现有 shape 集重新校验(输入=输出即原地刷新)
python ops/gen_inventory.py --migrate-shapes ops/shapes_single.yaml \
--shapes-out ops/shapes_single.yaml
"""
from __future__ import annotations
import argparse
import ast
import csv
import os
import subprocess
import sys
from pathlib import Path
from typing import Dict, List, Optional, Set, Tuple
import yaml
OPS_DIR = Path(__file__).resolve().parent
DEFAULT_REPOS = [
("main", Path("/workspace/dev/FlagGems")),
("vllm", Path("/workspace/dev/FlagGems-vllm")),
]
INVENTORY_FIELDS = ["repo", "op", "op_file", "file", "marked_skip", "collected", "op_names"]
def parse_args() -> argparse.Namespace:
p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--repo", action="append", default=None, metavar="LABEL=PATH",
help="要扫描的仓库(可多次),默认 main=/workspace/dev/FlagGems "
"与 vllm=/workspace/dev/FlagGems-vllm")
p.add_argument("--verify-collect", action="store_true",
help="用 pytest --collect-only 核实每个测试函数可被收集(慢,"
"需要完整运行环境;结果写进 collected 列)")
p.add_argument("--out-dir", type=Path, default=OPS_DIR,
help="inventory CSV 输出目录(默认 ops/")
p.add_argument("--migrate-shapes", type=Path, default=None,
help="旧 shape yaml 路径;按扫描出的 op_name 并集过滤后迁移")
p.add_argument("--shapes-out", type=Path, default=None,
help="迁移后 yaml 输出路径(--migrate-shapes 时必填)")
return p.parse_args()
def _resolve_repos(raw: Optional[List[str]]) -> List[Tuple[str, Path]]:
if not raw:
return [(label, path) for label, path in DEFAULT_REPOS if path.is_dir()]
repos = []
for item in raw:
label, _, path = item.partition("=")
if not path:
raise SystemExit(f"--repo 需要 LABEL=PATH 形式,收到: {item!r}")
repos.append((label, Path(path)))
return repos
def _is_unconditional_skip(dec: ast.AST) -> bool:
"""pytest.mark.skip(非 skipif)——带不带 reason 实参都算无条件。"""
target = dec.func if isinstance(dec, ast.Call) else dec
if not (isinstance(target, ast.Attribute) and target.attr == "skip"):
return False
mark = target.value
return (isinstance(mark, ast.Attribute) and mark.attr == "mark"
and isinstance(mark.value, ast.Name) and mark.value.id == "pytest")
def _op_names_in(node: ast.AST) -> Set[str]:
"""节点范围内的 op_name 字符串常量,覆盖三种上游写法:
1. 任意调用的 op_name= 关键字实参;
2. super().__init__ 的首个字符串位置实参(FusedDeepseekV4... 模式);
3. *Benchmark 类实例化的首个字符串位置实参(ScaledMMBenchmark("scaled_mm",
...) 这类经形参转发、常量提取够不到 super() 调用的模式)。
误报只会让 shape yaml 多留一个无人读取的键,无害。"""
names: Set[str] = set()
for sub in ast.walk(node):
if not isinstance(sub, ast.Call):
continue
for kw in sub.keywords:
if kw.arg == "op_name" and isinstance(kw.value, ast.Constant) \
and isinstance(kw.value.value, str):
names.add(kw.value.value)
first_str = (sub.args[0].value
if sub.args and isinstance(sub.args[0], ast.Constant)
and isinstance(sub.args[0].value, str) else None)
if first_str is None:
continue
func = sub.func
if (isinstance(func, ast.Attribute) and func.attr == "__init__"
and isinstance(func.value, ast.Call)
and isinstance(func.value.func, ast.Name)
and func.value.func.id == "super"):
names.add(first_str)
callee = func.id if isinstance(func, ast.Name) else (
func.attr if isinstance(func, ast.Attribute) else "")
if callee.endswith("Benchmark"):
names.add(first_str)
return names
def scan_repo(label: str, root: Path) -> Tuple[List[Dict[str, str]], Set[str]]:
"""返回 (inventory 行, 该仓库全部 op_name 集合)。"""
bench = root / "benchmark"
rows: List[Dict[str, str]] = []
op_names_all: Set[str] = set()
for path in sorted(bench.rglob("test_*.py")):
rel = path.relative_to(bench)
# 仅收 test_ 前缀链路(根级 test_x.py 与 test_XXX/ 子目录),保证
# "benchmark/test_" + op_file + ".py" 能原样重建路径。
if not str(rel).startswith("test_"):
continue
try:
tree = ast.parse(path.read_text())
except SyntaxError as exc:
print(f"[gen] warning: {path} 解析失败,跳过: {exc}", file=sys.stderr)
continue
op_names_all |= _op_names_in(tree)
op_file = str(rel)[len("test_"):-len(".py")]
for node in tree.body:
if not isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
continue
if not node.name.startswith("test_"):
continue
marked = any(_is_unconditional_skip(d) for d in node.decorator_list)
func_ops = sorted(_op_names_in(node))
rows.append({
"repo": label,
"op": node.name[len("test_"):],
"op_file": op_file,
"file": f"benchmark/{rel}",
"marked_skip": "yes" if marked else "",
"collected": "",
"op_names": ";".join(func_ops),
})
return rows, op_names_all
def verify_collect(label: str, root: Path, rows: List[Dict[str, str]]) -> None:
"""跑 pytest --collect-only 核实测试函数可被收集,写 collected 列。"""
env = os.environ.copy()
env.setdefault("GEMS_VENDOR", "nvidia") # 跳过 nvidia-smi 子进程探测(可能挂死)
proc = subprocess.run(
[sys.executable, "-m", "pytest", "--collect-only", "-q",
"--continue-on-collection-errors", "benchmark/"],
cwd=str(root), text=True, capture_output=True, env=env, timeout=1200,
)
collected_funcs: Set[Tuple[str, str]] = set()
for line in proc.stdout.splitlines():
line = line.strip()
if "::" not in line or not line.startswith("benchmark/"):
continue
file_part, _, rest = line.partition("::")
func = rest.split("::")[0].split("[")[0]
collected_funcs.add((file_part, func))
if not collected_funcs:
print(f"[gen] warning: {label} collect 无结果 (rc={proc.returncode})"
f"stderr 尾部: {proc.stderr[-500:]}", file=sys.stderr)
return
for row in rows:
key = (row["file"], f"test_{row['op']}")
row["collected"] = "yes" if key in collected_funcs else "no"
n_missing = sum(1 for r in rows if r["collected"] == "no")
if n_missing:
print(f"[gen] {label}: {n_missing} 个 AST 可见但 pytest 未收集到的函数"
f"import 失败/条件定义),inventory 中 collected=no")
def carry_over_collected(out_dir: Path, label: str,
rows: List[Dict[str, str]]) -> None:
"""未跑 --verify-collect 时,从已有 CSV 继承 collected 列,避免重写清单时
把上次核实的结果冲掉。"""
path = out_dir / f"inventory_{label}.csv"
if not path.is_file():
return
with path.open(newline="") as f:
prev = {(r["op"], r["op_file"]): r.get("collected", "")
for r in csv.DictReader(f)}
for row in rows:
row["collected"] = prev.get((row["op"], row["op_file"]), "")
def write_inventory(out_dir: Path, label: str, rows: List[Dict[str, str]]) -> Path:
out = out_dir / f"inventory_{label}.csv"
tmp = out.with_name(out.name + ".tmp")
with tmp.open("w", newline="") as f:
writer = csv.DictWriter(f, fieldnames=INVENTORY_FIELDS)
writer.writeheader()
writer.writerows(rows)
os.replace(tmp, out)
return out
def migrate_shapes(old_path: Path, out_path: Path, valid_op_names: Set[str],
repos: List[Tuple[str, Path]]) -> None:
"""迁移旧 shape yaml,并以上游 core_shapes.yaml 为底座合并。
上游 set_shapes 的回退链是 op_name 键 → MRO 类名键(BlasBenchmark 等)→
基类 DEFAULT_SHAPES1 维)。--shape_file 是整体替换而非叠加,若我们的
yaml 缺少类名级条目,BLAS 这类需要 (B,M,N,K) 的算子会跌到 1 维默认值上
解包崩溃。因此把各仓库 core_shapes.yaml 里我们没有的键(含类名键与
op 级键)原样并入,迁移条目优先,先并 main 后并 vllm(同名键先到先得)。"""
data = yaml.safe_load(old_path.read_text()) or {}
core_maps = []
for label, root in repos:
core = root / "benchmark" / "core_shapes.yaml"
if core.is_file():
core_maps.append((label, yaml.safe_load(core.read_text()) or {}))
core_keys = {k for _, m in core_maps for k in m}
kept = {k: v for k, v in data.items() if k in valid_op_names}
# core_shapes 来源的键(类名键等)不算"未匹配"——它们每次都从上游取新版,
# 这样输入=输出的在地刷新也能同步上游 core_shapes 的变动。
dropped = {k: v for k, v in data.items()
if k not in valid_op_names and k not in core_keys}
n_curated = len(kept)
merged_from = []
for label, m in core_maps:
n_before = len(kept)
for key, val in m.items():
kept.setdefault(key, val)
merged_from.append(f"{label}:+{len(kept) - n_before}")
header = (f"# 由 gen_inventory.py 从 {old_path.name} 迁移:保留新上游仍存在的"
f" op_name 键 {n_curated}/{len(data)}"
f"并合并上游 core_shapes.yaml 缺失键({' '.join(merged_from)})。\n"
"# 顶层键 = FlagGems op_name 或 Benchmark 类名(上游回退链需要);\n"
"# 两者都无键的算子使用其基类 DEFAULT_SHAPES。\n")
out_path.write_text(header + yaml.safe_dump(
kept, sort_keys=True, default_flow_style=None, allow_unicode=True))
if dropped:
unmatched = out_path.with_suffix(".unmatched.yaml")
unmatched.write_text(
f"# {old_path.name} 中未匹配新上游任何 op_name 的键({len(dropped)} 个),"
"供人工复核后手动挪回。\n"
+ yaml.safe_dump(dropped, sort_keys=True, default_flow_style=None,
allow_unicode=True))
print(f"[gen] shapes: 保留 {len(kept)},剔除 {len(dropped)} -> {unmatched}")
else:
print(f"[gen] shapes: 全部 {len(kept)} 个键有效")
# 覆盖率:有效 op_name 里有多少没有 shape 条目(用上游默认 shape,仅提示)
uncovered = sorted(valid_op_names - set(kept))
print(f"[gen] shapes: 新上游 {len(valid_op_names)} 个 op_name 中 "
f"{len(uncovered)} 个无 shape 条目(将用上游默认 shape)")
def main() -> None:
args = parse_args()
repos = _resolve_repos(args.repo)
if not repos:
raise SystemExit("没有可扫描的仓库")
args.out_dir.mkdir(parents=True, exist_ok=True)
all_op_names: Set[str] = set()
for label, root in repos:
rows, op_names = scan_repo(label, root)
all_op_names |= op_names
if args.verify_collect:
verify_collect(label, root, rows)
else:
carry_over_collected(args.out_dir, label, rows)
out = write_inventory(args.out_dir, label, rows)
n_skip = sum(1 for r in rows if r["marked_skip"])
print(f"[gen] {label}: {len(rows)} 个测试函数 -> {out}"
f"(其中 {n_skip} 个带无条件 skip 标记)")
if args.migrate_shapes:
if not args.shapes_out:
raise SystemExit("--migrate-shapes 需要同时给 --shapes-out")
migrate_shapes(args.migrate_shapes, args.shapes_out, all_op_names, repos)
if __name__ == "__main__":
main()
File diff suppressed because it is too large Load Diff
+109
View File
@@ -0,0 +1,109 @@
repo,op,op_file,file,marked_skip,collected,op_names
vllm,perf_chunk_gated_delta_rule_fwd,FLA/test_chunk_gated_delta_rule_fwd,benchmark/test_FLA/test_chunk_gated_delta_rule_fwd.py,,yes,chunk_gated_delta_rule_fwd
vllm,perf_chunk_gated_delta_rule,FLA/test_chunk_gated_delta_rule_perf,benchmark/test_FLA/test_chunk_gated_delta_rule_perf.py,,yes,chunk_gated_delta_rule
vllm,chunk_gdn2,FLA/test_chunk_gdn2,benchmark/test_FLA/test_chunk_gdn2.py,,yes,chunk_gdn2
vllm,perf_chunk_gla,FLA/test_chunk_gla_perf,benchmark/test_FLA/test_chunk_gla_perf.py,,yes,
vllm,chunk_kda,FLA/test_chunk_kda,benchmark/test_FLA/test_chunk_kda.py,,yes,chunk_kda
vllm,perf_fused_recurrent_gated_delta_rule,FLA/test_fused_recurrent_gated_delta_rule_perf,benchmark/test_FLA/test_fused_recurrent_gated_delta_rule_perf.py,,yes,fused_recurrent_gated_delta_rule
vllm,act_quant_perf,act_quant,benchmark/test_act_quant.py,,yes,act_quant_triton
vllm,add_rms_norm,add_rms_norm,benchmark/test_add_rms_norm.py,,yes,add_rms_norm
vllm,apply_repetition_penalties,apply_repetition_penalties,benchmark/test_apply_repetition_penalties.py,,yes,apply_repetition_penalties
vllm,apply_rotary_pos_emb,apply_rotary_pos_emb,benchmark/test_apply_rotary_pos_emb.py,,yes,apply_rotary_pos_emb
vllm,beam_search_score,beam_search_score,benchmark/test_beam_search_score.py,,yes,beam_search_score
vllm,beam_search_score_,beam_search_score,benchmark/test_beam_search_score.py,,yes,beam_search_score_
vllm,bincount,bincount,benchmark/test_bincount.py,,yes,bincount
vllm,bincount_weighted,bincount,benchmark/test_bincount.py,,yes,bincount_weighted
vllm,blas_benchmark,blas_perf_parallel,benchmark/test_blas_perf_parallel.py,,yes,
vllm,perf_w8a8_block_fp8_matmul,blas_perf_parallel,benchmark/test_blas_perf_parallel.py,,yes,w8a8_block_fp8_matmul
vllm,perf_w8a8_block_fp8_matmul_deepgemm,blas_perf_parallel,benchmark/test_blas_perf_parallel.py,,yes,w8a8_block_fp8_matmul_deepgemm
vllm,perf_sparse_attention,blas_perf_parallel,benchmark/test_blas_perf_parallel.py,,yes,sparse_attention
vllm,mv_and_outer_benchmark,blas_perf_parallel,benchmark/test_blas_perf_parallel.py,,yes,
vllm,addmv_benchmark,blas_perf_parallel,benchmark/test_blas_perf_parallel.py,,yes,
vllm,vdot_benchmark,blas_perf_parallel,benchmark/test_blas_perf_parallel.py,,yes,vdot
vllm,addr_benchmark,blas_perf_parallel,benchmark/test_blas_perf_parallel.py,,yes,addr
vllm,perf_router_gemm,blas_perf_parallel,benchmark/test_blas_perf_parallel.py,,yes,router_gemm
vllm,concat_and_cache_mla,concat_and_cache_mla,benchmark/test_concat_and_cache_mla.py,,yes,concat_and_cache_mla
vllm,cp_gather_indexer_k_quant_cache_benchmark,cp_gather_indexer_k_quant_cache,benchmark/test_cp_gather_indexer_k_quant_cache.py,,yes,
vllm,cross_entropy_loss,cross_entropy_loss,benchmark/test_cross_entropy_loss.py,,yes,cross_entropy_loss
vllm,cutlass_scaled_mm_benchmark,cutlass_scaled_mm,benchmark/test_cutlass_scaled_mm.py,,yes,
vllm,combine_topk_swa_indices_benchmark,deepseek_v4_attention_combine_topk_swa_indices,benchmark/test_deepseek_v4_attention_combine_topk_swa_indices.py,,yes,
vllm,compute_global_topk_indices_and_lens_benchmark,deepseek_v4_attention_compute_global_topk_indices_and_lens,benchmark/test_deepseek_v4_attention_compute_global_topk_indices_and_lens.py,,yes,
vllm,dequantize_and_gather_k_cache_benchmark,deepseek_v4_attention_dequantize_and_gather_k_cache,benchmark/test_deepseek_v4_attention_dequantize_and_gather_k_cache.py,,yes,
vllm,fused_q_kv_rmsnorm_benchmark,deepseek_v4_attention_fused_q_kv_rmsnorm,benchmark/test_deepseek_v4_attention_fused_q_kv_rmsnorm.py,,yes,
vllm,dgeglu,dgeglu,benchmark/test_dgeglu.py,,yes,dgeglu
vllm,dreglu,dreglu,benchmark/test_dreglu.py,,yes,dreglu
vllm,dswiglu,dswiglu,benchmark/test_dswiglu.py,,yes,dswiglu
vllm,flash_attention_forward,flash_attention_forward,benchmark/test_flash_attention_forward.py,,yes,flash_attention_forward
vllm,flash_attn_varlen_func,flash_attn_varlen_func,benchmark/test_flash_attn_varlen_func.py,,yes,flash_attn_varlen_func
vllm,flash_attn_varlen_opt_func,flash_attn_varlen_opt_init_func,benchmark/test_flash_attn_varlen_opt_init_func.py,,yes,flash_attn_varlen_func
vllm,flash_mla,flash_mla,benchmark/test_flash_mla.py,,yes,flash_mla
vllm,flash_mla_sparse_fwd,flash_mla_sparse_fwd,benchmark/test_flash_mla_sparse_fwd.py,,yes,
vllm,flash_mla_with_kvcache,flash_mla_with_kvcache,benchmark/test_flash_mla_with_kvcache.py,,yes,
vllm,fp8_fp4_mqa_logits,fp8_fp4_mqa_logits,benchmark/test_fp8_fp4_mqa_logits.py,,yes,fp8_fp4_mqa_logits
vllm,fp8_fp4_paged_mqa_logits,fp8_fp4_paged_mqa_logits,benchmark/test_fp8_fp4_paged_mqa_logits.py,,yes,fp8_fp4_paged_mqa_logits
vllm,fused_add_rms_norm,fused_add_rms_norm,benchmark/test_fused_add_rms_norm.py,,yes,fused_add_rms_norm
vllm,fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert,fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert,benchmark/test_fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert.py,,yes,
vllm,fused_indexer_q_rope_quant,fused_indexer_q_rope_quant,benchmark/test_fused_indexer_q_rope_quant.py,,yes,fused_indexer_q_rope_quant
vllm,fused_inv_rope_fp8_quant,fused_inv_rope_fp8_quant,benchmark/test_fused_inv_rope_fp8_quant.py,,yes,fused_inv_rope_fp8_quant
vllm,fused_marlin_moe,fused_marlin_moe,benchmark/test_fused_marlin_moe.py,,yes,fused_marlin_moe
vllm,fused_moe_impl_gems_vs_vllm,fused_moe,benchmark/test_fused_moe.py,,yes,fused_experts_impl
vllm,fused_moe_fp8,fused_moe_fp8,benchmark/test_fused_moe_fp8.py,,yes,fused_experts_impl
vllm,fused_moe_fp8_blockwise,fused_moe_fp8_blockwise,benchmark/test_fused_moe_fp8_blockwise.py,,yes,fused_experts_impl
vllm,fused_moe_int4_w4a16,fused_moe_int4_w4a16,benchmark/test_fused_moe_int4_w4a16.py,,yes,fused_experts_impl
vllm,fused_experts_impl_int8,fused_moe_int8,benchmark/test_fused_moe_int8.py,,yes,fused_experts_impl
vllm,fused_experts_impl_int8_w8a16,fused_moe_int8_w8a16,benchmark/test_fused_moe_int8_w8a16.py,,yes,fused_moe_int8_w8a16_gems_vs_bf16_deq
vllm,fused_moe_w8a16_mxq,fused_moe_w8a16,benchmark/test_fused_moe_w8a16.py,,yes,fused_moe_w8a16_mxq_gems_vs_bf16_deq
vllm,fused_moe_w8a16_mxq_gems_vs_vllm,fused_moe_w8a16,benchmark/test_fused_moe_w8a16.py,,yes,fused_moe_w8a16_mxq_gems_vs_vllm
vllm,geglu,geglu,benchmark/test_geglu.py,,yes,geglu
vllm,gelu_and_mul,gelu_and_mul,benchmark/test_gelu_and_mul.py,,yes,gelu_and_mul
vllm,grouped_topk_no_renorm,grouped_topk,benchmark/test_grouped_topk.py,,yes,grouped_topk
vllm,grouped_topk_score_0,grouped_topk,benchmark/test_grouped_topk.py,,yes,grouped_topk
vllm,grouped_topk_score_1,grouped_topk,benchmark/test_grouped_topk.py,,yes,grouped_topk
vllm,indexer_k_quant_and_cache_benchmark,indexer_k_quant_and_cache,benchmark/test_indexer_k_quant_and_cache.py,,yes,
vllm,inplace_fused_experts_gems_vs_vllm,inplace_fused_experts,benchmark/test_inplace_fused_experts.py,,yes,inplace_fused_experts
vllm,instance_norm,instance_norm,benchmark/test_instance_norm.py,,yes,instance_norm
vllm,mhc_post,mhc,benchmark/test_mhc.py,,yes,mhc_post
vllm,mhc_pre,mhc,benchmark/test_mhc.py,,yes,mhc_pre
vllm,hc_split_sinkhorn_forward,mhc,benchmark/test_mhc.py,,yes,hc_split_sinkhorn_forward
vllm,mhc_bwd,mhc,benchmark/test_mhc.py,,yes,mhc_bwd
vllm,hc_head_fused_kernel,mhc,benchmark/test_mhc.py,,yes,hc_head_fused_kernel
vllm,moe_align_block_size_triton,moe_align_block_size_triton,benchmark/test_moe_align_block_size_triton.py,,yes,moe_align_block_size_triton
vllm,moe_sum,moe_sum,benchmark/test_moe_sum.py,,yes,moe_sum
vllm,mrope_gems_vs_torch,mrope,benchmark/test_mrope.py,,yes,mrope
vllm,outer,outer,benchmark/test_outer.py,,yes,outer
vllm,outplace_fused_experts_gems_vs_vllm,outplace_fused_experts,benchmark/test_outplace_fused_experts.py,,yes,outplace_fused_experts
vllm,pack_seq,pack_seq,benchmark/test_pack_seq.py,,yes,pack_seq_triton
vllm,pack_seq_fp8,pack_seq,benchmark/test_pack_seq.py,,yes,pack_seq_triton
vllm,perf_parallel_nsa,parallel_nsa,benchmark/test_parallel_nsa.py,,yes,parallel_nsa
vllm,perf_parallel_nsa_compression,parallel_nsa_compression,benchmark/test_parallel_nsa_compression.py,,yes,parallel_nsa_compression
vllm,per_token_group_quant_fp8,per_token_group_quant_fp8,benchmark/test_per_token_group_quant_fp8.py,,yes,per_token_group_quant_fp8
vllm,persistent_topk,persistent_topk,benchmark/test_persistent_topk.py,,yes,persistent_topk
vllm,reglu,reglu,benchmark/test_reglu.py,,yes,reglu
vllm,reshape_and_cache,reshape_and_cache,benchmark/test_reshape_and_cache.py,,yes,reshape_and_cache
vllm,reshape_and_cache_flash,reshape_and_cache_flash,benchmark/test_reshape_and_cache_flash.py,,yes,reshape_and_cache_flash
vllm,perf_router_gemm,router_gemm,benchmark/test_router_gemm.py,,yes,router_gemm
vllm,rwkv_ka_fusion,rwkv_ka_fusion,benchmark/test_rwkv_ka_fusion.py,,yes,rwkv_ka_fusion
vllm,rwkv_mm_sparsity,rwkv_mm_sparsity,benchmark/test_rwkv_mm_sparsity.py,,yes,rwkv_mm_sparsity
vllm,dynamic_scaled_int8_quant,scaled_int8_quant,benchmark/test_scaled_int8_quant.py,,yes,dynamic_scaled_int8_quant
vllm,static_scaled_int8_quant,scaled_int8_quant,benchmark/test_scaled_int8_quant.py,,yes,static_scaled_int8_quant
vllm,silu_and_mul,silu_and_mul,benchmark/test_silu_and_mul.py,,yes,silu_and_mul
vllm,silu_and_mul_out,silu_and_mul,benchmark/test_silu_and_mul.py,,yes,silu_and_mul_out
vllm,silu_and_mul_with_clamp,silu_and_mul_with_clamp,benchmark/test_silu_and_mul_with_clamp.py,,yes,silu_and_mul_with_clamp
vllm,silu_and_mul_with_clamp_out,silu_and_mul_with_clamp,benchmark/test_silu_and_mul_with_clamp.py,,yes,silu_and_mul_with_clamp_out
vllm,skip_layernorm,skip_layer_norm,benchmark/test_skip_layer_norm.py,,yes,skip_layer_norm
vllm,sparse_attn_triton,sparse_attention,benchmark/test_sparse_attention.py,yes,yes,sparse_attention
vllm,sparse_mla_fwd_interface,sparse_mla_fwd_interface,benchmark/test_sparse_mla_fwd_interface.py,,yes,sparse_mla_fwd_interface
vllm,stage_deepseek_v4_mega_moe_inputs_benchmark,stage_deepseek_v4_mega_moe_inputs,benchmark/test_stage_deepseek_v4_mega_moe_inputs.py,,yes,
vllm,swiglu,swiglu,benchmark/test_swiglu.py,,yes,swiglu
vllm,top_k_per_row_decode,top_k_per_row_decode,benchmark/test_top_k_per_row_decode.py,,yes,top_k_per_row_decode
vllm,top_k_per_row_prefill,top_k_per_row_prefill,benchmark/test_top_k_per_row_prefill.py,,yes,top_k_per_row_prefill
vllm,topk_softmax,topk_softmax,benchmark/test_topk_softmax.py,,yes,topk_softmax
vllm,topk_softplus_sqrt,topk_softplus_sqrt,benchmark/test_topk_softplus_sqrt.py,,yes,topk_softplus_sqrt
vllm,triton_scaled_mm_benchmark,triton_scaled_mm,benchmark/test_triton_scaled_mm.py,,yes,triton_scaled_mm
vllm,triton_unified_attention_perf,triton_unified_attention_perf,benchmark/test_triton_unified_attention_perf.py,,yes,
vllm,unpack_seq,unpack_seq,benchmark/test_unpack_seq.py,,yes,unpack_seq_triton
vllm,unpack_seq_fp8,unpack_seq,benchmark/test_unpack_seq.py,,yes,unpack_seq_triton
vllm,weight_norm_dim0,weight_norm,benchmark/test_weight_norm.py,,yes,weight_norm
vllm,weight_norm_dim_last,weight_norm,benchmark/test_weight_norm.py,,yes,weight_norm
vllm,weight_norm_interface,weight_norm_interface,benchmark/test_weight_norm_interface.py,,yes,weight_norm_interface
vllm,weight_norm_interface_backward,weight_norm_interface,benchmark/test_weight_norm_interface.py,,yes,weight_norm_interface_backward
1 repo op op_file file marked_skip collected op_names
2 vllm perf_chunk_gated_delta_rule_fwd FLA/test_chunk_gated_delta_rule_fwd benchmark/test_FLA/test_chunk_gated_delta_rule_fwd.py yes chunk_gated_delta_rule_fwd
3 vllm perf_chunk_gated_delta_rule FLA/test_chunk_gated_delta_rule_perf benchmark/test_FLA/test_chunk_gated_delta_rule_perf.py yes chunk_gated_delta_rule
4 vllm chunk_gdn2 FLA/test_chunk_gdn2 benchmark/test_FLA/test_chunk_gdn2.py yes chunk_gdn2
5 vllm perf_chunk_gla FLA/test_chunk_gla_perf benchmark/test_FLA/test_chunk_gla_perf.py yes
6 vllm chunk_kda FLA/test_chunk_kda benchmark/test_FLA/test_chunk_kda.py yes chunk_kda
7 vllm perf_fused_recurrent_gated_delta_rule FLA/test_fused_recurrent_gated_delta_rule_perf benchmark/test_FLA/test_fused_recurrent_gated_delta_rule_perf.py yes fused_recurrent_gated_delta_rule
8 vllm act_quant_perf act_quant benchmark/test_act_quant.py yes act_quant_triton
9 vllm add_rms_norm add_rms_norm benchmark/test_add_rms_norm.py yes add_rms_norm
10 vllm apply_repetition_penalties apply_repetition_penalties benchmark/test_apply_repetition_penalties.py yes apply_repetition_penalties
11 vllm apply_rotary_pos_emb apply_rotary_pos_emb benchmark/test_apply_rotary_pos_emb.py yes apply_rotary_pos_emb
12 vllm beam_search_score beam_search_score benchmark/test_beam_search_score.py yes beam_search_score
13 vllm beam_search_score_ beam_search_score benchmark/test_beam_search_score.py yes beam_search_score_
14 vllm bincount bincount benchmark/test_bincount.py yes bincount
15 vllm bincount_weighted bincount benchmark/test_bincount.py yes bincount_weighted
16 vllm blas_benchmark blas_perf_parallel benchmark/test_blas_perf_parallel.py yes
17 vllm perf_w8a8_block_fp8_matmul blas_perf_parallel benchmark/test_blas_perf_parallel.py yes w8a8_block_fp8_matmul
18 vllm perf_w8a8_block_fp8_matmul_deepgemm blas_perf_parallel benchmark/test_blas_perf_parallel.py yes w8a8_block_fp8_matmul_deepgemm
19 vllm perf_sparse_attention blas_perf_parallel benchmark/test_blas_perf_parallel.py yes sparse_attention
20 vllm mv_and_outer_benchmark blas_perf_parallel benchmark/test_blas_perf_parallel.py yes
21 vllm addmv_benchmark blas_perf_parallel benchmark/test_blas_perf_parallel.py yes
22 vllm vdot_benchmark blas_perf_parallel benchmark/test_blas_perf_parallel.py yes vdot
23 vllm addr_benchmark blas_perf_parallel benchmark/test_blas_perf_parallel.py yes addr
24 vllm perf_router_gemm blas_perf_parallel benchmark/test_blas_perf_parallel.py yes router_gemm
25 vllm concat_and_cache_mla concat_and_cache_mla benchmark/test_concat_and_cache_mla.py yes concat_and_cache_mla
26 vllm cp_gather_indexer_k_quant_cache_benchmark cp_gather_indexer_k_quant_cache benchmark/test_cp_gather_indexer_k_quant_cache.py yes
27 vllm cross_entropy_loss cross_entropy_loss benchmark/test_cross_entropy_loss.py yes cross_entropy_loss
28 vllm cutlass_scaled_mm_benchmark cutlass_scaled_mm benchmark/test_cutlass_scaled_mm.py yes
29 vllm combine_topk_swa_indices_benchmark deepseek_v4_attention_combine_topk_swa_indices benchmark/test_deepseek_v4_attention_combine_topk_swa_indices.py yes
30 vllm compute_global_topk_indices_and_lens_benchmark deepseek_v4_attention_compute_global_topk_indices_and_lens benchmark/test_deepseek_v4_attention_compute_global_topk_indices_and_lens.py yes
31 vllm dequantize_and_gather_k_cache_benchmark deepseek_v4_attention_dequantize_and_gather_k_cache benchmark/test_deepseek_v4_attention_dequantize_and_gather_k_cache.py yes
32 vllm fused_q_kv_rmsnorm_benchmark deepseek_v4_attention_fused_q_kv_rmsnorm benchmark/test_deepseek_v4_attention_fused_q_kv_rmsnorm.py yes
33 vllm dgeglu dgeglu benchmark/test_dgeglu.py yes dgeglu
34 vllm dreglu dreglu benchmark/test_dreglu.py yes dreglu
35 vllm dswiglu dswiglu benchmark/test_dswiglu.py yes dswiglu
36 vllm flash_attention_forward flash_attention_forward benchmark/test_flash_attention_forward.py yes flash_attention_forward
37 vllm flash_attn_varlen_func flash_attn_varlen_func benchmark/test_flash_attn_varlen_func.py yes flash_attn_varlen_func
38 vllm flash_attn_varlen_opt_func flash_attn_varlen_opt_init_func benchmark/test_flash_attn_varlen_opt_init_func.py yes flash_attn_varlen_func
39 vllm flash_mla flash_mla benchmark/test_flash_mla.py yes flash_mla
40 vllm flash_mla_sparse_fwd flash_mla_sparse_fwd benchmark/test_flash_mla_sparse_fwd.py yes
41 vllm flash_mla_with_kvcache flash_mla_with_kvcache benchmark/test_flash_mla_with_kvcache.py yes
42 vllm fp8_fp4_mqa_logits fp8_fp4_mqa_logits benchmark/test_fp8_fp4_mqa_logits.py yes fp8_fp4_mqa_logits
43 vllm fp8_fp4_paged_mqa_logits fp8_fp4_paged_mqa_logits benchmark/test_fp8_fp4_paged_mqa_logits.py yes fp8_fp4_paged_mqa_logits
44 vllm fused_add_rms_norm fused_add_rms_norm benchmark/test_fused_add_rms_norm.py yes fused_add_rms_norm
45 vllm fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert benchmark/test_fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert.py yes
46 vllm fused_indexer_q_rope_quant fused_indexer_q_rope_quant benchmark/test_fused_indexer_q_rope_quant.py yes fused_indexer_q_rope_quant
47 vllm fused_inv_rope_fp8_quant fused_inv_rope_fp8_quant benchmark/test_fused_inv_rope_fp8_quant.py yes fused_inv_rope_fp8_quant
48 vllm fused_marlin_moe fused_marlin_moe benchmark/test_fused_marlin_moe.py yes fused_marlin_moe
49 vllm fused_moe_impl_gems_vs_vllm fused_moe benchmark/test_fused_moe.py yes fused_experts_impl
50 vllm fused_moe_fp8 fused_moe_fp8 benchmark/test_fused_moe_fp8.py yes fused_experts_impl
51 vllm fused_moe_fp8_blockwise fused_moe_fp8_blockwise benchmark/test_fused_moe_fp8_blockwise.py yes fused_experts_impl
52 vllm fused_moe_int4_w4a16 fused_moe_int4_w4a16 benchmark/test_fused_moe_int4_w4a16.py yes fused_experts_impl
53 vllm fused_experts_impl_int8 fused_moe_int8 benchmark/test_fused_moe_int8.py yes fused_experts_impl
54 vllm fused_experts_impl_int8_w8a16 fused_moe_int8_w8a16 benchmark/test_fused_moe_int8_w8a16.py yes fused_moe_int8_w8a16_gems_vs_bf16_deq
55 vllm fused_moe_w8a16_mxq fused_moe_w8a16 benchmark/test_fused_moe_w8a16.py yes fused_moe_w8a16_mxq_gems_vs_bf16_deq
56 vllm fused_moe_w8a16_mxq_gems_vs_vllm fused_moe_w8a16 benchmark/test_fused_moe_w8a16.py yes fused_moe_w8a16_mxq_gems_vs_vllm
57 vllm geglu geglu benchmark/test_geglu.py yes geglu
58 vllm gelu_and_mul gelu_and_mul benchmark/test_gelu_and_mul.py yes gelu_and_mul
59 vllm grouped_topk_no_renorm grouped_topk benchmark/test_grouped_topk.py yes grouped_topk
60 vllm grouped_topk_score_0 grouped_topk benchmark/test_grouped_topk.py yes grouped_topk
61 vllm grouped_topk_score_1 grouped_topk benchmark/test_grouped_topk.py yes grouped_topk
62 vllm indexer_k_quant_and_cache_benchmark indexer_k_quant_and_cache benchmark/test_indexer_k_quant_and_cache.py yes
63 vllm inplace_fused_experts_gems_vs_vllm inplace_fused_experts benchmark/test_inplace_fused_experts.py yes inplace_fused_experts
64 vllm instance_norm instance_norm benchmark/test_instance_norm.py yes instance_norm
65 vllm mhc_post mhc benchmark/test_mhc.py yes mhc_post
66 vllm mhc_pre mhc benchmark/test_mhc.py yes mhc_pre
67 vllm hc_split_sinkhorn_forward mhc benchmark/test_mhc.py yes hc_split_sinkhorn_forward
68 vllm mhc_bwd mhc benchmark/test_mhc.py yes mhc_bwd
69 vllm hc_head_fused_kernel mhc benchmark/test_mhc.py yes hc_head_fused_kernel
70 vllm moe_align_block_size_triton moe_align_block_size_triton benchmark/test_moe_align_block_size_triton.py yes moe_align_block_size_triton
71 vllm moe_sum moe_sum benchmark/test_moe_sum.py yes moe_sum
72 vllm mrope_gems_vs_torch mrope benchmark/test_mrope.py yes mrope
73 vllm outer outer benchmark/test_outer.py yes outer
74 vllm outplace_fused_experts_gems_vs_vllm outplace_fused_experts benchmark/test_outplace_fused_experts.py yes outplace_fused_experts
75 vllm pack_seq pack_seq benchmark/test_pack_seq.py yes pack_seq_triton
76 vllm pack_seq_fp8 pack_seq benchmark/test_pack_seq.py yes pack_seq_triton
77 vllm perf_parallel_nsa parallel_nsa benchmark/test_parallel_nsa.py yes parallel_nsa
78 vllm perf_parallel_nsa_compression parallel_nsa_compression benchmark/test_parallel_nsa_compression.py yes parallel_nsa_compression
79 vllm per_token_group_quant_fp8 per_token_group_quant_fp8 benchmark/test_per_token_group_quant_fp8.py yes per_token_group_quant_fp8
80 vllm persistent_topk persistent_topk benchmark/test_persistent_topk.py yes persistent_topk
81 vllm reglu reglu benchmark/test_reglu.py yes reglu
82 vllm reshape_and_cache reshape_and_cache benchmark/test_reshape_and_cache.py yes reshape_and_cache
83 vllm reshape_and_cache_flash reshape_and_cache_flash benchmark/test_reshape_and_cache_flash.py yes reshape_and_cache_flash
84 vllm perf_router_gemm router_gemm benchmark/test_router_gemm.py yes router_gemm
85 vllm rwkv_ka_fusion rwkv_ka_fusion benchmark/test_rwkv_ka_fusion.py yes rwkv_ka_fusion
86 vllm rwkv_mm_sparsity rwkv_mm_sparsity benchmark/test_rwkv_mm_sparsity.py yes rwkv_mm_sparsity
87 vllm dynamic_scaled_int8_quant scaled_int8_quant benchmark/test_scaled_int8_quant.py yes dynamic_scaled_int8_quant
88 vllm static_scaled_int8_quant scaled_int8_quant benchmark/test_scaled_int8_quant.py yes static_scaled_int8_quant
89 vllm silu_and_mul silu_and_mul benchmark/test_silu_and_mul.py yes silu_and_mul
90 vllm silu_and_mul_out silu_and_mul benchmark/test_silu_and_mul.py yes silu_and_mul_out
91 vllm silu_and_mul_with_clamp silu_and_mul_with_clamp benchmark/test_silu_and_mul_with_clamp.py yes silu_and_mul_with_clamp
92 vllm silu_and_mul_with_clamp_out silu_and_mul_with_clamp benchmark/test_silu_and_mul_with_clamp.py yes silu_and_mul_with_clamp_out
93 vllm skip_layernorm skip_layer_norm benchmark/test_skip_layer_norm.py yes skip_layer_norm
94 vllm sparse_attn_triton sparse_attention benchmark/test_sparse_attention.py yes yes sparse_attention
95 vllm sparse_mla_fwd_interface sparse_mla_fwd_interface benchmark/test_sparse_mla_fwd_interface.py yes sparse_mla_fwd_interface
96 vllm stage_deepseek_v4_mega_moe_inputs_benchmark stage_deepseek_v4_mega_moe_inputs benchmark/test_stage_deepseek_v4_mega_moe_inputs.py yes
97 vllm swiglu swiglu benchmark/test_swiglu.py yes swiglu
98 vllm top_k_per_row_decode top_k_per_row_decode benchmark/test_top_k_per_row_decode.py yes top_k_per_row_decode
99 vllm top_k_per_row_prefill top_k_per_row_prefill benchmark/test_top_k_per_row_prefill.py yes top_k_per_row_prefill
100 vllm topk_softmax topk_softmax benchmark/test_topk_softmax.py yes topk_softmax
101 vllm topk_softplus_sqrt topk_softplus_sqrt benchmark/test_topk_softplus_sqrt.py yes topk_softplus_sqrt
102 vllm triton_scaled_mm_benchmark triton_scaled_mm benchmark/test_triton_scaled_mm.py yes triton_scaled_mm
103 vllm triton_unified_attention_perf triton_unified_attention_perf benchmark/test_triton_unified_attention_perf.py yes
104 vllm unpack_seq unpack_seq benchmark/test_unpack_seq.py yes unpack_seq_triton
105 vllm unpack_seq_fp8 unpack_seq benchmark/test_unpack_seq.py yes unpack_seq_triton
106 vllm weight_norm_dim0 weight_norm benchmark/test_weight_norm.py yes weight_norm
107 vllm weight_norm_dim_last weight_norm benchmark/test_weight_norm.py yes weight_norm
108 vllm weight_norm_interface weight_norm_interface benchmark/test_weight_norm_interface.py yes weight_norm_interface
109 vllm weight_norm_interface_backward weight_norm_interface benchmark/test_weight_norm_interface.py yes weight_norm_interface_backward
File diff suppressed because it is too large Load Diff
File diff suppressed because it is too large Load Diff
+511
View File
@@ -0,0 +1,511 @@
#!/usr/bin/env python3
"""多算子批量测试驱动:按算子清单逐个调 run_pytest.sh,多卡分片、串行测量。
每个算子一个独立子进程(CUDA crash 只废单个算子)、独立产物目录
<batch>/<repo>/<op>/(与单算子跑 run_pytest.sh 的产物结构完全一致,含
run.log / shapes.yaml / autotune_records / ttgir)。批量层只做编排:
- 分片:按算子名稳定哈希到 GPU(每卡内部串行,测量口径与单算子一致;
跨卡并行测量存在功耗耦合噪声,本模式定位为筛查口径,可疑算子回
单算子串行流程确认);
- 超时:--op-timeout 杀死挂死算子(记 TIMEOUT,不重试,继续跑);
- 重试:pytest rc=1(用例失败/flaky)与信号杀(rc>=128)最多重试 2 次,
连续两次同 rc 视为确定性失败提前止损;--dtypes 下两级降级——遇上游
"can't be supported by this op" 即去掉限制重跑,其余失败且无成功行时
也去掉限制最后救一次(dtype_rescue);
- 断点续跑:--resume 复用带 .complete 标记的算子(PASS/SKIP 才打标记,
FAIL/TIMEOUT 不打、下次重跑);
- replay--replay-root 指向上一次批量目录,逐算子传 REPLAY_FROM
<root>/<repo>/<op>,语义与单算子 REPLAY_FROM 相同。
产物:<batch>/ops_status.csv(每算子一行)+ summary.csv(每表行一行,
由 batch_summary.py 生成)。
示例:
# 全量筛查(两仓库全部可收集算子,排除 marked_skip 与 dismiss
python run_batch.py --all --gpus 0,1,2,3,4,5,6,7 --op-timeout 1800
# 指定算子子集(repo 前缀默认 main;重名用 @文件名 消歧)
python run_batch.py --ops softmax vllm:fused_marlin_moe nextafter_@nextafter_
# 锁 config 复测(口径同单算子 REPLAY_FROM
python run_batch.py --ops-file my_ops.txt --replay-root runs/batch_xxx
"""
from __future__ import annotations
import argparse
import csv
import hashlib
import os
import signal
import subprocess
import sys
import threading
import time
from collections import Counter
from concurrent.futures import ThreadPoolExecutor
from datetime import datetime
from pathlib import Path
from typing import Dict, List, Optional, Set, Tuple
import batch_summary
ZL_BENCH = Path(__file__).resolve().parent
OPS_DIR = ZL_BENCH / "ops"
LOG = "[batch]"
MAX_RETRIES = 2
REPO_DIRS = {
"main": "/workspace/dev/FlagGems",
"vllm": "/workspace/dev/FlagGems-vllm",
}
STATUS_FIELDS = [
"repo", "op", "op_file", "run_name", "gpu", "status", "return_code",
"elapsed_ms", "rows_total", "rows_success", "notes", "out_dir",
]
DTYPE_UNSUPPORTED_MARK = "can't be supported by this op"
# ---------------- 清单与选择 ----------------
class Case:
__slots__ = ("repo", "op", "op_file", "run_name")
def __init__(self, repo: str, op: str, op_file: str, run_name: str):
self.repo, self.op, self.op_file, self.run_name = repo, op, op_file, run_name
@property
def key(self) -> str:
return f"{self.repo}:{self.op}@{self.op_file}"
def load_inventory() -> Dict[str, List[Dict[str, str]]]:
inv: Dict[str, List[Dict[str, str]]] = {}
for label in REPO_DIRS:
path = OPS_DIR / f"inventory_{label}.csv"
if path.is_file():
with path.open(newline="") as f:
inv[label] = list(csv.DictReader(f))
if not inv:
raise SystemExit(f"{LOG} ops/inventory_*.csv 不存在,先跑 ops/gen_inventory.py")
return inv
def _run_name(row: Dict[str, str], dup_ops: Set[Tuple[str, str]]) -> str:
# 目录名默认用 op;仓库内重名时追加 @op_file 消歧(按 inventory 全局判定,
# 与选择子集无关,保证 resume/replay 路径稳定)。op_file 可能带子目录
# (test_FLA/ 等),做目录名时把 "/" 压成 "_"。
op = row["op"]
if (row["repo"], op) not in dup_ops:
return op
return f"{op}@{row['op_file'].replace('/', '_')}"
def _dup_ops(inv: Dict[str, List[Dict[str, str]]]) -> Set[Tuple[str, str]]:
dup = set()
for label, rows in inv.items():
counts = Counter(r["op"] for r in rows)
dup |= {(label, op) for op, n in counts.items() if n > 1}
return dup
def load_dismiss(path: Path) -> Set[Tuple[str, str]]:
if not path.is_file():
return set()
out = set()
for line in path.read_text().splitlines():
line = line.strip()
if not line or line.startswith("#"):
continue
repo, _, op = line.rpartition(":")
out.add((repo or "main", op))
return out
def resolve_selection(args: argparse.Namespace,
inv: Dict[str, List[Dict[str, str]]]) -> List[Case]:
dup = _dup_ops(inv)
index: Dict[Tuple[str, str, str], Dict[str, str]] = {}
by_op: Dict[Tuple[str, str], List[Dict[str, str]]] = {}
for label, rows in inv.items():
for r in rows:
index[(label, r["op"], r["op_file"])] = r
by_op.setdefault((label, r["op"]), []).append(r)
dismissed = load_dismiss(args.dismiss)
def make_case(row: Dict[str, str]) -> Case:
return Case(row["repo"], row["op"], row["op_file"], _run_name(row, dup))
if args.all:
cases = []
for label, rows in inv.items():
if args.repos and label not in args.repos:
continue
for r in rows:
if r.get("collected") == "no":
continue
if r.get("marked_skip") == "yes" and not args.include_marked_skip:
continue
if (label, r["op"]) in dismissed:
continue
cases.append(make_case(r))
return cases
specs: List[str] = list(args.ops or [])
if args.ops_file:
for line in args.ops_file.read_text().splitlines():
line = line.strip()
if line and not line.startswith("#"):
specs.append(line)
if not specs:
raise SystemExit(f"{LOG} 未选择算子:用 --all / --ops / --ops-file 之一")
cases, errors = [], []
for spec in specs:
repo, _, rest = spec.rpartition(":")
repo = repo or "main"
op, _, op_file = rest.partition("@")
if repo not in inv:
errors.append(f"{spec}: 未知仓库 {repo!r}")
continue
if op_file:
row = index.get((repo, op, op_file))
if row is None:
errors.append(f"{spec}: inventory 中无此条目")
continue
cases.append(make_case(row))
continue
rows = by_op.get((repo, op)) or []
if not rows:
errors.append(f"{spec}: inventory 中无此算子")
elif len(rows) > 1:
files = ", ".join(r["op_file"] for r in rows)
errors.append(f"{spec}: 重名,需用 @op_file 消歧(候选: {files}")
else:
cases.append(make_case(rows[0]))
if errors:
for e in errors:
print(f"{LOG} error: {e}", file=sys.stderr)
raise SystemExit(1)
return cases
# ---------------- GPU 分片 ----------------
def visible_gpus(arg: Optional[str]) -> List[str]:
if arg:
return [g.strip() for g in arg.split(",") if g.strip()]
env = os.environ.get("CUDA_VISIBLE_DEVICES", "").strip()
if env:
return [g.strip() for g in env.split(",") if g.strip()]
try:
out = subprocess.run(["nvidia-smi", "--list-gpus"], text=True,
capture_output=True, timeout=30).stdout
n = sum(1 for line in out.splitlines() if line.startswith("GPU "))
if n:
return [str(i) for i in range(n)]
except Exception:
pass
return ["0"]
def partition(cases: List[Case], gpus: List[str]) -> Dict[str, List[Case]]:
"""按 (repo, run_name) 稳定哈希分片:与选择顺序无关,resume 时同卡复用。"""
parts: Dict[str, List[Case]] = {g: [] for g in gpus}
for c in cases:
digest = hashlib.sha256(f"{c.repo}/{c.run_name}".encode()).digest()
parts[gpus[int.from_bytes(digest[:8], "big") % len(gpus)]].append(c)
return parts
# ---------------- 单算子执行 ----------------
def _spawn(case: Case, out_dir: Path, gpu: str, args: argparse.Namespace,
dtypes: str) -> subprocess.Popen:
env = os.environ.copy()
env.pop("PARALLEL_WARMUP_GPUS", None) # 批量层已分片,算子内不再并行
env.pop("REPLAY_FROM", None)
env.pop("OUT_DIR", None)
env.update({
"OP": case.op,
"OP_FILE": case.op_file,
"FLAGGEMS_DIR": REPO_DIRS[case.repo],
"OUT_DIR": str(out_dir),
"CUDA_VISIBLE_DEVICES": gpu,
"FLAGGEMS_PERF_COLOR": "never",
"NO_COLOR": "1",
"USE_FLAGTUNE": str(args.use_flagtune),
})
if args.shape_file:
env["SHAPE_FILE"] = str(args.shape_file)
if dtypes:
env["DTYPES"] = dtypes
if args.replay_root:
env["REPLAY_FROM"] = str(args.replay_root / case.repo / case.run_name)
return subprocess.Popen(
["bash", str(ZL_BENCH / "run_pytest.sh")],
cwd=str(ZL_BENCH), env=env, start_new_session=True,
stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL,
)
def _wait(proc: subprocess.Popen, timeout: int) -> Tuple[int, bool]:
"""返回 (rc, timed_out);超时对整个进程组 SIGTERM→SIGKILL。"""
try:
return proc.wait(timeout=timeout or None), False
except subprocess.TimeoutExpired:
try:
pgid = os.getpgid(proc.pid)
os.killpg(pgid, signal.SIGTERM)
try:
proc.wait(timeout=10)
except subprocess.TimeoutExpired:
os.killpg(pgid, signal.SIGKILL)
except ProcessLookupError:
pass # 组内进程恰好在超时瞬间退出
proc.wait()
return 124, True
def run_case(case: Case, batch_dir: Path, gpu: str,
args: argparse.Namespace) -> Dict[str, str]:
out_dir = batch_dir / case.repo / case.run_name
marker = out_dir / ".complete"
if args.resume and marker.is_file() and (out_dir / "run.log").is_file():
row = _classify(case, out_dir, gpu, 0, 0, "", timed_out=False)
row["notes"] = _join("resumed", row["notes"])
return row
out_dir.mkdir(parents=True, exist_ok=True)
marker.unlink(missing_ok=True)
dtypes = args.dtypes or ""
t0 = time.time()
rcs: List[int] = []
notes = ""
timed_out = False
for attempt in range(MAX_RETRIES + 1):
if attempt:
# 保留上一轮日志再重试
log = out_dir / "run.log"
if log.is_file():
log.rename(out_dir / f"run.attempt{attempt - 1}.log")
rc, timed_out = _wait(_spawn(case, out_dir, gpu, args, dtypes),
args.op_timeout)
rcs.append(rc)
if timed_out:
break
log_text = _read_log(out_dir)
# dtype 不支持:去掉 --dtypes 降级重试(只降一次,不计入重试预算)
if rc != 0 and dtypes and DTYPE_UNSUPPORTED_MARK in log_text:
dtypes = ""
notes = _join(notes, "dtype_fallback")
continue
if rc == 0 or not (rc == 1 or rc >= 128):
break
if len(rcs) >= 2 and rcs[-1] == rcs[-2]:
break # 确定性失败:同 rc 连续两次,不再烧预算
# dtype 营救:限制 dtype 时整体失败且没有任何成功行,去掉限制最后救一次。
# 覆盖 baseline 对受限 dtype 编译失败(torch jiterator 的 special 函数族等)
# 但报错文案不是上游 "can't be supported by this op" 的场景。
if dtypes and not timed_out:
probe = _classify(case, out_dir, gpu, rcs[-1], 0, "", False)
if probe["status"] == "FAIL" and probe["rows_success"] == "0":
log = out_dir / "run.log"
if log.is_file():
log.rename(out_dir / f"run.attempt{len(rcs) - 1}.log")
notes = _join(notes, "dtype_rescue")
rc, timed_out = _wait(_spawn(case, out_dir, gpu, args, ""),
args.op_timeout)
rcs.append(rc)
elapsed_ms = int((time.time() - t0) * 1000)
if len(rcs) > 1:
notes = _join(notes, "retry_rc=" + "->".join(map(str, rcs)))
row = _classify(case, out_dir, gpu, rcs[-1], elapsed_ms, notes, timed_out)
if row["status"] in ("PASS", "SKIP"):
tmp = marker.with_name(".complete.tmp")
tmp.write_text(row["status"] + "\n")
os.replace(tmp, marker)
return row
def _read_log(out_dir: Path) -> str:
log = out_dir / "run.log"
return log.read_text(errors="replace") if log.is_file() else ""
def _join(*notes: str) -> str:
return ";".join(n for n in notes if n)
def _classify(case: Case, out_dir: Path, gpu: str, rc: int, elapsed_ms: int,
notes: str, timed_out: bool) -> Dict[str, str]:
text = _read_log(out_dir)
rows = batch_summary.parse_run_log(text)
if timed_out:
status = "TIMEOUT"
else:
status = batch_summary.classify(rc, rows, text)
n_skipped = batch_summary.pytest_skipped_count(text)
if status == "PASS" and n_skipped:
notes = _join(notes, f"partial_skip={n_skipped}")
return {
"repo": case.repo,
"op": case.op,
"op_file": case.op_file,
"run_name": case.run_name,
"gpu": gpu,
"status": status,
"return_code": str(rc),
"elapsed_ms": str(elapsed_ms),
"rows_total": str(len(rows)),
"rows_success": str(sum(1 for r in rows if r["row_status"] == "SUCCESS")),
"notes": _join(notes, batch_summary.marker_notes(text)),
"out_dir": str(out_dir),
}
# ---------------- 主流程 ----------------
def parse_args() -> argparse.Namespace:
p = argparse.ArgumentParser(
description=__doc__.splitlines()[0],
formatter_class=argparse.RawDescriptionHelpFormatter)
sel = p.add_argument_group("算子选择")
sel.add_argument("--all", action="store_true",
help="跑 inventory 全量(排除 collected=no / marked_skip / dismiss")
sel.add_argument("--ops", nargs="*", default=None, metavar="[repo:]op[@op_file]",
help="直接指定算子(默认 main 仓;重名需 @op_file 消歧)")
sel.add_argument("--ops-file", type=Path, default=None,
help="算子清单文件,每行一个 [repo:]op[@op_file]# 注释")
sel.add_argument("--repos", nargs="*", default=None, choices=list(REPO_DIRS),
help="--all 时限定仓库(默认全部)")
sel.add_argument("--dismiss", type=Path, default=OPS_DIR / "dismiss.txt",
help="排除清单(默认 ops/dismiss.txt,仅 --all 生效)")
sel.add_argument("--include-marked-skip", action="store_true",
help="--all 时不剔除带无条件 skip 标记的算子")
run = p.add_argument_group("执行")
run.add_argument("--gpus", default=None,
help="逗号分隔 GPU 列表(默认 CUDA_VISIBLE_DEVICES 或全部)")
run.add_argument("--batch-dir", type=Path, default=None,
help="产物目录(默认 runs/batch_<时间戳>resume 时必须指定)")
run.add_argument("--op-timeout", type=int, default=1800,
help="单算子超时秒数,0=不限(默认 %(default)s")
run.add_argument("--resume", action="store_true",
help="复用 --batch-dir 里带 .complete 标记的算子结果")
run.add_argument("--replay-root", type=Path, default=None,
help="上次批量目录:逐算子 REPLAY_FROM <root>/<repo>/<op>")
run.add_argument("--shape-file", type=Path, default=OPS_DIR / "shapes_single.yaml",
help="shape yaml(默认 %(default)s;深查用 ops/shapes_multi.yaml"
"传空串则用各算子上游默认 shape")
run.add_argument("--dtypes", default="",
help='限制 dtype 集(空格分隔,如 "bfloat16");空=上游默认扫描')
run.add_argument("--use-flagtune", type=int, default=0, choices=(0, 1),
help="透传 USE_FLAGTUNE(默认 0;批量 A/B 两侧必须同值)")
args = p.parse_args()
if args.resume and not args.batch_dir:
raise SystemExit(f"{LOG} --resume 需要显式 --batch-dir")
if args.shape_file and str(args.shape_file).strip():
args.shape_file = args.shape_file.resolve()
if not args.shape_file.is_file():
raise SystemExit(f"{LOG} shape 文件不存在: {args.shape_file}")
else:
args.shape_file = None
if args.replay_root:
args.replay_root = args.replay_root.resolve()
if not args.replay_root.is_dir():
raise SystemExit(f"{LOG} replay 目录不存在: {args.replay_root}")
return args
def main() -> None:
args = parse_args()
inv = load_inventory()
cases = resolve_selection(args, inv)
if not cases:
raise SystemExit(f"{LOG} 选择结果为空")
gpus = visible_gpus(args.gpus)
batch_dir = (args.batch_dir or
ZL_BENCH / "runs" / f"batch_{datetime.now():%Y%m%d-%H%M%S}")
batch_dir = batch_dir.resolve()
batch_dir.mkdir(parents=True, exist_ok=True)
parts = {g: cs for g, cs in partition(cases, gpus).items() if cs}
print(f"{LOG} {len(cases)} 个算子 -> {len(parts)} 张卡 "
f"(shards: {[len(cs) for cs in parts.values()]}) out={batch_dir}")
if args.replay_root:
print(f"{LOG} replay 模式: REPLAY_FROM={args.replay_root}/<repo>/<op>")
lock = threading.Lock()
done = [0]
results: Dict[str, Dict[str, str]] = {}
status_csv = batch_dir / "ops_status.csv"
wname = max((len(c.run_name) for c in cases), default=8)
def flush_status() -> None:
rows = [results[c.key] for c in cases if c.key in results]
tmp = status_csv.with_name(status_csv.name + ".tmp")
with tmp.open("w", newline="") as f:
writer = csv.DictWriter(f, fieldnames=STATUS_FIELDS)
writer.writeheader()
writer.writerows(rows)
os.replace(tmp, status_csv)
def run_partition(gpu: str, part: List[Case]) -> None:
for case in part:
try:
row = run_case(case, batch_dir, gpu, args)
except Exception as exc: # 编排层异常不拖垮整卡分片
row = {
"repo": case.repo, "op": case.op, "op_file": case.op_file,
"run_name": case.run_name, "gpu": gpu, "status": "FAIL",
"return_code": "EXC", "elapsed_ms": "0", "rows_total": "0",
"rows_success": "0", "notes": f"harness_exc={exc}",
"out_dir": str(batch_dir / case.repo / case.run_name),
}
with lock:
done[0] += 1
results[case.key] = row
flush_status()
print(f" ({done[0]:>4}/{len(cases)}) {row['status']:<7} "
f"{row['elapsed_ms']:>8}ms rows={row['rows_success']}/"
f"{row['rows_total']:<4} {case.repo}:{case.run_name:<{wname}} "
f"[gpu{gpu}]"
+ (f" {row['notes']}" if row["notes"] else ""),
flush=True)
if len(parts) == 1:
gpu, part = next(iter(parts.items()))
run_partition(gpu, part)
else:
with ThreadPoolExecutor(max_workers=len(parts)) as pool:
futures = [pool.submit(run_partition, g, p) for g, p in parts.items()]
for f in futures:
f.result()
summary = batch_summary.summarize(batch_dir)
rows = [results[c.key] for c in cases if c.key in results]
counts = {s: sum(1 for r in rows if r["status"] == s)
for s in ("PASS", "SKIP", "FAIL", "TIMEOUT")}
print(f"{LOG} 完成: " + " ".join(f"{k}={v}" for k, v in counts.items()))
print(f"{LOG} 状态表: {status_csv}")
print(f"{LOG} 总表: {summary}")
bad = [r for r in rows if r["status"] in ("FAIL", "TIMEOUT")]
if bad:
print(f"{LOG} 失败算子: " +
", ".join(f"{r['repo']}:{r['run_name']}" for r in bad))
sys.exit(1)
if __name__ == "__main__":
main()
+60 -47
View File
@@ -1,65 +1,62 @@
#!/bin/bash
# 单算子性能测试入口(手动调试用)。
# 用法:改下方【配置区】,或用环境变量覆盖,例如:
# 单算子性能测试入口。改配置区或用同名环境变量覆盖:
# OP=softmax OP_FILE=softmax SHAPE_FILE=my_shapes.yaml bash run_pytest.sh
# 产物统一落在 runs/<op>_<时间戳>/run.log、shapes.yaml、autotune_records/、ttgir/
# 产物落在 runs/<op>_<时间戳>/run.log、shapes.yaml、autotune_records/、ttgir/
set -euo pipefail
# ============================== 配置区 ==============================
# 每项均可用同名环境变量覆盖,详见 README「环境变量一览」。
# 被测算子:OP=测试函数名(去掉 test_ 前缀);OP_FILE=benchmark 文件名
# (对应 $FLAGGEMS_DIR/benchmark/test_<OP_FILE>.py::test_<OP>
# 被测算子,对应 $FLAGGEMS_DIR/benchmark/test_<OP_FILE>.py::test_<OP>
OP="${OP:-fused_marlin_moe_mxfp4}"
OP_FILE="${OP_FILE:-fused_marlin_moe}"
# FlagGems 仓库路径
FLAGGEMS_DIR="${FLAGGEMS_DIR:-/workspace/FlagGems-dev}"
# 上游已拆分为两个仓库:主仓 /workspace/dev/FlagGems(包名 flag_gems)与
# /workspace/dev/FlagGems-vllm(包名 flaggems_vllm)。测 vllm 仓库算子时覆盖本变量即可。
FLAGGEMS_DIR="${FLAGGEMS_DIR:-/workspace/dev/FlagGems}"
# 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:
# 4 个 MoE 模型 x 4 档 token 数(M=1,16,64,256= 16 组 shape
shapes:
# Mixtral (E=8)
- [1, 8, 4096, 14336, 2]
- [16, 8, 4096, 14336, 2]
- [64, 8, 4096, 14336, 2]
- [256, 8, 4096, 14336, 2]
# DeepSeek-V3 (E=256, H=7168)
- [1, 256, 7168, 2048, 8]
- [16, 256, 7168, 2048, 8]
- [64, 256, 7168, 2048, 8]
- [256, 256, 7168, 2048, 8]
# Qwen3 (E=512)
- [1, 512, 4096, 1024, 10]
- [16, 512, 4096, 1024, 10]
- [64, 512, 4096, 1024, 10]
- [256, 512, 4096, 1024, 10]
# DeepSeek-V4-Flash (E=256, H=4096)
- [1, 256, 4096, 2048, 6]
- [16, 256, 4096, 2048, 6]
- [64, 256, 4096, 2048, 6]
- [256, 256, 4096, 2048, 6]
- [1, 256, 4096, 256, 6]
- [2, 256, 4096, 256, 6]
- [4, 256, 4096, 256, 6]
- [8, 256, 4096, 256, 6]
- [16, 256, 4096, 256, 6]
- [32, 256, 4096, 256, 6]
- [64, 256, 4096, 256, 6]
- [128, 256, 4096, 256, 6]
- [256, 256, 4096, 256, 6]
- [512, 256, 4096, 256, 6]
- [1024, 256, 4096, 256, 6]
- [2048, 256, 4096, 256, 6]
- [4096, 256, 4096, 256, 6]
- [8192, 256, 4096, 256, 6]
- [16384, 256, 4096, 256, 6]
- [32768, 256, 4096, 256, 6]
shape_desc: "num_tokens, num_experts, hidden_size, intermediate_size, topk"
YAML
# 调优空间:0=普通 autotune(默认,快速验证);1=FlagTune 扩展空间(首跑全量搜索、慢)
USE_FLAGTUNE="${USE_FLAGTUNE:-0}"
# autotune record/replay(见 _autotune_record_plugin.py):
# - 空:record 模式,本次选中的 config 记录到 $OUT_DIR/autotune_records/<op>.json
# - 指向某次历史 runs/<op>_<时间戳> 目录:replay 该次记录(A/B 两侧同 config)。
# 可选:限制 dtype 集(空格分隔,如 "bfloat16 float16")。空=上游默认 dtype 扫描。
# 算子不支持指定 dtype 时上游会报 "can't be supported by this op"(批量驱动据此降级重试)。
DTYPES="${DTYPES:-}"
# 空=record 模式,把本次选中的 config 记入 autotune_records/<op>.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
@@ -76,12 +73,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)"
LOG_FILE="$OUT_DIR/run.log"
# 产物目录,可用 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"
# 颜色决策:下方 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
@@ -99,7 +108,13 @@ export PYTHONPATH="$SCRIPT_DIR${PYTHONPATH:+:$PYTHONPATH}"
export FLAGGEMS_PERF_CURRENT_OP="$OP"
export PYTHONUNBUFFERED=1 # 实时输出不缓冲
# record/replay 二选一,通过环境变量激活对应模式
# DTYPES 非空时逐个转为上游 --dtypes 选项(action=append,每个 dtype 一次)
DTYPE_ARGS=()
for _dt in $DTYPES; do
DTYPE_ARGS+=(--dtypes "$_dt")
done
# 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" ]] || \
@@ -111,7 +126,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"
@@ -124,9 +139,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 后已编译部分仍可拿到)。
# 目录结构与命名图例见 <dump>/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
@@ -135,7 +149,7 @@ status=0
env "$AUTOTUNE_ENV" \
USE_FLAGTUNE=$USE_FLAGTUNE python -u -m pytest -s "$TEST_FILE" \
"${PLUGINS[@]}" "${PYTEST_COLOR[@]}" \
--shape_file "$SHAPE_FILE" \
--shape_file "$SHAPE_FILE" ${DTYPE_ARGS[@]+"${DTYPE_ARGS[@]}"} \
--level core --mode kernel || status=$?
rm -rf "$CACHE_DIR"
@@ -148,7 +162,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"
+19
View File
@@ -0,0 +1,19 @@
fused_marlin_moe_mxfp4:
shapes:
- [1, 256, 4096, 256, 6]
- [2, 256, 4096, 256, 6]
- [4, 256, 4096, 256, 6]
- [8, 256, 4096, 256, 6]
- [16, 256, 4096, 256, 6]
- [32, 256, 4096, 256, 6]
- [64, 256, 4096, 256, 6]
- [128, 256, 4096, 256, 6]
- [256, 256, 4096, 256, 6]
- [512, 256, 4096, 256, 6]
- [1024, 256, 4096, 256, 6]
- [2048, 256, 4096, 256, 6]
- [4096, 256, 4096, 256, 6]
- [8192, 256, 4096, 256, 6]
- [16384, 256, 4096, 256, 6]
- [32768, 256, 4096, 256, 6]
shape_desc: "num_tokens, num_experts, hidden_size, intermediate_size, topk"