forked from zhoulin/zl_bench
AharrypotterandClaude Opus 5
9665d36a7c
do_bench already returns a median over the iterations inside one call, which suppresses iteration-level jitter. It cannot see drift across a call: clocks ramping, another process arriving. Direct timing is where that shows, because it walks the full Python launch path once per iteration, so on small kernels the number carries launch overhead and CPU jitter that graph capture removes. ZL_DIRECT_REPEAT=N repeats the call and takes the median over calls, which is the layer the inner median misses. It defaults to 1, so nothing changes unless asked. The knob covers both ways a row ends up on the direct timer: a run that selected timer=direct outright, and the rows a captured run falls back on, whose marker now carries repeat=N since those rows were already a different caliber from the captured ones and can now also be a different sample count. 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 undefined for it. On MThreads this is not a refinement. A convert-reduction screen over 21 operators puts the largest reductions at sort (65->15), dropout (48->0) and rand (24->0), and all three are exactly the operators that cannot be graph-captured there -- RNG seeding and a capture-illegal stream op. The headline numbers on that device are produced by the noisiest timer, which until now had no repetition at all. _cudagraph_plugin imports triton at module scope, so the tests load it against a stub: the repeat logic is arithmetic and should stay testable without a GPU. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
66 lines
2.9 KiB
Python
66 lines
2.9 KiB
Python
"""Select the timing path from backend capabilities, before any shape runs.
|
|
|
|
CUDA keeps the existing `_cudagraph_plugin` behaviour (graph capture, with
|
|
per-shape fallback to plain do_bench). Other backends do not attempt CUDA
|
|
Graph at all: they leave `triton.testing.do_bench` in place and record
|
|
`timer=direct`. Trying capture-then-fallback on a non-CUDA device would
|
|
produce FATAL_DEVICE noise and mixed-caliber rows.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import os
|
|
import sys
|
|
|
|
from backends import get_backend
|
|
from _run_identity import update_identity
|
|
|
|
|
|
def _graph_requested(backend) -> bool:
|
|
raw = os.environ.get("ZL_ENABLE_CUDAGRAPH", "auto").strip().lower()
|
|
if raw in ("0", "false", "no", "off", "never", "direct"):
|
|
return False
|
|
if raw in ("1", "true", "yes", "on", "always"):
|
|
return bool(backend.supports_cuda_graph)
|
|
return bool(backend.supports_cuda_graph)
|
|
|
|
|
|
def pytest_configure(config):
|
|
from _term_style import tag
|
|
backend = get_backend()
|
|
raw = os.environ.get("ZL_ENABLE_CUDAGRAPH", "auto").strip().lower()
|
|
force_on = raw in ("1", "true", "yes", "on", "always")
|
|
if force_on and not backend.supports_cuda_graph:
|
|
print(f"{tag('[timing-plugin]')} ZL_ENABLE_CUDAGRAPH=1 ignored on "
|
|
f"backend={backend.name} (no CUDA Graph); using timer=direct",
|
|
file=sys.stderr, flush=True)
|
|
if _graph_requested(backend):
|
|
import _cudagraph_plugin as _cg
|
|
_cg.pytest_configure(config)
|
|
# Recorded on this branch too: the fallback rows a captured run produces
|
|
# are timed by the direct path and take the same repeat count.
|
|
update_identity(timer="cuda_graph", graph_capture="attempted",
|
|
direct_repeat=_cg.direct_repeat())
|
|
print(f"{tag('[timing-plugin]')} timer=cuda_graph backend={backend.name}",
|
|
file=sys.stderr, flush=True)
|
|
return
|
|
# A wholly-direct run gets the same repetition the fallback rows get, so the
|
|
# "put everything on one timer" advice does not silently cost the samples.
|
|
# Patching only when repeat > 1 keeps the default path untouched: upstream's
|
|
# do_bench stays in place exactly as before.
|
|
import _cudagraph_plugin as _cg
|
|
repeat = _cg.direct_repeat()
|
|
if repeat > 1:
|
|
def _repeated(fn, warmup=25, rep=100, grad_to_none=None,
|
|
quantiles=None, return_mode="mean"):
|
|
return _cg.run_direct(fn, warmup, rep, grad_to_none,
|
|
quantiles, return_mode)
|
|
import triton.testing as _tt
|
|
_tt.do_bench = _repeated
|
|
update_identity(timer="direct", graph_capture="disabled",
|
|
direct_repeat=repeat)
|
|
print(f"{tag('[timing-plugin]')} timer=direct backend={backend.name} "
|
|
f"(upstream do_bench; graph capture not selected)"
|
|
+ (f"; median of {repeat} calls" if repeat > 1 else ""),
|
|
file=sys.stderr, flush=True)
|