From 88bf0f4a35bcd9aca3393c49e627c55a5080348c Mon Sep 17 00:00:00 2001 From: Jin Pan Date: Wed, 29 Jul 2026 23:01:48 +0000 Subject: [PATCH] fix(tests): exclude GPU drain from dispatch timing Move the final synchronization outside the measured window so the benchmark matches its documented host-dispatch semantics, and lock the ordering with a regression test. --- tests/unit/test_launch_overhead.py | 19 ++++++++++++++++++- 1 file changed, 18 insertions(+), 1 deletion(-) diff --git a/tests/unit/test_launch_overhead.py b/tests/unit/test_launch_overhead.py index a1a4c44cb..3f4e7deab 100644 --- a/tests/unit/test_launch_overhead.py +++ b/tests/unit/test_launch_overhead.py @@ -112,12 +112,29 @@ def bench_wallclock(fn, n_warmup=20, n_iters=1000): t0 = time.perf_counter() for _ in range(n_iters): fn() - torch.cuda.synchronize() t1 = time.perf_counter() + torch.cuda.synchronize() return (t1 - t0) / n_iters * 1e6 # µs +def test_bench_wallclock_excludes_final_gpu_drain(monkeypatch): + """The final drain must not be part of the host-dispatch window.""" + sync_count = 0 + + def fake_synchronize(): + nonlocal sync_count + sync_count += 1 + + monkeypatch.setattr(torch.cuda, "synchronize", fake_synchronize) + monkeypatch.setattr(time, "perf_counter", lambda: float(sync_count)) + + measured_us = bench_wallclock(lambda: None, n_warmup=0, n_iters=1) + + assert measured_us == 0.0 + assert sync_count == 2 + + def main(): SIZE = 1024 * 256 # 256K elements — small enough that GPU time is trivial BLOCK = 256