zl_bench/_cudagraph_plugin.py

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