2 Commits
Author SHA1 Message Date
AharrypotterandClaude Opus 5 ac818291a6 feat: turn on graph capture for MUSA
supports_cuda_graph was False for MUSA, so every run there was pinned to
timer=direct. That was correct when it was written -- the capture path named
torch.cuda and could not have worked -- but it is not correct now, and the two
timers are not interchangeable: direct timing walks the full Python launch path
per iteration, so on small kernels it measures mostly launch overhead. Measured
on an S5000, capture against direct: batch_norm 16x16x8x48 0.0758ms -> 0.0089ms,
tanh and gelu at 4096x4096 about 3.3x, GEMMs unchanged at 1.00x. Leaving it off
discards that on every operator that can capture, and it also puts MUSA on a
different caliber from the CUDA numbers the harness is meant to compare against.

torch_musa supplies Stream, Event, synchronize, graph and MUSAGraph, so capture
works through the device interface. Verified on an S5000 over five operators,
one per path:

  tanh      captured, no marker
  dropout   15/15 rows fall back: MUSAGeneratorImpl::current_seed is not
            callable during capture
  rand      same
  perf_sort 25/25 rows fall back: operation not permitted when stream is
            capturing
  unique2   denied up front; rc=0 with 10 rows

MUSA gets its own error wording; none of the CUDA markers match it. "MUSA
error: unknown error" during capture is classified as poisoning rather than
recoverable, because the fallback runs and the process then dies in teardown:
before the deny entry was keyed correctly, unique2 reproduced exactly that,
ending at rc=1 with no rows, but now under an explicit BENCHMARK_CAPTURE_POISONED
marker instead of a bare ExceptionGroup.

capture_deny_ops is keyed by FLAGGEMS_PERF_CURRENT_OP, which holds the pytest
FUNCTION name minus `test_` -- unique2, perf_sort -- not the benchmark file
name. Getting that wrong fails open silently, which is how the first attempt
still hit rc=1 with the entry written as "unique"; a test now pins the key
space. The plugin-plan test is renamed because its assertion that
_cudagraph_plugin is absent from the plan holds on CUDA too and says nothing
about capture support.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
2026-09-09 17:56:56 +08:00
AharrypotterandClaude Opus 5 6fa5f5d7a2 refactor: let any backend own graph capture and its failure modes
Two things kept graph-capture timing tied to NVIDIA, and both are mechanism
rather than policy, so neither backend's behaviour changes here.

The capture path named torch.cuda directly for stream, event, synchronize and
the graph type. FlagTree's do_bench already resolves those through
driver.active.get_device_interface(); its do_bench_cudagraph does not. On a
non-CUDA device every capture therefore raised and every row silently took the
fallback while the log still reported capture as enabled. The path now takes
those calls from the interface, with _graph_class picking CUDAGraph or
MUSAGraph. CUDA still calls upstream's do_bench_cudagraph verbatim, so its
numbers do not move; the port only matters where upstream cannot run.

Failure classification moves onto DeviceBackend, because the text it matches is
vendor wording that has no meaning on another runtime, and gains a third
verdict. A kernel refusing capture is ordinary and falls back. A kernel failing
on the device is fatal and re-raises, since a fallback latency there would
report a broken kernel as a measurement. New: the capture ATTEMPT can leave the
context unusable, and falling back then measures inside a dead process -- the
observed shape of that is a fallback marker followed by a teardown crash and
zero rows. Those raise under BENCHMARK_CAPTURE_POISONED instead, naming the
operator as capture-incompatible rather than broken. capture_deny_ops skips the
attempt entirely for operators already known to do this, because for them the
attempt is the damage and rediscovering it per shape costs the run.

CUDA's marker list is carried over unchanged and no poisoning wording is
guessed for it. The scratch-RNG guard around capture is skipped when an
interface exposes no default_generators, since there is then no registration to
protect.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
2026-09-09 17:56:32 +08:00
5 changed files with 323 additions and 73 deletions

No files matched your search

+176 -71
View File
@@ -1,18 +1,25 @@
"""pytest plugin: swap `triton.testing.do_bench` for cudagraph-based timing. """pytest plugin: swap `triton.testing.do_bench` for graph-capture timing.
Invoked by `_timing_plugin` when the active backend declares Invoked by `_timing_plugin` when the active backend declares
`supports_cuda_graph` (CUDA only). Do not put this module on the default `-p` `supports_cuda_graph`. Do not put this module on the default `-p` list; the
list for MUSA/ROCm/Ascend: those backends must start on plain do_bench, not timing plugin owns the decision. Monkey-patches do_bench →
attempt capture-then-fallback. Monkey-patches do_bench →
do_bench_cudagraph (kernel-only latency) at session configure, after Triton is do_bench_cudagraph (kernel-only latency) at session configure, after Triton is
loaded but before any benchmark calls it. Kernels that can't be graph-captured loaded but before any benchmark calls it. Kernels that can't be graph-captured
fall back to plain do_bench and print a BENCHMARK_DIRECT_NO_CUDAGRAPH marker fall back to plain do_bench and print a BENCHMARK_DIRECT_NO_CUDAGRAPH marker
into the log. Before capturing, fn is explicitly warmed up for the caller's into the log. Before capturing, fn is explicitly warmed up for the caller's
warmup time budget (do_bench_cudagraph's own 5-iter warmup is too short to warmup time budget (do_bench_cudagraph's own 5-iter warmup is too short to
settle autotune/JIT for the MoE kernels, which makes the first captured settle autotune/JIT for the MoE kernels, which makes the first captured
measurement unstable run-to-run). The capture itself runs under a scratch CUDA measurement unstable run-to-run). The capture itself runs under a scratch
RNG state so that an aborted capture cannot poison the process-wide generator RNG state so that an aborted capture cannot poison the process-wide generator
(rationale: _capture_with_scratch_rng). (rationale: _capture_with_scratch_rng).
Stream/event/graph calls come from Triton's active device interface, not from
torch.cuda. FlagTree's do_bench is already written that way; its
do_bench_cudagraph is not, so on a non-CUDA backend every capture raised and
every row silently took the fallback while the log still claimed capture was on.
CUDA still calls upstream's do_bench_cudagraph verbatim, so its numbers do not
move. Which failures are recoverable is likewise a backend question, because the
error wording is vendor text; see DeviceBackend.classify_capture_error.
""" """
from __future__ import annotations from __future__ import annotations
@@ -26,31 +33,33 @@ import triton.testing as _tt
_ORIGINAL_DO_BENCH = _tt.do_bench _ORIGINAL_DO_BENCH = _tt.do_bench
def _is_fatal_cuda_error(exc: Exception) -> bool: def _device_interface():
"""True for CUDA context errors (illegal address, OOM, device-side assert): """Triton's active device interface: torch.cuda on NVIDIA, torch_musa on MThreads."""
the context is poisoned, so re-raise rather than fall back and record a bogus from triton.runtime import driver
latency as success.
Markers must match the error text itself, not torch's generic footer: newer return driver.active.get_device_interface()
torch appends "Compile with `TORCH_USE_CUDA_DSA` to enable device-side
assertions." to EVERY AcceleratorError, so a loose "device-side assert"
marker turned all capture-incompatible ops into fatal re-raises and killed
the documented do_bench fallback path entirely."""
msg = f"{type(exc).__name__}: {exc}".lower()
markers = (
"illegal memory access",
"cudaerrorillegaladdress",
"out of memory",
"device-side assert triggered",
"an illegal instruction",
"misaligned address",
"uncorrectable ecc",
)
return any(m in msg for m in markers)
def _warmup_before_capture(fn, warmup_ms, grad_to_none): def _is_cuda(di) -> bool:
"""Warm fn up before cudagraph capture. import torch
return di is torch.cuda
def _graph_class(di):
"""The interface's graph type: CUDAGraph, MUSAGraph, or whatever it names."""
for name in ("CUDAGraph", "MUSAGraph"):
graph = getattr(di, name, None)
if graph is not None:
return graph
for name in dir(di):
if name.endswith("Graph"):
return getattr(di, name)
raise RuntimeError("device interface exposes no graph type")
def _warmup_before_capture(di, fn, warmup_ms, grad_to_none):
"""Warm fn up before graph capture.
do_bench_cudagraph only warms up 5 iterations internally — too few for the do_bench_cudagraph only warms up 5 iterations internally — too few for the
MoE kernels, where the first calls trigger Triton autotune config MoE kernels, where the first calls trigger Triton autotune config
@@ -68,8 +77,6 @@ def _warmup_before_capture(fn, warmup_ms, grad_to_none):
""" """
import time import time
import torch
def _run(): def _run():
if grad_to_none is not None: if grad_to_none is not None:
for x in grad_to_none: for x in grad_to_none:
@@ -78,23 +85,23 @@ def _warmup_before_capture(fn, warmup_ms, grad_to_none):
# First call absorbs autotune/JIT compilation (can take seconds); discard it. # First call absorbs autotune/JIT compilation (can take seconds); discard it.
_run() _run()
torch.cuda.synchronize() di.synchronize()
# Size the warmup loop from one timed steady-state run. dt_ms includes # Size the warmup loop from one timed steady-state run. dt_ms includes
# launch overhead, so n errs low for tiny kernels; warmup needn't be exact. # launch overhead, so n errs low for tiny kernels; warmup needn't be exact.
t0 = time.perf_counter() t0 = time.perf_counter()
_run() _run()
torch.cuda.synchronize() di.synchronize()
dt_ms = (time.perf_counter() - t0) * 1000.0 dt_ms = (time.perf_counter() - t0) * 1000.0
budget = warmup_ms if (warmup_ms and warmup_ms > 0) else 25.0 budget = warmup_ms if (warmup_ms and warmup_ms > 0) else 25.0
n = int(budget / dt_ms) if dt_ms > 0 else 25 n = int(budget / dt_ms) if dt_ms > 0 else 25
n = max(5, min(n, 200)) n = max(5, min(n, 200))
for _ in range(n): for _ in range(n):
_run() _run()
torch.cuda.synchronize() di.synchronize()
def _capture_with_scratch_rng(call): def _capture_with_scratch_rng(di, call):
"""Run a graph capture under a throwaway CUDA RNG state. """Run a graph capture under a throwaway device RNG state.
torch.cuda.graph's capture_begin unconditionally registers the default CUDA torch.cuda.graph's capture_begin unconditionally registers the default CUDA
RNG generator and sets its capturing flag; when a capture aborts midway the RNG generator and sets its capturing flag; when a capture aborts midway the
@@ -106,12 +113,20 @@ def _capture_with_scratch_rng(call):
before capturing and restore the original state object afterwards, success before capturing and restore the original state object afterwards, success
or failure. The real generator never takes part in a capture, so poisoning or failure. The real generator never takes part in a capture, so poisoning
can only land on the discarded scratch state, and _seed_plugin determinism can only land on the discarded scratch state, and _seed_plugin determinism
is left untouched.""" is left untouched.
An interface that exposes no default_generators has no such registration to
protect against (torch_musa is one), so the capture simply runs."""
import torch import torch
dev = torch.cuda.current_device()
gen = torch.cuda.default_generators[dev] generators = getattr(di, "default_generators", None)
if not generators:
return call()
gen = generators[di.current_device()]
if not hasattr(gen, "graphsafe_get_state"):
return call()
saved = gen.graphsafe_get_state() saved = gen.graphsafe_get_state()
scratch = torch.Generator(device="cuda") scratch = torch.Generator(device=torch.device(gen.device))
scratch.manual_seed(gen.initial_seed()) scratch.manual_seed(gen.initial_seed())
gen.graphsafe_set_state(scratch) gen.graphsafe_set_state(scratch)
try: try:
@@ -120,43 +135,67 @@ def _capture_with_scratch_rng(call):
gen.graphsafe_set_state(saved) gen.graphsafe_set_state(saved)
def _do_bench_graph(di, fn, rep, grad_to_none, quantiles, return_mode, n_retries=10):
"""do_bench_cudagraph, taken from the device interface instead of torch.cuda.
A step-for-step port of upstream's implementation: estimate one iteration,
unroll that many calls into a captured graph so host launch overhead drops
out, then time independent replays. Upstream is used verbatim on CUDA; this
exists because upstream names torch.cuda directly and so cannot run anywhere
else.
Replay timings go to _summarize_statistics as a plain list of floats, which
is what both do_bench and upstream's do_bench_cudagraph pass: this fork's
helper runs Python's statistics module over them, so wrapping them in a
tensor makes "mean" raise and makes a quantile request return tensors that
later fail to serialize.
"""
with di.stream(di.Stream()):
fn()
if grad_to_none is not None:
for x in grad_to_none:
x.detach_()
x.requires_grad_(True)
x.grad = None
start_event = di.Event(enable_timing=True)
end_event = di.Event(enable_timing=True)
start_event.record()
for _ in range(5):
fn()
end_event.record()
di.synchronize()
estimate_ms = start_event.elapsed_time(end_event) / 5
n_repeat = 1000 if estimate_ms == 0 else max(1, int(rep / estimate_ms))
graph = _graph_class(di)()
with di.graph(graph):
for _ in range(n_repeat):
if grad_to_none is not None:
for x in grad_to_none:
x.grad = None
fn()
di.synchronize()
ret = []
for _ in range(n_retries):
start_event = di.Event(enable_timing=True)
end_event = di.Event(enable_timing=True)
start_event.record()
graph.replay()
end_event.record()
di.synchronize()
ret.append(start_event.elapsed_time(end_event) / n_repeat)
del graph
return _tt._summarize_statistics(ret, quantiles, return_mode)
def _patched_do_bench(fn, warmup=25, rep=100, grad_to_none=None, def _patched_do_bench(fn, warmup=25, rep=100, grad_to_none=None,
quantiles=None, return_mode="mean"): quantiles=None, return_mode="mean"):
"""Drop-in replacement for triton.testing.do_bench using cudagraph capture.""" """Drop-in replacement for triton.testing.do_bench using graph capture."""
op = os.environ.get("FLAGGEMS_PERF_CURRENT_OP", "<unknown>") op = os.environ.get("FLAGGEMS_PERF_CURRENT_OP", "<unknown>")
phase = "sync" # which step rejected cudagraph, for the fallback marker
try: def _direct(note):
# do_bench_cudagraph runs fn on a fresh side stream without syncing
# with the default stream first; input tensors produced just before
# (quantize/topk in the input iter) may still be in flight, and racing
# on them intermittently kills the context (illegal instruction from a
# device trap on garbage indices). Sync before switching streams.
import torch
torch.cuda.synchronize()
phase = "warmup"
# Settle autotune/JIT before capture (rationale: _warmup_before_capture).
_warmup_before_capture(fn, warmup, grad_to_none)
phase = "capture"
return _capture_with_scratch_rng(lambda: _tt.do_bench_cudagraph(
fn,
rep=rep,
grad_to_none=grad_to_none,
quantiles=quantiles,
return_mode=return_mode,
))
except Exception as exc:
# Re-raise genuine device failures; anything else falls back to plain
# do_bench, with the failing phase (sync/warmup/capture) in the marker.
if _is_fatal_cuda_error(exc):
raise
reason = type(exc).__name__
detail = " ".join(str(exc).split())[:160]
from _term_style import YELLOW, paint from _term_style import YELLOW, paint
print( print(paint(f"BENCHMARK_DIRECT_NO_CUDAGRAPH op={op} {note}", YELLOW),
paint(f"BENCHMARK_DIRECT_NO_CUDAGRAPH op={op} phase={phase} " flush=True)
f"reason={reason} detail={detail!r}", YELLOW),
flush=True,
)
return _ORIGINAL_DO_BENCH( return _ORIGINAL_DO_BENCH(
fn, fn,
warmup=warmup, warmup=warmup,
@@ -166,6 +205,72 @@ def _patched_do_bench(fn, warmup=25, rep=100, grad_to_none=None,
return_mode=return_mode, return_mode=return_mode,
) )
backend = None
try:
from backends import get_backend
backend = get_backend()
except Exception:
pass
# Known context-poisoning operators never get a capture attempt: the point
# of the deny list is that for these the attempt itself is the damage, so
# discovering it again per shape would cost the run, not just the precision.
if backend is not None and backend.denies_capture(op):
return _direct("phase=denylist reason=KnownPoisoningOp "
f"detail='{backend.name} denies capture for this op'")
phase = "sync" # which step rejected capture, for the marker
try:
di = _device_interface()
# The graph path runs fn on a fresh side stream without syncing with the
# default stream first; input tensors produced just before (quantize/topk
# in the input iter) may still be in flight, and racing on them
# intermittently kills the context (illegal instruction from a device
# trap on garbage indices). Sync before switching streams.
di.synchronize()
phase = "warmup"
# Settle autotune/JIT before capture (rationale: _warmup_before_capture).
_warmup_before_capture(di, fn, warmup, grad_to_none)
phase = "capture"
if _is_cuda(di):
return _capture_with_scratch_rng(di, lambda: _tt.do_bench_cudagraph(
fn,
rep=rep,
grad_to_none=grad_to_none,
quantiles=quantiles,
return_mode=return_mode,
))
return _capture_with_scratch_rng(di, lambda: _do_bench_graph(
di,
fn,
rep=rep,
grad_to_none=grad_to_none,
quantiles=quantiles,
return_mode=return_mode,
))
except Exception as exc:
verdict = (backend.classify_capture_error(exc)
if backend is not None else "recoverable")
if verdict == "fatal":
# The kernel failed on the device. A fallback latency here would
# report a broken kernel as a measurement.
raise
if verdict == "poisoned":
# The attempt, not the kernel, killed the context. Falling back now
# would measure inside a dead process -- observed as a fallback
# marker followed by a teardown crash and zero rows. Surface it
# under its own marker and stop, so the operator can be re-run with
# ZL_ENABLE_CUDAGRAPH=0 and added to the backend's deny list.
from _term_style import RED, paint
detail = " ".join(str(exc).split())[:160]
print(paint(f"BENCHMARK_CAPTURE_POISONED op={op} phase={phase} "
f"reason={type(exc).__name__} detail={detail!r}", RED),
flush=True)
raise
detail = " ".join(str(exc).split())[:160]
return _direct(f"phase={phase} reason={type(exc).__name__} "
f"detail={detail!r}")
def pytest_configure(config): def pytest_configure(config):
from _term_style import tag from _term_style import tag
+38
View File
@@ -52,6 +52,44 @@ class DeviceBackend:
supports_cuda_graph: bool = False supports_cuda_graph: bool = False
supports_mm_cluster_fix: bool = False supports_mm_cluster_fix: bool = False
supports_triton_hooks: bool = True supports_triton_hooks: bool = True
# --- graph-capture failure classification -----------------------------
# A kernel that refuses capture is normal and falls back. A device error is
# not, and the two are told apart by vendor error text, so the wording lives
# on the backend rather than in the timing plugin.
#
# fatal : the kernel itself failed on the device. Re-raise; recording a
# fallback latency here would report a broken kernel as a result.
# poisoning: the capture ATTEMPT left the context unusable. Falling back runs
# the measurement inside a dead process, which is how an operator
# ends up emitting a fallback marker and then still dying at
# teardown with zero rows. Re-raise too, but under its own marker,
# because the operator is capture-incompatible rather than broken:
# the actionable response is to re-run it with ZL_ENABLE_CUDAGRAPH=0.
fatal_error_markers: Tuple[str, ...] = ()
poisoning_error_markers: Tuple[str, ...] = ()
# Operators whose capture attempt is known to poison the context on this
# backend. Skipped up front so one of them cannot cost a whole run.
capture_deny_ops: Tuple[str, ...] = ()
def classify_capture_error(self, exc: BaseException) -> str:
"""-> "fatal" | "poisoned" | "recoverable".
Matched against the exception text itself, never a generic footer: torch
appends "Compile with `TORCH_USE_CUDA_DSA` to enable device-side
assertions." to every AcceleratorError, so a loose "device-side assert"
marker would classify every capture refusal as fatal and disable the
documented fallback path outright.
"""
msg = f"{type(exc).__name__}: {exc}".lower()
if any(m in msg for m in self.fatal_error_markers):
return "fatal"
if any(m in msg for m in self.poisoning_error_markers):
return "poisoned"
return "recoverable"
def denies_capture(self, op: str) -> bool:
return bool(op) and op in self.capture_deny_ops
# Optional listing command; (argv, line_prefix). prefix="" means count # Optional listing command; (argv, line_prefix). prefix="" means count
# non-empty lines. Must be a real vendor CLI; do not invent flags. # non-empty lines. Must be a real vendor CLI; do not invent flags.
list_command: Optional[Tuple[Sequence[str], str]] = None list_command: Optional[Tuple[Sequence[str], str]] = None
+14
View File
@@ -12,3 +12,17 @@ class CudaBackend(DeviceBackend):
supports_mm_cluster_fix = True supports_mm_cluster_fix = True
supports_triton_hooks = True supports_triton_hooks = True
list_command = (("nvidia-smi", "--list-gpus"), "GPU ") list_command = (("nvidia-smi", "--list-gpus"), "GPU ")
# Unchanged from when this list lived in _cudagraph_plugin; CUDA numbers do
# not move because of the classification refactor.
fatal_error_markers = (
"illegal memory access",
"cudaerrorillegaladdress",
"out of memory",
"device-side assert triggered",
"an illegal instruction",
"misaligned address",
"uncorrectable ecc",
)
# No CUDA wording is known to poison the context without also matching one
# of the fatal markers above, so nothing is listed here rather than guessed.
poisoning_error_markers = ()
+32 -1
View File
@@ -8,9 +8,40 @@ class MusaBackend(DeviceBackend):
visibility_env_keys = ("MUSA_VISIBLE_DEVICES",) visibility_env_keys = ("MUSA_VISIBLE_DEVICES",)
torch_namespace = "musa" torch_namespace = "musa"
flaggems_vendor = "musa" flaggems_vendor = "musa"
supports_cuda_graph = False # 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_mm_cluster_fix = False
supports_triton_hooks = True 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",)
# Device count comes from torch.musa or MUSA_VISIBLE_DEVICES. Do not guess # 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. # a vendor CLI here; confirm the real listing tool before adding it.
list_command = None list_command = None
+63 -1
View File
@@ -173,7 +173,9 @@ class PluginPlanTests(unittest.TestCase):
self.assertNotIn("_cudagraph_plugin", names) self.assertNotIn("_cudagraph_plugin", names)
self.assertNotIn("_mm_cluster_fix_plugin", names) self.assertNotIn("_mm_cluster_fix_plugin", names)
def test_musa_same_core_without_graph_or_cluster(self): def test_musa_same_core_without_cluster_fix(self):
# _cudagraph_plugin is absent from the plan on every backend: the timing
# plugin loads it, so its absence here says nothing about capture support.
names = self._plan("musa") names = self._plan("musa")
self.assertIn("_timing_plugin", names) self.assertIn("_timing_plugin", names)
self.assertNotIn("_cudagraph_plugin", names) self.assertNotIn("_cudagraph_plugin", names)
@@ -206,6 +208,66 @@ class PluginPlanTests(unittest.TestCase):
os.environ["ZL_EXTRA_PLUGINS"] = old os.environ["ZL_EXTRA_PLUGINS"] = old
class CaptureClassificationTests(unittest.TestCase):
"""Graph-capture failures split three ways; the wording is vendor text."""
def test_cuda_and_musa_both_capture(self):
self.assertTrue(get_backend("cuda").supports_cuda_graph)
# torch_musa supplies Stream/Event/graph/MUSAGraph; measured on S5000.
self.assertTrue(get_backend("musa").supports_cuda_graph)
def test_recoverable_is_the_default(self):
# A kernel refusing capture is ordinary and must reach the fallback.
for name in ("cuda", "musa"):
v = get_backend(name).classify_capture_error(
RuntimeError("operation not permitted when stream is capturing"))
self.assertEqual(v, "recoverable", name)
def test_cuda_fatal_markers(self):
b = get_backend("cuda")
self.assertEqual(
b.classify_capture_error(RuntimeError(
"CUDA error: an illegal memory access was encountered")), "fatal")
self.assertEqual(
b.classify_capture_error(RuntimeError("CUDA out of memory")), "fatal")
def test_torch_dsa_footer_does_not_make_everything_fatal(self):
# torch appends this to every AcceleratorError; matching it loosely
# would classify each capture refusal as fatal and kill the fallback.
exc = RuntimeError("Cannot call current_seed during capture. Compile with "
"`TORCH_USE_CUDA_DSA` to enable device-side assertions.")
self.assertEqual(get_backend("cuda").classify_capture_error(exc),
"recoverable")
def test_musa_unknown_error_is_poisoning_not_recoverable(self):
b = get_backend("musa")
exc = RuntimeError("MUSA error: unknown error")
self.assertEqual(b.classify_capture_error(exc), "poisoned")
# The same text carries no meaning for CUDA, so it must not leak across.
self.assertEqual(get_backend("cuda").classify_capture_error(exc),
"recoverable")
def test_musa_rng_seed_failure_stays_recoverable(self):
exc = RuntimeError("Cannot call MUSAGeneratorImpl::current_seed during "
"MUSA graph capture")
self.assertEqual(get_backend("musa").classify_capture_error(exc),
"recoverable")
def test_deny_list_is_per_backend(self):
self.assertTrue(get_backend("musa").denies_capture("unique2"))
self.assertFalse(get_backend("musa").denies_capture("softmax"))
self.assertFalse(get_backend("cuda").denies_capture("unique2"))
self.assertFalse(get_backend("musa").denies_capture(""))
def test_deny_list_is_keyed_by_pytest_function_name(self):
# FLAGGEMS_PERF_CURRENT_OP holds the pytest function name minus `test_`,
# so a deny entry written as the benchmark FILE name never matches and
# fails open. Observed once: "unique" listed, op=unique2 at runtime.
musa = get_backend("musa")
self.assertTrue(musa.denies_capture("unique2"))
self.assertFalse(musa.denies_capture("unique"))
class DeviceArgTests(unittest.TestCase): class DeviceArgTests(unittest.TestCase):
def test_devices_preferred(self): def test_devices_preferred(self):
self.assertEqual(resolve_device_arg("0,1", None), "0,1") self.assertEqual(resolve_device_arg("0,1", None), "0,1")