Files
zl_bench/backends/musa.py
T
Aharrypotter ee5f70fa89 bench: record device clocks, temperature and power for the whole campaign
A 6-round A/B here reports side noise at the 0.01% level. Clock throttling and
thermal drift move latency by far more than that and leave no trace in
run.log, so the report comes out looking exactly as clean as a run on a stable
machine. The NV run that reproduced PR763 sampled nvidia-smi every 5s
alongside the campaign; the MUSA campaigns had no equivalent, and the parallel
plan makes it worse by heating the whole box.

telemetry.py samples through the backend, so cuda and musa emit the same
columns (index, sm_clock_mhz, mem_clock_mhz, temp_c, power_w, util_pct) and
two machines' campaigns can be read side by side. mthreads-gmi offers no CSV
query mode, so the MUSA path parses --json and strips the units glued onto
each value ("36C", "1888MHz", "477.68W"); its "Power Draw " key carries a
trailing space, matched by stripping rather than by literal.

The device comes from the backend's own visibility variable when not pinned. A
list of devices resolves to no -i at all, sampling every one, rather than
silently picking the first. Dropping the placeholder drops its flag too: a
bare trailing -i fails the sampler instead of widening it.

Telemetry is evidence, not a gate. A missing or unpermitted sampler exits
quietly with its reason in telemetry.err and the campaign runs unchanged.
2026-09-10 16:00:01 +08:00

102 lines
4.3 KiB
Python

from __future__ import annotations
from typing import Dict, List
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
# mthreads-gmi has no CSV query mode; --json is the only machine-readable
# form it offers. Values arrive with their units glued on ("36C", "1888MHz",
# "477.68W", "87%"), so the parser strips trailing non-numerics rather than
# trusting a fixed suffix.
telemetry_argv = (
"mthreads-gmi", "-q",
"-d", "TEMPERATURE,CLOCK,POWER,UTILIZATION",
"--json",
"-i", "{dev}",
)
def parse_telemetry(self, raw: str) -> List[Dict[str, str]]:
import json
import re
def num(v):
m = re.match(r"^\s*(-?[\d.]+)", str(v or ""))
return m.group(1) if m else ""
start = raw.find("{")
if start < 0:
return []
try:
data = json.loads(raw[start:])
except Exception:
return []
gpus = data.get("GPU") or []
if isinstance(gpus, dict):
gpus = [gpus]
out: List[Dict[str, str]] = []
for g in gpus:
clocks = g.get("Clocks") or {}
power = g.get("Power Readings") or {}
# 键名带尾随空格("Power Draw "),别按字面写死
draw = next((v for k, v in power.items()
if k.strip() == "Power Draw"), "")
out.append({
"index": str(g.get("Index", "")).strip(),
"sm_clock_mhz": num(clocks.get("Graphics")),
"mem_clock_mhz": num(clocks.get("Memory")),
"temp_c": num((g.get("Temperature") or {}).get("GPU Current Temp")),
"power_w": num(draw),
"util_pct": num((g.get("Utilization") or {}).get("Gpu")),
})
return out