Skip to content

[V1] [Hybrid] Add new test to verify that hybrid views into KVCacheTensor are compatible #21300

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 2 commits into from
Jul 22, 2025

Conversation

tdoublep
Copy link
Member

@tdoublep tdoublep commented Jul 21, 2025

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.

Purpose

When running hybrid (attention/mamba) models in V1, the different types of layers share the same KVCacheTensor, by having different views into it. When implementing the support, care had to be taken to ensure that these two views are "compatible" in the sense that writing to a block in the mamba view will not corrupt the data in a different block in the attention view. I have found mistakes here are a bit tricky to catch since they don't necessarily show up until you start filling up all of the blocks. This PR adds a test to explicitly verify the compatibility of these views.

I would like to have this test in place before we verify the solution in #21123 to the large block size problem. It will also be helpful to have the test in place when we add the additional support for attention backends with (2, num_blocks) layout.

cc @heheda12345 @tlrmchlsmth

Test Plan

See diff

Test Result

Passing locally

(Optional) Documentation Update

Signed-off-by: Thomas Parnell <tpa@zurich.ibm.com>
Copy link

👋 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 fastcheck CI which starts running only a small and essential subset of CI tests to quickly catch errors. You can run other CI tests on top of those by going to your fastcheck build on Buildkite UI (linked in the PR checks section) and unblock them. If you do not have permission to unblock, ping simon-mo or khluu to add you in our Buildkite org.

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 ready label to the PR or enable auto-merge.

🚀

@mergify mergify bot added the v1 label Jul 21, 2025
Copy link
Contributor

@gemini-code-assist gemini-code-assist bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Code Review

This pull request adds a new test to verify the compatibility of hybrid attention/mamba views into the KVCacheTensor. To ensure full compatibility, the test should also verify that the mamba blocks retain their expected values after the attention blocks are written.

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>

Signed-off-by: Thomas Parnell <tpa@zurich.ibm.com>
@tdoublep tdoublep force-pushed the hybrid-mamba-test-layout branch from eaba4fc to a95f4b5 Compare July 21, 2025 10:50
Copy link
Collaborator

@heheda12345 heheda12345 left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thank you very much.

@heheda12345 heheda12345 enabled auto-merge (squash) July 21, 2025 23:51
@github-actions github-actions bot added the ready ONLY add when PR is ready to merge/full CI is needed label Jul 21, 2025
@vllm-bot vllm-bot merged commit 488d8a9 into vllm-project:main Jul 22, 2025
51 of 54 checks passed
yeqcharlotte pushed a commit to yeqcharlotte/vllm that referenced this pull request Jul 23, 2025
…nsor are compatible (vllm-project#21300)

Signed-off-by: Thomas Parnell <tpa@zurich.ibm.com>
zixi-qi pushed a commit to zixi-qi/vllm that referenced this pull request Jul 23, 2025
…nsor are compatible (vllm-project#21300)

Signed-off-by: Thomas Parnell <tpa@zurich.ibm.com>
Signed-off-by: qizixi <qizixi@meta.com>
LyrisZhong pushed a commit to LyrisZhong/vllm that referenced this pull request Jul 23, 2025
…nsor are compatible (vllm-project#21300)

Signed-off-by: Thomas Parnell <tpa@zurich.ibm.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
ready ONLY add when PR is ready to merge/full CI is needed v1
Projects
None yet
Development

Successfully merging this pull request may close these issues.

3 participants