[executorch][cuda] Compile an autotuned quantized GEMM's configs in parallel before benchmarking them

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):

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