[Perf] Parallelize mHC pre-norm JIT warmup

[WoosukKwon]

Purpose

Compile independent mHC pre-norm warmup variants concurrently. Add a compile_many() hook to VllmJitKernel, retain sequential compilation by default, and use TileLang's par_compile() for missing MHCPreNormKernel variants with at most four compiler workers per process. Cached variants are retained and duplicate keys are removed. TileLang elaborates the modules serially before compiling them in parallel.

The existing open PRs #51802 and #49707 address mHC warmup coverage. This change preserves the existing registered keys and runtime dispatch, and changes how those keys are compiled. Searches for parallel JIT compilation and par_compile found no open PR implementing the same change. This PR is independent of the dummy-initialization change in #56682.

Test Plan

PYTEST_DISABLE_PLUGIN_AUTOLOAD=1 VLLM_USE_RUST_FRONTEND=0 .venv/bin/python -m pytest tests/model_executor/test_jit_warmup.py tests/kernels/test_mhc_jit_warmup.py -q
.venv/bin/pre-commit run --files vllm/model_executor/warmup/jit_warmup.py vllm/model_executor/kernels/mhc/warmup.py

Also exercised the existing test_deepseek_v41_mhc_pre_delayed reference and CUDA graph replay assertions after batch warmup, for broadcast, identity, and carried pre-mix (33 tokens, fused normalization, DeepGEMM enabled, 20 Sinkhorn iterations).

Test Result

60 existing warmup tests passed in 18.86 seconds. All three GPU reference/CUDA graph replay cases passed on GB200 after batch warmup. Applicable pre-commit hooks and git diff --check passed. No test additions are included.

On GB200, nine cold-cache DSV4.1 pre-norm variants (hidden size 5120, hc_mult 4, three input/pre-mix configurations and split counts 1/4/16) compiled in:

Compilation	Wall time
Sequential	95.03s
Four compiler workers	43.39s

This is a 2.19x speedup for the isolated compilation stage. These are single-run measurements excluding imports/key enumeration, with separate empty TileLang caches; host-side engine initialization was concurrent, so they are exploratory measurements rather than a controlled end-to-end startup comparison. Generated CUDA source sets matched between the two runs.

A later TP4 real-weight DSV4.1 startup and greedy completion succeeded with this change. Its registry mHC phase had five variants already compiled by profiling and compiled the remaining four in parallel in approximately 14 seconds. No model-quality benchmark was run; this changes compilation scheduling, not kernel math. Full-engine startup speedup is not claimed.

AI assistance was used for diagnosis, implementation, and validation.

[WoosukKwon]

/ci run

[github-actions]

✅ Triggered Buildkite CI #88687 for commit 293c2f5bda57.