[3/N][Feat][Perf] Add new warmup infrastructure for JITs. Add provider registry and orchestration for JIT warmup - #50174
Conversation
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
|
Documentation preview: https://vllm--50174.org.readthedocs.build/en/50174/ |
|
This pull request has merge conflicts that must be resolved before it can be |
Signed-off-by: LopezCastroRoberto <rocastro@redhat.com>
61e7a7f to
6c6505e
Compare
Signed-off-by: LopezCastroRoberto <rocastro@redhat.com>
Signed-off-by: LopezCastroRoberto <rocastro@redhat.com>
Signed-off-by: LopezCastroRoberto <rocastro@redhat.com>
Signed-off-by: LopezCastroRoberto <rocastro@redhat.com>
Signed-off-by: LopezCastroRoberto <rocastro@redhat.com>
Signed-off-by: LopezCastroRoberto <rocastro@redhat.com>
Signed-off-by: LopezCastroRoberto <rocastro@redhat.com>
Signed-off-by: LopezCastroRoberto <rocastro@redhat.com>
Signed-off-by: LopezCastroRoberto <rocastro@redhat.com>
LucasWilkinson
left a comment
LucasWilkinson
left a comment
There was a problem hiding this comment.
Overall looks good to me thank you! left one more comment
| if -(1 << 31) <= value < (1 << 31): | ||
| divisible_rep = 16 | ||
| generic_rep = 2 | ||
| elif -(1 << 63) <= value < (1 << 63): | ||
| divisible_rep = 1 << 31 | ||
| generic_rep = (1 << 31) + 1 | ||
| elif 0 <= value < (1 << 64): | ||
| divisible_rep = 1 << 63 | ||
| generic_rep = (1 << 63) + 1 | ||
| else: | ||
| raise OverflowError(f"Integer {value} is outside Triton's scalar range") |
There was a problem hiding this comment.
can you please document or clarify these magic values?
There was a problem hiding this comment.
@LucasWilkinson This follows Triton’s handle_long_type() implementation. It processes Python int arguments passed to a Triton JIT kernel, and returns:
- The kernel argument type: i32, i64, or u64.
- Its specialization class: constant 1, divisible by 16 (D), or generic.
These values are representative inputs for each (integer type, specialization class) pair. For i64 and u64, the type boundary and boundary-plus-one (i guess these are the magic values you refer to) represent the divisible and generic classes, respectively
LucasWilkinson
added
the
ready
|
/ci run |
|
✅ @LopezCastroRoberto, CI is now available for this PR.
|
|
✅ Triggered Buildkite CI #83781 for commit |
|
/ci run |
|
✅ Triggered Buildkite CI #83796 for commit |
Signed-off-by: LopezCastroRoberto <rocastro@redhat.com>
|
/ci retry |
|
✅ Triggered Buildkite CI #83871 for commit |
|
/ci run |
|
✅ Triggered Buildkite CI #83873 for commit |
|
@LucasWilkinson the CI failure seems unrelated |
|
Hi @LopezCastroRoberto @LucasWilkinson @mgoin While testing DSV4 over the past couple of days, I noticed that a few kernels were not being warmup. Based on your design document, I asked Codex to make some changes to address this issue. The changes are working as expected in my local tests. Could you please help take a look when you have a chance? Thanks! #52740 |
|
@chaunceyjiang These first PRs were just defining/implementing the warmup infra. For full DSv4 de-JITification we need to land #49627, which at the same time has been broken down into four independent PRs to make the review process easier: #50175, #50176, #50177, and #50178. Full details in #49349 Only after all these PRs land, we can consider DSv4 de-JITification completed. The kernels warmed up in your PR will be mostly covered by: #50178 Thanks |
|
@LopezCastroRoberto Thanks for the quick response! Looking forward to seeing those changes merged soon. I’ve closed my PR. |
…r registry and orchestration for JIT warmup (vllm-project#50174) Signed-off-by: LopezCastroRoberto <rocastro@redhat.com> Co-authored-by: Codex <codex@openai.com>
…r registry and orchestration for JIT warmup (vllm-project#50174) Signed-off-by: LopezCastroRoberto <rocastro@redhat.com> Co-authored-by: Codex <codex@openai.com> Signed-off-by: Wyett <wyettzeng@gmail.com>
Upstream's JIT warmup infrastructure (vllm-project#50174) absorbed the standalone v1_block_table_warmup module into BlockTable's registered warmup, and the worker now reads kv_cache_config.kv_cache_layout when adopting the engine core's layout. Drop the stale import/call and teach the gpu_worker test fake about the new field.
Upstream's warmup infrastructure (vllm-project#50174) registers these kernels through its own provider registry, so the fork-side warmup hooks are redundant. Reverse-apply the deltas of the upstream warmup PRs (vllm-project#42193, vllm-project#42215, vllm-project#43642, vllm-project#46446) so every file they touched matches origin/main: - Drop the fused MoE, TurboQuant, hybrid GDN/Mamba/MRoPE and block-table warmup modules plus their tests and kernel_warmup wiring. - Restore triton_decode_attention, triton_turboquant_decode, fused_recurrent and fused_moe to upstream (removes the VllmJitKernel wrappers those PRs introduced). - Keep `import math` in mrope.py: it is used by the bounded M-RoPE cache work, not by the reverted warmup code.
5 participants
Description
This PR extends the shared JIT warmup infrastructure with provider registration and centralized orchestration. It builds on #49315 and the contract described in #47456.
For more details, see parent (draft) PR: #49627 and tracking list issue #49349
Motivation
kernel_warmup()logging, ordering, progress reporting, and exception handling.enable_jit_warmup.What Changed