feat: allow the direct timer to repeat and take the median

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>
This commit is contained in:
AharrypotterandClaude Opus 5 committed 2026-09-10 12:56:08 +08:00
1 parent d7cb7efc94
commit 9665d36a7c
4 files changed
+147 -12

No files matched your search

+21 -3
View File
@@ -37,11 +37,29 @@ def pytest_configure(config):
if _graph_requested(backend):
import _cudagraph_plugin as _cg
_cg.pytest_configure(config)
update_identity(timer="cuda_graph", graph_capture="attempted")
# 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
update_identity(timer="direct", graph_capture="disabled")
# 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} "
"(upstream do_bench; graph capture not selected)",
f"(upstream do_bench; graph capture not selected)"
+ (f"; median of {repeat} calls" if repeat > 1 else ""),
file=sys.stderr, flush=True)