Repository navigation
Conversation
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>
|
👋 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! |
|
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>
09e25ec to
be493cb
Compare
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 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
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
|
There was a problem hiding this comment.
Code Review
Suggested PR Title:
[Attention][Feature] Refactor SFA indexer metadata building and support DCP/DSA-CP transfersSuggested 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)
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)
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)
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_sizeexplicitly. 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, andtest_attn_utils_v2.py, using paired vLLMb2f685834a6456197e7033966fdef52a23f1abcdin the remote CPU-mock UT environment.252/252 real NPU transfer checks passed using Mooncake's
ascendbackend 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