Skip to content

[2/N][Feat][Perf] Add new warmup infrastructure for JITs. Add predicate filtering for JIT warmup, and migrate Inkling FA4 - #49315

Merged
MatthewBonanni merged 21 commits into
vllm-project:mainfrom
LopezCastroRoberto:feat/inking_warmup
Aug 11, 2026
Merged

MatthewBonanni merged 21 commits into
vllm-project:mainfrom
LopezCastroRoberto:feat/inking_warmup

Conversation

@LopezCastroRoberto

@LopezCastroRoberto LopezCastroRoberto commented Jul 21, 2026

Copy link
Copy Markdown
Contributor

Description

This PR migrates Inkling FA4 attention warmup to the shared JIT warmup contract introduced in #47451, and deprecates the legacy CuTeDSL warmup path (there is a conflict with the Kimi K3 integration #50089 that prevents this deprecation from being completed. It will be addressed in a future PR).

See #49349 for more context

Motivation

  • Respect enable_jit_warmup and run Inkling FA4 warmup through kernel_warmup().
  • Use the shared kernel-owned warmup contract from RFC [RFC]: Design a shared warmup infrastructure for JITs #47456.
  • Run under the standard warmup logging, ordering, and exception-handling path.
  • Keep model construction cheap and side-effect free.
  • Avoid running compilation once per layer during model build (even if cached).
  • Run warmup after config/model setup is complete.
  • Keep runtime execution and startup compilation clearly separated.

What Changed

  • Migrated Inkling FA4 relative attention to a VllmJitKernel wrapper.
  • Deprecated the legacy CuTeDSL warmup registry (there is a conflict with the Kimi K3 integration [Model] Add Kimi K3 support: model files and kernels [1/N] #50089 that prevents this deprecation from being completed. It will be addressed in a future PR).
  • Extended _trace_dispatch(...) with _when predicate support.
  • Simplified FA4 MLA warmup by using zip_inputs(...) plus _when.

AST Warmup Engine Update

_trace_dispatch(...) now accepts _when=<predicate>.

This lets warmup definitions attach a small AST-traced predicate that filters generated dispatch points before they become compile keys. The predicate uses the same simple Python subset as dispatch(...): local assignments, helper calls, arithmetic, comparisons, and boolean expressions.

LopezCastroRoberto and others added 10 commits July 21, 2026 05:52
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>

@claude claude Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

LopezCastroRoberto and others added 2 commits July 21, 2026 15:58
Signed-off-by: LopezCastroRoberto <rocastro@redhat.com>
@mergify mergify Bot added the v1 label Jul 21, 2026
@LopezCastroRoberto LopezCastroRoberto changed the title [Feat] Migrate Inkling FA4 warmup and deprecate legacy CuTeDSL warmup [Feat] Add predicate filtering to JIT warmup and migrate Inkling FA4 Jul 21, 2026
@LopezCastroRoberto

Copy link
Copy Markdown
Contributor Author

@LucasWilkinson

LopezCastroRoberto and others added 2 commits July 29, 2026 11:13
@LopezCastroRoberto LopezCastroRoberto changed the title [2/N][Feat][Perf] Add new warmup infrastructure for JITs. Add predicate filtering, migrate Inkling FA4, and deprecate legacy CuTeDSL warmup path [2/N][Feat][Perf] Add new warmup infrastructure for JITs. Add predicate filtering, and migrate Inkling FA4 Jul 29, 2026
Signed-off-by: LopezCastroRoberto <rocastro@redhat.com>
@LopezCastroRoberto LopezCastroRoberto changed the title [2/N][Feat][Perf] Add new warmup infrastructure for JITs. Add predicate filtering, and migrate Inkling FA4 [2/N][Feat][Perf] Add new warmup infrastructure for JITs. Add predicate filtering for JIT warmup, and migrate Inkling FA4 Jul 29, 2026
Comment thread vllm/model_executor/warmup/cutedsl_warmup.py Outdated
Comment thread vllm/model_executor/warmup/jit_warmup.py Outdated
Comment thread vllm/model_executor/warmup/jit_warmup.py Outdated
Signed-off-by: LopezCastroRoberto <rocastro@redhat.com>

@LucasWilkinson LucasWilkinson left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM thanks!

@LucasWilkinson LucasWilkinson added the ready ONLY add when PR is ready to merge/full CI is needed label Jul 30, 2026
@LucasWilkinson

Copy link
Copy Markdown
Collaborator

/ci run

@github-actions

Copy link
Copy Markdown

✅ Triggered Buildkite CI #81339 for commit 141de4a4dfc3.

@LucasWilkinson

Copy link
Copy Markdown
Collaborator

/ci run

@github-actions

Copy link
Copy Markdown

✅ CI is already running for this commit: https://buildkite.com/vllm/ci/builds/81533

@LopezCastroRoberto

Copy link
Copy Markdown
Contributor Author

/ci run

@github-actions

Copy link
Copy Markdown

✅ Triggered Buildkite CI #83330 for commit 9fa74cf009e1.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

ready ONLY add when PR is ready to merge/full CI is needed v1

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants