4 Commits
Author SHA1 Message Date
AharrypotterandClaude Opus 5 783c0655bc fix: never let the parallel sweep cover only part of the shape set
The sharded autotune sweep could publish a config set with holes, and a run
using it reports itself as replaying while every uncovered key silently
self-tunes. That is the exact failure the record/replay caliber exists to
prevent, and it was reachable two ways.

Shape lookup was op_name-only, but BLAS operators keep their shapes under the
class name: core_shapes.yaml has BlasBenchmark, not mm. Every BLAS operator
therefore resolved to zero shapes and the plugin disabled itself with a message
that read like a normal skip. It now walks the benchmark class chain the way
FlagGems' set_shapes does. The chain is recovered statically -- the class comes
from the op_name= keyword of the constructor call in the test module, its bases
from that module plus benchmark/base.py -- because the plugin runs at
pytest_configure, before any Benchmark instance exists. Resolved keys are named
in the log, and a total miss lists the candidates it tried.

Two cases can still produce partial coverage, and both now refuse rather than
shard a subset:

- --level comprehensive, where set_shapes() merges set_more_shapes() on top of
  the yaml. That addition is arbitrary per-subclass Python and invisible here.
  Measured on mm: 5 of 29 shapes swept, 90 AUTOTUNE_REPLAY_FALLBACK
  reason=key_missing in the run that consumed the result.
- a shard exiting non-zero, whose shapes were never swept. This used to append a
  WARNING to a success message and publish the merge anyway; observed with 3 of
  5 shards dead. The whole sweep is now discarded, worker logs kept.

Both fall back to the normal serial sweep, which costs wall time rather than the
same-config premise.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
2026-09-09 16:40:12 +08:00
AharrypotterandClaude Opus 5 92dae82c7f feat: expose the upstream benchmark level as BENCH_LEVEL
run_pytest.sh hardcoded --level core, so the only shapes ever reachable were
each operator's core_shapes entry. FlagGems merges a subclass's
set_more_shapes() on top of that at --level comprehensive, and for BLAS that
addition is where the large GEMMs live: mm goes from 5 shapes to 29, taking in
four large-K tuples plus models_shapes/*.yaml, with N/K up to 152064. The
comprehensive path also yields each shape a second time (transposed, for BLAS),
so the case count doubles. None of that was reachable without editing the driver.

BENCH_LEVEL defaults to core, so existing runs are unchanged, and rejects
anything other than core/comprehensive rather than passing it through to pytest.
The chosen level is echoed in the run header, written into the archived
shapes.yaml, and recorded in run_identity.json, because an A/B whose two sides
used different levels is not comparable and is otherwise invisible afterwards.

Sweep cost scales with the shape count, so comprehensive on a BLAS operator
wants PARALLEL_WARMUP_GPUS.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
2026-09-09 16:39:45 +08:00
AharrypotterandClaude Opus 5 c8ea586d0f feat: make the harness backend-agnostic
Everything about which accelerator a run targets was previously spelled CUDA:
device discovery went through nvidia-smi, pinning through CUDA_VISIBLE_DEVICES,
the vendor guard only recognised NVIDIA, and the -p plugin list was hardcoded in
run_pytest.sh with no way to drop the CUDA-only members. Running the same
operator corpus on another vendor's device therefore meant editing the drivers.

The selection now hangs off one ZL_BACKEND variable (cuda/musa/rocm/ascend, with
the usual aliases):

- backends/ owns per-backend device enumeration, visibility env keys and
  capability flags, so nothing above it names a vendor;
- plugin_plan.py derives the -p list from those capabilities instead of a
  hardcoded array, which is what lets a non-CUDA run omit the graph-capture and
  cluster-fix plugins rather than watch them fail;
- _timing_plugin selects the timer from the same capabilities: CUDA keeps
  graph capture with its per-shape fallback, other backends stay on upstream
  do_bench and record timer=direct, because capture-then-fallback on a device
  without graph support only produces noise;
- _device_setup_plugin replaces the NVIDIA-only vendor guard;
- providers/flaggems.py resolves the pytest node, so a caller may pass TEST_NODE
  directly instead of the OP/OP_FILE pair, and SHAPE_MODE=upstream can defer to
  FlagGems' own core_shapes;
- _run_identity.py writes run_identity.json next to each run, recording backend,
  device, timer, autotune mode and the resolved test node. Two runs whose
  calibers differ are otherwise indistinguishable after the fact.

run_batch gains --backend and --devices (--gpus kept as an alias) and carries
backend/device/timer into ops_status.csv. test_backends.py covers the registry,
alias resolution, visibility env and plugin planning.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
2026-09-09 16:39:14 +08:00
AharrypotterandClaude Opus 5 35a928a516 chore: pin shipped text formats to LF
Everything in this repository runs on Linux accelerator hosts, and campaign
receipts record the SHA-256 of the harness files that produced them. With Git's
default autocrlf on Windows a checkout rewrites the shell drivers with CRLF,
which breaks them outright on the target host and changes every recorded digest.

The two inventory CSVs are renormalized here because their content is unchanged;
the remaining files normalize as part of the commits that actually modify them.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
2026-09-09 16:38:33 +08:00
27 changed files with 2764 additions and 1301 deletions

No files matched your search

+12
View File
@@ -0,0 +1,12 @@
# Everything here runs on Linux accelerator hosts, and campaign receipts record
# the SHA-256 of the harness files they were produced by. A CRLF checkout on
# Windows would both break the shell drivers and change every recorded digest,
# so pin the text formats this repository ships.
*.py text eol=lf
*.sh text eol=lf
*.json text eol=lf
*.yaml text eol=lf
*.csv text eol=lf
*.tsv text eol=lf
*.txt text eol=lf
*.md text eol=lf
+50 -19
View File
@@ -20,6 +20,11 @@ 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
# 对齐 FlagGems 原生节点(不传 --shape_file,沿用 --level core 的 core_shapes)
ZL_BACKEND=cuda \
TEST_NODE=/workspace/FlagGems/benchmark/test_fused_marlin_moe.py::test_fused_marlin_moe_mxfp4 \
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)。
@@ -35,7 +40,7 @@ softmax:
上面这条命令是 **record 模式**:现场 autotune、把选中的 config 存档。它用于单次摸底或给 replay 提供基准,**两次 record 的数字不能互相比**——做对比走下面的「A/B 对比测试标准流程」。
跑之前先 `nvidia-smi` 确认目标卡空闲:共享机器上别的任务会把两侧一起等比拖慢,出一份看似自洽实则作废的数据。多卡机器用 `CUDA_VISIBLE_DEVICES=<空闲卡号>` 选卡。**测量始终单进程单卡串行**,这是计时可比的前提;多卡仅用于可选的 sweep 加速(见「加速:多卡并行 sweep」)。
跑之前确认目标卡空闲:共享机器上别的任务会把两侧一起等比拖慢,出一份看似自洽实则作废的数据。用 `ZL_BACKEND` 选运行时(默认 `cuda`),再用该后端的可见性变量选卡(CUDA 为 `CUDA_VISIBLE_DEVICES`,也可设 `ZL_DEVICE`)。**测量始终单进程单设备串行**,这是计时可比的前提;多卡仅用于可选的 sweep 加速(见「加速:多卡并行 sweep」)。非 CUDA 后端默认走上游 `do_bench`,不会尝试 CUDA Graph。
脚本内固定了 `--level core --mode kernel`,其余可调项都走环境变量(见下);需要改动这两个口径本身时直接编辑 `run_pytest.sh` 中的 pytest 命令行。
@@ -45,10 +50,14 @@ softmax:
| 变量 | 默认 | 说明 |
|-----|------|------|
| `OP` | `fused_marlin_moe_mxfp4` | 测试函数名(`test_` 之后的部分) |
| `OP_FILE` | `fused_marlin_moe` | benchmark 文件名(`test_` 与 `.py` 之间的部分) |
| `SHAPE_FILE` | 空(用脚本内置 yaml) | shape yaml 路径 |
| `FLAGGEMS_DIR` | `/workspace/dev/FlagGems` | FlagGems 仓库路径(测 vllm 仓库时指向 `/workspace/dev/FlagGems-vllm`) |
| `OP` | `fused_marlin_moe_mxfp4` | 测试函数名(`test_` 之后的部分);设了 `TEST_NODE` 时由节点覆盖 |
| `OP_FILE` | `fused_marlin_moe` | benchmark 文件名(`test_` 与 `.py` 之间的部分);设了 `TEST_NODE` 时由节点覆盖 |
| `TEST_NODE` | 空 | FlagGems pytest 节点(`…/benchmark/test_<op_file>.py::test_<op>`)。设了且未指定 `SHAPE_FILE` 时默认不传 `--shape_file` |
| `SHAPE_FILE` | 空 | shape yaml 路径;空且无 `TEST_NODE` 时用脚本内置 yaml;`-`/`upstream`/`none` 表示沿用上游在 `BENCH_LEVEL` 下的原生 shape 集 |
| `SHAPE_MODE` | 空 | `inline` / `file` / `upstream`;通常不用手设,`TEST_NODE` 会把空 `SHAPE_FILE` 推成 `upstream` |
| `FLAGGEMS_DIR` | `/workspace/dev/FlagGems` | FlagGems 仓库路径。测试文件缺失时回退 `/workspace/FlagGems`;`TEST_NODE` 也可直接带出仓库根 |
| `ZL_BACKEND` | `cuda` | 设备后端:`cuda` / `musa` / `rocm` / `ascend`(别名 `nvidia`/`amd`/`npu` 等;未知名直接失败) |
| `ZL_DEVICE` | 空 | 父进程看到的设备 id,写入 `run_identity.json`;批量层绑卡后也会设 |
**怎么测**
@@ -56,8 +65,14 @@ softmax:
|-----|------|------|
| `REPLAY_FROM` | 空(record 模式) | 指向某次历史 run 目录,replay 其 autotune 选择。做 A/B 时 B 侧必须设(见「A/B 对比测试标准流程」) |
| `USE_FLAGTUNE` | `0` | `0` 走普通 autotune(跳过搜索、快速验证);`1` 用 FlagTune 扩展调优空间(首跑全量搜索、慢)。A/B 两侧须取同值 |
| `BENCH_LEVEL` | `core` | 透传上游 `--level`。`comprehensive` 相对 `core` 做两件事:① 在 `set_shapes()` 里并入子类 `set_more_shapes()`(`mm` 因此从 5 个 shape 变 29 个,含 N/K 到 152064 的大 K 与 `models_shapes/*.yaml` 的模型 shape);② 在 `get_input_iter()` 里对每个 shape 再产出一次变体(BLAS 是转置侧),用例数翻倍。注意 ① 只在 `SHAPE_MODE=upstream` 下可见——传了 `SHAPE_FILE` 时 `_shape_inject_plugin` 会在上游合并之后再覆盖一次、把 shape 集拉回 yaml 内容(BLAS 那类 arity 不一致的情形除外,插件会退让);② 不经插件,任何 shape 来源下都生效。sweep 成本同步放大,shape 多时配合 `PARALLEL_WARMUP_GPUS`。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」) |
| `ZL_ENABLE_CUDAGRAPH` | `auto` | `auto`/`1` 仅在 CUDA 上走 Graph 计时;非 CUDA 即使设 `1` 也拒绝并保持 `timer=direct` |
| `ZL_ENABLE_AUTOTUNE` | `1` | `0` 关闭 autotune record/replay 插件 |
| `ZL_ENABLE_IR` | `1` | `0` 关闭 TTGIR / 编译元数据插件 |
| `ZL_ENABLE_MM_CLUSTER` | `0` | `1` 且后端支持时加载 Hopper mm cluster 规避 |
| `ZL_EXTRA_PLUGINS` | 空 | 空格分隔、追加到 `-p` 列表的插件名 |
**输出**
@@ -72,9 +87,11 @@ softmax:
```
runs/softmax_20260716-031752/
├── run.log # 完整日志(含 SUCCESS 行的 latency/speedup 表;已去 ANSI 色的纯文本)
├── run_identity.json # provider/backend/device/timer/vendor(插件合并写入)
├── shapes.yaml # 本次实际使用的 shape(存档)
├── autotune_records/ # Triton autotune 选中的 config(供 replay)
│ └── softmax.json
│ ├── softmax.json
│ └── meta.json # provider/backend;跨后端 replay 会被拒绝
└── ttgir/ # 按 shape 分组的 IR 落盘
├── 1024x1024/
│ └── softmax_kernel_inner/
@@ -152,7 +169,7 @@ shape 多时墙上时间几乎全在 autotune sweep,而不在测量——被
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`,分片临时目录跑完即删。
加速比取决于 sweep 占比:shape 多、config 空间大的算子收益最明显,可达数倍。不设该变量(或设 <2)时插件静默 no-op。**产物结构、`REPLAY_FROM` 用法、latency 表格式全不变**——合并后的 config 就写进本次 run 的 `autotune_records/<op>.json`,分片临时目录跑完即删。
**测量本身不并行**:N 个进程压满同机多卡会通过功耗墙/散热耦合,单卡 latency 被邻居拖慢且不可复现。并行只用于"决定哪个 config 胜出",测量仍是单进程单卡、与不开插件走同一条路径。
@@ -162,7 +179,17 @@ PARALLEL_WARMUP_GPUS=8 SHAPE_FILE=shape/DeepSeek-V4-Flash-p32768d1024.yaml bash
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`——用它选过卡时,分片只会落在你选的那几张上。
**冲突数就是这次加速的可信度指标**:0 可放心用;偏多说明 config 选择已被污染,出正式结论前不设该变量重跑一遍。
插件会在三种情况下**主动放弃分片、退回串行 sweep**,日志各有独立措辞,都不是静默的:
- **`BENCH_LEVEL` 不是 `core`**。`comprehensive` 会让上游在 `set_shapes()` 里并入 `set_more_shapes()`,那部分是各子类自己的 Python(硬编码元组加 `models_shapes/*.yaml`),插件只看得到 yaml 里的 shape 列表。按子集分片会让未覆盖的 key 全部现场自调优,而这一轮却自称在 replay——`mm` 上实测是 29 个 shape 只扫到 5 个、90 次 `AUTOTUNE_REPLAY_FALLBACK reason=key_missing`。
- **任一分片非零退出**。死掉的分片的 shape 从未被扫过,合并出来的 config 集是带洞的,publish 之后每个洞都会静默自调优。因此整轮 sweep 直接作废,worker 日志保留在临时目录里备查。
- yaml 里查不到该算子的 shape(见下)、`REPLAY_FROM` 已设、或可见卡不足 2 张。
yaml key 的查找与上游对齐:先按 `op_name`,查不到再沿 Benchmark 类链逐级尝试类名(`mm` 因此命中 `BlasBenchmark`)。类链是静态解析测试模块里 `op_name=` 那次构造调用得到的,不 import FlagGems。日志会打出命中的 key;一个都没命中时会列出尝试过的候选。
其余行为(轮转分片而非按块切、用子进程而非 xdist 的原因)见插件 docstring。分片绑卡遵守当前后端的可见性变量(CUDA 为 `CUDA_VISIBLE_DEVICES`)——用它选过卡时,分片只会落在你选的那几张上。
## 批量多算子测试(run_batch.py)
@@ -170,7 +197,8 @@ WARNING 1 config key(s) got different winners across shards — parallel interfe
```bash
# 全量筛查(双仓库全部可收集算子,排除带无条件 skip 标记的与 ops/dismiss.txt 里的)
python run_batch.py --all --gpus 0,1,2,3,4,5,6,7 --dtypes bfloat16 --op-timeout 1800
python run_batch.py --all --devices 0,1,2,3,4,5,6,7 --backend cuda --dtypes bfloat16 --op-timeout 1800
# --gpus 仍是 --devices 的别名;两者同时给出会报错
# 指定子集:默认 main 仓,vllm 仓加前缀;同名测试函数用 @文件名 消歧
python run_batch.py --ops softmax vllm:fused_marlin_moe nextafter_@nextafter_
@@ -195,7 +223,7 @@ python run_batch.py --ops-file my_ops.txt --replay-root runs/batch_xxx
**执行语义**:每算子独立子进程(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` 计数,解读口径时先看这列。
**产物**:`<batch>/<repo>/<op>/` 与单算子 run 目录结构完全一致(run.log、run_identity.json、autotune_records/、ttgir/……),所以任何一个算子都可以事后单独 `REPLAY_FROM` 复测。批量层额外产出 `ops_status.csv`(每算子一行:backend/device/timer、状态/耗时/行数/备注,边跑边原子更新;`gpu` 列与 `device` 同值)与 `summary.csv`(每测量行一行:backend/device/timer、dtype、latency、speedup、Size Detail,来源是 run.log 的结果表,`batch_summary.py` 也可单独对旧批次重跑)。备注列聚合 `no_cudagraph_rows`/`replay_fallback` 计数,解读口径时先看这列。不同后端、不同 timer 的行不要当成同口径。
## 计时口径:cudagraph 的预热、回退与精度
@@ -208,27 +236,28 @@ python run_batch.py --ops-file my_ops.txt --replay-root runs/batch_xxx
回退本身不影响 A/B 公平性(两侧同一算子回退行为一致),但**回退口径的测量误差更大**:do_bench 每次迭代都走完整的 Python → launch 路径,kernel 越小,launch 开销和 CPU 侧抖动在数字里占比越高——亚毫秒级 kernel 上两种口径可差 2 倍以上,且行间波动更明显。解读这类算子的结果时:
- 小 shape 行的绝对值和小幅(<10%)差异不要过度解读,优先看大 shape 行;
- 需要更高置信度时,同一配置多跑几次取中位数,或对该算子直接注释掉 `_cudagraph_plugin` 统一用 do_bench 口径(消除同表混两种口径的问题)。
- 需要更高置信度时,同一配置多跑几次取中位数,或设 `ZL_ENABLE_CUDAGRAPH=0` 统一用 `timer=direct`(消除同表混两种口径的问题)。
回退是按测量点发生的,同一份 latency 表里可能混有两种口径的行;若不确定,先 `grep BENCHMARK_DIRECT_NO_CUDAGRAPH run.log` 确认哪些行是回退口径再下结论。
## 插件说明
脚本通过 `-p` 加载以下插件(`run_pytest.sh` 的 `PLUGINS` 数组,可按需注释):
`-p` 列表由 `plugin_plan.py` 按 `ZL_BACKEND` 生成,不再在 `run_pytest.sh` 里写死。`_cudagraph_plugin` **不**出现在默认列表上:只有 CUDA 且 `_timing_plugin` 选中 Graph 时才会 import 它。
| 插件 | 作用 | 何时关闭 |
|-----|------|---------|
| `_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 |
| `_device_setup_plugin` | 按后端设 `GEMS_VENDOR`(cuda→nvidia),跳过 flag_gems 启动时可能挂死的 `nvidia-smi` 探测,并写入 `run_identity.json`。旧名 `_device_guard_plugin` 仍是兼容 shim | 已设 `GEMS_VENDOR` 时不覆盖 |
| `_parallel_warmup_plugin` | 由 `PARALLEL_WARMUP_GPUS=N` 激活(未设则静默 no-op):autotune sweep 分片到 N 卡并行,测量仍单卡串行、产物结构不变(见「加速:多卡并行 sweep」) | 不设该变量即关闭 |
| `_seed_plugin` | 固定 random/numpy/torch 种子,数据相关算子(sort/topk 等)输入逐字节一致 | 不关 |
| `_shape_inject_plugin` | 让 shape yaml 覆盖子类硬编码的 `set_shapes()` | 不关 |
| `_shape_iter_inject_plugin` | 覆盖在 `get_input_iter` 里硬编码 shape 的类(conv/pool 等) | 不关 |
| `_bespoke_shape_plugin` | 覆盖特殊输入构造的算子(upsample/flash_mla/cutlass 等,按类名注册) | 测这些算子之外可关 |
| `_autotune_record_plugin` | record/replay Triton autotune 选择(由 RECORD/REPLAY 环境变量二选一激活) | 不关 |
| `_cudagraph_plugin` | `do_bench` → `do_bench_cudagraph`(kernel 纯耗时;内部先做跨流同步,修过一个间歇性 illegal instruction)。**capture 前按 warmup 时间预算显式预热**(先一次丢弃跑吸收 autotune/JIT 编译,再稳态预热),消除首轮抖动。无法 graph capture 的 kernel 自动回退并打 `BENCHMARK_DIRECT_NO_CUDAGRAPH` 标记 | 需要与他人的普通 do_bench 数据对齐时注释掉 |
| `_pretty_report_plugin` | 结果表整理 + 着色:全表相同的输入折叠成表头下一行图例,每行 Size Detail 只留随行变化的部分(shape 保持 `torch.Size([...])` 原样,MoE 类算子单行从 ~400 字符缩到一屏内);SUCCESS 绿 / FAILED 红。列名与 `SUCCESS` 字样保持上游原文,`run.log` 的 grep/解析不受影响 | 需要与上游原始表格逐字对齐时注释掉 |
| `_ir_meta_plugin` | 编译期记录每个变体的 constexpr/签名/特化;launch 钩子统计每个 (kernel, shape) 的实际使用;退出时按上述结构 dump ttgir | 不关 |
| `_mm_cluster_fix_plugin` | Hopper fp16 mm cluster kernel 越界崩溃的运行时规避 | 默认注释;测 fp16 mm 崩溃时打开 |
| `_autotune_record_plugin` | record/replay Triton autotune 选择;`autotune_records/meta.json` 记录 backend,跨后端 replay 拒绝 | `ZL_ENABLE_AUTOTUNE=0` |
| `_timing_plugin` | 按后端选择计时路径:CUDA 可走 Graph(内部调用 `_cudagraph_plugin`),其余后端一开始就是上游 `do_bench`(`timer=direct`) | `ZL_ENABLE_CUDAGRAPH=0` 强制 direct |
| `_cudagraph_plugin` | 仅由 `_timing_plugin` 在 CUDA 上调用。`do_bench` → `do_bench_cudagraph`,capture 前按 warmup 预算显式预热;无法 capture 时回退并打 `BENCHMARK_DIRECT_NO_CUDAGRAPH` | 不要直接 `-p` 加到非 CUDA 运行 |
| `_pretty_report_plugin` | 结果表整理 + 着色:全表相同的输入折叠成表头下一行图例,每行 Size Detail 只留随行变化的部分(shape 保持 `torch.Size([...])` 原样,MoE 类算子单行从 ~400 字符缩到一屏内);SUCCESS 绿 / FAILED 红。列名与 `SUCCESS` 字样保持上游原文,`run.log` 的 grep/解析不受影响 | 需要与上游原始表格逐字对齐时去掉 |
| `_ir_meta_plugin` | 编译期记录每个变体的 constexpr/签名/特化;launch 钩子统计每个 (kernel, shape) 的实际使用;退出时按上述结构 dump ttgir。Triton knobs 缺失时 fail-soft | `ZL_ENABLE_IR=0` |
| `_mm_cluster_fix_plugin` | Hopper fp16 mm cluster kernel 越界崩溃的运行时规避 | 默认关;`ZL_ENABLE_MM_CLUSTER=1` 或 `ZL_EXTRA_PLUGINS` |
## 常见问题
@@ -257,5 +286,7 @@ shape 多时这部分会主导墙上时间,可设 `PARALLEL_WARMUP_GPUS=N` 并
## 依赖假设
- 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 满足);
- 设备后端由 `ZL_BACKEND` / `--backend` 选择(`cuda` / `musa` / `rocm` / `ascend`)。inventory 仍是 FlagGems 逻辑算子清单,同一 case 可在多个后端上跑,不要按硬件复制 CSV;
- Triton 需支持 `knobs.compilation.listener`、`kernel_load_end_hook`、`launch_enter_hook`(当前 FlagTree 的 triton 3.6 满足)。缺 hook 时 IR / autotune 插件 fail-soft,不会把进程打挂;
- 插件通过 monkeypatch 挂钩上游内部结构,FlagGems/Triton 大版本升级后若行为异常,优先检查各插件 pytest_configure 输出的注册日志是否还正常打印。
- 无 GPU 的适配器自检:`python test_backends.py`。
+71 -1
View File
@@ -214,6 +214,57 @@ def _replay_run(original):
return runner
def _write_record_meta(record_dir: str) -> None:
"""Sidecar next to <op>.json. Replay refuses a different backend.
The config JSON itself stays `{kernel: {key: cfg}}` so existing REPLAY_FROM
directories keep working.
"""
try:
from backends import get_backend
backend_name = get_backend().name
except Exception:
backend_name = os.environ.get("ZL_BACKEND", "")
meta = {
"schema_version": 1,
"provider": "flaggems",
"backend": backend_name,
"op": _op_name(),
}
path = Path(record_dir) / "meta.json"
try:
path.parent.mkdir(parents=True, exist_ok=True)
path.write_text(json.dumps(meta, indent=2, ensure_ascii=False) + "\n")
except Exception as exc:
from _term_style import tag
print(f"{tag('[autotune-record-plugin]')} meta.json write failed: {exc}",
file=sys.stderr, flush=True)
def _replay_backend_ok(replay_dir: str) -> bool:
path = Path(replay_dir) / "meta.json"
if not path.is_file():
return True
try:
meta = json.loads(path.read_text())
except Exception:
return True
if not isinstance(meta, dict):
return True
recorded = str(meta.get("backend") or "").strip()
if not recorded:
return True
try:
from backends import get_backend
current = get_backend().name
except Exception:
current = os.environ.get("ZL_BACKEND", "").strip()
if recorded != current:
_emit_marker(f"backend_mismatch recorded={recorded} current={current}")
return False
return True
def _load_replay_dir(replay_dir: str) -> int:
path = _record_path(replay_dir)
if not path.is_file():
@@ -256,29 +307,48 @@ def _dump_record(record_dir: str) -> None:
def pytest_configure(config):
global _wrap
from _term_style import tag
from _run_identity import update_identity
record_dir = os.environ.get(_RECORD_DIR_ENV, "").strip()
replay_dir = os.environ.get(_REPLAY_DIR_ENV, "").strip()
if record_dir and replay_dir:
# Mutually exclusive: recording while replaying is pointless. Surface loudly.
print(f"{tag('[autotune-record-plugin]')} error: both RECORD_DIR and "
"REPLAY_DIR set; ignoring both (no-op)", file=sys.stderr, flush=True)
update_identity(autotune_mode="uncontrolled")
return
if not record_dir and not replay_dir:
update_identity(autotune_mode="uncontrolled")
return
try:
import triton.runtime.autotuner as _autotuner
except Exception as exc:
print(f"{tag('[autotune-record-plugin]')} Triton Autotuner unavailable "
f"({exc}); autotune_mode=uncontrolled", file=sys.stderr, flush=True)
update_identity(autotune_mode="uncontrolled")
return
if replay_dir and not _replay_backend_ok(replay_dir):
print(f"{tag('[autotune-record-plugin]')} refusing replay across backends",
file=sys.stderr, flush=True)
update_identity(autotune_mode="replay_fallback")
return
import triton.runtime.autotuner as _autotuner
_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.
_write_record_meta(record_dir)
atexit.register(_dump_record, record_dir)
update_identity(autotune_mode="record")
print(f"{tag('[autotune-record-plugin]')} recording autotune configs to "
f"{_record_path(record_dir)} (Autotuner)",
file=sys.stderr, flush=True)
else:
loaded = _load_replay_dir(replay_dir)
update_identity(autotune_mode="replay")
print(f"{tag('[autotune-record-plugin]')} replaying {loaded} recorded entries from "
f"{_record_path(replay_dir)} (Autotuner)",
file=sys.stderr, flush=True)
+15 -2
View File
@@ -1,7 +1,9 @@
"""pytest plugin: swap `triton.testing.do_bench` for cudagraph-based timing.
Loaded via `-p _cudagraph_plugin` in run_pytest.sh's PLUGINS list (comment that
line out to time with plain do_bench). Monkey-patches do_bench →
Invoked by `_timing_plugin` when the active backend declares
`supports_cuda_graph` (CUDA only). Do not put this module on the default `-p`
list for MUSA/ROCm/Ascend: those backends must start on plain do_bench, not
attempt capture-then-fallback. Monkey-patches do_bench →
do_bench_cudagraph (kernel-only latency) at session configure, after Triton is
loaded but before any benchmark calls it. Kernels that can't be graph-captured
fall back to plain do_bench and print a BENCHMARK_DIRECT_NO_CUDAGRAPH marker
@@ -167,6 +169,17 @@ def _patched_do_bench(fn, warmup=25, rep=100, grad_to_none=None,
def pytest_configure(config):
from _term_style import tag
try:
from backends import get_backend
if not get_backend().supports_cuda_graph:
print(f"{tag('[cudagraph-plugin]')} skipped: backend "
f"{get_backend().name!r} does not support CUDA Graph",
file=sys.stderr, flush=True)
return
except Exception as exc:
print(f"{tag('[cudagraph-plugin]')} backend probe failed ({exc}); "
"not patching do_bench", file=sys.stderr, flush=True)
return
_tt.do_bench = _patched_do_bench
print(f"{tag('[cudagraph-plugin]')} triton.testing.do_bench → do_bench_cudagraph",
file=sys.stderr, flush=True)
+5 -43
View File
@@ -1,45 +1,7 @@
"""Guard against FlagGems' device-detection hang (benchmark-side fix).
"""Backward-compatible alias of `_device_setup_plugin`.
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 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
device_finder takes its env fast-path (`_get_vendor_from_env`) and never
reaches the hang-prone subprocess probe. No change to the FlagGems repo.
Side effect: get_device_properties initializes the CUDA context very early in
the pytest process. Fine for the current single-process runs; revisit if
fork-based parallelism (e.g. pytest-xdist) is ever introduced.
Kept so an old `-p _device_guard_plugin` line still sets GEMS_VENDOR. New
plugin lists should load `_device_setup_plugin` instead.
"""
import os
import sys
_VENDOR_ENV_KEYS = ("GEMS_VENDOR", "FLAGGEMS_VENDOR", "GEMS_BACKEND", "FLAGGEMS_BACKEND")
def _guard_device_vendor():
# Respect an explicit choice if the user already set one.
if any(k in os.environ for k in _VENDOR_ENV_KEYS):
return
try:
import torch
if torch.cuda.is_available():
name = torch.cuda.get_device_properties(0).name.upper()
if "NVIDIA" in name:
os.environ["GEMS_VENDOR"] = "nvidia"
from _term_style import tag
print(f"{tag('[device-guard-plugin]')} set GEMS_VENDOR=nvidia "
"(skip flag_gems nvidia-smi subprocess probe, avoids import hang)",
file=sys.stderr, flush=True)
except Exception as e: # pragma: no cover
print(f"[device-guard-plugin] torch vendor probe failed ({e}); "
"leaving detection to flag_gems",
file=sys.stderr, flush=True)
_guard_device_vendor()
from _device_setup_plugin import * # noqa: F401,F403
from _device_setup_plugin import pytest_configure # noqa: F401
+69
View File
@@ -0,0 +1,69 @@
"""Set FlagGems vendor + record backend identity before flag_gems is imported.
Replaces the NVIDIA-only `_device_guard_plugin`. Setting GEMS_VENDOR from the
zl_bench backend name skips FlagGems' hang-prone nvidia-smi subprocess probe
on CUDA, and maps musa/rocm/ascend onto the matching FlagGems vendor.
Import-time side effect is intentional: FlagGems reads GEMS_VENDOR while the
benchmark module is being collected. Device name / arch are filled later in
pytest_configure so import does not initialize a device context.
"""
from __future__ import annotations
import os
import sys
from backends import get_backend
from providers.flaggems import apply_vendor_env
from _run_identity import update_identity
def _setup_at_import() -> None:
try:
backend = get_backend()
except Exception as exc:
print(f"[device-setup-plugin] cannot resolve backend ({exc}); "
"leaving GEMS_VENDOR to flag_gems",
file=sys.stderr, flush=True)
return
vendor = apply_vendor_env(backend)
os.environ["ZL_BACKEND"] = backend.name
device = os.environ.get("ZL_DEVICE", "").strip()
if not device:
for key in backend.visibility_env_keys:
raw = os.environ.get(key, "").strip()
if raw:
device = raw.split(",")[0].strip()
break
update_identity(
backend=backend.name,
device=device,
flaggems_vendor=vendor or "",
visibility_env=backend.visibility_env_keys[0] if backend.visibility_env_keys else "",
)
if vendor:
from _term_style import tag
print(f"{tag('[device-setup-plugin]')} backend={backend.name} "
f"GEMS_VENDOR={vendor}",
file=sys.stderr, flush=True)
def pytest_configure(config):
try:
backend = get_backend()
ident = backend.device_identity()
update_identity(
backend=backend.name,
device_name=ident.get("device_name", ""),
device_arch=ident.get("device_arch", ""),
timer_api=ident.get("timer_api", ""),
flaggems_vendor=ident.get("flaggems_vendor", ""),
)
except Exception as exc:
from _term_style import tag
print(f"{tag('[device-setup-plugin]')} identity probe failed ({exc})",
file=sys.stderr, flush=True)
_setup_at_import()
+21 -5
View File
@@ -402,17 +402,33 @@ def _dump_ttgir(cache_dir: str, dump_dir: str) -> None:
def pytest_configure(config):
global _dump_enabled
import triton
import triton.runtime.autotuner as _autotuner
from _term_style import tag
try:
import triton
import triton.runtime.autotuner as _autotuner
except Exception as exc:
print(f"{tag('[ir-meta-plugin]')} Triton knobs unavailable ({exc}); "
"IR collection disabled", file=sys.stderr, flush=True)
return
if not hasattr(triton, "knobs"):
print(f"{tag('[ir-meta-plugin]')} triton.knobs missing; "
"IR collection disabled", file=sys.stderr, flush=True)
return
prev = triton.knobs.compilation.listener
try:
compilation = triton.knobs.compilation
prev = compilation.listener
except Exception as exc:
print(f"{tag('[ir-meta-plugin]')} compilation listener unavailable "
f"({exc}); IR collection disabled", file=sys.stderr, flush=True)
return
if prev is None:
triton.knobs.compilation.listener = _compile_listener
compilation.listener = _compile_listener
else:
def chained(**kwargs):
prev(**kwargs)
_compile_listener(**kwargs)
triton.knobs.compilation.listener = chained
compilation.listener = chained
dump_dir = os.environ.get(_DUMP_DIR_ENV, "").strip()
if dump_dir:
+12
View File
@@ -105,6 +105,18 @@ def _warmup_load_backend():
def pytest_configure(config):
from _term_style import tag
try:
from backends import get_backend
if not get_backend().supports_mm_cluster_fix:
print(f"{tag('[mm-cluster-fix-plugin]')} skipped: backend "
f"{get_backend().name!r} is not the CUDA Hopper path",
file=sys.stderr, flush=True)
return
except Exception as exc:
print(f"{tag('[mm-cluster-fix-plugin]')} backend probe failed ({exc}); "
"leaving cluster_remote_mm_scenario unpatched",
file=sys.stderr, flush=True)
return
_patch_loaded()
_warmup_load_backend()
_patch_loaded()
+171 -30
View File
@@ -33,7 +33,7 @@ import sys
import tempfile
import time
from pathlib import Path
from typing import Any, Dict, List, Tuple
from typing import Any, Dict, List, Optional, Tuple
import pytest
@@ -60,8 +60,8 @@ def _requested_gpus() -> int:
if want < 2:
return 0
try:
import torch
have = torch.cuda.device_count()
from backends import get_backend
have = get_backend().probe_device_count()
except Exception:
return 0
return max(0, min(want, have))
@@ -70,30 +70,132 @@ def _requested_gpus() -> int:
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.
Must respect inherited visibility env for the *active* backend. 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()]
from backends import get_backend
backend = get_backend()
for key in backend.visibility_env_keys:
raw = os.environ.get(key, "").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())]
n = backend.probe_device_count()
return [str(i) for i in range(n)] if n else []
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."""
def _attr_name(node: Any) -> str:
"""Last component of a Name/Attribute node: `base.BlasBenchmark` -> `BlasBenchmark`."""
import ast
if isinstance(node, ast.Name):
return node.id
if isinstance(node, ast.Attribute):
return node.attr
return ""
def _class_bases(py_path: Path) -> Dict[str, List[str]]:
"""{class name: [base class names]} parsed statically from one module."""
import ast
try:
tree = ast.parse(py_path.read_text(errors="replace"))
except Exception:
return {}
return {
node.name: [n for n in (_attr_name(b) for b in node.bases) if n]
for node in ast.walk(tree)
if isinstance(node, ast.ClassDef)
}
def _class_chain(test_py: Path, op: str) -> List[str]:
"""Benchmark class names to try for `op`, nearest ancestor first.
FlagGems' `Benchmark.set_shapes` looks up `op_name` and, failing that, walks
`type(self).__mro__` trying each class name as a yaml key -- which is how
BLAS ops reach the shared `BlasBenchmark` entry in core_shapes.yaml. This
plugin runs at pytest_configure, before any Benchmark exists, so the real
MRO is unavailable; the class is recovered statically instead, from the
`op_name=` keyword of the constructor call inside the test module, then its
bases from that module plus benchmark/base.py.
Breadth-first over the inheritance graph, not a true C3 linearisation: the
benchmark classes are single-inheritance in practice, and a miss only costs
the sharding (see the caller's skip path), never correctness.
"""
import ast
try:
tree = ast.parse(test_py.read_text(errors="replace"))
except Exception:
return []
start = ""
for node in ast.walk(tree):
if not isinstance(node, ast.Call):
continue
for kw in node.keywords:
if (kw.arg == "op_name" and isinstance(kw.value, ast.Constant)
and kw.value.value == op):
start = _attr_name(node.func)
break
if start:
break
if not start:
return []
graph = _class_bases(test_py)
for name, bases in _class_bases(test_py.parent / "base.py").items():
graph.setdefault(name, bases)
chain: List[str] = []
seen: set = set()
queue = [start]
while queue:
name = queue.pop(0)
if not name or name in seen:
continue
seen.add(name)
chain.append(name)
queue.extend(graph.get(name, []))
return chain
def _test_module_path(argv: List[str]) -> Optional[Path]:
"""The benchmark module under test, from pytest's own argv."""
for a in argv:
if a.startswith("-"):
continue
p = Path(a.split("::", 1)[0])
if p.suffix == ".py" and p.is_file():
return p
return None
def _read_shape_yaml(path: str, op: str, test_py: Optional[Path] = None
) -> Tuple[List[Any], Dict[str, Any], str]:
"""Shapes for this op, the sibling keys to carry into each shard, and the
yaml key they came from ("" when nothing matched).
Tries `op` first, then the benchmark class chain, mirroring FlagGems. Without
the class fallback every BLAS op (mm/addmm/bmm/baddbmm, whose shapes live
under `BlasBenchmark`) resolved to zero shapes and sharding silently turned
itself off.
"""
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
keys = [op]
if not isinstance(doc.get(op), dict) and test_py is not None:
keys += _class_chain(test_py, op)
for key in keys:
entry = doc.get(key)
if isinstance(entry, dict) and entry.get("shapes"):
extra = {k: v for k, v in entry.items() if k != "shapes"}
return list(entry["shapes"]), extra, key
return [], {}, ""
def _shard(shapes: List[Any], n: int) -> List[List[Any]]:
@@ -111,8 +213,8 @@ def _spawn(shards: List[List[Any]], op: str, extra: Dict[str, Any],
) -> 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.
Subprocesses rather than pytest-xdist: `_device_setup_plugin` may initialize
a device context at import time, which does not survive fork-based parallelism.
"""
import yaml
procs: List[Tuple[int, subprocess.Popen, Any]] = []
@@ -143,7 +245,8 @@ def _spawn(shards: List[List[Any]], op: str, extra: Dict[str, Any],
child_argv += ["--shape_file", str(w / "shapes.yaml")]
env = dict(os.environ)
env["CUDA_VISIBLE_DEVICES"] = devices[i]
from backends import get_backend
env = get_backend().apply_visibility(env, 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
@@ -222,6 +325,23 @@ def pytest_configure(config):
return
op = _op_name()
# Only --level core is shardable. At comprehensive, FlagGems' set_shapes()
# merges set_more_shapes() on top of whatever the yaml holds, and that
# addition is arbitrary per-subclass Python (hardcoded tuples plus
# models_shapes/*.yaml) which this plugin cannot see. Sharding the yaml list
# alone would sweep a strict subset -- measured on mm: 5 of 29 shapes, 90
# AUTOTUNE_REPLAY_FALLBACK reason=key_missing -- and every uncovered key
# then self-tunes while the run reports itself as replaying. Refusing costs
# wall time; sharding a subset costs the same-config premise.
level = str(getattr(config.option, "level", "") or "").strip().lower()
if level and level != "core":
print(f"{tag('[parallel-warmup-plugin]')} --level {level} merges shapes "
"this plugin cannot enumerate (set_more_shapes); sharding a subset "
"would break the same-config premise, so sweeping serially instead",
flush=True)
return
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; "
@@ -232,15 +352,24 @@ def pytest_configure(config):
"already pinned, nothing to sweep, skipping", flush=True)
return
test_py = _test_module_path(list(config.invocation_params.args))
try:
shapes, extra = _read_shape_yaml(shape_file, op)
shapes, extra, key = _read_shape_yaml(shape_file, op, test_py)
except Exception as exc:
print(f"{tag('[parallel-warmup-plugin]')} cannot read shapes "
f"({type(exc).__name__}: {exc}); skipping", flush=True)
return
if not key:
# Distinct from "found a short list": nothing in the yaml describes this
# op at all, so the shard files would be built from an empty set.
print(f"{tag('[parallel-warmup-plugin]')} no yaml entry for {op!r} "
f"(tried: {', '.join([op] + _class_chain(test_py, op)) if test_py else op}); "
"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)
print(f"{tag('[parallel-warmup-plugin]')} only {len(shapes)} shape(s) "
f"under yaml key {key!r}; sharding would not pay off, skipping",
flush=True)
return
devices = _visible_devices()
@@ -251,8 +380,9 @@ def pytest_configure(config):
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 "
print(f"{tag('[parallel-warmup-plugin]')} sweeping {len(shapes)} shapes "
f"(yaml key {key!r}) on {len(shards)} GPUs "
f"(shards: {[len(s) for s in shards]}); this phase's "
"latency is discarded, only configs are kept", flush=True)
keep_scratch = False
@@ -260,9 +390,23 @@ def pytest_configure(config):
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 bad:
# A dead shard means its shapes were never swept, so publishing the
# merge would hand the measurement a config set with holes; every
# hole then silently self-tunes under a run that claims to replay.
# Discard the whole sweep rather than publish a partial one.
keep_scratch = True
print(f"{tag('[parallel-warmup-plugin]')} shard(s) {bad} exited "
f"non-zero after {elapsed:.0f}s; their shapes were never swept, "
"so the merged configs would have holes. Discarding this sweep "
f"and falling back to normal serial autotune (worker logs kept "
f"in {scratch})", flush=True)
return
merged, total, conflicts = _merge(recs, op)
if not total:
keep_scratch = True
print(f"{tag('[parallel-warmup-plugin]')} no configs recovered in "
@@ -287,9 +431,6 @@ def pytest_configure(config):
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 "
+91
View File
@@ -0,0 +1,91 @@
"""Merge-update `$OUT_DIR/run_identity.json` so plugins can record backend facts.
Shell writes a stub before pytest; plugins add device_name / timer / vendor
without clobbering earlier keys. Missing OUT_DIR is a silent no-op so unit
tests and ad-hoc pytest invocations do not fail.
"""
from __future__ import annotations
import argparse
import json
import os
from pathlib import Path
from typing import Any, Dict, Optional
IDENTITY_NAME = "run_identity.json"
def identity_path(out_dir: Optional[str] = None) -> Optional[Path]:
root = out_dir or os.environ.get("OUT_DIR", "").strip()
if not root:
return None
return Path(root) / IDENTITY_NAME
def load_identity(out_dir: Optional[str] = None) -> Dict[str, Any]:
path = identity_path(out_dir)
if path is None or not path.is_file():
return {}
try:
data = json.loads(path.read_text(encoding="utf-8"))
except Exception:
return {}
return data if isinstance(data, dict) else {}
def update_identity(out_dir: Optional[str] = None, **fields: Any) -> Dict[str, Any]:
path = identity_path(out_dir)
if path is None:
return dict(fields)
path.parent.mkdir(parents=True, exist_ok=True)
current = load_identity(out_dir)
for key, value in fields.items():
if value is None:
continue
current[key] = value
tmp = path.with_name(path.name + ".tmp")
tmp.write_text(json.dumps(current, indent=2, ensure_ascii=False) + "\n",
encoding="utf-8")
os.replace(tmp, path)
return current
def init_identity(out_dir: str, **fields: Any) -> Dict[str, Any]:
payload = {"schema_version": 1, "provider": "flaggems"}
payload.update(fields)
os.environ["OUT_DIR"] = str(out_dir)
return update_identity(out_dir, **payload)
def main(argv=None) -> int:
p = argparse.ArgumentParser()
p.add_argument("out_dir")
p.add_argument("--backend", default=os.environ.get("ZL_BACKEND", "cuda"))
p.add_argument("--op", default=os.environ.get("OP", ""))
p.add_argument("--op-file", default=os.environ.get("OP_FILE", ""))
p.add_argument("--flaggems-dir", default=os.environ.get("FLAGGEMS_DIR", ""))
p.add_argument("--autotune-mode", default="")
p.add_argument("--device", default=os.environ.get("ZL_DEVICE", ""))
p.add_argument("--test-file", default=os.environ.get("TEST_FILE", ""))
p.add_argument("--shape-mode", default=os.environ.get("SHAPE_MODE", ""))
p.add_argument("--bench-level", default=os.environ.get("BENCH_LEVEL", "core"))
args = p.parse_args(argv)
init_identity(
args.out_dir,
backend=args.backend,
op=args.op,
op_file=args.op_file,
flaggems_dir=args.flaggems_dir,
autotune_mode=args.autotune_mode,
device=args.device,
test_file=args.test_file,
shape_mode=args.shape_mode,
bench_level=args.bench_level,
)
return 0
if __name__ == "__main__":
raise SystemExit(main())
+47
View File
@@ -0,0 +1,47 @@
"""Select the timing path from backend capabilities, before any shape runs.
CUDA keeps the existing `_cudagraph_plugin` behaviour (graph capture, with
per-shape fallback to plain do_bench). Other backends do not attempt CUDA
Graph at all: they leave `triton.testing.do_bench` in place and record
`timer=direct`. Trying capture-then-fallback on a non-CUDA device would
produce FATAL_DEVICE noise and mixed-caliber rows.
"""
from __future__ import annotations
import os
import sys
from backends import get_backend
from _run_identity import update_identity
def _graph_requested(backend) -> bool:
raw = os.environ.get("ZL_ENABLE_CUDAGRAPH", "auto").strip().lower()
if raw in ("0", "false", "no", "off", "never", "direct"):
return False
if raw in ("1", "true", "yes", "on", "always"):
return bool(backend.supports_cuda_graph)
return bool(backend.supports_cuda_graph)
def pytest_configure(config):
from _term_style import tag
backend = get_backend()
raw = os.environ.get("ZL_ENABLE_CUDAGRAPH", "auto").strip().lower()
force_on = raw in ("1", "true", "yes", "on", "always")
if force_on and not backend.supports_cuda_graph:
print(f"{tag('[timing-plugin]')} ZL_ENABLE_CUDAGRAPH=1 ignored on "
f"backend={backend.name} (no CUDA Graph); using timer=direct",
file=sys.stderr, flush=True)
if _graph_requested(backend):
import _cudagraph_plugin as _cg
_cg.pytest_configure(config)
update_identity(timer="cuda_graph", graph_capture="attempted")
print(f"{tag('[timing-plugin]')} timer=cuda_graph backend={backend.name}",
file=sys.stderr, flush=True)
return
update_identity(timer="direct", graph_capture="disabled")
print(f"{tag('[timing-plugin]')} timer=direct backend={backend.name} "
"(upstream do_bench; graph capture not selected)",
file=sys.stderr, flush=True)
+76
View File
@@ -0,0 +1,76 @@
"""Registry of zl_bench device backends.
Aliases map common vendor spellings onto the canonical backend name. Unknown
names fail closed rather than defaulting to CUDA: a silent CUDA fallback would
load the cudagraph plugin on the wrong device.
"""
from __future__ import annotations
import os
from typing import Dict, Optional
from .ascend import AscendBackend
from .base import DeviceBackend, parse_device_list
from .cuda import CudaBackend
from .musa import MusaBackend
from .rocm import RocmBackend
ALIASES: Dict[str, str] = {
"nvidia": "cuda",
"nv": "cuda",
"amd": "rocm",
"hip": "rocm",
"npu": "ascend",
"huawei": "ascend",
}
_INSTANCES: Dict[str, DeviceBackend] = {
"cuda": CudaBackend(),
"musa": MusaBackend(),
"rocm": RocmBackend(),
"ascend": AscendBackend(),
}
class UnknownBackendError(ValueError):
pass
def known_backends() -> tuple[str, ...]:
return tuple(sorted(_INSTANCES))
def normalize_backend_name(name: str) -> str:
key = (name or "").strip().lower()
if not key:
return "cuda"
return ALIASES.get(key, key)
def get_backend(name: Optional[str] = None) -> DeviceBackend:
raw = name if name is not None else os.environ.get("ZL_BACKEND", "cuda")
key = normalize_backend_name(raw)
backend = _INSTANCES.get(key)
if backend is None:
raise UnknownBackendError(
f"unknown backend {raw!r}; known: {', '.join(known_backends())} "
f"(aliases: {', '.join(sorted(ALIASES))})"
)
return backend
def enumerate_devices(backend: Optional[DeviceBackend] = None,
explicit: Optional[str] = None) -> list[str]:
return (backend or get_backend()).enumerate_devices(explicit)
__all__ = [
"DeviceBackend",
"UnknownBackendError",
"enumerate_devices",
"get_backend",
"known_backends",
"normalize_backend_name",
"parse_device_list",
]
+17
View File
@@ -0,0 +1,17 @@
from __future__ import annotations
from .base import DeviceBackend
class AscendBackend(DeviceBackend):
name = "ascend"
visibility_env_keys = (
"ASCEND_RT_VISIBLE_DEVICES",
"ASCEND_VISIBLE_DEVICES",
)
torch_namespace = "npu"
flaggems_vendor = "ascend"
supports_cuda_graph = False
supports_mm_cluster_fix = False
supports_triton_hooks = True
list_command = None
+148
View File
@@ -0,0 +1,148 @@
"""Device backend contract for zl_bench screening runs.
This is the orchestration-layer adapter: device enumeration, process visibility
env, and capability flags that decide which pytest plugins are safe to load.
It does not implement kernel timing receipts or compiler-artifact identity;
those stay with the caller (FlagGems pytest for screening, a separate harness
for qualification).
"""
from __future__ import annotations
import os
import subprocess
from typing import Dict, Iterable, List, Optional, Sequence, Tuple
# Every vendor visibility variable zl_bench knows about. apply_visibility()
# clears the whole set before writing the active backend's keys so a parent
# CUDA_VISIBLE_DEVICES cannot leak onto a MUSA/Ascend child.
ALL_VISIBILITY_KEYS: Tuple[str, ...] = (
"CUDA_VISIBLE_DEVICES",
"MUSA_VISIBLE_DEVICES",
"HIP_VISIBLE_DEVICES",
"ROCR_VISIBLE_DEVICES",
"ASCEND_RT_VISIBLE_DEVICES",
"ASCEND_VISIBLE_DEVICES",
)
def parse_device_list(raw: Optional[str]) -> List[str]:
if not raw:
return []
return [tok.strip() for tok in raw.split(",") if tok.strip()]
class DeviceBackend:
"""One hardware runtime (CUDA, MUSA, ROCm, Ascend, ...).
Subclasses only fill class attributes and optional probe hooks. Callers
should not branch on ``backend.name`` for visibility or plugin policy;
they should read these attributes / methods.
"""
name: str = ""
# First key is the canonical one written into run_identity.json.
visibility_env_keys: Tuple[str, ...] = ()
# torch.<namespace> used for optional device_count / identity probes.
torch_namespace: str = ""
# FlagGems GEMS_VENDOR value. Lives on the backend object because zl_bench
# still only drives FlagGems, but the mapping is owned by providers.flaggems
# and copied here so DeviceBackend.prepare_environment() can stay generic.
flaggems_vendor: str = ""
supports_cuda_graph: bool = False
supports_mm_cluster_fix: bool = False
supports_triton_hooks: bool = True
# Optional listing command; (argv, line_prefix). prefix="" means count
# non-empty lines. Must be a real vendor CLI; do not invent flags.
list_command: Optional[Tuple[Sequence[str], str]] = None
def visibility_env(self, device: str) -> Dict[str, str]:
return {key: device for key in self.visibility_env_keys}
def apply_visibility(self, env: Dict[str, str], device: str) -> Dict[str, str]:
out = dict(env)
for key in ALL_VISIBILITY_KEYS:
out.pop(key, None)
out.update(self.visibility_env(device))
return out
def enumerate_devices(self, explicit: Optional[str] = None) -> List[str]:
listed = parse_device_list(explicit)
if listed:
return listed
for key in self.visibility_env_keys:
listed = parse_device_list(os.environ.get(key, ""))
if listed:
return listed
n = self.probe_device_count()
if n > 0:
return [str(i) for i in range(n)]
return ["0"]
def probe_device_count(self) -> int:
n = self._probe_list_command()
if n > 0:
return n
return self._probe_torch_count()
def _probe_list_command(self) -> int:
if not self.list_command:
return 0
argv, prefix = self.list_command
try:
proc = subprocess.run(
list(argv), text=True, capture_output=True, timeout=30,
)
except Exception:
return 0
if proc.returncode != 0:
return 0
lines = [ln for ln in proc.stdout.splitlines() if ln.strip()]
if prefix:
lines = [ln for ln in lines if ln.startswith(prefix)]
return len(lines)
def _probe_torch_count(self) -> int:
if not self.torch_namespace:
return 0
try:
import torch
api = getattr(torch, self.torch_namespace, None)
if api is None or not hasattr(api, "device_count"):
return 0
return int(api.device_count())
except Exception:
return 0
def device_identity(self) -> Dict[str, str]:
"""Best-effort identity for screening logs. Empty strings are allowed."""
info = {
"backend": self.name,
"device_name": "",
"device_arch": "",
"timer_api": f"torch.{self.torch_namespace}" if self.torch_namespace else "",
"flaggems_vendor": self.flaggems_vendor,
}
if not self.torch_namespace:
return info
try:
import torch
api = getattr(torch, self.torch_namespace, None)
if api is None:
return info
count = int(api.device_count()) if hasattr(api, "device_count") else 0
if count <= 0:
return info
index = int(api.current_device()) if hasattr(api, "current_device") else 0
if hasattr(api, "get_device_name"):
info["device_name"] = str(api.get_device_name(index))
if hasattr(api, "get_device_capability"):
cap = api.get_device_capability(index)
info["device_arch"] = f"{cap[0]}.{cap[1]}"
elif hasattr(api, "get_device_properties"):
props = api.get_device_properties(index)
arch = getattr(props, "name", "") or getattr(props, "arch", "")
info["device_arch"] = str(arch)
except Exception:
return info
return info
+14
View File
@@ -0,0 +1,14 @@
from __future__ import annotations
from .base import DeviceBackend
class CudaBackend(DeviceBackend):
name = "cuda"
visibility_env_keys = ("CUDA_VISIBLE_DEVICES",)
torch_namespace = "cuda"
flaggems_vendor = "nvidia"
supports_cuda_graph = True
supports_mm_cluster_fix = True
supports_triton_hooks = True
list_command = (("nvidia-smi", "--list-gpus"), "GPU ")
+16
View File
@@ -0,0 +1,16 @@
from __future__ import annotations
from .base import DeviceBackend
class MusaBackend(DeviceBackend):
name = "musa"
visibility_env_keys = ("MUSA_VISIBLE_DEVICES",)
torch_namespace = "musa"
flaggems_vendor = "musa"
supports_cuda_graph = False
supports_mm_cluster_fix = False
supports_triton_hooks = True
# Device count comes from torch.musa or MUSA_VISIBLE_DEVICES. Do not guess
# a vendor CLI here; confirm the real listing tool before adding it.
list_command = None
+16
View File
@@ -0,0 +1,16 @@
from __future__ import annotations
from .base import DeviceBackend
class RocmBackend(DeviceBackend):
name = "rocm"
# HIP runtimes honor both; write both so a parent CUDA/HIP mix cannot
# leave the child looking at the wrong cards.
visibility_env_keys = ("HIP_VISIBLE_DEVICES", "ROCR_VISIBLE_DEVICES")
torch_namespace = "cuda" # ROCm PyTorch still exposes torch.cuda
flaggems_vendor = "amd"
supports_cuda_graph = False
supports_mm_cluster_fix = False
supports_triton_hooks = True
list_command = None
+24 -3
View File
@@ -39,12 +39,31 @@ 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",
"repo", "op", "backend", "device", "timer",
"fg_op_name", "dtype", "case_id", "row_status",
"torch_lat", "gems_lat", "speedup", "extra_metrics", "size_detail",
"op_status", "notes",
]
def _identity_from_status(status_row: Dict[str, str], op_dir: Path) -> Dict[str, str]:
"""Backend/device/timer from ops_status.csv, falling back to run_identity.json."""
backend = (status_row.get("backend") or "").strip()
device = (status_row.get("device") or status_row.get("gpu") or "").strip()
timer = (status_row.get("timer") or "").strip()
ident_path = op_dir / "run_identity.json"
if ident_path.is_file() and not (backend and device and timer):
try:
from _run_identity import load_identity
ident = load_identity(str(op_dir))
backend = backend or str(ident.get("backend") or "")
device = device or str(ident.get("device") or ident.get("gpu") or "")
timer = timer or str(ident.get("timer") or "")
except Exception:
pass
return {"backend": backend, "device": device, "timer": timer}
def parse_run_log(text: str) -> List[Dict[str, str]]:
"""把 run.log 里的全部结果表解析为行列表(跨多个 dtype 表)。"""
rows: List[Dict[str, str]] = []
@@ -133,10 +152,11 @@ def summarize(batch_dir: Path, out_path: Optional[Path] = None) -> Path:
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
ident = _identity_from_status(status_row, op_dir)
table_rows = parse_run_log(text)
if not table_rows:
all_rows.append({
"repo": repo, "op": op, "fg_op_name": "", "dtype": "",
"repo": repo, "op": op, **ident, "fg_op_name": "", "dtype": "",
"case_id": "", "row_status": "", "torch_lat": "",
"gems_lat": "", "speedup": "", "extra_metrics": "",
"size_detail": "",
@@ -150,7 +170,8 @@ def summarize(batch_dir: Path, out_path: Optional[Path] = None) -> Path:
idx = case_counter.get(key, 0)
case_counter[key] = idx + 1
all_rows.append({
"repo": repo, "op": op, "case_id": f"case_{idx:03d}",
"repo": repo, "op": op, **ident,
"case_id": f"case_{idx:03d}",
"op_status": status_row.get("status", ""),
"notes": status_row.get("notes", ""),
**row,
+2
View File
@@ -41,6 +41,8 @@ import yaml
OPS_DIR = Path(__file__).resolve().parent
# Inventory is FlagGems-logical (repo/op/op_file), not per-hardware.
# One case is reused across backends; do not duplicate CSVs per vendor.
DEFAULT_REPOS = [
("main", Path("/workspace/dev/FlagGems")),
("vllm", Path("/workspace/dev/FlagGems-vllm")),
+1016 -1016
View File
File diff suppressed because it is too large. Load diff
+109 -109
View File
@@ -1,109 +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
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
+107
View File
@@ -0,0 +1,107 @@
"""Decide which pytest plugins a zl_bench run should load.
CUDA-specific plugins (cudagraph patch, Hopper mm cluster workaround) are not
placed on the default list; `_timing_plugin` enables graph capture only when
the backend declares `supports_cuda_graph`. IR / autotune plugins stay on for
Triton-compatible compilers and can be turned off with env flags.
"""
from __future__ import annotations
import argparse
import os
from typing import List, Optional, Sequence
from backends import get_backend
CORE_PLUGINS: tuple[str, ...] = (
"_device_setup_plugin",
"_parallel_warmup_plugin",
"_seed_plugin",
"_shape_inject_plugin",
"_shape_iter_inject_plugin",
"_bespoke_shape_plugin",
)
ALWAYS_AFTER: tuple[str, ...] = (
"_pretty_report_plugin",
)
def _env_flag(name: str, default: str = "1") -> bool:
raw = os.environ.get(name, default).strip().lower()
return raw not in ("0", "false", "no", "off", "")
def plan_plugins(
backend_name: Optional[str] = None,
extra: Optional[Sequence[str]] = None,
enable_autotune: Optional[bool] = None,
enable_ir: Optional[bool] = None,
enable_mm_cluster: Optional[bool] = None,
) -> List[str]:
backend = get_backend(backend_name)
plugins: List[str] = list(CORE_PLUGINS)
autotune = (
backend.supports_triton_hooks
if enable_autotune is None
else enable_autotune
)
if enable_autotune is None:
autotune = autotune and _env_flag("ZL_ENABLE_AUTOTUNE", "1")
if autotune:
plugins.append("_autotune_record_plugin")
plugins.append("_timing_plugin")
plugins.extend(ALWAYS_AFTER)
ir = backend.supports_triton_hooks if enable_ir is None else enable_ir
if enable_ir is None:
ir = ir and _env_flag("ZL_ENABLE_IR", "1")
if ir:
plugins.append("_ir_meta_plugin")
mm = enable_mm_cluster
if mm is None:
mm = backend.supports_mm_cluster_fix and _env_flag("ZL_ENABLE_MM_CLUSTER", "0")
elif mm:
mm = backend.supports_mm_cluster_fix
if mm:
plugins.append("_mm_cluster_fix_plugin")
if extra is None:
extras = [p for p in os.environ.get("ZL_EXTRA_PLUGINS", "").split() if p]
else:
extras = [p for p in extra if p]
for name in extras:
if name not in plugins:
plugins.append(name)
return plugins
def pytest_args(plugins: Optional[Sequence[str]] = None, **kwargs) -> List[str]:
names = list(plugins) if plugins is not None else plan_plugins(**kwargs)
args: List[str] = []
for name in names:
args.extend(("-p", name))
return args
def main(argv: Optional[Sequence[str]] = None) -> int:
p = argparse.ArgumentParser(description="Print zl_bench pytest -p args")
p.add_argument("--backend", default=None)
p.add_argument("--pytest-args", action="store_true",
help="print `-p name` tokens on one line")
p.add_argument("--one-per-line", action="store_true")
args = p.parse_args(argv)
names = plan_plugins(backend_name=args.backend)
if args.one_per_line:
print("\n".join(names))
else:
print(" ".join(pytest_args(names)))
return 0
if __name__ == "__main__":
raise SystemExit(main())
+1
View File
@@ -0,0 +1 @@
"""Benchmark providers. Phase 1 only implements FlagGems."""
+219
View File
@@ -0,0 +1,219 @@
"""FlagGems-specific mapping kept out of the generic scheduler.
GEMS_VENDOR values are FlagGems' vendor names, not zl_bench backend names.
cuda → nvidia is the important one: FlagGems' device_finder otherwise shells
out to nvidia-smi without a timeout.
"""
from __future__ import annotations
import argparse
import os
import shlex
import sys
from pathlib import Path
from typing import Dict, Optional
_ROOT = Path(__file__).resolve().parent.parent
if str(_ROOT) not in sys.path:
sys.path.insert(0, str(_ROOT))
from backends.base import DeviceBackend
# Live checkouts used by the NV FlagGems product path. Order is the fallback
# sequence when the configured root does not contain the requested test file.
KNOWN_FLAGGEMS_ROOTS = (
"/workspace/FlagGems",
"/workspace/dev/FlagGems",
)
VENDOR_ENV_KEYS = (
"GEMS_VENDOR", "FLAGGEMS_VENDOR", "GEMS_BACKEND", "FLAGGEMS_BACKEND",
)
VENDOR_BY_BACKEND: Dict[str, str] = {
"cuda": "nvidia",
"musa": "musa",
"rocm": "amd",
"ascend": "ascend",
}
def flaggems_vendor(backend: DeviceBackend) -> str:
return backend.flaggems_vendor or VENDOR_BY_BACKEND[backend.name]
def apply_vendor_env(backend: DeviceBackend) -> Optional[str]:
"""Set GEMS_VENDOR unless the user already chose one. Returns the value."""
for key in VENDOR_ENV_KEYS:
if key in os.environ and os.environ[key].strip():
return os.environ[key].strip()
vendor = flaggems_vendor(backend)
if vendor:
os.environ["GEMS_VENDOR"] = vendor
return vendor or None
def make_test_node(flaggems_dir: str, op_file: str, op: str) -> str:
root = flaggems_dir.replace("\\", "/").rstrip("/")
return f"{root}/benchmark/test_{op_file}.py::test_{op}"
# Back-compat alias. Do not call this from emit_run_env: the parameter of the
# same name would shadow it and raise TypeError.
test_node = make_test_node
def parse_test_node(node: str) -> Dict[str, Optional[str]]:
"""Parse a FlagGems pytest node into op / op_file / optional repo root.
Example:
/workspace/FlagGems/benchmark/test_fused_marlin_moe.py::test_fused_marlin_moe_mxfp4
"""
raw = (node or "").strip()
if "::" not in raw:
raise ValueError(f"pytest node must contain '::': {node!r}")
file_part, func = raw.rsplit("::", 1)
if not func.startswith("test_"):
raise ValueError(f"test function must start with 'test_': {func!r}")
stem = Path(file_part.replace("\\", "/").rsplit("/", 1)[-1]).stem
if not stem.startswith("test_"):
raise ValueError(f"benchmark filename must start with 'test_': {file_part!r}")
posix_file = file_part.replace("\\", "/")
parent = posix_file.rsplit("/", 1)[0] if "/" in posix_file else ""
parent_name = parent.rsplit("/", 1)[-1] if parent else ""
flaggems_dir = (
parent.rsplit("/", 1)[0] if parent_name == "benchmark" and "/" in parent else None
)
return {
"test_file": f"{file_part}::{func}",
"file": file_part,
"op": func[len("test_"):],
"op_file": stem[len("test_"):],
"flaggems_dir": flaggems_dir,
}
def _test_py(flaggems_dir: str, op_file: str) -> Path:
return Path(flaggems_dir) / "benchmark" / f"test_{op_file}.py"
def resolve_flaggems_dir(configured: str, op_file: str) -> str:
"""Return the first root that actually contains benchmark/test_<op_file>.py."""
candidates = []
for item in (configured, *KNOWN_FLAGGEMS_ROOTS):
if item and item not in candidates:
candidates.append(item)
for root in candidates:
if _test_py(root, op_file).is_file():
return root
return configured
def resolve_existing_test_file(test_file: str, op_file: str) -> str:
"""If the node path is missing, retry the same relative path under known roots."""
if "::" not in test_file:
return test_file
file_part, func = test_file.rsplit("::", 1)
path = Path(file_part)
if path.is_file():
return test_file
for root in KNOWN_FLAGGEMS_ROOTS:
candidate = _test_py(root, op_file)
if candidate.is_file():
return f"{candidate.as_posix()}::{func}"
return test_file
def emit_run_env(
test_node: str = "",
op: str = "",
op_file: str = "",
flaggems_dir: str = "",
shape_file: str = "",
shape_mode: str = "",
) -> Dict[str, str]:
"""Resolve OP / OP_FILE / TEST_FILE / SHAPE_MODE for run_pytest.sh."""
parsed = parse_test_node(test_node) if test_node.strip() else None
if parsed:
op = parsed["op"] or op
op_file = parsed["op_file"] or op_file
if parsed["flaggems_dir"]:
flaggems_dir = parsed["flaggems_dir"]
test_file = parsed["test_file"] or ""
else:
flaggems_dir = resolve_flaggems_dir(flaggems_dir, op_file)
test_file = make_test_node(flaggems_dir, op_file, op)
flaggems_dir = resolve_flaggems_dir(flaggems_dir, op_file)
test_file = resolve_existing_test_file(test_file, op_file)
shape_file = shape_file.strip()
shape_mode = shape_mode.strip().lower()
if shape_file in ("-", "upstream", "none"):
shape_file = ""
shape_mode = "upstream"
if not shape_mode:
shape_mode = "upstream" if (test_node.strip() and not shape_file) else "inline"
if shape_file:
shape_mode = "file"
pass_shape_file = "1"
elif shape_mode == "upstream":
pass_shape_file = "0"
else:
pass_shape_file = "1"
shape_mode = "inline"
return {
"OP": op,
"OP_FILE": op_file,
"FLAGGEMS_DIR": flaggems_dir,
"TEST_FILE": test_file,
"SHAPE_FILE": shape_file,
"SHAPE_MODE": shape_mode,
"PASS_SHAPE_FILE": pass_shape_file,
}
def _print_shell(env: Dict[str, str]) -> None:
for key, value in env.items():
print(f"{key}={shlex.quote(str(value))}")
def main(argv: Optional[list] = None) -> int:
p = argparse.ArgumentParser(description="FlagGems pytest node helpers")
p.add_argument("--parse-node", default="")
p.add_argument("--emit-shell", action="store_true")
p.add_argument("--op", default="")
p.add_argument("--op-file", default="")
p.add_argument("--flaggems-dir", default="")
p.add_argument("--shape-file", default="")
p.add_argument("--shape-mode", default="")
p.add_argument("--shell", action="store_true",
help="with --parse-node, print OP/OP_FILE/FLAGGEMS_DIR/TEST_FILE")
args = p.parse_args(argv)
if args.emit_shell:
_print_shell(emit_run_env(
test_node=args.parse_node,
op=args.op,
op_file=args.op_file,
flaggems_dir=args.flaggems_dir,
shape_file=args.shape_file,
shape_mode=args.shape_mode,
))
return 0
if args.parse_node:
parsed = parse_test_node(args.parse_node)
if args.shell:
_print_shell({k: v or "" for k, v in parsed.items()
if k in ("op", "op_file", "flaggems_dir", "test_file")})
return 0
for key, value in parsed.items():
print(f"{key}={value or ''}")
return 0
p.error("pass --emit-shell or --parse-node")
return 2
if __name__ == "__main__":
raise SystemExit(main())
+84 -47
View File
@@ -1,13 +1,16 @@
#!/usr/bin/env python3
"""多算子批量测试驱动:按算子清单逐个调 run_pytest.sh,多卡分片、串行测量。
每个算子一个独立子进程(CUDA crash 只废单个算子)、独立产物目录
每个算子一个独立子进程(设备 crash 只废单个算子)、独立产物目录
<batch>/<repo>/<op>/(与单算子跑 run_pytest.sh 的产物结构完全一致,含
run.log / shapes.yaml / autotune_records / ttgir)。批量层只做编排:
run.log / shapes.yaml / autotune_records / ttgir / run_identity.json)。
批量层只做编排:
- 分片:按算子名稳定哈希到 GPU(每卡内部串行,测量口径与单算子一致;
- 分片:按算子名稳定哈希到设备(每卡内部串行,测量口径与单算子一致;
跨卡并行测量存在功耗耦合噪声,本模式定位为筛查口径,可疑算子回
单算子串行流程确认);
- 后端:--backend / ZL_BACKEND 选择 cuda/musa/rocm/ascend,设备可见性
由 backends/ 写入,不再写死 CUDA_VISIBLE_DEVICES;
- 超时:--op-timeout 杀死挂死算子(记 TIMEOUT,不重试,继续跑);
- 重试:pytest rc=1(用例失败/flaky)与信号杀(rc>=128)最多重试 2 次,
连续两次同 rc 视为确定性失败提前止损;--dtypes 下两级降级——遇上游
@@ -23,7 +26,7 @@ run.log / shapes.yaml / autotune_records / ttgir)。批量层只做编排:
示例:
# 全量筛查(两仓库全部可收集算子,排除 marked_skip 与 dismiss)
python run_batch.py --all --gpus 0,1,2,3,4,5,6,7 --op-timeout 1800
python run_batch.py --all --devices 0,1,2,3,4,5,6,7 --backend cuda --op-timeout 1800
# 指定算子子集(repo 前缀默认 main;重名用 @文件名 消歧)
python run_batch.py --ops softmax vllm:fused_marlin_moe nextafter_@nextafter_
@@ -50,6 +53,8 @@ from pathlib import Path
from typing import Dict, List, Optional, Set, Tuple
import batch_summary
from backends import UnknownBackendError, get_backend
from _run_identity import load_identity
ZL_BENCH = Path(__file__).resolve().parent
OPS_DIR = ZL_BENCH / "ops"
@@ -62,7 +67,9 @@ REPO_DIRS = {
}
STATUS_FIELDS = [
"repo", "op", "op_file", "run_name", "gpu", "status", "return_code",
"repo", "op", "op_file", "run_name",
"backend", "device", "gpu", "timer",
"status", "return_code",
"elapsed_ms", "rows_total", "rows_success", "notes", "out_dir",
]
@@ -194,48 +201,44 @@ def resolve_selection(args: argparse.Namespace,
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 resolve_device_arg(devices: Optional[str], gpus: Optional[str]) -> Optional[str]:
"""Prefer --devices; --gpus is an alias. Passing both is an error."""
if devices is not None and gpus is not None:
raise ValueError("do not pass both --devices and --gpus")
return devices if devices is not None else gpus
def partition(cases: List[Case], gpus: List[str]) -> Dict[str, List[Case]]:
def visible_devices(backend, explicit: Optional[str] = None) -> List[str]:
return list(backend.enumerate_devices(explicit))
def partition(cases: List[Case], devices: List[str]) -> Dict[str, List[Case]]:
"""按 (repo, run_name) 稳定哈希分片:与选择顺序无关,resume 时同卡复用。"""
parts: Dict[str, List[Case]] = {g: [] for g in gpus}
parts: Dict[str, List[Case]] = {d: [] for d in devices}
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)
parts[devices[int.from_bytes(digest[:8], "big") % len(devices)]].append(c)
return parts
# ---------------- 单算子执行 ----------------
def _spawn(case: Case, out_dir: Path, gpu: str, args: argparse.Namespace,
def _spawn(case: Case, out_dir: Path, device: 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 = args.backend_obj.apply_visibility(env, device)
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,
"ZL_BACKEND": args.backend,
"ZL_DEVICE": device,
"FLAGGEMS_PERF_COLOR": "never",
"NO_COLOR": "1",
"USE_FLAGTUNE": str(args.use_flagtune),
@@ -271,13 +274,14 @@ def _wait(proc: subprocess.Popen, timeout: int) -> Tuple[int, bool]:
return 124, True
def run_case(case: Case, batch_dir: Path, gpu: str,
def run_case(case: Case, batch_dir: Path, device: 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 = _classify(case, out_dir, device, 0, 0, "", timed_out=False,
args=args)
row["notes"] = _join("resumed", row["notes"])
return row
@@ -295,7 +299,7 @@ def run_case(case: Case, batch_dir: Path, gpu: str,
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),
rc, timed_out = _wait(_spawn(case, out_dir, device, args, dtypes),
args.op_timeout)
rcs.append(rc)
if timed_out:
@@ -315,20 +319,22 @@ def run_case(case: Case, batch_dir: Path, gpu: str,
# 覆盖 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)
probe = _classify(case, out_dir, device, rcs[-1], 0, "", False,
args=args)
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, ""),
rc, timed_out = _wait(_spawn(case, out_dir, device, 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)
row = _classify(case, out_dir, device, rcs[-1], elapsed_ms, notes, timed_out,
args=args)
if row["status"] in ("PASS", "SKIP"):
tmp = marker.with_name(".complete.tmp")
tmp.write_text(row["status"] + "\n")
@@ -345,8 +351,21 @@ 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]:
def _identity_fields(out_dir: Path, device: str, backend: str) -> Dict[str, str]:
ident = load_identity(str(out_dir))
backend_v = str(ident.get("backend") or backend or "")
device_v = str(ident.get("device") or device or "")
timer_v = str(ident.get("timer") or "")
return {
"backend": backend_v,
"device": device_v,
"gpu": device_v,
"timer": timer_v,
}
def _classify(case: Case, out_dir: Path, device: str, rc: int, elapsed_ms: int,
notes: str, timed_out: bool, args: argparse.Namespace) -> Dict[str, str]:
text = _read_log(out_dir)
rows = batch_summary.parse_run_log(text)
if timed_out:
@@ -356,12 +375,13 @@ def _classify(case: Case, out_dir: Path, gpu: str, rc: int, elapsed_ms: int,
n_skipped = batch_summary.pytest_skipped_count(text)
if status == "PASS" and n_skipped:
notes = _join(notes, f"partial_skip={n_skipped}")
ident = _identity_fields(out_dir, device, args.backend)
return {
"repo": case.repo,
"op": case.op,
"op_file": case.op_file,
"run_name": case.run_name,
"gpu": gpu,
**ident,
"status": status,
"return_code": str(rc),
"elapsed_ms": str(elapsed_ms),
@@ -393,8 +413,12 @@ def parse_args() -> argparse.Namespace:
help="--all 时不剔除带无条件 skip 标记的算子")
run = p.add_argument_group("执行")
run.add_argument("--backend", default=os.environ.get("ZL_BACKEND", "cuda"),
help="设备后端:cuda/musa/rocm/ascend(默认 ZL_BACKEND 或 cuda)")
run.add_argument("--devices", default=None,
help="逗号分隔设备列表(默认该后端可见设备;优先于 --gpus)")
run.add_argument("--gpus", default=None,
help="逗号分隔 GPU 列表(默认 CUDA_VISIBLE_DEVICES 或全部)")
help="--devices 的别名(两者同时给出则报错)")
run.add_argument("--batch-dir", type=Path, default=None,
help="产物目录(默认 runs/batch_<时间戳>;resume 时必须指定)")
run.add_argument("--op-timeout", type=int, default=1800,
@@ -429,19 +453,30 @@ def parse_args() -> argparse.Namespace:
def main() -> None:
args = parse_args()
try:
backend = get_backend(args.backend)
except UnknownBackendError as exc:
raise SystemExit(f"{LOG} {exc}") from exc
args.backend_obj = backend
args.backend = backend.name
try:
explicit = resolve_device_arg(args.devices, args.gpus)
except ValueError as exc:
raise SystemExit(f"{LOG} {exc}") from exc
devices = visible_devices(backend, explicit)
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)} 张卡 "
parts = {d: cs for d, cs in partition(cases, devices).items() if cs}
print(f"{LOG} backend={args.backend} {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>")
@@ -461,17 +496,19 @@ def main() -> None:
writer.writerows(rows)
os.replace(tmp, status_csv)
def run_partition(gpu: str, part: List[Case]) -> None:
def run_partition(device: str, part: List[Case]) -> None:
for case in part:
try:
row = run_case(case, batch_dir, gpu, args)
row = run_case(case, batch_dir, device, args)
except Exception as exc: # 编排层异常不拖垮整卡分片
out_dir = batch_dir / case.repo / case.run_name
ident = _identity_fields(out_dir, device, args.backend)
row = {
"repo": case.repo, "op": case.op, "op_file": case.op_file,
"run_name": case.run_name, "gpu": gpu, "status": "FAIL",
"run_name": case.run_name, **ident, "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),
"out_dir": str(out_dir),
}
with lock:
done[0] += 1
@@ -480,16 +517,16 @@ def main() -> None:
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"[device{device}]"
+ (f" {row['notes']}" if row["notes"] else ""),
flush=True)
if len(parts) == 1:
gpu, part = next(iter(parts.items()))
run_partition(gpu, part)
device, part = next(iter(parts.items()))
run_partition(device, part)
else:
with ThreadPoolExecutor(max_workers=len(parts)) as pool:
futures = [pool.submit(run_partition, g, p) for g, p in parts.items()]
futures = [pool.submit(run_partition, d, p) for d, p in parts.items()]
for f in futures:
f.result()
+66 -26
View File
@@ -8,15 +8,21 @@ set -euo pipefail
# 每项均可用同名环境变量覆盖,详见 README「环境变量一览」。
# 被测算子,对应 $FLAGGEMS_DIR/benchmark/test_<OP_FILE>.py::test_<OP>
# 也可直接传 FlagGems pytest 节点,例如:
# TEST_NODE=/workspace/FlagGems/benchmark/test_fused_marlin_moe.py::test_fused_marlin_moe_mxfp4
OP="${OP:-fused_marlin_moe_mxfp4}"
OP_FILE="${OP_FILE:-fused_marlin_moe}"
# 上游已拆分为两个仓库:主仓 /workspace/dev/FlagGems(包名 flag_gems)与
# /workspace/dev/FlagGems-vllm(包名 flaggems_vllm)。测 vllm 仓库算子时覆盖本变量即可。
# NV 产品路径也常见 /workspace/FlagGems;TEST_NODE 或缺失测试文件时会回退到该路径。
FLAGGEMS_DIR="${FLAGGEMS_DIR:-/workspace/dev/FlagGems}"
TEST_NODE="${TEST_NODE:-${ZL_TEST_NODE:-}}"
# shape 来源:非空则用该 yaml,否则用 INLINE_YAML。顶层 key 必须是 op 名。
# shape 来源:非空则用该 yaml;SHAPE_MODE=upstream 则不传 --shape_file,沿用 FlagGems
# 在 $BENCH_LEVEL 下的原生 shape 集。设了 TEST_NODE 且未指定 SHAPE_FILE 时默认 upstream。
SHAPE_FILE="${SHAPE_FILE:-}"
SHAPE_MODE="${SHAPE_MODE:-}"
read -r -d '' INLINE_YAML <<'YAML' || true
fused_marlin_moe_mxfp4:
shapes:
@@ -53,25 +59,34 @@ REPLAY_FROM="${REPLAY_FROM:-}"
# always/never 强制开关终端颜色,空则按 tty 判断(run.log 始终去色)
FLAGGEMS_PERF_COLOR="${FLAGGEMS_PERF_COLOR:-}"
# 各插件作用见 README「插件说明」;注释掉某行即停用该插件
PLUGINS=(
-p _device_guard_plugin
-p _parallel_warmup_plugin
-p _seed_plugin
-p _shape_inject_plugin
-p _shape_iter_inject_plugin
-p _bespoke_shape_plugin
# -p _mm_cluster_fix_plugin
-p _autotune_record_plugin
-p _cudagraph_plugin
-p _pretty_report_plugin
-p _ir_meta_plugin
)
# 设备后端:cuda / musa / rocm / ascend(别名见 backends/__init__.py)。
# 插件列表由 plugin_plan.py 按后端能力生成,不再在此硬编码 -p。
# 非 CUDA 默认走上游 do_bench,不会加载 _cudagraph_plugin。
ZL_BACKEND="${ZL_BACKEND:-cuda}"
# 上游 --level:core(默认)只用各算子的 core_shapes;comprehensive 额外并入
# 子类的 set_more_shapes(),并对同一 shape 再产出一次变体(BLAS 是转置侧)。
# 差距可以很大:mm 在 core 下是 5 个 shape,comprehensive 下并入 4 个大 K
# (N/K 到 152064)与 models_shapes/*.yaml 的模型 shape,共 29 个,且用例翻倍。
# 代价同步放大——sweep 是墙上时间的大头,shape 多时配合 PARALLEL_WARMUP_GPUS。
BENCH_LEVEL="${BENCH_LEVEL:-core}"
case "$BENCH_LEVEL" in
core|comprehensive) ;;
*) echo "!!! BENCH_LEVEL 只能是 core 或 comprehensive(当前:$BENCH_LEVEL)" >&2; exit 2 ;;
esac
# ============================== 执行逻辑 ==============================
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
TEST_FILE="$FLAGGEMS_DIR/benchmark/test_${OP_FILE}.py::test_${OP}"
export PYTHONPATH="$SCRIPT_DIR${PYTHONPATH:+:$PYTHONPATH}"
eval "$(python "$SCRIPT_DIR/providers/flaggems.py" --emit-shell \
--parse-node "$TEST_NODE" \
--op "$OP" \
--op-file "$OP_FILE" \
--flaggems-dir "$FLAGGEMS_DIR" \
--shape-file "$SHAPE_FILE" \
--shape-mode "$SHAPE_MODE")"
: "${TEST_FILE:=$FLAGGEMS_DIR/benchmark/test_${OP_FILE}.py::test_${OP}}"
# 产物目录,可用 OUT_DIR 指定。时间戳只到秒,并发启动会撞名,故用不带 -p 的
# mkdir 抢占(已存在即失败),撞上就退让到 -2、-3……
@@ -104,9 +119,12 @@ else
PYTEST_COLOR=()
fi
export PYTHONPATH="$SCRIPT_DIR${PYTHONPATH:+:$PYTHONPATH}"
export FLAGGEMS_PERF_CURRENT_OP="$OP"
export PYTHONUNBUFFERED=1 # 实时输出不缓冲
export ZL_BACKEND
# 按 ZL_BACKEND / ZL_ENABLE_* 生成 -p 列表(见 plugin_plan.py)
read -r -a PLUGINS <<< "$(python "$SCRIPT_DIR/plugin_plan.py" --backend "$ZL_BACKEND" --pytest-args)"
# DTYPES 非空时逐个转为上游 --dtypes 选项(action=append,每个 dtype 一次)
DTYPE_ARGS=()
@@ -126,17 +144,39 @@ else
MODE_DESC=record
fi
# 未指定 shape 文件时,把 INLINE_YAML 落到临时文件供 pytest 读取
if [[ -z "$SHAPE_FILE" ]]; then
SHAPE_FILE="$(mktemp --suffix=.yaml)"
printf '%s\n' "$INLINE_YAML" > "$SHAPE_FILE"
trap 'rm -f "$SHAPE_FILE"' EXIT
# 未指定 shape 文件时:inline 用 INLINE_YAML;upstream 不传 --shape_file。
SHAPE_ARGS=()
if [[ "${PASS_SHAPE_FILE:-1}" == "0" ]]; then
if [[ -f "$FLAGGEMS_DIR/benchmark/core_shapes.yaml" ]]; then
cp -f "$FLAGGEMS_DIR/benchmark/core_shapes.yaml" "$OUT_DIR/shapes.yaml"
else
printf '%s\n' "source: upstream" "level: $BENCH_LEVEL" > "$OUT_DIR/shapes.yaml"
fi
else
if [[ -z "$SHAPE_FILE" ]]; then
SHAPE_FILE="$(mktemp --suffix=.yaml)"
printf '%s\n' "$INLINE_YAML" > "$SHAPE_FILE"
trap 'rm -f "$SHAPE_FILE"' EXIT
fi
cp -f "$SHAPE_FILE" "$OUT_DIR/shapes.yaml" # 存档本次实际使用的 shape
SHAPE_ARGS=(--shape_file "$SHAPE_FILE")
fi
cp -f "$SHAPE_FILE" "$OUT_DIR/shapes.yaml" # 存档本次实际使用的 shape
python "$SCRIPT_DIR/_run_identity.py" "$OUT_DIR" \
--backend "$ZL_BACKEND" \
--op "$OP" \
--op-file "$OP_FILE" \
--flaggems-dir "$FLAGGEMS_DIR" \
--autotune-mode "$MODE_DESC" \
--device "${ZL_DEVICE:-}" \
--test-file "$TEST_FILE" \
--shape-mode "$SHAPE_MODE" \
--bench-level "$BENCH_LEVEL"
status=0
{
echo "${C_BOLD}>>> op=$OP mode=$MODE_DESC USE_FLAGTUNE=$USE_FLAGTUNE${C_RESET}"
echo "${C_BOLD}>>> op=$OP backend=$ZL_BACKEND mode=$MODE_DESC USE_FLAGTUNE=$USE_FLAGTUNE shape=$SHAPE_MODE level=$BENCH_LEVEL${C_RESET}"
echo "${C_DIM}>>> node=$TEST_FILE${C_RESET}"
echo "${C_DIM}>>> out=$OUT_DIR${C_RESET}"
# 每次用独立的 Triton 缓存目录,跑完即删:保证编译过程可复现,且 ttgir 落盘
@@ -149,8 +189,8 @@ status=0
env "$AUTOTUNE_ENV" \
USE_FLAGTUNE=$USE_FLAGTUNE python -u -m pytest -s "$TEST_FILE" \
"${PLUGINS[@]}" "${PYTEST_COLOR[@]}" \
--shape_file "$SHAPE_FILE" ${DTYPE_ARGS[@]+"${DTYPE_ARGS[@]}"} \
--level core --mode kernel || status=$?
${SHAPE_ARGS[@]+"${SHAPE_ARGS[@]}"} ${DTYPE_ARGS[@]+"${DTYPE_ARGS[@]}"} \
--level "$BENCH_LEVEL" --mode kernel || status=$?
rm -rf "$CACHE_DIR"
+285
View File
@@ -0,0 +1,285 @@
#!/usr/bin/env python3
"""No-GPU unit tests for zl_bench backend adapters and plugin policy.
FlagGems pytest invocations pass an explicit node, so this file is not
collected during operator runs. Run it with `python test_backends.py`.
"""
from __future__ import annotations
import json
import os
import sys
import tempfile
import unittest
from pathlib import Path
ROOT = Path(__file__).resolve().parent
if str(ROOT) not in sys.path:
sys.path.insert(0, str(ROOT))
from backends import ( # noqa: E402
UnknownBackendError,
get_backend,
normalize_backend_name,
)
from backends.base import ALL_VISIBILITY_KEYS # noqa: E402
from batch_summary import SUMMARY_FIELDS # noqa: E402
from plugin_plan import plan_plugins # noqa: E402
from providers.flaggems import ( # noqa: E402
VENDOR_BY_BACKEND,
emit_run_env,
flaggems_vendor,
parse_test_node,
)
from run_batch import STATUS_FIELDS, resolve_device_arg # noqa: E402
from _run_identity import init_identity, load_identity, update_identity # noqa: E402
class BackendResolveTests(unittest.TestCase):
def test_empty_and_alias_resolve_to_cuda(self):
self.assertEqual(normalize_backend_name(""), "cuda")
self.assertEqual(get_backend("nvidia").name, "cuda")
self.assertEqual(get_backend("nv").name, "cuda")
self.assertEqual(get_backend("").name, "cuda")
def test_other_aliases(self):
self.assertEqual(get_backend("amd").name, "rocm")
self.assertEqual(get_backend("hip").name, "rocm")
self.assertEqual(get_backend("npu").name, "ascend")
self.assertEqual(get_backend("huawei").name, "ascend")
def test_unknown_backend_fails_closed(self):
with self.assertRaises(UnknownBackendError):
get_backend("tpu")
with self.assertRaises(UnknownBackendError):
get_backend("not-a-backend")
class VisibilityTests(unittest.TestCase):
def test_musa_drops_leftover_cuda(self):
env = {"CUDA_VISIBLE_DEVICES": "7", "PATH": "/bin"}
out = get_backend("musa").apply_visibility(env, "1")
self.assertNotIn("CUDA_VISIBLE_DEVICES", out)
self.assertEqual(out["MUSA_VISIBLE_DEVICES"], "1")
self.assertEqual(out["PATH"], "/bin")
for key in ALL_VISIBILITY_KEYS:
if key != "MUSA_VISIBLE_DEVICES":
self.assertNotIn(key, out)
def test_rocm_writes_hip_and_rocr(self):
out = get_backend("rocm").apply_visibility(
{"CUDA_VISIBLE_DEVICES": "0"}, "3")
self.assertEqual(out["HIP_VISIBLE_DEVICES"], "3")
self.assertEqual(out["ROCR_VISIBLE_DEVICES"], "3")
self.assertNotIn("CUDA_VISIBLE_DEVICES", out)
def test_ascend_writes_both_keys(self):
out = get_backend("ascend").apply_visibility({}, "2")
self.assertEqual(out["ASCEND_RT_VISIBLE_DEVICES"], "2")
self.assertEqual(out["ASCEND_VISIBLE_DEVICES"], "2")
self.assertNotIn("CUDA_VISIBLE_DEVICES", out)
def test_cuda_only_sets_cuda(self):
out = get_backend("cuda").apply_visibility(
{"MUSA_VISIBLE_DEVICES": "9"}, "0")
self.assertEqual(out["CUDA_VISIBLE_DEVICES"], "0")
self.assertNotIn("MUSA_VISIBLE_DEVICES", out)
class FlagGemsVendorTests(unittest.TestCase):
def test_vendor_map(self):
self.assertEqual(VENDOR_BY_BACKEND["cuda"], "nvidia")
self.assertEqual(VENDOR_BY_BACKEND["rocm"], "amd")
self.assertEqual(flaggems_vendor(get_backend("cuda")), "nvidia")
self.assertEqual(flaggems_vendor(get_backend("rocm")), "amd")
self.assertEqual(flaggems_vendor(get_backend("musa")), "musa")
self.assertEqual(flaggems_vendor(get_backend("ascend")), "ascend")
class FlagGemsNodeTests(unittest.TestCase):
_NODE = (
"/workspace/FlagGems/benchmark/test_fused_marlin_moe.py"
"::test_fused_marlin_moe_mxfp4"
)
def test_parse_user_command_node(self):
parsed = parse_test_node(self._NODE)
self.assertEqual(parsed["op"], "fused_marlin_moe_mxfp4")
self.assertEqual(parsed["op_file"], "fused_marlin_moe")
self.assertEqual(parsed["flaggems_dir"], "/workspace/FlagGems")
self.assertTrue(parsed["test_file"].endswith(
"test_fused_marlin_moe.py::test_fused_marlin_moe_mxfp4"))
def test_test_node_defaults_to_upstream_shapes(self):
env = emit_run_env(
test_node=self._NODE,
op="softmax",
op_file="softmax",
flaggems_dir="/workspace/dev/FlagGems",
)
self.assertEqual(env["OP"], "fused_marlin_moe_mxfp4")
self.assertEqual(env["OP_FILE"], "fused_marlin_moe")
self.assertEqual(env["FLAGGEMS_DIR"], "/workspace/FlagGems")
self.assertEqual(env["SHAPE_MODE"], "upstream")
self.assertEqual(env["PASS_SHAPE_FILE"], "0")
self.assertEqual(env["TEST_FILE"], self._NODE)
def test_plain_run_keeps_inline_shapes(self):
env = emit_run_env(
op="fused_marlin_moe_mxfp4",
op_file="fused_marlin_moe",
flaggems_dir="/workspace/dev/FlagGems",
)
self.assertEqual(env["SHAPE_MODE"], "inline")
self.assertEqual(env["PASS_SHAPE_FILE"], "1")
self.assertIn("test_fused_marlin_moe.py::test_fused_marlin_moe_mxfp4",
env["TEST_FILE"])
def test_explicit_shape_file_wins(self):
env = emit_run_env(
test_node=self._NODE,
shape_file="/tmp/my_shapes.yaml",
)
self.assertEqual(env["SHAPE_MODE"], "file")
self.assertEqual(env["PASS_SHAPE_FILE"], "1")
def test_resolve_uses_existing_local_root(self):
with tempfile.TemporaryDirectory() as tmp:
root = Path(tmp) / "FlagGems"
py = root / "benchmark" / "test_fused_marlin_moe.py"
py.parent.mkdir(parents=True)
py.write_text("# stub\n", encoding="utf-8")
env = emit_run_env(
op="fused_marlin_moe_mxfp4",
op_file="fused_marlin_moe",
flaggems_dir=str(root),
)
self.assertEqual(env["FLAGGEMS_DIR"], str(root))
self.assertTrue(Path(env["TEST_FILE"].split("::", 1)[0]).is_file())
class PluginPlanTests(unittest.TestCase):
def _plan(self, backend, **kwargs):
kwargs.setdefault("extra", [])
kwargs.setdefault("enable_autotune", True)
kwargs.setdefault("enable_ir", True)
kwargs.setdefault("enable_mm_cluster", False)
return plan_plugins(backend, **kwargs)
def test_cuda_does_not_put_cudagraph_on_the_list(self):
names = self._plan("cuda")
self.assertIn("_timing_plugin", names)
self.assertIn("_device_setup_plugin", names)
self.assertNotIn("_cudagraph_plugin", names)
self.assertNotIn("_mm_cluster_fix_plugin", names)
def test_musa_same_core_without_graph_or_cluster(self):
names = self._plan("musa")
self.assertIn("_timing_plugin", names)
self.assertNotIn("_cudagraph_plugin", names)
self.assertNotIn("_mm_cluster_fix_plugin", names)
def test_mm_cluster_ignored_on_non_cuda(self):
names = self._plan("musa", enable_mm_cluster=True)
self.assertNotIn("_mm_cluster_fix_plugin", names)
names = self._plan("cuda", enable_mm_cluster=True)
self.assertIn("_mm_cluster_fix_plugin", names)
def test_extra_plugins_isolated(self):
names = self._plan("cuda", extra=["_cudagraph_plugin"])
self.assertIn("_cudagraph_plugin", names)
names = self._plan("cuda", extra=[])
self.assertNotIn("_cudagraph_plugin", names)
def test_extra_none_reads_env(self):
old = os.environ.get("ZL_EXTRA_PLUGINS")
os.environ["ZL_EXTRA_PLUGINS"] = "_cudagraph_plugin"
try:
names = plan_plugins(
"cuda", extra=None, enable_autotune=True,
enable_ir=True, enable_mm_cluster=False)
self.assertIn("_cudagraph_plugin", names)
finally:
if old is None:
os.environ.pop("ZL_EXTRA_PLUGINS", None)
else:
os.environ["ZL_EXTRA_PLUGINS"] = old
class DeviceArgTests(unittest.TestCase):
def test_devices_preferred(self):
self.assertEqual(resolve_device_arg("0,1", None), "0,1")
self.assertEqual(resolve_device_arg(None, "2"), "2")
self.assertIsNone(resolve_device_arg(None, None))
def test_both_set_is_error(self):
with self.assertRaises(ValueError):
resolve_device_arg("0", "1")
class CsvFieldTests(unittest.TestCase):
def test_status_fields(self):
for name in ("backend", "device", "gpu", "timer"):
self.assertIn(name, STATUS_FIELDS)
def test_summary_fields(self):
for name in ("backend", "device", "timer"):
self.assertIn(name, SUMMARY_FIELDS)
self.assertLess(SUMMARY_FIELDS.index("backend"),
SUMMARY_FIELDS.index("fg_op_name"))
class IdentityMergeTests(unittest.TestCase):
def test_init_then_update(self):
old_out = os.environ.get("OUT_DIR")
try:
with tempfile.TemporaryDirectory() as tmp:
init_identity(tmp, backend="cuda", device="0", op="softmax")
update_identity(tmp, timer="direct")
ident = load_identity(tmp)
self.assertEqual(ident["schema_version"], 1)
self.assertEqual(ident["provider"], "flaggems")
self.assertEqual(ident["backend"], "cuda")
self.assertEqual(ident["device"], "0")
self.assertEqual(ident["timer"], "direct")
self.assertEqual(ident["op"], "softmax")
finally:
if old_out is None:
os.environ.pop("OUT_DIR", None)
else:
os.environ["OUT_DIR"] = old_out
class AutotuneReplayBackendTests(unittest.TestCase):
def setUp(self):
self._old = os.environ.get("ZL_BACKEND")
os.environ["ZL_BACKEND"] = "cuda"
def tearDown(self):
if self._old is None:
os.environ.pop("ZL_BACKEND", None)
else:
os.environ["ZL_BACKEND"] = self._old
def test_missing_meta_still_replays(self):
from _autotune_record_plugin import _replay_backend_ok
with tempfile.TemporaryDirectory() as tmp:
self.assertTrue(_replay_backend_ok(tmp))
def test_matching_backend_ok(self):
from _autotune_record_plugin import _replay_backend_ok
with tempfile.TemporaryDirectory() as tmp:
Path(tmp, "meta.json").write_text(
json.dumps({"backend": "cuda"}), encoding="utf-8")
self.assertTrue(_replay_backend_ok(tmp))
def test_mismatch_refuses_replay(self):
from _autotune_record_plugin import _replay_backend_ok
with tempfile.TemporaryDirectory() as tmp:
Path(tmp, "meta.json").write_text(
json.dumps({"backend": "musa"}), encoding="utf-8")
self.assertFalse(_replay_backend_ok(tmp))
if __name__ == "__main__":
unittest.main()