forked from zhoulin/zl_bench
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:
1 parent
d7cb7efc94
commit
9665d36a7c
4 files changed
+147
-12
No files matched your search
@@ -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")
|
||||
|
||||
Reference in new issue
Block a user