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

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."""
    msg = f"{type(exc).__name__}: {exc}".lower()
    markers = (
        "illegal memory access",
        "cudaerrorillegaladdress",
        "out of memory",
        "device-side assert",
        "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 _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 _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)