[executorch][cuda] Shorter autotune benchmark window and one compile pool for every split candidate

backends/cuda/autotune/launch_params.py

@@ -49,6 +49,7 @@ def _launch(bucket, x, qdata, ..., *, split_k):
 from __future__ import annotations
 
 import contextlib
+import contextvars
 import functools
 import gc
 import inspect
@@ -57,7 +58,7 @@ def _launch(bucket, x, qdata, ..., *, split_k):
 import statistics
 import threading
 from dataclasses import dataclass, field
-from typing import Any, Callable, Iterable, Iterator, Sequence, TypeVar
+from typing import Any, Callable, ContextManager, Iterable, Iterator, Sequence, TypeVar
 
 import torch
 from executorch.backends.cuda.autotune.cuda_graph_timing import (
@@ -114,6 +115,9 @@ def __init__(self) -> None:
 # Held while timing, so concurrent compiles do not time against each other
 # (re-entrant: a decorated function may call another one).
 _TIMING_LOCK = threading.RLock()
+_IN_PREPASS: contextvars.ContextVar[bool] = contextvars.ContextVar(
+    "autotune_in_prepass", default=False
+)
 
 
 @contextlib.contextmanager
@@ -193,12 +197,33 @@ def _random_like(value: Any) -> Any:
     return flat.as_strided(shape, stride)
 
 
+def _run_prepass(
+    fn: Callable[..., Any],
+    param: str,
+    candidates: tuple[Any, ...],
+    args: tuple,
+    kwargs: dict,
+    candidate_prepass: Callable[[], ContextManager[Any]],
+) -> None:
+    token = _IN_PREPASS.set(True)
+    try:
+        with candidate_prepass():
+            for value in candidates:
+                try:
+                    fn(*args, **kwargs, **{param: value})
+                except InvalidLaunchParam:
+                    pass
+    finally:
+        _IN_PREPASS.reset(token)
+
+
 def _time_candidates(
     fn: Callable[..., Any],
     param: str,
     candidates: tuple[Any, ...],
     args: tuple,
     kwargs: dict,
+    candidate_prepass: Callable[[], ContextManager[Any]] | None,
 ) -> tuple[dict[Any, float], dict[Any, str]]:
     # Leave export's fake / functional / proxy modes, and launch the Triton
     # kernels directly, as triton_op's eager path does: through wrap_triton's
@@ -216,6 +241,10 @@ def _time_candidates(
         with torch.cuda.device(device), torch.cuda.stream(stream):
             real_args = tuple(_random_like(a) for a in args)
             real_kwargs = {k: _random_like(v) for k, v in kwargs.items()}
+            if candidate_prepass is not None:
+                _run_prepass(
+                    fn, param, candidates, real_args, real_kwargs, candidate_prepass
+                )
             for value in candidates:
 
                 def run(value: Any = value) -> Any:
@@ -249,23 +278,26 @@ def _time_candidates_within_memory(
     args: tuple,
     kwargs: dict,
     stats: LaunchParamStats,
+    candidate_prepass: Callable[[], ContextManager[Any]] | None,
 ) -> tuple[dict[Any, float] | None, dict[Any, str]]:
     """``_time_candidates``, retried after freeing the device's cache and then
     with the offload tensors parked on the CPU. Returns no times when the device
     still runs out of memory."""
     try:
-        return _time_candidates(fn, param, candidates, args, kwargs)
+        return _time_candidates(fn, param, candidates, args, kwargs, candidate_prepass)
     except torch.OutOfMemoryError:
         _free_device_memory()
     try:
-        return _time_candidates(fn, param, candidates, args, kwargs)
+        return _time_candidates(fn, param, candidates, args, kwargs, candidate_prepass)
     except torch.OutOfMemoryError:
         _free_device_memory()
     if _STATE.offload is not None:
         try:
             with _parked_on_cpu(_STATE.offload()):
                 stats.offloads += 1
-                return _time_candidates(fn, param, candidates, args, kwargs)
+                return _time_candidates(
+                    fn, param, candidates, args, kwargs, candidate_prepass
+                )
         except torch.OutOfMemoryError:
             _free_device_memory()
     logger.warning(
@@ -278,9 +310,18 @@ def _time_candidates_within_memory(
     return None, {}
 
 
-def autotune_launch_param(param: str, candidates: Sequence[Any]) -> Callable[[F], F]:
+def autotune_launch_param(
+    param: str,
+    candidates: Sequence[Any],
+    *,
+    candidate_prepass: Callable[[], ContextManager[Any]] | None = None,
+) -> Callable[[F], F]:
     """Fills keyword argument ``param`` of the decorated launch function with
-    the fastest of ``candidates`` during AOTI compile (see the module doc)."""
+    the fastest of ``candidates`` during AOTI compile (see the module doc).
+
+    If provided, ``candidate_prepass`` scopes one untimed call of every
+    candidate before normal launch and timing.
+    """
     candidates = tuple(candidates)
     if not candidates:
         raise ValueError("autotune_launch_param needs at least one candidate")
@@ -294,6 +335,11 @@ def decorate(fn: F) -> F:
 
         @functools.wraps(fn)
         def wrapper(*args: Any, **kwargs: Any) -> Any:
+            if _IN_PREPASS.get():
+                raise RuntimeError(
+                    f"{name} runs inside another autotuned launch's candidate prepass; "
+                    "nested autotuning is not supported"
+                )
             stats = _STATE.stats
             if param in kwargs or stats is None or len(candidates) == 1:
                 kwargs.setdefault(param, candidates[0])
@@ -314,7 +360,14 @@ def wrapper(*args: Any, **kwargs: Any) -> Any:
                     times = None
                 else:
                     times, rejected = _time_candidates_within_memory(
-                        fn, name, param, candidates, args, kwargs, stats
+                        fn,
+                        name,
+                        param,
+                        candidates,
+                        args,
+                        kwargs,
+                        stats,
+                        candidate_prepass,
                     )
                     if times is None:
                         stats.out_of_memory += 1

backends/cuda/tests/test_autotune_launch_params.py

@@ -6,6 +6,7 @@
 
 """Tests for @autotune_launch_param (autotune/launch_params.py)."""
 
+import contextlib
 import unittest
 from unittest import mock
 
@@ -160,6 +161,57 @@ def launch(x, *, split_k):
             second(self.x)
         self.assertEqual([p.choice for p in stats.picks], [1, 2])
 
+    def test_a_prepass_calls_every_legal_candidate_once_inside_its_scope(self) -> None:
+        scoped: list[list[int]] = []
+
+        class _Scope:
+            def __enter__(self):
+                self.start = len(_CALLS)
+                return self
+
+            def __exit__(self, *exc):
+                scoped.append(_CALLS[self.start :])
+                return False
+
+        launch = autotune_launch_param("split_k", (1, 2, 4), candidate_prepass=_Scope)(
+            _toy_launch.__wrapped__
+        )
+        with autotune_launch_params() as stats:
+            launch(self.x, 2)
+        # split_k=4 is illegal for max_split=2, so the scope runs 1 and 2 only.
+        self.assertEqual(scoped, [[1, 2]])
+        self.assertEqual(stats.picks[0].choice, 2)
+
+    def test_an_autotuned_launch_inside_a_prepass_is_refused(self) -> None:
+        @autotune_launch_param("split_k", (1, 2))
+        def inner(x, *, split_k):
+            return x @ x
+
+        @autotune_launch_param(
+            "split_k", (1, 2), candidate_prepass=contextlib.nullcontext
+        )
+        def outer(x, *, split_k):
+            return inner(x)
+
+        with autotune_launch_params():
+            with self.assertRaisesRegex(RuntimeError, "nested autotuning"):
+                outer(self.x)
+
+    def test_an_explicit_nested_launch_inside_a_prepass_is_refused(self) -> None:
+        @autotune_launch_param("split_k", (1, 2))
+        def inner(x, *, split_k):
+            return x @ x
+
+        @autotune_launch_param(
+            "split_k", (1, 2), candidate_prepass=contextlib.nullcontext
+        )
+        def outer(x, *, split_k):
+            return inner(x, split_k=2)
+
+        with autotune_launch_params():
+            with self.assertRaisesRegex(RuntimeError, "nested autotuning"):
+                outer(self.x)
+
     def test_parameter_must_exist(self) -> None:
         with self.assertRaisesRegex(TypeError, "has no parameter 'num_splits'"):
 

backends/cuda/tests/test_int4_quantized_gemm.py

@@ -16,6 +16,7 @@
     python -m pytest backends/cuda/tests/test_int4_quantized_gemm.py -v
 """
 
+import contextlib
 import statistics
 import types
 import unittest
@@ -148,14 +149,46 @@ def test_an_untuned_shape_compiles_every_config_in_the_pool_first(self) -> None:
             pooled, compiled = counts["pooled"], self._compiled(bucket)
             again = int4_kernel._launch(bucket, x, *weights, GROUP_SIZE, split_k=1)
         self.assertGreater(pooled, 1)
-        # Every kernel the benchmark ran was compiled in the pool.
-        self.assertEqual(compiled - before, pooled)
+        # Every kernel the benchmark ran was compiled in the pool (which may
+        # merge identical requests): a compile outside it adds a kernel the pool
+        # was never asked for.
+        self.assertGreater(compiled, before)
+        self.assertLessEqual(compiled - before, pooled)
         # A tuned key is not compiled again.
         self.assertEqual(counts["pooled"], pooled)
         self.assertEqual(self._compiled(bucket), compiled)
         _check_close(self, out, _dequant_matmul(x, *weights, GROUP_SIZE))
         self.assertTrue(torch.equal(out, again))
 
+    def test_every_split_candidate_compiles_on_one_pool(self) -> None:
+        # A shape no other test tunes; the split is timed, not given.
+        n, k, bucket = 320, 2560, 2
+        weights = _packed(n, k, seed=11)
+        x = torch.randn(bucket, k, dtype=torch.bfloat16, device="cuda")
+        compile_pool = quantized_gemm_family._compile_pool
+        compiled_at_pool_exit: list[int] = []
+
+        @contextlib.contextmanager
+        def recording_pool():
+            with compile_pool():
+                yield
+            compiled_at_pool_exit.append(self._compiled(bucket))
+
+        counts, patch = self._count_pooled()
+        before = self._compiled(bucket)
+        with patch, mock.patch.object(
+            quantized_gemm_family, "_compile_pool", recording_pool
+        ), autotune_launch_params() as stats:
+            out = int4_kernel._launch(bucket, x, *weights, GROUP_SIZE)
+        self.assertEqual(stats.measured, 1)
+        self.assertGreater(len(stats.picks[0].times_us), 1)
+        # The first pool, which holds every split candidate, compiles every
+        # kernel the timed launches then benchmark.
+        self.assertGreater(compiled_at_pool_exit[0], before)
+        self.assertEqual(self._compiled(bucket), compiled_at_pool_exit[0])
+        self.assertLessEqual(self._compiled(bucket) - before, counts["pooled"])
+        _check_close(self, out, _dequant_matmul(x, *weights, GROUP_SIZE))
+
     def test_fake_tensors_are_not_precompiled(self) -> None:
         # A full launch of a shape no other test tunes, so only the fake-tensor
         # check can keep it off the pool: the same launch on real tensors uses it.
@@ -252,6 +285,14 @@ def test_cpython_with_the_ast_fix_compiles_in_parallel(self) -> None:
             with self.subTest(version=version):
                 self.assertTrue(self._available(version=version))
 
+    def test_a_gated_off_prepass_neither_pools_nor_queues(self) -> None:
+        with mock.patch.object(
+            quantized_gemm_family, "_parallel_compile_available", return_value=False
+        ), mock.patch.object(quantized_gemm_family, "_compile_pool") as pool:
+            with quantized_gemm_family._compile_every_candidate_at_once():
+                self.assertFalse(quantized_gemm_family._QUEUING_ONLY.get())
+        pool.assert_not_called()
+
     def test_the_gate_decides_whether_an_eligible_launch_precompiles(self) -> None:
         for gate, precompiled in ((False, False), (True, True)):
             with self.subTest(gate=gate):

backends/cuda/tests/test_quantized_gemm_family.py

@@ -12,6 +12,7 @@
 import unittest
 
 import torch
+import triton
 from executorch.backends.cuda.triton.kernels.quantized_gemm_family import (
     QuantizedGemmFamily,
 )
@@ -117,6 +118,35 @@ def test_prototype_with_postponed_annotations_from_its_own_module(self) -> None:
         x, w = torch.randn(1, 4), torch.randn(3, 4)
         torch.testing.assert_close(family.op(1)(x, w), x @ w.t())
 
+    def test_every_autotuned_kernel_uses_the_shared_benchmark_window(self) -> None:
+        from executorch.backends.cuda.triton.kernels import (
+            int4_quantized_gemm,
+            int5_quantized_gemm,
+            int6_quantized_gemm,
+            int8_quantized_gemm,
+        )
+        from executorch.backends.cuda.triton.kernels.quantized_gemm_utils import (
+            autotune_do_bench,
+        )
+
+        modules = (
+            int4_quantized_gemm,
+            int5_quantized_gemm,
+            int6_quantized_gemm,
+            int8_quantized_gemm,
+        )
+        autotuners = [
+            value
+            for module in modules
+            for table in vars(module).values()
+            if isinstance(table, dict)
+            for value in table.values()
+            if isinstance(value, triton.runtime.Autotuner)
+        ]
+        self.assertGreater(len(autotuners), len(modules))
+        for autotuner in autotuners:
+            self.assertIs(autotuner.do_bench, autotune_do_bench)
+
     def test_rejects_invalid_buckets(self) -> None:
         for buckets in ((), (0, 1), (1, 1)):
             with self.assertRaises(ValueError):

backends/cuda/triton/kernels/int4_quantized_gemm.py

@@ -40,13 +40,15 @@
 import triton.language as tl
 from executorch.backends.cuda.autotune.launch_params import autotune_launch_param
 from executorch.backends.cuda.triton.kernels.quantized_gemm_family import (
+    _compile_every_candidate_at_once,
     launch_split_k_gemm,
     QuantizedGemmFamily,
 )
 from executorch.backends.cuda.triton.kernels.quantized_gemm_utils import (
     _dp4a_u8_s8,
     _warp_sum_f32,
     autotune_configs,
+    autotune_do_bench,
     check_contiguous,
     check_device,
     check_dtypes,
@@ -1664,6 +1666,7 @@ def _prune_tiles(configs, named_args, **kwargs):
         configs=int4_autotune_configs(bucket),
         key=["N", "K", "SPLIT_K", "DYNAMIC_M"],
         prune_configs_by={"early_config_prune": _prune},
+        do_bench=autotune_do_bench,
     )(_int4_w4a8_bucket_kernel)
     for bucket in _SMALL_BUCKETS
 }
@@ -1672,6 +1675,7 @@ def _prune_tiles(configs, named_args, **kwargs):
         configs=int4_autotune_configs(bucket),
         key=["M", "N", "K", "SPLIT_K"],
         prune_configs_by={"early_config_prune": _prune_tiles},
+        do_bench=autotune_do_bench,
     )(_int4_w4a16_direct_kernel)
     for bucket in _LARGE_BUCKETS
 }
@@ -1680,6 +1684,7 @@ def _prune_tiles(configs, named_args, **kwargs):
         configs=int4_autotune_configs(bucket),
         key=["M", "N", "K", "SPLIT_K"],
         prune_configs_by={"early_config_prune": _prune_tiles},
+        do_bench=autotune_do_bench,
     )(_int4_w4a16_splitk_kernel)
     for bucket in _LARGE_BUCKETS
 }
@@ -1768,7 +1773,11 @@ def _launch_w4a8(
     )
 
 
-@autotune_launch_param("split_k", SPLIT_K_CANDIDATES)
+@autotune_launch_param(
+    "split_k",
+    SPLIT_K_CANDIDATES,
+    candidate_prepass=_compile_every_candidate_at_once,
+)
 def _launch_rows(
     bucket: int,
     x: torch.Tensor,
@@ -1827,7 +1836,11 @@ def _launch_rows(
     )
 
 
-@autotune_launch_param("split_k", TILE_SPLIT_K_CANDIDATES)
+@autotune_launch_param(
+    "split_k",
+    TILE_SPLIT_K_CANDIDATES,
+    candidate_prepass=_compile_every_candidate_at_once,
+)
 def _launch_w4a16(
     bucket: int,
     x: torch.Tensor,