From de7eeaa87bc19ecf088472d5f4471cdc9a5ab995 Mon Sep 17 00:00:00 2001 From: Ayushman Singh <40520701+ayush1399@users.noreply.github.com> Date: Fri, 16 Oct 2026 08:59:42 -0400 Subject: [PATCH] [Perf] Share Triton JIT warmup compilation pool across registered kernels (#61422) --- .../test_jit_warmup_triton_launcher.py | 28 +++++---- vllm/model_executor/warmup/jit_warmup.py | 39 ++++++++++-- .../warmup/jit_warmup_triton_helper.py | 62 ++++++++++++------- 3 files changed, 90 insertions(+), 39 deletions(-) diff --git a/tests/model_executor/test_jit_warmup_triton_launcher.py b/tests/model_executor/test_jit_warmup_triton_launcher.py index c19804eae4..1c02e18b96 100644 --- a/tests/model_executor/test_jit_warmup_triton_launcher.py +++ b/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 diff --git a/vllm/model_executor/warmup/jit_warmup.py b/vllm/model_executor/warmup/jit_warmup.py index 84f713aa64..823c485665 100644 --- a/vllm/model_executor/warmup/jit_warmup.py +++ b/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) diff --git a/vllm/model_executor/warmup/jit_warmup_triton_helper.py b/vllm/model_executor/warmup/jit_warmup_triton_helper.py index 70ec8cd70f..a03fe2f20b 100644 --- a/vllm/model_executor/warmup/jit_warmup_triton_helper.py +++ b/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, ...]: