[Perf] Parallelize registered CUDA Triton kernel warmup at startup (#58582)
Signed-off-by: mgoin <mgoin64@gmail.com> Co-authored-by: Codex <noreply@openai.com>
Misha Goin committed
Sep 24, 2026 at 19:52 UTC
71891bdb94d79533049c9eb314884f53480eea3d
3 files changed
+139
tests/model_executor/test_jit_warmup_triton_launcher.py
+102
@@ -1,14 +1,17 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
+import threading
5
from dataclasses import dataclass
6
from typing import Any
7
8
import pytest
9
import torch
10
11
+from vllm import envs
12
from vllm.model_executor.warmup import jit_warmup_triton_helper
13
from vllm.model_executor.warmup.jit_warmup import (
14
+ JitWarmupRegistry,
15
WarmupChoices,
16
WarmupIntRange,
17
kernel_launcher,
@@ -116,6 +119,105 @@ def test_triton_launcher_supports_compile_and_runtime_adapters() -> None:
119
assert runtime_calls == [(owner.kernel, (2,), ("runtime", 2), {"CONST": 7})]
120
121
122
+@pytest.mark.skipif(
123
+ not current_platform.is_cuda() or not hasattr(triton, "AsyncCompileMode"),
124
+ reason="Requires CUDA Triton async compilation",
125
+)
126
+def test_failed_parallel_warmup_does_not_leak_into_runtime(monkeypatch) -> None:
127
+ from triton.runtime._async_compile import active_mode
128
+
129
+ monkeypatch.setattr(envs, "VLLM_TRITON_JIT_WARMUP_NUM_THREADS", 2)
130
+ owner = _TestTritonKernel()
131
+ owner.kernel = _FakeTritonKernel()
132
+
133
+ def fail():
134
+ raise RuntimeError("warmup failed")
135
+
136
+ def warmup(**kwargs):
137
+ active_mode.get().submit(kwargs["second"], fail, lambda result: None)
138
+
139
+ monkeypatch.setattr(owner.kernel, "warmup", warmup)
140
+ with pytest.raises(RuntimeError, match="warmup failed"):
141
+ owner.compile_many(owner.CompileKey(value=i) for i in range(2))
142
+
143
+ assert active_mode.get() is None
144
+ owner("runtime", 3, None)
145
+ assert len(owner.kernel.runtime_calls) == 1
146
+
147
+
148
+def test_autotuning_warmup_stays_serial(monkeypatch) -> None:
149
+ owner = _TestTritonKernel()
150
+ owner.kernel = _FakeTritonKernel()
151
+ owner._run_autotune = True
152
+ monkeypatch.setattr(envs, "VLLM_TRITON_JIT_WARMUP_NUM_THREADS", 4)
153
+
154
+ def unexpected_async(*args):
155
+ pytest.fail("Autotuning must not use asynchronous compilation")
156
+
157
+ monkeypatch.setattr(triton, "AsyncCompileMode", unexpected_async, raising=False)
158
+ owner.compile_many(owner.CompileKey(value=i) for i in range(2))
159
+ assert len(owner.kernel.runtime_calls) == 2
160
+
161
+
162
+@triton.jit
163
+def _warmup_store_kernel(out, VALUE: tl.constexpr):
164
+ tl.store(out, VALUE)
165
+
166
+
167
+@pytest.mark.skipif(
168
+ not current_platform.is_cuda() or not hasattr(triton, "AsyncCompileMode"),
169
+ reason="Requires CUDA Triton async compilation",
170
+)
171
+def test_registered_warmup_parallelizes_compile_only_and_runtime_stays_serial(
172
+ monkeypatch,
173
+) -> None:
174
+ from triton.runtime._async_compile import active_mode
175
+
176
+ monkeypatch.setattr(envs, "VLLM_TRITON_JIT_WARMUP_NUM_THREADS", 2)
177
+
178
+ def warmup_inputs():
179
+ return dict(out=TritonWarmupTensor(torch.int32), value=WarmupChoices(1, 2))
180
+
181
+ @triton_kernel_dispatcher_with_warmup(
182
+ kernel=_warmup_store_kernel, warmup_inputs=warmup_inputs
183
+ )
184
+ def store(out, value):
185
+ return (1,), dict(VALUE=value)
186
+
187
+ store.get_warmup_keys() # Initialize Triton's binder before wrapping its compiler.
188
+ caller = threading.get_ident()
189
+ barrier = threading.Barrier(2, timeout=10)
190
+ compiled_on = []
191
+ compile_kernel = _warmup_store_kernel.compile
192
+
193
+ def record_compile(*args, **kwargs):
194
+ if threading.get_ident() != caller:
195
+ barrier.wait()
196
+ result = compile_kernel(*args, **kwargs)
197
+ compiled_on.append(threading.get_ident())
198
+ return result
199
+
200
+ monkeypatch.setattr(_warmup_store_kernel, "compile", record_compile)
201
+ registry = JitWarmupRegistry(None)
202
+ with registry.activate():
203
+ store.register_warmup()
204
+ registry.warmup()
205
+
206
+ assert len(compiled_on) == 2
207
+ assert all(thread != caller for thread in compiled_on)
208
+ assert active_mode.get() is None
209
+ out = torch.empty(1, dtype=torch.int32, device=current_platform.device_type)
210
+ for value in (1, 2):
211
+ store(out, value)
212
+ assert out.item() == value
213
+ assert len(compiled_on) == 2
214
+
215
+ # A specialization first seen during inference uses the normal blocking JIT.
216
+ store(out, 3)
217
+ assert out.item() == 3
218
+ assert compiled_on[2:] == [caller]
219
+
220
+
221
def test_triton_launcher_supports_cpu_function_wrappers() -> None:
222
calls: list[tuple[Any, ...]] = []
223
vllm/envs.py
+7
@@ -123,6 +123,7 @@ if TYPE_CHECKING:
123
VLLM_USE_TRITON_AWQ: bool = False
124
VLLM_FASTSAFETENSORS_QUEUE_SIZE: int = 0
125
VLLM_TRITON_FORCE_FIRST_CONFIG: bool = False
126
+ VLLM_TRITON_JIT_WARMUP_NUM_THREADS: int = 4
127
VLLM_ALLOW_RUNTIME_LORA_UPDATING: bool = False
128
VLLM_SKIP_P2P_CHECK: bool = False
129
VLLM_DISABLED_KERNELS: list[str] = []
@@ -1181,6 +1182,11 @@ environment_variables: dict[str, Callable[[], Any]] = {
1182
os.environ.get("VLLM_TRITON_FORCE_FIRST_CONFIG", "0").strip().lower()
1183
in ("1", "true")
1184
),
1185
+ # Maximum compiler threads per worker for registered CUDA Triton JIT warmup.
1186
+ # Set to 1 for serial warmup. Does not affect runtime JIT or autotuning.
1187
+ "VLLM_TRITON_JIT_WARMUP_NUM_THREADS": lambda: int(
1188
+ os.getenv("VLLM_TRITON_JIT_WARMUP_NUM_THREADS", "4")
1189
+ ),
1190
# If set, allow loading or unloading lora adapters in runtime,
1191
"VLLM_ALLOW_RUNTIME_LORA_UPDATING": lambda: (
1192
os.environ.get("VLLM_ALLOW_RUNTIME_LORA_UPDATING", "0").strip().lower()
@@ -2272,6 +2278,7 @@ def compile_factors() -> dict[str, object]:
2278
hash everything else. This keeps the cache key aligned across workers."""
2279
ignored_factors: set[str] = {
2280
"MAX_JOBS",
2281
+ "VLLM_TRITON_JIT_WARMUP_NUM_THREADS",
2282
"VLLM_RPC_BASE_PATH",
2283
"VLLM_USE_MODELSCOPE",
2284
"VLLM_RINGBUFFER_WARNING_INTERVAL",
vllm/model_executor/warmup/jit_warmup_triton_helper.py
+30
@@ -4,11 +4,14 @@ import ast
4
import inspect
5
from abc import abstractmethod
6
from collections.abc import Callable, Hashable, Iterable, Mapping
7
+from concurrent.futures import ThreadPoolExecutor
8
from contextlib import nullcontext, suppress
9
+from contextvars import copy_context
10
from dataclasses import dataclass, field
11
from functools import cache, cached_property, update_wrapper
12
from typing import Any, Generic, ParamSpec, Protocol, TypeVar, cast, overload
13
14
+from vllm import envs
15
from vllm.model_executor.warmup.jit_warmup import (
16
VllmJitKernel,
17
get_ast_full_name,
@@ -227,6 +230,33 @@ class VllmTritonJitKernel(VllmJitKernel[CompileKeyT], Generic[CompileKeyT]):
230
self._warming = False
231
self._warming_compile_key = None
232
233
+ def compile_many(self, compile_keys: Iterable[CompileKeyT]) -> None:
234
+ """Compile CUDA startup warmup variants in parallel, then wait."""
235
+ keys = list(compile_keys)
236
+ num_threads = min(envs.VLLM_TRITON_JIT_WARMUP_NUM_THREADS, len(keys))
237
+ async_compile = getattr(triton, "AsyncCompileMode", None)
238
+ # AMD LLVM code generation can abort under concurrent compilation.
239
+ if (
240
+ not current_platform.is_cuda()
241
+ or self._run_autotune
242
+ or async_compile is None
243
+ or num_threads <= 1
244
+ ):
245
+ return super().compile_many(keys)
246
+
247
+ def compile_all() -> None:
248
+ with (
249
+ ThreadPoolExecutor(max_workers=num_threads) as executor,
250
+ async_compile(executor),
251
+ ):
252
+ # Dispatch stays on this thread: compile() mutates owner state.
253
+ # Only Triton's underlying compiler work runs in the pool.
254
+ super(VllmTritonJitKernel, self).compile_many(keys)
255
+
256
+ # Older Triton versions leave async mode active on compilation failure.
257
+ # Isolate that state so it cannot affect later runtime JIT compilation.
258
+ copy_context().run(compile_all)
259
+
260
@cached_property
261
def _kernel_arg_names(self) -> tuple[str, ...]:
262
arg_names = getattr(self.kernel, "arg_names", None)