diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..2fcb51f --- /dev/null +++ b/.gitignore @@ -0,0 +1,2 @@ +runs/ +__pycache__/ diff --git a/README.md b/README.md index 60d34d3..c34c64a 100644 --- a/README.md +++ b/README.md @@ -1,2 +1,143 @@ -# zl_bench +# zl_bench — FlagGems 单算子性能测试工具 +针对编译器(FlagTree)改动做算子级 A/B 性能对比的 pytest 插件集 + 驱动脚本。在 FlagGems benchmark 体系之上解决三个问题: + +1. **可复现**:固定随机种子、指定 shape、固定 autotune config,把 A/B 两次运行之间的差异收敛到"编译器改动"这一个变量; +2. **可解释**:自动按 shape 收集每次运行实际使用的 ttgir(带可读的变体命名),供 IR 级 diff; +3. **口径统一**:cudagraph 计时消除 launch 开销,小 kernel 的对比不被 CPU 侧噪声淹没。 + +## 快速开始 + +```bash +# 默认算子(fused_marlin_moe_mxfp4,内置 16 组 MoE shape) +bash run_pytest.sh + +# 换算子:OP=测试函数名去掉 test_ 前缀,OP_FILE=benchmark 文件名去掉 test_ 前缀/.py 后缀 +# (benchmark 文件为 $FLAGGEMS_DIR/benchmark/test_.py::test_) +OP=softmax OP_FILE=softmax SHAPE_FILE=my_shapes.yaml bash run_pytest.sh +``` + +shape yaml 顶层 key 必须是 op 名: + +```yaml +softmax: + shapes: + - [1024, 1024] + - [64, 512, 512] +``` + +跑之前先 `nvidia-smi` 确认目标卡空闲——共享机器上别的任务会把 baseline 和被测两边一起等比拖慢,出一份看似自洽实则作废的数据。多卡机器上用 `CUDA_VISIBLE_DEVICES=<空闲卡号>` 明确选卡。 + +脚本内固定了 `--level core --mode kernel` 和 `USE_FLAGTUNE=1`(FlagTune 扩展调优空间,见"为什么需要 replay"一节);需要改动这些口径时直接编辑 `run_pytest.sh` 中的 pytest 命令行。 + +### 环境变量一览 + +| 变量 | 默认 | 说明 | +|----------------|---------------------------|--------------------------------------------------| +| `OP` | `fused_marlin_moe_mxfp4` | 测试函数名(`test_` 之后的部分) | +| `OP_FILE` | `fused_marlin_moe` | benchmark 文件名(`test_` 与 `.py` 之间的部分) | +| `SHAPE_FILE` | 空(用脚本内置 yaml) | shape yaml 路径 | +| `FLAGGEMS_DIR` | `/workspace/FlagGems-dev` | FlagGems 仓库路径 | +| `REPLAY_FROM` | 空(record 模式) | 指向某次历史 run 目录,replay 其 autotune 选择(见下) | + +## 输出目录结构 + +每次运行产出 `runs/_<时间戳>/`: + +``` +runs/softmax_20260716-031752/ +├── run.log # 完整日志(含 SUCCESS 行的 latency/speedup 表) +├── shapes.yaml # 本次实际使用的 shape(存档) +├── autotune_records/ # Triton autotune 选中的 config(供 replay) +│ └── softmax.json +└── ttgir/ # 按 shape 分组的 IR 落盘 + ├── 1024x1024/ + │ └── softmax_kernel_inner/ + │ ├── w4s3__bf16.ttgir + │ ├── w4s3__fp16.ttgir + │ └── w4s3__fp32.ttgir + ├── 4096x4096/... + ├── index.tsv # 每个编译变体的完整参数、启动次数、cache hash + └── naming.md # 文件名缩写图例 +``` + +ttgir 关键设计: + +- **只落盘"实际使用"的变体**:autotune sweep 中测过但落选的 config 不拷贝(只在 index.tsv 里留 `sweep loser` 记录)。mm 这类 sweep 上千个变体的算子,最终只留真正被选中的几个文件。 +- **命名 = 组内有区分度的参数**:同一 (shape, kernel) 组内取值相同的 constexpr 不进文件名;多词参数缩写为首字母(`BLOCK_SIZE_M` → `BSM`,图例见 naming.md);同名冲突依次用 dtype(`__fp16`)、参数对齐特化(`__EMdiv16`)、hash 前缀消歧。 +- **崩溃也能拿到 IR**:dump 挂在进程 atexit 上,CUDA crash 后仍会落盘已编译部分。 + +## A/B 对比测试标准流程 + +```bash +# A 侧(基线编译器):正常跑,自动 record autotune 选择 +bash run_pytest.sh # -> runs/_/ + +# B 侧(改动后编译器):replay A 侧的 config,保证两侧同 config +REPLAY_FROM=$PWD/runs/_ bash run_pytest.sh +``` + +对比 `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 回退与测量精度 + +部分算子本身不支持 CUDA graph capture(测量函数内含 host 同步、动态显存分配、不合法的流操作等),这类算子会自动回退到普通 do_bench 计时,`run.log` 中打 `BENCHMARK_DIRECT_NO_CUDAGRAPH` 标记(含具体原因)。 + +回退本身不影响 A/B 公平性(两侧同一算子回退行为一致),但**回退口径的测量误差更大**:do_bench 每次迭代都走完整的 Python → launch 路径,kernel 越小,launch 开销和 CPU 侧抖动在数字里占比越高——亚毫秒级 kernel 上两种口径可差 2 倍以上,且行间波动更明显。解读这类算子的结果时: + +- 小 shape 行的绝对值和小幅(<10%)差异不要过度解读,优先看大 shape 行; +- 需要更高置信度时,同一配置多跑几次取中位数,或对该算子直接注释掉 `_cudagraph_plugin` 统一用 do_bench 口径(消除同表混两种口径的问题)。 + +回退是按测量点发生的,同一份 latency 表里可能混有两种口径的行;若不确定,先 `grep BENCHMARK_DIRECT_NO_CUDAGRAPH run.log` 确认哪些行是回退口径再下结论。 + +### 为什么需要 replay(以及它管不到什么) + +本项目涉及两层调优机制,对 A/B 的影响不同: + +| 机制 | 结果存储 | 编译器改动后 | A/B 风险 | +|-----|---------|------------|----------| +| Triton `@triton.autotune` | 仅进程内存 | 每次进程重新 sweep | 计时噪声可能让 A/B 选中**不同 config** → 用 REPLAY_FROM 固定 | +| FlagGems `@libtuner`(含 FlagTune 扩展空间) | `~/.flaggems/config_cache/*.db`(sqlite,跨进程持久) | **不失效**(表名只含 kernel 源码与 config 空间的 hash),A/B 自动命中同一 winner | 反向风险:B 侧沿用 A 侧选的旧 winner,测的是"旧 config 下的编译器差异"而非"各自最优" | + +libtuner 的持久缓存何时失效:FlagGems kernel 源码改动、tune_configs.yaml / expand yaml / `USE_FLAGTUNE` 开关变化、Triton 大版本或 GPU 型号变化。如果需要"各自最优"口径(让 B 侧重新 sweep),删除 sqlite 中对应表,或设 `FLAGGEMS_DB_URL` 指向一次性文件。**两种口径都合理,报告结论时注明用的哪种。** + +## 插件说明 + +脚本通过 `-p` 加载以下插件(`run_pytest.sh` 的 `PLUGINS` 数组,可按需注释): + +| 插件 | 作用 | 何时关闭 | +|-----|------|---------| +| `_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 数据对齐时注释掉 | +| `_ir_meta_plugin` | 编译期记录每个变体的 constexpr/签名/特化;launch 钩子统计每个 (kernel, shape) 的实际使用;退出时按上述结构 dump ttgir | 不关 | +| `_mm_cluster_fix_plugin` | Hopper fp16 mm cluster kernel 越界崩溃的运行时规避 | 默认注释;测 fp16 mm 崩溃时打开 | + +## 常见问题 + +**Q: latency 和别人跑的差很多?** + +先检查 GPU 是否被共占(`nvidia-smi`);再确认对方是否开了 cudagraph——M=1 这类小 shape 下 launch 开销占比大,两种口径可差 2 倍以上,大 shape 基本一致。 + +**Q: 第一次跑某算子特别慢?** + +libtuner 算子首跑要做全量 sweep(mm 约 50 分钟),winner 持久化到 `~/.flaggems/config_cache/` 后,之后同 shape 秒级命中。 + +**Q: `--warmup/--iter` 要设吗?** + +不用。cudagraph 计时路径下 warmup 参数本就不生效(内部自带预热),iter 默认 100ms 预算按 kernel 耗时自适应换算次数。 + +**Q: ttgir 目录里某个 shape 少了文件?** + +看 `index.tsv` 的 `launches` 列——没在该 shape 下真正启动过的变体不落盘。`run.log` 里的 `BENCHMARK_DIRECT_NO_CUDAGRAPH` / `AUTOTUNE_REPLAY_FALLBACK` 标记可解释异常回退。 + +## 依赖假设 + +- FlagGems benchmark 体系(`benchmark/base.py` 的 `Benchmark` 类、conftest 的 `--shape_file/--level/--mode` 选项); +- Triton 需支持 `knobs.compilation.listener`、`kernel_load_end_hook`、`launch_enter_hook`(当前 FlagTree 的 triton 3.6 满足); +- 插件通过 monkeypatch 挂钩上游内部结构,FlagGems/Triton 大版本升级后若行为异常,优先检查各插件 pytest_configure 输出的注册日志是否还正常打印。 diff --git a/_autotune_record_plugin.py b/_autotune_record_plugin.py new file mode 100644 index 0000000..7824d16 --- /dev/null +++ b/_autotune_record_plugin.py @@ -0,0 +1,275 @@ +"""pytest plugin: record or replay autotune configs across benchmark runs. + +When comparing two runs (e.g. two compiler builds), replay the first run's +recorded configs in the second: the autotuner's choice is pinned, so the +latency delta reflects the change under test rather than the autotuner picking +a different config under different IR. + +Modes (mutually exclusive, selected by env var; op identity from +FLAGGEMS_PERF_CURRENT_OP): +- FLAGGEMS_PERF_AUTOTUNE_RECORD_DIR: run full autotune, dump the chosen Config + per (kernel, key) to /.json at session end. +- FLAGGEMS_PERF_AUTOTUNE_REPLAY_DIR: pre-populate Autotuner.cache from + /.json, skipping the sweep. Missing key / non-compiling config / + missing json fall back to full autotune and emit an AUTOTUNE_REPLAY_FALLBACK + marker into the log. +""" + +from __future__ import annotations + +import atexit +import json +import os +import sys +import threading +from pathlib import Path +from typing import Any, Dict, Optional, Tuple + + +_RECORD_DIR_ENV = "FLAGGEMS_PERF_AUTOTUNE_RECORD_DIR" +_REPLAY_DIR_ENV = "FLAGGEMS_PERF_AUTOTUNE_REPLAY_DIR" +_OP_ENV = "FLAGGEMS_PERF_CURRENT_OP" + +# Module-global aggregates, accumulated as run() is called and dumped at exit. +# Lock guards merges in case a kernel uses threads internally. +_record_map: Dict[str, Dict[str, Dict[str, Any]]] = {} +_record_lock = threading.Lock() +_replay_map: Dict[str, Dict[str, Dict[str, Any]]] = {} +_dump_done = False + + +def _op_name() -> str: + return os.environ.get(_OP_ENV, "").strip() or "default" + + +def _record_path(record_dir: str) -> Path: + return Path(record_dir) / f"{_op_name()}.json" + + +def _kernel_id(tuner: Any) -> str: + # Stable across runs: source module + qualified name. Shapes/keys split + # by the inner key dict. + name = ( + getattr(tuner, "__name__", None) + or getattr(getattr(tuner, "base_fn", None), "__name__", None) + or getattr(getattr(tuner, "fn", None), "__name__", None) + or "unknown_kernel" + ) + module = getattr(getattr(tuner, "fn", None), "__module__", "") or "" + return f"{module}::{name}" + + +def _serialize_key(key: Tuple[Any, ...]) -> str: + # Per-element repr, JSON-encoded into a flat string usable as a JSON key; + # stable across runs given seeded RNG + yaml-fixed shapes. + return json.dumps([repr(x) for x in key]) + + +def _serialize_config(cfg: Any) -> Optional[Dict[str, Any]]: + # Drop configs with a pre_hook (un-serializable callable) rather than replay + # without it and compute incorrectly. + if getattr(cfg, "pre_hook", None) is not None: + return None + try: + return { + "kwargs": dict(cfg.kwargs), + "num_warps": cfg.num_warps, + "num_stages": cfg.num_stages, + "num_ctas": cfg.num_ctas, + "maxnreg": cfg.maxnreg, + } + except Exception: + return None + + +def _deserialize_config(entry: Dict[str, Any]) -> Any: + # Late import (triton loads during collection, after plugin registration). + from triton import Config + return Config( + kwargs=dict(entry["kwargs"]), + num_warps=entry.get("num_warps", 4), + num_stages=entry.get("num_stages", 3), + num_ctas=entry.get("num_ctas", 1), + maxnreg=entry.get("maxnreg"), + ) + + +def _compute_key(tuner: Any, args: Tuple[Any, ...], kwargs: Dict[str, Any]) -> Optional[Tuple[Any, ...]]: + # Reproduce Autotuner/LibTuner key computation; return None on any internals + # shift (caller then just loses replay for that call). + try: + arg_names = tuner.arg_names + nargs = dict(zip(arg_names, args)) + all_args = {**nargs, **kwargs} + _args = {k: v for k, v in all_args.items() if k in arg_names} + if hasattr(tuner, "get_key"): + return tuner.get_key(_args) + key = [_args[name] for name in tuner.keys if name in _args] + for _, arg in _args.items(): + if hasattr(arg, "dtype"): + key.append(str(arg.dtype)) + return tuple(key) + except Exception: + return None + + +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) + + +def _record_run(original): + # Snapshot self.cache before/after run() to capture the chosen config (works + # for both Autotuner and LibTuner). Single-config kernels skip the cache + # write, so record configs[0] explicitly for uniform replay. + def runner(self, *args, **kwargs): + try: + keys_before = set(self.cache.keys()) if hasattr(self.cache, "keys") else set() + except Exception: + keys_before = set() + result = original(self, *args, **kwargs) + try: + keys_after = set(self.cache.keys()) if hasattr(self.cache, "keys") else set() + except Exception: + keys_after = set() + new_keys = keys_after - keys_before + if not new_keys: + # No new entry: record configs[0] for single-config kernels (cache + # write bypassed); otherwise nothing to record (disk-cache hit). + if len(getattr(self, "configs", []) or []) == 1: + key = _compute_key(self, args, kwargs) + if key is not None and key not in self.cache: + cfg = self.configs[0] + entry = _serialize_config(cfg) + if entry is not None: + kid = _kernel_id(self) + with _record_lock: + bucket = _record_map.setdefault(kid, {}) + bucket[_serialize_key(key)] = entry + return result + kid = _kernel_id(self) + with _record_lock: + bucket = _record_map.setdefault(kid, {}) + for k in new_keys: + try: + cfg = self.cache[k] + except Exception: + continue + entry = _serialize_config(cfg) + if entry is None: + continue + bucket[_serialize_key(k)] = entry + return result + return runner + + +def _replay_run(original): + # Inject recorded configs before benchmarking; on key-missing or a recorded + # config failing to compile/launch, fall back to full autotune (+ marker). + def runner(self, *args, **kwargs): + injected_key = None + if len(getattr(self, "configs", []) or []) > 1: + kid = _kernel_id(self) + bucket = _replay_map.get(kid) + if bucket: + key = _compute_key(self, args, kwargs) + if key is not None and key not in self.cache: + rec = bucket.get(_serialize_key(key)) + if rec is None: + _emit_marker("key_missing") + else: + try: + cfg = _deserialize_config(rec) + self.cache[key] = cfg + injected_key = key + except Exception as exc: + _emit_marker(f"deserialize_{type(exc).__name__}") + try: + return original(self, *args, **kwargs) + except Exception as exc: + if injected_key is not None: + # Recorded config failed under the current build: drop it and + # let run() autotune. + try: + del self.cache[injected_key] + except Exception: + pass + _emit_marker(f"compile_{type(exc).__name__}") + return original(self, *args, **kwargs) + raise + return runner + + +def _load_replay_dir(replay_dir: str) -> int: + path = _record_path(replay_dir) + if not path.is_file(): + _emit_marker("record_missing") + return 0 + try: + data = json.loads(path.read_text()) + except Exception as exc: + _emit_marker(f"load_{type(exc).__name__}") + return 0 + if not isinstance(data, dict): + return 0 + count = 0 + for kid, bucket in data.items(): + if not isinstance(bucket, dict): + continue + _replay_map[kid] = dict(bucket) + count += len(bucket) + return count + + +def _dump_record(record_dir: str) -> None: + global _dump_done + if _dump_done: + return + _dump_done = True + path = _record_path(record_dir) + try: + path.parent.mkdir(parents=True, exist_ok=True) + with _record_lock: + payload = {kid: dict(bucket) for kid, bucket in _record_map.items()} + 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) + + +def pytest_configure(config): + 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) + return + if not record_dir and not replay_dir: + return + + import triton.runtime.autotuner as _autotuner + wrap = _record_run if record_dir else _replay_run + _autotuner.Autotuner.run = wrap(_autotuner.Autotuner.run) + patched = ["Autotuner"] + try: + from flag_gems.utils.libentry import LibTuner + if "run" in LibTuner.__dict__: + LibTuner.run = wrap(LibTuner.__dict__["run"]) + patched.append("LibTuner") + except Exception: + pass + + if record_dir: + # atexit (not sessionfinish): persist whatever was recorded even if an op + # crash kills the session; a later replay run falls back for missing keys. + atexit.register(_dump_record, record_dir) + print(f"[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 " + f"{_record_path(replay_dir)} ({'+'.join(patched)})", + file=sys.stderr, flush=True) diff --git a/_bespoke_shape_plugin.py b/_bespoke_shape_plugin.py new file mode 100644 index 0000000..bf32525 --- /dev/null +++ b/_bespoke_shape_plugin.py @@ -0,0 +1,150 @@ +"""pytest plugin: make the shape yaml authoritative for ops whose inputs are +*bespoke* (built inline / from TestParam objects / a sampled kit) and so are +unreachable by `_shape_inject_plugin` and `_shape_iter_inject_plugin`. + +A per-op registry; each handler REUSES the class's own input construction and +only swaps the shape source for our yaml (so we never measure a different +kernel). Covered: upsample_bicubic2d_aa_backward, flash_mla_sparse_fwd, +fused_deepseek_..._quant_insert, cutlass_scaled_mm. Ops with no yaml entry fall +through unchanged. Couples to upstream internals; a mismatch surfaces as a loud +benchmark FAIL, not silent corruption. +""" + +from __future__ import annotations + +import os +import sys + +import yaml as _yaml + +_YAML_CACHE: dict | None = None + + +def _load_yaml(path: str) -> dict: + if not path or not os.path.isfile(path): + return {} + try: + 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) + return {} + + +def _yaml_shapes(op_name): + entry = (_YAML_CACHE or {}).get(op_name) + if isinstance(entry, dict): + shapes = entry.get("shapes") + if shapes: + return shapes + return None + + +# --- per-op builders: reuse the class's own construction, swap the shape source --- + +def _upsample_builder(self, dtype, shapes, original): + # _cfgs entries are (N,C,Hi,Wi,Ho,Wo,ac,label); yaml carries the 7 numerics. + self._cfgs = [tuple(shape) + ("yaml",) for shape in shapes] + yield from original(self, dtype) + + +def _fused_deepseek_builder(self, dtype, shapes, original): + TestParam = original.__globals__["TestParam"] + for shape in shapes: + num_tokens, num_heads, num_tokens_insert, block_size, max_pos, eps = shape + param = TestParam( + num_tokens=num_tokens, + num_heads=num_heads, + num_tokens_insert=num_tokens_insert, + block_size=block_size, + max_pos=max_pos, + eps=eps, + ) + yield from type(self).make_input(param) + + +def _flash_mla_sparse_builder(self, dtype, shapes, original): + TestParam = original.__globals__["TestParam"] + for shape in shapes: + s_q, s_kv, topk, h_q, d_qk = shape + param = TestParam( + s_q=s_q, s_kv=s_kv, topk=topk, h_q=h_q, d_qk=d_qk, + have_attn_sink=True, # every hardcoded case uses True + ) + yield from type(self).make_input_flashmla(param) + + +_WRAP_BUILDERS = { + "UpsampleBicubic2dAaBackwardBenchmark": _upsample_builder, + "FusedDeepseekV4QnormRopeKVRopeQuantInsertBenchmark": _fused_deepseek_builder, + "FlashmlaSparseBenchmark": _flash_mla_sparse_builder, +} + + +def _wrap_get_input_iter(cls, builder) -> None: + original = cls.get_input_iter + + def patched(self, dtype): + shapes = _yaml_shapes(getattr(self, "op_name", None)) + if shapes: + yield from builder(self, dtype, shapes, original) + else: + yield from original(self, dtype) + + cls.get_input_iter = patched + + +def _patch_cutlass(cls) -> bool: + """Redirect CutlassScaledMMPerfKit's hardcoded mnk to the yaml M,N,K, reusing + the kit's own combination/sampling pipeline (the quant-mode sweep stays + intrinsic).""" + shapes = _yaml_shapes("cutlass_scaled_mm") + if not shapes: + return False + kit = getattr(sys.modules.get(cls.__module__), "CutlassScaledMMPerfKit", None) + if kit is None: + return False + import torch + from itertools import product + + mnk = [tuple(shape) for shape in shapes] + + def _get_all_combinations(): + # Mirror upstream verbatim except mnk (from yaml); keep in sync with + # CutlassScaledMMPerfKit._get_all_combinations. + scale_shape_types = ["scalar", "vector", "matrix"] + if_use_bias = [True, False] + dtypes = [(torch.int8, torch.float16), (torch.float8_e4m3fn, torch.bfloat16)] + return product(mnk, scale_shape_types, scale_shape_types, if_use_bias, dtypes) + + kit._get_all_combinations = staticmethod(_get_all_combinations) + return True + + +def pytest_collection_finish(session): + global _YAML_CACHE + from benchmark import base as fg_base + from benchmark.conftest import Config as fg_Config + + _YAML_CACHE = _load_yaml(getattr(fg_Config, "shape_file", "") or "") + + def _subclasses(cls): + for sub in cls.__subclasses__(): + yield sub + yield from _subclasses(sub) + + covered = [] + seen = set() + for cls in _subclasses(fg_base.Benchmark): + if cls in seen: + continue + seen.add(cls) + builder = _WRAP_BUILDERS.get(cls.__name__) + if builder is not None and "get_input_iter" in cls.__dict__: + _wrap_get_input_iter(cls, builder) + covered.append(cls.__name__) + 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'}", + file=sys.stderr, flush=True) diff --git a/_cudagraph_plugin.py b/_cudagraph_plugin.py new file mode 100644 index 0000000..ac75b77 --- /dev/null +++ b/_cudagraph_plugin.py @@ -0,0 +1,82 @@ +"""pytest plugin: swap `triton.testing.do_bench` for cudagraph-based timing. + +Loaded via `-p _cudagraph_plugin` in run_pytest.sh's PLUGINS list (comment that +line out to time with plain do_bench). Monkey-patches do_bench → +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). +""" + +from __future__ import annotations + +import os +import sys + +import triton.testing as _tt + + +_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.""" + msg = f"{type(exc).__name__}: {exc}".lower() + markers = ( + "illegal memory access", + "cudaerrorillegaladdress", + "out of memory", + "device-side assert", + "an illegal instruction", + "misaligned address", + "uncorrectable ecc", + ) + return any(m in msg for m in markers) + + +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", "") + try: + # do_bench_cudagraph runs fn on a fresh side stream without syncing + # with the default stream first; input tensors produced just before + # (quantize/topk in the input iter) may still be in flight, and racing + # on them intermittently kills the context (illegal instruction from a + # device trap on garbage indices). Sync before switching streams. + import torch + torch.cuda.synchronize() + return _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; only graph-capture rejections fall back. + if _is_fatal_cuda_error(exc): + raise + reason = type(exc).__name__ + detail = " ".join(str(exc).split())[:160] + print( + f"BENCHMARK_DIRECT_NO_CUDAGRAPH op={op} reason={reason} detail={detail!r}", + flush=True, + ) + return _ORIGINAL_DO_BENCH( + fn, + warmup=warmup, + rep=rep, + grad_to_none=grad_to_none, + quantiles=quantiles, + return_mode=return_mode, + ) + + +def pytest_configure(config): + _tt.do_bench = _patched_do_bench + print("[cudagraph-plugin] triton.testing.do_bench → do_bench_cudagraph", + file=sys.stderr, flush=True) diff --git a/_ir_meta_plugin.py b/_ir_meta_plugin.py new file mode 100644 index 0000000..9f09b7e --- /dev/null +++ b/_ir_meta_plugin.py @@ -0,0 +1,428 @@ +"""pytest plugin: annotate compiled kernels and dump their ttgir, grouped by +benchmark shape, at session end. + +Pieces, all keyed by the kernel's Triton cache dir: + +1. Compilation listener — writes a __constexprs.json sidecar (constexpr values + + per-arg specializations like tt.divisibility) next to each compiled + kernel. The ttgir alone doesn't say which autotune variant it is. +2. Shape tracking — wraps every Benchmark subclass's get_input_iter (after the + shape-inject plugins have done their patching) to publish the benchmark + shape currently being fed, e.g. "16x256x7168x2048x8". +3. Launch tracking — kernel_load_end_hook maps a loaded function handle to its + cache dir (must be the END hook: at load-start the handle is still None); + launch_enter_hook counts launches per (kernel cache dir, current shape). + Launches made inside Autotuner._bench (the tuning sweep; LibTuner inherits + it) are counted separately, so "actually used" = launched at least once + OUTSIDE the sweep. +4. Dump at exit (only when FLAGGEMS_PERF_TTGIR_DUMP_DIR is set) — walks + TRITON_CACHE_DIR and copies each ttgir into + ///_w{warps}s{stages}.ttgir + for every shape that actually used it. Only constexprs that VARY within a + (shape, kernel) group are named, abbreviated per-word (BLOCK_SIZE_M -> BSM; + legend in naming.md). Same-name collisions are disambiguated by signature + dtype (dtype sweeps, e.g. __fp16), then varying per-arg specializations + (e.g. __EMdiv16), then a cache-hash prefix. index.tsv lists every compiled + variant per shape with launch counts, including never-used sweep losers + (not copied). atexit (not sessionfinish) so a mid-run CUDA crash still + dumps whatever was compiled. +""" + +from __future__ import annotations + +import atexit +import json +import os +import re +import sys +from collections import defaultdict + +import pytest + +_DUMP_DIR_ENV = "FLAGGEMS_PERF_TTGIR_DUMP_DIR" + +_fn_to_cachedir: dict = {} # GPU function handle -> cache-dir basename +_launches: dict = {} # cache-dir basename -> {shape label: [real, sweep]} +_in_bench = 0 # >0 while inside Autotuner._bench (autotune sweep) +_current_shape: str | None = None +_dump_enabled = False + + +# --- compile-time sidecar ------------------------------------------------- + +def _plain(value): + value = getattr(value, "value", value) # unwrap tl.constexpr + if value is None or isinstance(value, (bool, int, float, str)): + return value + return repr(value) + + +def _compile_listener(*, src, metadata, metadata_group, times, cache_hit): + try: + fn = getattr(src, "fn", None) + arg_names = getattr(fn, "arg_names", None) + constants = getattr(src, "constants", None) + if not arg_names or constants is None: + return # IRSource or unexpected layout: nothing to record + paths = list(metadata_group.values()) + if not paths: + return + path = os.path.join(os.path.dirname(paths[0]), "__constexprs.json") + if os.path.exists(path): + return # cache hit on a dir we already annotated + + def arg_name(key): + # ASTSource keys constants/attrs by arg-index tuples; map back to names. + if isinstance(key, tuple): + return ".".join( + arg_names[i] if isinstance(i, int) and i < len(arg_names) else str(i) + for i in key + ) + return str(key) + + out = {arg_name(k): _plain(v) for k, v in constants.items()} + + # Argument signature (dtypes): benchmarks sweep dtypes, and dtype is + # not a constexpr — without this, dtype variants collide into hash + # suffixes. Values look like "*fp16", "i32". + sig = {} + try: + sig = {str(k): str(v) for k, v in (getattr(src, "signature", None) or {}).items() + if str(v) != "constexpr"} + except Exception: + pass + + # Per-arg specializations, e.g. ("tt.divisibility", 16) -> "div16". + specs = {} + for key, props in (getattr(src, "attrs", None) or {}).items(): + encoded = [] + for p in props or []: + try: + pname, pval = p[0], p[1] + except (TypeError, IndexError): + encoded.append(str(p)) + continue + encoded.append(f"div{pval}" if "divisibility" in str(pname) + else f"{pname}={pval}") + if encoded: + specs[arg_name(key)] = sorted(encoded) + + tmp = f"{path}.tmp.pid{os.getpid()}" + with open(tmp, "w") as f: + json.dump({"name": getattr(src, "name", "unknown"), + "constexprs": out, "attrs": specs, "signature": sig}, + 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) + + +# --- shape + launch tracking ---------------------------------------------- + +def _load_hook(module, function, name, metadata_group, hash): + try: + paths = list(metadata_group.values()) + if paths: + _fn_to_cachedir[function] = os.path.basename(os.path.dirname(paths[0])) + except Exception: + pass + + +def _launch_hook(md): + try: + cachedir = _fn_to_cachedir.get(md.data.get("function")) + if cachedir is None: + return + rec = _launches.setdefault(cachedir, {}).setdefault( + _current_shape or "shape_unknown", [0, 0]) + rec[1 if _in_bench else 0] += 1 + except Exception: + pass + + +def _wrap_bench(original): + def wrapped(self, *args, **kwargs): + global _in_bench + _in_bench += 1 + try: + return original(self, *args, **kwargs) + finally: + _in_bench -= 1 + return wrapped + + +def _shape_label(bench, idx): + shapes = getattr(bench, "shapes", None) or [] + if idx < len(shapes): + s = shapes[idx] + if isinstance(s, (list, tuple)): + return "x".join(re.sub(r"\W", "", str(v)) for v in s) + return re.sub(r"\W", "", str(s)) + return f"input{idx}" + + +def _wrap_input_iter(original): + # Publish the shape label BEFORE resuming the generator, so kernels + # launched while building the inputs attribute to the right shape too. + def patched(self, dtype): + global _current_shape + it = original(self, dtype) + idx = 0 + while True: + _current_shape = _shape_label(self, idx) + try: + item = next(it) + except StopIteration: + _current_shape = None + return + yield item + idx += 1 + return patched + + +@pytest.hookimpl(trylast=True) +def pytest_collection_finish(session): + # trylast: run after the shape-inject plugins have re-pointed + # get_input_iter, so we wrap the version that will actually execute. + if not _dump_enabled: + return + from benchmark import base as fg_base + + def _subclasses(cls): + for sub in cls.__subclasses__(): + yield sub + yield from _subclasses(sub) + + seen, wrapped = set(), 0 + for cls in (fg_base.Benchmark, *_subclasses(fg_base.Benchmark)): + if cls in seen: + continue + seen.add(cls) + own = cls.__dict__.get("get_input_iter") + 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", + file=sys.stderr, flush=True) + + +# --- naming ---------------------------------------------------------------- + +_SAN = re.compile(r"[^A-Za-z0-9.-]+") + + +def _fmt_value(v): + # isinstance check first: 1 == True in Python, a plain dict lookup would + # render GROUP_SIZE_M=1 as "T". + if isinstance(v, bool): + return "T" if v else "F" + if v is None: + return "-" + s = str(v) + if "." in s and not re.fullmatch(r"-?\d+(\.\d+)?", s): + s = s.split(".")[-1] # dotted repr like triton.language.bfloat16 + s = _SAN.sub("", s) + return s[:24] or "x" + + +def _abbrev(name): + words = [w for w in name.split("_") if w] + if len(name) <= 4 or len(words) < 2: + return name + out = [] + for w in words: + m = re.match(r"^([A-Za-z])[A-Za-z]*?(\d*)$", w) + out.append((m.group(1) + m.group(2)) if m else w[0]) + return "".join(out) + + +def _disambiguate(group): + """Suffixes for variants whose constexprs + launch config coincide: name by + the signature dtypes that differ (dtype sweeps), else by differing per-arg + specializations, else fall back to a hash prefix.""" + sig_keys = sorted({k for v in group for k in (v.get("signature") or {})}) + sig_varying = [ + k for k in sig_keys + if len({(v.get("signature") or {}).get(k) for v in group}) > 1 + ] + if sig_varying: + # One dtype sweep usually shifts every tensor arg together; the first + # varying arg's dtype identifies the variant. Sanitize "*fp16" -> fp16. + sufs = [] + for v in group: + toks = [re.sub(r"\W", "", (v.get("signature") or {}).get(k) or "none") + for k in sig_varying] + uniq = sorted(set(toks)) + sufs.append("__" + (uniq[0] if len(uniq) == 1 else "_".join( + f"{_abbrev(k)}{t}" for k, t in zip(sig_varying, toks)))) + if len(set(sufs)) == len(group): + return sufs + attr_keys = sorted({k for v in group for k in (v.get("attrs") or {})}) + varying = [ + k for k in attr_keys + if len({tuple((v.get("attrs") or {}).get(k, [])) for v in group}) > 1 + ] + if varying: + sufs = [] + for v in group: + toks = [ + f"{_abbrev(k)}{'.'.join((v.get('attrs') or {}).get(k) or ['none'])}" + for k in varying + ] + sufs.append("__" + "_".join(toks)) + if len(set(sufs)) == len(group): + return sufs + return ["__" + v["hash"][:8] for v in group] + + +def _name_group(items): + """Filenames for one (shape, kernel) group: only constexprs whose value + varies within the group, abbreviated. Returns ([(item, filename)], legend).""" + keys = sorted({k for it in items for k in it["constexprs"]}) + varying = [ + k for k in keys + if len({json.dumps(it["constexprs"].get(k), sort_keys=True) for it in items}) > 1 + ] + by_ab = defaultdict(list) + for k in varying: + by_ab[_abbrev(k)].append(k) + ab = {k: (a if len(ks) == 1 else k) for a, ks in by_ab.items() for k in ks} + + named = defaultdict(list) + for it in items: + parts = [f"{ab[k]}{_fmt_value(it['constexprs'][k])}" + for k in varying if k in it["constexprs"]] + w, s = it.get("warps"), it.get("stages") + parts.append(f"w{w}s{s}" if w not in (None, "") else "cfg-unknown") + named["_".join(parts)].append(it) + + results = [] + for base, group in sorted(named.items()): + sufs = _disambiguate(group) if len(group) > 1 else [""] + for it, suf in zip(group, sufs): + results.append((it, f"{base}{suf}.ttgir")) + return results, {k: a for k, a in ab.items() if a != k} + + +# --- dump ------------------------------------------------------------------- + +def _dump_ttgir(cache_dir: str, dump_dir: str) -> None: + import pathlib + import shutil + + 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) + return + dest.mkdir(parents=True, exist_ok=True) + + def load(path): + try: + return json.loads(path.read_text()) + except Exception: + return {} + + infos = [] + for ttgir in sorted(cache.rglob("*.ttgir")): + meta = load(ttgir.with_suffix(".json")) + side = load(ttgir.parent / "__constexprs.json") + infos.append({ + "src": ttgir, "hash": ttgir.parent.name, "kernel": ttgir.stem, + "warps": meta.get("num_warps"), "stages": meta.get("num_stages"), + "constexprs": side.get("constexprs", {}), "attrs": side.get("attrs", {}), + "signature": side.get("signature", {}), + "per_shape": _launches.get(ttgir.parent.name, {}), + }) + + # (shape, kernel) -> variants really used there (launched outside the sweep) + groups = defaultdict(list) + for info in infos: + for shape, (real, _sweep) in info["per_shape"].items(): + if real > 0: + groups[(shape, info["kernel"])].append(info) + + rows, legends, copied = [], defaultdict(dict), 0 + for (shape, kernel), items in sorted(groups.items()): + named, legend = _name_group(items) + legends[kernel].update(legend) + for it, fname in named: + target = dest / shape / kernel / fname + target.parent.mkdir(parents=True, exist_ok=True) + shutil.copyfile(it["src"], target) + copied += 1 + real, sweep = it["per_shape"][shape] + rows.append((shape, kernel, fname, real, sweep, + it["warps"] or "", it["stages"] or "", + json.dumps(it["constexprs"], sort_keys=True), it["hash"])) + + # Compiled but never really used anywhere (autotune losers): index-only. + for info in infos: + if not any(real > 0 for real, _ in info["per_shape"].values()): + sweep = sum(s for _, s in info["per_shape"].values()) + rows.append(("-", info["kernel"], "(sweep loser, not dumped)", 0, sweep, + info["warps"] or "", info["stages"] or "", + json.dumps(info["constexprs"], sort_keys=True), info["hash"])) + + rows.sort() + with open(dest / "index.tsv", "w") as f: + f.write("shape\tkernel\tfile\tlaunches\tsweep_launches\t" + "num_warps\tnum_stages\tconstexprs\tcache_hash\n") + for r in rows: + f.write("\t".join(str(x) for x in r) + "\n") + + lines = ["# ttgir naming legend", "", + "Layout: `//_w{warps}s{stages}[__spec|__hash8].ttgir`", + "Shape dirs mirror the benchmark shape yaml. Only constexprs that vary", + "within a (shape, kernel) group appear; `index.tsv` has the full map.", + "Values: `T`/`F` = true/false, `-` = none/null."] + for kernel in sorted(legends): + if not legends[kernel]: + continue + lines += ["", f"## {kernel}"] + width = max(len(a) for a in legends[kernel].values()) + for full, a in sorted(legends[kernel].items(), key=lambda kv: kv[1]): + lines.append(f"- `{a:<{width}}` = {full}") + (dest / "naming.md").write_text("\n".join(lines) + "\n") + + shapes = sorted({r[0] for r in rows if r[0] != "-"}) + losers = sum(1 for r in rows if r[0] == "-") + 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) + + +# --- registration ----------------------------------------------------------- + +def pytest_configure(config): + global _dump_enabled + import triton + import triton.runtime.autotuner as _autotuner + + prev = triton.knobs.compilation.listener + if prev is None: + triton.knobs.compilation.listener = _compile_listener + else: + def chained(**kwargs): + prev(**kwargs) + _compile_listener(**kwargs) + triton.knobs.compilation.listener = chained + + dump_dir = os.environ.get(_DUMP_DIR_ENV, "").strip() + if dump_dir: + cache_dir = os.environ.get("TRITON_CACHE_DIR", "").strip() + if cache_dir: + _dump_enabled = True + triton.knobs.runtime.kernel_load_end_hook.add(_load_hook) + triton.knobs.runtime.launch_enter_hook.add(_launch_hook) + _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 " + "is not; ttgir dump disabled", file=sys.stderr, flush=True) + + print("[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 new file mode 100644 index 0000000..de453c4 --- /dev/null +++ b/_mm_cluster_fix_plugin.py @@ -0,0 +1,137 @@ +"""pytest plugin: work around the Hopper fp16 cluster-remote GEMM OOB crash via +runtime monkeypatch (no upstream source change). + +Root cause: fp16 mm uses `_cluster_remote_gemm_kernel` (bf16 doesn't). When +num_pid_n = cdiv(N,BN) is odd (not a multiple of CLUSTER_SIZE=2) AND the GEMM is +perfectly tiled (USE_MASK off), the tail cluster's extra CTA reads out of range +-> cudaErrorIllegalAddress (e.g. N=128256 Llama-3 vocab). + +Fix: wrap `cluster_remote_mm_scenario` so only that unsafe combination falls +back to the general/splitk path; all other fp16 shapes keep the cluster kernel. +The real backend mm module loads lazily under a synthetic name, so we scan +sys.modules, force a warmup mm to trigger the load, patch, and re-scan per test. +""" +from __future__ import annotations + +import sys +import types + + +def _patch_loaded() -> list: + import triton + + patched = [] + for nm, m in list(sys.modules.items()): + if not isinstance(m, types.ModuleType): + continue + fn = getattr(m, "cluster_remote_mm_scenario", None) + if not isinstance(fn, types.FunctionType): + continue + if getattr(fn, "_mm_cluster_guarded", False): + continue + bn = getattr(m, "TLE_REMOTE_BN", 256) + bm = getattr(m, "TLE_REMOTE_BM", 64) + bk = getattr(m, "TLE_REMOTE_BK", 64) + cs = getattr(m, "TLE_CLUSTER_SIZE", 2) + + def make(orig, bm, bn, bk, cs): + def wrapped(a, b, c, M, N, K): + if not orig(a, b, c, M, N, K): + return False + use_mask = (M % bm != 0) or (N % bn != 0) or (K % bk != 0) + # odd N-tile count + no mask -> tail cluster over-runs N + if (triton.cdiv(N, bn) % cs != 0) and not use_mask: + return False + return True + wrapped._mm_cluster_guarded = True + return wrapped + + m.cluster_remote_mm_scenario = make(fn, bm, bn, bk, cs) + patched.append(nm) + return patched + + +def _guarded_modules() -> list: + """Loaded modules whose cluster_remote_mm_scenario is already our wrapper + (cumulative state, unlike _patch_loaded()'s newly-wrapped-only return).""" + out = [] + for nm, m in list(sys.modules.items()): + if not isinstance(m, types.ModuleType): + continue + fn = getattr(m, "cluster_remote_mm_scenario", None) + if isinstance(fn, types.FunctionType) and getattr(fn, "_mm_cluster_guarded", False): + out.append(nm) + return out + + +def _is_hopper(): + """True iff a CUDA device with capability >= 9 (the only HW that hits the + cluster-remote path) is visible. Returns False on any error / no CUDA.""" + try: + import torch + if not torch.cuda.is_available(): + return False + return torch.cuda.get_device_capability()[0] >= 9 + except Exception: # noqa: BLE001 + return False + + +def _warmup_load_backend(): + """Trigger the lazy backend mm module load via one tiny 64x64 fp16 mm (too + small for the cluster path, so it can't trip the bug). A CUDA error here + means a pre-broken context -> surface loudly; only import/attr errors are + swallowed.""" + try: + import torch + import flag_gems + if not torch.cuda.is_available(): + return + a = torch.randn(64, 64, dtype=torch.float16, device="cuda") + b = torch.randn(64, 64, dtype=torch.float16, device="cuda") + with flag_gems.use_gems(): + torch.mm(a, b) + torch.cuda.synchronize() + except (ImportError, AttributeError) as e: + print(f"[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}", + file=sys.stderr, flush=True) + + +def pytest_configure(config): + _patch_loaded() + _warmup_load_backend() + _patch_loaded() + guarded = _guarded_modules() + if guarded: + print( + "[mm-cluster-fix-plugin] guarded cluster_remote_mm_scenario " + f"(odd N-tile unmasked OOB) in: {guarded}", + file=sys.stderr, + flush=True, + ) + elif _is_hopper(): + # On Hopper with nothing patched, the guard is a silent no-op and fp16 mm + # will crash; fail loudly rather than risk an illegal-access. + raise RuntimeError( + "[mm-cluster-fix-plugin] FATAL: on a Hopper (cap>=9) device but " + "cluster_remote_mm_scenario was not found in any loaded module after " + "warmup. The fp16 mm OOB guard is NOT active; aborting rather than " + "risk an illegal-memory-access crash. (FlagGems mm module layout may " + "have changed.)" + ) + 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; " + "non-Hopper device, guard not needed (no-op).", + file=sys.stderr, + flush=True, + ) + + +def pytest_runtest_setup(item): + # backend module may load lazily on the first gems call; re-scan (idempotent). + _patch_loaded() diff --git a/_seed_plugin.py b/_seed_plugin.py new file mode 100644 index 0000000..d0e7144 --- /dev/null +++ b/_seed_plugin.py @@ -0,0 +1,39 @@ +"""pytest plugin: seed random/numpy/torch RNG for reproducible benchmark inputs. + +Data-dependent kernels (sort, topk, nonzero, ...) have value-dependent latency; +a fixed seed makes every run generate byte-identical inputs, so the latency +delta between two runs reflects the change under test, not the input data. +""" + +from __future__ import annotations + +import sys + +_SEED = 0 + + +def pytest_configure(config): + seed = _SEED + # random.seed too: some kits (cutlass_scaled_mm) pick cases via random.shuffle. + import random + random.seed(seed) + seeded = ["random"] + try: + import numpy as _np + _np.random.seed(seed) + seeded.append("numpy") + except Exception: + pass + 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}", + 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)} " + "— reproducible benchmark inputs/cases", + file=sys.stderr, flush=True) diff --git a/_shape_inject_plugin.py b/_shape_inject_plugin.py new file mode 100644 index 0000000..b145250 --- /dev/null +++ b/_shape_inject_plugin.py @@ -0,0 +1,99 @@ +"""pytest plugin: make our shape yaml win over subclasses that override +`set_shapes()` with hardcoded shapes (~9% of upstream Benchmark subclasses). + +Wraps `Benchmark.init_user_config` (runs right after `set_shapes()`): if the +op_name has a yaml entry, overwrite `self.shapes` with the yaml shapes. Ops with +no yaml entry keep their own shapes. An arity guard skips ops whose subclass +normalizes shape arity after reading yaml (e.g. BLAS (B,M,N,K)↔(M,N,K)). +""" + +from __future__ import annotations + +import os +import sys + +import yaml as _yaml + + +_YAML_CACHE: dict | None = None # filled lazily from Config.shape_file +_PATCHED = False + + +def _load_yaml(path: str) -> dict: + if not path or not os.path.isfile(path): + return {} + try: + 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) + return {} + + +def _normalize_shapes(raw): + """Mirror base.py's `[tuple(shape) for shape in self.shapes]`, defensively + handling scalar entries (`5` → `(5,)`).""" + out = [] + for shape in raw: + if isinstance(shape, (list, tuple)): + out.append(tuple(shape)) + else: + out.append((shape,)) + return out + + +def _patch_init_user_config(): + global _PATCHED + if _PATCHED: + return + from benchmark import base as fg_base + from benchmark.conftest import Config as fg_Config + + original = fg_base.Benchmark.init_user_config + + def _shape_arity(shape): + if isinstance(shape, (list, tuple)): + return len(shape) + return 1 + + def patched(self): + original(self) + global _YAML_CACHE + if _YAML_CACHE is None: + _YAML_CACHE = _load_yaml(getattr(fg_Config, "shape_file", "") or "") + entry = _YAML_CACHE.get(getattr(self, "op_name", None)) + if not isinstance(entry, dict): + return + shapes = entry.get("shapes") + if not shapes: + return + # Skip ops whose subclass normalizes shape arity after reading yaml (e.g. + # BLAS (B,M,N,K)↔(M,N,K)): detected as both sides being homogeneous in + # arity but differing. Heterogeneous subclass shapes are safe to override. + if getattr(self, "shapes", None) and len(self.shapes) > 0: + cur_arities = {_shape_arity(s) for s in self.shapes} + yaml_arities = {_shape_arity(s) for s in shapes} + cur_homogeneous = len(cur_arities) == 1 + yaml_homogeneous = len(yaml_arities) == 1 + if cur_homogeneous and yaml_homogeneous and cur_arities != yaml_arities: + print( + f"[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", + file=sys.stderr, + ) + return + self.shapes = _normalize_shapes(shapes) + if "shape_desc" in entry: + self.shape_desc = entry["shape_desc"] + + fg_base.Benchmark.init_user_config = patched + _PATCHED = True + + +def pytest_configure(config): + _patch_init_user_config() + print("[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 new file mode 100644 index 0000000..2e4f0c4 --- /dev/null +++ b/_shape_iter_inject_plugin.py @@ -0,0 +1,100 @@ +"""pytest plugin: make our yaml shapes win over subclasses that hardcode their +shape list inside an overridden `get_input_iter` (instead of `self.shapes`), +which `_shape_inject_plugin` can't reach (conv padding / pool / nll / +scaled_softmax families, ~17 ops). + +At `pytest_collection_finish` (all Benchmark subclasses defined), redirect +`get_input_iter` ONLY on classes whose own implementation does not reference +`self.shapes` (the source-level guard that protects well-behaved BLAS/generic +classes). The redirect yields one input per yaml shape via the op's standard +`input_fn(shape, dtype, device)`; ops with no yaml entry or a non-standard +input_fn fall through unchanged. This also bounds compile cost vs the large +hardcoded lists. +""" + +from __future__ import annotations + +import inspect +import os +import sys + +import yaml as _yaml + + +_YAML_CACHE: dict | None = None + + +def _load_yaml(path: str) -> dict: + if not path or not os.path.isfile(path): + return {} + try: + 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) + return {} + + +def _reads_self_shapes(func) -> bool: + """True if the method references self.shapes (already covered by + _shape_inject_plugin, so don't redirect). Defaults True on introspection + failure (safe: leave the method alone).""" + try: + return "self.shapes" in inspect.getsource(func) + except (OSError, TypeError): + return True + + +def _accepts_shape_dtype_device(fn) -> bool: + """True if `fn` can be called as `fn(shape, dtype, device)` (the standard + input_fn signature). Guards against blas-style `fn(b,m,n,k,dtype,device,...)`. + Defaults False on introspection failure (safe: fall back to original).""" + if not callable(fn): + return False + try: + inspect.signature(fn).bind(None, None, None) + except (TypeError, ValueError): + return False + return True + + +def _make_patched(original): + def patched(self, dtype): + entry = (_YAML_CACHE or {}).get(getattr(self, "op_name", None)) + shapes = entry.get("shapes") if isinstance(entry, dict) else None + input_fn = getattr(self, "input_fn", None) + if shapes and _accepts_shape_dtype_device(input_fn): + for shape in shapes: + yield from input_fn(tuple(shape), dtype, self.device) + else: + yield from original(self, dtype) + return patched + + +def pytest_collection_finish(session): + global _YAML_CACHE + from benchmark import base as fg_base + from benchmark.conftest import Config as fg_Config + + _YAML_CACHE = _load_yaml(getattr(fg_Config, "shape_file", "") or "") + + def _subclasses(cls): + for sub in cls.__subclasses__(): + yield sub + yield from _subclasses(sub) + + patched = 0 + seen = set() + for cls in _subclasses(fg_base.Benchmark): + if cls in seen: + continue + seen.add(cls) + own = cls.__dict__.get("get_input_iter") + # Redirect only classes with their own get_input_iter that ignore self.shapes. + if own is not None and not _reads_self_shapes(own): + cls.get_input_iter = _make_patched(own) + patched += 1 + + print(f"[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/run_pytest.sh b/run_pytest.sh new file mode 100644 index 0000000..daf9b82 --- /dev/null +++ b/run_pytest.sh @@ -0,0 +1,108 @@ +#!/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/. +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:-}" + +# Inline shape yaml (used when SHAPE_FILE is empty); top-level key must be the op name. +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 + shapes: + # Mixtral (E=8) + - [1, 8, 4096, 14336, 2] + - [16, 8, 4096, 14336, 2] + - [64, 8, 4096, 14336, 2] + - [256, 8, 4096, 14336, 2] + # DeepSeek-V3 (E=256, H=7168) + - [1, 256, 7168, 2048, 8] + - [16, 256, 7168, 2048, 8] + - [64, 256, 7168, 2048, 8] + - [256, 256, 7168, 2048, 8] + # Qwen3 (E=512) + - [1, 512, 4096, 1024, 10] + - [16, 512, 4096, 1024, 10] + - [64, 512, 4096, 1024, 10] + - [256, 512, 4096, 1024, 10] + # DeepSeek-V4-Flash (E=256, H=4096) + - [1, 256, 4096, 2048, 6] + - [16, 256, 4096, 2048, 6] + - [64, 256, 4096, 2048, 6] + - [256, 256, 4096, 2048, 6] + 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}" + +OUT_DIR="$SCRIPT_DIR/runs/${OP}_$(date +%Y%m%d-%H%M%S)" +LOG_FILE="$OUT_DIR/run.log" +mkdir -p "$OUT_DIR" + +export PYTHONPATH="$SCRIPT_DIR${PYTHONPATH:+:$PYTHONPATH}" +PLUGINS=( + -p _seed_plugin + -p _shape_inject_plugin + -p _shape_iter_inject_plugin + -p _bespoke_shape_plugin + # -p _mm_cluster_fix_plugin + -p _autotune_record_plugin + -p _cudagraph_plugin + -p _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:-}" +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 +else + AUTOTUNE_ENV="FLAGGEMS_PERF_AUTOTUNE_RECORD_DIR=$OUT_DIR/autotune_records" + mkdir -p "$OUT_DIR/autotune_records" +fi +export FLAGGEMS_PERF_CURRENT_OP="$OP" + +# Resolve the shape file: use SHAPE_FILE, or write the inline yaml to a temp file. +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 + +{ + # _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). + 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[@]}" \ + --shape_file "$SHAPE_FILE" \ + --level core --mode kernel || status=$? + + rm -rf "$CACHE_DIR" + + if (( status == 0 )); then + echo ">>> done. outputs in $OUT_DIR" + else + echo ">>> FAILED (pytest exit $status). partial outputs in $OUT_DIR" + fi + exit "$status" +} 2>&1 | tee "$LOG_FILE"