forked from zhoulin/zl_bench
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>
53 lines
2.6 KiB
Python
53 lines
2.6 KiB
Python
from __future__ import annotations
|
|
|
|
from .base import DeviceBackend
|
|
|
|
|
|
class MusaBackend(DeviceBackend):
|
|
name = "musa"
|
|
visibility_env_keys = ("MUSA_VISIBLE_DEVICES",)
|
|
torch_namespace = "musa"
|
|
flaggems_vendor = "musa"
|
|
# torch_musa supplies Stream, Event, synchronize, graph and MUSAGraph, so
|
|
# capture works once the timing path takes them from Triton's device
|
|
# interface instead of torch.cuda. Measured on an S5000 over 15 operators:
|
|
# 11 captured, 3 fell back cleanly with a marker (dropout/rand on the RNG
|
|
# seed, sort on a capture-illegal stream op), 1 poisoned the context (see
|
|
# capture_deny_ops). Capture is worth having because it removes launch
|
|
# overhead the direct timer cannot: batch_norm 16x16x8x48 went 0.0758ms ->
|
|
# 0.0089ms, tanh/gelu 4096x4096 about 3.3x, while GEMMs stayed at 1.00x.
|
|
supports_cuda_graph = True
|
|
supports_mm_cluster_fix = False
|
|
supports_triton_hooks = True
|
|
# MUSA reports its own wording; none of the CUDA markers match it.
|
|
fatal_error_markers = (
|
|
"musa error: an illegal memory access",
|
|
"out of memory",
|
|
"device-side assert triggered",
|
|
"misaligned address",
|
|
"uncorrectable ecc",
|
|
)
|
|
# "unknown error" during capture is not recoverable on this runtime: the
|
|
# fallback runs, and the process then dies in teardown with zero rows. Treat
|
|
# it as poisoning so the operator is reported as capture-incompatible rather
|
|
# than silently losing its measurement.
|
|
poisoning_error_markers = ("musa error: unknown error",)
|
|
# Keyed by FLAGGEMS_PERF_CURRENT_OP, which is the pytest FUNCTION name minus
|
|
# `test_` -- `unique2`, `perf_sort` -- not the benchmark file name. Getting
|
|
# that wrong is silent: the entry simply never matches.
|
|
#
|
|
# Observed on an S5000: capturing test_unique2 takes two recoverable
|
|
# fallbacks and then raises "MUSA error: unknown error", ending the run at
|
|
# rc=1 with no measured rows. Skipping the attempt costs that operator's
|
|
# timing precision instead of its measurement. Revisit if a newer torch_musa
|
|
# changes the behaviour.
|
|
capture_deny_ops = ("unique2",)
|
|
# Several FlagGems benchmarks skip on `not torch.cuda.is_available()` while
|
|
# the code under test dispatches through flag_gems.device. On MThreads that
|
|
# turns a working operator into "not measured"; the plugin drops those
|
|
# guards, pinned to one file and one exact skip reason each.
|
|
compat_plugins = ("_musa_benchmark_compat_plugin",)
|
|
# Device count comes from torch.musa or MUSA_VISIBLE_DEVICES. Do not guess
|
|
# a vendor CLI here; confirm the real listing tool before adding it.
|
|
list_command = None
|