[Attention][MLA] Per-request scheduling for MLA chunked context - #50613
Merged
MatthewBonanni merged 14 commits intoAug 6, 2026
Merged
MatthewBonanni merged 14 commits into
MatthewBonanni merged 14 commits into
Conversation
Replace the batch-column chunked-context schedule with a flat list of per-request chunks (RFC vllm-project#50497). Previously the workspace was split evenly across every prefill with context and each iteration processed the same context window of every prefill, so a heterogeneous batch wasted most of the workspace per launch and paid attention/merge over the whole prefill batch on every iteration. The workspace also had to be floored at `max_num_seqs * block_size`, which overrides the deliberate 64k cap on models whose page size is inflated to cover a hybrid state page. Chunks now greedily pack whole requests and split a request only when its context alone exceeds the workspace, so a chunk covers a contiguous run of prefills and charges attention, up-projection and merging only to those requests. Because chunks stay in request order, only a chunk's first request can continue an earlier chunk: accumulating the context partial is one request-slice merge plus one bulk write. - `mask_empty_context` and `has_empty_context` are deleted; a chunk cannot contain an empty context span by construction. - The workspace floor drops from `max_num_seqs * block_size` to one page, making workspace memory independent of `max_num_seqs`. - `prefill_tokens_with_context` is now the query end of the last prefill with context, and the builder reports the token ranges of context-free prefills inside it so the partial is neutralized there. This fixes a latent mismatch when the prefills with context were not a contiguous prefix of the batch. - `run_prefill_context_chunk` takes the `ContextChunk` instead of a chunk index, since every sequence-length field is now chunk-local. Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>
Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>
Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>
Co-authored-by: OpenAI Codex <noreply@openai.com> Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>
MatthewBonanni
marked this pull request as ready for review
August 5, 2026 19:18
MatthewBonanni
requested review from
AndreasKaratzas,
LucasWilkinson,
WoosukKwon,
mgoin,
pavanimajety,
tjtanaa,
tlrmchlsmth,
yewentao256 and
zyongye
as code owners
August 5, 2026 19:18
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>
LucasWilkinson
approved these changes
Aug 6, 2026
LucasWilkinson
left a comment
LucasWilkinson
left a comment
Collaborator
There was a problem hiding this comment.
LGTM thank you!!!
LucasWilkinson
added
the
ready
Collaborator
|
/run ci |
Member
Author
|
/ci run |
|
✅ Triggered Buildkite CI #82693 for commit |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Implements #50497. MLA prefill chunks are fit into the available workspace rather than forced to be the same size. This can reduce the overall number of chunks and improve prefill latency.
Validation
Correctness
pytest tests/v1/attention/test_mla_context_chunks.py -qpassesGSM8K with
DeepSeek-V2-Lite-Chat, TP=4, DCP=2, FlashMLA, prefix caching:mainPerformance
This is a worst-case scenario but it highlights the optimization. We run a batch of size 32: 1 request has 60k context and rest have have 16. (60,496 total context tokens)
On main, the workspace (65,536 tok) gets divided evenly among the requests, so it uses 30 chunks for the batch. On PR, all context fits in the workspace, so it's all processed in a single chunk.
Run on H100:
main:PR:
AI assistance was used to develop and validate this change.