diff --git a/_cudagraph_plugin.py b/_cudagraph_plugin.py index 9c5564b..4334c03 100644 --- a/_cudagraph_plugin.py +++ b/_cudagraph_plugin.py @@ -1,4 +1,4 @@ -"""pytest plugin: swap `triton.testing.do_bench` for cudagraph-based timing. +"""pytest plugin: swap `triton.testing.do_bench` for graph-capture timing. Loaded via `-p _cudagraph_plugin` in run_pytest.sh's PLUGINS list (comment that line out to time with plain do_bench). Monkey-patches do_bench → @@ -8,9 +8,18 @@ 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 -RNG state so that an aborted capture cannot poison the process-wide generator +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). + +Everything here goes through Triton's active device interface rather than +torch.cuda. FlagTree's do_bench is already written that way, but its +do_bench_cudagraph is not: it hardcodes torch.cuda.stream/Event/CUDAGraph, so +on a non-CUDA backend the graph path raised and every measurement silently fell +back to plain do_bench. On MThreads the interface is torch_musa, which supplies +Stream, Event, synchronize, graph and MUSAGraph, so the same capture works once +the calls are taken from the interface. CUDA still calls upstream's +do_bench_cudagraph, so its numbers are unchanged. """ from __future__ import annotations @@ -24,8 +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): +def _device_interface(): + """Triton's active device interface: torch.cuda on NVIDIA, torch_musa on MThreads.""" + from triton.runtime import driver + + return driver.active.get_device_interface() + + +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 _is_fatal_device_error(exc: Exception) -> bool: + """True for device 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. @@ -47,8 +81,12 @@ def _is_fatal_cuda_error(exc: Exception) -> bool: return any(m in msg for m in markers) -def _warmup_before_capture(fn, warmup_ms, grad_to_none): - """Warm fn up before cudagraph capture. +# Retained under its historical name; the marker list is backend independent. +_is_fatal_cuda_error = _is_fatal_device_error + + +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 @@ -66,8 +104,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: @@ -76,23 +112,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 @@ -104,12 +140,21 @@ 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 does not expose 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: @@ -118,24 +163,85 @@ 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 on any + other backend. + + The replay timings are handed 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 + phase = "sync" # which step rejected graph capture, 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() + 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(fn, warmup, grad_to_none) + _warmup_before_capture(di, fn, warmup, grad_to_none) phase = "capture" - return _capture_with_scratch_rng(lambda: _tt.do_bench_cudagraph( + 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, @@ -145,7 +251,7 @@ def _patched_do_bench(fn, warmup=25, rep=100, grad_to_none=None, 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): + if _is_fatal_device_error(exc): raise reason = type(exc).__name__ detail = " ".join(str(exc).split())[:160] @@ -168,5 +274,10 @@ def _patched_do_bench(fn, warmup=25, rep=100, grad_to_none=None, def pytest_configure(config): from _term_style import tag _tt.do_bench = _patched_do_bench - print(f"{tag('[cudagraph-plugin]')} triton.testing.do_bench → do_bench_cudagraph", + try: + backend = getattr(_device_interface(), "__name__", "?") + except Exception: # pragma: no cover - reported, never fatal at configure + backend = "?" + print(f"{tag('[cudagraph-plugin]')} triton.testing.do_bench → graph capture " + f"via {backend}", file=sys.stderr, flush=True)