[Feature][Model] Support Kimi K3 PP+DSpark+SP on MRV2 - #17157
Conversation
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 enables pipeline parallelism (PP) combined with sequence parallelism (SP) for Kimi-K3 models, further extending support to include DSpark on Model Runner V2. The changes include logic for auxiliary state relay, host-position synchronization for speculative execution, and updated buffer management for PP+SP pipelines. These enhancements ensure consistent model execution across stages and are validated by new CPU regression tests. 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
|
github-actions
Bot
added
documentation
|
👋 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 [Feature] 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
This pull request implements pipeline parallelism (PP) combined with sequence parallelism (SP) and DSpark speculative decoding for Kimi-K3 models on Ascend NPUs. The changes enable token sharding across PP stages, cumulative auxiliary state relay, and host-position synchronization for speculative PP stages on both Model Runner V1 and V2. Additionally, the documentation has been updated, and comprehensive CPU unit tests have been added. The reviewer provided a suggested PR title and summary adhering to the repository style guide, and recommended adding a defensive check in _slice_kimi_sp_intermediate_tensors to prevent potential crashes when intermediate tensors are not yet allocated.
66d1731 to
0db08fa
Compare
LostFox11
changed the title
0db08fa to
2e0d494
Compare
github-actions
Bot
removed
the
documentation
ce92d80 to
18693de
Compare
|
This pull request has conflicts, please resolve those before we can evaluate the pull request. |
66f0b17 to
dcf6275
Compare
dcf6275 to
af80082
Compare
|
This pull request has conflicts, please resolve those before we can evaluate the pull request. |
0dd2c0c to
7999859
Compare
|
|
/rerun Rerun (failed jobs only):
|
### 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: - 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 #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>
|
This pull request has conflicts, please resolve those before we can evaluate the pull request. |
92178cc to
caf8303
Compare
|
/rerun Failed:
|
|
/rerun Rerun (failed jobs only):
|
…-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>
7cb44c6 to
5de1b2c
Compare
Keep full-token tensors at PP boundaries and shard them only inside the K3 model. Relay raw or materialized auxiliary states with capture-point-aware boundary slots, and preserve draft embedding ownership on the final stage. Reuse upstream weight mappings and auxiliary capture logic. Refresh Ascend CPU token counts on every speculative PP stage after rejection and prefill advancement, including native upstream PP. FIA still requires exact host lengths and the existing D2H completion wait. Cover sequence padding, empty stages, auxiliary boundaries, host token counts and draft weight alignment with CPU regressions. Real-weight PP3 TP16 graph smoke validation passed before the behavior-preserving cleanup; full GPQA Diamond evaluation remains running on that service snapshot. Signed-off-by: LostFox11 <wangziyue17@huawei.com>
Cache a K3-only opt-in for the additional speculative PP host-count refresh. Preserve the existing speculator-owned synchronization for other models and leave their non-last PP stages on upstream host upper bounds. Add CPU regressions for target aliases, PP/speculative enablement and opted-out native/legacy rejection and prefill paths. All 106 isolated CPU tests pass; the running GPQA service is unchanged. Signed-off-by: LostFox11 <wangziyue17@huawei.com>
Signed-off-by: LostFox11 <wangziyue17@huawei.com>
Initialize the CPU mock decoder's new main-branch fusion flag so PP/SP transport tests continue to exercise the unfused path without changing model behavior. Signed-off-by: LostFox11 <wangziyue17@huawei.com>
Adapt CPU PP/SP regressions to main's fused add, AttnRes and RMSNorm contract. Cover raw and materialized auxiliary states at every layer boundary, with SP on or off and empty first, middle and last stages. Signed-off-by: LostFox11 <wangziyue17@huawei.com>
5de1b2c to
bb65a08
Compare
|
/rerun Rerun (failed jobs only):
|
1 similar comment
|
/rerun Rerun (failed jobs only):
|
4 participants
What this PR does / why we need it?
Support Kimi K3 PP + SP + DSpark on Model Runner V2, keeping tensor layout changes in the model and existing PP utilities.
sync_spec_pp_cpu_countsat initialization, enabled only for K3 + PP + speculative decoding. It refreshes the Ascend host token-count snapshot after rejection and prefill advancement on non-last stages as well, for native and legacy PP. Other models keep the main-branch behavior: no additional non-last-stage synchronization, and the existing last-stage/speculator synchronization remains unchanged.K3 MLA and the tested GQA draft use FIA, whose CPU
actual_seq_kvlenmust be exact. The host-count fix reuses the existing side-stream D2H copy and event wait. Native upstream PP corrects device counters but does not replace this Ascend host snapshot. This is a correctness fix, not a claim to remove CPU synchronization or improve throughput.No changes to the MRV1 runner,
worker.py,aclgraph_utils.py, or the upstream PP token-broadcast protocol. This does not import #16511's runner/graph receive-buffer slicing.Cleanup removes the K3-only SP helper/allowlist, the duplicate upstream aux-capture override, and unused runner/test fields. Draft weight loading derives its PP mapper from the upstream mapper instead of copying the mapping table. The cleanup at
7999859e8removed 65 net production lines relative to0dd2c0c12; the follow-up only adds the K3 synchronization opt-in and its regression coverage.Does this PR introduce any user-facing change?
Kimi K3 can compose PP, SP and DSpark on MRV2 without a new option. Existing context-parallel restrictions remain. Hardware coverage below uses the GQA draft checkpoint; the separate K3 MLA draft is not hardware-validated by this run.
How was this patch tested?
tests/ut/models/test_kimi_k3_pp_sp.py. Covers SP token padding, dense MLP collectives, raw/materialized aux relay, PP cuts including an empty first stage, K3-only opt-in, native/legacy host-count synchronization, prefill advancement and non-speculative paths. An isolated CPU harness also passed 32 draft weight-alignment cases fromtest_dspark_model_loading.py, covering PP on/off, rotation on/off, and embedding/head ownership. These harnesses execute AST-extracted production methods with CPU tensors; they are not full repository-import tests and do not initialize NPU.git diff --checkpassed. Scoped mypy passed for the latest PP CPU regression suite (Python 3.10); the preceding cleanup also passed Python 3.12 checks for that suite and the DSpark loading tests.bash format.sh cicould not complete becausepre-commitis unavailable locally.32,32,29, MRV2, model-side SP, EP, DSpark 7, targetFULL_DECODE_ONLY, draft graphs enabled, asynchronous scheduling, prefix cache and chunked prefill enabled. Context 81920; max sequences 16; batched tokens 8192; memory utilization 0.92. Target and draft graph capture completed and/v1/modelsconfirmed context 81920.stop, used 50-349 output tokens, and had finite logprobs with semantically correct final answers. The previous NaN/repeated-!failure was not reproduced. Logs confirm active draft generation and accepted tokens.Answer: A/B/C/Dextraction; complete responses are retained. No final dataset accuracy is claimed yet.cee499a23, which differs from the pre-cleanup head0dd2c0c12only by import/format cleanup. The final cleanup has CPU coverage but has not been redeployed to the live service, so evaluation can finish without another large-weight reload. Containers, installed libraries and the existing deployment-only network patch were reused. No eager fallback, reinstall or rebuild.Follow-up