[Attention][Spec Decode] FlashMLA spec decode support #26541

MatthewBonanni · 2025-10-09T22:29:20Z

Purpose

This PR implements speculative decoding support for the FlashMLA backend.

NOTE: the comment about the intermittent test failure was true prior to this PR, I just made a note of it.

cc @LucasWilkinson

Test Plan

pytest tests/v1/attention/test_mla_backends.py

Test Result

Passes

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.

gemini-code-assist

Code Review

This pull request introduces speculative decoding support for the FlashMLA backend, which is a significant feature enhancement. The changes are well-structured, particularly the introduction of the QueryLenSupport enum to clearly define the capabilities of different backends. The test cases have been updated appropriately to cover speculative decoding scenarios. My main feedback is a minor performance improvement in the test suite by moving a repeated import out of a loop.

tests/v1/attention/test_mla_backends.py

chatgpt-codex-connector

💡 Codex Review

Here are some automated review suggestions for this pull request.

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

Open a pull request for review
Mark a draft as ready
Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

vllm/v1/attention/backends/mla/common.py

Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>

LucasWilkinson

Nice! LGTM but lets have @benchislett look too

mergify · 2025-10-13T01:41:11Z

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

Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>

vllm/v1/attention/backends/mla/common.py

vllm/v1/attention/backends/mla/flashmla.py

benchislett

minor issue flagged. otherwise LGTM, please ping me when resolved

Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>

MatthewBonanni · 2025-10-14T17:01:44Z

@benchislett thanks for your review! I've addressed your comments

benchislett

LGTM, Thanks!

…6541) Signed-off-by: Matthew Bonanni <mbonanni@redhat.com> Signed-off-by: Jonah Bernard <jb2528@cornell.edu>

…6541) Signed-off-by: Matthew Bonanni <mbonanni@redhat.com> Signed-off-by: bbartels <benjamin@bartels.dev>

…6541) Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>

…6541) Signed-off-by: Matthew Bonanni <mbonanni@redhat.com> Signed-off-by: xuebwang-amd <xuebwang@amd.com>

MatthewBonanni requested a review from LucasWilkinson as a code owner October 9, 2025 22:29

mergify bot added the v1 label Oct 9, 2025

gemini-code-assist bot reviewed Oct 9, 2025

View reviewed changes

tests/v1/attention/test_mla_backends.py Show resolved Hide resolved

chatgpt-codex-connector bot reviewed Oct 9, 2025

View reviewed changes

vllm/v1/attention/backends/mla/common.py Show resolved Hide resolved

MatthewBonanni added 10 commits October 10, 2025 16:02

first pass

3a58cfb

Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>

fix test_mla_backends - vllm_config context and flashmla support

79bb4c0

Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>

fix allocated block count

a76ad79

Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>

fix index collision

3115b2d

Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>

add note about flakiness

e3edffe

Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>

mark flashattention mla as supporting spec decode

6ccfaa3

Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>

add speculative config

e8269da

Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>

implement QueryLenSupport, tests pass

b8303e9

Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>

update comment

fce2c63

Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>

address comment

b11a56e

Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>

MatthewBonanni force-pushed the flashmla_spec_decode branch from dae6cb0 to b11a56e Compare October 10, 2025 16:06

add query_len_support to flashinfer_mla

1eaab57

Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>

LucasWilkinson approved these changes Oct 10, 2025

View reviewed changes

LucasWilkinson enabled auto-merge (squash) October 13, 2025 01:40

LucasWilkinson disabled auto-merge October 13, 2025 01:40

github-actions bot added the ready ONLY add when PR is ready to merge/full CI is needed label Oct 13, 2025

mergify bot added the needs-rebase label Oct 13, 2025

Merge branch 'main' into flashmla_spec_decode

c9da1f0

Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>

mergify bot removed the needs-rebase label Oct 13, 2025

benchislett reviewed Oct 14, 2025

View reviewed changes

vllm/v1/attention/backends/mla/common.py Outdated Show resolved Hide resolved

benchislett reviewed Oct 14, 2025

View reviewed changes

vllm/v1/attention/backends/mla/flashmla.py Outdated Show resolved Hide resolved

benchislett requested changes Oct 14, 2025

View reviewed changes

benchislett mentioned this pull request Oct 14, 2025

[Bug]: Encountered AssertionError and precision issues when enabling MTP in deepseek v3.1 #26621

Open

1 task

MatthewBonanni added 2 commits October 14, 2025 12:59

remove assertion

1a7ca49

Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>

comments

1afd4e0

Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>

benchislett approved these changes Oct 14, 2025

View reviewed changes

benchislett enabled auto-merge (squash) October 14, 2025 17:10

benchislett merged commit 82af928 into vllm-project:main Oct 14, 2025
48 checks passed

MatthewBonanni deleted the flashmla_spec_decode branch October 14, 2025 19:42

Jonahcb pushed a commit to Jonahcb/vllm that referenced this pull request Oct 15, 2025

[Attention][Spec Decode] FlashMLA spec decode support (vllm-project#2…

a053c02

…6541) Signed-off-by: Matthew Bonanni <mbonanni@redhat.com> Signed-off-by: Jonah Bernard <jb2528@cornell.edu>

benchislett mentioned this pull request Oct 15, 2025

[Bugfix] [Performance]Better MTP Support when use flashmla #24045

Closed

bbartels pushed a commit to bbartels/vllm that referenced this pull request Oct 16, 2025

[Attention][Spec Decode] FlashMLA spec decode support (vllm-project#2…

6c130f8

…6541) Signed-off-by: Matthew Bonanni <mbonanni@redhat.com> Signed-off-by: bbartels <benjamin@bartels.dev>

FENP mentioned this pull request Oct 16, 2025

[Bugfix][Attention][DCP] Set reorder_batch_threshold back to 1 when using FlashMLA and enable DCP #27023

Open

5 tasks

lywa1998 pushed a commit to lywa1998/vllm that referenced this pull request Oct 20, 2025

[Attention][Spec Decode] FlashMLA spec decode support (vllm-project#2…

2c6de2f

…6541) Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>

MatthewBonanni mentioned this pull request Oct 22, 2025

[Attention] Fix FlashMLA metadata builder arguments for q_len > 1 #27368

Merged

5 tasks

alhridoy pushed a commit to alhridoy/vllm that referenced this pull request Oct 24, 2025

[Attention][Spec Decode] FlashMLA spec decode support (vllm-project#2…

c6ffc1c

…6541) Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>

xuebwang-amd pushed a commit to xuebwang-amd/vllm that referenced this pull request Oct 24, 2025

[Attention][Spec Decode] FlashMLA spec decode support (vllm-project#2…

17580a9

…6541) Signed-off-by: Matthew Bonanni <mbonanni@redhat.com> Signed-off-by: xuebwang-amd <xuebwang@amd.com>

xuebwang-amd pushed a commit to xuebwang-amd/vllm that referenced this pull request Oct 24, 2025

[Attention][Spec Decode] FlashMLA spec decode support (vllm-project#2…

6872684

…6541) Signed-off-by: Matthew Bonanni <mbonanni@redhat.com> Signed-off-by: xuebwang-amd <xuebwang@amd.com>

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

Uh oh!

[Attention][Spec Decode] FlashMLA spec decode support #26541

[Attention][Spec Decode] FlashMLA spec decode support #26541

MatthewBonanni commented Oct 9, 2025 •

edited by github-actions bot

Loading

Uh oh!

gemini-code-assist bot left a comment

Uh oh!

Uh oh!

chatgpt-codex-connector bot left a comment

Uh oh!

Uh oh!

LucasWilkinson left a comment

Uh oh!

mergify bot commented Oct 13, 2025

Uh oh!

Uh oh!

Uh oh!

benchislett left a comment

Uh oh!

MatthewBonanni commented Oct 14, 2025

Uh oh!

benchislett left a comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Uh oh!

[Attention][Spec Decode] FlashMLA spec decode support #26541

[Attention][Spec Decode] FlashMLA spec decode support #26541

Conversation

MatthewBonanni commented Oct 9, 2025 • edited by github-actions bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Purpose

Test Plan

Test Result

Uh oh!

gemini-code-assist bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

Uh oh!

chatgpt-codex-connector bot left a comment

Choose a reason for hiding this comment

💡 Codex Review

Uh oh!

Uh oh!

LucasWilkinson left a comment

Choose a reason for hiding this comment

Uh oh!

mergify bot commented Oct 13, 2025

Uh oh!

Uh oh!

Uh oh!

benchislett left a comment

Choose a reason for hiding this comment

Uh oh!

MatthewBonanni commented Oct 14, 2025

Uh oh!

benchislett left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

MatthewBonanni commented Oct 9, 2025 •

edited by github-actions bot

Loading