[Perf] Share Triton JIT warmup compilation pool across registered kernels (#61422)

tests/model_executor/test_jit_warmup_triton_launcher.py

@@ -168,23 +168,28 @@ def _warmup_store_kernel(out, VALUE: tl.constexpr):
     not current_platform.is_cuda() or not hasattr(triton, "AsyncCompileMode"),
     reason="Requires CUDA Triton async compilation",
 )
-def test_registered_warmup_parallelizes_compile_only_and_runtime_stays_serial(
+def test_registered_warmup_overlaps_kernels_and_runtime_stays_serial(
     monkeypatch,
 ) -> None:
     from triton.runtime._async_compile import active_mode
 
     monkeypatch.setattr(envs, "VLLM_TRITON_JIT_WARMUP_NUM_THREADS", 2)
 
-    def warmup_inputs():
-        return dict(out=TritonWarmupTensor(torch.int32), value=WarmupChoices(1, 2))
+    def make_store(value):
+        def warmup_inputs():
+            return dict(out=TritonWarmupTensor(torch.int32), value=value)
 
-    @triton_kernel_dispatcher_with_warmup(
-        kernel=_warmup_store_kernel, warmup_inputs=warmup_inputs
-    )
-    def store(out, value):
-        return (1,), dict(VALUE=value)
+        @triton_kernel_dispatcher_with_warmup(
+            kernel=_warmup_store_kernel, warmup_inputs=warmup_inputs
+        )
+        def store(out, value):
+            return (1,), dict(VALUE=value)
 
-    store.get_warmup_keys()  # Initialize Triton's binder before wrapping its compiler.
+        # Initialize Triton's binder before wrapping its compiler.
+        store.get_warmup_keys()
+        return store
+
+    stores = [make_store(value) for value in (1, 2)]
     caller = threading.get_ident()
     barrier = threading.Barrier(2, timeout=10)
     compiled_on = []
@@ -200,14 +205,15 @@ def test_registered_warmup_parallelizes_compile_only_and_runtime_stays_serial(
     monkeypatch.setattr(_warmup_store_kernel, "compile", record_compile)
     registry = JitWarmupRegistry(None)
     with registry.activate():
-        store.register_warmup()
+        for store in stores:
+            store.register_warmup()
     registry.warmup()
 
     assert len(compiled_on) == 2
     assert all(thread != caller for thread in compiled_on)
     assert active_mode.get() is None
     out = torch.empty(1, dtype=torch.int32, device=current_platform.device_type)
-    for value in (1, 2):
+    for store, value in zip(stores, (1, 2)):
         store(out, value)
         assert out.item() == value
     assert len(compiled_on) == 2

vllm/model_executor/warmup/jit_warmup.py

@@ -1082,6 +1082,10 @@ class JitWarmupRegistry:
         from tqdm import tqdm
 
         from vllm.distributed import is_global_first_rank
+        from vllm.model_executor.warmup.jit_warmup_triton_helper import (
+            VllmTritonJitKernel,
+            compile_triton_warmup,
+        )
 
         kernel_items: list[tuple[VllmJitKernel[Any], dict[Any, None]]] = []
         for kernel, registrations in self._registrations.items():
@@ -1102,17 +1106,40 @@ class JitWarmupRegistry:
         if not kernel_items:
             return
 
+        parallel_items: list[tuple[VllmJitKernel[Any], dict[Any, None]]] = []
+        serial_items = []
+        for kernel, compile_keys in kernel_items:
+            if (
+                isinstance(kernel, VllmTritonJitKernel)
+                and kernel.supports_parallel_compile
+            ):
+                parallel_items.append((kernel, compile_keys))
+            else:
+                serial_items.append((kernel, compile_keys))
+
         total_keys = sum(len(compile_keys) for _, compile_keys in kernel_items)
         with tqdm(
-            kernel_items,
+            total=len(kernel_items),
             desc=f"JIT kernel warmup ({total_keys} compile keys)",
             disable=not is_global_first_rank(),
             dynamic_ncols=True,
             unit="kernel",
         ) as progress:
-            for kernel, compile_keys in progress:
-                progress.set_postfix_str(
-                    f"{kernel.__class__.__name__} ({len(compile_keys)} keys)",
-                    refresh=False,
+
+            def compile_kernels(
+                items: list[tuple[VllmJitKernel[Any], dict[Any, None]]],
+            ) -> None:
+                for kernel, compile_keys in items:
+                    progress.set_postfix_str(
+                        f"{kernel.__class__.__name__} ({len(compile_keys)} keys)",
+                        refresh=False,
+                    )
+                    kernel.compile_many(compile_keys)
+                    progress.update()
+
+            if parallel_items:
+                compile_triton_warmup(
+                    lambda: compile_kernels(parallel_items),
+                    sum(len(keys) for _, keys in parallel_items),
                 )
-                kernel.compile_many(compile_keys)
+            compile_kernels(serial_items)

vllm/model_executor/warmup/jit_warmup_triton_helper.py

@@ -6,7 +6,7 @@ from abc import abstractmethod
 from collections.abc import Callable, Hashable, Iterable, Mapping
 from concurrent.futures import ThreadPoolExecutor
 from contextlib import nullcontext, suppress
-from contextvars import copy_context
+from contextvars import ContextVar, copy_context
 from dataclasses import dataclass, field
 from functools import cache, cached_property, update_wrapper
 from typing import Any, Generic, ParamSpec, Protocol, TypeVar, cast, overload
@@ -207,6 +207,31 @@ class TritonWarmupTensor:
         return result if dim is None else result[dim]
 
 
+_shared_compile_pool: ContextVar[bool] = ContextVar(
+    "shared_triton_warmup_compile_pool", default=False
+)
+
+
+def compile_triton_warmup(compile_all: Callable[[], None], num_keys: int) -> None:
+    """Share one compiler pool across all eligible warmup kernels."""
+    num_threads = min(envs.VLLM_TRITON_JIT_WARMUP_NUM_THREADS, num_keys)
+    if num_threads <= 1:
+        compile_all()
+        return
+
+    def compile_in_context() -> None:
+        _shared_compile_pool.set(True)
+        with (
+            ThreadPoolExecutor(max_workers=num_threads) as executor,
+            triton.AsyncCompileMode(executor),
+        ):
+            # Dispatch mutates owner state and must stay on this thread.
+            compile_all()
+
+    # Older Triton versions leave async mode active on compilation failure.
+    copy_context().run(compile_in_context)
+
+
 class VllmTritonJitKernel(VllmJitKernel[CompileKeyT], Generic[CompileKeyT]):
     """Triton owner whose runtime launch specification is reused for warmup."""
 
@@ -230,32 +255,25 @@ class VllmTritonJitKernel(VllmJitKernel[CompileKeyT], Generic[CompileKeyT]):
             self._warming = False
             self._warming_compile_key = None
 
+    @property
+    def supports_parallel_compile(self) -> bool:
+        # AMD LLVM code generation can abort under concurrent compilation.
+        return (
+            current_platform.is_cuda()
+            and not self._run_autotune
+            and getattr(triton, "AsyncCompileMode", None) is not None
+            and envs.VLLM_TRITON_JIT_WARMUP_NUM_THREADS > 1
+        )
+
     def compile_many(self, compile_keys: Iterable[CompileKeyT]) -> None:
         """Compile CUDA startup warmup variants in parallel, then wait."""
         keys = list(compile_keys)
-        num_threads = min(envs.VLLM_TRITON_JIT_WARMUP_NUM_THREADS, len(keys))
-        async_compile = getattr(triton, "AsyncCompileMode", None)
-        # AMD LLVM code generation can abort under concurrent compilation.
-        if (
-            not current_platform.is_cuda()
-            or self._run_autotune
-            or async_compile is None
-            or num_threads <= 1
-        ):
+        if not self.supports_parallel_compile or _shared_compile_pool.get():
             return super().compile_many(keys)
 
-        def compile_all() -> None:
-            with (
-                ThreadPoolExecutor(max_workers=num_threads) as executor,
-                async_compile(executor),
-            ):
-                # Dispatch stays on this thread: compile() mutates owner state.
-                # Only Triton's underlying compiler work runs in the pool.
-                super(VllmTritonJitKernel, self).compile_many(keys)
-
-        # Older Triton versions leave async mode active on compilation failure.
-        # Isolate that state so it cannot affect later runtime JIT compilation.
-        copy_context().run(compile_all)
+        compile_triton_warmup(
+            lambda: super(VllmTritonJitKernel, self).compile_many(keys), len(keys)
+        )
 
     @cached_property
     def _kernel_arg_names(self) -> tuple[str, ...]: