zhoulin
8ace0d790c
pytest plugin set + driver script for compiler A/B perf comparison: - reproducible runs: fixed seed, yaml-driven shapes, autotune record/replay - per-shape ttgir dump of actually-used variants with readable naming - cudagraph-based timing with documented fallback semantics Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
83 lines
2.9 KiB
Python
83 lines
2.9 KiB
Python
"""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. The warmup kwarg is dropped (do_bench_cudagraph warms up
|
|
internally).
|
|
"""
|
|
|
|
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 _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>")
|
|
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()
|
|
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; only graph-capture rejections fall back.
|
|
if _is_fatal_cuda_error(exc):
|
|
raise
|
|
reason = type(exc).__name__
|
|
detail = " ".join(str(exc).split())[:160]
|
|
print(
|
|
f"BENCHMARK_DIRECT_NO_CUDAGRAPH op={op} reason={reason} detail={detail!r}",
|
|
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):
|
|
_tt.do_bench = _patched_do_bench
|
|
print("[cudagraph-plugin] triton.testing.do_bench → do_bench_cudagraph",
|
|
file=sys.stderr, flush=True)
|