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>

README.md

@@ -67,6 +67,7 @@ softmax:
 | `USE_FLAGTUNE` | `0` | `0` 走普通 autotune（跳过搜索、快速验证）；`1` 用 FlagTune 扩展调优空间（首跑全量搜索、慢）。A/B 两侧须取同值 |
 | `BENCH_LEVEL` | `core` | 透传上游 `--level`。`comprehensive` 相对 `core` 做两件事：① 在 `set_shapes()` 里并入子类 `set_more_shapes()`（`mm` 因此从 5 个 shape 变 29 个，含 N/K 到 152064 的大 K 与 `models_shapes/*.yaml` 的模型 shape）；② 在 `get_input_iter()` 里对每个 shape 再产出一次变体（BLAS 是转置侧），用例数翻倍。注意 ① 只在 `SHAPE_MODE=upstream` 下可见——传了 `SHAPE_FILE` 时 `_shape_inject_plugin` 会在上游合并之后再覆盖一次、把 shape 集拉回 yaml 内容（BLAS 那类 arity 不一致的情形除外，插件会退让）；② 不经插件，任何 shape 来源下都生效。sweep 成本同步放大，shape 多时配合 `PARALLEL_WARMUP_GPUS`。A/B 两侧须取同值 |
 | `ZL_ISOLATE_TUNER_CACHE` | `0` | FlagGems 的 libtuner winner 存在 `$FLAGGEMS_CACHE_DIR/config_cache` 的 sqlite 里（默认 `~/.flaggems`），跨 run 长存。`1` 给本次 run 一个私有缓存目录，两侧各自冷调优（README「各自最优」口径），跑完随 Triton 缓存一起删——选中的 config 已在 `autotune_records/` 里，DB 本身无留存价值。共享缓存的风险是单向的：先跑的一侧写入 winner，后跑的一侧直接命中、只编译那几个胜出变体。实测 `router_gemm` 共享模式下第二次运行编译变体从 234 掉到 10、耗时 52s→17s，隔离模式下两次都是 234 个变体。走 `REPLAY_FROM` 的标准 A/B 不受影响（config 本就锁死），record 对 record 才会中招。取值无论如何都写进 `run_identity.json` 的 `tuner_cache` |
+| `ZL_DIRECT_REPEAT` | `1` | direct 计时的重复次数，取中位。只作用于 direct 路径：整体 `timer=direct` 的运行，以及图捕获运行里逐行回退的那些点（标记会带上 `repeat=N`）；捕获成功的行不受影响。`do_bench` 自己已在调用内取中位，压住的是迭代间抖动；跨调用重复补的是它看不见的那层——时钟爬坡、邻居进程进来这类漂移。调用方要 `quantiles` 或 `return_mode="all"` 时返回序列、跨调用中位无定义，自动不重复。在 MThreads 上这不是可选项：cvt 消除量最大的 `sort`/`dropout`/`rand` 恰好都无法捕获，那批结论正是由这条路径产出的。取值写进 `run_identity.json` 的 `direct_repeat` |
 | `DTYPES` | 空（上游默认 dtype 扫描） | 空格分隔的 dtype 白名单（如 `"bfloat16 float16"`），逐个转为上游 `--dtypes`。算子不支持时上游报 `can't be supported by this op` |
 | `PARALLEL_WARMUP_GPUS` | 空（串行） | 设为 `N`（≥2）时把 autotune sweep 分片到 N 卡并行，测量仍单卡串行；不设或 `<2` 则完全串行（见「加速：多卡并行 sweep」） |
 | `ZL_ENABLE_CUDAGRAPH` | `auto` | `auto`/`1` 仅在 CUDA 上走 Graph 计时；非 CUDA 即使设 `1` 也拒绝并保持 `timer=direct` |

_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:

_timing_plugin.py

@@ -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)

test_backends.py

@@ -10,6 +10,7 @@ import json
 import os
 import sys
 import tempfile
+import types
 import unittest
 from pathlib import Path
 
@@ -279,6 +280,82 @@ class CaptureClassificationTests(unittest.TestCase):
         self.assertFalse(musa.denies_capture("unique"))
 
 
+class DirectRepeatTests(unittest.TestCase):
+    """Repetition on the direct timer.
+
+    _cudagraph_plugin imports triton at module scope, so it is loaded here
+    against a stub: the repeat logic is plain arithmetic and should stay
+    testable on a machine with no GPU and no Triton.
+    """
+
+    @classmethod
+    def setUpClass(cls):
+        cls._saved = {k: sys.modules.get(k) for k in ("triton", "triton.testing")}
+        triton = types.ModuleType("triton")
+        testing = types.ModuleType("triton.testing")
+        testing.do_bench = lambda *a, **k: 0.0
+        testing.do_bench_cudagraph = lambda *a, **k: 0.0
+        testing._summarize_statistics = lambda *a, **k: 0.0
+        triton.testing = testing
+        sys.modules["triton"] = triton
+        sys.modules["triton.testing"] = testing
+        import _cudagraph_plugin  # noqa: E402
+        cls.cg = _cudagraph_plugin
+
+    @classmethod
+    def tearDownClass(cls):
+        for name, mod in cls._saved.items():
+            if mod is None:
+                sys.modules.pop(name, None)
+            else:
+                sys.modules[name] = mod
+        sys.modules.pop("_cudagraph_plugin", None)
+
+    def _run(self, value, **kwargs):
+        """-> (result, how many times do_bench was called)"""
+        seq, calls = [3.0, 1.0, 2.0, 9.0, 5.0], []
+        self.cg._ORIGINAL_DO_BENCH = (
+            lambda fn, **kw: (calls.append(1), seq[(len(calls) - 1) % len(seq)])[1])
+        old = os.environ.get("ZL_DIRECT_REPEAT")
+        os.environ["ZL_DIRECT_REPEAT"] = value
+        try:
+            out = self.cg.run_direct(lambda: None, 25, 100, None,
+                                     kwargs.get("quantiles"),
+                                     kwargs.get("return_mode", "mean"))
+        finally:
+            if old is None:
+                os.environ.pop("ZL_DIRECT_REPEAT", None)
+            else:
+                os.environ["ZL_DIRECT_REPEAT"] = old
+        return out, len(calls)
+
+    def test_default_is_one_call(self):
+        self.assertEqual(self._run("1"), (3.0, 1))
+
+    def test_unset_is_one_call(self):
+        old = os.environ.pop("ZL_DIRECT_REPEAT", None)
+        try:
+            self.assertEqual(self.cg.direct_repeat(), 1)
+        finally:
+            if old is not None:
+                os.environ["ZL_DIRECT_REPEAT"] = old
+
+    def test_median_over_calls(self):
+        # 3, 1, 2, 9, 5 -> 3
+        self.assertEqual(self._run("5"), (3.0, 5))
+
+    def test_garbage_and_zero_fall_back_to_one(self):
+        self.assertEqual(self._run("abc"), (3.0, 1))
+        self.assertEqual(self._run("0"), (3.0, 1))
+        self.assertEqual(self._run("-4"), (3.0, 1))
+
+    def test_sequence_results_are_not_repeated(self):
+        # quantiles / return_mode="all" make do_bench return a sequence, and a
+        # median across calls is undefined for it; take the single call.
+        self.assertEqual(self._run("5", quantiles=[0.5])[1], 1)
+        self.assertEqual(self._run("5", return_mode="all")[1], 1)
+
+
 class DeviceArgTests(unittest.TestCase):
     def test_devices_preferred(self):
         self.assertEqual(resolve_device_arg("0,1", None), "0,1")