From 6fa5f5d7a2b52549c3abc7bd7575035d664653aa Mon Sep 17 00:00:00 2001 From: Aharrypotter Date: Wed, 9 Sep 2026 17:56:32 +0800 Subject: [PATCH] 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 --- _cudagraph_plugin.py | 247 ++++++++++++++++++++++++++++++------------- backends/base.py | 38 +++++++ backends/cuda.py | 14 +++ test_backends.py | 30 ++++++ 4 files changed, 258 insertions(+), 71 deletions(-) diff --git a/_cudagraph_plugin.py b/_cudagraph_plugin.py index 60e470f..c647432 100644 --- a/_cudagraph_plugin.py +++ b/_cudagraph_plugin.py @@ -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 -`supports_cuda_graph` (CUDA only). Do not put this module on the default `-p` -list for MUSA/ROCm/Ascend: those backends must start on plain do_bench, not -attempt capture-then-fallback. Monkey-patches do_bench → +`supports_cuda_graph`. Do not put this module on the default `-p` list; the +timing plugin owns the decision. Monkey-patches do_bench → 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 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 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 -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 (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 @@ -26,31 +33,33 @@ import triton.testing as _tt _ORIGINAL_DO_BENCH = _tt.do_bench -def _is_fatal_cuda_error(exc: Exception) -> bool: - """True for CUDA context errors (illegal address, OOM, device-side assert): - the context is poisoned, so re-raise rather than fall back and record a bogus - latency as success. +def _device_interface(): + """Triton's active device interface: torch.cuda on NVIDIA, torch_musa on MThreads.""" + from triton.runtime import driver - Markers must match the error text itself, not torch's generic footer: newer - 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) + return driver.active.get_device_interface() -def _warmup_before_capture(fn, warmup_ms, grad_to_none): - """Warm fn up before cudagraph capture. +def _is_cuda(di) -> bool: + 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 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 torch - def _run(): if grad_to_none is not 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. _run() - torch.cuda.synchronize() + di.synchronize() # 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. t0 = time.perf_counter() _run() - torch.cuda.synchronize() + di.synchronize() dt_ms = (time.perf_counter() - t0) * 1000.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 = max(5, min(n, 200)) for _ in range(n): _run() - torch.cuda.synchronize() + di.synchronize() -def _capture_with_scratch_rng(call): - """Run a graph capture under a throwaway CUDA RNG state. +def _capture_with_scratch_rng(di, call): + """Run a graph capture under a throwaway device RNG state. torch.cuda.graph's capture_begin unconditionally registers the default CUDA 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 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 - 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 - 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() - scratch = torch.Generator(device="cuda") + scratch = torch.Generator(device=torch.device(gen.device)) scratch.manual_seed(gen.initial_seed()) gen.graphsafe_set_state(scratch) try: @@ -120,43 +135,67 @@ def _capture_with_scratch_rng(call): 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, 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", "") - phase = "sync" # which step rejected cudagraph, for the fallback marker - try: - # 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] + + def _direct(note): from _term_style import YELLOW, paint - print( - paint(f"BENCHMARK_DIRECT_NO_CUDAGRAPH op={op} phase={phase} " - f"reason={reason} detail={detail!r}", YELLOW), - flush=True, - ) + print(paint(f"BENCHMARK_DIRECT_NO_CUDAGRAPH op={op} {note}", YELLOW), + flush=True) return _ORIGINAL_DO_BENCH( fn, warmup=warmup, @@ -166,6 +205,72 @@ def _patched_do_bench(fn, warmup=25, rep=100, grad_to_none=None, 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): from _term_style import tag diff --git a/backends/base.py b/backends/base.py index 9948eaf..afaef00 100644 --- a/backends/base.py +++ b/backends/base.py @@ -52,6 +52,44 @@ class DeviceBackend: supports_cuda_graph: bool = False supports_mm_cluster_fix: bool = False 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 # non-empty lines. Must be a real vendor CLI; do not invent flags. list_command: Optional[Tuple[Sequence[str], str]] = None diff --git a/backends/cuda.py b/backends/cuda.py index a366728..0c3ae54 100644 --- a/backends/cuda.py +++ b/backends/cuda.py @@ -12,3 +12,17 @@ class CudaBackend(DeviceBackend): supports_mm_cluster_fix = True supports_triton_hooks = True 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 = () diff --git a/test_backends.py b/test_backends.py index e96ddeb..e105d2d 100644 --- a/test_backends.py +++ b/test_backends.py @@ -206,6 +206,36 @@ class PluginPlanTests(unittest.TestCase): os.environ["ZL_EXTRA_PLUGINS"] = old +class CaptureClassificationTests(unittest.TestCase): + """Graph-capture failures split three ways; the wording is vendor text.""" + + def test_recoverable_is_the_default(self): + # A kernel refusing capture is ordinary and must reach the fallback. + v = get_backend("cuda").classify_capture_error( + RuntimeError("operation not permitted when stream is capturing")) + self.assertEqual(v, "recoverable") + + 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_deny_list_defaults_to_empty(self): + self.assertFalse(get_backend("cuda").denies_capture("anything")) + self.assertFalse(get_backend("cuda").denies_capture("")) + + class DeviceArgTests(unittest.TestCase): def test_devices_preferred(self): self.assertEqual(resolve_device_arg("0,1", None), "0,1")