[Future] [P/D] support hybrid attention for mooncake connector by liziyu179 · Pull Request #8850 · vllm-project/vllm-ascend

liziyu179 · 2026-04-30T08:32:52Z

What this PR does / why we need it?

support hybrid attention for mooncake connector

Does this PR introduce any user-facing change?

How was this patch tested?

by nightly

vLLM version: v0.20.2
vLLM main: vllm-project/vllm@1ac10f1

gemini-code-assist · 2026-04-30T08:33:01Z

Summary of Changes

Hello, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request introduces support for hybrid attention and Mamba models within the Mooncake connector. The changes enable the connector to handle diverse KV cache specifications and improve the robustness of data transfers by accounting for varying block lengths and specific model requirements. These updates ensure better compatibility with advanced model architectures while maintaining existing functionality.

Highlights

Hybrid Attention Support: Implemented hybrid attention (HMA) support for the Mooncake connector, allowing for more flexible KV cache management across different attention specifications.
Mamba Integration: Added support for Mamba models within the Mooncake connector, including specific metadata handling for SSM sizes and specialized transfer logic.
Connector Refactoring: Refactored the Mooncake connector to support multiple KV cache groups and updated internal data structures to handle block lengths per address.

🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	`/gemini review`	Performs a code review for the current pull request in its current state.
Pull Request Summary	`/gemini summary`	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	`/gemini help`	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for GitHub and other Google products, sign up here.

Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

github-actions · 2026-04-30T08:33:20Z

👋 Hi! Thank you for contributing to the vLLM Ascend project. The following points will speed up your PR merge:‌‌

A PR should do only one thing, smaller PRs enable faster reviews.
Every PR should include unit tests and end-to-end tests ‌to ensure it works and is not broken by other future PRs.
Write the commit message by fulfilling the PR description to help reviewer and future developers understand.

If CI fails, you can run linting and testing checks locally according Contributing and Testing.

gemini-code-assist

Code Review

This pull request implements support for hybrid attention (HMA), specifically integrating Mamba into the Mooncake connector for KV cache transfer. It updates metadata structures, introduces multi-group transfer logic, and refines cache registration. Feedback highlights critical bugs including a TypeError in Mamba block initialization, swapped source and destination addresses in transfer logic, and overly restrictive HMA checks that could crash non-Mamba hybrid models. Suggested PR Title: [Ops][Feature] Support hybrid attention for mooncake connector Suggested PR Summary: ### What this PR does / why we need it? This PR adds support for hybrid attention in the Mooncake connector for KV cache transfer. It updates metadata and implements a new transfer method for multiple KV cache groups. ### Does this PR introduce _any_ user-facing change? No. ### How was this patch tested? By nightly.

gemini-code-assist · 2026-04-30T08:35:40Z

+                grouped_remote_block_ids = [remote_block_ids[i][transfer_block_idx]]
+                grouped_local_block_ids = [local_block_ids[i][0]]


Suggested PR Title:

[Future][Ops][Feature] Support hybrid attention for mooncake connector

Suggested PR Summary:

### What this PR does / why we need it? This PR adds support for hybrid attention (e.g., Attention + Mamba) in the Mooncake connector for KV cache transfer. It updates the metadata exchange to include Mamba-specific information and implements a new transfer method that can handle multiple KV cache groups simultaneously. ### Does this PR introduce _any_ user-facing change? No. ### How was this patch tested? By nightly.

CRITICAL BUG: The initialization of grouped_remote_block_ids and grouped_local_block_ids for Mamba groups will cause a TypeError at runtime. They are assigned as lists of integers, but the subsequent logic at lines 545-547 expects them to be lists of lists (e.g., it accesses local_block_id[0] and calls len(local_block_id)).

Suggested change

grouped_remote_block_ids = [remote_block_ids[i][transfer_block_idx]]

grouped_local_block_ids = [local_block_ids[i][0]]

grouped_remote_block_ids = [[remote_block_ids[i][transfer_block_idx]]]

grouped_local_block_ids = [[local_block_ids[i][0]]]

gemini-code-assist · 2026-04-30T08:35:40Z

+            for k, (src_layer_base_addr, dst_layer_base_addr) in enumerate(
+                zip(local_kv_caches_base_addrs, remote_kv_caches_base_addrs)


CRITICAL BUG: The source and destination addresses for the KV transfer appear to be swapped. In the KVCacheRecvingThread (consumer side), the read operation should pull data from the remote producer into the local consumer. However, the code currently assigns local addresses to src and remote addresses to dst in the zip and subsequent calculations. This would mean the consumer is attempting to read from its own memory and write to the producer's memory.

gemini-code-assist · 2026-04-30T08:35:40Z

        for k, (src_layer_base_addr, dst_layer_base_addr) in enumerate(
            zip(local_kv_caches_base_addrs, remote_kv_caches_base_addrs)


CRITICAL BUG: Similar to _transfer_kv_cache_all_groups, the source and destination addresses in _transfer_kv_cache appear to be swapped. The consumer should be reading from the remote producer's memory into its local memory, but the zip assigns local addresses to the source side of the transfer.

gemini-code-assist · 2026-04-30T08:35:40Z

+        remote_block_ids = req_meta["remote_block_ids"][0]
+        local_block_ids = req_meta["local_block_ids"][0]


HIGH SEVERITY: The modification to use [0] for remote_block_ids and local_block_ids assumes there is only one KV cache group when is_hma_required is false. If multiple FullAttentionSpec groups exist, this will only transfer the first group, leading to incomplete KV cache transfer for the other groups.

gemini-code-assist · 2026-04-30T08:35:40Z

+        else:
+            raise TypeError("Mooncake connector does not support this type kv_cache now.")


HIGH SEVERITY: This implementation raises a TypeError if hybrid attention is required but Mamba is not present. This unnecessarily restricts hybrid attention support to only Mamba-based models and will crash for other hybrid configurations.

github-actions · 2026-05-08T13:40:41Z

This pull request has conflicts, please resolve those before we can evaluate the pull request.

Ensure consistent return statement formatting.

Removed license comment and multiple import statements.

Signed-off-by: liziyu <liziyu16@huawei.com>

liziyu179 requested review from LCAIZJ and MengqingCao as code owners April 30, 2026 08:32

gemini-code-assist Bot reviewed Apr 30, 2026

View reviewed changes

liziyu179 force-pushed the support_hybrid_atten_mooncake_connector branch from 5e527fb to 5d4c64b Compare May 8, 2026 08:32

github-actions Bot added the merge-conflicts label May 8, 2026

liziyu179 force-pushed the support_hybrid_atten_mooncake_connector branch from bd74e07 to 81da92b Compare May 9, 2026 09:57

liziyu179 and others added 11 commits May 10, 2026 23:15

Create tmp.log

d6539af

Merge branch 'vllm-project:main' into main

68bcd55

Update tmp.log

d6b42b1

Merge branch 'vllm-project:main' into main

981ff14

support qwen 3.5 tp

4d9f3c4

Ensure consistent return statement formatting.

Create 123.py

b420009

Update 123.py

0aac34e

Clean up imports and remove SPDX license comment

f40bd47

Removed license comment and multiple import statements.

mooncake connector support hybrid attention

f0d8a8a

Signed-off-by: liziyu <liziyu16@huawei.com>

support tp not equal

386a472

Signed-off-by: liziyu <liziyu16@huawei.com>

support tp not equal

fd6657d

Signed-off-by: liziyu <liziyu16@huawei.com>

liziyu179 force-pushed the support_hybrid_atten_mooncake_connector branch from 81da92b to fd6657d Compare May 21, 2026 02:57

github-actions Bot removed the merge-conflicts label May 21, 2026

refactor

ef13fe3

Signed-off-by: liziyu <liziyu16@huawei.com>

liziyu179 force-pushed the support_hybrid_atten_mooncake_connector branch 5 times, most recently from 5ded65a to f4cee00 Compare May 25, 2026 09:52

support mamba tp

e38e282

Signed-off-by: liziyu <liziyu16@huawei.com>

liziyu179 force-pushed the support_hybrid_atten_mooncake_connector branch from f4cee00 to e38e282 Compare May 25, 2026 13:02

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[Future] [P/D] support hybrid attention for mooncake connector#8850

[Future] [P/D] support hybrid attention for mooncake connector#8850
liziyu179 wants to merge 13 commits into
vllm-project:mainfrom
liziyu179:support_hybrid_atten_mooncake_connector

liziyu179 commented Apr 30, 2026 •

edited by github-actions Bot

Loading

Uh oh!

gemini-code-assist Bot commented Apr 30, 2026

Uh oh!

github-actions Bot commented Apr 30, 2026

Uh oh!

gemini-code-assist Bot left a comment

Uh oh!

gemini-code-assist Bot Apr 30, 2026

Uh oh!

gemini-code-assist Bot Apr 30, 2026

Uh oh!

gemini-code-assist Bot Apr 30, 2026

Uh oh!

gemini-code-assist Bot Apr 30, 2026

Uh oh!

gemini-code-assist Bot Apr 30, 2026

Uh oh!

github-actions Bot commented May 8, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

		grouped_remote_block_ids = [remote_block_ids[i][transfer_block_idx]]
		grouped_local_block_ids = [local_block_ids[i][0]]

		for k, (src_layer_base_addr, dst_layer_base_addr) in enumerate(
		zip(local_kv_caches_base_addrs, remote_kv_caches_base_addrs)

		remote_block_ids = req_meta["remote_block_ids"][0]
		local_block_ids = req_meta["local_block_ids"][0]

		else:
		raise TypeError("Mooncake connector does not support this type kv_cache now.")

Conversation

liziyu179 commented Apr 30, 2026 • edited by github-actions Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What this PR does / why we need it?

Does this PR introduce any user-facing change?

How was this patch tested?

Uh oh!

gemini-code-assist Bot commented Apr 30, 2026

Summary of Changes

Highlights

Footnotes

Uh oh!

github-actions Bot commented Apr 30, 2026

Uh oh!

gemini-code-assist Bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

gemini-code-assist Bot Apr 30, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Apr 30, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Apr 30, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Apr 30, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Apr 30, 2026

Choose a reason for hiding this comment

Uh oh!

github-actions Bot commented May 8, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

liziyu179 commented Apr 30, 2026 •

edited by github-actions Bot

Loading