Repository navigation
[Feature][P/D] Support FullAttn use TP parallel when DCP>1 and num_kv_heads > 1 - #17524
Conversation
…s_tp>1 Signed-off-by: nwpu-zxr <zhouxuerong2@huawei.com>
Summary of ChangesHello, I'm Gemini Code Assist1! 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 enhances the Mooncake KV transfer mechanism by allowing FullAttention layers to leverage Tensor Parallelism (TP) sharding for KV cache transmission even when Data Cache Parallelism (DCP) is active. By dynamically adjusting DCP sizes based on head counts and PCP (Pipeline Cache Parallelism) status, the system can now maintain efficient TP-based KV distribution in scenarios where it was previously restricted. This change ensures better compatibility and performance for models like Kimi K3 when using specific draft model configurations. Highlights
New Features🧠 You can now enable Memory (public preview) 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 AssistThe 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
Customization To customize the 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 Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counterproductive. 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. Footnotes
|
|
👋 Hi! Thank you for contributing to the vLLM Ascend project. The following points will speed up your PR merge:
If CI fails, you can run linting and testing checks locally according Contributing and Testing. Tip 💡 Consider Linking a Related Issue or RFCYour PR title contains the [Feature] tag, indicating a bug fix or new feature. Linking a related issue or RFC in the PR description is strongly encouraged — it gives reviewers helpful context and speeds up the review. You can use any of these keywords:
🙏 Thanks for helping us keep the project well-organized! |
![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
Suggested PR Title:
[Attention][Feature] Support TP head shards for multi-head FA with global DCP when PCP is disabledSuggested PR Summary:
### What this PR does / why we need it?
This pull request implements support for using TP head shards for multi-head Full Attention (FA) when PCP is disabled. It introduces the `_get_full_attention_dcp_sizes` helper method to determine local and remote DCP sizes and integrates it into the rank group retrieval, block bucket building, and transfer address generation processes. Additionally, a redundant check for `MLAAttentionSpec` in `_get_full_attention_dcp_sizes` was identified and should be removed to clean up dead code.
### Does this PR introduce _any_ user-facing change?
No.
### How was this patch tested?
The changes are covered by new unit tests added in `tests/ut/kv_offload/mooncake_v2/test_pull_worker.py`, specifically verifying address generation, transfer bucket building, and TP head shard usage with global DCP.
Signed-off-by: nwpu-zxr <zhouxuerong2@huawei.com>
…_heads > 1 (vllm-project#17524) ### What this PR does / why we need it? Support FullAttention use TP parallel when DCP>1 and `num_kv_heads > 1`. For example, Kimi K3 with MHA dspark draft model, if DCP is enabled, the MLA layer will use DCP to extend block_size and reuse the Redundant kv_cache across DP, but Mamba and MHA dspark drafter will still use TP to split their caches. This PR supports GQA/MHA layers to still use TP sharding for KV transmission even when DCP > 1, as long as `local_num_kv_heads > 1` and `remote_num_kv_heads > 1`. ### Does this PR introduce _any_ user-facing change? No. ### How was this patch tested? By CI. - vLLM main: vllm-project/vllm@ced6857 --------- Signed-off-by: nwpu-zxr <zhouxuerong2@huawei.com>
…ention_dcp_sizes and _compute_group_block_ids local_dcp params The previous manual merge of upstream vllm-project#17524 (3a52c55) carried only the three call sites into pull_worker.py but dropped: - the _get_full_attention_dcp_sizes helper definition (18 lines) - the local_dcp_size parameter and local_dcp_rank binding of _compute_group_block_ids That left the tree in a guaranteed NameError/TypeError state on the first Mooncake pull of a FullAttention layer. This commit restores both pieces verbatim from upstream 3a52c55. Verified: AST parse OK; upstream-added substantive lines present 39/39 (100%); UT tests/ut/kv_offload/mooncake_v2/test_pull_worker.py already in tree covers _compute_group_block_ids signatures. Signed-off-by: xiangyongzh <ascend-operator@localhost>
…-indexer → glm5.3-flash-0.29) 含 45ee010 因果可见性修复(LoadRowMeta 改读绝对 positions, 修复 prefix-cache 命中 prefill/每步 decode 的池可见性坍缩, gsm8k 6.00→98.00 与 triton 持平) 及 PR#17542/vllm-project#17524 backport、c5182d3d4 padded-rows 掩码、2fbcf3eea dispatcher 契约测试。 029-pd-v2 转正回归证据 (experiments/029-pd-v2-split-gate-20260930): - gsm8k-20=95.00 / gsm8k-100=98.00 (=triton) - TTFT 64k 9.65s vs triton 12.45s (-22.5%); 192k 28.50s vs 45.91s (-37.9%) - bundle 31ae34e694bd 内容物与源仓一致 (.o b79cb7e6… = dispatcher 281b30a9… = HEAD) Signed-off-by: xiangyongzh <ascend-operator@localhost>
合并 feat/glm5-kpool-ascendc-indexer 44 commits + 队友线 2 commits: - glm5_kpool_split AscendC 内核转正(split 达成: 64k TTFT -37.5% / 192k -59% vs triton) - glm5_kpool_indexer: dispatcher PR#17542 契约补全 + padded rows mask + absolute positions - 上游 backport: PR#17115/vllm-project#17613/vllm-project#17542/vllm-project#17524(P/D FullAttn TP parallel / mooncake) - Restore causal_conv1d_update_npu Triton kernel (dropped by vllm-project#14620) - memcache: filter cacheable KV groups + rewarm SSD on load - ascend_store: port vllm-project#16945 empty RoPE view regression tests 原始历史: backup/glm5.3-flash-0.29-pre-squash (bd6c295)
- glm5_kpool_split AscendC 内核(64k TTFT -37.5% / 192k -59% vs triton)+ dispatcher PR#17542 契约 - 上游 backport: PR#17115/vllm-project#17613/vllm-project#17542/vllm-project#17524 - causal_conv1d_update_npu Triton 恢复 + memcache rewarm + vllm-project#16945 回归测试 squash 自 47 commits; 原始历史: backup/glm5.3-flash-0.29-pre-squash
The glm5.3-flash-0.29 merge of vllm-project#17524 lost method definitions in conflict resolution: re-add _get_full_attention_dcp_sizes (upstream 3a52c55) and the selection_index/local_dcp_size parameters of _compute_group_block_ids; make remote_pcp_size lookups tolerate missing metadata fields. Without this the first decode-side request fails with AttributeError. Verified on herd v029-base 39P/40D.
The glm5.3-flash-0.29 merge of vllm-project#17524 lost method definitions in conflict resolution: re-add _get_full_attention_dcp_sizes (upstream 3a52c55) and the selection_index/local_dcp_size parameters of _compute_group_block_ids; make remote_pcp_size lookups tolerate missing metadata fields. Without this the first decode-side request fails with AttributeError. Verified on herd v029-base 39P/40D.
3 participants
What this PR does / why we need it?
Support FullAttention use TP parallel when DCP>1 and
num_kv_heads > 1. For example, Kimi K3 with MHA dspark draft model, if DCP is enabled, the MLA layer will use DCP to extend block_size and reuse the Redundant kv_cache across DP, but Mamba and MHA dspark drafter will still use TP to split their caches. This PR supports GQA/MHA layers to still use TP sharding for KV transmission even when DCP > 1, as long aslocal_num_kv_heads > 1andremote_num_kv_heads > 1.Does this PR introduce any user-facing change?
No.
How was this patch tested?
By CI.