From 6e79d582e3c4a4da4a159904125f3de7a6cbae2b Mon Sep 17 00:00:00 2001 From: zyz111222 Date: Wed, 9 Sep 2026 19:03:09 +0800 Subject: [PATCH] [Bugfix][Ascend] fix npu 310/a5 bugs (#6685) Signed-off-by: zouyizhou --- tests/platforms/npu/test_310p_patches.py | 69 ++--- tests/platforms/npu/test_a5_patches.py | 259 ++++++++++++++++++ .../models/qwen3_tts/prompt_embeds_builder.py | 10 +- vllm_omni/platforms/npu/_310p/__init__.py | 7 +- .../platforms/npu/_310p/patch/qwen3_tts.py | 43 +-- vllm_omni/platforms/npu/__init__.py | 16 ++ .../{qwen3_tts_code2wav.py => qwen3_tts.py} | 59 +++- vllm_omni/platforms/npu/platform.py | 6 +- 8 files changed, 380 insertions(+), 89 deletions(-) create mode 100644 tests/platforms/npu/test_a5_patches.py rename vllm_omni/platforms/npu/models/{qwen3_tts_code2wav.py => qwen3_tts.py} (54%) diff --git a/tests/platforms/npu/test_310p_patches.py b/tests/platforms/npu/test_310p_patches.py index 08cbb2adc..c19dcd997 100644 --- a/tests/platforms/npu/test_310p_patches.py +++ b/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() diff --git a/tests/platforms/npu/test_a5_patches.py b/tests/platforms/npu/test_a5_patches.py new file mode 100644 index 000000000..0d6648235 --- /dev/null +++ b/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 diff --git a/vllm_omni/model_executor/models/qwen3_tts/prompt_embeds_builder.py b/vllm_omni/model_executor/models/qwen3_tts/prompt_embeds_builder.py index ddb55cd51..609117ce3 100644 --- a/vllm_omni/model_executor/models/qwen3_tts/prompt_embeds_builder.py +++ b/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( diff --git a/vllm_omni/platforms/npu/_310p/__init__.py b/vllm_omni/platforms/npu/_310p/__init__.py index 129cf8e20..605a1aa53 100644 --- a/vllm_omni/platforms/npu/_310p/__init__.py +++ b/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: diff --git a/vllm_omni/platforms/npu/_310p/patch/qwen3_tts.py b/vllm_omni/platforms/npu/_310p/patch/qwen3_tts.py index dd03ba7d2..4ec4b6b72 100644 --- a/vllm_omni/platforms/npu/_310p/patch/qwen3_tts.py +++ b/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 diff --git a/vllm_omni/platforms/npu/__init__.py b/vllm_omni/platforms/npu/__init__.py index 74dfc419c..3df9b5ca3 100644 --- a/vllm_omni/platforms/npu/__init__.py +++ b/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}") diff --git a/vllm_omni/platforms/npu/models/qwen3_tts_code2wav.py b/vllm_omni/platforms/npu/models/qwen3_tts.py similarity index 54% rename from vllm_omni/platforms/npu/models/qwen3_tts_code2wav.py rename to vllm_omni/platforms/npu/models/qwen3_tts.py index bd74e9dc6..a91d2cf80 100644 --- a/vllm_omni/platforms/npu/models/qwen3_tts_code2wav.py +++ b/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: diff --git a/vllm_omni/platforms/npu/platform.py b/vllm_omni/platforms/npu/platform.py index cf65da414..2c8d42f60 100644 --- a/vllm_omni/platforms/npu/platform.py +++ b/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()