Skip to content

[2/N][Attention] Enable masked MHA for sparse MLA prefills - #48770

Merged
vllm-bot merged 28 commits into
vllm-project:mainfrom
MatthewBonanni:fa4_masked_mha_remainder
Jul 31, 2026
Merged

[2/N][Attention] Enable masked MHA for sparse MLA prefills#48770
vllm-bot merged 28 commits into
vllm-project:mainfrom
MatthewBonanni:fa4_masked_mha_remainder

Conversation

@MatthewBonanni

@MatthewBonanni MatthewBonanni commented Jul 15, 2026 •
edited
Loading

Copy link
Copy Markdown
Member

Depends on vllm-project/flash-attention#155

Purpose

For pure prefills, there is a threshold of sequence length below which a masked MHA pathway is faster than sparse MQA. The DeepSeek V3.2 paper mentions this:

Note that for short-sequence prefilling, we specially implement a masked MHA mode to
simulate DSA, which can achieve higher efficiency under short-context conditions.

#47327 implemented a dense MHA shortcut for when seq_len < topk. This PR extends that optimization by implementing the masked MHA pathway.

Test Plan

Test Result

mha_vs_mqa_dp_heads128 mha_vs_mqa_tp8_heads16 speedup_heatmap_tp1_tp2_tp4_tp8
Essential Elements of an Effective PR Description Checklist
  • The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
  • The test plan, such as providing test command.
  • The test results, such as pasting the results comparison before and after, or e2e results
  • (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.
  • (Optional) Release notes update. If your change is user facing, please update the release notes draft in the Google Doc.

@MatthewBonanni @codex
[Attention] Add masked MHA path for sparse MLA
2721e26 
Build on the dense-MHA shortcut from vllm-project#47327 and route eligible pure prefills through FA4 masked attention.

Co-authored-by: OpenAI Codex <noreply@openai.com>
Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>
@mergify mergify Bot added ci/build performance Performance-related issues v1 labels Jul 15, 2026
@MatthewBonanni MatthewBonanni mentioned this pull request Jul 15, 2026
Closed
5 tasks
@MatthewBonanni
Fix IMA
e2b53b6 
Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>
@MatthewBonanni
Clean up padding
3008fcc 
Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>
@qianlihuang qianlihuang mentioned this pull request Jul 16, 2026
Merged
4 tasks
@mergify

mergify Bot commented Jul 16, 2026

Copy link
Copy Markdown
Contributor

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

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

@mergify mergify Bot added the needs-rebase label Jul 16, 2026
@MatthewBonanni
Merge branch 'main' into fa4_masked_mha_remainder
b47873f 
Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>
@mergify mergify Bot removed the needs-rebase label Jul 16, 2026
@MatthewBonanni
Chunked context support
15d6999 
Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>
@MatthewBonanni
Reuse global topk mask
a3ca0e2 
Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>
@MatthewBonanni
Remove block sparse stuff
1a37b3c 
Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>
@MatthewBonanni
Add heuristic
c19aa02 
Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>
@mergify

mergify Bot commented Jul 16, 2026

Copy link
Copy Markdown
Contributor

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

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

@mergify mergify Bot added the needs-rebase label Jul 16, 2026
@MatthewBonanni
Merge branch 'main' into fa4_masked_mha_remainder
81ec3e4 
Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>
@mergify mergify Bot removed the needs-rebase label Jul 17, 2026
@MatthewBonanni
MatthewBonanni marked this pull request as ready for review July 17, 2026 16:30
claude[bot]
claude Bot reviewed Jul 17, 2026
View reviewed changes

@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.

@MatthewBonanni
Merge branch 'main' into fa4_masked_mha_remainder
08524a0 
Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>
@mergify

mergify Bot commented Jul 22, 2026

Copy link
Copy Markdown
Contributor

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

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

@mergify mergify Bot added the needs-rebase label Jul 22, 2026
@MatthewBonanni
Merge branch 'main' into fa4_masked_mha_remainder
6562310 
Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>
@mergify mergify Bot removed the needs-rebase label Jul 23, 2026
LucasWilkinson
LucasWilkinson approved these changes Jul 28, 2026
View reviewed changes

@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.

Overall looks pretty good to me; left some comments for cleanup/simplification

@MatthewBonanni
Merge branch 'main' into fa4_masked_mha_remainder
f17b1a7 
Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>
@MatthewBonanni
Cleanup
7a0c45a 
Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>
@MatthewBonanni
Cleanup
c68f1a1 
Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>
@MatthewBonanni
Cleanup, break out helper
8e1241a 
Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>
@MatthewBonanni
Preallocate topk mask buffer
002ed43 
Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>
@mergify mergify Bot removed the needs-rebase label Jul 28, 2026
@MatthewBonanni MatthewBonanni added ready ONLY add when PR is ready to merge/full CI is needed needs-rebase labels Jul 29, 2026
@mergify mergify Bot removed the needs-rebase label Jul 29, 2026
@MatthewBonanni MatthewBonanni changed the title [2/N][Attention] Enable masked MHA for sparse MLA prefills [Do Not Merge][2/N][Attention] Enable masked MHA for sparse MLA prefills Jul 29, 2026
@MatthewBonanni

MatthewBonanni commented Jul 29, 2026

Copy link
Copy Markdown
Member Author

Marking as [Do Not Merge] because we still need to land FA once CI is green

@MatthewBonanni
Gate workspace allocation on masked mha availability
247048e 
Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>
@MatthewBonanni
Clean up gate and heuristic
c266472 
Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>
@MatthewBonanni
Fix tests
e8f3507 
Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>
@MatthewBonanni
update FA
34a709f
@MatthewBonanni
Zero in kernel
5722c83 
Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>
@MatthewBonanni
update FA
8d44a98
@MatthewBonanni
update FA after landing
469862e
@MatthewBonanni
Merge branch 'main' into fa4_masked_mha_remainder
6338c21
@MatthewBonanni
MatthewBonanni enabled auto-merge (squash) July 30, 2026 13:57
@MatthewBonanni MatthewBonanni changed the title [Do Not Merge][2/N][Attention] Enable masked MHA for sparse MLA prefills [2/N][Attention] Enable masked MHA for sparse MLA prefills Jul 30, 2026
@vllm-bot
vllm-bot merged commit 82ae416 into vllm-project:main Jul 31, 2026
247 of 249 checks passed
@yewentao256 yewentao256 mentioned this pull request Aug 3, 2026
Open
21 tasks
@MatthewBonanni MatthewBonanni mentioned this pull request Aug 25, 2026
Merged
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@claude claude[bot] claude[bot] left review comments
@mgoin mgoin mgoin approved these changes
@LucasWilkinson LucasWilkinson LucasWilkinson approved these changes
@tlrmchlsmth tlrmchlsmth Awaiting requested review from tlrmchlsmth tlrmchlsmth is a code owner
@Harry-Chen Harry-Chen Awaiting requested review from Harry-Chen Harry-Chen is a code owner

Assignees

No one assigned

Labels

ci/build performance Performance-related issues ready ONLY add when PR is ready to merge/full CI is needed v1

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

4 participants

@MatthewBonanni @mgoin @LucasWilkinson @vllm-bot