Share Triton JIT warmup compilation pool across registered kernels

[ayush1399]

Overview

Follow-up to #58582. Eligible Triton kernels now share one AsyncCompileMode pool during startup warmup, instead of each kernel draining its own pool in turn.

Claims

• Compilation now overlaps across kernels, not just within one kernel. Single-key kernels previously always compiled serially.
• 9–12% faster cold JIT warmup at default threads, 23–28% at 16 threads, when several Triton warmup kernels are registered (MRV1 sampler/spec-decode kernels, DSv4, GLM-5 Next).
• No change for ROCm, autotuned kernels, non-Triton owners, or configs with a single registered kernel.

Validation

pytest tests/model_executor/test_jit_warmup_triton_launcher.py: 15 passed. The updated test (two single-key kernels must compile concurrently) fails on the base commit.

Cold-cache JIT kernel warmup finished in …s on H100, VLLM_USE_V2_MODEL_RUNNER=0, mean of 3 runs per arm:

Config	Threads	Before	After
Gemma 4 31B + EAGLE3, TP=8 (12 kernels)	4	20.72s	18.86s (−9%)
	16	15.25s	11.69s (−23%)
Qwen3-0.6B-FP8 + ngram (9 kernels)	4	17.99s	15.89s (−12%)
	16	11.55s	8.37s (−28%)

Gemma 4 on default MRV2 registers only the FA4 kernel, so there's nothing to overlap: 5.74s → 5.73s. No model evals: only compile scheduling changes, not the compiled kernels.

Repro

VLLM_USE_V2_MODEL_RUNNER=0 VLLM_TRITON_JIT_WARMUP_NUM_THREADS=16 \
TRITON_CACHE_DIR=$(mktemp -d) VLLM_CACHE_ROOT=$(mktemp -d) python -c "
from vllm import LLM
LLM('google/gemma-4-31b-it', tensor_parallel_size=8, max_model_len=8192,
    kernel_config={'enable_flashinfer_autotune': False},
    speculative_config={'method': 'eagle3', 'num_speculative_tokens': 3,
                        'model': 'RedHatAI/gemma-4-31B-it-speculator.eagle3'})
" 2>&1 | grep 'JIT kernel warmup finished'

Details

JitWarmupRegistry.warmup() runs all kernels with supports_parallel_compile inside one compile_triton_warmup() pool, then the rest serially. A ContextVar stops compile_many from nesting pools; called directly, it still creates its own pool. Dispatch stays on the calling thread, as in #58582.

Limitations:

• The tqdm bar advances on dispatch, so it jumps ahead and then waits for the pool to drain.
• Compile errors surface when the shared pool exits, not inside the failing kernel's compile_many.

Not a duplicate: no open PR shares a pool across kernels. #56154 (compile-key derivation) and #56761 (TileLang) are the closest.

---

Pull Request Checklist

• I used vLLM's /pr-checklist skill. (Mandatory for agents, optional for humans).
• AI assistance was used during the creation of this PR.
• Design Fit: Minimizes impact on core components, reuses existing functionality, and justifies added complexity.
• Testing and Validation: Validates the change and ensures any added tests are meaningful and reliable, with CI coverage or documented CI resource constraints and validation performed outside CI.
• Code Quality and Style: Keeps code and comments clear and concise, and updates relevant documentation and examples.
• Pull Request Contents: Includes a brief summary and relevant links, supports claims with evidence, explains root causes and implementation trade-offs, and follows the contributing guide.

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.