diff --git a/vllm/model_executor/kernels/mhc/warmup.py b/vllm/model_executor/kernels/mhc/warmup.py index 6a42ca622270..0abe1a2161c2 100644 --- a/vllm/model_executor/kernels/mhc/warmup.py +++ b/vllm/model_executor/kernels/mhc/warmup.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from collections.abc import Iterable from dataclasses import asdict, dataclass from typing import Any @@ -61,6 +62,21 @@ def compile(self, compile_key: CompileKey) -> None: **asdict(compile_key) ) + def compile_many(self, compile_keys: Iterable[CompileKey]) -> None: + missing = list( + dict.fromkeys( + key for key in compile_keys if key not in self._compiled_cache + ) + ) + if len(missing) < 2: + return super().compile_many(missing) + + # TileLang elaborates serially before compiling independent TIR modules. + compiled = self.kernel.par_compile( + [asdict(key) for key in missing], num_workers=min(4, len(missing)) + ) + self._compiled_cache.update(zip(missing, compiled, strict=True)) + def __call__(self, *tensors, **fields): compile_key = self.dispatch( n_splits=tensors[0].shape[0], diff --git a/vllm/model_executor/warmup/jit_warmup.py b/vllm/model_executor/warmup/jit_warmup.py index 72626964ae67..84f713aa6489 100644 --- a/vllm/model_executor/warmup/jit_warmup.py +++ b/vllm/model_executor/warmup/jit_warmup.py @@ -963,14 +963,18 @@ def compile(self, compile_key: CompileKeyT) -> None: """Compile one warmup key.""" raise NotImplementedError + def compile_many(self, compile_keys: Iterable[CompileKeyT]) -> None: + """Compile a batch of warmup keys, allowing backend-specific scheduling.""" + for compile_key in compile_keys: + self.compile(compile_key) + def register_warmup(self, *args: Any, **kwargs: Any) -> None: """Register this kernel with the active runner's warmup registry.""" JitWarmupRegistry.register(self, *args, **kwargs) def warmup(self, *args: Any, **kwargs: Any) -> None: """Compile this kernel's warmup keys.""" - for compile_key in self.get_warmup_keys(*args, **kwargs): - self.compile(compile_key) + self.compile_many(self.get_warmup_keys(*args, **kwargs)) def _same_value(left: Any, right: Any) -> bool: @@ -1111,5 +1115,4 @@ def warmup(self) -> None: f"{kernel.__class__.__name__} ({len(compile_keys)} keys)", refresh=False, ) - for compile_key in compile_keys: - kernel.compile(compile_key) + kernel.compile_many(compile_keys)