[BugFix] fix dflash spec decode rope and graph error

[HF-001]

What this PR does / why we need it?

Fix the bug in the profile of dummy_run. If the subsequent parameters are unreasonable, vllm will verify and prompt.
It is recommended to set reasonable parameters --max-num-batched-tokens when users start the model on Ascend currently.
Resolve the issue of error during dflash startup, specifically:
error 1:

File "<eval_with_key>.900 from /usr/local/python3.12.13/lib/python3.12/site-packages/torch/fx/experimental/proxy_tensor.py:1507 in wrapped", line 32, in forward
(Worker_TP0 pid=544) ERROR 07-27 10:06:06 [multiproc_executor.py:1000]     npu_rotary_embedding = torch.ops.vllm.npu_rotary_embedding.default(arg7_1, view_1, view_3, arg9_1, 128, 128, True);  arg7_1 = view_1 = view_3 = arg9_1 = None
(Worker_TP0 pid=544) ERROR 07-27 10:06:06 [multiproc_executor.py:1000]                            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
(Worker_TP0 pid=544) ERROR 07-27 10:06:06 [multiproc_executor.py:1000]   File "/usr/local/python3.12.13/lib/python3.12/site-packages/torch/_ops.py", line 819, in __call__
(Worker_TP0 pid=544) ERROR 07-27 10:06:06 [multiproc_executor.py:1000]     return self._op(*args, **kwargs)
(Worker_TP0 pid=544) ERROR 07-27 10:06:06 [multiproc_executor.py:1000]            ^^^^^^^^^^^^^^^^^^^^^^^^^
(Worker_TP0 pid=544) ERROR 07-27 10:06:06 [multiproc_executor.py:1000]   File "/vllm-workspace/vllm-ascend/vllm_ascend/ops/rotary_embedding.py", line 169, in rope_forward_oot
(Worker_TP0 pid=544) ERROR 07-27 10:06:06 [multiproc_executor.py:1000]     query, key = rope_forward_triton(
(Worker_TP0 pid=544) ERROR 07-27 10:06:06 [multiproc_executor.py:1000]                  ^^^^^^^^^^^^^^^^^^^^
(Worker_TP0 pid=544) ERROR 07-27 10:06:06 [multiproc_executor.py:1000]   File "/vllm-workspace/vllm-ascend/vllm_ascend/ops/triton/rope.py", line 287, in rope_forward_triton
(Worker_TP0 pid=544) ERROR 07-27 10:06:06 [multiproc_executor.py:1000]     assert positions.shape[0] == num_tokens

error 2:

(Worker_TP0 pid=18307) ERROR 07-28 02:12:55 [multiproc_executor.py:1000]   File "/vllm-workspace/vllm/vllm/model_executor/models/qwen3_dflash.py", line 438, in forward
(Worker_TP0 pid=18307) ERROR 07-28 02:12:55 [multiproc_executor.py:1000]     def forward(
(Worker_TP0 pid=18307) ERROR 07-28 02:12:55 [multiproc_executor.py:1000]   File "/vllm-workspace/vllm/vllm/compilation/caching.py", line 217, in __call__
(Worker_TP0 pid=18307) ERROR 07-28 02:12:55 [multiproc_executor.py:1000]     return self.optimized_call(*args, **kwargs)
(Worker_TP0 pid=18307) ERROR 07-28 02:12:55 [multiproc_executor.py:1000]            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
(Worker_TP0 pid=18307) ERROR 07-28 02:12:55 [multiproc_executor.py:1000]   File "<string>", line 25, in execution_fn
(Worker_TP0 pid=18307) ERROR 07-28 02:12:55 [multiproc_executor.py:1000]   File "/vllm-workspace/vllm-ascend/vllm_ascend/compilation/acl_graph.py", line 145, in __call__
(Worker_TP0 pid=18307) ERROR 07-28 02:12:55 [multiproc_executor.py:1000]     return self.runnable(*args, **kwargs)
(Worker_TP0 pid=18307) ERROR 07-28 02:12:55 [multiproc_executor.py:1000]            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
(Worker_TP0 pid=18307) ERROR 07-28 02:12:55 [multiproc_executor.py:1000]   File "/vllm-workspace/vllm/vllm/compilation/piecewise_backend.py", line 362, in __call__
(Worker_TP0 pid=18307) ERROR 07-28 02:12:55 [multiproc_executor.py:1000]     assert range_entry is not None, (
(Worker_TP0 pid=18307) ERROR 07-28 02:12:55 [multiproc_executor.py:1000]            ^^^^^^^^^^^^^^^^^^^^^^^
(Worker_TP0 pid=18307) ERROR 07-28 02:12:55 [multiproc_executor.py:1000] AssertionError: Shape: 2304 out of considered ranges: [(1, 2048)]

How was this patch tested?

--max-num-batched-tokens default value is 2048 ，which caused error1 and error2, Now there will be errors in the verification layer

export ASCEND_RT_VISIBLE_DEVICES=0,1

vllm serve /nas/disk1/Qwen3-32B --port 8898 --dtype bfloat16  \
--tensor-parallel-size 2 --gpu-memory-utilization 0.8 \
--max-model-len 32768 --trust-remote-code \
--no-enable-prefix-caching \
--async-scheduling \
--compilation-config '{"max_cudagraph_capture_size": 128}' \
--speculative_config '{"method":"dflash", "model":"/nas/disk1/Qwen3-32B-DFlash","num_speculative_tokens":8,"draft_tensor_parallel_size":2}''

change --max-num-batched-tokens value greater than or equal to 2304，
2304 = 256*(1+8)=max_num_reqs*(1+k)

export ASCEND_RT_VISIBLE_DEVICES=0,1

vllm serve /nas/disk1/Qwen3-32B --port 8898 --dtype bfloat16  \
--tensor-parallel-size 2 --gpu-memory-utilization 0.8 \
--max-model-len 32768 --trust-remote-code \
--no-enable-prefix-caching \
--async-scheduling \
--max-num-batched-tokens 2304 \
--compilation-config '{"max_cudagraph_capture_size": 128}' \
--speculative_config '{"method":"dflash", "model":"/nas/disk1/Qwen3-32B-DFlash","num_speculative_tokens":8,"draft_tensor_parallel_size":2}''

• vLLM version: v0.26.0
• vLLM main: vllm-project/vllm@0351e9a

[github-actions]

👋 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 [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:

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

Summary of Changes

Hello, I'm Gemini Code Assist¹! 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 runtime error encountered during dflash startup where RoPE operations failed due to buffer size mismatches. By dynamically calculating the required buffer length as the maximum of query tokens and total tokens, the changes ensure that buffers are sufficiently sized to handle padded inputs, preventing truncation and subsequent assertion failures.

Highlights

• Buffer Size Adjustment: Increased the size of position and slot mapping buffers in both DFlash and DSpark proposers to accommodate padded query token counts.
• RoPE Error Resolution: Prevented silent truncation of position buffers during RoPE operations by ensuring buffer lengths are at least as large as the maximum number of tokens.

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]

Code Review

Suggested PR Title:

[spec_decode][BugFix] Prevent RoPE failure by sizing query buffers to cover padded graph size

Suggested PR Summary:

### What this PR does / why we need it?

This PR fixes a potential RoPE failure during speculative decoding on Ascend. The query-side buffers `self.positions` and `self._slot_mapping_buffer` in both `AscendDflashProposer` and `AscendDsparkProposer` were previously sized to `max_query_tokens`. However, during execution, these buffers are sliced by `num_input_tokens`, which is padded up to the runner's CUDA graph capture size (bounded by `max_num_tokens`). If `num_input_tokens` exceeds `max_query_tokens`, slicing would silently truncate the buffers, causing RoPE to fail its shape validation check (`positions.shape[0] == num_tokens`).

To resolve this, the buffers are now sized to `query_buffer_len = max(self.max_query_tokens, self.max_num_tokens)`, ensuring they can safely accommodate both valid query data and any padded graph size.

### Does this PR introduce _any_ user-facing change?

No. This is an internal bug fix for speculative decoding on Ascend devices.

### How was this patch tested?

No specific tests were added in this PR, but the changes prevent runtime crashes during CUDA graph execution with speculative decoding.

I have no further feedback to provide as there are no review comments to evaluate.

[HF-001]

@wangxiyuan @yiz-liu @drslark hi, this pr is ready ,cloud you please take a look?

[ZT-AIA]

@lilinsiman Please review this PR. @yiz-liu Please confirm whether not specifying max_cudagraph_capture_size will result in compiling many graphs, consuming a large amount of resources, and ultimately causing an error.

[HF-001]

@yiz-liu Recently, I have been running mtp+tp2, dflash+tp2, minimax+tp8, and qwen3+tp8 on vllm-ascend v0.23.0rc1-a3+910c, and I have encountered errors. After troubleshooting, I found that the graph complex stage consumed too many resources, causing the service to crash. Currently, reducing resource consumption through the -- compilation config '{"max_cudagraph_capture size": 128}' parameter can solve the errors in these scenarios, but adjusting the parameters is only a temporary compromise solution. It should be caused by an unreasonable default -- compilation config parameter.This issue is not directly related to this PR

[HF-001]

@ZT-AIA @lilinsiman Error 1 and Error 2 are both caused by the --max-num-batched-tokens being too small, so they can be resolved by adding parameter validation logic.

[HF-001]

@yiz-liu @drslark hi, this pr is ready ,cloud you please take a look?

[github-actions]

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

[slippersss]

Can you post the complete error logs? Is it only about dflash? I'm not sure about this root cause.

[drslark]

Can this problem reproduce in gpu vllm?

[HF-001]

@slippersss @drslark On GPU, the default max_num_batched_tokens value is 8k, while on Ascend it is 2k, so there is an issue with Ascend. only dflash has this problem. The log is not easy to export, but it can be reproduced through the following script：

export ASCEND_RT_VISIBLE_DEVICES=0,1

vllm serve /nas/disk1/Qwen3-32B --port 8898 --dtype bfloat16
--tensor-parallel-size 2 --gpu-memory-utilization 0.8
--max-model-len 32768 --trust-remote-code
--no-enable-prefix-caching
--async-scheduling
--compilation-config '{"max_cudagraph_capture_size": 128}' \ ## can del this line
--speculative_config '{"method":"dflash", "model":"/nas/disk1/Qwen3-32B-DFlash","num_speculative_tokens":8,"draft_tensor_parallel_size":2}''

[HF-001]

@slippersss thank you, you are right. The best solution is num_query_total to num_query_tokens in dummy_run is_profile branch

[HF-001]

1. If it can reproduce in vllm. I think you should contribute to vllm.
2. If num_query_total is incorrect, it is a bug; we should fix it rather than skip it.
   So, I don't think put this validation in vllm-ascend is necessary in any case.

@drslark has changed to fix this bug in dummy_run

[slippersss]

LGTM. Thanks for your contribution!

[HF-001]

@slippersss @drslark I have fixed the verification logic of dflash in VLLM, and PR has been approved. After testing, it has been confirmed that only dummy_run needs to be modified, and there is no need to add warnings. vllm-project/vllm#51256

[drslark]

It is the root cause.
@slippersss @HF-001
Thanks for your efforts.
LGTM.

[slippersss]

@slippersss @drslark I have fixed the verification logic of dflash in VLLM, and PR has been approved. After testing, it has been confirmed that only dummy_run needs to be modified, and there is no need to add warnings. vllm-project/vllm#51256

Can we also remain similar modification in dspark? I think there is still semantically error in dspark.

[HF-001]

@slippersss @drslark I have fixed the verification logic of dflash in VLLM, and PR has been approved. After testing, it has been confirmed that only dummy_run needs to be modified, and there is no need to add warnings. vllm-project/vllm#51256

Can we also remain similar modification in dspark? I think there is still semantically error in dspark.

@slippersss I checked and found no need to modify in dspark. The num_input_tokens is indeed the standard forward length of the upstream proposer. However, in the current single dummy_run pipeline of Ascend DSpark, the runner has already passed num_reqs_padded, and DSpark operates in eager-only mode. Therefore, num_input_tokens typically equals num_query_total

[HF-001]

@ningjingbengxiaohai @slippersss @drslark @zzzzwwjj hi, this pr is ready. cloud you help merge it?