[BugFix][Attention] Isolate PCP global RoPE runtime buffers - #16311
zhao-stack wants to merge 2 commits into
Conversation
Keep global KV-update RoPE results separate from rank-local attention metadata while preserving fixed-capacity buffers and stable graph replay addresses. Add numerical, multi-group, capacity and address-stability regression coverage. Signed-off-by: shenzhao <shenzhao9@huawei.com>
|
/e2e tests/e2e/pull_request/four_card/context_parallel/test_deepseek_v4.py::test_deepseek_v4_dsa_pcp_dspark |
|
👋 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 addresses a critical bug where RoPE runtime buffers were being aliased between PCP's global KV-update metadata and rank-local attention metadata. By isolating these buffers through caller-owned persistent maps, the change ensures that global position values do not corrupt local attention metadata, thereby fixing issues with generated tokens and speculative acceptance in affected PCP + DSpark batches. 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] Isolate RoPE runtime buffers for context parallel metadata builders to prevent aliasingSuggested PR Summary:
### What this PR does / why we need it?
This pull request introduces isolated RoPE runtime buffers for Context Parallel (PCP) metadata builders to prevent aliasing between global KV updates and rank-local attention positions during the same step. By passing a `rope_runtime_buffer` to `get_cos_and_sin_dsa`, caller-owned buffers remain stable, which is essential for graph replay.
### Does this PR introduce _any_ user-facing change?
No.
### How was this patch tested?
New unit tests were added in `tests/ut/attention/test_dsa_v1.py` and `tests/ut/ops/test_rope_dsv4.py` to verify that global and local RoPE buffers are kept separate and survive interleaved cache groups.Add concrete dictionary types required by the repository mypy checks. Production code and the original e2e test remain byte-identical to the fix commit. Signed-off-by: shenzhao <shenzhao9@huawei.com>
1 participant
What this PR does / why we need it?
Fix RoPE runtime-buffer aliasing between PCP's canonical global KV-update metadata and rank-local attention metadata.
With DSpark verification, short query rows can put both builds on the cached RoPE path even though their positions differ. Each PCP cache-group builder constructs global metadata, then local metadata. Local RoPE proxies are reused through the step's shared metadata, but a later cache group's global build writes the same process-wide runtime buffers again. The local proxy can then contain global-position values, corrupting target hidden states, generated tokens and speculative acceptance.
Each PCP global metadata builder now owns a persistent RoPE buffer map keyed by configuration and group. Each cos/sin pair is allocated once at the full registered capacity; subsequent calls use the existing in-place gather. Addresses remain stable from small warmup/capture batches through larger replay batches. Local/default, uncached prefill and speculative draft-buffer behavior is preserved; no global cache disabling, per-step cloning or version-specific branch is needed.
Buffer lifetime follows the global builder. Extra storage per requested configuration/group per global builder is
2 * max_num_batched_tokens * rotary_dim * element_size; the global Tensor-position path requests only the default group. Distinct PCP cache-group builders own separate maps. Non-PCP paths allocate no extra buffers.This independent fix branch starts from upstream main
f4b6bd05f05717e10e4521a78f4b3beba5bb5a99. It is separate from main2main PR #16009 and changes none of #16009, #15627, #15965 or #16216. The PR-wide diff contains only three production files and two UT files; pins, release tags, skips, thresholds, golden data and e2e sources are unchanged. The shared-buffer structure predates the historical upgrade experiment; there is no evidence that #16009 or6f3c34aintroduced it.Does this PR introduce any user-facing change?
Correct position encodings and generation/acceptance behavior for affected PCP + DSpark batches. No configuration or public API change.
How was this patch tested?
Current branch validation — actual pytest logs inspected:
b2f685834a6456197e7033966fdef52a23f1abcdv0.28.0/2cf0a6915ce544dc493a0990f2ea38d81601128aThe e2e run was triggered once with the repository's supported comment syntax:
The logs confirm the exact pytest node, model, version checkouts, full graph capture and final passing pytest summaries. Original model, prompts, parameters, prefix assertions, acceptance thresholds and 0.03 tolerance remain intact. No probes, controlled admission, extra test plugins or assertion changes were committed. The passing test does not print final aggregate acceptance rates; periodic metric-log rates are not substituted for the assertion's aggregate metrics.
Commit provenance: the e2e workflow pins fix commit
c14993202b111e5c496dc4ffa8ce946e226f4fe7and rebases it on fixed upstream snapshotfa9f8a7f3d57d28122109dd51a50d91d9f53a805. Final headca169f36aef921203a9857cb15cbf12a4ef695a1only adds UT dictionary annotations required by mypy. All production files and the original e2e source are byte-identical between these commits; this is not presented as a final-head e2e invocation. The two 56-test suites were rerun after that annotation follow-up, and the remote tested tracked tree was verified against the final signed commit.The new numerical builder regression fails on an independent, unmodified upstream-base checkout: 12/16 elements differ after global/local builds (maximum absolute difference 1.088031530380249). It passes with the fix. Coverage includes global → local → next global cache-group builds, preservation of both outputs, repeated calls, full-capacity growth, address stability, multiple RoPE configurations/groups, and unchanged uncached/draft paths.
Final-head repository CI:
Historical causal evidence, distinct from this branch's CI: an earlier isolated Ascend experiment with vLLM
a97dacb7106ee49f39f3d1fc6ae1800ff724e01dreproduced third-position acceptance 0.4933095450490633, matching the failing CI's complete outputs. With fixed diagnostic admission conditions in the same directory, the ownership fix passed, reverting it restored the exact failure, and reapplying it passed with identical tokens/counts to the first fixed run. Two subsequent original uninstrumented main singletons passed; historical release attempts were resource-blocked. Immutable causal archive SHA256:da75a996eba465e79837c834a5b21d72e9f4b0ae811163593c5eb6ba1f4484f6. The same Triton disk-cache seed produced both passes and failures, so cache contents are not claimed as the root cause.The separate local four-card attempt on this branch did not complete: main stopped before generation because ModelScope tried to write
.mdlthrough a read-only model-cache link; release was blocked before pytest by external device occupancy. Writable private metadata was prepared without downloading/modifying weights, and no external processes were stopped. These local attempts are not counted as passes or precision failures; the completed current-version device results above come from CI.Remaining merge gate:
ci-gaterequires a maintainer to addready-preciseorready-allfor the repository's broader selected/full suite. The targeted dual-version/e2erun does not automatically satisfy that label-controlled gate. No merge is requested by this task.