[Perf] Parallelize mHC pre-norm JIT warmup (#56683)
Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>
Woosuk Kwon committed
Sep 13, 2026 at 18:58 UTC
3b533197a9419302987c12bcea5e4e06c9d40858
2 files changed
+23
-4
vllm/model_executor/kernels/mhc/warmup.py
+16
@@ -1,6 +1,7 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
+from collections.abc import Iterable
5
from dataclasses import asdict, dataclass
6
from typing import Any
7
@@ -61,6 +62,21 @@ class MHCPreNormKernel(VllmJitKernel["MHCPreNormKernel.CompileKey"]):
62
**asdict(compile_key)
63
)
64
65
+ def compile_many(self, compile_keys: Iterable[CompileKey]) -> None:
66
+ missing = list(
67
+ dict.fromkeys(
68
+ key for key in compile_keys if key not in self._compiled_cache
69
+ )
70
+ )
71
+ if len(missing) < 2:
72
+ return super().compile_many(missing)
73
+
74
+ # TileLang elaborates serially before compiling independent TIR modules.
75
+ compiled = self.kernel.par_compile(
76
+ [asdict(key) for key in missing], num_workers=min(4, len(missing))
77
+ )
78
+ self._compiled_cache.update(zip(missing, compiled, strict=True))
79
+
80
def __call__(self, *tensors, **fields):
81
compile_key = self.dispatch(
82
n_splits=tensors[0].shape[0],
vllm/model_executor/warmup/jit_warmup.py
+7
-4
@@ -963,14 +963,18 @@ class VllmJitKernel(Generic[CompileKeyT], ABC):
963
"""Compile one warmup key."""
964
raise NotImplementedError
965
966
+ def compile_many(self, compile_keys: Iterable[CompileKeyT]) -> None:
967
+ """Compile a batch of warmup keys, allowing backend-specific scheduling."""
968
+ for compile_key in compile_keys:
969
+ self.compile(compile_key)
970
+
971
def register_warmup(self, *args: Any, **kwargs: Any) -> None:
972
"""Register this kernel with the active runner's warmup registry."""
973
JitWarmupRegistry.register(self, *args, **kwargs)
974
975
def warmup(self, *args: Any, **kwargs: Any) -> None:
976
"""Compile this kernel's warmup keys."""
972
- for compile_key in self.get_warmup_keys(*args, **kwargs):
973
- self.compile(compile_key)
977
+ self.compile_many(self.get_warmup_keys(*args, **kwargs))
978
979
980
def _same_value(left: Any, right: Any) -> bool:
@@ -1111,5 +1115,4 @@ class JitWarmupRegistry:
1115
f"{kernel.__class__.__name__} ({len(compile_keys)} keys)",
1116
refresh=False,
1117
)
1114
- for compile_key in compile_keys:
1115
- kernel.compile(compile_key)
1118
+ kernel.compile_many(compile_keys)