diff --git a/_musa_benchmark_compat_plugin.py b/_musa_benchmark_compat_plugin.py new file mode 100644 index 0000000..6a8ab87 --- /dev/null +++ b/_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, + )