feat: absorb the MThreads compat plugin and the A/B range report

Three pieces from the fork's feature/musa-benchmark-support branch that never
reached main, rebased onto the backend abstraction rather than copied verbatim.

_musa_benchmark_compat_plugin drops FlagGems benchmark skip guards written as
`torch.cuda.is_available()` where the code under test dispatches through
flag_gems.device. On MThreads those turn a working operator into "not measured":
test_fft.py skipped entirely, and unskipped on an S5000 it produces 25 rows. The
plugin pins each override to one file and one exact skip reason and only acts
when flag_gems has dispatched to the mthreads vendor; two guards were probed on
device and deliberately left in place, which is why they are named in its
docstring rather than silently absent.

It is wired through a new DeviceBackend.compat_plugins rather than a name test
in the planner, so the policy stays where every other backend difference lives.
Compat plugins are planned first because they rewrite collection-time markers
and nothing may read the item list before them; a test pins that position and
another pins that they do not leak to other backends.

batch_summary gains --ab: the speedup column compares gems against torch, not
one build against another, so an A/B needs the two run.logs paired by operator,
dtype and size detail. Only rows that succeeded on both sides with a positive
latency count, so an incomplete pairing cannot look like a complete range, and
the range is annotated when either side carries a BENCHMARK_DIRECT_NO_CUDAGRAPH
marker, since a fallback row was timed under a different regime.

ab_rlc_test.sh is the same caliber as ab_fold_test.sh with the comparison axis
moved to FLAGTREE_<VENDOR>_RLC_ENHANCE. Both are caliber examples, not
regression scripts: if the switch under test no longer exists the script still
runs every round and prints a full table showing no difference, so confirm the
axis is live before reusing either.

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

README.md

@@ -161,7 +161,7 @@ benchmark 目录，注册 `--level`/`--mode` 的 `conftest.py` 就没被加载
 只校验节点的文件那一半；`::test_<fn>` 交给 pytest，函数不存在时它自己会清楚地报
 `not found: ...`。
 
-### ab_fold_test.sh：A/B 测试参考示例
+### ab_fold_test.sh / ab_rlc_test.sh：A/B 测试参考示例
 
 `ab_fold_test.sh` 是上述 A/B 流程的现成参考：以 FlagGems `FLAGGEMS_MXFP4_FOLDSCALE`（fold_scale 优化）开关为对比对象，对 4 个 DeepSeek-V4-Flash 真实 trace shape 集（`shape/DeepSeek-V4-Flash-p{1024,4096,32768,65536}d1024.yaml`，随仓库提供）各跑一对 FOLD=0（基线，record）→ FOLD=1（replay 同 config），共 8 轮：
 
@@ -170,6 +170,8 @@ bash ab_fold_test.sh
 # 汇总 log：runs/ab_fold_<时间戳>.log（已去色）；快速对账：grep -E '^#####|FAILED' <log>
 ```
 
+`ab_rlc_test.sh` 是同一口径的另一个实例，对比轴换成 FlagTree 的 `FLAGTREE_<VENDOR>_RLC_ENHANCE`（`VENDOR` 默认 `MUSA`），算子表由 `OPS` 指定，跑完调 `batch_summary.py --ab <before> <after>` 出该算子的 speedup 区间。
+
 可复用的是**口径**——同 config replay、cudagraph 计时、真实 trace shape 集、逐对成组——对比其它开关 / 编译器改动 / 算子，改循环变量与环境变量即可，口径部分不用动。
 
 但**它是口径示例，不是可长期直接执行的回归脚本**：对比轴是 FlagGems 侧的算子开关，会随上游演进被改名、改语义或移除，本仓库不跟随同步更新（该脚本引用的开关目前就已被上游移除）。失效后果是**静默的**：脚本照样跑完全部轮次、照样输出完整对比表，只是两侧执行同一份代码，得出"无差异"的假结论。所以套用前先确认对比轴仍有效——`grep -rn '<开关名>' $FLAGGEMS_DIR/src` 要有命中，且两侧 latency 表确有差异。
@@ -268,6 +270,7 @@ python run_batch.py --ops-file my_ops.txt --replay-root runs/batch_xxx
 | `_autotune_record_plugin` | record/replay Triton autotune 选择；`autotune_records/meta.json` 记录 backend，跨后端 replay 拒绝 | `ZL_ENABLE_AUTOTUNE=0` |
 | `_timing_plugin` | 按后端选择计时路径：CUDA 可走 Graph（内部调用 `_cudagraph_plugin`），其余后端一开始就是上游 `do_bench`（`timer=direct`） | `ZL_ENABLE_CUDAGRAPH=0` 强制 direct |
 | `_cudagraph_plugin` | 仅由 `_timing_plugin` 在 CUDA 上调用。`do_bench` → `do_bench_cudagraph`，capture 前按 warmup 预算显式预热；无法 capture 时回退并打 `BENCHMARK_DIRECT_NO_CUDAGRAPH` | 不要直接 `-p` 加到非 CUDA 运行 |
+| `_musa_benchmark_compat_plugin` | 剥掉 FlagGems 上游若干 benchmark 上过时的 `torch.cuda.is_available()` 跳过门——被测代码本就经 `flag_gems.device` 分发，在 MThreads 上那些门把可用算子变成「未测量」。每条覆盖都钉死到一个文件与一条确切的 skip 理由，且仅在 `flag_gems.vendor_name == "mthreads"` 时生效 | 由 `backends/musa.py` 的 `compat_plugins` 声明，其他后端不加载 |
 | `_pretty_report_plugin` | 结果表整理 + 着色：全表相同的输入折叠成表头下一行图例，每行 Size Detail 只留随行变化的部分（shape 保持 `torch.Size([...])` 原样，MoE 类算子单行从 ~400 字符缩到一屏内）；SUCCESS 绿 / FAILED 红。列名与 `SUCCESS` 字样保持上游原文，`run.log` 的 grep/解析不受影响 | 需要与上游原始表格逐字对齐时去掉 |
 | `_ir_meta_plugin` | 编译期记录每个变体的 constexpr/签名/特化；launch 钩子统计每个 (kernel, shape) 的实际使用；退出时按上述结构 dump ttgir。Triton knobs 缺失时 fail-soft | `ZL_ENABLE_IR=0` |
 | `_mm_cluster_fix_plugin` | Hopper fp16 mm cluster kernel 越界崩溃的运行时规避 | 默认关；`ZL_ENABLE_MM_CLUSTER=1` 或 `ZL_EXTRA_PLUGINS` |

_musa_benchmark_compat_plugin.py

@@ -0,0 +1,89 @@
+"""Narrow compatibility fixes for the frozen MUSA product corpus.
+
+Several upstream FlagGems tests gate themselves on ``torch.cuda.is_available()``
+even though the code under test runs on any accelerator FlagGems dispatches to.
+On MThreads those guards skip the test instead of exercising the operator, which
+turns a working path into a verdict of "not measured".
+
+Every override here is pinned to one file and one exact skip reason, applies
+only when FlagGems has dispatched to the MThreads vendor, and is only added
+after the unskipped test was observed to pass on the device.
+
+Two files were probed on S5000 and deliberately left out, because their guards
+are not stale:
+
+- ``test_cutlass_scaled_mm.py`` gates on vLLM as well as CUDA, so removing the
+  guard would raise on the missing dependency instead of measuring anything.
+- ``test_grid_sample.py`` builds its inputs through a helper whose device
+  argument defaults to the literal ``"cuda"``. Its guard is consistent with its
+  own body: unskipped, all 90 collected cases fail in ``aten::empty`` on a CUDA
+  backend that does not exist here. Making that file portable means editing
+  upstream test source, which is outside what a compatibility plugin may do.
+"""
+
+from __future__ import annotations
+
+import sys
+from pathlib import Path
+
+
+# file name -> exact skipif reasons that are stale on a dispatching backend.
+_STALE_CUDA_GUARDS: dict[str, frozenset[str]] = {
+    "test_fft.py": frozenset({"CUDA is not available"}),
+    "test_top_k_per_row_decode.py": frozenset({"CUDA device required"}),
+}
+
+# Guards that are only stale for one test in the file stay pinned to it.
+_PINNED_TESTS: dict[str, str] = {
+    "test_fft.py": "test_perf_fft",
+}
+
+
+def _is_stale_cuda_guard(item, mark) -> bool:
+    if mark.name != "skipif":
+        return False
+    name = Path(str(item.path)).name
+    reasons = _STALE_CUDA_GUARDS.get(name)
+    if reasons is None or mark.kwargs.get("reason") not in reasons:
+        return False
+    pinned = _PINNED_TESTS.get(name)
+    return pinned is None or item.name == pinned
+
+
+# Retained under its historical name: the sealed campaign's tests reference it.
+def _is_stale_fft_cuda_guard(item, mark) -> bool:
+    return Path(str(item.path)).name == "test_fft.py" and _is_stale_cuda_guard(
+        item, mark
+    )
+
+
+def pytest_collection_modifyitems(items):
+    import flag_gems
+
+    if flag_gems.vendor_name != "mthreads":
+        return
+    removed: dict[str, int] = {}
+    # A guard can sit on the test, its class, or the module's pytestmark, and
+    # skipif is resolved by walking that whole chain. Strip every level.
+    seen: set[int] = set()
+    for item in items:
+        for node in item.listchain():
+            if id(node) in seen:
+                continue
+            seen.add(id(node))
+            markers = getattr(node, "own_markers", None)
+            if not markers:
+                continue
+            kept = [mark for mark in markers if not _is_stale_cuda_guard(item, mark)]
+            dropped = len(markers) - len(kept)
+            if dropped:
+                markers[:] = kept
+                name = Path(str(item.path)).name
+                removed[name] = removed.get(name, 0) + dropped
+    for name, count in sorted(removed.items()):
+        print(
+            f"[musa-benchmark-compat] removed {count} stale CUDA-only skip "
+            f"guard(s) from {name}",
+            file=sys.stderr,
+            flush=True,
+        )

ab_rlc_test.sh

@@ -0,0 +1,68 @@
+#!/bin/bash
+# A/B 对比：FlagTree RemoveLayoutConversions 增强关 vs 开，逐算子两轮。
+#
+# 口径：每个算子先跑关（record autotune config），再用 REPLAY_FROM 回放同一套
+# config 跑开，两侧只差 FLAGTREE_<VENDOR>_RLC_ENHANCE 这一个开关。产物沿用
+# run_pytest.sh 的目录结构，所以每个算子恰好留下两份 run.log：
+#
+#   runs/<op>_<时间戳>/run.log   关（before）
+#   runs/<op>_<时间戳>/run.log   开（after）
+#
+# 跑完用 batch_summary.py --ab <before> <after> 得到该算子的 speedup 区间。
+#
+# 用法：
+#   OPS="softmax gelu_backward" bash ab_rlc_test.sh
+#   OPS="sort:perf_sort topk" bash ab_rlc_test.sh     # 冒号分隔 OP:OP_FILE
+#
+# 注意：RLC 开关名带厂商前缀，用 VENDOR 指定（默认 MUSA）。套用前先确认开关有效：
+# grep 该开关名于 FlagTree 源码应有命中，否则两侧跑的是同一份代码。
+set -uo pipefail
+
+cd "$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+
+VENDOR="${VENDOR:-MUSA}"
+ENHANCE_VAR="FLAGTREE_${VENDOR}_RLC_ENHANCE"
+
+# 空格分隔的算子表。元素为 OP 或 OP:OP_FILE（省略时 OP_FILE=OP）。
+OPS="${OPS:-softmax}"
+
+if [[ -t 1 && -z "${NO_COLOR:-}" ]]; then
+  export FLAGGEMS_PERF_COLOR=always
+fi
+
+LOG="runs/ab_rlc_$(date +%Y%m%d-%H%M%S).log"
+mkdir -p runs
+
+{
+  echo ">>> vendor=$VENDOR  switch=$ENHANCE_VAR  ops=$OPS"
+  for entry in $OPS; do
+    op="${entry%%:*}"
+    op_file="${entry#*:}"
+    [[ "$op_file" == "$entry" ]] && op_file="$op"
+
+    echo "########## $op | $ENHANCE_VAR=0 before (record) ##########"
+    if ! OP="$op" OP_FILE="$op_file" env "$ENHANCE_VAR=0" bash run_pytest.sh; then
+      echo "!!! FAILED before: $op"
+      continue
+    fi
+    BEFORE=$(ls -dt "runs/${op}_"* | head -1)
+
+    echo "########## $op | $ENHANCE_VAR=1 after (replay $BEFORE) ##########"
+    if ! OP="$op" OP_FILE="$op_file" REPLAY_FROM="$BEFORE" env "$ENHANCE_VAR=1" bash run_pytest.sh; then
+      echo "!!! FAILED after: $op"
+      continue
+    fi
+    AFTER=$(ls -dt "runs/${op}_"* | head -1)
+
+    echo "########## $op | summary ##########"
+    python batch_summary.py --ab "$BEFORE/run.log" "$AFTER/run.log" --label "$op" \
+      || echo "!!! summary failed: $op"
+  done
+  echo "########## ALL DONE ##########"
+} 2>&1 | tee "$LOG"
+
+# 去掉 ANSI 转义以便 grep/diff（终端输出保留颜色）
+sed -i -E $'s/\x1b\\[[0-9;]*[A-Za-z]//g' "$LOG"
+
+echo ">>> 总 log: $(pwd)/$LOG"
+echo ">>> 每算子两份 run.log 在 runs/<op>_<时间戳>/ 下，按时间先后即 before / after"

backends/base.py

@@ -90,6 +90,13 @@ class DeviceBackend:
 
     def denies_capture(self, op: str) -> bool:
         return bool(op) and op in self.capture_deny_ops
+
+    # Plugins that patch around upstream assumptions this backend violates, e.g.
+    # FlagGems benchmarks that gate themselves on torch.cuda.is_available() even
+    # though the operator under test runs on whatever device FlagGems dispatched
+    # to. Loaded first so collection-time overrides land before anything reads
+    # the item list. Each such plugin must scope itself to its own vendor.
+    compat_plugins: Tuple[str, ...] = ()
     # Optional listing command; (argv, line_prefix). prefix="" means count
     # non-empty lines. Must be a real vendor CLI; do not invent flags.
     list_command: Optional[Tuple[Sequence[str], str]] = None

backends/musa.py

@@ -42,6 +42,11 @@ class MusaBackend(DeviceBackend):
     # timing precision instead of its measurement. Revisit if a newer torch_musa
     # changes the behaviour.
     capture_deny_ops = ("unique2",)
+    # Several FlagGems benchmarks skip on `not torch.cuda.is_available()` while
+    # the code under test dispatches through flag_gems.device. On MThreads that
+    # turns a working operator into "not measured"; the plugin drops those
+    # guards, pinned to one file and one exact skip reason each.
+    compat_plugins = ("_musa_benchmark_compat_plugin",)
     # Device count comes from torch.musa or MUSA_VISIBLE_DEVICES. Do not guess
     # a vendor CLI here; confirm the real listing tool before adding it.
     list_command = None

batch_summary.py

@@ -186,12 +186,82 @@ def summarize(batch_dir: Path, out_path: Optional[Path] = None) -> Path:
     return out_path
 
 
+def ab_rows(before_text: str, after_text: str) -> List[Dict[str, str]]:
+    """把两次 run.log 的测量行按 (op, dtype, size_detail) 配对。
+
+    speedup 列是 gems 对 torch，不是前后对比，所以这里用 gems_lat 自己算
+    before/after - 1。只有两侧都 SUCCESS 且延迟为正的行才计入；对不上的行
+    单独返回，避免用不完整的配对报一个看起来完整的区间。
+    """
+    def index(text: str) -> Dict[Tuple[str, str, str], float]:
+        out: Dict[Tuple[str, str, str], float] = {}
+        for row in parse_run_log(text):
+            if row.get("row_status") != "SUCCESS":
+                continue
+            try:
+                lat = float(row.get("gems_lat", ""))
+            except ValueError:
+                continue
+            if lat > 0:
+                out[(row.get("fg_op_name", ""), row.get("dtype", ""),
+                     row.get("size_detail", ""))] = lat
+        return out
+
+    before, after = index(before_text), index(after_text)
+    paired: List[Dict[str, str]] = []
+    for key in sorted(before.keys() & after.keys()):
+        op, dtype, size = key
+        b, a = before[key], after[key]
+        paired.append({
+            "fg_op_name": op, "dtype": dtype, "size_detail": size,
+            "before_lat": f"{b:.6f}", "after_lat": f"{a:.6f}",
+            "speedup_percent": f"{(b / a - 1.0) * 100.0:.2f}",
+        })
+    return paired
+
+
+def report_ab(before: Path, after: Path, label: str = "") -> List[Dict[str, str]]:
+    """打印一个算子的前后对比：逐行 speedup 与整体区间。"""
+    before_text = before.read_text(encoding="utf-8", errors="replace")
+    after_text = after.read_text(encoding="utf-8", errors="replace")
+    rows = ab_rows(before_text, after_text)
+    name = label or (rows[0]["fg_op_name"] if rows else "?")
+    if not rows:
+        print(f"[ab] {name}: 没有可配对的测量行"
+              f"（before={before} after={after}）")
+        return rows
+    values = [float(r["speedup_percent"]) for r in rows]
+    for r in rows:
+        print(f"  {r['dtype']:<16} {float(r['speedup_percent']):+7.2f}%  "
+              f"{r['before_lat']} -> {r['after_lat']}  {r['size_detail']}")
+    fallback = ""
+    if DIRECT_NO_CUDAGRAPH_RE.search(before_text) or \
+            DIRECT_NO_CUDAGRAPH_RE.search(after_text):
+        # 回退行用的是 do_bench 口径，与 cudagraph 行不同表可比性
+        fallback = "  [含 BENCHMARK_DIRECT_NO_CUDAGRAPH 回退行]"
+    print(f"[ab] {name}: {len(rows)} 行配对，speedup "
+          f"{min(values):+.2f}% ~ {max(values):+.2f}%{fallback}")
+    return rows
+
+
 def main() -> None:
     p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
-    p.add_argument("batch_dir", type=Path, help="run_batch.py 的产物目录")
+    p.add_argument("batch_dir", type=Path, nargs="?",
+                   help="run_batch.py 的产物目录")
     p.add_argument("-o", "--output", type=Path, default=None,
                    help="输出 CSV（默认 <batch_dir>/summary.csv）")
+    p.add_argument("--ab", nargs=2, type=Path, metavar=("BEFORE", "AFTER"),
+                   help="两份 run.log，按 shape 配对输出该算子的 speedup 区间")
+    p.add_argument("--label", default="", help="--ab 报告里的算子名")
     args = p.parse_args()
+    if args.ab:
+        for path in args.ab:
+            if not path.is_file():
+                raise SystemExit(f"run.log 不存在: {path}")
+        report_ab(args.ab[0], args.ab[1], args.label)
+        return
+    if args.batch_dir is None:
+        raise SystemExit("需要 batch_dir，或用 --ab BEFORE AFTER")
     if not args.batch_dir.is_dir():
         raise SystemExit(f"目录不存在: {args.batch_dir}")
     out = summarize(args.batch_dir, args.output)

plugin_plan.py

@@ -41,7 +41,9 @@ def plan_plugins(
     enable_mm_cluster: Optional[bool] = None,
 ) -> List[str]:
     backend = get_backend(backend_name)
-    plugins: List[str] = list(CORE_PLUGINS)
+    # Compat first: these override collection-time markers, so they have to be
+    # in place before anything else inspects the collected items.
+    plugins: List[str] = list(backend.compat_plugins) + list(CORE_PLUGINS)
 
     autotune = (
         backend.supports_triton_hooks

test_backends.py

@@ -173,6 +173,17 @@ class PluginPlanTests(unittest.TestCase):
         self.assertNotIn("_cudagraph_plugin", names)
         self.assertNotIn("_mm_cluster_fix_plugin", names)
 
+    def test_musa_loads_its_compat_plugin_first(self):
+        names = self._plan("musa")
+        self.assertIn("_musa_benchmark_compat_plugin", names)
+        # It rewrites collection-time markers, so nothing may inspect the item
+        # list before it runs.
+        self.assertEqual(names[0], "_musa_benchmark_compat_plugin")
+
+    def test_compat_plugins_do_not_leak_to_other_backends(self):
+        for name in ("cuda", "rocm", "ascend"):
+            self.assertNotIn("_musa_benchmark_compat_plugin", self._plan(name), name)
+
     def test_musa_same_core_without_cluster_fix(self):
         # _cudagraph_plugin is absent from the plan on every backend: the timing
         # plugin loads it, so its absence here says nothing about capture support.