[3/N][warmup][DSv4] Migrate FA4 MLA and shared CuTeDSL kernels by LopezCastroRoberto · Pull Request #53565 · vllm-project/vllm

LopezCastroRoberto · 2026-08-24T10:39:45Z

Depends on: #50175

For more details, see parent (draft) PR: #49627 and tracking list issue #49349

Description

This PR migrates the FA4 MLA prefill and Inkling FA4 relative-attention paths to the shared JIT warmup contract.

What Changed

Migrated FA4 MLA prefill specialization and compile-only warmup.
Migrated Inkling FA4 relative-attention warmup.
Removed the corresponding legacy external FA4 warmup orchestration.
Kept compile-key generation and warmup search spaces next to their runtime kernels.
Preserved the existing FA4 compile-only adapters and specialization behavior.
Added tests comparing migrated warmup keys with the pre-migration logic.

Co-authored-by: Codex <codex@openai.com> Signed-off-by: LopezCastroRoberto <rocastro@redhat.com>

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.

mergify · 2026-09-01T15:51:32Z

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @LopezCastroRoberto.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

…armup_migration/dsv4_7 Signed-off-by: LopezCastroRoberto <rocastro@redhat.com> # Conflicts: # vllm/model_executor/layers/attention/mla_attention.py # vllm/model_executor/warmup/kernel_warmup.py

LucasWilkinson

LGTM! thanks!

github-actions · 2026-09-03T21:14:35Z

✅ @LopezCastroRoberto, CI is now available for this PR.

/ci run starts upstream CI; /amd-ci run starts AMD CI only.
/ci retry retries failed jobs in the CI build for the current PR head. If the current head has no CI build, it starts a new CI build for the current head containing only jobs that failed in the latest earlier CI build for this PR.
/amd-ci retry retries failed jobs in AMD CI for the current PR head. Use /amd-ci run when the current head has no AMD CI build.
/ci cancel cancels scheduled or running CI builds for this PR branch; /amd-ci cancel does the same for AMD CI only.

coderabbitai · 2026-09-04T08:20:57Z

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Repository UI

Review profile: CHILL

Plan: Team

Run ID: d9497c44-7465-4e1e-8704-efec1a843182

📥 Commits

Reviewing files that changed from the base of the PR and between 9cc7793 and 9ccd8bf.

📒 Files selected for processing (8)

tests/models/inkling/test_fa4_rel_attention.py
tests/v1/attention/test_mla_prefill_quant_output.py
vllm/model_executor/warmup/fa4_cutedsl_warmup.py
vllm/model_executor/warmup/kernel_warmup.py
vllm/models/inkling/nvidia/attention.py
vllm/models/inkling/nvidia/ops/__init__.py
vllm/models/inkling/nvidia/ops/fa4_rel_attention.py
vllm/v1/attention/backends/mla/prefill/flash_attn.py

💤 Files with no reviewable changes (2)

vllm/model_executor/warmup/kernel_warmup.py
vllm/model_executor/warmup/fa4_cutedsl_warmup.py

🚧 Files skipped from review as they are similar to previous changes (6)

vllm/models/inkling/nvidia/ops/fa4_rel_attention.py
tests/models/inkling/test_fa4_rel_attention.py
vllm/models/inkling/nvidia/ops/init.py
tests/v1/attention/test_mla_prefill_quant_output.py
vllm/v1/attention/backends/mla/prefill/flash_attn.py
vllm/models/inkling/nvidia/attention.py

Included review availability: Your plan provides up to 10 included reviews per hour; 9 remain after this review.

📝 Summary

Summary by CodeRabbit

Performance
- FA4 MLA prefill and Inkling relative-attention kernels now register JIT warmup directly when enabled, helping reduce initial execution overhead.
Refactor
- Consolidated FA4 kernel warmup behavior within the relevant attention components.
- Removed the standalone FA4 warmup pathway without changing supported attention functionality.
Tests
- Added coverage verifying that warmup registration occurs only for FA4 when JIT warmup is enabled.
- Updated relative-attention test references to match the current kernel naming.

Walkthrough

The change makes FA4 kernel singletons private, registers Inkling and MLA JIT warmup during attention construction, removes the standalone FA4 warmup module, and updates related tests.

Changes

FA4 kernel updates

Layer / File(s)	Summary
Private kernel exports and references `vllm/models/inkling/nvidia/ops/__init__.py`, `vllm/models/inkling/nvidia/ops/fa4_rel_attention.py`, `vllm/models/inkling/nvidia/attention.py`, `vllm/v1/attention/backends/mla/prefill/flash_attn.py`, `tests/models/inkling/test_fa4_rel_attention.py`	The Inkling and MLA FA4 kernel singletons use private names. Imports, calls, exports, and test references use the renamed identifiers.
Inkling warmup integration `vllm/models/inkling/nvidia/attention.py`, `vllm/model_executor/warmup/fa4_cutedsl_warmup.py`	Inkling attention registers the private FA4 relative-attention kernel for JIT warmup when enabled. The standalone FA4 warmup module is removed.
FA4 MLA prefill warmup integration `vllm/v1/attention/backends/mla/prefill/flash_attn.py`, `tests/v1/attention/test_mla_prefill_quant_output.py`	The backend registers MLA prefill warmup only when JIT warmup is enabled and FlashAttention version 4 is active. A parametrized test validates these conditions.

Estimated code review effort: 3 (Moderate) | ~20 minutes

Merge Risk: ⚪ Minimal · up to 33e62

This change moves FA4 warmup registration into the attention paths and removes legacy orchestration. No concrete merge-blocking risk remains in the supplied evidence.

Sequence Diagram(s)

sequenceDiagram
  participant InklingAttention
  participant KernelConfig
  participant InklingKernel
  InklingAttention->>KernelConfig: read enable_jit_warmup
  KernelConfig-->>InklingAttention: enabled flag
  InklingAttention->>InklingKernel: register_warmup()

sequenceDiagram
  participant FlashAttnPrefillBackend
  participant KernelConfig
  participant FA4MLAPrefillKernel
  FlashAttnPrefillBackend->>KernelConfig: read enable_jit_warmup
  KernelConfig-->>FlashAttnPrefillBackend: enabled flag
  FlashAttnPrefillBackend->>FA4MLAPrefillKernel: register_warmup() when version is 4

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 18.18% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 11 functions across 6 files.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly identifies the main change: migrating FA4 MLA and shared CuTeDSL kernel warmup paths.
Description check	✅ Passed	The description directly explains the FA4 MLA and Inkling warmup migration, removal of legacy orchestration, and related tests.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

_{Comment @coderabbitai help to get the list of available commands.}

Signed-off-by: LopezCastroRoberto <rocastro@redhat.com>

LopezCastroRoberto · 2026-09-04T11:35:32Z

/ci run

github-actions · 2026-09-04T11:35:51Z

✅ Triggered Buildkite CI #87242 for commit 30c7286384a3.

Signed-off-by: LopezCastroRoberto <rocastro@redhat.com>

LopezCastroRoberto · 2026-09-04T12:00:01Z

/ci cancel

LopezCastroRoberto · 2026-09-04T12:00:37Z

/ci run

github-actions · 2026-09-04T12:01:06Z

✅ Triggered Buildkite CI #87246 for commit 3581dc1710de.

LucasWilkinson · 2026-09-04T15:58:40Z

/ci retry

github-actions · 2026-09-04T15:58:57Z

✅ Queued 7 failed job(s) for retry in Buildkite CI #87246.

LopezCastroRoberto · 2026-09-04T16:57:22Z

/ci retry

github-actions · 2026-09-04T16:57:53Z

✅ Queued 2 failed job(s) for retry in Buildkite CI #87246.

LopezCastroRoberto · 2026-09-04T17:51:56Z

/ci run

github-actions · 2026-09-04T17:52:15Z

✅ Triggered Buildkite CI #87296 for commit ae8142eaf56f.

coderabbitai · 2026-09-07T04:56:59Z

Note

GitHub couldn't provide a complete incremental comparison for this pull request, so CodeRabbit is performing a full review instead. This review may take a little longer.

LopezCastroRoberto · 2026-09-07T08:21:34Z

/ci run

github-actions · 2026-09-07T08:21:55Z

✅ Triggered Buildkite CI #87515 for commit 33e6239d3aeb.

…project#53565) Signed-off-by: LopezCastroRoberto <rocastro@redhat.com> Co-authored-by: Codex <codex@openai.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> Signed-off-by: Jyotirmoy Roy <jyotirmoyroy649@gmail.com>

Resolve conflict in kernel_warmup.py: main (vllm-project#53565) migrated the FA4 CuTeDSL warmup into the JIT warmup registry, so the fa4_cutedsl_warmup call site this PR gated on has_cutedsl() no longer exists. FA4 warmup is now registered only by the FA4 MLA prefill backend, which XPU never selects, so the XPU-safety intent is preserved. Dropped the obsolete call and its now-unused has_cutedsl import; kept main's watermark_sample_warmup. The is_xpu() gate on the sparse-MLA warmups is unaffected. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

…project#53565) Signed-off-by: LopezCastroRoberto <rocastro@redhat.com> Co-authored-by: Codex <codex@openai.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>

feat(warmup): migrate FA4 MLA and shared CuTeDSL kernels

efbd12b

Co-authored-by: Codex <codex@openai.com> Signed-off-by: LopezCastroRoberto <rocastro@redhat.com>

LopezCastroRoberto requested review from AndreasKaratzas, DarkLight1337, LucasWilkinson, MatthewBonanni, pavanimajety and ywang96 as code owners August 24, 2026 10:39

claude Bot reviewed Aug 24, 2026

View reviewed changes

mergify Bot added inkling DSv4 labels Aug 24, 2026

This was referenced Aug 24, 2026

Zero JIT compilation during runtime #49349

Open

[DO NOT MERGE][Perf][Feat] DSv4 de-JITification #49627

Open

mergify Bot added the needs-rebase label Sep 1, 2026

Merge remote-tracking branch 'refs/remotes/vllm-upstream/main' into w…

2ece2e1

…armup_migration/dsv4_7 Signed-off-by: LopezCastroRoberto <rocastro@redhat.com> # Conflicts: # vllm/model_executor/layers/attention/mla_attention.py # vllm/model_executor/warmup/kernel_warmup.py

LucasWilkinson approved these changes Sep 3, 2026

View reviewed changes

LucasWilkinson added the ready ONLY add when PR is ready to merge/full CI is needed label Sep 3, 2026

Merge branch 'main' into warmup_migration/dsv4_7

0263851

mergify Bot removed the needs-rebase label Sep 4, 2026

LopezCastroRoberto and others added 2 commits September 4, 2026 10:23

fix(warmup): register FA4 MLA prefill provider

74c23b3

Signed-off-by: LopezCastroRoberto <rocastro@redhat.com>

Merge branch 'main' into warmup_migration/dsv4_7

30c7286

LopezCastroRoberto and others added 2 commits September 4, 2026 11:58

fix CI

3b85c79

Signed-off-by: LopezCastroRoberto <rocastro@redhat.com>

Merge branch 'main' into warmup_migration/dsv4_7

3581dc1

Merge branch 'main' into warmup_migration/dsv4_7

ae8142e

Merge branch 'main' into warmup_migration/dsv4_7

9ccd8bf

Merge branch 'main' into warmup_migration/dsv4_7

33e6239

LucasWilkinson enabled auto-merge (squash) September 7, 2026 13:36

LucasWilkinson merged commit 42801b3 into vllm-project:main Sep 7, 2026
132 checks passed

drkm-eng mentioned this pull request Sep 7, 2026

[Feature][Model][Kernel] Inkling: fp8 e4m3 KV cache with per-tensor scales on SM100 #54705

Draft

Conversation

LopezCastroRoberto commented Aug 24, 2026

Description

What Changed

Uh oh!

claude Bot left a comment

Choose a reason for hiding this comment

Claude Code Review

Uh oh!

mergify Bot commented Sep 1, 2026

Uh oh!

LucasWilkinson left a comment

Choose a reason for hiding this comment

Uh oh!

github-actions Bot commented Sep 3, 2026

Uh oh!

coderabbitai Bot commented Sep 4, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary by CodeRabbit

Walkthrough

Changes

Sequence Diagram(s)

❌ Failed checks (1 warning)

Uh oh!

LopezCastroRoberto commented Sep 4, 2026

Uh oh!

github-actions Bot commented Sep 4, 2026

Uh oh!

LopezCastroRoberto commented Sep 4, 2026

Uh oh!

LopezCastroRoberto commented Sep 4, 2026

Uh oh!

github-actions Bot commented Sep 4, 2026

Uh oh!

LucasWilkinson commented Sep 4, 2026

Uh oh!

github-actions Bot commented Sep 4, 2026

Uh oh!

LopezCastroRoberto commented Sep 4, 2026

Uh oh!

github-actions Bot commented Sep 4, 2026

Uh oh!

LopezCastroRoberto commented Sep 4, 2026

Uh oh!

github-actions Bot commented Sep 4, 2026

Uh oh!

coderabbitai Bot commented Sep 7, 2026

Uh oh!

LopezCastroRoberto commented Sep 7, 2026

Uh oh!

github-actions Bot commented Sep 7, 2026

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

coderabbitai Bot commented Sep 4, 2026 •

edited

Loading