Aharrypotter/zl_bench
1
0 Fork 0
forked from zhoulin/zl_bench
Code Pull Requests Activity

fix: take graph-capture timing from Triton's device interface

Browse Source 
The cudagraph plugin swapped do_bench for do_bench_cudagraph, but upstream's
do_bench_cudagraph names torch.cuda directly: torch.cuda.stream, Event,
synchronize, CUDAGraph, graph. FlagTree's do_bench already resolves the device
through driver.active.get_device_interface(), its cudagraph twin does not. On a
non-CUDA backend every capture therefore raised and every measurement silently
fell back to plain do_bench, so the plugin looked enabled while nothing was
being captured.

The plugin now resolves the interface itself and takes the stream, event,
synchronize and graph calls from it. CUDA still calls upstream's
do_bench_cudagraph verbatim, so its numbers do not move; other backends get an
otherwise step-for-step port. The scratch-RNG guard around capture is skipped
when the interface exposes no default_generators, which is the case on MThreads
and where there is no generator registration to protect against.

The replay timings are handed to _summarize_statistics as a plain list of
floats, matching do_bench and upstream's cudagraph path. 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.

Verified on a Moore Threads S5000 through torch_musa, which supplies Stream,
Event, synchronize, graph and MUSAGraph. softmax and gelu_backward capture and
replay with no fallback marker. RNG operators cannot be captured there yet:
dropout aborts with "Cannot call MUSAGeneratorImpl::current_seed during MUSA
graph capture" and takes the documented BENCHMARK_DIRECT_NO_CUDAGRAPH fallback,
so a corpus mixing RNG and non-RNG operators mixes two timing regimes and the
marker is the thing to grep before comparing them.
This commit is contained in:
Aharrypotter committed 2026-09-02 12:19:07 +08:00
1 parent 64baa11a23
commit 3c9529dab2
1 file changed
+142 -31
Download Patch File Download Diff File Expand all files Collapse all files
_cudagraph_plugin.py
+142 -31
View File
@@ -1,4 +1,4 @@
"""pytest plugin: swap `triton.testing.do_bench` for cudagraph-based timing.
"""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 →
@@ -8,9 +8,18 @@ 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
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
@@ -24,8 +33,33 @@ 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):
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.
@@ -47,8 +81,12 @@ def _is_fatal_cuda_error(exc: Exception) -> bool:
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.
# 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
@@ -66,8 +104,6 @@ def _warmup_before_capture(fn, warmup_ms, grad_to_none):
"""
import time
import torch
def _run():
if grad_to_none is not None:
for x in grad_to_none:
@@ -76,23 +112,23 @@ def _warmup_before_capture(fn, warmup_ms, grad_to_none):
# First call absorbs autotune/JIT compilation (can take seconds); discard it.
_run()
torch.cuda.synchronize()
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()
torch.cuda.synchronize()
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()
torch.cuda.synchronize()
di.synchronize()
def _capture_with_scratch_rng(call):
"""Run a graph capture under a throwaway CUDA RNG state.
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
@@ -104,12 +140,21 @@ def _capture_with_scratch_rng(call):
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."""
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
dev = torch.cuda.current_device()
gen = torch.cuda.default_generators[dev]
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="cuda")
scratch = torch.Generator(device=torch.device(gen.device))
scratch.manual_seed(gen.initial_seed())
gen.graphsafe_set_state(scratch)
try:
@@ -118,24 +163,85 @@ def _capture_with_scratch_rng(call):
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 cudagraph capture."""
"""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 cudagraph, for the fallback marker
phase = "sync" # which step rejected graph capture, 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()
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(fn, warmup, grad_to_none)
_warmup_before_capture(di, fn, warmup, grad_to_none)
phase = "capture"
return _capture_with_scratch_rng(lambda: _tt.do_bench_cudagraph(
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,
@@ -145,7 +251,7 @@ def _patched_do_bench(fn, warmup=25, rep=100, grad_to_none=None,
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):
if _is_fatal_device_error(exc):
raise
reason = type(exc).__name__
detail = " ".join(str(exc).split())[:160]
@@ -168,5 +274,10 @@ def _patched_do_bench(fn, warmup=25, rep=100, grad_to_none=None,
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",
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)