[BugFix][Mamba] Fix staged-write kernel and duplicate is_prefilling argument - #14649
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 addresses critical stability and runtime issues in the Ascend V2 Runner inference path. By correcting memory access patterns in the Mamba state copy kernel and resolving argument collisions in attention metadata, the changes ensure reliable execution across different inference workloads. Additionally, it integrates NPU-optimized sampling operations to improve performance. 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
|
![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:
[Ops][Feature] Implement Ascend-safe Mamba pre-copy kernel and patch sampler/attention metadataSuggested PR Summary:
### What this PR does / why we need it?
This PR introduces several compatibility and safety improvements for running vLLM on Ascend NPU:
1. **Ascend-Safe Mamba Pre-copy Kernel**: Adds a custom Triton kernel `precopy_mamba_align_fused_kernel` in `vllm_ascend/ops/triton/mamba/precopy.py` and patches it in `patch_mamba_utils.py`. This implementation hoists pointer casts out of copy loops to prevent Triton-Ascend's AxisInfo analysis from aborting.
2. **Top-K/Top-P Sampler Patching**: Patches `topk_topp_sampler.apply_top_k_top_p` to use the NPU-specific `_apply_top_k_top_p_torch_npu` implementation.
3. **Attention Metadata Fix**: Safely extracts `is_prefilling` from `common_attn_metadata_extra_kwargs` in `build_attn_metadata` to avoid potential duplicate argument errors or incorrect states.
### Does this PR introduce _any_ user-facing change?
No. These are internal backend patches and kernel implementations for Ascend NPU support.
### How was this patch tested?
No specific tests were added in this PR. The changes are expected to be verified via existing integration tests on Ascend hardware.
|
👋 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! |
| from vllm.v1.sample.ops import topk_topp_sampler | ||
| from vllm_ascend.sample.sampler import _apply_top_k_top_p_torch_npu | ||
|
|
||
| topk_topp_sampler.apply_top_k_top_p = _apply_top_k_top_p_torch_npu |
There was a problem hiding this comment.
apply_top_k_top_p has been fixed on ascend. No need to do patch here.
There was a problem hiding this comment.
Thanks for pointing this out. I removed the redundant patch and pushed a new commit.
fc99029 to
3fdcbf4
Compare
LoganJane
added
ready
ready-precise
| common_is_prefilling = common_attn_metadata_extra_kwargs.pop( | ||
| "is_prefilling", | ||
| is_prefilling, | ||
| ) | ||
| common_attn_metadata = AscendCommonAttentionMetadata( |
There was a problem hiding this comment.
Please confirm the reason why the value is repeatedly transferred.
There was a problem hiding this comment.
Thanks for pointing this out. This follows the same pattern already used upstream in vllm/v1/worker/gpu/attn_utils.py around line 545:
group_is_prefilling = common_attn_metadata_extra_kwargs.pop(
"is_prefilling", is_prefilling
)
It allows the model-specific, padding-aware is_prefilling value to take precedence over the default, so I believe this change is still necessary.
| is_prefilling: torch.Tensor | None = None, | ||
| pcp_context: "AscendPCPAttentionContext | None" = None, | ||
| model_specific_attn_metadata: ModelSpecificAttnMetadata | None = None, |
There was a problem hiding this comment.
Maybe we should check out why there is is_prefilling in both here and model_specific_attn_metadata.
There was a problem hiding this comment.
Thanks for pointing this out. This follows the same pattern already used upstream in vllm/v1/worker/gpu/attn_utils.py around line 545:
group_is_prefilling = common_attn_metadata_extra_kwargs.pop(
"is_prefilling", is_prefilling
)
It allows the model-specific, padding-aware is_prefilling value to take precedence over the default, so I believe this change is still necessary.
There was a problem hiding this comment.
LGTM, thanks for your effort. CC @weijinqian0.
There was a problem hiding this comment.
BTW, we'd better change common_is_prefilling to group_is_prefilling, to align with vllm.
|
This pull request has conflicts, please resolve those before we can evaluate the pull request. |
c928cf3 to
4133f47
Compare
Signed-off-by: Lethobenthos20 <3205604419@qq.com>
Signed-off-by: Lethobenthos20 <3205604419@qq.com>
Signed-off-by: Lethobenthos20 <3205604419@qq.com>
Signed-off-by: Lethobenthos20 <3205604419@qq.com>
4133f47 to
a50b9b0
Compare
…rgument (vllm-project#14649) ## Summary This PR fixes two issues in the Ascend V2 Runner inference path: 1. Fixes incorrect index/address offset calculation in the block-table staged-write Triton kernel, which caused: ```text PtroffsetInfo::AxisInfo::operator[]: Assertion `idx < size()` failed 2. Fixes the duplicate is_prefilling argument passed to AscendCommonAttentionMetadata, which caused: ```text got multiple values for keyword argument 'is_prefilling' ``` After these changes, V2 Runner inference works correctly with curl and aisbench. - vLLM version: v0.27.1 - vLLM main: vllm-project/vllm@ba07e4a --------- Signed-off-by: Lethobenthos20 <3205604419@qq.com> Signed-off-by: QiuChunshuo <qiuchunshuo@huawei.com>
…rgument (vllm-project#14649) ## Summary This PR fixes two issues in the Ascend V2 Runner inference path: 1. Fixes incorrect index/address offset calculation in the block-table staged-write Triton kernel, which caused: ```text PtroffsetInfo::AxisInfo::operator[]: Assertion `idx < size()` failed 2. Fixes the duplicate is_prefilling argument passed to AscendCommonAttentionMetadata, which caused: ```text got multiple values for keyword argument 'is_prefilling' ``` After these changes, V2 Runner inference works correctly with curl and aisbench. - vLLM version: v0.27.1 - vLLM main: vllm-project/vllm@ba07e4a --------- Signed-off-by: Lethobenthos20 <3205604419@qq.com>
…rgument (vllm-project#14649) ## Summary This PR fixes two issues in the Ascend V2 Runner inference path: 1. Fixes incorrect index/address offset calculation in the block-table staged-write Triton kernel, which caused: ```text PtroffsetInfo::AxisInfo::operator[]: Assertion `idx < size()` failed 2. Fixes the duplicate is_prefilling argument passed to AscendCommonAttentionMetadata, which caused: ```text got multiple values for keyword argument 'is_prefilling' ``` After these changes, V2 Runner inference works correctly with curl and aisbench. - vLLM version: v0.27.1 - vLLM main: vllm-project/vllm@ba07e4a --------- Signed-off-by: Lethobenthos20 <3205604419@qq.com> Signed-off-by: d30086105 <denghaojie1@h-partners.com>
### What this PR does / why we need it? This PR fixes the following Triton Ascend assertion failure in the Mamba state pre-copy kernel: ```text PtrOffsetInfo::AxisInfo::operator[]: Assertion `idx < size()` failed ``` This issue was also observed in #7103 and #10888 . This PR also: - Adds an end-to-end test covering the supported layouts and execution modes. - Adds documentation for the kernel behavior, parameters, constraints, and usage. This is a follow-up to #14649. After these changes, V2 Runner inference works correctly. - vLLM version: v0.27.1 - vLLM main: vllm-project/vllm@ba07e4a --------- Signed-off-by: Lethobenthos20 <3205604419@qq.com>
…oject#15032) ### What this PR does / why we need it? This PR fixes the following Triton Ascend assertion failure in the Mamba state pre-copy kernel: ```text PtrOffsetInfo::AxisInfo::operator[]: Assertion `idx < size()` failed ``` This issue was also observed in vllm-project#7103 and vllm-project#10888 . This PR also: - Adds an end-to-end test covering the supported layouts and execution modes. - Adds documentation for the kernel behavior, parameters, constraints, and usage. This is a follow-up to vllm-project#14649. After these changes, V2 Runner inference works correctly. - vLLM version: v0.27.1 - vLLM main: vllm-project/vllm@ba07e4a --------- Signed-off-by: Lethobenthos20 <3205604419@qq.com> Signed-off-by: d30086105 <denghaojie1@h-partners.com>
…rgument (vllm-project#14649) ## Summary This PR fixes two issues in the Ascend V2 Runner inference path: 1. Fixes incorrect index/address offset calculation in the block-table staged-write Triton kernel, which caused: ```text PtroffsetInfo::AxisInfo::operator[]: Assertion `idx < size()` failed 2. Fixes the duplicate is_prefilling argument passed to AscendCommonAttentionMetadata, which caused: ```text got multiple values for keyword argument 'is_prefilling' ``` After these changes, V2 Runner inference works correctly with curl and aisbench. - vLLM version: v0.27.1 - vLLM main: vllm-project/vllm@ba07e4a --------- Signed-off-by: Lethobenthos20 <3205604419@qq.com>
…oject#15032) ### What this PR does / why we need it? This PR fixes the following Triton Ascend assertion failure in the Mamba state pre-copy kernel: ```text PtrOffsetInfo::AxisInfo::operator[]: Assertion `idx < size()` failed ``` This issue was also observed in vllm-project#7103 and vllm-project#10888 . This PR also: - Adds an end-to-end test covering the supported layouts and execution modes. - Adds documentation for the kernel behavior, parameters, constraints, and usage. This is a follow-up to vllm-project#14649. After these changes, V2 Runner inference works correctly. - vLLM version: v0.27.1 - vLLM main: vllm-project/vllm@ba07e4a --------- Signed-off-by: Lethobenthos20 <3205604419@qq.com>
5 participants
Summary
This PR fixes two issues in the Ascend V2 Runner inference path:
Fixes incorrect index/address offset calculation in the block-table staged-write Triton kernel, which caused:
Fixes the duplicate is_prefilling argument passed to AscendCommonAttentionMetadata, which caused:
After these changes, V2 Runner inference works correctly with curl and aisbench.