[perf] Fused operator SplitMrope used in the Qwen2.5-Omni-7B model #31763
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 a fused operator npu_split_mrope for Qwen2 models to improve performance on Ascend NPUs. However, the implementation has several critical issues. It unconditionally uses an Ascend-specific operator, which will cause failures on other hardware like NVIDIA GPUs. Additionally, the QK normalization feature has been removed, which is a correctness issue for models that use it. There is also a redundant qkv.split operation that should be removed. The changes need to be reworked to be conditional on the hardware and to preserve existing functionality.
| # q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) | ||
|
|
||
| # Apply QK normalization if enabled (before RoPE) | ||
| if self.qk_norm: | ||
| # Reshape to apply per-head normalization | ||
| # q shape: (total_tokens, q_size) -> (total_tokens, num_heads, head_dim) | ||
| total_tokens = q.shape[0] | ||
| q = q.view(total_tokens, self.num_heads, self.head_dim) | ||
| k = k.view(total_tokens, self.num_kv_heads, self.head_dim) | ||
|
|
||
| # Apply normalization | ||
| q = self.q_norm(q) | ||
| k = self.k_norm(k) | ||
|
|
||
| # Reshape back | ||
| q = q.view(total_tokens, self.q_size) | ||
| k = k.view(total_tokens, self.kv_size) | ||
| # if self.qk_norm: | ||
| # # Reshape to apply per-head normalization | ||
| # # q shape: (total_tokens, q_size) -> (total_tokens, num_heads, head_dim) | ||
| # total_tokens = q.shape[0] | ||
| # q = q.view(total_tokens, self.num_heads, self.head_dim) | ||
| # k = k.view(total_tokens, self.num_kv_heads, self.head_dim) | ||
| # | ||
| # # Apply normalization | ||
| # q = self.q_norm(q) | ||
| # k = self.k_norm(k) | ||
| # | ||
| # # Reshape back | ||
| # q = q.view(total_tokens, self.q_size) | ||
| # k = k.view(total_tokens, self.kv_size) | ||
| # | ||
| # q, k = self.rotary_emb(positions, q, k) | ||
|
|
||
| q, k, v = torch.ops._C_ascend.npu_split_mrope( | ||
| positions.contiguous(), | ||
| qkv, | ||
| self.rotary_emb.cos_sin_cache, | ||
| self.q_size, | ||
| self.kv_size, | ||
| self.num_heads, | ||
| self.num_kv_heads, | ||
| self.head_dim, | ||
| mrope_section=self.rotary_emb.mrope_section, | ||
| rotary_mode="half" | ||
| ) |
There was a problem hiding this comment.
This change introduces a fused operator for performance on Ascend NPUs, but it has a few issues:
- Redundant Operation: The
qkv.split()on line 207 is unnecessary as the returnedq,k, andvare immediately overwritten. Thenpu_split_mropeoperator likely handles this split internally. - Correctness Bug: The QK normalization logic (
if self.qk_norm: ...) is commented out. This will produce incorrect results for models that require it. The new fused operator does not seem to support this feature. - Portability Issue: The code unconditionally calls an Ascend NPU-specific operator (
torch.ops._C_ascend.npu_split_mrope), which will cause vLLM to fail on other hardware like NVIDIA GPUs.
To address these issues, the implementation should be conditional. It should check for the specific hardware and whether qk_norm is disabled before using the fused operator. For all other cases, it should fall back to the original, general implementation. This also resolves the redundant split.
if qkv.device.type == "npu" and not self.qk_norm:
q, k, v = torch.ops._C_ascend.npu_split_mrope(
positions.contiguous(),
qkv,
self.rotary_emb.cos_sin_cache,
self.q_size,
self.kv_size,
self.num_heads,
self.num_kv_heads,
self.head_dim,
mrope_section=self.rotary_emb.mrope_section,
rotary_mode="half"
)
else:
q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
# Apply QK normalization if enabled (before RoPE)
if self.qk_norm:
# Reshape to apply per-head normalization
# q shape: (total_tokens, q_size) -> (total_tokens, num_heads, head_dim)
total_tokens = q.shape[0]
q = q.view(total_tokens, self.num_heads, self.head_dim)
k = k.view(total_tokens, self.num_kv_heads, self.head_dim)
# Apply normalization
q = self.q_norm(q)
k = self.k_norm(k)
# Reshape back
q = q.view(total_tokens, self.q_size)
k = k.view(total_tokens, self.kv_size)
q, k = self.rotary_emb(positions, q, k)
|
👋 Hi! Thank you for contributing to the vLLM project. 💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in #pr-reviews, coordinate on features in #feat- channels, or join special interest groups in #sig- channels. Just a reminder: PRs would not trigger full CI run by default. Instead, it would only run You ask your reviewers to trigger select CI tests on top of Once the PR is approved and ready to go, your PR reviewer(s) can run CI to test the changes comprehensively before merging. To run CI, PR reviewers can either: Add If you have any questions, please reach out to us on Slack at https://slack.vllm.ai. 🚀 |
1 participant
Purpose
Test Plan
Test Result
Essential Elements of an Effective PR Description Checklist
supported_models.mdandexamplesfor a new model.