From 6de0ca36bb3f60f7f9aefd62bf7fcc930cf56c0b Mon Sep 17 00:00:00 2001 From: zyz111222 <2542678383@qq.com> Date: Fri, 12 Jun 2026 17:39:21 +0800 Subject: [PATCH] [Hardware][Ascend] Adapt Qwen3 TTS for 310P (#4283) Signed-off-by: zouyizhou Co-authored-by: Canlin Guo --- requirements/npu.txt | 2 +- tests/platforms/npu/test_310p_patches.py | 338 ++++++++++++++++++ .../models/qwen3_tts/prompt_embeds_builder.py | 19 +- .../models/qwen3_tts/qwen3_tts_talker.py | 21 +- vllm_omni/platforms/npu/_310p/__init__.py | 40 +++ .../platforms/npu/_310p/patch/__init__.py | 30 ++ .../platforms/npu/_310p/patch/qwen3_tts.py | 227 ++++++++++++ vllm_omni/platforms/npu/_310p/patch/worker.py | 27 ++ vllm_omni/platforms/npu/platform.py | 2 + vllm_omni/platforms/npu/worker/base.py | 7 +- .../platforms/npu/worker/npu_model_runner.py | 13 +- 11 files changed, 706 insertions(+), 20 deletions(-) create mode 100644 tests/platforms/npu/test_310p_patches.py create mode 100644 vllm_omni/platforms/npu/_310p/__init__.py create mode 100644 vllm_omni/platforms/npu/_310p/patch/__init__.py create mode 100644 vllm_omni/platforms/npu/_310p/patch/qwen3_tts.py create mode 100644 vllm_omni/platforms/npu/_310p/patch/worker.py diff --git a/requirements/npu.txt b/requirements/npu.txt index 955742757..954e983d8 100644 --- a/requirements/npu.txt +++ b/requirements/npu.txt @@ -1,3 +1,3 @@ -r common.txt onnxruntime-cann>=1.23.2 -torchaudio==2.9.0 +torchaudio==2.10.0 diff --git a/tests/platforms/npu/test_310p_patches.py b/tests/platforms/npu/test_310p_patches.py new file mode 100644 index 000000000..255b092c8 --- /dev/null +++ b/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] 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 b227be19d..b4f1d478c 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 @@ -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: - self._speaker_encoder.to(device=dev, dtype=torch.bfloat16) + if spk_param.device != dev or spk_param.dtype != dtype: + 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] - return spk.to(dtype=torch.bfloat16) + spk = self._speaker_encoder(mels.to(dtype=dtype))[0] + 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: diff --git a/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_talker.py b/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_talker.py index 541efb580..bf8ad5608 100644 --- a/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_talker.py +++ b/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) - past_hidden = last_talker_hidden.reshape(bsz, 1, -1).to(dtype=torch.bfloat16, device=dev) - text_step = text_step.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=dtype, 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 diff --git a/vllm_omni/platforms/npu/_310p/__init__.py b/vllm_omni/platforms/npu/_310p/__init__.py new file mode 100644 index 000000000..129cf8e20 --- /dev/null +++ b/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() diff --git a/vllm_omni/platforms/npu/_310p/patch/__init__.py b/vllm_omni/platforms/npu/_310p/patch/__init__.py new file mode 100644 index 000000000..b3ab637ff --- /dev/null +++ b/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() diff --git a/vllm_omni/platforms/npu/_310p/patch/qwen3_tts.py b/vllm_omni/platforms/npu/_310p/patch/qwen3_tts.py new file mode 100644 index 000000000..c43cdd3cf --- /dev/null +++ b/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 diff --git a/vllm_omni/platforms/npu/_310p/patch/worker.py b/vllm_omni/platforms/npu/_310p/patch/worker.py new file mode 100644 index 000000000..3116b54da --- /dev/null +++ b/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 diff --git a/vllm_omni/platforms/npu/platform.py b/vllm_omni/platforms/npu/platform.py index 777efac20..6f2eb99aa 100644 --- a/vllm_omni/platforms/npu/platform.py +++ b/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: diff --git a/vllm_omni/platforms/npu/worker/base.py b/vllm_omni/platforms/npu/worker/base.py index d617ed5bc..ff1a860b0 100644 --- a/vllm_omni/platforms/npu/worker/base.py +++ b/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): diff --git a/vllm_omni/platforms/npu/worker/npu_model_runner.py b/vllm_omni/platforms/npu/worker/npu_model_runner.py index 3108b68a8..5a50f523d 100644 --- a/vllm_omni/platforms/npu/worker/npu_model_runner.py +++ b/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.