[Bugfix][NPU] Fallback to BSND RoPE when Qwen3-TTS short-seq BNSD li… (#5608)

Signed-off-by: XIN GAO <1037396230@qq.com>

tests/platforms/npu/test_310p_patches.py

@@ -489,6 +489,14 @@ def test_qwen3_tts_tokenizer_npu_patch_dispatches_fused_ops(monkeypatch: pytest.
     _install_fake_module(monkeypatch, "vllm")
     _install_fake_module(monkeypatch, "vllm.logger", init_logger=lambda _name: SimpleNamespace(debug=lambda *_: None))
 
+    rotary_path = _repo_root() / "vllm_omni" / "platforms" / "npu" / "layers" / "rotary_embedding.py"
+    rotary_module = _load_source_module("vllm_omni_test_npu_rotary_embedding", rotary_path)
+    monkeypatch.setitem(
+        sys.modules,
+        "vllm_omni.platforms.npu.layers.rotary_embedding",
+        rotary_module,
+    )
+
     class FakeRMSNorm(torch.nn.Module):
         def __init__(self):
             super().__init__()
@@ -536,7 +544,7 @@ def test_qwen3_tts_tokenizer_npu_patch_dispatches_fused_ops(monkeypatch: pytest.
     norm_out = norm(torch.ones(1, 2))
 
     assert len(rotary_calls) == 2
-    assert rotary_calls[0][1].shape == (1, 1, 3, 2)
+    assert rotary_calls[0][1].shape == (1, 3, 1, 2)
     torch.testing.assert_close(q_out, q + 1)
     torch.testing.assert_close(k_out, k + 1)
     assert len(rms_calls) == 1

vllm_omni/platforms/npu/layers/rotary_embedding.py

@@ -0,0 +1,57 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""Shared NPU rotary embedding helpers."""
+
+from __future__ import annotations
+
+import torch
+import torch_npu
+
+
+def _bnsd_rotary_shape_is_supported(hidden_states: torch.Tensor) -> bool:
+    """Return whether the fused half-mode BNSD tiler supports the shape.
+
+    For the aligned-head-dimension branch, CANN requires
+    ``batch * num_heads <= seq_len * 8``. Other alignments are handled with
+    the BSND layout, which does not have this short-sequence limit.
+    """
+    if hidden_states.ndim != 4:
+        return False
+
+    batch_size, num_heads, seq_len, head_dim = hidden_states.shape
+    element_size = hidden_states.element_size()
+    if head_dim % 2 != 0 or element_size <= 0 or 32 % element_size != 0:
+        return False
+
+    half_dim_alignment = 32 // element_size
+    return head_dim // 2 % half_dim_alignment == 0 and batch_size * num_heads <= seq_len * 8
+
+
+def npu_rotary_mul_with_bsnd_fallback(
+    hidden_states: torch.Tensor,
+    cos: torch.Tensor,
+    sin: torch.Tensor,
+    unsqueeze_dim: int,
+) -> torch.Tensor:
+    """Apply fused rotary multiplication, falling back from BNSD to BSND.
+
+    ``unsqueeze_dim=1`` denotes BNSD input following the Transformers rotary
+    embedding convention. Supported BNSD shapes stay on the fast path. Shapes
+    rejected by the CANN BNSD tiler are transposed to BSND for the fused op and
+    transposed back before returning.
+    """
+    if unsqueeze_dim != 1 or _bnsd_rotary_shape_is_supported(hidden_states):
+        return torch_npu.npu_rotary_mul(
+            hidden_states,
+            cos.unsqueeze(unsqueeze_dim),
+            sin.unsqueeze(unsqueeze_dim),
+        )
+
+    hidden_states_bsnd = hidden_states.transpose(1, 2).contiguous()
+    output_bsnd = torch_npu.npu_rotary_mul(
+        hidden_states_bsnd,
+        cos.unsqueeze(2),
+        sin.unsqueeze(2),
+    )
+    return output_bsnd.transpose(1, 2)

vllm_omni/platforms/npu/models/qwen3_tts_tokenizer_v2.py

@@ -9,6 +9,10 @@ import torch
 import torch_npu
 from vllm.logger import init_logger
 
+from vllm_omni.platforms.npu.layers.rotary_embedding import (
+    npu_rotary_mul_with_bsnd_fallback,
+)
+
 logger = init_logger(__name__)
 
 _PATCHED = False
@@ -23,9 +27,10 @@ def _apply_rotary_pos_emb_npu(
     unsqueeze_dim: int = 1,
 ) -> tuple[torch.Tensor, torch.Tensor]:
     del position_ids
-    cos = cos.unsqueeze(unsqueeze_dim)
-    sin = sin.unsqueeze(unsqueeze_dim)
-    return torch_npu.npu_rotary_mul(q, cos, sin), torch_npu.npu_rotary_mul(k, cos, sin)
+    return (
+        npu_rotary_mul_with_bsnd_fallback(q, cos, sin, unsqueeze_dim),
+        npu_rotary_mul_with_bsnd_fallback(k, cos, sin, unsqueeze_dim),
+    )
 
 
 def _rms_norm_forward_npu(self, hidden_states: torch.Tensor) -> torch.Tensor: