forked from zhoulin/zl_bench
AharrypotterandClaude Opus 5
29af59be9b
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>
90 lines
3.4 KiB
Python
90 lines
3.4 KiB
Python
"""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,
|
|
)
|