diff --git a/README.md b/README.md index a524eee..4da1c0a 100644 --- a/README.md +++ b/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` | diff --git a/_cudagraph_plugin.py b/_cudagraph_plugin.py index c647432..ff8ccaf 100644 --- a/_cudagraph_plugin.py +++ b/_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: diff --git a/_timing_plugin.py b/_timing_plugin.py index 0fb0401..7013bd6 100644 --- a/_timing_plugin.py +++ b/_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) diff --git a/test_backends.py b/test_backends.py index 943d50a..9af9f51 100644 --- a/test_backends.py +++ b/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")