Skip to content

[BugFix][Mooncake] Align MRV2 PCP and DCP KV transfers - #16454

Draft
recky-c wants to merge 4 commits into
vllm-project:mainfrom
recky-c:fix/mooncake-v1-on-pr16325
Draft

recky-c wants to merge 4 commits into
vllm-project:mainfrom
recky-c:fix/mooncake-v1-on-pr16325

Conversation

@recky-c

@recky-c recky-c commented Sep 13, 2026 •

Copy link
Copy Markdown
Contributor

What this PR does / why we need it?

Depends on #16325 and is based on its head b43f34281887e689ae978c54c0b7ff7f246f078d. Until that PR lands, this comparison also includes its commits.

Align Mooncake transfers with MRV2 DCP groups, which already span PCP ranks. Handle SFA prefill without DCP and decode with DCP in a separate entry branch, selecting an existing PCP replica and mapping global prompt blocks to each decode rank. Preserve the original SFA scope, equal-block-size requirement, and transfer interfaces.

Keep replicated indexer KV on one source port, account for PCP when decoding PP ports, and require callers to pass remote_pcp_size explicitly. Forward the MRV2 indexer graph-capture entry point to the implementation from #16325.

Does this PR introduce any user-facing change?

Fixes KV transfer for the affected SFA/MRV2 PCP and DCP configurations. No new configuration options.

How was this patch tested?

Regression results for be493cbba7a841a0d1cc0567b3b805275e762349 (2026-09-13):

  • Verified all 3,880 tracked source files in the isolated test directory against the commit snapshot.

  • 182 passed across test_mooncake_connector.py, test_indexer.py, test_sfa_cp.py, and test_attn_utils_v2.py, using paired vLLM b2f685834a6456197e7033966fdef52a23f1abcd in the remote CPU-mock UT environment.

  • 252/252 real NPU transfer checks passed using Mooncake's ascend backend on two NPUs. The harness extracts the current production entry/mapping methods and checks every destination element, including untouched regions. It covers DCP2/4/8 ranks, PCP1/2 source replicas, short/tail prompts, cached prefixes, and replicated indexer transfers.

  • This synthetic-buffer transfer check is not an end-to-end model or cross-host P/D serving test. Full model P/D regression for this revision remains pending. Test processes have exited and the two NPUs have been released.

  • The transfer-engine wheel has a separately reproduced import-only shutdown allocator failure in this environment. The harness performs assertions and memory unregistration before explicit process exit; normal library destruction is not claimed as validated.

  • Production Mooncake file SHA256: 37c45fb33ba57949236c082345cc3a1e583ff9d13820b6b50ceb334de2bc2a4c.

  • vLLM main: vllm-project/vllm@a97dacb

Construct replicated DCP block tables and slot mappings inside the indexer metadata builder, and pass PCP context explicitly without reading SFA metadata.

Signed-off-by: Biuapha <1731372716@qq.com>
Key mutable graph buffers by each common slot mapping and build RoPE, DSA-CP, PCP, and LI-C8 metadata inside the indexer builder. Preserve split indexer metadata groups through MTP and Eagle draft steps without advancing shared token state twice.

Signed-off-by: Biuapha <1731372716@qq.com>
Compile fused MoE with the decode-capture reduction contract and bypass graph calls whose runtime contract differs. Serialize non-latent shared-expert TP collectives before routed EP MoE so multistream execution cannot deadlock.

Signed-off-by: Biuapha <1731372716@qq.com>
@github-actions

Copy link
Copy Markdown
Contributor

👋 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.


Tip

💡 Consider Linking a Related Issue or RFC

Your 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:

  • Fixes #<issue_number>
  • Closes #<issue_number>
  • Resolves #<issue_number>
  • Refs #<rfc_or_issue_number> (for RFCs)

🙏 Thanks for helping us keep the project well-organized!

@github-actions

Copy link
Copy Markdown
Contributor

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

Preserve the original Mooncake adaptation, separate SFA decode-only DCP from the non-DCP entry, and select the existing PCP replica without restricting PCP size. Keep the original mapping and equal-block-size requirement.

Signed-off-by: recky-c <ruiqicheng510@gmail.com>
Signed-off-by: chengruiqi (C) <c00913489@china.huawei.com>
@recky-c
recky-c force-pushed the fix/mooncake-v1-on-pr16325 branch from 09e25ec to be493cb Compare September 13, 2026 10:19
@gemini-code-assist

Copy link
Copy Markdown
Contributor

Summary of Changes

Hello, 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 aligns Mooncake KV transfers with MRV2 DCP groups to ensure correct operation within PCP and DCP configurations. It introduces logic to handle SFA prefill and decode operations independently, ensuring proper replica selection and block mapping. Additionally, it improves support for replicated indexer KV transfers and enhances test coverage for complex SFA/MRV2 scenarios.

Highlights

  • MRV2 DCP Alignment: Aligns Mooncake KV transfers with MRV2 DCP groups, ensuring consistent handling across PCP ranks.
  • SFA Logic Refinement: Implements independent entry branches for SFA prefill and decode, improving replica selection and block mapping.
  • Replicated Indexer Support: Ensures replicated indexer KV transfers are correctly managed across PCP and DCP ports, preserving SFA scope.
  • Test Coverage: Updates regression tests to validate asymmetric DCP mappings and SFA decode-only configurations.
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 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 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 .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 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

  1. 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. ↩

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

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

Code Review

Suggested PR Title:

[Attention][Feature] Refactor SFA indexer metadata building and support DCP/DSA-CP transfers

Suggested PR Summary:

### What this PR does / why we need it?
This PR refactors the SFA indexer metadata building to make it independent of the SFA attention layer. It introduces support for Decode Context Parallel (DCP) and DSA-CP in the indexer, extracting common DCP utilities into a separate module. Additionally, it updates the Mooncake connector to handle SFA decode-only DCP transfers and aligns CP group metadata. For MoE, it adds detection for communication contract mismatches between runtime and compiled graphs, executing calls eagerly when a mismatch is found, and serializes shared-input gather before routed MoE in sequence parallel multistream mode. Speculative decoding proposers are also updated to support independent metadata groups for SFA and its split indexer.

Feedback on the changes:
1. In `mooncake_connector.py`, there is a potential `IndexError` or incorrect block mapping when `meta.local_full_block_ids` is `None` and `local_blocks` falls back to `meta.local_block_ids`. We should conditionally offset the index using `(block - first_block) // self.dcp_size` when `meta.local_full_block_ids` is not provided.
2. In `indexer.py`, we should explicitly cast the results of `torch.cumsum` and `torch.where` to `torch.int32` before assigning them to the pre-allocated `int32` buffers to prevent potential runtime type mismatch or promotion errors on Ascend NPU.

### Does this PR introduce _any_ user-facing change?
No.

### How was this patch tested?
CI passed with updated and newly added unit tests in `test_indexer.py`, `test_sfa_v1.py`, `test_mooncake_connector.py`, `test_fused_moe.py`, `test_mla.py`, `test_eagle_proposer.py`, `test_ascend_forward_context.py`, and `test_attn_utils_v2.py`.
I am having trouble creating individual review comments. Click here to see my feedback.

vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake_connector.py (2976-2986)

high

When meta.local_full_block_ids is None, local_blocks falls back to meta.local_block_ids, which only contains the newly allocated blocks (excluding prefix blocks). In this case, indexing with block // self.dcp_size will result in an IndexError or incorrect block mapping because the indices are not offset by first_block. We should conditionally offset the index using (block - first_block) // self.dcp_size when meta.local_full_block_ids is not provided.

                is_full = meta.local_full_block_ids is not None
                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)
                block_id_idx = group_idx if use_transfer_group_block_ids else group_id
                local_block_ids[block_id_idx] = self._expand_block_ids(
                    [local_blocks[block // self.dcp_size if is_full else (block - first_block) // self.dcp_size] for block in global_blocks], scale
                )

vllm_ascend/attention/indexer.py (844-853)

high

To prevent potential runtime type mismatch or promotion errors on Ascend NPU (which is strict about tensor dtypes), we should explicitly cast the results of torch.cumsum and torch.where to torch.int32 before assigning them to the pre-allocated int32 buffers actual_seq_lengths_query and actual_seq_lengths_key.

        actual_seq_lengths_query[:num_segs] = torch.cumsum(
            num_local_tokens.clamp(min=0),
            dim=0,
        ).to(dtype=torch.int32)
        offset = cum_query_lens - req_local_end
        actual_seq_lengths_key[:num_segs] = torch.where(
            num_local_tokens > 0,
            torch.clamp_min(seq_lens - offset, 0),
            0,
        ).to(dtype=torch.int32)

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants