Repository navigation
[BugFix][V2] Fix speculative PP CPU token counts for FIA models - #17675
Conversation
Extract the Kimi speculative PP count refresh and enable it for dense Qwen3.5. Reuse the pinned asynchronous D2H staging after rejection and unsampled prefill chunks, then wait before building CPU attention lengths. Preserve existing behavior for other models and non-speculative PP. Signed-off-by: LostFox11 <wangziyue17@huawei.com>
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 an issue where non-last speculative PP stages failed to correctly update device-side request states, leading to incorrect or repeated output in FIA models. By ensuring host-count refreshes occur on all relevant PP stages—not just those owning a speculator—and extending this to unsampled prefill chunks, the changes maintain consistency in KV lengths and request states during speculative decoding. 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
|
|
👋 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! |
![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:
[Worker][Feature] Support speculative PP host-count updates for FIA modelsSuggested PR Summary:
### What this PR does / why we need it?
This PR introduces speculative PP host-count updates on non-last PP stages for specific FIA models (such as Kimi and Qwen3.5 architectures). It adds the `sync_spec_pp_cpu_counts` flag to synchronize and copy corrected computed tokens to the CPU after rejection or prefill chunks.
However, accessing `self.model_config.architecture` directly may raise an `AttributeError` because standard vLLM uses `architectures` (a list of strings). A fallback mechanism is recommended to safely retrieve the architecture name while maintaining compatibility with the AST-manipulating test fixture.
### Does this PR introduce _any_ user-facing change?
No.
### How was this patch tested?
A new CPU-only unit test suite `tests/ut/worker/test_spec_pp_cpu_counts.py` has been added to verify exact counts after rejection or chunking.Signed-off-by: LostFox11 <wangziyue17@huawei.com>
Signed-off-by: LostFox11 <wangziyue17@huawei.com>
|
Add the following information in the comments:
|
|
/rerun Rerun (failed jobs only):
|
Signed-off-by: LostFox11 <wangziyue17@huawei.com>
…-project#17675) ### What this PR does / why we need it? Non-last speculative PP stages receive rejected-token counts and update device-side request state, but they do not own a speculator. Ascend currently refreshes the CPU counts only on the speculator-owning rank. As a result, FIA can consume optimistic KV lengths after rejection, producing incorrect/repeated output. Extract the host-count refresh from vllm-project#17157 into a standalone fix: - Enable the extra refresh for Kimi/K3 and dense Qwen3.5 speculative PP stages. - Reuse the existing pinned D2H buffer, side stream and event after sampled-output processing. - Also refresh after unsampled prefill chunks, so the next chunk cannot consume a stale snapshot. - Wait before deriving CPU attention lengths, updating cached requests only. Preserve new-request slot initialization and existing behavior for other models and non-speculative PP. No changes to PP communication, rejection sampling, model forward, graph dispatch or operators. This does not include the separate K3 PP/SP model adaptation. ### Does this PR introduce _any_ user-facing change? Fixes incorrect/repeated output with the covered speculative PP configurations. No new CLI/environment option. Reuses asynchronous staging but still waits for the event before consuming host counts; this is not a zero-synchronization performance change. ### How was this patch tested? - CPU-only tests: `PYTEST_DISABLE_PLUGIN_AUTOLOAD=1 python -m pytest --noconftest -q tests/ut/worker/test_spec_pp_cpu_counts.py`: **72 passed**. Covers model opt-in, PP/speculative enablement, non-last/last stages, rejected tokens, unsampled prefill chunks and newly reused slots, without importing NPU workers. - Ascend real-weight regression: Qwen3.5-27B, MRV2, PP2TP8, async scheduling, MTP with 3 draft tokens and draft eager, target `FULL_DECODE_ONLY`, prefix cache and chunked prefill. Max model length 32768, max sequences 64, max batched tokens 8192, graph capture maximum 256, memory utilization 0.85. - Qwen3.5: 128 deterministic arithmetic/format requests **at each concurrency, 8 and 64**. Both runs: **128/128 correct, no repeated exclamation marks, truncations or request errors**. Includes approximately 11K-token prompts and checks the complete ordered number list. These are regression probes, not a benchmark dataset accuracy score. - Also tested the separate local Qwen3.8-27B checkpoint, which declares `Qwen3_5ForConditionalGeneration`, with the same configuration: three readable single-request responses and **128/128 correct at both concurrency 8 and 64**, with no repetition, truncation or errors. Metrics confirmed actual running-request peaks of 64. Concurrency 64 followed concurrency 8 with warm prefix cache; no performance comparison is claimed. - K3 real-weight testing was not repeated for this standalone extraction; its PP/SP model adaptation remains in vllm-project#17157. - Targeted mypy (Python 3.10), Ruff and `git diff --check` passed. `bash format.sh ci` was attempted; code-related checks passed, but the full local run was blocked by missing gitleaks/wget and shellcheck binaries. - vLLM main: vllm-project/vllm@ced6857 --------- Signed-off-by: LostFox11 <wangziyue17@huawei.com> Co-authored-by: LostFox11 <wangziyue17@huawei.com>
5 participants
What this PR does / why we need it?
Non-last speculative PP stages receive rejected-token counts and update device-side request state, but they do not own a speculator. Ascend currently refreshes the CPU counts only on the speculator-owning rank. As a result, FIA can consume optimistic KV lengths after rejection, producing incorrect/repeated output.
Extract the host-count refresh from #17157 into a standalone fix:
No changes to PP communication, rejection sampling, model forward, graph dispatch or operators. This does not include the separate K3 PP/SP model adaptation.
Does this PR introduce any user-facing change?
Fixes incorrect/repeated output with the covered speculative PP configurations. No new CLI/environment option. Reuses asynchronous staging but still waits for the event before consuming host counts; this is not a zero-synchronization performance change.
How was this patch tested?
CPU-only tests:
PYTEST_DISABLE_PLUGIN_AUTOLOAD=1 python -m pytest --noconftest -q tests/ut/worker/test_spec_pp_cpu_counts.py: 72 passed. Covers model opt-in, PP/speculative enablement, non-last/last stages, rejected tokens, unsampled prefill chunks and newly reused slots, without importing NPU workers.Ascend real-weight regression: Qwen3.5-27B, MRV2, PP2TP8, async scheduling, MTP with 3 draft tokens and draft eager, target
FULL_DECODE_ONLY, prefix cache and chunked prefill. Max model length 32768, max sequences 64, max batched tokens 8192, graph capture maximum 256, memory utilization 0.85.Qwen3.5: 128 deterministic arithmetic/format requests at each concurrency, 8 and 64. Both runs: 128/128 correct, no repeated exclamation marks, truncations or request errors. Includes approximately 11K-token prompts and checks the complete ordered number list. These are regression probes, not a benchmark dataset accuracy score.
Also tested the separate local Qwen3.8-27B checkpoint, which declares
Qwen3_5ForConditionalGeneration, with the same configuration: three readable single-request responses and 128/128 correct at both concurrency 8 and 64, with no repetition, truncation or errors. Metrics confirmed actual running-request peaks of 64. Concurrency 64 followed concurrency 8 with warm prefix cache; no performance comparison is claimed.K3 real-weight testing was not repeated for this standalone extraction; its PP/SP model adaptation remains in [Feature][Model] Support Kimi K3 PP+DSpark+SP on MRV2 #17157.
Targeted mypy (Python 3.10), Ruff and
git diff --checkpassed.bash format.sh ciwas attempted; code-related checks passed, but the full local run was blocked by missing gitleaks/wget and shellcheck binaries.vLLM main: vllm-project/vllm@ced6857