diff --git a/README.md b/README.md index c34c64a..8ddbedc 100644 --- a/README.md +++ b/README.md @@ -4,7 +4,7 @@ 1. **可复现**:固定随机种子、指定 shape、固定 autotune config,把 A/B 两次运行之间的差异收敛到"编译器改动"这一个变量; 2. **可解释**:自动按 shape 收集每次运行实际使用的 ttgir(带可读的变体命名),供 IR 级 diff; -3. **口径统一**:cudagraph 计时消除 launch 开销,小 kernel 的对比不被 CPU 侧噪声淹没。 +3. **口径统一**:cudagraph 计时消除 launch 开销,小 kernel 的对比不被 CPU 侧噪声淹没;capture 前按 warmup 预算显式预热吸收 autotune/JIT,首轮即稳态、多轮一致。 ## 快速开始 @@ -28,7 +28,7 @@ softmax: 跑之前先 `nvidia-smi` 确认目标卡空闲——共享机器上别的任务会把 baseline 和被测两边一起等比拖慢,出一份看似自洽实则作废的数据。多卡机器上用 `CUDA_VISIBLE_DEVICES=<空闲卡号>` 明确选卡。 -脚本内固定了 `--level core --mode kernel` 和 `USE_FLAGTUNE=1`(FlagTune 扩展调优空间,见"为什么需要 replay"一节);需要改动这些口径时直接编辑 `run_pytest.sh` 中的 pytest 命令行。 +脚本内固定了 `--level core --mode kernel`;`USE_FLAGTUNE` 默认 `1`(FlagTune 扩展调优空间,见"为什么需要 replay"一节),可用环境变量覆盖(`USE_FLAGTUNE=0 bash run_pytest.sh` 走普通 autotune,跳过全量搜索,快速验证时用);需要改动其它口径时直接编辑 `run_pytest.sh` 中的 pytest 命令行。 ### 环境变量一览 @@ -39,6 +39,8 @@ softmax: | `SHAPE_FILE` | 空(用脚本内置 yaml) | shape yaml 路径 | | `FLAGGEMS_DIR` | `/workspace/FlagGems-dev` | FlagGems 仓库路径 | | `REPLAY_FROM` | 空(record 模式) | 指向某次历史 run 目录,replay 其 autotune 选择(见下) | +| `USE_FLAGTUNE` | `1` | `1` 用 FlagTune 扩展调优空间(首跑全量搜索、慢);`0` 走普通 autotune,跳过搜索、快速验证 | +| `FLAGGEMS_PERF_COLOR` | 空(按 tty 自动判断) | `always`/`never` 强制开/关终端颜色;`run.log` 始终为去色纯文本 | ## 输出目录结构 @@ -46,7 +48,7 @@ softmax: ``` runs/softmax_20260716-031752/ -├── run.log # 完整日志(含 SUCCESS 行的 latency/speedup 表) +├── run.log # 完整日志(含 SUCCESS 行的 latency/speedup 表;已去 ANSI 色的纯文本) ├── shapes.yaml # 本次实际使用的 shape(存档) ├── autotune_records/ # Triton autotune 选中的 config(供 replay) │ └── softmax.json @@ -77,13 +79,17 @@ bash run_pytest.sh # -> runs/_/ REPLAY_FROM=$PWD/runs/_ bash run_pytest.sh ``` +A/B 两侧的 `USE_FLAGTUNE` 必须取同值:该开关决定调优空间,两侧不一致时 replay 会大量 fallback,libtuner 持久缓存也各自独立命中,"同 config"前提不再成立。 + 对比 `run.log` 的 latency 表看性能差异;diff 两侧 `ttgir///` 下的同名文件看 IR 差异。 replay 的兜底行为:B 侧遇到记录中没有的 key、或记录的 config 在新编译器下编译失败时,自动回退到现场 autotune 并在 `run.log` 打 `AUTOTUNE_REPLAY_FALLBACK reason=...` 标记——出现该标记的测量点不再满足"同 config"前提,解读时注意。另外 replay 模式的 run 目录不产生 `autotune_records/`,后续 run 的 `REPLAY_FROM` 应始终指向最初 record 的那次 A 侧目录,不要链式指向 replay 产物。 -### cudagraph 回退与测量精度 +### cudagraph 计时的 warmup、回退与精度 -部分算子本身不支持 CUDA graph capture(测量函数内含 host 同步、动态显存分配、不合法的流操作等),这类算子会自动回退到普通 do_bench 计时,`run.log` 中打 `BENCHMARK_DIRECT_NO_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 被拒)。 回退本身不影响 A/B 公平性(两侧同一算子回退行为一致),但**回退口径的测量误差更大**:do_bench 每次迭代都走完整的 Python → launch 路径,kernel 越小,launch 开销和 CPU 侧抖动在数字里占比越高——亚毫秒级 kernel 上两种口径可差 2 倍以上,且行间波动更明显。解读这类算子的结果时: @@ -109,12 +115,14 @@ libtuner 的持久缓存何时失效:FlagGems kernel 源码改动、tune_confi | 插件 | 作用 | 何时关闭 | |-----|------|---------| +| `_device_guard_plugin` | 在 import flag_gems 之前直接用 torch 探测到 NVIDIA 卡就设 `GEMS_VENDOR=nvidia`,跳过 flag_gems 启动时 `nvidia-smi` 子进程探测(该探测在部分 fork 环境下会挂在 `wait4` 上导致 import 卡死) | 一般无需关:已设 `GEMS_VENDOR`/`FLAGGEMS_VENDOR` 等 env 时自动跳过,非 NVIDIA 卡上自动 no-op | | `_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)。无法 graph capture 的 kernel 自动回退并打 `BENCHMARK_DIRECT_NO_CUDAGRAPH` 标记 | 需要与他人的普通 do_bench 数据对齐时注释掉 | +| `_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 崩溃时打开 | @@ -126,11 +134,11 @@ libtuner 的持久缓存何时失效:FlagGems kernel 源码改动、tune_confi **Q: 第一次跑某算子特别慢?** -libtuner 算子首跑要做全量 sweep(mm 约 50 分钟),winner 持久化到 `~/.flaggems/config_cache/` 后,之后同 shape 秒级命中。 +两个来源:(1)默认 `USE_FLAGTUNE=1`,首跑要在 FlagTune 扩展空间做全量搜索(fused_marlin_moe_mxfp4 单进程可达数分钟甚至十几分钟,期间 GPU 满载、`run.log` 停在测试名不动属正常,不是卡死);(2)libtuner 算子首跑的全量 sweep(mm 约 50 分钟)。winner 持久化到 `~/.flaggems/config_cache/` 后同 shape 秒级命中(但该缓存会因源码/开关/GPU 变化失效,失效后又需重搜)。快速验证(不追求"各自最优 config"、只看功能/稳定性)时直接 `USE_FLAGTUNE=0 bash run_pytest.sh` 走普通 autotune,单 shape 通常 10 秒级出结果。 **Q: `--warmup/--iter` 要设吗?** -不用。cudagraph 计时路径下 warmup 参数本就不生效(内部自带预热),iter 默认 100ms 预算按 kernel 耗时自适应换算次数。 +不用。cudagraph 计时路径下,传入的 warmup 时间预算会被用来在 capture 前显式预热(先吸收 autotune/JIT 编译再稳态预热,稳定首轮,见"cudagraph 计时的 warmup、回退与精度"),iter 默认 100ms 预算按 kernel 耗时自适应换算次数。 **Q: ttgir 目录里某个 shape 少了文件?** diff --git a/_autotune_record_plugin.py b/_autotune_record_plugin.py index 7824d16..19f6782 100644 --- a/_autotune_record_plugin.py +++ b/_autotune_record_plugin.py @@ -115,8 +115,10 @@ def _compute_key(tuner: Any, args: Tuple[Any, ...], kwargs: Dict[str, Any]) -> O def _emit_marker(reason: str) -> None: # Marker lands in run.log (grep it to audit replay coverage); flush so it - # appears before the kernel runs. - print(f"AUTOTUNE_REPLAY_FALLBACK reason={reason}", flush=True) + # appears before the kernel runs. Yellow on the live terminal only — + # run_pytest.sh strips ANSI from run.log. + from _term_style import YELLOW, paint + print(paint(f"AUTOTUNE_REPLAY_FALLBACK reason={reason}", YELLOW), flush=True) def _record_run(original): @@ -235,16 +237,19 @@ def _dump_record(record_dir: str) -> None: path.write_text(json.dumps(payload, indent=2, ensure_ascii=False) + "\n") except Exception as exc: # Don't break the run if dump fails; the perf row is still produced. - print(f"[autotune-record-plugin] dump failed: {exc}", file=sys.stderr, flush=True) + from _term_style import tag + print(f"{tag('[autotune-record-plugin]')} dump failed: {exc}", + file=sys.stderr, flush=True) def pytest_configure(config): + from _term_style import tag 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("[autotune-record-plugin] error: both RECORD_DIR and REPLAY_DIR set; " - "ignoring both (no-op)", file=sys.stderr, flush=True) + print(f"{tag('[autotune-record-plugin]')} error: both RECORD_DIR and " + "REPLAY_DIR set; ignoring both (no-op)", file=sys.stderr, flush=True) return if not record_dir and not replay_dir: return @@ -265,11 +270,11 @@ def pytest_configure(config): # atexit (not sessionfinish): persist whatever was recorded even if an op # crash kills the session; a later replay run falls back for missing keys. atexit.register(_dump_record, record_dir) - print(f"[autotune-record-plugin] recording autotune configs to " + print(f"{tag('[autotune-record-plugin]')} recording autotune configs to " f"{_record_path(record_dir)} ({'+'.join(patched)})", file=sys.stderr, flush=True) else: loaded = _load_replay_dir(replay_dir) - print(f"[autotune-record-plugin] replaying {loaded} recorded entries from " + print(f"{tag('[autotune-record-plugin]')} replaying {loaded} recorded entries from " f"{_record_path(replay_dir)} ({'+'.join(patched)})", file=sys.stderr, flush=True) diff --git a/_bespoke_shape_plugin.py b/_bespoke_shape_plugin.py index bf32525..6e86748 100644 --- a/_bespoke_shape_plugin.py +++ b/_bespoke_shape_plugin.py @@ -27,7 +27,9 @@ def _load_yaml(path: str) -> dict: with open(path, "r") as f: return _yaml.safe_load(f) or {} except Exception as exc: - print(f"[bespoke-shape-plugin] failed to load {path}: {exc}", file=sys.stderr) + from _term_style import tag + print(f"{tag('[bespoke-shape-plugin]')} failed to load {path}: {exc}", + file=sys.stderr) return {} @@ -146,5 +148,7 @@ def pytest_collection_finish(session): if cls.__name__ == "CutlassScaledMMBenchmark" and _patch_cutlass(cls): covered.append("CutlassScaledMMBenchmark(mnk)") - print(f"[bespoke-shape-plugin] yaml-driven inputs for: {', '.join(covered) or 'none'}", + from _term_style import tag + print(f"{tag('[bespoke-shape-plugin]')} yaml-driven inputs for: " + f"{', '.join(covered) or 'none'}", file=sys.stderr, flush=True) diff --git a/_cudagraph_plugin.py b/_cudagraph_plugin.py index ac75b77..3b54a08 100644 --- a/_cudagraph_plugin.py +++ b/_cudagraph_plugin.py @@ -5,8 +5,10 @@ line out to time with plain do_bench). 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 -into the log. The warmup kwarg is dropped (do_bench_cudagraph warms up -internally). +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). """ from __future__ import annotations @@ -37,10 +39,55 @@ def _is_fatal_cuda_error(exc: Exception) -> bool: return any(m in msg for m in markers) +def _warmup_before_capture(fn, warmup_ms, grad_to_none): + """Warm fn up before cudagraph capture. + + do_bench_cudagraph only warms up 5 iterations internally — too few for the + MoE kernels, where the first calls trigger Triton autotune config + compilation, libtuner selection and lazy JIT/init. If that work leaks into + the captured graph or the first timed iteration the latency is unstable + run-to-run (most visibly on the small-M multi-kernel path). Absorb the + one-off compile with a throwaway run, then loop for the caller's warmup time + budget so the graph is captured at steady state. + + The loop is clamped to 5..200 iterations: the cap keeps sub-ms kernels from + spinning through a large budget (FlagGems' Config.warm_up defaults to + 1000 ms) — 200 iterations settles them fine in practice — and the floor + covers kernels whose single run exceeds the budget. If first-run jitter + reappears on a new op, raising the cap is the first knob to try. + """ + import time + + import torch + + def _run(): + if grad_to_none is not None: + for x in grad_to_none: + x.grad = None + fn() + + # First call absorbs autotune/JIT compilation (can take seconds); discard it. + _run() + torch.cuda.synchronize() + # Size the warmup loop from one timed steady-state run. dt_ms includes + # launch overhead, so n errs low for tiny kernels; warmup needn't be exact. + t0 = time.perf_counter() + _run() + torch.cuda.synchronize() + dt_ms = (time.perf_counter() - t0) * 1000.0 + budget = warmup_ms if (warmup_ms and warmup_ms > 0) else 25.0 + n = int(budget / dt_ms) if dt_ms > 0 else 25 + n = max(5, min(n, 200)) + for _ in range(n): + _run() + torch.cuda.synchronize() + + 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.""" op = os.environ.get("FLAGGEMS_PERF_CURRENT_OP", "") + phase = "sync" # which step rejected cudagraph, for the fallback marker try: # do_bench_cudagraph runs fn on a fresh side stream without syncing # with the default stream first; input tensors produced just before @@ -49,6 +96,10 @@ def _patched_do_bench(fn, warmup=25, rep=100, grad_to_none=None, # device trap on garbage indices). Sync before switching streams. import torch torch.cuda.synchronize() + phase = "warmup" + # Settle autotune/JIT before capture (rationale: _warmup_before_capture). + _warmup_before_capture(fn, warmup, grad_to_none) + phase = "capture" return _tt.do_bench_cudagraph( fn, rep=rep, @@ -57,13 +108,16 @@ def _patched_do_bench(fn, warmup=25, rep=100, grad_to_none=None, return_mode=return_mode, ) except Exception as exc: - # Re-raise genuine device failures; only graph-capture rejections fall back. + # Re-raise genuine device failures; anything else falls back to plain + # do_bench, with the failing phase (sync/warmup/capture) in the marker. if _is_fatal_cuda_error(exc): raise reason = type(exc).__name__ detail = " ".join(str(exc).split())[:160] + from _term_style import YELLOW, paint print( - f"BENCHMARK_DIRECT_NO_CUDAGRAPH op={op} reason={reason} detail={detail!r}", + paint(f"BENCHMARK_DIRECT_NO_CUDAGRAPH op={op} phase={phase} " + f"reason={reason} detail={detail!r}", YELLOW), flush=True, ) return _ORIGINAL_DO_BENCH( @@ -77,6 +131,7 @@ def _patched_do_bench(fn, warmup=25, rep=100, grad_to_none=None, def pytest_configure(config): + from _term_style import tag _tt.do_bench = _patched_do_bench - print("[cudagraph-plugin] triton.testing.do_bench → do_bench_cudagraph", + print(f"{tag('[cudagraph-plugin]')} triton.testing.do_bench → do_bench_cudagraph", file=sys.stderr, flush=True) diff --git a/_device_guard_plugin.py b/_device_guard_plugin.py new file mode 100644 index 0000000..f735a1a --- /dev/null +++ b/_device_guard_plugin.py @@ -0,0 +1,45 @@ +"""Guard against FlagGems' device-detection hang (benchmark-side fix). + +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%). + +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. +""" +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() diff --git a/_ir_meta_plugin.py b/_ir_meta_plugin.py index 9f09b7e..ebe3017 100644 --- a/_ir_meta_plugin.py +++ b/_ir_meta_plugin.py @@ -114,7 +114,8 @@ def _compile_listener(*, src, metadata, metadata_group, times, cache_hit): f, indent=1, sort_keys=True) os.replace(tmp, path) except Exception as exc: # never break compilation over a metadata dump - print(f"[ir-meta-plugin] sidecar dump failed: {exc}", file=sys.stderr) + from _term_style import tag + print(f"{tag('[ir-meta-plugin]')} sidecar dump failed: {exc}", file=sys.stderr) # --- shape + launch tracking ---------------------------------------------- @@ -202,7 +203,8 @@ def pytest_collection_finish(session): if own is not None: cls.get_input_iter = _wrap_input_iter(own) wrapped += 1 - print(f"[ir-meta-plugin] shape tracking wrapped on {wrapped} Benchmark classes", + from _term_style import tag + print(f"{tag('[ir-meta-plugin]')} shape tracking wrapped on {wrapped} Benchmark classes", file=sys.stderr, flush=True) @@ -312,7 +314,9 @@ def _dump_ttgir(cache_dir: str, dump_dir: str) -> None: cache, dest = pathlib.Path(cache_dir), pathlib.Path(dump_dir) if not cache.is_dir(): - print(f"[ir-meta-plugin] no cache dir {cache}; nothing to dump", file=sys.stderr) + from _term_style import tag + print(f"{tag('[ir-meta-plugin]')} no cache dir {cache}; nothing to dump", + file=sys.stderr) return dest.mkdir(parents=True, exist_ok=True) @@ -386,11 +390,12 @@ def _dump_ttgir(cache_dir: str, dump_dir: str) -> None: shapes = sorted({r[0] for r in rows if r[0] != "-"}) losers = sum(1 for r in rows if r[0] == "-") + from _term_style import DIM, paint print(f">>> [Dump] ttgir -> {dest} ({copied} files across {len(shapes)} shapes; " f"{losers} unused variants index-only; legend: naming.md)", flush=True) for s in shapes: n = sum(1 for r in rows if r[0] == s) - print(f">>> [Dump] {s}: {n}", flush=True) + print(paint(f">>> [Dump] {s}: {n}", DIM), flush=True) # --- registration ----------------------------------------------------------- @@ -419,10 +424,12 @@ def pytest_configure(config): _autotuner.Autotuner._bench = _wrap_bench(_autotuner.Autotuner._bench) atexit.register(_dump_ttgir, cache_dir, dump_dir) else: - print("[ir-meta-plugin] warning: dump dir set but TRITON_CACHE_DIR " + from _term_style import tag + print(f"{tag('[ir-meta-plugin]')} warning: dump dir set but TRITON_CACHE_DIR " "is not; ttgir dump disabled", file=sys.stderr, flush=True) - print("[ir-meta-plugin] compilation listener registered" + from _term_style import tag + print(f"{tag('[ir-meta-plugin]')} compilation listener registered" + (f"; shape/launch tracking on, ttgir dump -> {dump_dir}" if _dump_enabled else ""), file=sys.stderr, flush=True) diff --git a/_mm_cluster_fix_plugin.py b/_mm_cluster_fix_plugin.py index de453c4..8fe0dbf 100644 --- a/_mm_cluster_fix_plugin.py +++ b/_mm_cluster_fix_plugin.py @@ -92,22 +92,26 @@ def _warmup_load_backend(): torch.mm(a, b) torch.cuda.synchronize() except (ImportError, AttributeError) as e: - print(f"[mm-cluster-fix-plugin] warmup skipped (benign): {e!r}", + from _term_style import tag + print(f"{tag('[mm-cluster-fix-plugin]')} warmup skipped (benign): {e!r}", file=sys.stderr, flush=True) except Exception as e: # noqa: BLE001 # e.g. a CUDA error — do NOT hide it; the harness needs to see it. - print(f"[mm-cluster-fix-plugin] WARNING: warmup mm failed unexpectedly: {e!r}", + from _term_style import tag + print(f"{tag('[mm-cluster-fix-plugin]')} WARNING: warmup mm failed " + f"unexpectedly: {e!r}", file=sys.stderr, flush=True) def pytest_configure(config): + from _term_style import tag _patch_loaded() _warmup_load_backend() _patch_loaded() guarded = _guarded_modules() if guarded: print( - "[mm-cluster-fix-plugin] guarded cluster_remote_mm_scenario " + f"{tag('[mm-cluster-fix-plugin]')} guarded cluster_remote_mm_scenario " f"(odd N-tile unmasked OOB) in: {guarded}", file=sys.stderr, flush=True, @@ -125,7 +129,7 @@ def pytest_configure(config): else: # Non-Hopper: the cluster path never runs, so an unpatched state is fine. print( - "[mm-cluster-fix-plugin] no cluster_remote_mm_scenario found; " + f"{tag('[mm-cluster-fix-plugin]')} no cluster_remote_mm_scenario found; " "non-Hopper device, guard not needed (no-op).", file=sys.stderr, flush=True, diff --git a/_pretty_report_plugin.py b/_pretty_report_plugin.py new file mode 100644 index 0000000..70459b0 --- /dev/null +++ b/_pretty_report_plugin.py @@ -0,0 +1,148 @@ +"""pytest plugin: readable, colorized benchmark result table. + +Replaces FlagGems' BenchmarkResult.__str__ (benchmark/consts.py), whose rows +print every input's torch.Size(...) on every line — for the MoE ops that is 11 +tensors and ~400 chars per row, ~90% of it identical across rows. This plugin: + +- folds inputs identical across all rows into one legend line above the table, + so each row's Size Detail keeps only what varies between rows; +- shapes keep their original torch.Size([...]) spelling (no compaction); +- colors the table title and Status (SUCCESS green / FAILED red); +- keeps the literal "SUCCESS"/"FAILED" words and upstream column titles, so + run.log greps and __str__-wrapping hooks (vllm column renames) still work. + +Color only reaches the live terminal: run_pytest.sh strips ANSI from run.log, +and _term_style disables color for non-tty stdout unless forced. +""" + +from __future__ import annotations + +import sys + +from _term_style import BOLD, CYAN, DIM, GREEN, RED, paint, tag + + +def _shape_rows(metrics_list) -> list: + """One list of per-input cell strings per metrics row (upstream spelling).""" + rows = [] + for m in metrics_list: + sd = m.shape_detail + if isinstance(sd, (list, tuple)): + rows.append([str(e) for e in sd]) + else: + rows.append([str(sd) if sd is not None else "N/A"]) + return rows + + +def _factor_static(rows): + """Fold columns identical across all rows into a legend. + + Only kicks in when it actually helps: >=2 rows of equal arity >=3, with + >=2 static positions and at least one varying position left. Returns + (legend_cells_or_None, per_row_varying_strings_or_None). + """ + if len(rows) >= 2: + arities = {len(r) for r in rows} + if len(arities) == 1: + n = next(iter(arities)) + if n >= 3: + static = {i for i in range(n) if len({r[i] for r in rows}) == 1} + if 2 <= len(static) < n: + legend = [rows[0][i] for i in sorted(static)] + kept = [ + ", ".join(r[i] for i in range(n) if i not in static) + for r in rows + ] + return legend, kept + return None, None + + +def _num_cell(value, fmt: str, width: int) -> str: + s = f"{value:{fmt}}" if value is not None else "N/A" + return f"{s:>{width}}" + + +def _pretty_str(self) -> str: + metrics = self.result or [] + title = ( + "\n" + + paint(f"Operator: {self.op_name}", BOLD, CYAN) + + paint(f" (dtype={self.dtype}, mode={self.mode}, level={self.level})", DIM) + + "\n" + ) + if not metrics: + return title + "(no results)\n" + + # Same optional-column conditions as upstream. + with_tflops = bool(metrics[0].tflops) + with_gbps = metrics[0].gbps is not None + + legend, varying = _factor_static(_shape_rows(metrics)) + legend_line = "" + if legend: + legend_line = ( + paint( + f"{len(legend)} inputs identical across all rows: " + + ", ".join(legend), + DIM, + ) + + "\n" + ) + shape_cells = varying + else: + # No folding possible: print the full shape_detail, upstream style. + shape_cells = [ + str(m.shape_detail) if m.shape_detail is not None else "N/A" + for m in metrics + ] + + cols = [ + ("Status", 8, "<"), + ("Torch Latency (ms)", 19, ">"), + ("Gems Latency (ms)", 18, ">"), + ("Gems Speedup", 13, ">"), + ] + if with_tflops: + cols.append(("TFLOPS", 13, ">")) + if with_gbps: + cols.append(("Torch GBPS", 12, ">")) + cols.append(("Gems GBPS", 12, ">")) + head_plain = ( + " ".join(f"{name:{align}{width}}" for name, width, align in cols) + + " Size Detail" + ) + header = paint(head_plain, BOLD) + "\n" + paint("-" * len(head_plain), DIM) + "\n" + + lines = [] + for m, shape_cell in zip(metrics, shape_cells): + ok = m.error_msg is None + cells = [paint(f"{'SUCCESS' if ok else 'FAILED':<8}", GREEN if ok else RED)] + cells.append(_num_cell(m.latency_base, ".6f", 19)) + cells.append(_num_cell(m.latency, ".6f", 18)) + cells.append(_num_cell(m.speedup, ".3f", 13)) + if with_tflops: + cells.append(_num_cell(m.tflops, ".3f", 13)) + if with_gbps: + cells.append(_num_cell(m.gbps_base, ".3f", 12)) + cells.append(_num_cell(m.gbps, ".3f", 12)) + lines.append(" ".join(cells) + " " + shape_cell) + + return title + legend_line + header + "\n".join(lines) + "\n" + + +def pytest_configure(config): + try: + from benchmark import consts as fg_consts + except Exception as exc: + print( + f"{tag('[pretty-report-plugin]')} disabled " + f"(cannot import benchmark.consts: {exc})", + file=sys.stderr, flush=True, + ) + return + fg_consts.BenchmarkResult.__str__ = _pretty_str + print( + f"{tag('[pretty-report-plugin]')} BenchmarkResult table -> folded + color " + "(SUCCESS/FAILED words and column titles unchanged)", + file=sys.stderr, flush=True, + ) diff --git a/_seed_plugin.py b/_seed_plugin.py index d0e7144..9da42d6 100644 --- a/_seed_plugin.py +++ b/_seed_plugin.py @@ -27,13 +27,16 @@ def pytest_configure(config): try: import torch except Exception as exc: # torch missing should never happen here, stay safe - print(f"[seed-plugin] torch unavailable, seeded {'+'.join(seeded)} only: {exc}", + from _term_style import tag + print(f"{tag('[seed-plugin]')} torch unavailable, " + f"seeded {'+'.join(seeded)} only: {exc}", file=sys.stderr) return torch.manual_seed(seed) if torch.cuda.is_available(): torch.cuda.manual_seed_all(seed) seeded.append("torch") - print(f"[seed-plugin] manual_seed({seed}) for {'+'.join(seeded)} " + from _term_style import tag + print(f"{tag('[seed-plugin]')} manual_seed({seed}) for {'+'.join(seeded)} " "— reproducible benchmark inputs/cases", file=sys.stderr, flush=True) diff --git a/_shape_inject_plugin.py b/_shape_inject_plugin.py index b145250..e6f4847 100644 --- a/_shape_inject_plugin.py +++ b/_shape_inject_plugin.py @@ -26,7 +26,9 @@ def _load_yaml(path: str) -> dict: with open(path, "r") as f: return _yaml.safe_load(f) or {} except Exception as exc: - print(f"[shape-inject-plugin] failed to load {path}: {exc}", file=sys.stderr) + from _term_style import tag + print(f"{tag('[shape-inject-plugin]')} failed to load {path}: {exc}", + file=sys.stderr) return {} @@ -76,8 +78,9 @@ def _patch_init_user_config(): cur_homogeneous = len(cur_arities) == 1 yaml_homogeneous = len(yaml_arities) == 1 if cur_homogeneous and yaml_homogeneous and cur_arities != yaml_arities: + from _term_style import tag print( - f"[shape-inject-plugin] skip override for op_name={self.op_name!r}: " + f"{tag('[shape-inject-plugin]')} skip override for op_name={self.op_name!r}: " f"subclass produced homogeneous arity {next(iter(cur_arities))}, " f"yaml has homogeneous arity {next(iter(yaml_arities))} " f"— subclass normalization preserved", @@ -93,7 +96,8 @@ def _patch_init_user_config(): def pytest_configure(config): + from _term_style import tag _patch_init_user_config() - print("[shape-inject-plugin] Benchmark.init_user_config patched: " + print(f"{tag('[shape-inject-plugin]')} Benchmark.init_user_config patched: " "yaml shapes now win over hardcoded subclass set_shapes()", file=sys.stderr, flush=True) diff --git a/_shape_iter_inject_plugin.py b/_shape_iter_inject_plugin.py index 2e4f0c4..77a02ce 100644 --- a/_shape_iter_inject_plugin.py +++ b/_shape_iter_inject_plugin.py @@ -31,7 +31,9 @@ def _load_yaml(path: str) -> dict: with open(path, "r") as f: return _yaml.safe_load(f) or {} except Exception as exc: - print(f"[shape-iter-inject-plugin] failed to load {path}: {exc}", file=sys.stderr) + from _term_style import tag + print(f"{tag('[shape-iter-inject-plugin]')} failed to load {path}: {exc}", + file=sys.stderr) return {} @@ -95,6 +97,7 @@ def pytest_collection_finish(session): cls.get_input_iter = _make_patched(own) patched += 1 - print(f"[shape-iter-inject-plugin] get_input_iter redirected on {patched} " + from _term_style import tag + print(f"{tag('[shape-iter-inject-plugin]')} get_input_iter redirected on {patched} " f"hardcoded-shape Benchmark classes; yaml shapes now win", file=sys.stderr, flush=True) diff --git a/_term_style.py b/_term_style.py new file mode 100644 index 0000000..c948fee --- /dev/null +++ b/_term_style.py @@ -0,0 +1,61 @@ +"""Shared ANSI color helper for zl_bench plugins and scripts. + +Color policy, decided once at import (env is fixed before python starts): +1. FLAGGEMS_PERF_COLOR=always|never wins — run_pytest.sh sets `always` when + *its* stdout is a tty, because the tee pipeline hides the tty from python; +2. FORCE_COLOR forces on (or off when set to 0/false, node semantics), taking + precedence over NO_COLOR; +3. NO_COLOR (non-empty) forces off; +4. otherwise: on iff stdout is a tty. + +run.log stays plain either way: run_pytest.sh strips ANSI from the log file +after the run. +""" + +from __future__ import annotations + +import os +import sys + +_FALSY = ("", "0", "false", "no", "off", "none") + + +def _color_enabled() -> bool: + mode = os.environ.get("FLAGGEMS_PERF_COLOR", "").strip().lower() + if mode in ("always", "1", "yes", "on"): + return True + if mode in ("never", "0", "no", "off"): + return False + force = os.environ.get("FORCE_COLOR") + if force is not None: + return force.strip().lower() not in _FALSY + if os.environ.get("NO_COLOR"): + return False + try: + return sys.stdout.isatty() + except Exception: + return False + + +COLOR = _color_enabled() + +RESET = "\033[0m" +BOLD = "\033[1m" +DIM = "\033[2m" +RED = "\033[31m" +GREEN = "\033[32m" +YELLOW = "\033[33m" +CYAN = "\033[36m" +BRIGHT_BLUE = "\033[94m" + + +def paint(text: str, *codes: str) -> str: + """Wrap text in ANSI codes when color is enabled; identity otherwise.""" + if not COLOR or not codes: + return text + return "".join(codes) + text + RESET + + +def tag(label: str) -> str: + """Bright-blue '[xxx-plugin]' prefix for plugin log lines.""" + return paint(label, BRIGHT_BLUE) diff --git a/run_pytest.sh b/run_pytest.sh index daf9b82..79017bc 100644 --- a/run_pytest.sh +++ b/run_pytest.sh @@ -1,18 +1,27 @@ #!/bin/bash -# Single-operator perf benchmark (manual debug). -# Edit OP/OP_FILE/INLINE_YAML below (or override via env: OP=... SHAPE_FILE=... ./run_pytest.sh). -# All artifacts go under runs/_/: run.log, shapes.yaml, -# autotune_records/, ttgir/. +# 单算子性能测试入口(手动调试用)。 +# 用法:改下方【配置区】,或用环境变量覆盖,例如: +# OP=softmax OP_FILE=softmax SHAPE_FILE=my_shapes.yaml bash run_pytest.sh +# 产物统一落在 runs/_<时间戳>/:run.log、shapes.yaml、autotune_records/、ttgir/。 set -euo pipefail -OP="${OP:-fused_marlin_moe_mxfp4}" -OP_FILE="${OP_FILE:-fused_marlin_moe}" # benchmark file is test_.py -SHAPE_FILE="${SHAPE_FILE:-}" +# ============================== 配置区 ============================== +# 每项均可用同名环境变量覆盖,详见 README「环境变量一览」。 -# Inline shape yaml (used when SHAPE_FILE is empty); top-level key must be the op name. +# 被测算子:OP=测试函数名(去掉 test_ 前缀);OP_FILE=benchmark 文件名 +# (对应 $FLAGGEMS_DIR/benchmark/test_.py::test_) +OP="${OP:-fused_marlin_moe_mxfp4}" +OP_FILE="${OP_FILE:-fused_marlin_moe}" + +# FlagGems 仓库路径 +FLAGGEMS_DIR="${FLAGGEMS_DIR:-/workspace/FlagGems-dev}" + +# shape 来源:SHAPE_FILE 非空则用该 yaml;为空则用下方内置 yaml。 +# yaml 顶层 key 必须是 op 名。 +SHAPE_FILE="${SHAPE_FILE:-}" read -r -d '' INLINE_YAML <<'YAML' || true fused_marlin_moe_mxfp4: - # 4 MoE models x 4 token counts (M=1,16,64,256) = 16 shapes + # 4 个 MoE 模型 x 4 档 token 数(M=1,16,64,256)= 16 组 shape shapes: # Mixtral (E=8) - [1, 8, 4096, 14336, 2] @@ -37,16 +46,20 @@ fused_marlin_moe_mxfp4: shape_desc: "num_tokens, num_experts, hidden_size, intermediate_size, topk" YAML -SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" -FLAGGEMS_DIR="${FLAGGEMS_DIR:-/workspace/FlagGems-dev}" -TEST_FILE="$FLAGGEMS_DIR/benchmark/test_${OP_FILE}.py::test_${OP}" +# 调优空间:1=FlagTune 扩展空间(首跑全量搜索、慢);0=普通 autotune(快速验证) +USE_FLAGTUNE="${USE_FLAGTUNE:-1}" -OUT_DIR="$SCRIPT_DIR/runs/${OP}_$(date +%Y%m%d-%H%M%S)" -LOG_FILE="$OUT_DIR/run.log" -mkdir -p "$OUT_DIR" +# autotune record/replay(见 _autotune_record_plugin.py): +# - 空:record 模式,本次选中的 config 记录到 $OUT_DIR/autotune_records/.json; +# - 指向某次历史 runs/_<时间戳> 目录:replay 该次记录(A/B 两侧同 config)。 +REPLAY_FROM="${REPLAY_FROM:-}" -export PYTHONPATH="$SCRIPT_DIR${PYTHONPATH:+:$PYTHONPATH}" +# 终端颜色:always/never 强制开/关;为空则按 tty 自动判断(run.log 始终去色) +FLAGGEMS_PERF_COLOR="${FLAGGEMS_PERF_COLOR:-}" + +# pytest 插件列表(可按需注释停用;各插件作用见 README「插件说明」) PLUGINS=( + -p _device_guard_plugin -p _seed_plugin -p _shape_inject_plugin -p _shape_iter_inject_plugin @@ -54,55 +67,88 @@ PLUGINS=( # -p _mm_cluster_fix_plugin -p _autotune_record_plugin -p _cudagraph_plugin + -p _pretty_report_plugin -p _ir_meta_plugin ) -# Autotune record/replay (see _autotune_record_plugin.py): -# - default: record chosen configs to $OUT_DIR/autotune_records/.json -# - REPLAY_FROM=_ dir>: replay that run's recorded -# configs instead (for A/B runs where both sides must use the same config). -REPLAY_FROM="${REPLAY_FROM:-}" + +# ============================== 执行逻辑 ============================== + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +TEST_FILE="$FLAGGEMS_DIR/benchmark/test_${OP_FILE}.py::test_${OP}" + +OUT_DIR="$SCRIPT_DIR/runs/${OP}_$(date +%Y%m%d-%H%M%S)" +LOG_FILE="$OUT_DIR/run.log" +mkdir -p "$OUT_DIR" + +# 颜色决策:下方 tee 管道会让 python 侧看不到 tty,所以在这里判断, +# 并通过 FLAGGEMS_PERF_COLOR 下传给各插件(_term_style.py 消费)。 +if [[ -z "$FLAGGEMS_PERF_COLOR" && -t 1 && -z "${NO_COLOR:-}" ]]; then + FLAGGEMS_PERF_COLOR=always +fi +export FLAGGEMS_PERF_COLOR +if [[ "$FLAGGEMS_PERF_COLOR" == "always" ]]; then + C_RED=$'\033[31m'; C_GREEN=$'\033[32m'; C_YELLOW=$'\033[33m' + C_DIM=$'\033[2m'; C_BOLD=$'\033[1m'; C_RESET=$'\033[0m' + PYTEST_COLOR=(--color=yes) +else + C_RED='' C_GREEN='' C_YELLOW='' C_DIM='' C_BOLD='' C_RESET='' + PYTEST_COLOR=() +fi + +export PYTHONPATH="$SCRIPT_DIR${PYTHONPATH:+:$PYTHONPATH}" +export FLAGGEMS_PERF_CURRENT_OP="$OP" +export PYTHONUNBUFFERED=1 # 实时输出不缓冲 + +# record/replay 二选一,通过环境变量激活对应模式 if [[ -n "$REPLAY_FROM" ]]; then AUTOTUNE_ENV="FLAGGEMS_PERF_AUTOTUNE_REPLAY_DIR=$REPLAY_FROM/autotune_records" [[ -f "$REPLAY_FROM/autotune_records/$OP.json" ]] || \ - echo ">>> warning: $REPLAY_FROM/autotune_records/$OP.json not found; replay will fall back to autotune" >&2 + echo "${C_YELLOW}>>> warning: $REPLAY_FROM/autotune_records/$OP.json 不存在;replay 将回退为现场 autotune${C_RESET}" >&2 + MODE_DESC="replay($REPLAY_FROM)" else AUTOTUNE_ENV="FLAGGEMS_PERF_AUTOTUNE_RECORD_DIR=$OUT_DIR/autotune_records" mkdir -p "$OUT_DIR/autotune_records" + MODE_DESC=record fi -export FLAGGEMS_PERF_CURRENT_OP="$OP" -# Resolve the shape file: use SHAPE_FILE, or write the inline yaml to a temp file. +# 解析 shape 文件:SHAPE_FILE 为空时把内置 yaml 写到临时文件 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" # archive the shape used - -export PYTHONUNBUFFERED=1 # unbuffered live output +cp -f "$SHAPE_FILE" "$OUT_DIR/shapes.yaml" # 存档本次实际使用的 shape +status=0 { - # _ir_meta_plugin dumps organized ttgir (only variants actually launched - # outside the autotune sweep) into TTGIR_DUMP_DIR at process exit, so the - # dump happens even if a CUDA crash kills the run. Layout/legend: see - # /naming.md and index.tsv (which also lists unused sweep losers). + echo "${C_BOLD}>>> op=$OP mode=$MODE_DESC USE_FLAGTUNE=$USE_FLAGTUNE${C_RESET}" + echo "${C_DIM}>>> out=$OUT_DIR${C_RESET}" + + # _ir_meta_plugin 在进程退出时把"实际被使用"的变体的 ttgir 按 shape 整理落盘 + # (挂在 atexit 上,CUDA crash 后已编译部分仍可拿到)。 + # 目录结构与命名图例见 /naming.md 和 index.tsv(后者也记录落选的 sweep 变体)。 CACHE_DIR="$OUT_DIR/.triton_cache" rm -rf "$CACHE_DIR"; mkdir -p "$CACHE_DIR" status=0 TRITON_CACHE_DIR="$CACHE_DIR" \ FLAGGEMS_PERF_TTGIR_DUMP_DIR="$OUT_DIR/ttgir" \ env "$AUTOTUNE_ENV" \ - USE_FLAGTUNE=1 python -u -m pytest -s "$TEST_FILE" \ - "${PLUGINS[@]}" \ + USE_FLAGTUNE=$USE_FLAGTUNE python -u -m pytest -s "$TEST_FILE" \ + "${PLUGINS[@]}" "${PYTEST_COLOR[@]}" \ --shape_file "$SHAPE_FILE" \ --level core --mode kernel || status=$? rm -rf "$CACHE_DIR" if (( status == 0 )); then - echo ">>> done. outputs in $OUT_DIR" + echo "${C_GREEN}>>> done. outputs in $OUT_DIR${C_RESET}" else - echo ">>> FAILED (pytest exit $status). partial outputs in $OUT_DIR" + echo "${C_RED}>>> FAILED (pytest exit $status). partial outputs in $OUT_DIR${C_RESET}" fi exit "$status" -} 2>&1 | tee "$LOG_FILE" +} 2>&1 | tee "$LOG_FILE" || status=$? + +# 终端保留颜色;落盘的 run.log 去掉 ANSI 转义,保证 grep/diff 面对纯文本 +# (Ctrl-C 中断时会跳过去色,仅影响观感)。 +sed -i -E $'s/\x1b\\[[0-9;]*[A-Za-z]//g' "$LOG_FILE" +exit "$status"