zl_bench/_cudagraph_plugin.py

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

import os
import sys

import triton.testing as _tt


_ORIGINAL_DO_BENCH = _tt.do_bench


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.

    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)


# 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
    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

    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()
    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()
    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()
    di.synchronize()


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

    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

    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=torch.device(gen.device))
    scratch.manual_seed(gen.initial_seed())
    gen.graphsafe_set_state(scratch)
    try:
        return call()
    finally:
        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 graph capture."""
    op = os.environ.get("FLAGGEMS_PERF_CURRENT_OP", "<unknown>")
    phase = "sync"  # which step rejected graph capture, for the fallback 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:
        # 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_device_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
    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)