[v0.26.0rc][BugFix][Mooncake] Fix DCP mapping and support P/D-only DCP in PD disaggregation (#16492) - #16844
[v0.26.0rc][BugFix][Mooncake] Fix DCP mapping and support P/D-only DCP in PD disaggregation (#16492)#16844lllyys wants to merge 2 commits into
Conversation
…vllm-project#16487) Fix shape mismatches when DCP is enabled with PD-disaggregated recomputation. Short recomputation requests, including last-token recomputation, are now correctly classified as decode requests in the attention, MLA, and SFA metadata builders. Cache DCP state during initialization and use the builder's configuration so metadata construction works outside the current-config context. No API or configuration changes. This fixes failures in DCP-enabled PD-disaggregated recomputation. - Added regression tests covering PD recomputation, mixed query lengths, and behavior when the override does not apply. - Added coverage for metadata construction without an active current-config context. - vLLM main: vllm-project/vllm@a97dacb --------- Signed-off-by: weiguihua2 <weiguihua2@huawei.com> (cherry picked from commit cab3719)
|
👋 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 [BugFix] 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! |
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 backports critical fixes and enhancements for the Mooncake connector to the v0.26.0rc release branch. It introduces support for P/D-only DCP disaggregation, improves block and port mapping for unequal DCP configurations, and ensures robust indexer cache transfers by refining replicate-K filtering. These changes enhance the reliability and flexibility of KV offloading in context-parallel setups. 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
|
![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 decode-only DCP and adapt CP layout for PCP shard axisSuggested PR Summary:
### What this PR does / why we need it?
This PR introduces support for decode-only DCP (Distributed Context Parallel) transfers and adapts the context parallel layout to handle configurations where PCP (Pipeline Context Parallel) remains a shard axis. Specifically, it:
- Implements `_get_decode_only_dcp_metadata` in `MooncakeConnectorWorker` to map global blocks to each rank during decode-only DCP.
- Updates metadata builders (`AscendAttentionDCPMetadataBuilder`, `AscendSFADCPMetadataBuilder`, and `AscendMLAMetadataBuilder`) to treat short extends as decodes when DCP and the PD decode recompute scheduler are enabled.
- Refines CP layout validation and parameter checks to support bidirectional divisibility of P/D CP sizes.
- Adds comprehensive unit tests for these changes.
Feedback:
- In `_get_decode_only_dcp_metadata`, a potential `IndexError` or incorrect block mapping may occur if `meta.local_full_block_ids` is `None` and `meta.num_computed_tokens > 0`. A safer fallback to `meta.local_block_ids` with proper index offset handling is recommended.
- In `start_load_kv`, a guard should be added to prevent a potential `KeyError` when accessing `self.kv_group2layeridx[pull.group_id]`.
### Does this PR introduce _any_ user-facing change?
No.
### How was this patch tested?
Tested with new and updated unit tests in `test_attention_cp.py`, `test_mla_v1.py`, `test_sfa_cp.py`, and `test_mooncake_connector.py`.| local_blocks = (meta.local_full_block_ids or meta.local_block_ids)[group_id] | ||
| remote_blocks = meta.remote_block_ids[group_id] | ||
| first_block = meta.num_computed_tokens // self.block_size | ||
| first_block += (self.dcp_rank - first_block) % self.dcp_size | ||
| # P owns the full sequence; D rank r owns r, r + DCP, ... . | ||
| global_blocks = range(first_block, min(meta.num_prompt_blocks, len(remote_blocks)), self.dcp_size) | ||
| scale = self._get_kernel_block_scale(layer_indices) | ||
| local_block_ids[block_id_idx] = self._expand_block_ids( | ||
| [local_blocks[block // self.dcp_size] for block in global_blocks], scale | ||
| ) |
There was a problem hiding this comment.
In _get_decode_only_dcp_metadata, if meta.local_full_block_ids is None (e.g., when prefix caching is disabled or not triggered), local_blocks falls back to meta.local_block_ids. Since meta.local_block_ids only contains the blocks for the external tokens, indexing it with block // self.dcp_size (which is a global block index) will lead to an IndexError or incorrect block mapping when meta.num_computed_tokens > 0. We should safely handle both full and incremental block ID lists and guard against potential None or out-of-bounds accesses.
if meta.local_full_block_ids and group_id < len(meta.local_full_block_ids):
local_blocks = meta.local_full_block_ids[group_id]
is_full = True
elif meta.local_block_ids and group_id < len(meta.local_block_ids):
local_blocks = meta.local_block_ids[group_id]
is_full = False
else:
raise ValueError(f"No local block IDs found for group {group_id}")
remote_blocks = meta.remote_block_ids[group_id] if meta.remote_block_ids and group_id < len(meta.remote_block_ids) else []
first_block = meta.num_computed_tokens // self.block_size
first_block += (self.dcp_rank - first_block) % self.dcp_size
# P owns the full sequence; D rank r owns r, r + DCP, ... .
global_blocks = range(first_block, min(meta.num_prompt_blocks, len(remote_blocks)), self.dcp_size)
scale = self._get_kernel_block_scale(layer_indices)
local_indices = [
block // self.dcp_size if is_full else (block - first_block) // self.dcp_size
for block in global_blocks
]
local_block_ids[block_id_idx] = self._expand_block_ids(
[local_blocks[idx] for idx in local_indices if idx < len(local_blocks)], scale
)| group_pulls = [ | ||
| pull | ||
| for pull in group_pulls | ||
| if self.kv_group2layeridx[pull.group_id][0]["kv_cache_spec_type"] | ||
| != "AscendSFAIndexerCacheSpec" | ||
| ] |
There was a problem hiding this comment.
Add a guard to ensure pull.group_id exists in self.kv_group2layeridx before accessing it, preventing potential KeyError exceptions during group pulls filtering.
group_pulls = [
pull
for pull in group_pulls
if pull.group_id in self.kv_group2layeridx
and self.kv_group2layeridx[pull.group_id][0]["kv_cache_spec_type"]
!= "AscendSFAIndexerCacheSpec"
]
…CP in PD disaggregation (vllm-project#16492) Mooncake's MRV2 context-parallel transfers can select incorrect source ports and KV blocks when DCP already spans PCP ranks. A decode worker with DCP also needs to select its local shards from an unsharded prefill cache, while the replicated indexer must be transferred only once. This patch corrects the port/block mapping and keeps decode-only DCP in an explicit branch: `self.dcp_size > 1 and meta.remote_dcp_size == 1` calls `_get_decode_only_dcp_metadata`. The existing non-DCP branch and general CP flow remain separate. The helper handles MLA blocks, replicated Indexer caches and KDA/Mamba state explicitly, and retains the equal P/D block-size requirement. This avoids combining the existing transfer paths into a new global block-mapping abstraction. Based on upstream commit cdad5a3. This PR contains the standalone Mooncake connector fix and its tests; it does not include the separate SFA operator or attention-backend experiments. Fixes MRV2 DCP port/block selection for P/D disaggregation, including one-sided DCP and PCP+DCP configurations. No new configuration or public interface. - Restored the connector and its regression tests to the exact contents of bfea331, before the combined block-mapping refactors. The restoration commit is 037f514. - Verified that the restored tree has no file differences from that historical commit; `git diff --check` passed. - Runtime tests have not been rerun for this restoration. Earlier regression results do not establish full-model or all-model coverage for the current PR. Fresh CI and P/D regression remain required. - No test logs, reports, experimental operator patches or model weights are committed. Paired vLLM source: vllm-project/vllm@84030bb - vLLM main: vllm-project/vllm@84030bb Signed-off-by: chengruiqi (C) <c00913489@china.huawei.com> Co-authored-by: chengruiqi (C) <c00913489@china.huawei.com> (cherry picked from commit 00b0b97) Signed-off-by: lllyys <221728326+lllyys@users.noreply.github.com>
f380184 to
9646377
Compare
3 participants
What this PR does / why we need it?
Cherry-pick of #16492 to
releases/v0.26.0rc. Stacked on #16833 — pleasemerge that one first (this PR shows both commits until then).
Fix DCP mapping in the Mooncake connector and add P/D-only DCP support in PD
disaggregation: a decode-only DCP fast path (
_get_decode_only_dcp_metadata)for P without CP → D with DCP, corrected unequal-DCP block/port mapping
(
min()-based shard mapping that also supports local DCP > remote DCP forpure-DCP setups), and replicate-K indexer pull filtering so non-replicate-K
ports do not overwrite the full indexer cache transfer.
Original PR: #16492 (main commit 00b0b97)
Backport scope note (PCP): on main, #16492 also reworks PCP+DCP semantics
on top of the MRV2 invariant that the DCP group spans PCP (PCP ranks hold
complete KV replicas, from #15809/#16080 — not on this branch). This branch
still treats PCP as a KV shard axis. This backport therefore keeps the
existing product-based CP geometry wherever PCP is involved — numerically
identical to main's DCP-only forms when PCP == 1 — and gates the new
decode-only DCP path to PCP == 1 on both sides. Result: pure-DCP setups get
all of #16492's fixes; existing PCP transfer behavior on this branch is
unchanged. The PCP+DCP portion of #16492 is intentionally not ported.
Backport conflicts and how they were resolved
All 6 conflict hunks are in
mooncake_connector.py. The root cause is thesemantic divergence above, so each resolution follows one rule: main's
DCP-only expressions are adopted where they are equivalent or purely additive
for PCP == 1, and this branch's product-based forms are kept wherever PCP > 1
behavior would change.
_get_kv_split_metadatagate: added main's decode-only-DCP earlyreturn, additionally gated to
self.pcp_size == 1 and meta.remote_pcp_size == 1; kept this branch's product fast-path gate(
remote_pcp * remote_dcp * pcp * dcp == 1) so remote-PCP transfers keeptaking the CP path that populates
remote_port_send_num(completionaccounting).
_get_cp_shard_pulls→_get_dcp_shard_pulls: thisbranch's body already matched main's, including the
remote_pcp_sizeparameter; only name and docstring changed.
pp_rankformula: took main's(port - base) // (ptp * remote_pcp),a strict generalization of this branch's
0 if remote_pcp > 1 else (port - base) // ptp(PCP and PP are mutuallyexclusive, so results agree on all legal layouts).
cp_transfergate,plus a decode-only-DCP exception that routes those ports to the hybrid
rank-table lookup — their ports are selected via the rank table, and the
shard-pulls builder would drop attention pulls for later PP stages.
add_requestloop: took main's side (replicate-K-filteredgroup_pulls,shard_idxindexing) and renamed this branch's loopvariable
pcp_dcp_rank→shard_idxto match.Other adaptations:
because this branch lacks them:
_get_selected_pcp_rank(origin [Feature][MRV2][P/D] Support PCP KV transfer #16080)and
_get_kernel_block_scale(origin [Feature] Mooncake_connector support MinmaxM3 #13457).only; whenever either side has PCP > 1 the parent's one-direction rule
(
remote % local == 0) is kept, because completion accounting for alarger local CP size is only implemented for pure DCP.
(remote PCP=1, remote DCP bounded by prefill TP, port layout
source_cp_rank = offset % remote_dcp); the SFA decode-only test isrestricted to PCP=1 (matching the gated path); the replicate-K ratio test
fixture sets
vllm_config(this branch readsvllm_config.model_configdirectly).
Independent review (Codex)
Reviewed pre-submission with Codex over three rounds; it executed extracted
connector methods with real CPU PyTorch against both the parent branch and
the backport:
broke remote-PCP transfers: the DCP-only fast path skipped
remote_port_send_numpopulation while the retained product gate indexesit (reproduced KeyError for P TP=1/PCP=2 → D TP=1, including a
previously-working one-block prompt), and unselected P-side PCP ranks
would stay pending until the abort timeout; the decode-only path's
pcp_offsetis never removed by this branch's table lookup (reproducedKeyError). This review drove the resolution model above.
decode-only-DCP ports through the shard-pulls builder, dropping attention
pulls for later pipeline stages (reproduced with P TP=2/PP=2 → D
TP=2/DCP=2, MLA+Mamba groups); the either-direction divisibility newly
accepted PCP configs whose completion accounting under-counts (sender
frees KV after the first pull); two added tests still encoded main's PCP
topology. All fixed as described above.
isolated execution (67 subcases, including all 18 ownership cases);
1,224 parent-vs-backport comparisons match, including 624 accepted and
456 rejected PCP cases plus 27 replicate-K comparisons; no remaining
inconsistency between path selection, port mapping, and completion
accounting.
Does this PR introduce any user-facing change?
No API changes. P/D-only DCP (P without CP → D with DCP) PD disaggregation
now works on this branch; previously-rejected unequal pure-DCP
configurations (local DCP > remote DCP) are now supported.
How was this patch tested?
Unit tests from the original PR adapted to this branch in
tests/ut/kv_offload/test_mooncake_connector.py(decode-only DCP cachegroup handling, KDA state/TP owners, SFA decode-only block mapping,
replicate-K ratio, DCP ownership coverage).
Method-level isolated execution of the added tests and 1,224 before/after
mapping/count comparisons against the parent branch (accepted and rejected
PCP cases match the parent exactly) — see the review notes above.
CI unit tests on this branch.
vLLM main: vllm-project/vllm@d02df74