zl_bench/_cudagraph_plugin.py

"""pytest plugin: swap `triton.testing.do_bench` for graph-capture timing.

Invoked by `_timing_plugin` when the active backend declares
`supports_cuda_graph`. Do not put this module on the default `-p` list; the
timing plugin owns the decision. 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).

Stream/event/graph calls come from Triton's active device interface, not from
torch.cuda. FlagTree's do_bench is already written that way; its
do_bench_cudagraph is not, so on a non-CUDA backend every capture raised and
every row silently took the fallback while the log still claimed capture was on.
CUDA still calls upstream's do_bench_cudagraph verbatim, so its numbers do not
move. Which failures are recoverable is likewise a backend question, because the
error wording is vendor text; see DeviceBackend.classify_capture_error.
"""

from __future__ import annotations

import os
import sys

import triton.testing as _tt


_ORIGINAL_DO_BENCH = _tt.do_bench


def direct_repeat() -> int:
    """ZL_DIRECT_REPEAT, clamped to >= 1. 1 keeps today's single call."""
    try:
        return max(1, int(os.environ.get("ZL_DIRECT_REPEAT", "1")))
    except ValueError:
        return 1


def run_direct(fn, warmup, rep, grad_to_none, quantiles, return_mode):
    """Upstream do_bench, optionally repeated with the median taken.

    Direct timing walks the full Python launch path once per iteration, so on
    small kernels it carries launch overhead and CPU jitter that graph capture
    does not. do_bench already returns a median over the iterations inside one
    call, which suppresses iteration-level jitter but not drift across a call:
    clocks ramping, another process arriving. Repeating the call and taking the
    median over calls is the part the inner median cannot see.

    It matters most where it is least optional. On MThreads the operators with
    the largest RLC convert reductions -- sort, dropout, rand -- are exactly the
    ones that cannot be graph-captured, so the headline numbers there are
    produced by this path.

    Repetition is skipped when the caller asks for quantiles or return_mode
    "all", because the result is then a sequence and a median across calls is
    not defined for it; the single call is returned unchanged.
    """
    n = direct_repeat()
    if n == 1 or quantiles is not None or return_mode == "all":
        return _ORIGINAL_DO_BENCH(
            fn, warmup=warmup, rep=rep, grad_to_none=grad_to_none,
            quantiles=quantiles, return_mode=return_mode)
    import statistics
    samples = [
        _ORIGINAL_DO_BENCH(fn, warmup=warmup, rep=rep, grad_to_none=grad_to_none,
                           quantiles=None, return_mode=return_mode)
        for _ in range(n)
    ]
    return statistics.median(samples)


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 _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 exposes no 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 anywhere
    else.

    Replay timings go 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>")

    def _direct(note):
        from _term_style import YELLOW, paint
        n = direct_repeat()
        # The repeat count belongs in the marker: these rows are already a
        # different caliber from the captured ones, and now they can also be a
        # different number of samples.
        rpt = "" if n == 1 else f" repeat={n}"
        print(paint(f"BENCHMARK_DIRECT_NO_CUDAGRAPH op={op} {note}{rpt}", YELLOW),
              flush=True)
        return run_direct(fn, warmup, rep, grad_to_none, quantiles, return_mode)

    backend = None
    try:
        from backends import get_backend
        backend = get_backend()
    except Exception:
        pass

    # Known context-poisoning operators never get a capture attempt: the point
    # of the deny list is that for these the attempt itself is the damage, so
    # discovering it again per shape would cost the run, not just the precision.
    if backend is not None and backend.denies_capture(op):
        return _direct("phase=denylist reason=KnownPoisoningOp "
                       f"detail='{backend.name} denies capture for this op'")

    phase = "sync"  # which step rejected capture, for the 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:
        verdict = (backend.classify_capture_error(exc)
                   if backend is not None else "recoverable")
        if verdict == "fatal":
            # The kernel failed on the device. A fallback latency here would
            # report a broken kernel as a measurement.
            raise
        if verdict == "poisoned":
            # The attempt, not the kernel, killed the context. Falling back now
            # would measure inside a dead process -- observed as a fallback
            # marker followed by a teardown crash and zero rows. Surface it
            # under its own marker and stop, so the operator can be re-run with
            # ZL_ENABLE_CUDAGRAPH=0 and added to the backend's deny list.
            from _term_style import RED, paint
            detail = " ".join(str(exc).split())[:160]
            print(paint(f"BENCHMARK_CAPTURE_POISONED op={op} phase={phase} "
                        f"reason={type(exc).__name__} detail={detail!r}", RED),
                  flush=True)
            raise
        detail = " ".join(str(exc).split())[:160]
        return _direct(f"phase={phase} reason={type(exc).__name__} "
                       f"detail={detail!r}")


def pytest_configure(config):
    from _term_style import tag
    try:
        from backends import get_backend
        if not get_backend().supports_cuda_graph:
            print(f"{tag('[cudagraph-plugin]')} skipped: backend "
                  f"{get_backend().name!r} does not support CUDA Graph",
                  file=sys.stderr, flush=True)
            return
    except Exception as exc:
        print(f"{tag('[cudagraph-plugin]')} backend probe failed ({exc}); "
              "not patching do_bench", file=sys.stderr, flush=True)
        return
    _tt.do_bench = _patched_do_bench
    print(f"{tag('[cudagraph-plugin]')} triton.testing.do_bench → do_bench_cudagraph",
          file=sys.stderr, flush=True)