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>

_cudagraph_plugin.py

@@ -33,6 +33,47 @@ import triton.testing as _tt
 _ORIGINAL_DO_BENCH = _tt.do_bench
 
 
+def direct_repeat() -> int:
+    """ZL_DIRECT_REPEAT, clamped to >= 1. 1 keeps today's single call."""
+    try:
+        return max(1, int(os.environ.get("ZL_DIRECT_REPEAT", "1")))
+    except ValueError:
+        return 1
+
+
+def run_direct(fn, warmup, rep, grad_to_none, quantiles, return_mode):
+    """Upstream do_bench, optionally repeated with the median taken.
+
+    Direct timing walks the full Python launch path once per iteration, so on
+    small kernels it carries launch overhead and CPU jitter that graph capture
+    does not. do_bench already returns a median over the iterations inside one
+    call, which suppresses iteration-level jitter but not drift across a call:
+    clocks ramping, another process arriving. Repeating the call and taking the
+    median over calls is the part the inner median cannot see.
+
+    It matters most where it is least optional. On MThreads the operators with
+    the largest RLC convert reductions -- sort, dropout, rand -- are exactly the
+    ones that cannot be graph-captured, so the headline numbers there are
+    produced by this path.
+
+    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
+    not defined for it; the single call is returned unchanged.
+    """
+    n = direct_repeat()
+    if n == 1 or quantiles is not None or return_mode == "all":
+        return _ORIGINAL_DO_BENCH(
+            fn, warmup=warmup, rep=rep, grad_to_none=grad_to_none,
+            quantiles=quantiles, return_mode=return_mode)
+    import statistics
+    samples = [
+        _ORIGINAL_DO_BENCH(fn, warmup=warmup, rep=rep, grad_to_none=grad_to_none,
+                           quantiles=None, return_mode=return_mode)
+        for _ in range(n)
+    ]
+    return statistics.median(samples)
+
+
 def _device_interface():
     """Triton's active device interface: torch.cuda on NVIDIA, torch_musa on MThreads."""
     from triton.runtime import driver
@@ -194,16 +235,14 @@ def _patched_do_bench(fn, warmup=25, rep=100, grad_to_none=None,
 
     def _direct(note):
         from _term_style import YELLOW, paint
-        print(paint(f"BENCHMARK_DIRECT_NO_CUDAGRAPH op={op} {note}", YELLOW),
+        n = direct_repeat()
+        # The repeat count belongs in the marker: these rows are already a
+        # different caliber from the captured ones, and now they can also be a
+        # different number of samples.
+        rpt = "" if n == 1 else f" repeat={n}"
+        print(paint(f"BENCHMARK_DIRECT_NO_CUDAGRAPH op={op} {note}{rpt}", YELLOW),
               flush=True)
-        return _ORIGINAL_DO_BENCH(
-            fn,
-            warmup=warmup,
-            rep=rep,
-            grad_to_none=grad_to_none,
-            quantiles=quantiles,
-            return_mode=return_mode,
-        )
+        return run_direct(fn, warmup, rep, grad_to_none, quantiles, return_mode)
 
     backend = None
     try: