add qwen3_next ops: fused_qkvzba_split_reshape #4747
Conversation
![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 introduces new Triton-based operators, fused_qkvzba_split_reshape and rope_forward_triton, to optimize performance for the qwen3_next model on Ascend hardware. The changes include a new Triton kernel for fused QKVZBA splitting and reshaping, and a new RoPE implementation using Triton. The review identifies two critical issues: a potential bug in rotary_embedding.py that might break graph fusion due to a removed return value assignment, and an import typo in patch_qwen3_next.py that will lead to a runtime error.
vllm_ascend/ops/rotary_embedding.py
Outdated
| key = key.contiguous().view(1, key.shape[0], -1, self.head_size) | ||
| # If cos and sin are generated outside, use npu_apply_rotary_pos_emb to avoid redundant calculation. | ||
| # This method requires head_size and rotary_dim equal 128 and neox_style is True | ||
| torch_npu.npu_apply_rotary_pos_emb(query, key, self.cos, self.sin) |
There was a problem hiding this comment.
The previous implementation assigned the result of torch_npu.npu_apply_rotary_pos_emb back to query, key. A comment in the original code indicated this was crucial for graph fusion to work correctly, even though the function modifies tensors in-place. This assignment has been removed. If this is no longer a concern, it would be helpful to add a comment explaining why. Otherwise, this could be a critical bug and the assignment should be restored.
| torch_npu.npu_apply_rotary_pos_emb(query, key, self.cos, self.sin) | |
| query, key = torch_npu.npu_apply_rotary_pos_emb(query, key, self.cos, self.sin) |
| from vllm.model_executor.models.qwen3_next import Qwen3NextGatedDeltaNet | ||
| from vllm_ascend.ops.triton.fla.fused_qkvzba_split_reshape import fused_qkvzba_split_reshape_cat | ||
| from vllm.triton_utils import tl, triton | ||
| from vll.config import (CUDAGraphMode, get_current_vllm_config) |
There was a problem hiding this comment.
There is a typo in the import path. It should be vllm.config instead of vll.config. This will cause an ImportError at runtime.
| from vll.config import (CUDAGraphMode, get_current_vllm_config) | |
| from vllm.config import (CUDAGraphMode, get_current_vllm_config) |
|
👋 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. |
ZT-AIA
changed the title
1 participant
What this PR does / why we need it?
Does this PR introduce any user-facing change?
How was this patch tested?