"""pytest plugin: swap `triton.testing.do_bench` for cudagraph-based 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 → 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 RNG state so that an aborted capture cannot poison the process-wide generator (rationale: _capture_with_scratch_rng). """ from __future__ import annotations import os import sys 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. 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) def _warmup_before_capture(fn, warmup_ms, grad_to_none): """Warm fn up before cudagraph capture. do_bench_cudagraph only warms up 5 iterations internally — too few for the MoE kernels, where the first calls trigger Triton autotune config compilation, libtuner selection and lazy JIT/init. If that work leaks into the captured graph or the first timed iteration the latency is unstable run-to-run (most visibly on the small-M multi-kernel path). Absorb the one-off compile with a throwaway run, then loop for the caller's warmup time budget so the graph is captured at steady state. The loop is clamped to 5..200 iterations: the cap keeps sub-ms kernels from spinning through a large budget (FlagGems' Config.warm_up defaults to 1000 ms) — 200 iterations settles them fine in practice — and the floor covers kernels whose single run exceeds the budget. If first-run jitter reappears on a new op, raising the cap is the first knob to try. """ import time import torch def _run(): if grad_to_none is not None: for x in grad_to_none: x.grad = None fn() # First call absorbs autotune/JIT compilation (can take seconds); discard it. _run() torch.cuda.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() 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() def _capture_with_scratch_rng(call): """Run a graph capture under a throwaway CUDA 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 epilogue may never run, the flag sticks, and every later CUDA RNG call in the process (the next shape's torch.randn included) raises "Offset increment outside graph capture". Repairing after the fact is unreliable — when the poisoning surfaces depends on which stage the capture died in — so prevent it instead: swap in a scratch state via graphsafe_set_state 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.""" import torch dev = torch.cuda.current_device() gen = torch.cuda.default_generators[dev] saved = gen.graphsafe_get_state() scratch = torch.Generator(device="cuda") scratch.manual_seed(gen.initial_seed()) gen.graphsafe_set_state(scratch) try: return call() finally: gen.graphsafe_set_state(saved) 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.""" 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] 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, ) return _ORIGINAL_DO_BENCH( fn, warmup=warmup, rep=rep, grad_to_none=grad_to_none, quantiles=quantiles, return_mode=return_mode, ) 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", file=sys.stderr, flush=True)