forked from zhoulin/zl_bench
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:
1 parent
c9f9b7885f
commit
29af59be9b
8 files changed
+258
-3
No files matched your search
@@ -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,
|
||||
)
|
||||
Reference in new issue
Block a user