[Feature] Support dsv4 mtp graph for dsa_cp - #12599
Conversation
|
👋 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! |
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 MTP=3 speculative decoding within the DSA-CP path, extending support to both eager and graph execution modes. By introducing pre-allocated metadata buffers and updated builder logic, the changes ensure deterministic graph capture and replay, addressing previous limitations with multi-step draft processing. 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] Optimize DSA-CP metadata building with Triton kernel and cachingSuggested PR Summary:
### What this PR does / why we need it?
This pull request optimizes the DSA-CP metadata building process for DeepSeek-V4 on Ascend NPU. It introduces a fused Triton kernel `build_local_metadata_triton` to compute local token metadata, reducing kernel launch overhead. Additionally, it implements caching for device-local, CPU-local, and RoPE-local metadata across KV-cache groups to avoid redundant computations. It also adds a `pad_to` method to `RopeDataProxy` to facilitate padding.
Feedback and issues identified:
- In `test_build_local_metadata_triton.py`, the rank filtering logic is inverted, skipping the boundary ranks (0 and tp_size - 1) instead of testing only them.
- In `dsa_cp.py`, slicing operations using `SAS_METADATA_SIZE` can cause `RuntimeError` due to shape mismatches if the metadata tensor size is smaller than `SAS_METADATA_SIZE` (e.g., during speculative decoding with small batch sizes).
### Does this PR introduce _any_ user-facing change?
No.
### How was this patch tested?
- Added unit tests in `tests/ut/ops/test_rope_proxy.py` to verify the equivalence of `RopeDataProxy.pad_to` and slicing against the original gather approach.
- Added Triton kernel tests in `tests/e2e/nightly/single_node/ops/singlecard_ops/triton/test_build_local_metadata_triton.py`.7e48fcd to
a928f51
Compare
|
This pull request has conflicts, please resolve those before we can evaluate the pull request. |
a928f51 to
c594885
Compare
|
This pull request has conflicts, please resolve those before we can evaluate the pull request. |
c594885 to
ed404f7
Compare
|
@pisceskkk PTAL This PR adds full ACL graph support for MTP=3 for DSA-CP |
|
Hi @pisceskkk, gentle ping on this PR, it adds full ACL graph support for MTP=3 for DSA-CP. Just wanted to check if you've had a chance to take a look. |
2b3ae50 to
caf8e36
Compare
|
This pull request has conflicts, please resolve those before we can evaluate the pull request. |
caf8e36 to
f58713c
Compare
|
Hi @pisceskkk, thanks for the review! Per your comment, the MTP+graph test is now folded into the existing deepseek_v4_w4a8_dsa_cp_full_features case, and I've rebased onto the latest main to resolve the conflicts. PTAL again, thanks! |
|
This pull request has conflicts, please resolve those before we can evaluate the pull request. |
939a2f5 to
a77e4cd
Compare
|
This pull request has conflicts, please resolve those before we can evaluate the pull request. |
8aeb69b to
c4f607a
Compare
|
This pull request has conflicts, please resolve those before we can evaluate the pull request. |
- Pad seq_lens_cpu to the graph-dispatch request count so batch dims stay stable across ACL graph capture/replay - Cache SAS metadata in per-draft-index preallocated buffers for address stability - Use per-draft-index RoPE cache (use_cache=True) so cos/sin tensor addresses stay stable across draft steps - Replace magic number 1024 with SAS_METADATA_SIZE constant - Add no-op update_graph_params hook for ACL graph dispatch Adapted to the cross-kv-cache-group drafting metadata cache from vllm-project#14916: per-draft stable buffers are filled after the shared-cache lookup so both cache-hit and cache-miss paths end up with stable tensor addresses. Signed-off-by: frankie <wangyongsheng686@gmail.com>
Enable 3-token MTP speculative decoding in the four-card DSA-CP accuracy test to cover the new acl graph path. Signed-off-by: frankie <wangyongsheng686@gmail.com>
The draft-step metadata builder returned freshly cloned local_query_start_loc / local_seq_lens (and start_pos) tensors. ACL graph capture bakes those tensor addresses into the draft graph, but the allocator freely reuses the clone's memory between capture and replay, so replays read arbitrary data as TND cumsum metadata and the attention kernel moves rows out of range (MTE illegal GM access, EZ9999). Return views into the per-draft-index persistent buffers instead: - cache-miss path fills spec_local_query_start_loc / spec_local_seq_lens / spec_start_pos and returns their slices (no clone) - cache-hit path copies the sibling group's snapshot into the same stable buffers before returning views - the cross-group cache stores value snapshots (clone) so sibling groups can copy them into their own stable buffers Signed-off-by: frankie <wangyongsheng686@gmail.com>
c4f607a to
8288a08
Compare
### What this PR does / why we need it? This PR enables MTP=3 speculative decoding in the DSA‑CP path for both eager and graph modes. It introduces full ACL graph support for the MTP draft steps, including pre‑allocated metadata buffers that allow deterministic graph capture and replay when enable_dsa_cp is true. ### Changes summary #### MTP > 1 support in DSA‑CP ~Previously build_for_drafting only handled MTP=1 and lacked the CPU‑side sequence length path, which caused tensor dimension errors for any MTP > 1.~ Supported by vllm-project#13249. This PR updates build_for_drafting to populate CPU sequence lengths, so that all downstream metadata builders can correctly process multiple draft steps. #### Full ACL graph support for MTP draft steps ACL graph capture and replay require stable tensor addresses across invocations. To guarantee this, the PR pre‑allocates per‑step buffers for all draft‑step metadata during builder initialization, and pads seq_lens_cpu to match the graph‑dispatched batch size, ensuring deterministic addresses throughout replay. In addition, a static `update_graph_params` method was added, so that the graph dispatch can correctly invoke the DSA-CP attention backend during graph replay. ~#### Depends on PR vllm-project#12193.~ ~vllm-project#12193 also fixes metadata mismatch issues in the DSA-CP builder, and this PR is directly based on the refactored code. Cherry-picking the commits onto main without vllm-project#12193 causes runtime errors under concurrent requests.~ ~Only the top commits (after Commits on Jul 21, 2026) belong to this PR. Please review by focusing on the top commits. Once it is merged, I'll rebase onto main quickly.`~ ### Does this PR introduce _any_ user-facing change? No. ### How was this patch tested? - Added an E2E accuracy test in `tests/e2e/pull_request/four_card/test_deepseek_v4.py` with exact `expected_token_ids` assertion to guard against regressions in the DSA-CP + MTP=3 + full graph path. - GMS8K accuracy evaluation for MTP+eager and MTP+graph is attached below. ``` mtp3+eager +---------+-----------+----------+----------+-------+---------+---------+ | Model | Dataset | Metric | Subset | Num | Score | Cat.0 | +=========+===========+==========+==========+=======+=========+=========+ | dsv4-f | gsm8k | mean_acc | main | 1319 | 0.9666 | default | +---------+-----------+----------+----------+-------+---------+---------+ mtp3+full graph +---------+-----------+----------+----------+-------+---------+---------+ | Model | Dataset | Metric | Subset | Num | Score | Cat.0 | +=========+===========+==========+==========+=======+=========+=========+ | dsv4-f | gsm8k | mean_acc | main | 1319 | 0.9659 | default | +---------+-----------+----------+----------+-------+---------+---------+ ``` - vLLM main: vllm-project/vllm@e6bfe03 --------- Signed-off-by: frankie <wangyongsheng686@gmail.com> Co-authored-by: frankie <wangyongsheng686@gmail.com>
…)" This reverts commit 32c28fb. Signed-off-by: hejianping-00178005 <44997374+winson-00178005@users.noreply.github.com>
Reverts #12599 - vLLM main: vllm-project/vllm@b2f6858 Signed-off-by: hejianping-00178005 <44997374+winson-00178005@users.noreply.github.com>
… sizing main reverted vllm-project#12599 (5d4294f); this re-applies the DSA-CP aclgraph feature on top of current main, keeping the graph-mode-aware per-request buffer sizing (max_padded_reqs) that fixes EZ1007 copy_ shape mismatches and out-of-bounds triton writes when num_reqs is padded beyond max_num_seqs during FULL_DECODE_ONLY aclgraph replay. Signed-off-by: frankie <wangyongsheng686@gmail.com>
…cp graph Re-applies the MTP speculative config removed by the vllm-project#12599 revert and adds a concurrent-request accuracy guard: max_num_seqs smaller than the cudagraph capture bucket with 4 concurrent MTP requests exercises the padded draft path. Signed-off-by: frankie <wangyongsheng686@gmail.com>
### What this PR does / why we need it? This PR enables MTP=3 speculative decoding in the DSA‑CP path for both eager and graph modes. It introduces full ACL graph support for the MTP draft steps, including pre‑allocated metadata buffers that allow deterministic graph capture and replay when enable_dsa_cp is true. ### Changes summary #### MTP > 1 support in DSA‑CP ~Previously build_for_drafting only handled MTP=1 and lacked the CPU‑side sequence length path, which caused tensor dimension errors for any MTP > 1.~ Supported by vllm-project#13249. This PR updates build_for_drafting to populate CPU sequence lengths, so that all downstream metadata builders can correctly process multiple draft steps. #### Full ACL graph support for MTP draft steps ACL graph capture and replay require stable tensor addresses across invocations. To guarantee this, the PR pre‑allocates per‑step buffers for all draft‑step metadata during builder initialization, and pads seq_lens_cpu to match the graph‑dispatched batch size, ensuring deterministic addresses throughout replay. In addition, a static `update_graph_params` method was added, so that the graph dispatch can correctly invoke the DSA-CP attention backend during graph replay. ~#### Depends on PR vllm-project#12193.~ ~vllm-project#12193 also fixes metadata mismatch issues in the DSA-CP builder, and this PR is directly based on the refactored code. Cherry-picking the commits onto main without vllm-project#12193 causes runtime errors under concurrent requests.~ ~Only the top commits (after Commits on Jul 21, 2026) belong to this PR. Please review by focusing on the top commits. Once it is merged, I'll rebase onto main quickly.`~ ### Does this PR introduce _any_ user-facing change? No. ### How was this patch tested? - Added an E2E accuracy test in `tests/e2e/pull_request/four_card/test_deepseek_v4.py` with exact `expected_token_ids` assertion to guard against regressions in the DSA-CP + MTP=3 + full graph path. - GMS8K accuracy evaluation for MTP+eager and MTP+graph is attached below. ``` mtp3+eager +---------+-----------+----------+----------+-------+---------+---------+ | Model | Dataset | Metric | Subset | Num | Score | Cat.0 | +=========+===========+==========+==========+=======+=========+=========+ | dsv4-f | gsm8k | mean_acc | main | 1319 | 0.9666 | default | +---------+-----------+----------+----------+-------+---------+---------+ mtp3+full graph +---------+-----------+----------+----------+-------+---------+---------+ | Model | Dataset | Metric | Subset | Num | Score | Cat.0 | +=========+===========+==========+==========+=======+=========+=========+ | dsv4-f | gsm8k | mean_acc | main | 1319 | 0.9659 | default | +---------+-----------+----------+----------+-------+---------+---------+ ``` - vLLM main: vllm-project/vllm@e6bfe03 --------- Signed-off-by: frankie <wangyongsheng686@gmail.com> Co-authored-by: frankie <wangyongsheng686@gmail.com>
) Reverts vllm-project#12599 - vLLM main: vllm-project/vllm@b2f6858 Signed-off-by: hejianping-00178005 <44997374+winson-00178005@users.noreply.github.com>
) Reverts vllm-project#12599 - vLLM main: vllm-project/vllm@b2f6858 Signed-off-by: hejianping-00178005 <44997374+winson-00178005@users.noreply.github.com> Signed-off-by: tianming2009 <13246728590@163.com>
### What this PR does / why we need it? This PR enables MTP=3 speculative decoding in the DSA‑CP path for both eager and graph modes. It introduces full ACL graph support for the MTP draft steps, including pre‑allocated metadata buffers that allow deterministic graph capture and replay when enable_dsa_cp is true. ### Changes summary #### MTP > 1 support in DSA‑CP ~Previously build_for_drafting only handled MTP=1 and lacked the CPU‑side sequence length path, which caused tensor dimension errors for any MTP > 1.~ Supported by vllm-project#13249. This PR updates build_for_drafting to populate CPU sequence lengths, so that all downstream metadata builders can correctly process multiple draft steps. #### Full ACL graph support for MTP draft steps ACL graph capture and replay require stable tensor addresses across invocations. To guarantee this, the PR pre‑allocates per‑step buffers for all draft‑step metadata during builder initialization, and pads seq_lens_cpu to match the graph‑dispatched batch size, ensuring deterministic addresses throughout replay. In addition, a static `update_graph_params` method was added, so that the graph dispatch can correctly invoke the DSA-CP attention backend during graph replay. ~#### Depends on PR vllm-project#12193.~ ~vllm-project#12193 also fixes metadata mismatch issues in the DSA-CP builder, and this PR is directly based on the refactored code. Cherry-picking the commits onto main without vllm-project#12193 causes runtime errors under concurrent requests.~ ~Only the top commits (after Commits on Jul 21, 2026) belong to this PR. Please review by focusing on the top commits. Once it is merged, I'll rebase onto main quickly.`~ ### Does this PR introduce _any_ user-facing change? No. ### How was this patch tested? - Added an E2E accuracy test in `tests/e2e/pull_request/four_card/test_deepseek_v4.py` with exact `expected_token_ids` assertion to guard against regressions in the DSA-CP + MTP=3 + full graph path. - GMS8K accuracy evaluation for MTP+eager and MTP+graph is attached below. ``` mtp3+eager +---------+-----------+----------+----------+-------+---------+---------+ | Model | Dataset | Metric | Subset | Num | Score | Cat.0 | +=========+===========+==========+==========+=======+=========+=========+ | dsv4-f | gsm8k | mean_acc | main | 1319 | 0.9666 | default | +---------+-----------+----------+----------+-------+---------+---------+ mtp3+full graph +---------+-----------+----------+----------+-------+---------+---------+ | Model | Dataset | Metric | Subset | Num | Score | Cat.0 | +=========+===========+==========+==========+=======+=========+=========+ | dsv4-f | gsm8k | mean_acc | main | 1319 | 0.9659 | default | +---------+-----------+----------+----------+-------+---------+---------+ ``` - vLLM main: vllm-project/vllm@e6bfe03 --------- Signed-off-by: frankie <wangyongsheng686@gmail.com> Co-authored-by: frankie <wangyongsheng686@gmail.com> Signed-off-by: like-0517 <ithwlike@126.com>
) Reverts vllm-project#12599 - vLLM main: vllm-project/vllm@b2f6858 Signed-off-by: hejianping-00178005 <44997374+winson-00178005@users.noreply.github.com> Signed-off-by: like-0517 <ithwlike@126.com>
…der full cudagraph padding (#16168) ### What this PR does / why we need it? This PR re-applies the DSA-CP + MTP aclgraph support for DeepSeek V4 that was reverted from main in #16181 (revert of #12599), together with a fix for the crash that motivated scrutiny of the original feature: serving concurrent requests with `cudagraph_mode=FULL_DECODE_ONLY` failed with `aclnnInplaceCopy` shape-mismatch errors (EZ1007) or illegal device memory accesses (EZ9999). **Root cause of the crash** In FULL-decode graph mode, the DSA-CP draft path receives `num_reqs` as the *padded* request count — the cudagraph capture bucket size plus the FIA dummy request from mixed-batch padding — which can exceed `scheduler_config.max_num_seqs`. The per-request metadata buffers were sized by `max_num_seqs`, so: 1. `[:num_reqs]` views of these buffers get silently truncated and `copy_()` into them fails with shape mismatches — e.g. `--max-num-seqs 6` with 3 concurrent MTP requests pads 12 tokens to bucket 16, producing `Shape [16] vs [6] do not meet the broadcast condition (EZ1007)`; 2. triton kernels receiving the whole buffer write past its end, corrupting adjacent device memory (`EZ9999: MTE accesses an invalid GM address`). **Fix** Compute one graph-mode-aware capacity up front and size all per-request buffers with it: max_padded_reqs = max(max_num_seqs, max_cudagraph_capture_size) + 1 # +1: FIA dummy request This covers the step-0 buffers (`start_pos_prefill`, `local_query_start_loc`, `local_seq_lens`), the draft buffers (`spec_local_query_start_loc`, `spec_local_seq_lens`, `spec_start_pos`), and the QLI buffers (`qli_seqused_k`, `qli_cmp_residual_k`) with a single capacity source. All usage sites are `[:num_reqs]` slices or whole-buffer kernel inputs, so the change is a pure capacity gain with negligible (int32-level) memory overhead. The re-applied feature is fully adapted to current main: it fuses with the sequence-parallel prefill rework (#15549), the dynamic DSA indexer quant_mode (#16224) and the DSA PCP + DSpark support (#15958) already on main. The feature enables aclgraph capture/replay for DSA-CP with MTP=1 and MTP=3, including stable per-draft-index metadata buffers (`spec_sas_metadata`, `spec_start_pos`, `spec_local_*`, per-draft RoPE cache) so tensor addresses stay fixed across graph capture and replay. ### Does this PR introduce _any_ user-facing change? Yes, it re-enables DSA-CP + MTP with FULL_DECODE_ONLY graphs for DeepSeek V4 (`enable_dsa_cp: true` + `speculative_config` + `cudagraph_mode: FULL_DECODE_ONLY`), which was available before the revert. Users with `max_num_seqs` smaller than the max cudagraph capture size no longer hit EZ1007 / EZ9999 errors under concurrent load. The gain show below: when use dsv4+mtp 3+eager: <img width="371" height="509" alt="c98d2620e0c57e8a2de08fd8281e6903" src="https://github.com/user-attachments/assets/78d1eeb8-e0ae-4421-b555-84ce5d6b8987" /> when use dsv4+ mtp 3 + graph: <img width="374" height="508" alt="74786d0aebc590d420bac3bfbf436f32" src="https://github.com/user-attachments/assets/7f3eb4d5-7bef-4b83-97f0-e9c825c2486e" /> GMS8K accuracy evaluation for MTP+eager and MTP+graph is attached below. dsv4-flash+mtp3+eager Overall report table: +---------+-----------+----------+----------+-------+---------+---------+ | Model | Dataset | Metric | Subset | Num | Score | Cat.0 | +=========+===========+==========+==========+=======+=========+=========+ | dsv4 | gsm8k | mean_acc | main | 1318 | 0.9734 | default | +---------+-----------+----------+----------+-------+---------+---------+ dsv4-flash+mtp3+graph Overall report table: +---------+-----------+----------+----------+-------+---------+---------+ | Model | Dataset | Metric | Subset | Num | Score | Cat.0 | +=========+===========+==========+==========+=======+=========+=========+ | dsv4 | gsm8k | mean_acc | main | 1319 | 0.9742 | default | +---------+-----------+----------+----------+-------+---------+---------+ glm5.2-w4a8c8 + mtp 5 +graph Overall report table: +---------+-----------+----------+----------+-------+---------+---------+ | Model | Dataset | Metric | Subset | Num | Score | Cat.0 | +=========+===========+==========+==========+=======+=========+=========+ | glm5.2 | gsm8k | mean_acc | main | 1319 | 0.978 | default | +---------+-----------+----------+----------+-------+---------+---------+ --------- Signed-off-by: frankie <wangyongsheng686@gmail.com>
What this PR does / why we need it?
This PR enables MTP=3 speculative decoding in the DSA‑CP path for both eager and graph modes.
It introduces full ACL graph support for the MTP draft steps, including pre‑allocated metadata buffers that allow deterministic graph capture and replay when enable_dsa_cp is true.
Changes summary
MTP > 1 support in DSA‑CP
Previously build_for_drafting only handled MTP=1 and lacked the CPU‑side sequence length path, which caused tensor dimension errors for any MTP > 1.Supported by #13249.This PR updates build_for_drafting to populate CPU sequence lengths, so that all downstream metadata builders can correctly process multiple draft steps.
Full ACL graph support for MTP draft steps
ACL graph capture and replay require stable tensor addresses across invocations. To guarantee this, the PR pre‑allocates per‑step buffers for all draft‑step metadata during builder initialization, and pads seq_lens_cpu to match the graph‑dispatched batch size, ensuring deterministic addresses throughout replay.
In addition, a static
update_graph_paramsmethod was added, so that the graph dispatch cancorrectly invoke the DSA-CP attention backend during graph replay.
#### Depends on PR #12193.#12193 also fixes metadata mismatch issues in the DSA-CP builder, and this PR is directly based on the refactored code. Cherry-picking the commits onto main without #12193 causes runtime errors under concurrent requests.Only the top commits (after Commits on Jul 21, 2026) belong to this PR. Please review by focusing on the top commits. Once it is merged, I'll rebase onto main quickly.`Does this PR introduce any user-facing change?
No.
How was this patch tested?
tests/e2e/pull_request/four_card/test_deepseek_v4.pywith exactexpected_token_idsassertion to guard against regressions in theDSA-CP + MTP=3 + full graph path.