diff --git a/README.md b/README.md index cc5b888..a524eee 100644 --- a/README.md +++ b/README.md @@ -161,7 +161,7 @@ benchmark 目录，注册 `--level`/`--mode` 的 `conftest.py` 就没被加载 只校验节点的文件那一半；`::test_` 交给 pytest，函数不存在时它自己会清楚地报 `not found: ...`。 -### ab_fold_test.sh：A/B 测试参考示例 +### ab_fold_test.sh / ab_rlc_test.sh：A/B 测试参考示例 `ab_fold_test.sh` 是上述 A/B 流程的现成参考：以 FlagGems `FLAGGEMS_MXFP4_FOLDSCALE`（fold_scale 优化）开关为对比对象，对 4 个 DeepSeek-V4-Flash 真实 trace shape 集（`shape/DeepSeek-V4-Flash-p{1024,4096,32768,65536}d1024.yaml`，随仓库提供）各跑一对 FOLD=0（基线，record）→ FOLD=1（replay 同 config），共 8 轮： @@ -170,6 +170,8 @@ bash ab_fold_test.sh # 汇总 log：runs/ab_fold_<时间戳>.log（已去色）；快速对账：grep -E '^#####|FAILED' ``` +`ab_rlc_test.sh` 是同一口径的另一个实例，对比轴换成 FlagTree 的 `FLAGTREE__RLC_ENHANCE`（`VENDOR` 默认 `MUSA`），算子表由 `OPS` 指定，跑完调 `batch_summary.py --ab ` 出该算子的 speedup 区间。 + 可复用的是**口径**——同 config replay、cudagraph 计时、真实 trace shape 集、逐对成组——对比其它开关 / 编译器改动 / 算子，改循环变量与环境变量即可，口径部分不用动。 但**它是口径示例，不是可长期直接执行的回归脚本**：对比轴是 FlagGems 侧的算子开关，会随上游演进被改名、改语义或移除，本仓库不跟随同步更新（该脚本引用的开关目前就已被上游移除）。失效后果是**静默的**：脚本照样跑完全部轮次、照样输出完整对比表，只是两侧执行同一份代码，得出"无差异"的假结论。所以套用前先确认对比轴仍有效——`grep -rn '<开关名>' $FLAGGEMS_DIR/src` 要有命中，且两侧 latency 表确有差异。 @@ -268,6 +270,7 @@ python run_batch.py --ops-file my_ops.txt --replay-root runs/batch_xxx | `_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 运行 | +| `_musa_benchmark_compat_plugin` | 剥掉 FlagGems 上游若干 benchmark 上过时的 `torch.cuda.is_available()` 跳过门——被测代码本就经 `flag_gems.device` 分发，在 MThreads 上那些门把可用算子变成「未测量」。每条覆盖都钉死到一个文件与一条确切的 skip 理由，且仅在 `flag_gems.vendor_name == "mthreads"` 时生效 | 由 `backends/musa.py` 的 `compat_plugins` 声明，其他后端不加载 | | `_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` | diff --git a/_musa_benchmark_compat_plugin.py b/_musa_benchmark_compat_plugin.py new file mode 100644 index 0000000..6a8ab87 --- /dev/null +++ b/_musa_benchmark_compat_plugin.py @@ -0,0 +1,89 @@ +"""Narrow compatibility fixes for the frozen MUSA product corpus. + +Several upstream FlagGems tests gate themselves on ``torch.cuda.is_available()`` +even though the code under test runs on any accelerator FlagGems dispatches to. +On MThreads those guards skip the test instead of exercising the operator, which +turns a working path into a verdict of "not measured". + +Every override here is pinned to one file and one exact skip reason, applies +only when FlagGems has dispatched to the MThreads vendor, and is only added +after the unskipped test was observed to pass on the device. + +Two files were probed on S5000 and deliberately left out, because their guards +are not stale: + +- ``test_cutlass_scaled_mm.py`` gates on vLLM as well as CUDA, so removing the + guard would raise on the missing dependency instead of measuring anything. +- ``test_grid_sample.py`` builds its inputs through a helper whose device + argument defaults to the literal ``"cuda"``. Its guard is consistent with its + own body: unskipped, all 90 collected cases fail in ``aten::empty`` on a CUDA + backend that does not exist here. Making that file portable means editing + upstream test source, which is outside what a compatibility plugin may do. +""" + +from __future__ import annotations + +import sys +from pathlib import Path + + +# file name -> exact skipif reasons that are stale on a dispatching backend. +_STALE_CUDA_GUARDS: dict[str, frozenset[str]] = { + "test_fft.py": frozenset({"CUDA is not available"}), + "test_top_k_per_row_decode.py": frozenset({"CUDA device required"}), +} + +# Guards that are only stale for one test in the file stay pinned to it. +_PINNED_TESTS: dict[str, str] = { + "test_fft.py": "test_perf_fft", +} + + +def _is_stale_cuda_guard(item, mark) -> bool: + if mark.name != "skipif": + return False + name = Path(str(item.path)).name + reasons = _STALE_CUDA_GUARDS.get(name) + if reasons is None or mark.kwargs.get("reason") not in reasons: + return False + pinned = _PINNED_TESTS.get(name) + return pinned is None or item.name == pinned + + +# Retained under its historical name: the sealed campaign's tests reference it. +def _is_stale_fft_cuda_guard(item, mark) -> bool: + return Path(str(item.path)).name == "test_fft.py" and _is_stale_cuda_guard( + item, mark + ) + + +def pytest_collection_modifyitems(items): + import flag_gems + + if flag_gems.vendor_name != "mthreads": + return + removed: dict[str, int] = {} + # A guard can sit on the test, its class, or the module's pytestmark, and + # skipif is resolved by walking that whole chain. Strip every level. + seen: set[int] = set() + for item in items: + for node in item.listchain(): + if id(node) in seen: + continue + seen.add(id(node)) + markers = getattr(node, "own_markers", None) + if not markers: + continue + kept = [mark for mark in markers if not _is_stale_cuda_guard(item, mark)] + dropped = len(markers) - len(kept) + if dropped: + markers[:] = kept + name = Path(str(item.path)).name + removed[name] = removed.get(name, 0) + dropped + for name, count in sorted(removed.items()): + print( + f"[musa-benchmark-compat] removed {count} stale CUDA-only skip " + f"guard(s) from {name}", + file=sys.stderr, + flush=True, + ) diff --git a/ab_rlc_test.sh b/ab_rlc_test.sh new file mode 100644 index 0000000..270e9f3 --- /dev/null +++ b/ab_rlc_test.sh @@ -0,0 +1,68 @@ +#!/bin/bash +# A/B 对比：FlagTree RemoveLayoutConversions 增强关 vs 开，逐算子两轮。 +# +# 口径：每个算子先跑关（record autotune config），再用 REPLAY_FROM 回放同一套 +# config 跑开，两侧只差 FLAGTREE__RLC_ENHANCE 这一个开关。产物沿用 +# run_pytest.sh 的目录结构，所以每个算子恰好留下两份 run.log： +# +# runs/_<时间戳>/run.log 关（before） +# runs/_<时间戳>/run.log 开（after） +# +# 跑完用 batch_summary.py --ab 得到该算子的 speedup 区间。 +# +# 用法： +# OPS="softmax gelu_backward" bash ab_rlc_test.sh +# OPS="sort:perf_sort topk" bash ab_rlc_test.sh # 冒号分隔 OP:OP_FILE +# +# 注意：RLC 开关名带厂商前缀，用 VENDOR 指定（默认 MUSA）。套用前先确认开关有效： +# grep 该开关名于 FlagTree 源码应有命中，否则两侧跑的是同一份代码。 +set -uo pipefail + +cd "$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" + +VENDOR="${VENDOR:-MUSA}" +ENHANCE_VAR="FLAGTREE_${VENDOR}_RLC_ENHANCE" + +# 空格分隔的算子表。元素为 OP 或 OP:OP_FILE（省略时 OP_FILE=OP）。 +OPS="${OPS:-softmax}" + +if [[ -t 1 && -z "${NO_COLOR:-}" ]]; then + export FLAGGEMS_PERF_COLOR=always +fi + +LOG="runs/ab_rlc_$(date +%Y%m%d-%H%M%S).log" +mkdir -p runs + +{ + echo ">>> vendor=$VENDOR switch=$ENHANCE_VAR ops=$OPS" + for entry in $OPS; do + op="${entry%%:*}" + op_file="${entry#*:}" + [[ "$op_file" == "$entry" ]] && op_file="$op" + + echo "########## $op | $ENHANCE_VAR=0 before (record) ##########" + if ! OP="$op" OP_FILE="$op_file" env "$ENHANCE_VAR=0" bash run_pytest.sh; then + echo "!!! FAILED before: $op" + continue + fi + BEFORE=$(ls -dt "runs/${op}_"* | head -1) + + echo "########## $op | $ENHANCE_VAR=1 after (replay $BEFORE) ##########" + if ! OP="$op" OP_FILE="$op_file" REPLAY_FROM="$BEFORE" env "$ENHANCE_VAR=1" bash run_pytest.sh; then + echo "!!! FAILED after: $op" + continue + fi + AFTER=$(ls -dt "runs/${op}_"* | head -1) + + echo "########## $op | summary ##########" + python batch_summary.py --ab "$BEFORE/run.log" "$AFTER/run.log" --label "$op" \ + || echo "!!! summary failed: $op" + done + echo "########## ALL DONE ##########" +} 2>&1 | tee "$LOG" + +# 去掉 ANSI 转义以便 grep/diff（终端输出保留颜色） +sed -i -E $'s/\x1b\\[[0-9;]*[A-Za-z]//g' "$LOG" + +echo ">>> 总 log: $(pwd)/$LOG" +echo ">>> 每算子两份 run.log 在 runs/_<时间戳>/ 下，按时间先后即 before / after" diff --git a/backends/base.py b/backends/base.py index afaef00..c1afae7 100644 --- a/backends/base.py +++ b/backends/base.py @@ -90,6 +90,13 @@ class DeviceBackend: def denies_capture(self, op: str) -> bool: return bool(op) and op in self.capture_deny_ops + + # Plugins that patch around upstream assumptions this backend violates, e.g. + # FlagGems benchmarks that gate themselves on torch.cuda.is_available() even + # though the operator under test runs on whatever device FlagGems dispatched + # to. Loaded first so collection-time overrides land before anything reads + # the item list. Each such plugin must scope itself to its own vendor. + compat_plugins: Tuple[str, ...] = () # 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 diff --git a/backends/musa.py b/backends/musa.py index 18f0add..066af8d 100644 --- a/backends/musa.py +++ b/backends/musa.py @@ -42,6 +42,11 @@ class MusaBackend(DeviceBackend): # timing precision instead of its measurement. Revisit if a newer torch_musa # changes the behaviour. capture_deny_ops = ("unique2",) + # Several FlagGems benchmarks skip on `not torch.cuda.is_available()` while + # the code under test dispatches through flag_gems.device. On MThreads that + # turns a working operator into "not measured"; the plugin drops those + # guards, pinned to one file and one exact skip reason each. + compat_plugins = ("_musa_benchmark_compat_plugin",) # 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 diff --git a/batch_summary.py b/batch_summary.py index c4ebb69..5d2c190 100644 --- a/batch_summary.py +++ b/batch_summary.py @@ -186,12 +186,82 @@ def summarize(batch_dir: Path, out_path: Optional[Path] = None) -> Path: return out_path +def ab_rows(before_text: str, after_text: str) -> List[Dict[str, str]]: + """把两次 run.log 的测量行按 (op, dtype, size_detail) 配对。 + + speedup 列是 gems 对 torch，不是前后对比，所以这里用 gems_lat 自己算 + before/after - 1。只有两侧都 SUCCESS 且延迟为正的行才计入；对不上的行 + 单独返回，避免用不完整的配对报一个看起来完整的区间。 + """ + def index(text: str) -> Dict[Tuple[str, str, str], float]: + out: Dict[Tuple[str, str, str], float] = {} + for row in parse_run_log(text): + if row.get("row_status") != "SUCCESS": + continue + try: + lat = float(row.get("gems_lat", "")) + except ValueError: + continue + if lat > 0: + out[(row.get("fg_op_name", ""), row.get("dtype", ""), + row.get("size_detail", ""))] = lat + return out + + before, after = index(before_text), index(after_text) + paired: List[Dict[str, str]] = [] + for key in sorted(before.keys() & after.keys()): + op, dtype, size = key + b, a = before[key], after[key] + paired.append({ + "fg_op_name": op, "dtype": dtype, "size_detail": size, + "before_lat": f"{b:.6f}", "after_lat": f"{a:.6f}", + "speedup_percent": f"{(b / a - 1.0) * 100.0:.2f}", + }) + return paired + + +def report_ab(before: Path, after: Path, label: str = "") -> List[Dict[str, str]]: + """打印一个算子的前后对比：逐行 speedup 与整体区间。""" + before_text = before.read_text(encoding="utf-8", errors="replace") + after_text = after.read_text(encoding="utf-8", errors="replace") + rows = ab_rows(before_text, after_text) + name = label or (rows[0]["fg_op_name"] if rows else "?") + if not rows: + print(f"[ab] {name}: 没有可配对的测量行" + f"（before={before} after={after}）") + return rows + values = [float(r["speedup_percent"]) for r in rows] + for r in rows: + print(f" {r['dtype']:<16} {float(r['speedup_percent']):+7.2f}% " + f"{r['before_lat']} -> {r['after_lat']} {r['size_detail']}") + fallback = "" + if DIRECT_NO_CUDAGRAPH_RE.search(before_text) or \ + DIRECT_NO_CUDAGRAPH_RE.search(after_text): + # 回退行用的是 do_bench 口径，与 cudagraph 行不同表可比性 + fallback = " [含 BENCHMARK_DIRECT_NO_CUDAGRAPH 回退行]" + print(f"[ab] {name}: {len(rows)} 行配对，speedup " + f"{min(values):+.2f}% ~ {max(values):+.2f}%{fallback}") + return rows + + def main() -> None: p = argparse.ArgumentParser(description=__doc__.splitlines()[0]) - p.add_argument("batch_dir", type=Path, help="run_batch.py 的产物目录") + p.add_argument("batch_dir", type=Path, nargs="?", + help="run_batch.py 的产物目录") p.add_argument("-o", "--output", type=Path, default=None, help="输出 CSV（默认 /summary.csv）") + p.add_argument("--ab", nargs=2, type=Path, metavar=("BEFORE", "AFTER"), + help="两份 run.log，按 shape 配对输出该算子的 speedup 区间") + p.add_argument("--label", default="", help="--ab 报告里的算子名") args = p.parse_args() + if args.ab: + for path in args.ab: + if not path.is_file(): + raise SystemExit(f"run.log 不存在: {path}") + report_ab(args.ab[0], args.ab[1], args.label) + return + if args.batch_dir is None: + raise SystemExit("需要 batch_dir，或用 --ab BEFORE AFTER") if not args.batch_dir.is_dir(): raise SystemExit(f"目录不存在: {args.batch_dir}") out = summarize(args.batch_dir, args.output) diff --git a/plugin_plan.py b/plugin_plan.py index d39ea3d..48b058a 100644 --- a/plugin_plan.py +++ b/plugin_plan.py @@ -41,7 +41,9 @@ def plan_plugins( enable_mm_cluster: Optional[bool] = None, ) -> List[str]: backend = get_backend(backend_name) - plugins: List[str] = list(CORE_PLUGINS) + # Compat first: these override collection-time markers, so they have to be + # in place before anything else inspects the collected items. + plugins: List[str] = list(backend.compat_plugins) + list(CORE_PLUGINS) autotune = ( backend.supports_triton_hooks diff --git a/test_backends.py b/test_backends.py index c7b600a..943d50a 100644 --- a/test_backends.py +++ b/test_backends.py @@ -173,6 +173,17 @@ class PluginPlanTests(unittest.TestCase): self.assertNotIn("_cudagraph_plugin", names) self.assertNotIn("_mm_cluster_fix_plugin", names) + def test_musa_loads_its_compat_plugin_first(self): + names = self._plan("musa") + self.assertIn("_musa_benchmark_compat_plugin", names) + # It rewrites collection-time markers, so nothing may inspect the item + # list before it runs. + self.assertEqual(names[0], "_musa_benchmark_compat_plugin") + + def test_compat_plugins_do_not_leak_to_other_backends(self): + for name in ("cuda", "rocm", "ascend"): + self.assertNotIn("_musa_benchmark_compat_plugin", self._plan(name), name) + def test_musa_same_core_without_cluster_fix(self): # _cudagraph_plugin is absent from the plan on every backend: the timing # plugin loads it, so its absence here says nothing about capture support.