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>
This commit is contained in:
AharrypotterandClaude Opus 5 committed 2026-09-10 10:31:59 +08:00
1 parent c9f9b7885f
commit 29af59be9b
8 files changed
+258 -3

No files matched your search

+89
View File
@@ -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,
)