diff --git a/backends/cuda/tests/test_int4_quantized_gemm.py b/backends/cuda/tests/test_int4_quantized_gemm.py index b80c041dfe1..f7ffc18d520 100644 --- a/backends/cuda/tests/test_int4_quantized_gemm.py +++ b/backends/cuda/tests/test_int4_quantized_gemm.py @@ -17,6 +17,7 @@ """ import statistics +import types import unittest from unittest import mock @@ -36,7 +37,10 @@ SaturatedCpu, summarize, ) -from executorch.backends.cuda.triton.kernels import int4_quantized_gemm as int4_kernel +from executorch.backends.cuda.triton.kernels import ( + int4_quantized_gemm as int4_kernel, + quantized_gemm_family, +) from executorch.backends.cuda.triton.kernels.int4_quantized_gemm import ( int4_autotune_configs, INT4_QUANTIZED_GEMM, @@ -49,6 +53,8 @@ from executorch.extension.llm.export.int4 import ExportableInt4Tensor from executorch.extension.llm.export.quant.quantize import quantize_weight from executorch.extension.llm.export.quant.recipe import QuantConfig +from torch._library.triton import set_wrap_triton_enabled +from torch._subclasses.fake_tensor import FakeTensorMode GROUP_SIZE = 32 MODEL_SHAPES = ( @@ -98,6 +104,173 @@ def _check_close(test, out, ref): test.assertLess(mean_rel.item(), 0.02) +@unittest.skipIf( + not quantized_gemm_family._parallel_compile_available(), + "this Triton or Python cannot compile in parallel", +) +class Int4ParallelPrecompileTest(unittest.TestCase): + """launch_split_k_gemm compiles an untuned shape's configs in parallel.""" + + @classmethod + def setUpClass(cls) -> None: + if not torch.cuda.is_available(): + raise unittest.SkipTest("CUDA required") + + def _count_pooled(self): + counts = {"pooled": 0} + submit = quantized_gemm_family._async_compile.AsyncCompileMode.submit + + def counting_submit(mode, *args, **kwargs): + counts["pooled"] += 1 + return submit(mode, *args, **kwargs) + + return counts, mock.patch.object( + quantized_gemm_family._async_compile.AsyncCompileMode, + "submit", + counting_submit, + ) + + @staticmethod + def _compiled(bucket: int) -> int: + fn = int4_kernel._BUCKET_KERNELS[bucket].fn + return len(fn.device_caches[torch.cuda.current_device()][0]) + + def test_an_untuned_shape_compiles_every_config_in_the_pool_first(self) -> None: + # A shape no other test tunes, so its key is not in the autotune cache. + n, k, bucket = 384, 2304, 2 + weights = _packed(n, k, seed=7) + x = torch.randn(bucket, k, dtype=torch.bfloat16, device="cuda") + counts, patch = self._count_pooled() + before = self._compiled(bucket) + # Launch the bucket's autotuner directly, so its cache gets the tuning key. + with patch, set_wrap_triton_enabled(False): + out = int4_kernel._launch(bucket, x, *weights, GROUP_SIZE, split_k=1) + 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) + # 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_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. + n, k, bucket = 384, 1792, 1 + weights = _packed(n, k, seed=19) + x = torch.randn(bucket, k, dtype=torch.bfloat16, device="cuda") + mode = FakeTensorMode() + fake = [mode.from_tensor(t) for t in (x, *weights)] + counts, patch = self._count_pooled() + with patch: + with mode: + int4_kernel._launch(bucket, *fake, GROUP_SIZE) + self.assertEqual(counts["pooled"], 0) + int4_kernel._launch(bucket, x, *weights, GROUP_SIZE) + self.assertGreater(counts["pooled"], 0) + + def test_a_failing_precompile_does_not_block_the_launch(self) -> None: + # A shape no other test tunes, so the launch tries to precompile it. + n, k, bucket = 320, 1792, 1 + weights = _packed(n, k, seed=29) + x = torch.randn(bucket, k, dtype=torch.bfloat16, device="cuda") + autotuner = int4_kernel._BUCKET_KERNELS[bucket] + with mock.patch.object( + autotuner, "warmup", side_effect=RuntimeError("no precompile") + ), self.assertLogs(quantized_gemm_family.logger, "WARNING") as logs: + out = int4_kernel._launch(bucket, x, *weights, GROUP_SIZE) + self.assertIn("no precompile", "\n".join(logs.output)) + _check_close(self, out, _dequant_matmul(x, *weights, GROUP_SIZE)) + + def test_the_tuning_key_is_the_autotuners_cache_key(self) -> None: + # _tuning_key mirrors Triton's Autotuner.run; check it against the key + # Triton stored its pick under. + n, k, bucket = 384, 1280, 1 + weights = _packed(n, k, seed=23) + x = torch.randn(bucket, k, dtype=torch.bfloat16, device="cuda") + autotuner = int4_kernel._BUCKET_KERNELS[bucket] + launches = [] + wrap = quantized_gemm_family.wrap_triton + + def recording(kernel): + launcher = wrap(kernel) + + class _Grid: + def __getitem__(self, grid): + def call(*args, **kwargs): + launches.append((kernel, args, kwargs)) + return launcher[grid](*args, **kwargs) + + return call + + return _Grid() + + # Through the op, whose eager path runs the autotuner itself (wrap_triton's + # eager path tunes a fresh copy on every call). + with mock.patch.object(quantized_gemm_family, "wrap_triton", recording): + INT4_QUANTIZED_GEMM.op(bucket)(x, *weights, GROUP_SIZE) + _, args, kwargs = next(launch for launch in launches if launch[0] is autotuner) + self.assertIn( + quantized_gemm_family._tuning_key(autotuner, args, kwargs), autotuner.cache + ) + + +class ParallelCompileGateTest(unittest.TestCase): + def _available(self, *, version: tuple) -> bool: + fake_sys = types.SimpleNamespace(version_info=version) + quantized_gemm_family._parallel_compile_available.cache_clear() + self.addCleanup(quantized_gemm_family._parallel_compile_available.cache_clear) + with mock.patch.object( + quantized_gemm_family, "_async_compile", types.SimpleNamespace() + ), mock.patch.object(quantized_gemm_family, "sys", fake_sys): + return quantized_gemm_family._parallel_compile_available() + + def test_a_triton_without_async_compile_compiles_serially(self) -> None: + self.addCleanup(quantized_gemm_family._parallel_compile_available.cache_clear) + quantized_gemm_family._parallel_compile_available.cache_clear() + with mock.patch.object(quantized_gemm_family, "_async_compile", None): + self.assertFalse(quantized_gemm_family._parallel_compile_available()) + + def test_cpython_before_the_ast_fix_compiles_serially(self) -> None: + for version in ( + (3, 11, 4, "final", 0), + (3, 11, 7, "final", 0), + (3, 12, 1, "final", 0), + ): + with self.subTest(version=version): + self.assertFalse(self._available(version=version)) + + def test_cpython_with_the_ast_fix_compiles_in_parallel(self) -> None: + for version in ( + (3, 11, 8, "final", 0), + (3, 12, 2, "final", 0), + (3, 13, 0, "final", 0), + ): + with self.subTest(version=version): + self.assertTrue(self._available(version=version)) + + 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): + kernel = mock.Mock( + spec=triton.runtime.Autotuner, + configs=[1, 2], + cache={}, + arg_names=[], + keys=[], + ) + with mock.patch.object( + quantized_gemm_family, + "_parallel_compile_available", + return_value=gate, + ): + quantized_gemm_family._precompile_untuned_configs(kernel, (), {}) + self.assertEqual(kernel.warmup.called, precompiled) + + class Int4QuantizedGemmRulesTest(unittest.TestCase): def test_split_k_needs_one_k_tile_per_split(self) -> None: for k in (256, 768, 4096): diff --git a/backends/cuda/triton/kernels/quantized_gemm_family.py b/backends/cuda/triton/kernels/quantized_gemm_family.py index beaff6b32a9..1337c22b016 100644 --- a/backends/cuda/triton/kernels/quantized_gemm_family.py +++ b/backends/cuda/triton/kernels/quantized_gemm_family.py @@ -24,7 +24,13 @@ from __future__ import annotations +import concurrent.futures +import functools import inspect +import logging +import os +import sys +import threading import typing from typing import Callable, Optional, Sequence @@ -33,6 +39,19 @@ from executorch.backends.cuda.triton.kernels.quantized_gemm_utils import splitk_reduce from torch.library import triton_op, wrap_triton +try: + from triton.runtime import _async_compile +except ( + ImportError +): # A Triton without parallel compile compiles each config on first use. + _async_compile = None + +# Threads that compile one kernel's autotune configs. Compiling is mostly ptxas +# and LLVM, which run outside the GIL. +_MAX_COMPILE_THREADS = 32 + +logger: logging.Logger = logging.getLogger(__name__) + def _resolved_signature(prototype) -> inspect.Signature: """The prototype's signature with its annotations evaluated in the @@ -126,6 +145,80 @@ def validate(self, bucket: int, *args) -> None: raise RuntimeError(f"{self.name}_m{bucket}: {reason}") +def _tuning_key(kernel, args: Sequence, kwargs: dict) -> tuple: + # Mirrors the key triton.runtime.Autotuner.run caches its pick under (built + # the same way in Triton 3.8). Should Triton change + # it, lookups miss and work is redone, never skipped; + # test_the_tuning_key_is_the_autotuners_cache_key catches the change. + named = {**dict(zip(kernel.arg_names, args)), **kwargs} + named = {k: v for k, v in named.items() if k in kernel.arg_names} + key = [named[k] for k in kernel.keys if k in named] + key += [str(v.dtype) for v in named.values() if hasattr(v, "dtype")] + return tuple(key) + + +_COMPILE_THREADS_LOCK = threading.Lock() +_compile_thread_pool: Optional[concurrent.futures.ThreadPoolExecutor] = None + + +def _compile_threads() -> concurrent.futures.ThreadPoolExecutor: + """The compile threads, started on first use and shared by every shape.""" + global _compile_thread_pool + with _COMPILE_THREADS_LOCK: + if _compile_thread_pool is None: + _compile_thread_pool = concurrent.futures.ThreadPoolExecutor( + min(_MAX_COMPILE_THREADS, os.cpu_count() or 1), + thread_name_prefix="quantized_gemm_compile", + ) + return _compile_thread_pool + + +@functools.cache +def _parallel_compile_available() -> bool: + if _async_compile is None: + return False + # CPython before 3.11.8 and 3.12.2 cannot run ast.parse from several threads. + version = sys.version_info[:3] + return (3, 11, 8) <= version < (3, 12) or version >= (3, 12, 2) + + +def _precompile_untuned_configs(kernel, args: Sequence, kwargs: dict) -> None: + """Compiles every config an autotuned kernel is about to benchmark, in parallel. + + triton.autotune compiles each config on its first benchmark call, one after + another, and at compile time that is most of a shape's autotune. Compiling + the pruned configs on a thread pool first leaves the benchmark, and so the + pick, unchanged; it then times kernels that are already compiled. A config + that fails to compile fails again in the benchmark, which skips it as before. + """ + 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: + 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 + ) + + def launch_split_k_gemm( kernel, *, @@ -176,16 +269,16 @@ def grid(meta): return (n_tiles, m_tiles) return (n_tiles, m_tiles, split_k) - wrap_triton(kernel)[grid]( + args = ( *inputs, partial, *shape_args, stride_os, partial.stride(-2), partial.stride(-1), - SPLIT_K=split_k, - **constexprs, ) + _precompile_untuned_configs(kernel, args, {"SPLIT_K": split_k, **constexprs}) + wrap_triton(kernel)[grid](*args, SPLIT_K=split_k, **constexprs) if split_k > 1: splitk_reduce(partial, out, block_m, reduce_block_n) return out