[Hardware][Ascend] Adapt Qwen3 TTS for 310P (#4283)

Signed-off-by: zouyizhou <zouyizhou@huawei.com>
Co-authored-by: Canlin Guo <canlinguosdu@gmail.com>

tests/platforms/npu/test_310p_patches.py

@@ -0,0 +1,338 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Unit tests for 310P patch wiring.
+
+The tests load patch modules from source with fake Qwen3-TTS dependencies, so
+they validate the patch contract without loading real model or NPU kernels.
+"""
+
+from __future__ import annotations
+
+import importlib.util
+import sys
+import types
+from pathlib import Path
+from types import SimpleNamespace
+
+import numpy as np
+import pytest
+import torch
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def _repo_root() -> Path:
+    marker = Path("vllm_omni") / "platforms" / "npu" / "_310p" / "patch"
+    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 _install_qwen3_tts_patch_fakes(monkeypatch: pytest.MonkeyPatch):
+    class FakeAudioResampler:
+        def __init__(self, *, target_sr: int):
+            self.target_sr = target_sr
+
+        def resample(self, wav, *, orig_sr: int):
+            del orig_sr
+            return wav
+
+    class FakeCodePredictorAttention(torch.nn.Module):
+        def __init__(self, *args, **kwargs):
+            del args, kwargs
+            super().__init__()
+            self.register_buffer("_fusion_causal_mask", torch.ones(1), persistent=False)
+
+    class FakeCodePredictorDecoderLayer(torch.nn.Module):
+        pass
+
+    class FakeCodePredictorBaseModel(torch.nn.Module):
+        pass
+
+    class FakeMimiEuclideanCodebook(torch.nn.Module):
+        @property
+        def embed(self):
+            return self._embed
+
+    class FakeEncoder:
+        def __init__(self):
+            self.to_calls: list[dict] = []
+            self.last_input_dtype = None
+
+        def to(self, **kwargs):
+            self.to_calls.append(kwargs)
+            return self
+
+        def encode(self, *, input_values, return_dict: bool):
+            assert return_dict
+            self.last_input_dtype = input_values.dtype
+            return SimpleNamespace(audio_codes=torch.arange(12, dtype=torch.long).reshape(1, 3, 4))
+
+    class FakeFeatureBatch(dict):
+        def to(self, target):
+            if isinstance(target, torch.dtype):
+                for key, value in list(self.items()):
+                    if torch.is_floating_point(value):
+                        self[key] = value.to(dtype=target)
+                self.dtype = target
+            else:
+                self.device = torch.device(target)
+                for key, value in list(self.items()):
+                    self[key] = value.to(device=self.device)
+            return self
+
+    class FakeFeatureExtractor:
+        sampling_rate = 24000
+
+        def __call__(self, *, raw_audio, sampling_rate: int, return_tensors: str):
+            assert len(raw_audio) == 1
+            assert sampling_rate == self.sampling_rate
+            assert return_tensors == "pt"
+            return FakeFeatureBatch(
+                input_values=torch.ones(1, 1, 8, dtype=torch.float32),
+                padding_mask=torch.ones(1, 1, 8, dtype=torch.long),
+            )
+
+    class FakeTalkerBase(torch.nn.Module):
+        def __init__(self, *, vllm_config, prefix: str = ""):
+            del vllm_config, prefix
+            super().__init__()
+            self._embedding_dtype = torch.bfloat16
+            self._prompt_builder = SimpleNamespace(_embedding_dtype=torch.bfloat16)
+            self.encoder = FakeEncoder()
+            self._encoder_feature_extractor = FakeFeatureExtractor()
+            self._encoder_valid_num_quantizers = 2
+            self._encoder_downsample_rate = 2
+
+        def load_weights(self, weights):
+            del weights
+            self.encoder.to(dtype=torch.bfloat16)
+            return {"loaded"}
+
+    class FakePromptEmbedsBuilder:
+        pass
+
+    fake_qwen3_code_predictor = _install_fake_module(
+        monkeypatch,
+        "vllm_omni.model_executor.models.common.qwen3_code_predictor",
+        CodePredictorAttention=FakeCodePredictorAttention,
+        CodePredictorDecoderLayer=FakeCodePredictorDecoderLayer,
+        CodePredictorBaseModel=FakeCodePredictorBaseModel,
+        _rotate_half=lambda x: x,
+    )
+    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",
+        Qwen3TTSTalkerForConditionalGeneration=FakeTalkerBase,
+        Qwen3TTSPromptEmbedsBuilder=FakePromptEmbedsBuilder,
+    )
+    fake_modeling_mimi = _install_fake_module(
+        monkeypatch,
+        "transformers.models.mimi.modeling_mimi",
+        MimiEuclideanCodebook=FakeMimiEuclideanCodebook,
+    )
+    fake_mimi = _install_fake_module(
+        monkeypatch,
+        "transformers.models.mimi",
+        modeling_mimi=fake_modeling_mimi,
+    )
+
+    _install_fake_module(monkeypatch, "vllm")
+    _install_fake_module(monkeypatch, "vllm.multimodal")
+    _install_fake_module(monkeypatch, "vllm.multimodal.audio", AudioResampler=FakeAudioResampler)
+    _install_fake_module(monkeypatch, "transformers")
+    _install_fake_module(monkeypatch, "transformers.models", mimi=fake_mimi)
+    _install_fake_module(monkeypatch, "vllm_omni")
+    _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.common",
+        qwen3_code_predictor=fake_qwen3_code_predictor,
+    )
+    _install_fake_module(
+        monkeypatch,
+        "vllm_omni.model_executor.models.qwen3_tts",
+        prompt_embeds_builder=fake_prompt_builder,
+        qwen3_tts_talker=fake_talker,
+    )
+    return fake_qwen3_code_predictor, fake_prompt_builder, fake_talker
+
+
+def _load_qwen3_tts_patch(monkeypatch: pytest.MonkeyPatch):
+    fakes = _install_qwen3_tts_patch_fakes(monkeypatch)
+    path = _repo_root() / "vllm_omni" / "platforms" / "npu" / "_310p" / "patch" / "qwen3_tts.py"
+    module = _load_source_module("vllm_omni_test_310p_qwen3_tts_patch", path)
+    return module, fakes
+
+
+def test_registry_applies_worker_once_and_model_patch_lazily(monkeypatch: pytest.MonkeyPatch) -> None:
+    registry_path = _repo_root() / "vllm_omni" / "platforms" / "npu" / "_310p" / "patch" / "__init__.py"
+    registry = _load_source_module("vllm_omni_test_310p_patch_registry", registry_path)
+    calls = {"worker": 0, "talker": 0}
+
+    _install_fake_module(
+        monkeypatch,
+        "vllm_omni.platforms.npu._310p.patch.worker",
+        apply_patch=lambda: calls.__setitem__("worker", calls["worker"] + 1),
+    )
+    _install_fake_module(
+        monkeypatch,
+        "vllm_omni.platforms.npu._310p.patch.qwen3_tts",
+        apply_talker_patches=lambda: calls.__setitem__("talker", calls["talker"] + 1),
+    )
+
+    registry.apply_patches()
+    registry.apply_patches()
+    registry.apply_model_patches(SimpleNamespace(model_arch="OtherModel"))
+    registry.apply_model_patches(SimpleNamespace(model_arch="Qwen3TTSTalkerForConditionalGeneration"))
+
+    assert calls == {"worker": 1, "talker": 1}
+
+
+def test_worker_patch_replaces_base_and_runs_disable_jit(monkeypatch: pytest.MonkeyPatch) -> None:
+    calls: list[str] = []
+
+    class FakeOmniNPUWorkerBase:
+        def _init_device(self):
+            calls.append("parent")
+            return "npu:0"
+
+    fake_worker_base = _install_fake_module(
+        monkeypatch,
+        "vllm_omni.platforms.npu.worker.base",
+        OmniNPUWorkerBase=FakeOmniNPUWorkerBase,
+    )
+    _install_fake_module(monkeypatch, "vllm_omni")
+    _install_fake_module(monkeypatch, "vllm_omni.platforms")
+    _install_fake_module(monkeypatch, "vllm_omni.platforms.npu")
+    _install_fake_module(
+        monkeypatch,
+        "vllm_omni.platforms.npu._310p",
+        disable_jit_compile=lambda: calls.append("disable_jit"),
+    )
+    _install_fake_module(monkeypatch, "vllm_omni.platforms.npu.worker", base=fake_worker_base)
+
+    path = _repo_root() / "vllm_omni" / "platforms" / "npu" / "_310p" / "patch" / "worker.py"
+    module = _load_source_module("vllm_omni_test_310p_worker_patch", path)
+    module.apply_patch()
+
+    assert fake_worker_base.OmniNPUWorkerBase is module._OmniNPUWorkerBase310P
+    assert fake_worker_base.OmniNPUWorkerBase()._init_device() == "npu:0"
+    assert calls == ["parent", "disable_jit"]
+
+
+def test_qwen3_tts_patch_replaces_target_classes(monkeypatch: pytest.MonkeyPatch) -> None:
+    module, (fake_code_predictor, fake_prompt_builder, fake_talker) = _load_qwen3_tts_patch(monkeypatch)
+
+    module.apply_talker_patches()
+
+    assert fake_talker.Qwen3TTSTalkerForConditionalGeneration is module._Qwen3TTSTalker310P
+    assert fake_talker.Qwen3TTSPromptEmbedsBuilder is module._Qwen3TTSPromptEmbedsBuilder310P
+    assert fake_prompt_builder.Qwen3TTSPromptEmbedsBuilder is module._Qwen3TTSPromptEmbedsBuilder310P
+    assert fake_code_predictor.CodePredictorAttention is module._Qwen3CodePredictorAttention310P
+    assert fake_code_predictor.CodePredictorDecoderLayer is module._Qwen3CodePredictorDecoderLayer310P
+    assert fake_code_predictor.CodePredictorBaseModel is module._Qwen3CodePredictorBaseModel310P
+
+
+def test_qwen3_tts_talker_patch_uses_fp16_runtime_dtype(monkeypatch: pytest.MonkeyPatch) -> None:
+    module, _ = _load_qwen3_tts_patch(monkeypatch)
+    talker = module._Qwen3TTSTalker310P(vllm_config=object())
+
+    assert talker._embedding_dtype is torch.float16
+    assert talker._prompt_builder._embedding_dtype is torch.float16
+    assert talker.load_weights([]) == {"loaded"}
+    assert talker.encoder.to_calls[-1] == {"dtype": torch.float16}
+
+    codes = talker._encode_ref_audio_batch([np.zeros(8, dtype=np.float32)], 24000, device=torch.device("cpu"))
+
+    assert talker.encoder.last_input_dtype is torch.float16
+    assert len(codes) == 1
+    assert codes[0].dtype is torch.long
+    assert codes[0].shape == (4, 2)
+
+
+def test_qwen3_tts_prompt_patch_runs_stft_frontend_on_cpu(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
+
+
+def test_qwen3_tts_mimi_codebook_quantize_uses_cpu_fp32_cdist(monkeypatch: pytest.MonkeyPatch) -> None:
+    module, _ = _load_qwen3_tts_patch(monkeypatch)
+    real_cdist = torch.cdist
+    captured = {}
+
+    def fake_cdist(x1, x2, p=2):
+        captured["x1_device"] = x1.device
+        captured["x2_device"] = x2.device
+        captured["x1_dtype"] = x1.dtype
+        captured["x2_dtype"] = x2.dtype
+        return real_cdist(x1, x2, p=p)
+
+    monkeypatch.setattr(torch, "cdist", fake_cdist)
+    codebook = object.__new__(module._MimiEuclideanCodebook310P)
+    codebook._embed = torch.tensor([[0.0, 0.0], [2.0, 0.0]], dtype=torch.float16)
+
+    indices = codebook.quantize(torch.tensor([[0.1, 0.0], [1.8, 0.0]], dtype=torch.float16))
+
+    assert captured == {
+        "x1_device": torch.device("cpu"),
+        "x2_device": torch.device("cpu"),
+        "x1_dtype": torch.float32,
+        "x2_dtype": torch.float32,
+    }
+    assert indices.tolist() == [0, 1]