fix: never let the parallel sweep cover only part of the shape set

The sharded autotune sweep could publish a config set with holes, and a run
using it reports itself as replaying while every uncovered key silently
self-tunes. That is the exact failure the record/replay caliber exists to
prevent, and it was reachable two ways.

Shape lookup was op_name-only, but BLAS operators keep their shapes under the
class name: core_shapes.yaml has BlasBenchmark, not mm. Every BLAS operator
therefore resolved to zero shapes and the plugin disabled itself with a message
that read like a normal skip. It now walks the benchmark class chain the way
FlagGems' set_shapes does. The chain is recovered statically -- the class comes
from the op_name= keyword of the constructor call in the test module, its bases
from that module plus benchmark/base.py -- because the plugin runs at
pytest_configure, before any Benchmark instance exists. Resolved keys are named
in the log, and a total miss lists the candidates it tried.

Two cases can still produce partial coverage, and both now refuse rather than
shard a subset:

- --level comprehensive, where set_shapes() merges set_more_shapes() on top of
  the yaml. That addition is arbitrary per-subclass Python and invisible here.
  Measured on mm: 5 of 29 shapes swept, 90 AUTOTUNE_REPLAY_FALLBACK
  reason=key_missing in the run that consumed the result.
- a shard exiting non-zero, whose shapes were never swept. This used to append a
  WARNING to a success message and publish the merge anyway; observed with 3 of
  5 shards dead. The whole sweep is now discarded, worker logs kept.

Both fall back to the normal serial sweep, which costs wall time rather than the
same-config premise.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

README.md

@@ -179,7 +179,17 @@ 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 为 `CUDA_VISIBLE_DEVICES`）——用它选过卡时，分片只会落在你选的那几张上。
+**冲突数就是这次加速的可信度指标**：0 可放心用；偏多说明 config 选择已被污染，出正式结论前不设该变量重跑一遍。
+
+插件会在三种情况下**主动放弃分片、退回串行 sweep**，日志各有独立措辞，都不是静默的：
+
+- **`BENCH_LEVEL` 不是 `core`**。`comprehensive` 会让上游在 `set_shapes()` 里并入 `set_more_shapes()`，那部分是各子类自己的 Python（硬编码元组加 `models_shapes/*.yaml`），插件只看得到 yaml 里的 shape 列表。按子集分片会让未覆盖的 key 全部现场自调优，而这一轮却自称在 replay——`mm` 上实测是 29 个 shape 只扫到 5 个、90 次 `AUTOTUNE_REPLAY_FALLBACK reason=key_missing`。
+- **任一分片非零退出**。死掉的分片的 shape 从未被扫过，合并出来的 config 集是带洞的，publish 之后每个洞都会静默自调优。因此整轮 sweep 直接作废，worker 日志保留在临时目录里备查。
+- yaml 里查不到该算子的 shape（见下）、`REPLAY_FROM` 已设、或可见卡不足 2 张。
+
+yaml key 的查找与上游对齐：先按 `op_name`，查不到再沿 Benchmark 类链逐级尝试类名（`mm` 因此命中 `BlasBenchmark`）。类链是静态解析测试模块里 `op_name=` 那次构造调用得到的，不 import FlagGems。日志会打出命中的 key；一个都没命中时会列出尝试过的候选。
+
+其余行为（轮转分片而非按块切、用子进程而非 xdist 的原因）见插件 docstring。分片绑卡遵守当前后端的可见性变量（CUDA 为 `CUDA_VISIBLE_DEVICES`）——用它选过卡时，分片只会落在你选的那几张上。
 
 ## 批量多算子测试（run_batch.py）
 

_parallel_warmup_plugin.py

@@ -33,7 +33,7 @@ import sys
 import tempfile
 import time
 from pathlib import Path
-from typing import Any, Dict, List, Tuple
+from typing import Any, Dict, List, Optional, Tuple
 
 import pytest
 
@@ -88,15 +88,114 @@ def _visible_devices() -> List[str]:
         return []
 
 
-def _read_shape_yaml(path: str, op: str) -> Tuple[List[Any], Dict[str, Any]]:
-    """Shapes for this op, plus the sibling keys to carry into each shard."""
+def _attr_name(node: Any) -> str:
+    """Last component of a Name/Attribute node: `base.BlasBenchmark` -> `BlasBenchmark`."""
+    import ast
+    if isinstance(node, ast.Name):
+        return node.id
+    if isinstance(node, ast.Attribute):
+        return node.attr
+    return ""
+
+
+def _class_bases(py_path: Path) -> Dict[str, List[str]]:
+    """{class name: [base class names]} parsed statically from one module."""
+    import ast
+    try:
+        tree = ast.parse(py_path.read_text(errors="replace"))
+    except Exception:
+        return {}
+    return {
+        node.name: [n for n in (_attr_name(b) for b in node.bases) if n]
+        for node in ast.walk(tree)
+        if isinstance(node, ast.ClassDef)
+    }
+
+
+def _class_chain(test_py: Path, op: str) -> List[str]:
+    """Benchmark class names to try for `op`, nearest ancestor first.
+
+    FlagGems' `Benchmark.set_shapes` looks up `op_name` and, failing that, walks
+    `type(self).__mro__` trying each class name as a yaml key -- which is how
+    BLAS ops reach the shared `BlasBenchmark` entry in core_shapes.yaml. This
+    plugin runs at pytest_configure, before any Benchmark exists, so the real
+    MRO is unavailable; the class is recovered statically instead, from the
+    `op_name=` keyword of the constructor call inside the test module, then its
+    bases from that module plus benchmark/base.py.
+
+    Breadth-first over the inheritance graph, not a true C3 linearisation: the
+    benchmark classes are single-inheritance in practice, and a miss only costs
+    the sharding (see the caller's skip path), never correctness.
+    """
+    import ast
+    try:
+        tree = ast.parse(test_py.read_text(errors="replace"))
+    except Exception:
+        return []
+    start = ""
+    for node in ast.walk(tree):
+        if not isinstance(node, ast.Call):
+            continue
+        for kw in node.keywords:
+            if (kw.arg == "op_name" and isinstance(kw.value, ast.Constant)
+                    and kw.value.value == op):
+                start = _attr_name(node.func)
+                break
+        if start:
+            break
+    if not start:
+        return []
+
+    graph = _class_bases(test_py)
+    for name, bases in _class_bases(test_py.parent / "base.py").items():
+        graph.setdefault(name, bases)
+
+    chain: List[str] = []
+    seen: set = set()
+    queue = [start]
+    while queue:
+        name = queue.pop(0)
+        if not name or name in seen:
+            continue
+        seen.add(name)
+        chain.append(name)
+        queue.extend(graph.get(name, []))
+    return chain
+
+
+def _test_module_path(argv: List[str]) -> Optional[Path]:
+    """The benchmark module under test, from pytest's own argv."""
+    for a in argv:
+        if a.startswith("-"):
+            continue
+        p = Path(a.split("::", 1)[0])
+        if p.suffix == ".py" and p.is_file():
+            return p
+    return None
+
+
+def _read_shape_yaml(path: str, op: str, test_py: Optional[Path] = None
+                     ) -> Tuple[List[Any], Dict[str, Any], str]:
+    """Shapes for this op, the sibling keys to carry into each shard, and the
+    yaml key they came from ("" when nothing matched).
+
+    Tries `op` first, then the benchmark class chain, mirroring FlagGems. Without
+    the class fallback every BLAS op (mm/addmm/bmm/baddbmm, whose shapes live
+    under `BlasBenchmark`) resolved to zero shapes and sharding silently turned
+    itself off.
+    """
     import yaml
     with open(path) as f:
         doc = yaml.safe_load(f) or {}
-    entry = doc.get(op) or {}
-    shapes = entry.get("shapes") or []
-    extra = {k: v for k, v in entry.items() if k != "shapes"}
-    return list(shapes), extra
+    keys = [op]
+    if not isinstance(doc.get(op), dict) and test_py is not None:
+        keys += _class_chain(test_py, op)
+    for key in keys:
+        entry = doc.get(key)
+        if isinstance(entry, dict) and entry.get("shapes"):
+            extra = {k: v for k, v in entry.items() if k != "shapes"}
+            return list(entry["shapes"]), extra, key
+    return [], {}, ""
 
 
 def _shard(shapes: List[Any], n: int) -> List[List[Any]]:
@@ -226,6 +325,23 @@ def pytest_configure(config):
         return
 
     op = _op_name()
+
+    # Only --level core is shardable. At comprehensive, FlagGems' set_shapes()
+    # merges set_more_shapes() on top of whatever the yaml holds, and that
+    # addition is arbitrary per-subclass Python (hardcoded tuples plus
+    # models_shapes/*.yaml) which this plugin cannot see. Sharding the yaml list
+    # alone would sweep a strict subset -- measured on mm: 5 of 29 shapes, 90
+    # AUTOTUNE_REPLAY_FALLBACK reason=key_missing -- and every uncovered key
+    # then self-tunes while the run reports itself as replaying. Refusing costs
+    # wall time; sharding a subset costs the same-config premise.
+    level = str(getattr(config.option, "level", "") or "").strip().lower()
+    if level and level != "core":
+        print(f"{tag('[parallel-warmup-plugin]')} --level {level} merges shapes "
+              "this plugin cannot enumerate (set_more_shapes); sharding a subset "
+              "would break the same-config premise, so sweeping serially instead",
+              flush=True)
+        return
+
     shape_file = getattr(config.option, "shape_file", "") or ""
     if not shape_file or not Path(shape_file).is_file():
         print(f"{tag('[parallel-warmup-plugin]')} no --shape_file; "
@@ -236,15 +352,24 @@ def pytest_configure(config):
               "already pinned, nothing to sweep, skipping", flush=True)
         return
 
+    test_py = _test_module_path(list(config.invocation_params.args))
     try:
-        shapes, extra = _read_shape_yaml(shape_file, op)
+        shapes, extra, key = _read_shape_yaml(shape_file, op, test_py)
     except Exception as exc:
         print(f"{tag('[parallel-warmup-plugin]')} cannot read shapes "
               f"({type(exc).__name__}: {exc}); skipping", flush=True)
         return
+    if not key:
+        # Distinct from "found a short list": nothing in the yaml describes this
+        # op at all, so the shard files would be built from an empty set.
+        print(f"{tag('[parallel-warmup-plugin]')} no yaml entry for {op!r} "
+              f"(tried: {', '.join([op] + _class_chain(test_py, op)) if test_py else op}); "
+              "skipping", flush=True)
+        return
     if len(shapes) < 2:
-        print(f"{tag('[parallel-warmup-plugin]')} only {len(shapes)} shape(s); "
-              "sharding would not pay off, skipping", flush=True)
+        print(f"{tag('[parallel-warmup-plugin]')} only {len(shapes)} shape(s) "
+              f"under yaml key {key!r}; sharding would not pay off, skipping",
+              flush=True)
         return
 
     devices = _visible_devices()
@@ -255,8 +380,9 @@ def pytest_configure(config):
 
     shards = _shard(shapes, min(gpus, len(shapes), len(devices)))
     scratch = Path(tempfile.mkdtemp(prefix=f"warmup_{op}_"))
-    print(f"{tag('[parallel-warmup-plugin]')} sweeping {len(shapes)} shapes on "
-          f"{len(shards)} GPUs (shards: {[len(s) for s in shards]}); this phase's "
+    print(f"{tag('[parallel-warmup-plugin]')} sweeping {len(shapes)} shapes "
+          f"(yaml key {key!r}) on {len(shards)} GPUs "
+          f"(shards: {[len(s) for s in shards]}); this phase's "
           "latency is discarded, only configs are kept", flush=True)
 
     keep_scratch = False
@@ -264,9 +390,23 @@ def pytest_configure(config):
     try:
         recs, bad = _spawn(shards, op, extra, scratch,
                            list(config.invocation_params.args), devices)
-        merged, total, conflicts = _merge(recs, op)
         elapsed = time.time() - t0
 
+        if bad:
+            # A dead shard means its shapes were never swept, so publishing the
+            # merge would hand the measurement a config set with holes; every
+            # hole then silently self-tunes under a run that claims to replay.
+            # Discard the whole sweep rather than publish a partial one.
+            keep_scratch = True
+            print(f"{tag('[parallel-warmup-plugin]')} shard(s) {bad} exited "
+                  f"non-zero after {elapsed:.0f}s; their shapes were never swept, "
+                  "so the merged configs would have holes. Discarding this sweep "
+                  f"and falling back to normal serial autotune (worker logs kept "
+                  f"in {scratch})", flush=True)
+            return
+
+        merged, total, conflicts = _merge(recs, op)
+
         if not total:
             keep_scratch = True
             print(f"{tag('[parallel-warmup-plugin]')} no configs recovered in "
@@ -291,9 +431,6 @@ def pytest_configure(config):
         msg = (f"{tag('[parallel-warmup-plugin]')} sweep done in {elapsed:.0f}s: "
                f"{len(merged)} kernels / {total} configs -> {out / f'{op}.json'}; "
                "measurement continues serially on one GPU")
-        if bad:
-            msg += f" (WARNING shard(s) {bad} exited non-zero; missing keys will "
-            msg += "fall back to live autotune)"
         print(msg, flush=True)
         if conflicts:
             print(f"{tag('[parallel-warmup-plugin]')} WARNING {conflicts} config "