[Perf] Parallelize mHC pre-norm JIT warmup (#56683)

Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>

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 @@ class MHCPreNormKernel(VllmJitKernel["MHCPreNormKernel.CompileKey"]):
                 **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],

vllm/model_executor/warmup/jit_warmup.py

@@ -963,14 +963,18 @@ class VllmJitKernel(Generic[CompileKeyT], ABC):
         """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 @@ class JitWarmupRegistry:
                     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)