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")