Add multi-op batch screening layer; fix cudagraph fallback regressions

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

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