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

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

requirements/npu.txt

@@ -1,3 +1,3 @@
 -r common.txt
 onnxruntime-cann>=1.23.2
-torchaudio==2.9.0
+torchaudio==2.10.0

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]

vllm_omni/model_executor/models/qwen3_tts/prompt_embeds_builder.py

@@ -329,6 +329,7 @@ class Qwen3TTSPromptEmbedsBuilder:
         self._tts_pad_embed_buffer = tts_pad_embed
         self._encode_ref_audio_batch_fn = encode_ref_audio_batch
         self._speaker_cache = speaker_cache
+        self._embedding_dtype = torch.bfloat16
 
         self._text_tokenizer: Any | None = None
 
@@ -598,7 +599,7 @@ class Qwen3TTSPromptEmbedsBuilder:
         if isinstance(ref_code, torch.Tensor):
             entry["ref_code"] = ref_code.detach().to("cpu", dtype=torch.long).contiguous()
         if isinstance(ref_spk_embedding, torch.Tensor):
-            entry["ref_spk_embedding"] = ref_spk_embedding.detach().to("cpu", dtype=torch.bfloat16).reshape(-1)
+            entry["ref_spk_embedding"] = ref_spk_embedding.detach().to("cpu", dtype=self._embedding_dtype).reshape(-1)
         if not entry:
             return
         self._ref_audio_artifact_cache[cache_key] = entry
@@ -613,10 +614,11 @@ class Qwen3TTSPromptEmbedsBuilder:
         # CUDA. Ensure the speaker encoder is on the same device/dtype as the
         # main model before running it.
         dev = self._device()
+        dtype = self._embedding_dtype
         try:
             spk_param = next(self._speaker_encoder.parameters())
-            if spk_param.device != dev or spk_param.dtype != torch.bfloat16:
+            if spk_param.device != dev or spk_param.dtype != dtype:
-                self._speaker_encoder.to(device=dev, dtype=torch.bfloat16)
+                self._speaker_encoder.to(device=dev, dtype=dtype)
         except StopIteration:
             pass
 
@@ -641,8 +643,8 @@ class Qwen3TTSPromptEmbedsBuilder:
             fmin=0,
             fmax=12000,
         ).transpose(1, 2)
-        spk = self._speaker_encoder(mels.to(dtype=torch.bfloat16))[0]
+        spk = self._speaker_encoder(mels.to(dtype=dtype))[0]
-        return spk.to(dtype=torch.bfloat16)
+        return spk.to(dtype=dtype)
 
     def encode_ref_audio_batch(
         self,
@@ -1187,17 +1189,18 @@ class Qwen3TTSPromptEmbedsBuilder:
             # NOTE: Do NOT use _as_singleton here — the embedding may be a plain
             # float list (from API via msgspec IPC) that _as_singleton would
             # destructively unwrap to a single scalar.
+            speaker_dtype = self._embedding_dtype
             spk = None
             if voice_clone_prompt is not None:
                 spk = voice_clone_prompt.get("ref_spk_embedding")
             if isinstance(spk, torch.Tensor):
-                speaker_embed = spk.to(device=input_ids.device, dtype=torch.bfloat16).view(1, 1, -1)
+                speaker_embed = spk.to(device=input_ids.device, dtype=speaker_dtype).view(1, 1, -1)
             elif isinstance(spk, (list, np.ndarray)):
-                speaker_embed = torch.tensor(spk, dtype=torch.bfloat16, device=input_ids.device).view(1, 1, -1)
+                speaker_embed = torch.tensor(spk, dtype=speaker_dtype, device=input_ids.device).view(1, 1, -1)
             elif cached_artifacts is not None and isinstance(cached_artifacts.get("ref_spk_embedding"), torch.Tensor):
                 speaker_embed = (
                     cached_artifacts["ref_spk_embedding"]
-                    .to(device=input_ids.device, dtype=torch.bfloat16)
+                    .to(device=input_ids.device, dtype=speaker_dtype)
                     .view(1, 1, -1)
                 )
             elif artifact_only:

vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_talker.py

@@ -400,6 +400,7 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module):
             torch.zeros(1, int(self.talker_config.hidden_size), dtype=model_dtype),
             persistent=False,
         )
+        self._embedding_dtype = torch.bfloat16
 
         tokenizer_config = Qwen3TTSTokenizerV2Config.from_pretrained(
             self.model_path,
@@ -642,7 +643,8 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module):
         # :meth:`_init_runtime_buffers`. Materialize once on the right
         # device/dtype and reuse for both the prefill placeholder padding
         # and the decode text-step fallback below.
-        tts_pad_embed = self._tts_pad_embed.to(device=input_ids.device, dtype=torch.bfloat16).reshape(1, -1)
+        dtype = self._embedding_dtype
+        tts_pad_embed = self._tts_pad_embed.to(device=input_ids.device, dtype=dtype).reshape(1, -1)
 
         if is_prefill:
             # Prefill (prompt embeddings)
@@ -689,7 +691,7 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module):
                 pad_n = int(span_len - int(take.shape[0]))
                 pad_rows = tts_pad_embed.reshape(1, -1).to("cpu").expand(pad_n, -1)
                 take = torch.cat([take, pad_rows], dim=0)
-            prompt_embeds = take.to(device=input_ids.device, dtype=torch.bfloat16)
+            prompt_embeds = take.to(device=input_ids.device, dtype=dtype)
             info_update["meta"]["talker_prefill_offset"] = int(offset + span_len)
 
             # When inputs_embeds is set, token ids are ignored by the model but must stay in-vocab for vLLM bookkeeping.
@@ -719,7 +721,7 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module):
                     tail[text_offset : text_offset + 1]
                     .to(
                         device=input_ids.device,
-                        dtype=torch.bfloat16,
+                        dtype=dtype,
                     )
                     .reshape(1, -1)
                 )
@@ -744,14 +746,14 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module):
 
         last_hidden = hs.get("last")
         if isinstance(last_hidden, torch.Tensor):
-            past_hidden = last_hidden.to(device=input_ids.device, dtype=torch.bfloat16).reshape(1, -1)
+            past_hidden = last_hidden.to(device=input_ids.device, dtype=dtype).reshape(1, -1)
         else:
             # Defensive: EOS step row is zeroed by the invalid-layer-0 mask and filtered downstream.
             past_hidden = torch.zeros_like(text_step)
 
         # Use OmniGPUModelRunner talker_mtp fast-path for residual codebooks and per-step inputs_embeds update.
         last_id_hidden = self.embed_input_ids(input_ids.reshape(1, 1).to(torch.long)).to(
-            device=input_ids.device, dtype=torch.bfloat16
+            device=input_ids.device, dtype=dtype
         )
         inputs_embeds_out = last_id_hidden.reshape(1, -1)
 
@@ -784,7 +786,7 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module):
             )
 
         device = input_ids_flat.device
-        dtype = torch.bfloat16
+        dtype = self._embedding_dtype
         # Request-independent constant (see :meth:`_init_runtime_buffers`) —
         # compute once for the batch instead of fetching it per request from
         # ``info_dict["embed"]["tts_pad"]``.
@@ -1043,11 +1045,12 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module):
         bsz = int(input_ids.shape[0])
         q = int(self.talker_config.num_code_groups)
         dev = input_embeds.device
+        dtype = self._embedding_dtype
 
         input_ids = input_ids.reshape(bsz, 1).to(dtype=torch.long, device=dev)
-        last_id_hidden = input_embeds.reshape(bsz, 1, -1).to(dtype=torch.bfloat16, device=dev)
+        last_id_hidden = input_embeds.reshape(bsz, 1, -1).to(dtype=dtype, device=dev)
-        past_hidden = last_talker_hidden.reshape(bsz, 1, -1).to(dtype=torch.bfloat16, device=dev)
+        past_hidden = last_talker_hidden.reshape(bsz, 1, -1).to(dtype=dtype, device=dev)
-        text_step = text_step.reshape(bsz, 1, -1).to(dtype=torch.bfloat16, device=dev)
+        text_step = text_step.reshape(bsz, 1, -1).to(dtype=dtype, device=dev)
 
         # Residual predictor runs fixed-length (Q-1) steps via the vLLM-native code_predictor.
         max_steps = q - 1

vllm_omni/platforms/npu/_310p/__init__.py

@@ -0,0 +1,40 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""310P platform helpers."""
+
+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:
+    if device is not None and device.type != "npu":
+        return False
+    try:
+        from vllm_ascend.utils import is_310p as ascend_is_310p
+
+        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:
+    import torch_npu
+
+    torch_npu.npu.set_compile_mode(jit_compile=False)
+
+
+def apply_patches() -> None:
+    if not is_310p():
+        return
+
+    from vllm_omni.platforms.npu._310p.patch import apply_patches as apply_310p_patches
+
+    apply_310p_patches()

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

@@ -0,0 +1,30 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""Apply 310P patches."""
+
+from __future__ import annotations
+
+_WORKER_PATCHED = False
+_QWEN3_TTS_TALKER_ARCH = "Qwen3TTSTalkerForConditionalGeneration"
+
+
+def apply_patches() -> None:
+    global _WORKER_PATCHED
+
+    if _WORKER_PATCHED:
+        return
+
+    from vllm_omni.platforms.npu._310p.patch.worker import apply_patch
+
+    apply_patch()
+    _WORKER_PATCHED = True
+
+
+def apply_model_patches(model_config) -> None:
+    if getattr(model_config, "model_arch", None) != _QWEN3_TTS_TALKER_ARCH:
+        return
+
+    from vllm_omni.platforms.npu._310p.patch.qwen3_tts import apply_talker_patches
+
+    apply_talker_patches()

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

@@ -0,0 +1,227 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""Patch Qwen3-TTS for the 310P NPU path."""
+
+from __future__ import annotations
+
+import numpy as np
+import torch
+from transformers.models.mimi import modeling_mimi
+from vllm.multimodal.audio import AudioResampler
+
+from vllm_omni.model_executor.models.common import qwen3_code_predictor
+from vllm_omni.model_executor.models.qwen3_tts import prompt_embeds_builder, qwen3_tts_talker
+
+_RUNTIME_DTYPE = torch.float16
+_CPU_DEVICE = torch.device("cpu")
+
+
+class _MimiEuclideanCodebook310P(modeling_mimi.MimiEuclideanCodebook):
+    def quantize(self, hidden_states):
+        # 310P does not support torch.cdist on NPU.
+        device = hidden_states.device
+        dists = torch.cdist(
+            hidden_states[None].to(_CPU_DEVICE, torch.float32), self.embed[None].to(_CPU_DEVICE, torch.float32), p=2
+        )[0]
+        return dists.argmin(dim=-1).to(device=device)
+
+
+class _Qwen3TTSTalker310P(qwen3_tts_talker.Qwen3TTSTalkerForConditionalGeneration):
+    def __init__(self, *, vllm_config, prefix: str = "") -> None:
+        super().__init__(vllm_config=vllm_config, prefix=prefix)
+        self._embedding_dtype = _RUNTIME_DTYPE
+        self._prompt_builder._embedding_dtype = _RUNTIME_DTYPE
+
+    def load_weights(self, weights):
+        loaded = super().load_weights(weights)
+        self.encoder.to(dtype=_RUNTIME_DTYPE)
+        return loaded
+
+    def _encode_ref_audio_batch(
+        self,
+        wavs: list[np.ndarray],
+        sr: int,
+        *,
+        device: torch.device,
+    ) -> list[torch.Tensor]:
+        fe = self._encoder_feature_extractor
+        target_sr = int(fe.sampling_rate)
+        if int(sr) != target_sr:
+            resampler = AudioResampler(target_sr=target_sr)
+            wavs = [resampler.resample(w.astype(np.float32), orig_sr=int(sr)) for w in wavs]
+
+        inputs = fe(raw_audio=wavs, sampling_rate=target_sr, return_tensors="pt").to(device).to(_RUNTIME_DTYPE)
+
+        with torch.inference_mode():
+            encoded = self.encoder.encode(
+                input_values=inputs["input_values"].squeeze(1).unsqueeze(1),
+                return_dict=True,
+            )
+
+        audio_codes = encoded.audio_codes[:, : self._encoder_valid_num_quantizers]
+        padding_mask = inputs["padding_mask"].squeeze(1)
+        downsample = self._encoder_downsample_rate
+        return [
+            code[..., : -(-mask.sum() // downsample)].transpose(0, 1).to(device=device, dtype=torch.long)
+            for code, mask in zip(audio_codes, padding_mask)
+        ]
+
+
+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
+
+        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)
+
+
+class _Qwen3CodePredictorAttention310P(qwen3_code_predictor.CodePredictorAttention):
+    def __init__(self, *args, **kwargs) -> None:
+        super().__init__(*args, **kwargs)
+        self._buffers.pop("_fusion_causal_mask", None)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        position_embeddings: tuple[torch.Tensor, torch.Tensor],
+        attention_mask: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        if hidden_states.device.type != "npu" or attention_mask is None:
+            return super().forward(hidden_states, position_embeddings)
+
+        bsz, seq_len, _ = hidden_states.shape
+        hidden_shape_q = (bsz, seq_len, self.num_heads, self.head_dim)
+        hidden_shape_kv = (bsz, seq_len, self.num_kv_heads, self.head_dim)
+
+        q = self.q_norm(self.q_proj(hidden_states).view(hidden_shape_q)).transpose(1, 2)
+        k = self.k_norm(self.k_proj(hidden_states).view(hidden_shape_kv)).transpose(1, 2)
+        v = self.v_proj(hidden_states).view(hidden_shape_kv).transpose(1, 2)
+
+        cos, sin = position_embeddings
+        cos = cos.unsqueeze(1)
+        sin = sin.unsqueeze(1)
+        q = (q * cos) + (qwen3_code_predictor._rotate_half(q) * sin)
+        k = (k * cos) + (qwen3_code_predictor._rotate_half(k) * sin)
+
+        real_tokens = int(bsz) * int(seq_len)
+        output_dtype = q.dtype
+
+        import torch_npu
+        from vllm_ascend.utils import aligned_16
+
+        q_f = aligned_16(q.to(torch.float16).transpose(1, 2).reshape(real_tokens, self.num_heads, self.head_dim))
+        k_f = aligned_16(k.to(torch.float16).transpose(1, 2).reshape(real_tokens, self.num_kv_heads, self.head_dim))
+        v_f = aligned_16(v.to(torch.float16).transpose(1, 2).reshape(real_tokens, self.num_kv_heads, self.head_dim))
+
+        aligned_tokens = int(q_f.shape[0])
+        seq_lens = torch.full((int(bsz),), int(seq_len), dtype=torch.int32, device="cpu")
+        if aligned_tokens > real_tokens:
+            seq_lens[-1] += aligned_tokens - real_tokens
+
+        out = torch.empty((aligned_tokens, self.num_heads, self.head_dim), dtype=torch.float16, device=q.device)
+        torch_npu._npu_flash_attention(
+            query=q_f.contiguous(),
+            key=k_f.contiguous(),
+            value=v_f.contiguous(),
+            mask=attention_mask,
+            seq_len=seq_lens,
+            scale_value=float(self.scaling),
+            num_heads=int(self.num_heads),
+            num_kv_heads=int(self.num_kv_heads),
+            out=out,
+        )
+        attn_out = out[:real_tokens].reshape(bsz, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
+        return self.o_proj(attn_out.to(output_dtype).transpose(1, 2).reshape(bsz, seq_len, -1))
+
+
+class _Qwen3CodePredictorDecoderLayer310P(qwen3_code_predictor.CodePredictorDecoderLayer):
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        position_embeddings: tuple[torch.Tensor, torch.Tensor],
+        attention_mask: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        residual = hidden_states
+        hidden_states = self.input_layernorm(hidden_states)
+        hidden_states = self.self_attn(hidden_states, position_embeddings, attention_mask=attention_mask)
+        hidden_states = residual + hidden_states
+
+        residual = hidden_states
+        hidden_states = self.post_attention_layernorm(hidden_states)
+        return residual + self.mlp(hidden_states)
+
+
+class _Qwen3CodePredictorBaseModel310P(qwen3_code_predictor.CodePredictorBaseModel):
+    def __init__(self, *args, **kwargs) -> None:
+        super().__init__(*args, **kwargs)
+        self._attention_mask_310p = None
+        self._attention_mask_310p_device = None
+        self._attention_mask_310p_max_seq = ((int(self.config.num_code_groups) + 16) // 16) * 16
+
+    def forward(
+        self,
+        inputs_embeds: torch.Tensor,
+        position_ids: torch.Tensor,
+    ) -> torch.Tensor:
+        if inputs_embeds.device.type != "npu":
+            return super().forward(inputs_embeds, position_ids)
+
+        if self._attention_mask_310p is None or self._attention_mask_310p_device != inputs_embeds.device:
+            import torch_npu
+            from vllm_ascend._310p.attention.attention_mask import AttentionMaskBuilder310
+            from vllm_ascend.utils import ACL_FORMAT_FRACTAL_NZ, nd_to_nz_2d
+
+            mask = AttentionMaskBuilder310.gen_causal_additive_mask(
+                self._attention_mask_310p_max_seq,
+                inputs_embeds.device,
+            )
+            self._attention_mask_310p = torch_npu.npu_format_cast(nd_to_nz_2d(mask), ACL_FORMAT_FRACTAL_NZ)
+            self._attention_mask_310p_device = inputs_embeds.device
+
+        input_dtype = inputs_embeds.dtype
+        hidden_states = inputs_embeds
+        with torch.amp.autocast(inputs_embeds.device.type, enabled=False, dtype=torch.float32):
+            position_embeddings = self.rotary_emb(hidden_states, position_ids)
+            for layer in self.layers:
+                hidden_states = layer(
+                    hidden_states,
+                    position_embeddings,
+                    attention_mask=self._attention_mask_310p,
+                )
+            hidden_states = self.norm(hidden_states)
+        return hidden_states.to(input_dtype)
+
+
+def apply_talker_patches() -> None:
+    modeling_mimi.MimiEuclideanCodebook = _MimiEuclideanCodebook310P
+    qwen3_tts_talker.Qwen3TTSTalkerForConditionalGeneration = _Qwen3TTSTalker310P
+    qwen3_tts_talker.Qwen3TTSPromptEmbedsBuilder = _Qwen3TTSPromptEmbedsBuilder310P
+    prompt_embeds_builder.Qwen3TTSPromptEmbedsBuilder = _Qwen3TTSPromptEmbedsBuilder310P
+    qwen3_code_predictor.CodePredictorAttention = _Qwen3CodePredictorAttention310P
+    qwen3_code_predictor.CodePredictorDecoderLayer = _Qwen3CodePredictorDecoderLayer310P
+    qwen3_code_predictor.CodePredictorBaseModel = _Qwen3CodePredictorBaseModel310P

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

@@ -0,0 +1,27 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""Patch Omni NPU worker device initialization for 310P.
+
+This patches:
+    vllm_omni.platforms.npu.worker.base.OmniNPUWorkerBase
+
+Omni AR/generation workers share this base initialization path, so the 310P
+compile-mode setup is applied here after device initialization.
+"""
+
+from __future__ import annotations
+
+from vllm_omni.platforms.npu._310p import disable_jit_compile
+from vllm_omni.platforms.npu.worker import base as worker_base
+
+
+class _OmniNPUWorkerBase310P(worker_base.OmniNPUWorkerBase):
+    def _init_device(self):
+        device = super()._init_device()
+        disable_jit_compile()
+        return device
+
+
+def apply_patch() -> None:
+    worker_base.OmniNPUWorkerBase = _OmniNPUWorkerBase310P

vllm_omni/platforms/npu/platform.py

@@ -33,11 +33,13 @@ class NPUOmniPlatform(OmniPlatform, NPUPlatform):
 
     # conv2d convolution operator in the code2wav module of Qwen3-TTS not being able to run on Aclnn
     def __init__(self) -> None:
+        from vllm_omni.platforms.npu._310p import apply_patches as apply_310p_patches
         from vllm_omni.platforms.npu.models.qwen3_tts_code2wav import (
             apply_qwen3_tts_code2wav_patch,
         )
 
         apply_qwen3_tts_code2wav_patch()
+        apply_310p_patches()
 
     @classmethod
     def set_device(cls, device: torch.device) -> None:

vllm_omni/platforms/npu/worker/base.py

@@ -5,7 +5,12 @@
 
 import time
 
-from vllm_ascend.worker.worker import NPUWorker
+from vllm_omni.platforms.npu._310p import is_310p
+
+if is_310p():
+    from vllm_ascend._310p.worker_310p import NPUWorker310 as NPUWorker
+else:
+    from vllm_ascend.worker.worker import NPUWorker
 
 
 class OmniNPUWorkerBase(NPUWorker):

vllm_omni/platforms/npu/worker/npu_model_runner.py

@@ -18,16 +18,27 @@ from vllm_ascend.attention.attention_v1 import AscendAttentionState
 from vllm_ascend.attention.utils import using_paged_attention
 from vllm_ascend.ops.rotary_embedding import update_cos_sin
 from vllm_ascend.utils import enable_sp, lmhead_tp_enable
-from vllm_ascend.worker.model_runner_v1 import SEQ_LEN_WITH_MAX_PA_WORKSPACE, NPUModelRunner
+from vllm_ascend.worker.model_runner_v1 import SEQ_LEN_WITH_MAX_PA_WORKSPACE
 
 from vllm_omni.model_executor.models.output_templates import OmniOutput
+from vllm_omni.platforms.npu._310p import is_310p
 from vllm_omni.worker.gpu_model_runner import OmniGPUModelRunner
 
 logger = init_logger(__name__)
 
 
+if is_310p():
+    from vllm_ascend._310p.model_runner_310p import NPUModelRunner310 as NPUModelRunner
+else:
+    from vllm_ascend.worker.model_runner_v1 import NPUModelRunner
+
+
 class OmniNPUModelRunner(OmniGPUModelRunner, NPUModelRunner):
     def load_model(self, *args, **kwargs) -> None:
+        if is_310p():
+            from vllm_omni.platforms.npu._310p.patch import apply_model_patches
+
+            apply_model_patches(self.model_config)
         NPUModelRunner.load_model(self, *args, **kwargs)
         # Initialize enable_sp cache to avoid get_current_vllm_config() error
         # in _pad_for_sequence_parallelism during execute_model.