forked from zhoulin/zl_bench
Aharrypotter
7522143747
Several upstream FlagGems tests gate themselves on torch.cuda.is_available() even though the code under test runs on whatever accelerator FlagGems dispatches to. On MThreads those guards skip the operator instead of exercising it, so a working path reports as "not measured". A guard is dropped only when it is pinned to one file and one exact skip reason, only when FlagGems has dispatched to the MThreads vendor, and only after the unskipped test was observed to pass on an S5000. The plugin walks the node chain rather than the item's own markers, because a module-level pytestmark does not appear on the item. Two guards were probed on device and deliberately left in place, which is why they are named in the docstring rather than silently absent. test_grid_sample.py builds its inputs through a helper whose device argument defaults to the literal "cuda": unskipped, all 90 collected cases fail in aten::empty on a CUDA backend that is not present. test_cutlass_scaled_mm.py gates on vLLM as well as CUDA, so removing its guard would raise on the missing dependency instead of measuring anything.
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,
|
|
)
|