[Perf] Parallelize registered CUDA Triton kernel warmup at startup

tests/model_executor/test_jit_warmup_triton_launcher.py

@@ -1,14 +1,17 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+import threading
 from dataclasses import dataclass
 from typing import Any
 
 import pytest
 import torch
 
+from vllm import envs
 from vllm.model_executor.warmup import jit_warmup_triton_helper
 from vllm.model_executor.warmup.jit_warmup import (
+    JitWarmupRegistry,
     WarmupChoices,
     WarmupIntRange,
     kernel_launcher,
@@ -116,6 +119,104 @@ def runtime_launcher(kernel: Any, grid: Any, *args: Any, **kwargs: Any) -> None:
     assert runtime_calls == [(owner.kernel, (2,), ("runtime", 2), {"CONST": 7})]
 
 
+@pytest.mark.skipif(
+    not hasattr(triton, "AsyncCompileMode"), reason="Requires Triton async compilation"
+)
+def test_failed_parallel_warmup_does_not_leak_into_runtime(monkeypatch) -> None:
+    from triton.runtime._async_compile import active_mode
+
+    monkeypatch.setattr(envs, "VLLM_TRITON_JIT_WARMUP_NUM_THREADS", 2)
+    owner = _TestTritonKernel()
+    owner.kernel = _FakeTritonKernel()
+
+    def fail():
+        raise RuntimeError("warmup failed")
+
+    def warmup(**kwargs):
+        active_mode.get().submit(kwargs["second"], fail, lambda result: None)
+
+    monkeypatch.setattr(owner.kernel, "warmup", warmup)
+    with pytest.raises(RuntimeError, match="warmup failed"):
+        owner.compile_many(owner.CompileKey(value=i) for i in range(2))
+
+    assert active_mode.get() is None
+    owner("runtime", 3, None)
+    assert len(owner.kernel.runtime_calls) == 1
+
+
+def test_autotuning_warmup_stays_serial(monkeypatch) -> None:
+    owner = _TestTritonKernel()
+    owner.kernel = _FakeTritonKernel()
+    owner._run_autotune = True
+    monkeypatch.setattr(envs, "VLLM_TRITON_JIT_WARMUP_NUM_THREADS", 4)
+
+    def unexpected_async(*args):
+        pytest.fail("Autotuning must not use asynchronous compilation")
+
+    monkeypatch.setattr(triton, "AsyncCompileMode", unexpected_async, raising=False)
+    owner.compile_many(owner.CompileKey(value=i) for i in range(2))
+    assert len(owner.kernel.runtime_calls) == 2
+
+
+@triton.jit
+def _warmup_store_kernel(out, VALUE: tl.constexpr):
+    tl.store(out, VALUE)
+
+
+@pytest.mark.skipif(
+    not current_platform.is_cuda_alike() or not hasattr(triton, "AsyncCompileMode"),
+    reason="Requires GPU Triton async compilation",
+)
+def test_registered_warmup_parallelizes_compile_only_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))
+
+    @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.
+    caller = threading.get_ident()
+    barrier = threading.Barrier(2, timeout=10)
+    compiled_on = []
+    compile_kernel = _warmup_store_kernel.compile
+
+    def record_compile(*args, **kwargs):
+        if threading.get_ident() != caller:
+            barrier.wait()
+        result = compile_kernel(*args, **kwargs)
+        compiled_on.append(threading.get_ident())
+        return result
+
+    monkeypatch.setattr(_warmup_store_kernel, "compile", record_compile)
+    registry = JitWarmupRegistry(None)
+    with registry.activate():
+        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):
+        store(out, value)
+        assert out.item() == value
+    assert len(compiled_on) == 2
+
+    # A specialization first seen during inference uses the normal blocking JIT.
+    store(out, 3)
+    assert out.item() == 3
+    assert compiled_on[2:] == [caller]
+
+
 def test_triton_launcher_supports_cpu_function_wrappers() -> None:
     calls: list[tuple[Any, ...]] = []
 

vllm/envs.py

@@ -123,6 +123,7 @@
     VLLM_USE_TRITON_AWQ: bool = False
     VLLM_FASTSAFETENSORS_QUEUE_SIZE: int = 0
     VLLM_TRITON_FORCE_FIRST_CONFIG: bool = False
+    VLLM_TRITON_JIT_WARMUP_NUM_THREADS: int = 4
     VLLM_ALLOW_RUNTIME_LORA_UPDATING: bool = False
     VLLM_SKIP_P2P_CHECK: bool = False
     VLLM_DISABLED_KERNELS: list[str] = []
@@ -1181,6 +1182,11 @@ def _resolve_rust_cli_path() -> str | None:
         os.environ.get("VLLM_TRITON_FORCE_FIRST_CONFIG", "0").strip().lower()
         in ("1", "true")
     ),
+    # Maximum compiler threads per worker for registered Triton JIT warmup.
+    # Set to 1 for serial warmup. Does not affect runtime JIT or autotuning.
+    "VLLM_TRITON_JIT_WARMUP_NUM_THREADS": lambda: int(
+        os.getenv("VLLM_TRITON_JIT_WARMUP_NUM_THREADS", "4")
+    ),
     # If set, allow loading or unloading lora adapters in runtime,
     "VLLM_ALLOW_RUNTIME_LORA_UPDATING": lambda: (
         os.environ.get("VLLM_ALLOW_RUNTIME_LORA_UPDATING", "0").strip().lower()
@@ -2272,6 +2278,7 @@ def compile_factors() -> dict[str, object]:
     hash everything else. This keeps the cache key aligned across workers."""
     ignored_factors: set[str] = {
         "MAX_JOBS",
+        "VLLM_TRITON_JIT_WARMUP_NUM_THREADS",
         "VLLM_RPC_BASE_PATH",
         "VLLM_USE_MODELSCOPE",
         "VLLM_RINGBUFFER_WARNING_INTERVAL",

vllm/model_executor/warmup/jit_warmup_triton_helper.py

@@ -4,11 +4,14 @@
 import inspect
 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 dataclasses import dataclass, field
 from functools import cache, cached_property, update_wrapper
 from typing import Any, Generic, ParamSpec, Protocol, TypeVar, cast, overload
 
+from vllm import envs
 from vllm.model_executor.warmup.jit_warmup import (
     VllmJitKernel,
     get_ast_full_name,
@@ -227,6 +230,27 @@ def compile(self, compile_key: CompileKeyT) -> None:
             self._warming = False
             self._warming_compile_key = None
 
+    def compile_many(self, compile_keys: Iterable[CompileKeyT]) -> None:
+        """Compile startup warmup variants in parallel, then wait for completion."""
+        keys = list(compile_keys)
+        num_threads = min(envs.VLLM_TRITON_JIT_WARMUP_NUM_THREADS, len(keys))
+        async_compile = getattr(triton, "AsyncCompileMode", None)
+        if self._run_autotune or async_compile is None or num_threads <= 1:
+            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)
+
     @cached_property
     def _kernel_arg_names(self) -> tuple[str, ...]:
         arg_names = getattr(self.kernel, "arg_names", None)