Skip to content

[Feature] Support dsv4 mtp graph for dsa_cp - #12599

Merged
weijinqian0 merged 3 commits into
vllm-project:mainfrom
Csrayz:feat/dsacp_mtp_graph
Sep 8, 2026
Merged

weijinqian0 merged 3 commits into
vllm-project:mainfrom
Csrayz:feat/dsacp_mtp_graph

Conversation

@Csrayz

@Csrayz Csrayz commented Jul 22, 2026 •

Copy link
Copy Markdown
Contributor

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_params method was added, so that the graph dispatch can
correctly 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?

  • 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 |
+---------+-----------+----------+----------+-------+---------+---------+ 

@github-actions

Copy link
Copy Markdown
Contributor

👋 Hi! Thank you for contributing to the vLLM Ascend project. The following points will speed up your PR merge:‌‌

  • A PR should do only one thing, smaller PRs enable faster reviews.
  • Every PR should include unit tests and end-to-end tests ‌to ensure it works and is not broken by other future PRs.
  • Write the commit message by fulfilling the PR description to help reviewer and future developers understand.

If CI fails, you can run linting and testing checks locally according Contributing and Testing.


Tip

💡 Consider Linking a Related Issue or RFC

Your 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:

  • Fixes #<issue_number>
  • Closes #<issue_number>
  • Resolves #<issue_number>
  • Refs #<rfc_or_issue_number> (for RFCs)

🙏 Thanks for helping us keep the project well-organized!

@gemini-code-assist

Copy link
Copy Markdown
Contributor

Summary of Changes

Hello, 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

  • MTP Support: Enabled MTP=3 speculative decoding in the DSA-CP path for both eager and graph execution modes.
  • ACL Graph Support: Implemented deterministic graph capture and replay by pre-allocating metadata buffers and adding a static update_graph_params method.
  • Metadata Handling: Updated build_for_drafting to correctly handle MTP > 1 by populating CPU sequence lengths and ensuring consistent metadata processing.
  • Performance Optimization: Introduced a fused Triton kernel for local token metadata computation to reduce kernel launch overhead.
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 Assist

The 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 /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature Command Description
Code Review /gemini review Performs a code review for the current pull request in its current state.
Pull Request Summary /gemini summary Provides a summary of the current pull request in its current state.
Comment @gemini-code-assist Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help /gemini help Displays a list of available commands.

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 .gemini/ folder in the base of the repository. Detailed instructions can be found here.

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

  1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

@gemini-code-assist gemini-code-assist Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Code Review

Suggested PR Title:

[Attention][Feature] Optimize DSA-CP metadata building with Triton kernel and caching

Suggested 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`.

Comment thread vllm_ascend/attention/context_parallel/dsa_cp.py
Comment thread vllm_ascend/attention/context_parallel/dsa_cp.py
Comment thread vllm_ascend/attention/context_parallel/dsa_cp.py
@github-actions

Copy link
Copy Markdown
Contributor

This pull request has conflicts, please resolve those before we can evaluate the pull request.

@github-actions

Copy link
Copy Markdown
Contributor

This pull request has conflicts, please resolve those before we can evaluate the pull request.

@Csrayz
Csrayz force-pushed the feat/dsacp_mtp_graph branch from c594885 to ed404f7 Compare July 25, 2026 01:27
@Csrayz
Csrayz marked this pull request as ready for review July 25, 2026 01:37
@Csrayz

Csrayz commented Jul 27, 2026

Copy link
Copy Markdown
Contributor Author

@pisceskkk PTAL This PR adds full ACL graph support for MTP=3 for DSA-CP

@Csrayz

Csrayz commented Jul 29, 2026

Copy link
Copy Markdown
Contributor Author

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.

Comment thread tests/e2e/pull_request/four_card/test_deepseek_v4.py Outdated
@Csrayz
Csrayz force-pushed the feat/dsacp_mtp_graph branch 3 times, most recently from 2b3ae50 to caf8e36 Compare July 31, 2026 15:15
@github-actions

github-actions Bot commented Aug 1, 2026

Copy link
Copy Markdown
Contributor

This pull request has conflicts, please resolve those before we can evaluate the pull request.

@Csrayz

Csrayz commented Aug 4, 2026

Copy link
Copy Markdown
Contributor Author

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!

@github-actions

github-actions Bot commented Sep 3, 2026

Copy link
Copy Markdown
Contributor

This pull request has conflicts, please resolve those before we can evaluate the pull request.

@github-actions

github-actions Bot commented Sep 4, 2026

Copy link
Copy Markdown
Contributor

This pull request has conflicts, please resolve those before we can evaluate the pull request.

@github-actions

github-actions Bot commented Sep 7, 2026

Copy link
Copy Markdown
Contributor

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>
@weijinqian0
weijinqian0 merged commit 32c28fb into vllm-project:main Sep 8, 2026
28 checks passed
jiangli221 pushed a commit to jiangli221/vllm-ascend that referenced this pull request Sep 9, 2026
### 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>
winson-00178005 added a commit to winson-00178005/vllm-ascend that referenced this pull request Sep 9, 2026
…)"

This reverts commit 32c28fb.

Signed-off-by: hejianping-00178005 <44997374+winson-00178005@users.noreply.github.com>
weijinqian0 pushed a commit that referenced this pull request Sep 10, 2026
Reverts #12599

- vLLM main:
vllm-project/vllm@b2f6858

Signed-off-by: hejianping-00178005 <44997374+winson-00178005@users.noreply.github.com>
frankie-ys added a commit to frankie-ys/vllm-ascend that referenced this pull request Sep 11, 2026
… 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>
frankie-ys added a commit to frankie-ys/vllm-ascend that referenced this pull request Sep 11, 2026
…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>
sunny-rain-63 pushed a commit to sunny-rain-63/vllm-ascend that referenced this pull request Sep 12, 2026
### 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>
sunny-rain-63 pushed a commit to sunny-rain-63/vllm-ascend that referenced this pull request Sep 12, 2026
)

Reverts vllm-project#12599

- vLLM main:
vllm-project/vllm@b2f6858

Signed-off-by: hejianping-00178005 <44997374+winson-00178005@users.noreply.github.com>
johnnysluckydays pushed a commit to johnnysluckydays/vllm-ascend that referenced this pull request Sep 14, 2026
)

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>
like-0517 pushed a commit to like-0517/vllm-ascend that referenced this pull request Sep 15, 2026
### 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>
like-0517 pushed a commit to like-0517/vllm-ascend that referenced this pull request Sep 15, 2026
)

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>
yiz-liu pushed a commit that referenced this pull request Sep 18, 2026
…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>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

module:tests ready-precise run selected e2e test for pr

Projects

None yet

Development

Successfully merging this pull request may close these issues.

7 participants