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