From 58376aabe7008df1d2487db1b44a1999c842e6c4 Mon Sep 17 00:00:00 2001 From: gasoonjia Date: Sat, 17 Oct 2026 16:56:19 -0700 Subject: [PATCH] [executorch][cuda] Shorter autotune benchmark window and one compile pool for every split candidate Pull Request resolved: https://github.com/pytorch/executorch/pull/23737 Builds on D124355785 (configs compiled in parallel before Triton benchmarks them). Two remaining costs of the compile-time quantized GEMM autotune: 1. **Benchmark window.** The family kernels' `triton.autotune` used Triton's default `do_bench(warmup=25, rep=100)`: about 125 ms per config whatever the kernel's length. These kernels run 10-100 us, so every config ran thousands of times; a bucket-64 shape with 288 (config, split) pairs spent ~36 s in it. They now use `quantized_gemm_utils.autotune_do_bench` (5 ms warmup, 20 ms timing), which still runs each config a few hundred times. Only the INT4/5/6/8 family kernels use it; SDPA, MoE and the rest keep Triton's default. 2. **One compile pool per shape.** `@autotune_launch_param` tries each split-K candidate in turn, and D124355785 opened one compile pool per candidate, one after another (e.g. 72, 72, 72, 48 and 24 configs at bucket 64). `@autotune_launch_param` now takes an optional `candidate_prepass`: a scope around one untimed call of every legal candidate before any is timed. The quantized GEMM launches pass `quantized_gemm_family._compile_every_candidate_at_once`, in which `launch_split_k_gemm` only queues its kernel's untuned configs on one shared pool and skips the launch. Leaving the scope waits for all compiles; timing then runs exactly as before. Each key compiled ahead is consumed once by its timed launch, which then does not open a pool of its own. Neither changes the search: the same configs and split candidates are compiled and timed, and the picks are the fastest by the same rules. **Does the shorter window still pick the best config?** For the models' linear shapes (INT4: 8448x6656, 6656x4096, 6656x19968, 256x6656, 1024x6656; INT6: 4096x5376, 5376x21504, 6656x19968, 1024x6656; INT5 and INT8: two each) at buckets 1/4/16/64 (INT5, INT8: 1/4) and every legal split, every pruned config was timed with the current window and with the new one; the union of each window's five fastest was then re-timed with a long window (100 ms warmup, 300 ms timing) as the reference. - 232 (shape, bucket, split) cases: the new window's pick is within 1% of the reference best in 199 cases (current window: 202), within 2% in 214 (213); mean regret 1.0054 (1.0050). - What ships is one split per shape: over the 44 (shape, bucket) pairs, with the split chosen as the fastest, the new window is more than 2% off the best in 2 pairs and the current window in 3. Worst: new 8.8% on INT4 256 x 6656 at M = 1 (an 8 us kernel, 0.7 us), current 6.0% on INT4 1024 x 6656 at M = 16. Both windows miss by a few percent on the shortest kernels, at the same rate. **Compile time.** One INT4 W4A16 bucket-64 shape (5376 x 4096) under `autotune_launch_params()` + `cuda_graph_timing()`, A100, empty Triton cache: 334.6 s before D124355785, 77.8 s with it, 58.0 s with this diff (26.7 s with a warm cache, was 44.8 s); split pick 4 in all, the same 292 kernels compiled, per-split times within 1%. Full Muse Glimmer 30B batching export (forward_{1,2,4,8,16,32,64} + forward_others + samplers; A100; empty Triton and Inductor caches), per method in minutes from the AOTI wrapper times: | method | before D124355785 | D124355785 | this diff | |---|---|---|---| | load + trace + forward_1 | 46.5 | 51.3 | 40.8 | | forward_2 | 21.4 | 15.3 | 14.5 | | forward_4 | 23.4 | 15.5 | 14.4 | | forward_8 | 22.3 | 12.2 | 10.3 | | forward_16 | 39.3 | 15.5 | 12.9 | | forward_32 | 52.8 | 18.8 | 14.7 | | forward_64 | 73.3 | 24.2 | 18.7 | | forward_others | 10.5 | 10.3 | 11.1 | | total | 4h55m | 2h50m | 2h25m | The export generates the same tokens as before (greedy; 512 and 4 x 2k prompts) at the same speed as with D124355785 (single_512 63.0 / 62.7 tok/s vs 62.5 / 62.9; 4 x 2k 111.3 vs 111.0). What is left per autotuned shape is mostly compiling: during forward_64 only ~5 of the 32 compile threads are busy on average, and 87% of their time is in Triton's MLIR/LLVM passes (make_llir, make_ttgir, make_ptx), which appear to serialize on the GIL. ghstack-source-id: 451021038 @exported-using-ghexport Differential Revision: [D124496888](https://our.internmc.facebook.com/intern/diff/D124496888/) --- backends/cuda/autotune/launch_params.py | 67 ++++++++++-- .../cuda/tests/test_autotune_launch_params.py | 52 +++++++++ .../cuda/tests/test_int4_quantized_gemm.py | 45 +++++++- .../cuda/tests/test_quantized_gemm_family.py | 30 ++++++ .../triton/kernels/int4_quantized_gemm.py | 17 ++- .../triton/kernels/int5_quantized_gemm.py | 11 +- .../triton/kernels/int6_quantized_gemm.py | 16 ++- .../triton/kernels/int8_quantized_gemm.py | 9 +- .../triton/kernels/quantized_gemm_family.py | 100 ++++++++++++++---- .../triton/kernels/quantized_gemm_utils.py | 20 ++++ 10 files changed, 330 insertions(+), 37 deletions(-) diff --git a/backends/cuda/autotune/launch_params.py b/backends/cuda/autotune/launch_params.py index ba54a0c8f09..02554c922b3 100644 --- a/backends/cuda/autotune/launch_params.py +++ b/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 diff --git a/backends/cuda/tests/test_autotune_launch_params.py b/backends/cuda/tests/test_autotune_launch_params.py index 777ed2f032f..65ba15b5121 100644 --- a/backends/cuda/tests/test_autotune_launch_params.py +++ b/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'"): diff --git a/backends/cuda/tests/test_int4_quantized_gemm.py b/backends/cuda/tests/test_int4_quantized_gemm.py index f7ffc18d520..03d40cfe345 100644 --- a/backends/cuda/tests/test_int4_quantized_gemm.py +++ b/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): diff --git a/backends/cuda/tests/test_quantized_gemm_family.py b/backends/cuda/tests/test_quantized_gemm_family.py index 109dc35215b..2186f26e62a 100644 --- a/backends/cuda/tests/test_quantized_gemm_family.py +++ b/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): diff --git a/backends/cuda/triton/kernels/int4_quantized_gemm.py b/backends/cuda/triton/kernels/int4_quantized_gemm.py index ef217e3c1bd..1d80c68b66a 100644 --- a/backends/cuda/triton/kernels/int4_quantized_gemm.py +++ b/backends/cuda/triton/kernels/int4_quantized_gemm.py @@ -40,6 +40,7 @@ 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, ) @@ -47,6 +48,7 @@ _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, diff --git a/backends/cuda/triton/kernels/int5_quantized_gemm.py b/backends/cuda/triton/kernels/int5_quantized_gemm.py index 9f540d9b381..7d0524267f4 100644 --- a/backends/cuda/triton/kernels/int5_quantized_gemm.py +++ b/backends/cuda/triton/kernels/int5_quantized_gemm.py @@ -38,6 +38,7 @@ 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, ) @@ -45,6 +46,7 @@ _dp4a_u8_s8, _warp_sum_f32, autotune_configs, + autotune_do_bench, check_contiguous, check_device, check_dtypes, @@ -1343,6 +1345,7 @@ def _prune(configs, named_args, **kwargs): configs=int5_autotune_configs(bucket), key=["N", "K", "SPLIT_K", "DYNAMIC_M"], prune_configs_by={"early_config_prune": _prune}, + do_bench=autotune_do_bench, )(_int5_w5a8_bucket_kernel) for bucket in _SMALL_BUCKETS } @@ -1361,6 +1364,7 @@ def _prune_tiles(configs, named_args, **kwargs): configs=_int5_large_autotune_configs(b), key=["M", "N", "K", "SPLIT_K"], prune_configs_by={"early_config_prune": _prune_tiles}, + do_bench=autotune_do_bench, )(_int5_w5a16_direct_kernel) for b in _LARGE_BUCKETS } @@ -1369,6 +1373,7 @@ def _prune_tiles(configs, named_args, **kwargs): configs=_int5_large_autotune_configs(b), key=["M", "N", "K", "SPLIT_K"], prune_configs_by={"early_config_prune": _prune_tiles}, + do_bench=autotune_do_bench, )(_int5_w5a16_splitk_kernel) for b in _LARGE_BUCKETS } @@ -1568,7 +1573,11 @@ def _launch( ) -@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, diff --git a/backends/cuda/triton/kernels/int6_quantized_gemm.py b/backends/cuda/triton/kernels/int6_quantized_gemm.py index cfdfef16ba2..5ddd6fe8ed3 100644 --- a/backends/cuda/triton/kernels/int6_quantized_gemm.py +++ b/backends/cuda/triton/kernels/int6_quantized_gemm.py @@ -30,6 +30,7 @@ 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, ) @@ -37,6 +38,7 @@ _dp4a_u8_s8, _warp_sum_f32, autotune_configs, + autotune_do_bench, check_contiguous, check_device, check_dtypes, @@ -1247,6 +1249,7 @@ def _prune_tiles(configs, named_args, **kwargs): configs=int6_autotune_configs(), key=["N", "K", "SPLIT_K", "DYNAMIC_M"], prune_configs_by={"early_config_prune": _prune}, + do_bench=autotune_do_bench, )(_int6_w6a8_bucket_kernel) for bucket in _SMALL_BUCKETS } @@ -1256,6 +1259,7 @@ def _prune_tiles(configs, named_args, **kwargs): configs=tile_autotune_configs(bucket), key=["M", "N", "K", "SPLIT_K"], prune_configs_by={"early_config_prune": _prune_tiles}, + do_bench=autotune_do_bench, )(_int6_w6a16_tile_kernel) for bucket in _LARGE_BUCKETS } @@ -1337,7 +1341,11 @@ def _launch_w6a8( ) -@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, @@ -1384,7 +1392,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_w6a16( bucket: int, x: torch.Tensor, diff --git a/backends/cuda/triton/kernels/int8_quantized_gemm.py b/backends/cuda/triton/kernels/int8_quantized_gemm.py index 40d420aa907..74de5ca374c 100644 --- a/backends/cuda/triton/kernels/int8_quantized_gemm.py +++ b/backends/cuda/triton/kernels/int8_quantized_gemm.py @@ -27,6 +27,7 @@ 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, ) @@ -34,6 +35,7 @@ _dp4a_u8_s8, _warp_sum_f32, autotune_configs, + autotune_do_bench, check_contiguous, check_device, check_dtypes, @@ -1098,6 +1100,7 @@ def _prune(configs, named_args, **kwargs): configs=int8_autotune_configs(bucket), key=["N", "K", "SPLIT_K", "DYNAMIC_M"], prune_configs_by={"early_config_prune": _prune}, + do_bench=autotune_do_bench, )(_int8_w8a8_bucket_kernel) for bucket in SUPPORTED_BUCKETS } @@ -1195,7 +1198,11 @@ def _launch( ) -@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, diff --git a/backends/cuda/triton/kernels/quantized_gemm_family.py b/backends/cuda/triton/kernels/quantized_gemm_family.py index 1337c22b016..e647dec808b 100644 --- a/backends/cuda/triton/kernels/quantized_gemm_family.py +++ b/backends/cuda/triton/kernels/quantized_gemm_family.py @@ -25,6 +25,8 @@ from __future__ import annotations import concurrent.futures +import contextlib +import contextvars import functools import inspect import logging @@ -157,6 +159,22 @@ def _tuning_key(kernel, args: Sequence, kwargs: dict) -> tuple: return tuple(key) +def _has_untuned_configs(kernel, args: Sequence, kwargs: dict) -> bool: + """Whether ``kernel`` is about to benchmark configs it could compile ahead: + an autotuner with several configs, a key it has not tuned, and tensors with + real CUDA storage (tracing passes fake or functional tensors, which the + launch itself handles).""" + if not isinstance(kernel, triton.runtime.Autotuner) or len(kernel.configs) <= 1: + return False + for arg in (*args, *kwargs.values()): + if isinstance(arg, torch.Tensor) and ( + type(arg) not in (torch.Tensor, torch.nn.Parameter) + or arg.device.type != "cuda" + ): + return False + return _tuning_key(kernel, args, kwargs) not in kernel.cache + + _COMPILE_THREADS_LOCK = threading.Lock() _compile_thread_pool: Optional[concurrent.futures.ThreadPoolExecutor] = None @@ -182,6 +200,25 @@ def _parallel_compile_available() -> bool: return (3, 11, 8) <= version < (3, 12) or version >= (3, 12, 2) +@contextlib.contextmanager +def _compile_pool(): + with _async_compile.AsyncCompileMode(_compile_threads(), ignore_errors=True): + yield + + +def _queue_untuned_configs(kernel, args: Sequence, kwargs: dict) -> None: + """Warms up an untuned key so its configs compile on the active compile pool.""" + if not _has_untuned_configs(kernel, args, kwargs): + return + try: + kernel.warmup(*args, grid=(1,), **kwargs) + except Exception as error: # noqa: BLE001 - precompiling only saves time + # The timed launch compiles this candidate's configs itself. + logger.warning( + "Not precompiling %s's configs: %s", kernel.base_fn.__name__, error + ) + + def _precompile_untuned_configs(kernel, args: Sequence, kwargs: dict) -> None: """Compiles every config an autotuned kernel is about to benchmark, in parallel. @@ -193,30 +230,39 @@ def _precompile_untuned_configs(kernel, args: Sequence, kwargs: dict) -> None: """ if ( not _parallel_compile_available() - or not isinstance(kernel, triton.runtime.Autotuner) - or len(kernel.configs) <= 1 or _async_compile.active_mode.get() is not None ): return - # Only real CUDA storage can be compiled against: tracing passes fake or - # functional tensors (subclasses), which the launch itself handles. - for arg in (*args, *kwargs.values()): - if isinstance(arg, torch.Tensor) and ( - type(arg) not in (torch.Tensor, torch.nn.Parameter) - or arg.device.type != "cuda" - ): - return - if _tuning_key(kernel, args, kwargs) in kernel.cache: + if not _has_untuned_configs(kernel, args, kwargs): return - try: - with _async_compile.AsyncCompileMode(_compile_threads(), ignore_errors=True): - kernel.warmup(*args, grid=(1,), **kwargs) - except Exception as error: # noqa: BLE001 - precompiling only saves time - # The launch compiles and benchmarks the configs itself, as without - # precompiling, and reports a real problem there. - logger.warning( - "Not precompiling %s's configs: %s", kernel.base_fn.__name__, error - ) + with _compile_pool(): + _queue_untuned_configs(kernel, args, kwargs) + + +# Set while every @autotune_launch_param candidate makes one untimed call: the +# launches then only queue their configs on one shared compile pool. +_QUEUING_ONLY: contextvars.ContextVar[bool] = contextvars.ContextVar( + "quantized_gemm_queuing_only", default=False +) + + +@contextlib.contextmanager +def _compile_every_candidate_at_once(): + """Compiles what all launch-parameter candidates (e.g. every split-K) will + benchmark on one pool, rather than one candidate's configs at a time; leaving + waits for all of them. Timing then runs as before on compiled kernels.""" + if ( + not _parallel_compile_available() + or _async_compile.active_mode.get() is not None + ): + yield + return + with _compile_pool(): + token = _QUEUING_ONLY.set(True) + try: + yield + finally: + _QUEUING_ONLY.reset(token) def launch_split_k_gemm( @@ -277,8 +323,18 @@ def grid(meta): partial.stride(-2), partial.stride(-1), ) - _precompile_untuned_configs(kernel, args, {"SPLIT_K": split_k, **constexprs}) - wrap_triton(kernel)[grid](*args, SPLIT_K=split_k, **constexprs) + kwargs = {"SPLIT_K": split_k, **constexprs} + if _QUEUING_ONLY.get(): + if ( + not _parallel_compile_available() + or _async_compile.active_mode.get() is None + ): + raise RuntimeError("candidate prepass lost its active compile pool") + _queue_untuned_configs(kernel, args, kwargs) + # The prepass discards the result, so the output is not written. + return out + _precompile_untuned_configs(kernel, args, kwargs) + wrap_triton(kernel)[grid](*args, **kwargs) if split_k > 1: splitk_reduce(partial, out, block_m, reduce_block_n) return out diff --git a/backends/cuda/triton/kernels/quantized_gemm_utils.py b/backends/cuda/triton/kernels/quantized_gemm_utils.py index 28a9e3f802f..535ce67388c 100644 --- a/backends/cuda/triton/kernels/quantized_gemm_utils.py +++ b/backends/cuda/triton/kernels/quantized_gemm_utils.py @@ -25,6 +25,7 @@ import torch import triton import triton.language as tl +import triton.testing from executorch.backends.cuda.autotune.launch_params import InvalidLaunchParam from torch.fx.experimental.symbolic_shapes import statically_known_true from torch.library import wrap_triton @@ -343,6 +344,25 @@ def check_split_k(split_k: int, k: int, k_per_split_unit: int = Q8_TILE) -> None ) +# The window triton.autotune times each config with. Triton's default (25 ms of +# warmup and 100 ms of timing, whatever the kernel's length) runs these +# 20-100 us kernels thousands of times per config; this still runs each a few +# hundred times, and its picks match the default's within 2% on the models' +# shapes. +AUTOTUNE_WARMUP_MS = 5 +AUTOTUNE_REP_MS = 20 + + +def autotune_do_bench(kernel_call, quantiles): + """``do_bench`` for the quantized GEMM kernels' ``triton.autotune``.""" + return triton.testing.do_bench( + kernel_call, + warmup=AUTOTUNE_WARMUP_MS, + rep=AUTOTUNE_REP_MS, + quantiles=quantiles, + ) + + def autotune_configs( implementations: Iterable[dict] = ({},), ) -> list[triton.Config]: