[Attention][DBO] Add support for "splitting" the CommonAttentionMetadata #21153
Conversation
Signed-off-by: Sage Moore <[email protected]>
Signed-off-by: Sage Moore <[email protected]>
Signed-off-by: Sage Moore <[email protected]>
Signed-off-by: Sage Moore <[email protected]>
Signed-off-by: Sage Moore <[email protected]>
Signed-off-by: Sage Moore <[email protected]>
|
👋 Hi! Thank you for contributing to the vLLM project. 💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in #pr-reviews, coordinate on features in #feat- channels, or join special interest groups in #sig- channels. Just a reminder: PRs would not trigger full CI run by default. Instead, it would only run Once the PR is approved and ready to go, your PR reviewer(s) can run CI to test the changes comprehensively before merging. To run CI, PR reviewers can either: Add 🚀 |
Signed-off-by: Sage Moore <[email protected]>
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces a mechanism to split CommonAttentionMetadata, which is a prerequisite for dual batch overlap. The changes are well-contained and include new unit tests. I've identified a critical issue where the code could crash due to an unhandled edge case when creating a metadata slice for zero requests. A fix is suggested to prevent this runtime error.
Signed-off-by: Sage Moore <[email protected]>
Signed-off-by: Sage Moore <[email protected]>
Signed-off-by: Sage Moore <[email protected]>
Signed-off-by: Sage Moore <[email protected]>
Signed-off-by: Sage Moore <[email protected]>
Signed-off-by: Sage Moore <[email protected]>
Signed-off-by: Sage Moore <[email protected]>
SageMoore
requested review from
WoosukKwon,
alexm-redhat,
comaniac,
njhill,
robertgshaw2-redhat and
ywang96
as code owners
SageMoore
changed the title
Signed-off-by: Sage Moore <[email protected]>
github-actions
bot
added
the
ready
|
Can you merge from main to see if the CI failures are fixed? |
…ata (vllm-project#21153) Signed-off-by: Sage Moore <[email protected]>
…ata (vllm-project#21153) Signed-off-by: Sage Moore <[email protected]> Signed-off-by: Jinzhen Lin <[email protected]>
…ata (vllm-project#21153) Signed-off-by: Sage Moore <[email protected]> Signed-off-by: Noam Gat <[email protected]>
…ata (vllm-project#21153) Signed-off-by: Sage Moore <[email protected]> Signed-off-by: Paul Pak <[email protected]>
…ata (vllm-project#21153) Signed-off-by: Sage Moore <[email protected]> Signed-off-by: Diego-Castan <[email protected]>
…ata (vllm-project#21153) Signed-off-by: Sage Moore <[email protected]>
…ata (vllm-project#21153) Signed-off-by: Sage Moore <[email protected]>
4 participants
Purpose
Note: the infrastructure in this PR is currently unused outside of unit tests.
The purpose of this PR is to add a "splitting" mechanism to the
CommonAttentionMetadataclass. This PR is a prerequisite for Dual Batch Overlap (#20448 ) support in vllm. Splitting, in this context, means slicing a batch of requests along some mid point and generating two newCommonAttentionMetadatainstances. One for each "micro" batch.Test Plan
I added a number of unit tests to
test_attention_splitting.py. These tests are largely focused on decode-only batches for now since that's all Dual Batch Overlap will initially support, but this infrastructure should support prefills and mixed batches without any issue.