[Bugfix][Ascend] fix npu 310/a5 bugs (#6685)

Signed-off-by: zouyizhou <zouyizhou@huawei.com>

tests/platforms/npu/test_310p_patches.py

@@ -431,9 +431,17 @@ def test_310p_attention_forward_runs_fused_qkv_runtime(monkeypatch: pytest.Monke
         trans_calls.append(weight.detach().clone())
         return weight + 1.0
 
-    def npu_rotary_mul(hidden_states: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
-        rotary_calls.append((hidden_states.detach().clone(), cos.detach().clone(), sin.detach().clone()))
-        return hidden_states + cos - sin
+    def npu_apply_rotary_pos_emb(
+        query: torch.Tensor,
+        key: torch.Tensor,
+        cos: torch.Tensor,
+        sin: torch.Tensor,
+        *,
+        layout: str = "BSND",
+        rotary_mode: str = "half",
+    ):
+        rotary_calls.append((query.detach().clone(), key.detach().clone(), cos.detach().clone(), sin.detach().clone()))
+        return query + cos - sin, key + cos - sin
 
     def npu_flash_attention(**kwargs) -> None:
         flash_calls.append(
@@ -443,7 +451,7 @@ def test_310p_attention_forward_runs_fused_qkv_runtime(monkeypatch: pytest.Monke
 
     monkeypatch.setattr(module, "maybe_trans_nz", maybe_trans_nz)
     monkeypatch.setattr(module, "aligned_16", lambda tensor: tensor)
-    monkeypatch.setattr(module.torch_npu, "npu_rotary_mul", npu_rotary_mul, raising=False)
+    monkeypatch.setattr(module.torch_npu, "npu_apply_rotary_pos_emb", npu_apply_rotary_pos_emb, raising=False)
     monkeypatch.setattr(module.torch_npu, "_npu_flash_attention", npu_flash_attention, raising=False)
 
     config = SimpleNamespace(
@@ -503,7 +511,12 @@ def test_310p_attention_forward_runs_fused_qkv_runtime(monkeypatch: pytest.Monke
     expected_v_f = expected_v.transpose(1, 2).reshape(3, attention.num_kv_heads, attention.head_dim)
     expected_flash = (expected_q_f + expected_k_f + expected_v_f).to(torch.float16).to(torch.float32)
 
-    assert len(rotary_calls) == 2
+    # One fused aclnn rope call covers q and k in the 4D BSND layout.
+    assert len(rotary_calls) == 1
+    assert rotary_calls[0][0].shape == (1, 3, attention.num_heads, attention.head_dim)
+    assert rotary_calls[0][1].shape == (1, 3, attention.num_kv_heads, attention.head_dim)
+    assert rotary_calls[0][2].shape == (1, 3, 1, attention.head_dim)
+    assert rotary_calls[0][3].shape == (1, 3, 1, attention.head_dim)
     assert len(flash_calls) == 1
     torch.testing.assert_close(flash_calls[0]["query"], expected_q_f)
     torch.testing.assert_close(flash_calls[0]["key"], expected_k_f)
@@ -535,39 +548,16 @@ def test_qwen3_tts_talker_patch_uses_fp16_runtime_dtype(monkeypatch: pytest.Monk
     assert codes[0].shape == (4, 2)
 
 
-def test_qwen3_tts_prompt_patch_runs_stft_frontend_on_cpu(monkeypatch: pytest.MonkeyPatch) -> None:
+def test_qwen3_tts_prompt_patch_sets_cpu_mel_front_end_flag(monkeypatch: pytest.MonkeyPatch) -> None:
     module, _ = _load_qwen3_tts_patch(monkeypatch)
-    captured = {}
 
-    def fake_mel_spectrogram(wav_tensor, **kwargs):
-        captured["wav_device"] = wav_tensor.device
-        captured["wav_dtype"] = wav_tensor.dtype
-        captured["kwargs"] = kwargs
-        return torch.ones(1, 128, 3, dtype=torch.float32)
-
-    class FakeSpeakerEncoder(torch.nn.Module):
-        def __init__(self):
-            super().__init__()
-            self.param = torch.nn.Parameter(torch.zeros(1, dtype=torch.float16))
-
-        def forward(self, mels):
-            captured["speaker_input_dtype"] = mels.dtype
-            return (torch.ones(4, dtype=mels.dtype),)
-
-    monkeypatch.setattr(module.prompt_embeds_builder, "mel_spectrogram", fake_mel_spectrogram)
-    builder = object.__new__(module._Qwen3TTSPromptEmbedsBuilder310P)
-    builder._device = lambda: torch.device("cpu")
-    builder._embedding_dtype = torch.float16
-    builder._speaker_encoder = FakeSpeakerEncoder()
-    builder._config = SimpleNamespace(speaker_encoder_config=SimpleNamespace(sample_rate=24000))
-
-    speaker = builder.extract_speaker_embedding(np.zeros(16, dtype=np.float32), 24000)
-
-    assert captured["wav_device"] == torch.device("cpu")
-    assert captured["wav_dtype"] is torch.float32
-    assert captured["kwargs"]["sampling_rate"] == 24000
-    assert captured["speaker_input_dtype"] is torch.float16
-    assert speaker.dtype is torch.float16
+    # The 310P builder keeps the base implementation and only opts into the
+    # CPU mel front-end, because the NPU has no torch.stft.
+    assert issubclass(
+        module._Qwen3TTSPromptEmbedsBuilder310P,
+        module.prompt_embeds_builder.Qwen3TTSPromptEmbedsBuilder,
+    )
+    assert module._Qwen3TTSPromptEmbedsBuilder310P._mel_spectrogram_on_cpu is True
 
 
 def test_qwen3_tts_tokenizer_npu_patch_dispatches_fused_ops(monkeypatch: pytest.MonkeyPatch) -> None:
@@ -706,11 +696,10 @@ def test_qwen3_tts_code2wav_npu_patch_prepares_loaded_decoder(monkeypatch: pytes
     _install_fake_module(monkeypatch, "vllm_ascend.utils", maybe_trans_nz=maybe_trans_nz)
     _install_fake_module(monkeypatch, "vllm_omni")
     _install_fake_module(monkeypatch, "vllm_omni.platforms", current_omni_platform=current_platform)
-    _install_fake_module(monkeypatch, "vllm_omni.model_executor")
-    _install_fake_module(monkeypatch, "vllm_omni.model_executor.models")
-    _install_fake_module(monkeypatch, "vllm_omni.model_executor.models.qwen3_tts")
+    # Not A5: keep the FRACTAL_Z conv weight layout for 310P.
+    _install_fake_module(monkeypatch, "vllm_omni.platforms.npu", is_a5=lambda: False)
 
-    path = _repo_root() / "vllm_omni" / "platforms" / "npu" / "models" / "qwen3_tts_code2wav.py"
+    path = _repo_root() / "vllm_omni" / "platforms" / "npu" / "models" / "qwen3_tts.py"
     module = _load_source_module("vllm_omni_test_qwen3_tts_code2wav_npu_patch", path)
     module.apply_qwen3_tts_code2wav_patch()
 

tests/platforms/npu/test_a5_patches.py

@@ -0,0 +1,259 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Unit tests for A5 (Ascend 950 PR) patch wiring.
+
+The A5 logic lives in the shared NPU platform files (``npu/__init__.py`` for
+device detection, ``npu/models/qwen3_tts.py`` for the Qwen3-TTS
+patch). The tests load those modules from source with fake Qwen3-TTS
+dependencies, so they validate the patch contract without loading real model
+or NPU kernels. A5 device detection is exercised with a fake
+``vllm_ascend.utils`` module.
+"""
+
+from __future__ import annotations
+
+import importlib.util
+import sys
+import types
+from pathlib import Path
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def _repo_root() -> Path:
+    marker = Path("vllm_omni") / "platforms" / "npu" / "models"
+    for parent in Path(__file__).resolve().parents:
+        if (parent / marker).is_dir():
+            return parent
+    raise FileNotFoundError(f"could not locate repo root containing {marker}")
+
+
+def _load_source_module(name: str, path: Path):
+    spec = importlib.util.spec_from_file_location(name, path)
+    if spec is None or spec.loader is None:
+        raise RuntimeError(f"cannot load import spec for {path}")
+    module = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(module)
+    return module
+
+
+def _install_fake_module(monkeypatch: pytest.MonkeyPatch, name: str, **attrs):
+    module = types.ModuleType(name)
+    for key, value in attrs.items():
+        setattr(module, key, value)
+    monkeypatch.setitem(sys.modules, name, module)
+    return module
+
+
+def _load_npu_init(monkeypatch: pytest.MonkeyPatch):
+    _install_fake_module(monkeypatch, "vllm")
+    _install_fake_module(
+        monkeypatch,
+        "vllm.logger",
+        init_logger=lambda _name: SimpleNamespace(debug=lambda *_: None),
+    )
+    path = _repo_root() / "vllm_omni" / "platforms" / "npu" / "__init__.py"
+    return _load_source_module("vllm_omni_test_npu_init", path)
+
+
+def _load_qwen3_tts_code2wav_patch(
+    monkeypatch: pytest.MonkeyPatch,
+    *,
+    a5: bool,
+):
+    class FakePromptEmbedsBuilder:
+        pass
+
+    class FakeCode2Wav:
+        def __init__(self, *, vllm_config, prefix=""):
+            pass
+
+        def load_weights(self, weights):
+            pass
+
+    fake_prompt_builder = _install_fake_module(
+        monkeypatch,
+        "vllm_omni.model_executor.models.qwen3_tts.prompt_embeds_builder",
+        Qwen3TTSPromptEmbedsBuilder=FakePromptEmbedsBuilder,
+        mel_spectrogram=lambda *_args, **_kwargs: torch.empty(0),
+    )
+    fake_talker = _install_fake_module(
+        monkeypatch,
+        "vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_talker",
+        Qwen3TTSPromptEmbedsBuilder=FakePromptEmbedsBuilder,
+    )
+    _install_fake_module(
+        monkeypatch,
+        "vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_code2wav",
+        Qwen3TTSCode2Wav=FakeCode2Wav,
+    )
+    _install_fake_module(monkeypatch, "vllm_omni")
+    _install_fake_module(monkeypatch, "vllm_omni.model_executor")
+    _install_fake_module(monkeypatch, "vllm_omni.model_executor.models")
+    pkg_qwen3_tts = _install_fake_module(
+        monkeypatch,
+        "vllm_omni.model_executor.models.qwen3_tts",
+        prompt_embeds_builder=fake_prompt_builder,
+        qwen3_tts_talker=fake_talker,
+    )
+    # Mark the package fakes as real packages so the patch's lazy submodule
+    # imports (e.g. ``qwen3_tts_code2wav``) resolve inside the harness.
+    pkg_qwen3_tts.__path__ = []
+    _install_fake_module(
+        monkeypatch,
+        "vllm_omni.platforms",
+        current_omni_platform=SimpleNamespace(is_npu=lambda: False),
+    )
+    _install_fake_module(monkeypatch, "vllm_omni.platforms.npu", is_a5=lambda: a5)
+    _install_fake_module(monkeypatch, "vllm")
+    _install_fake_module(monkeypatch, "vllm.config", VllmConfig=object)
+    _install_fake_module(
+        monkeypatch,
+        "vllm.logger",
+        init_logger=lambda _name: SimpleNamespace(info=lambda *_: None, debug=lambda *_: None),
+    )
+    _install_fake_module(monkeypatch, "vllm_ascend")
+    _install_fake_module(monkeypatch, "vllm_ascend.utils", maybe_trans_nz=lambda weight: weight)
+    _install_fake_module(monkeypatch, "torch_npu", npu_format_cast=lambda weight, _fmt: weight)
+
+    path = _repo_root() / "vllm_omni" / "platforms" / "npu" / "models" / "qwen3_tts.py"
+    module = _load_source_module("vllm_omni_test_qwen3_tts_code2wav_patch", path)
+    return module, fake_prompt_builder, fake_talker
+
+
+def test_is_a5_detects_a5_device(monkeypatch: pytest.MonkeyPatch) -> None:
+    module = _load_npu_init(monkeypatch)
+    _install_fake_module(monkeypatch, "vllm_ascend")
+    _install_fake_module(monkeypatch, "vllm_ascend.utils", is_950=lambda: True)
+
+    assert module.is_a5() is True
+
+    class NpuDevice:
+        type = "npu"
+
+    assert module.is_a5(NpuDevice()) is True
+    assert module.is_a5(torch.device("cpu")) is False
+
+
+def test_is_a5_falls_back_false_without_is_950(monkeypatch: pytest.MonkeyPatch) -> None:
+    module = _load_npu_init(monkeypatch)
+    # Older vllm-ascend releases have no `is_950`; detection must fall back.
+    _install_fake_module(monkeypatch, "vllm_ascend")
+    _install_fake_module(monkeypatch, "vllm_ascend.utils")
+
+    assert module.is_a5(torch.device("cuda")) is False
+    assert module.is_a5() is False
+
+
+def test_qwen3_tts_patch_swaps_prompt_builder_on_a5(monkeypatch: pytest.MonkeyPatch) -> None:
+    module, fake_prompt_builder, fake_talker = _load_qwen3_tts_code2wav_patch(
+        monkeypatch,
+        a5=True,
+    )
+    original = fake_prompt_builder.Qwen3TTSPromptEmbedsBuilder
+
+    module.apply_qwen3_tts_patches()
+
+    assert fake_prompt_builder.Qwen3TTSPromptEmbedsBuilder is not original
+    assert fake_prompt_builder.Qwen3TTSPromptEmbedsBuilder is fake_talker.Qwen3TTSPromptEmbedsBuilder
+
+    module.apply_qwen3_tts_patches()
+    assert fake_prompt_builder.Qwen3TTSPromptEmbedsBuilder is fake_talker.Qwen3TTSPromptEmbedsBuilder
+
+
+def test_qwen3_tts_patch_keeps_prompt_builder_off_a5(monkeypatch: pytest.MonkeyPatch) -> None:
+    module, fake_prompt_builder, fake_talker = _load_qwen3_tts_code2wav_patch(
+        monkeypatch,
+        a5=False,
+    )
+    original = fake_prompt_builder.Qwen3TTSPromptEmbedsBuilder
+
+    module.apply_qwen3_tts_patches()
+
+    assert fake_prompt_builder.Qwen3TTSPromptEmbedsBuilder is original
+    assert fake_talker.Qwen3TTSPromptEmbedsBuilder is original
+
+
+def test_a5_prompt_patch_sets_cpu_mel_front_end_flag(monkeypatch: pytest.MonkeyPatch) -> None:
+    module, fake_prompt_builder, fake_talker = _load_qwen3_tts_code2wav_patch(
+        monkeypatch,
+        a5=True,
+    )
+    original = fake_prompt_builder.Qwen3TTSPromptEmbedsBuilder
+
+    module.apply_qwen3_tts_patches()
+
+    # The A5 builder keeps the base implementation and only opts into the CPU
+    # mel front-end, because the NPU has no torch.stft.
+    builder_cls = fake_prompt_builder.Qwen3TTSPromptEmbedsBuilder
+    assert builder_cls is not original
+    assert builder_cls._mel_spectrogram_on_cpu is True
+    assert fake_talker.Qwen3TTSPromptEmbedsBuilder is builder_cls
+
+
+def test_qwen3_tts_code2wav_patch_skips_fractal_z_conv_cast_on_a5(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    module, _, _ = _load_qwen3_tts_code2wav_patch(monkeypatch, a5=True)
+    linear_weights = []
+    conv_weights = []
+
+    monkeypatch.setattr(
+        module,
+        "maybe_trans_nz",
+        lambda weight: linear_weights.append(weight) or weight,
+    )
+    monkeypatch.setattr(
+        module.torch_npu,
+        "npu_format_cast",
+        lambda weight, fmt: conv_weights.append((weight, fmt)) or weight,
+    )
+
+    class FakeDecoder(torch.nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.linear = torch.nn.Linear(4, 4)
+            self.conv = torch.nn.Conv1d(4, 4, 3)
+            self.deconv = torch.nn.ConvTranspose1d(4, 4, 4)
+            self.grouped_conv = torch.nn.Conv1d(4, 4, 3, groups=2)
+            self.cache_precompute_calls = 0
+
+        def precompute_snake_caches(self):
+            self.cache_precompute_calls += 1
+
+    class FakeCode2Wav:
+        def __init__(self, *, vllm_config, prefix=""):
+            self.vllm_config = vllm_config
+            self.prefix = prefix
+            self.decoder = FakeDecoder()
+
+        def _npu_decoder_runtime_dtype(self, _device):
+            return torch.float16
+
+        def load_weights(self, weights):
+            assert list(weights) == []
+            return {"loaded"}
+
+    target = _install_fake_module(
+        monkeypatch,
+        "vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_code2wav",
+        Qwen3TTSCode2Wav=FakeCode2Wav,
+    )
+    module.apply_qwen3_tts_code2wav_patch()
+
+    model = target.Qwen3TTSCode2Wav(
+        vllm_config=SimpleNamespace(device_config=SimpleNamespace(device=torch.device("cpu"))),
+        prefix="stage1",
+    )
+    assert model.load_weights(iter(())) == {"loaded"}
+
+    assert model.prefix == "stage1"
+    assert model.decoder.linear.weight.dtype is torch.float16
+    assert [weight.data_ptr() for weight in linear_weights] == [model.decoder.linear.weight.data_ptr()]
+    # On A5 convs keep the contiguous ND layout: no FRACTAL_Z cast is issued.
+    assert conv_weights == []
+    assert model.decoder.cache_precompute_calls == 1

vllm_omni/model_executor/models/qwen3_tts/prompt_embeds_builder.py

@@ -643,6 +643,10 @@ class Qwen3TTSPromptEmbedsBuilder:
 
     # -------------------- speaker encoder / codec encoder --------------------
 
+    # Some NPUs (310P, Ascend 950) do not support torch.stft; builders on
+    # such devices set this flag to compute the mel spectrogram on CPU.
+    _mel_spectrogram_on_cpu: bool = False
+
     def extract_speaker_embedding(self, wav: np.ndarray, sr: int) -> torch.Tensor:
         # vLLM workers do not automatically move arbitrary torch.nn.Modules to
         # CUDA. Ensure the speaker encoder is on the same device/dtype as the
@@ -666,7 +670,9 @@ class Qwen3TTSPromptEmbedsBuilder:
         # Follow official implementation: mel_spectrogram expects 24kHz. Move
         # the waveform first so STFT/mel computation stays on the model device
         # instead of materializing a CPU mel tensor and copying it per request.
-        wav_tensor = torch.from_numpy(wav).to(device=dev, dtype=torch.float32).unsqueeze(0)
+        # Devices without torch.stft set _mel_spectrogram_on_cpu instead.
+        mel_device = torch.device("cpu") if self._mel_spectrogram_on_cpu else dev
+        wav_tensor = torch.from_numpy(wav).to(device=mel_device, dtype=torch.float32).unsqueeze(0)
         mels = mel_spectrogram(
             wav_tensor,
             n_fft=1024,
@@ -677,7 +683,7 @@ class Qwen3TTSPromptEmbedsBuilder:
             fmin=0,
             fmax=12000,
         ).transpose(1, 2)
-        spk = self._speaker_encoder(mels.to(dtype=dtype))[0]
+        spk = self._speaker_encoder(mels.to(device=dev, dtype=dtype))[0]
         return spk.to(dtype=dtype)
 
     def encode_ref_audio_batch(

vllm_omni/platforms/npu/_310p/__init__.py

@@ -6,12 +6,10 @@
 from __future__ import annotations
 
 import torch
-from vllm.logger import init_logger
-
-logger = init_logger(__name__)
 
 
 def is_310p(device: torch.device | None = None) -> bool:
+    """Return True on Ascend 310P devices, False otherwise."""
     if device is not None and device.type != "npu":
         return False
     try:
@@ -20,9 +18,6 @@ def is_310p(device: torch.device | None = None) -> bool:
         return bool(ascend_is_310p())
     except (ImportError, AttributeError):
         return False
-    except Exception:
-        logger.debug("Failed to detect Ascend 310P device.", exc_info=True)
-        return False
 
 
 def disable_jit_compile() -> None:

vllm_omni/platforms/npu/_310p/patch/qwen3_tts.py

@@ -81,35 +81,9 @@ class _Qwen3TTSTalker310P(qwen3_tts_talker.Qwen3TTSTalkerForConditionalGeneratio
 
 
 class _Qwen3TTSPromptEmbedsBuilder310P(prompt_embeds_builder.Qwen3TTSPromptEmbedsBuilder):
-    def extract_speaker_embedding(self, wav: np.ndarray, sr: int) -> torch.Tensor:
-        dev = self._device()
-        dtype = self._embedding_dtype
-        try:
-            spk_param = next(self._speaker_encoder.parameters())
-            if spk_param.device != dev or spk_param.dtype != dtype:
-                self._speaker_encoder.to(device=dev, dtype=dtype)
-        except StopIteration:
-            pass
+    """Qwen3-TTS prompt-embeds builder specialized for the 310P NPU path."""
 
-        target_sr = int(getattr(self._config.speaker_encoder_config, "sample_rate", 24000))
-        if sr != target_sr:
-            resampler = self._get_resampler(int(sr), target_sr)
-            wav = resampler.resample(wav.astype(np.float32), orig_sr=int(sr))
-
-        # 310P does not support torch.stft on NPU.
-        wav_tensor = torch.from_numpy(wav).to(device=_CPU_DEVICE, dtype=torch.float32).unsqueeze(0)
-        mels = prompt_embeds_builder.mel_spectrogram(
-            wav_tensor,
-            n_fft=1024,
-            num_mels=128,
-            sampling_rate=24000,
-            hop_size=256,
-            win_size=1024,
-            fmin=0,
-            fmax=12000,
-        ).transpose(1, 2)
-        spk = self._speaker_encoder(mels.to(device=dev, dtype=dtype))[0]
-        return spk.to(dtype=dtype)
+    _mel_spectrogram_on_cpu = True
 
 
 # ===================================================================
@@ -192,10 +166,15 @@ class _Qwen3CodePredictorAttention310P(qwen3_code_predictor.CodePredictorAttenti
         cos, sin = position_embeddings
         cos = cos.unsqueeze(1)
         sin = sin.unsqueeze(1)
-        # Use the fused Ascend RoPE op instead of expanding RoPE into
-        # elementwise mul/add/rotate-half kernels.
-        q = torch_npu.npu_rotary_mul(q, cos, sin)
-        k = torch_npu.npu_rotary_mul(k, cos, sin)
+        # Use the fused aclnn rope op, which is NPUGraph-capturable; it
+        # requires the 4D BSND layout with head_dim 128 or 64.
+        q = q.transpose(1, 2).contiguous()
+        k = k.transpose(1, 2).contiguous()
+        cos = cos.transpose(1, 2).contiguous()
+        sin = sin.transpose(1, 2).contiguous()
+        q, k = torch_npu.npu_apply_rotary_pos_emb(q, k, cos, sin, rotary_mode="half")
+        q = q.transpose(1, 2)
+        k = k.transpose(1, 2)
 
         real_tokens = int(bsz) * int(seq_len)
         output_dtype = q.dtype

vllm_omni/platforms/npu/__init__.py

@@ -1,14 +1,30 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+from __future__ import annotations
+
 from typing import TYPE_CHECKING
 
+import torch
+
 if TYPE_CHECKING:
     from vllm_omni.platforms.npu.platform import NPUOmniPlatform
 
 __all__ = ["NPUOmniPlatform"]
 
 
+def is_a5(device: torch.device | None = None) -> bool:
+    """Return True on Ascend 950 (A5) devices, False otherwise."""
+    if device is not None and device.type != "npu":
+        return False
+    try:
+        from vllm_ascend.utils import is_950
+
+        return bool(is_950())
+    except (ImportError, AttributeError):
+        return False
+
+
 def __getattr__(name: str):
     if name != "NPUOmniPlatform":
         raise AttributeError(f"module {__name__!r} has no attribute {name!r}")

vllm_omni/platforms/npu/models/qwen3_tts.py

@@ -1,12 +1,16 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
-"""Patch Qwen3-TTS Code2Wav NPU runtime setup and weight preparation."""
+"""Patch Qwen3-TTS NPU runtime setup and weights.
+
+The code2wav runtime/weight fixes apply on every NPU. On A5 (Ascend 950 PR)
+``torch.stft`` is unavailable, so the speaker-encoder mel front-end of the
+prompt-embeds builder runs on CPU and the result is moved back to NPU
+(same trade-off as the 310P path).
+"""
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING
-
 import torch
 import torch.nn as nn
 import torch_npu
@@ -14,12 +18,12 @@ from vllm.config import VllmConfig
 from vllm.logger import init_logger
 from vllm_ascend.utils import maybe_trans_nz
 
-if TYPE_CHECKING:
-    pass
+from vllm_omni.platforms.npu import is_a5
 
 logger = init_logger(__name__)
 
 _PATCHED = False
+_A5_PATCHED = False
 _original_init = None
 _original_load_weights = None
 _ACL_FORMAT_FRACTAL_Z = 4
@@ -49,7 +53,12 @@ def _prepare_npu_decoder_weights(decoder: nn.Module) -> None:
                 module.weight.data = maybe_trans_nz(module.weight.data)
                 linear_count += 1
             elif isinstance(module, (nn.Conv1d, nn.ConvTranspose1d)) and module.groups == 1:
-                module.weight.data = torch_npu.npu_format_cast(module.weight.data.contiguous(), _ACL_FORMAT_FRACTAL_Z)
+                # A5 (Ascend 950 PR) does not support the FRACTAL_Z conv
+                # weight layout, keep the contiguous ND layout there.
+                if not is_a5():
+                    module.weight.data = torch_npu.npu_format_cast(
+                        module.weight.data.contiguous(), _ACL_FORMAT_FRACTAL_Z
+                    )
                 conv_count += 1
 
     logger.info("Prepared NPU Code2Wav weights: linear=%d conv=%d", linear_count, conv_count)
@@ -67,6 +76,44 @@ def _patched_load_weights(self, weights):
     return loaded
 
 
+def apply_qwen3_tts_patches() -> None:
+    """Install all Qwen3-TTS NPU patches.
+
+    The code2wav runtime/weight patch applies on every NPU; the A5
+    prompt-embeds builder swap is gated on ``is_a5()`` because it changes the
+    speaker-encoder mel front-end (``torch.stft`` is unsupported there).
+    """
+    apply_qwen3_tts_code2wav_patch()
+    if not is_a5():
+        return
+    _apply_a5_prompt_embeds_builder_patch()
+
+
+def _apply_a5_prompt_embeds_builder_patch() -> None:
+    global _A5_PATCHED
+    if _A5_PATCHED:
+        return
+
+    # Import and patch lazily so the Qwen3-TTS model modules are only loaded
+    # on A5 devices (same gate-then-import convention as the 310P patch).
+    from vllm_omni.model_executor.models.qwen3_tts import (
+        prompt_embeds_builder,
+        qwen3_tts_talker,
+    )
+
+    class _Qwen3TTSPromptEmbedsBuilderA5(prompt_embeds_builder.Qwen3TTSPromptEmbedsBuilder):
+        """Qwen3-TTS prompt-embeds builder specialized for the A5 NPU path."""
+
+        _mel_spectrogram_on_cpu = True
+
+    # The talker imports the builder by name at module scope, so its reference
+    # must be swapped as well (same as the 310P patch).
+    prompt_embeds_builder.Qwen3TTSPromptEmbedsBuilder = _Qwen3TTSPromptEmbedsBuilderA5
+    qwen3_tts_talker.Qwen3TTSPromptEmbedsBuilder = _Qwen3TTSPromptEmbedsBuilderA5
+    _A5_PATCHED = True
+    logger.debug("Applied A5 prompt-embeds builder patch for Qwen3-TTS")
+
+
 def apply_qwen3_tts_code2wav_patch() -> None:
     global _PATCHED, _original_init, _original_load_weights
     if _PATCHED:

vllm_omni/platforms/npu/platform.py

@@ -62,8 +62,8 @@ class NPUOmniPlatform(OmniPlatform, NPUPlatform):
         from vllm_omni.platforms.npu.models.minicpmo_4_5_code2wav import (
             apply_minicpmo_4_5_code2wav_patch,
         )
-        from vllm_omni.platforms.npu.models.qwen3_tts_code2wav import (
-            apply_qwen3_tts_code2wav_patch,
+        from vllm_omni.platforms.npu.models.qwen3_tts import (
+            apply_qwen3_tts_patches,
         )
         from vllm_omni.platforms.npu.models.qwen3_tts_tokenizer_v2 import (
             apply_qwen3_tts_tokenizer_v2_patch,
@@ -71,7 +71,7 @@ class NPUOmniPlatform(OmniPlatform, NPUPlatform):
 
         adapt_patch(is_global_patch=True)
         apply_minicpmo_4_5_code2wav_patch()
-        apply_qwen3_tts_code2wav_patch()
+        apply_qwen3_tts_patches()
         apply_qwen3_tts_tokenizer_v2_patch()
         apply_310p_patches()