[BugFix][GDN] Skip recurrent state updates for idle DP dummy runs - #15510
[BugFix][GDN] Skip recurrent state updates for idle DP dummy runs#15510qijiajin wants to merge 4 commits into
Conversation
Keep GDN recurrent metadata at zero length when a data-parallel rank executes a synchronization-only dummy batch. This preserves graph and collective participation without advancing conv or KDA state through stale cache indices. Signed-off-by: maoxx241 <maomaoyu870@gmail.com>
Gemini review flagged common_attn_metadata.replace(...) as an AttributeError on older vLLM pins whose CommonAttentionMetadata dataclass has no instance replace() method. Use the already-imported dataclasses.replace() helper, which works on any dataclass and matches current vLLM's replace() implementation. Signed-off-by: maoxx241 <maomaoyu870@gmail.com>
Replace the dummy-run assert on num_reqs_padded with a fallback to num_reqs so a missing pad count cannot crash the execute path. Document that num_actual_tokens=0 is the GDN kernel no-op while query_start_loc stays a zero-filled graph-sized prefix. Signed-off-by: maoxx241 <maomaoyu870@gmail.com>
Use the uniform decode query length and descriptor for idle data-parallel dummy runs. This prevents an idle rank from downgrading FULL_DECODE_ONLY across the DP group while preserving the GDN state-update guard from PR vllm-project#15253. Signed-off-by: q00852295 <qijiajin1@huawei.com> (cherry picked from commit 83dd6f2)
|
👋 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 bug where idle Data Parallel (DP) ranks would trigger recurrent state updates during dummy runs, potentially causing synchronization issues or performance degradation in graph mode. By introducing a mechanism to skip these updates for idle ranks, the changes ensure that dummy runs maintain consistency across the DP group without unnecessary overhead. 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
This pull request introduces a skip_gdn_state_update flag to _dummy_run and attention metadata building in the model runner. This allows idle Data Parallel (DP) ranks to perform dummy runs without updating the GDN recurrent state, preventing out-of-sync issues while maintaining correct coordinate batch behavior. The reviewer provided a suggested PR Title and PR Summary formatted according to the repository's style guide.
| self.model_runner._dummy_run( | ||
| num_tokens, | ||
| uniform_decode=True, | ||
| skip_gdn_state_update=True, | ||
| ) |
There was a problem hiding this comment.
Suggested PR Title:
[Attention][BugFix] Skip recurrent state updates for idle DP dummy runsSuggested PR Summary:
### What this PR does / why we need it?
This PR skips recurrent state updates (specifically GDN state updates) for idle Data Parallel (DP) dummy runs. During dummy runs on idle DP ranks, updating the GDN state can lead to out-of-sync issues or incorrect state propagation across the DP group. By introducing the `skip_gdn_state_update` flag, we ensure that idle ranks do not perform unnecessary state updates while still maintaining the correct coordinate batch behavior to avoid out-of-sync issues.
### Does this PR introduce _any_ user-facing change?
No.
### How was this patch tested?
- Unit tests in `tests/ut/worker/test_model_runner_v1.py` and `tests/ut/worker/test_worker_v1.py` were updated and verified.
- Validation with `git diff --check` and `python -m compileall` on the changed runtime files.References
- Generate PR Title and PR Summary in the specified markdown format as required by the Repository Style Guide. (link)
2 participants
What this PR does / why we need it
Data-parallel ranks with no scheduled requests still execute a one-token dummy batch to participate in graph replay and collectives. For GDN/Mamba layers, treating that synchronization-only token as a real decode token can advance causal-conv and KDA recurrent state through block-table entries left by an earlier request.
This PR keeps the model/collective dummy shape unchanged, but builds zero-length GDN metadata for idle DP dummies. Both captured recurrent branches become no-ops while other layers and distributed communication continue to execute normally.
The change restores the behavior previously validated on the v0.26 Kimi K3 path. It adds no device-wide synchronization and does not change GDN kernels.
Does this PR introduce any user-facing change?
No API change. Multi-node GDN/Mamba inference no longer depends on whether another DP rank was idle between requests.
How was this patch tested?
python -m compileallfor the changed runtime filesgit diff --checkExisting unit expectations updated for both idle-dummy entry points; the existing GDN builder regression covers zero-length recurrent metadata
Real-weight A/B on two A3 nodes, TP16 x DP2 / EP32, Kimi K3 GQA DSpark, async scheduling, prefix cache, FULL_DECODE_ONLY graph, 131072-input + 1024-output:
1.1.1...output while speculative acceptance remained 80.55%Real-weight four-node smoke on TP16 x DP4 / EP64 with full QuaRot weights and the same 128K-to-1K case:
vLLM main: vllm-project/vllm@ba07e4a