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

refactor: let any backend own graph capture and its failure modes

Browse Source 
Two things kept graph-capture timing tied to NVIDIA, and both are mechanism
rather than policy, so neither backend's behaviour changes here.

The capture path named torch.cuda directly for stream, event, synchronize and
the graph type. FlagTree's do_bench already resolves those through
driver.active.get_device_interface(); its do_bench_cudagraph does not. On a
non-CUDA device every capture therefore raised and every row silently took the
fallback while the log still reported capture as enabled. The path now takes
those calls from the interface, with _graph_class picking CUDAGraph or
MUSAGraph. CUDA still calls upstream's do_bench_cudagraph verbatim, so its
numbers do not move; the port only matters where upstream cannot run.

Failure classification moves onto DeviceBackend, because the text it matches is
vendor wording that has no meaning on another runtime, and gains a third
verdict. A kernel refusing capture is ordinary and falls back. A kernel failing
on the device is fatal and re-raises, since a fallback latency there would
report a broken kernel as a measurement. New: the capture ATTEMPT can leave the
context unusable, and falling back then measures inside a dead process -- the
observed shape of that is a fallback marker followed by a teardown crash and
zero rows. Those raise under BENCHMARK_CAPTURE_POISONED instead, naming the
operator as capture-incompatible rather than broken. capture_deny_ops skips the
attempt entirely for operators already known to do this, because for them the
attempt is the damage and rediscovering it per shape costs the run.

CUDA's marker list is carried over unchanged and no poisoning wording is
guessed for it. The scratch-RNG guard around capture is skipped when an
interface exposes no default_generators, since there is then no registration to
protect.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
This commit is contained in:
AharrypotterandClaude Opus 5 committed 2026-09-09 17:56:32 +08:00
1 parent 783c0655bc
commit 6fa5f5d7a2
4 files changed
+258 -71
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

_cudagraph_plugin.py
+176 -71
View File
@@ -1,18 +1,25 @@
"""pytest plugin: swap `triton.testing.do_bench` for cudagraph-based timing.
"""pytest plugin: swap `triton.testing.do_bench` for graph-capture timing.
Invoked by `_timing_plugin` when the active backend declares
`supports_cuda_graph` (CUDA only). Do not put this module on the default `-p`
list for MUSA/ROCm/Ascend: those backends must start on plain do_bench, not
attempt capture-then-fallback. Monkey-patches do_bench →
`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 CUDA
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
@@ -26,31 +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):
the context is poisoned, so re-raise rather than fall back and record a bogus
latency as success.
def _device_interface():
"""Triton's active device interface: torch.cuda on NVIDIA, torch_musa on MThreads."""
from triton.runtime import driver
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)
return driver.active.get_device_interface()
def _warmup_before_capture(fn, warmup_ms, grad_to_none):
"""Warm fn up before cudagraph capture.
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
@@ -68,8 +77,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:
@@ -78,23 +85,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
@@ -106,12 +113,20 @@ 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 exposes no 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:
@@ -120,43 +135,67 @@ 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 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 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
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]
def _direct(note):
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,
)
print(paint(f"BENCHMARK_DIRECT_NO_CUDAGRAPH op={op} {note}", YELLOW),
flush=True)
return _ORIGINAL_DO_BENCH(
fn,
warmup=warmup,
@@ -166,6 +205,72 @@ def _patched_do_bench(fn, warmup=25, rep=100, grad_to_none=None,
return_mode=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