diff --git a/tests/config/test_config_factory.py b/tests/config/test_config_factory.py index 82c1fb32e..d1dc0fd6d 100644 --- a/tests/config/test_config_factory.py +++ b/tests/config/test_config_factory.py @@ -2587,6 +2587,26 @@ class TestPlatformOverrides: assert rocm.stages[0].enforce_eager is None assert rocm.stages[1].enforce_eager is True + @pytest.mark.parametrize("deploy_name", ["qwen3_tts.yaml", "qwen3_tts_high_concurrency.yaml"]) + @pytest.mark.parametrize("platform", ["cuda", "npu", "rocm"]) + def test_qwen3_tts_default_code2wav_dtype_is_bf16(self, deploy_name, platform): + deploy = load_deploy_config(Path(get_deploy_config_path(deploy_name))) + deploy = _apply_platform_overrides(deploy, platform=platform) + stages = merge_pipeline_deploy(resolve_pipeline_config("qwen3_tts"), deploy) + + assert stages[1].yaml_engine_args["dtype"] == "bfloat16" + + @pytest.mark.parametrize("deploy_name", ["qwen3_tts.yaml", "qwen3_tts_high_concurrency.yaml"]) + @pytest.mark.parametrize("platform", ["cuda", "npu"]) + @pytest.mark.parametrize("dtype", ["float32", "bfloat16", "float16"]) + def test_qwen3_tts_propagates_explicit_code2wav_dtype(self, deploy_name, platform, dtype): + deploy = load_deploy_config(Path(get_deploy_config_path(deploy_name))) + deploy.stages[1].engine_extras["dtype"] = dtype + deploy = _apply_platform_overrides(deploy, platform=platform) + stages = merge_pipeline_deploy(resolve_pipeline_config("qwen3_tts"), deploy) + + assert stages[1].yaml_engine_args["dtype"] == dtype + def test_higgs_audio_v3_rocm_uses_triton_attention(self): deploy_path = Path(get_deploy_config_path("higgs_multimodal_qwen3.yaml")) diff --git a/tests/model_executor/models/qwen3_tts/test_qwen3_tts_code2wav.py b/tests/model_executor/models/qwen3_tts/test_qwen3_tts_code2wav.py index 64f9ac8de..8ad2b519f 100644 --- a/tests/model_executor/models/qwen3_tts/test_qwen3_tts_code2wav.py +++ b/tests/model_executor/models/qwen3_tts/test_qwen3_tts_code2wav.py @@ -23,7 +23,7 @@ class _FakeDecoder(nn.Module): def __init__(self, total_upsample: int = _TOTAL_UPSAMPLE): super().__init__() self.total_upsample = total_upsample - self.decode_calls: list[dict[str, int]] = [] + self.decode_calls: list[dict[str, int | tuple[int, ...]]] = [] self.batched_decode_calls: list[dict[str, object]] = [] self.decode_codes: list[torch.Tensor] = [] self.cudagraph_calls: list[dict[str, int | torch.device]] = [] @@ -63,7 +63,7 @@ class _FakeDecoder(nn.Module): chunk_size: int = 300, left_context_size: int = 25, max_batch_size: int = 0, - ) -> torch.Tensor: + ) -> list[torch.Tensor]: self.batched_decode_calls.append( { "chunk_size": chunk_size, @@ -91,8 +91,8 @@ class _FakeDecoder(nn.Module): wav_len = frames * self.total_upsample + 6 wav = torch.arange(wav_len, dtype=torch.float32).view(1, 1, -1) offsets = torch.arange(batch, dtype=torch.float32).view(batch, 1, 1) * 1000 - outputs = wav.expand(batch, 1, wav_len) + offsets - return [outputs[row, :, : lengths[row] * self.total_upsample] for row in range(batch)] + batched_wav = wav.expand(batch, 1, wav_len) + offsets + return [batched_wav[row, :, : lengths[row] * self.total_upsample] for row in range(batch)] def enable_cudagraph(self, **kwargs): self.cudagraph_calls.append(kwargs) @@ -131,6 +131,7 @@ def _make_model( load_config=SimpleNamespace(), model_config=SimpleNamespace( model="unused", + dtype=torch.bfloat16, revision=None, stage_connector_config=stage_connector_config, async_chunk=async_chunk, @@ -708,3 +709,16 @@ def test_invalid_decode_chunking_is_rejected(): with pytest.raises(ValueError, match="decode_chunk_frames=0"): _load_weights_noop(model) + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16, torch.float16]) +def test_load_weights_uses_model_dtype_before_precomputing_caches(dtype): + model = _make_model() + model.vllm_config.model_config.dtype = dtype + model.decoder = nn.Linear(2, 2) + cache_dtypes = [] + model.decoder.precompute_snake_caches = lambda: cache_dtypes.append(model.decoder.weight.dtype) + + assert _load_weights_noop(model) == {"decoder.fake_weight"} + assert model.decoder.weight.dtype is dtype + assert cache_dtypes == [dtype] diff --git a/tests/platforms/npu/test_310p_patches.py b/tests/platforms/npu/test_310p_patches.py index c19dcd997..0db4b9d41 100644 --- a/tests/platforms/npu/test_310p_patches.py +++ b/tests/platforms/npu/test_310p_patches.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Unit tests for 310P patch wiring. The tests load patch modules from source with fake Qwen3-TTS dependencies, so @@ -644,7 +644,11 @@ def test_qwen3_tts_tokenizer_npu_patch_dispatches_fused_ops(monkeypatch: pytest. torch.testing.assert_close(norm_out, torch.tensor([[2.0, 3.0]])) -def test_qwen3_tts_code2wav_npu_patch_prepares_loaded_decoder(monkeypatch: pytest.MonkeyPatch) -> None: +def _load_qwen3_tts_code2wav_npu_patch( + monkeypatch: pytest.MonkeyPatch, + *, + platform_dtype_override: bool = True, +): linear_weights = [] conv_weights = [] @@ -674,15 +678,22 @@ def test_qwen3_tts_code2wav_npu_patch_prepares_loaded_decoder(monkeypatch: pytes self.prefix = prefix self.decoder = FakeDecoder() - def _npu_decoder_runtime_dtype(self, _device): - return torch.float16 - def load_weights(self, weights): assert list(weights) == [] + self.decoder.to(dtype=self.vllm_config.model_config.dtype) + self.decoder.precompute_snake_caches() return {"loaded"} - logger = SimpleNamespace(info=lambda *_: None, debug=lambda *_: None) + if platform_dtype_override: + + def _npu_decoder_runtime_dtype(self, _device: torch.device) -> torch.dtype: + return torch.float16 + + FakeCode2Wav._npu_decoder_runtime_dtype = _npu_decoder_runtime_dtype # type: ignore[attr-defined] + + logger = SimpleNamespace(info=lambda *_: None, debug=lambda *_: None, warning=lambda *_: None) current_platform = SimpleNamespace(is_npu=lambda: False) + target = _install_fake_module( monkeypatch, "vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_code2wav", @@ -702,11 +713,23 @@ def test_qwen3_tts_code2wav_npu_patch_prepares_loaded_decoder(monkeypatch: pytes 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() + return module, target, linear_weights, conv_weights - model = target.Qwen3TTSCode2Wav( - vllm_config=SimpleNamespace(device_config=SimpleNamespace(device=torch.device("cpu"))), + +def _make_patched_code2wav(target, *, dtype=torch.bfloat16): + return target.Qwen3TTSCode2Wav( + vllm_config=SimpleNamespace( + device_config=SimpleNamespace(device=torch.device("cpu")), + model_config=SimpleNamespace(dtype=dtype), + ), prefix="stage1", ) + + +def test_qwen3_tts_code2wav_npu_patch_prepares_loaded_decoder(monkeypatch: pytest.MonkeyPatch) -> None: + module, target, linear_weights, conv_weights = _load_qwen3_tts_code2wav_npu_patch(monkeypatch) + model = _make_patched_code2wav(target, dtype=torch.bfloat16) + assert model.load_weights(iter(())) == {"loaded"} assert model.prefix == "stage1" @@ -717,4 +740,17 @@ def test_qwen3_tts_code2wav_npu_patch_prepares_loaded_decoder(monkeypatch: pytes model.decoder.deconv.weight.data_ptr(), } assert all(fmt == module._ACL_FORMAT_FRACTAL_Z for _, fmt in conv_weights) + assert model.decoder.cache_precompute_calls == 2 + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16, torch.float16]) +def test_qwen3_tts_code2wav_npu_patch_preserves_stage_dtype( + monkeypatch: pytest.MonkeyPatch, + dtype: torch.dtype, +) -> None: + _, target, _, _ = _load_qwen3_tts_code2wav_npu_patch(monkeypatch, platform_dtype_override=False) + model = _make_patched_code2wav(target, dtype=dtype) + + assert model.load_weights(iter(())) == {"loaded"} + assert all(parameter.dtype is dtype for parameter in model.decoder.parameters()) assert model.decoder.cache_precompute_calls == 1 diff --git a/vllm_omni/deploy/qwen3_tts.yaml b/vllm_omni/deploy/qwen3_tts.yaml index 747b68852..4019f10aa 100644 --- a/vllm_omni/deploy/qwen3_tts.yaml +++ b/vllm_omni/deploy/qwen3_tts.yaml @@ -95,6 +95,8 @@ stages: top_p: 1.0 - stage_id: 1 + # Decoder compute dtype on GPU/NPU; 310P requires FP16 and overrides this. + dtype: bfloat16 # Tuned for Base voice clone; CustomVoice / VoiceDesign are TTFA-optimal at 1. max_num_seqs: 64 gpu_memory_utilization: 0.3 diff --git a/vllm_omni/deploy/qwen3_tts_high_concurrency.yaml b/vllm_omni/deploy/qwen3_tts_high_concurrency.yaml index 7a73dd50a..7914ecd52 100644 --- a/vllm_omni/deploy/qwen3_tts_high_concurrency.yaml +++ b/vllm_omni/deploy/qwen3_tts_high_concurrency.yaml @@ -85,6 +85,8 @@ stages: top_p: 1.0 - stage_id: 1 + # Decoder compute dtype on GPU/NPU; 310P requires FP16 and overrides this. + dtype: bfloat16 max_num_seqs: 10 gpu_memory_utilization: 0.3 enforce_eager: true diff --git a/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_code2wav.py b/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_code2wav.py index f1692f73d..8ab5c5e34 100644 --- a/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_code2wav.py +++ b/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_code2wav.py @@ -499,7 +499,7 @@ class Qwen3TTSCode2Wav(nn.Module): if wav.numel() == 0: continue if wav.shape[0] > 0: - # Decoder already runs in fp32, so the .to(float32) is a redundant dispatch. + # Return FP32 audio regardless of the decoder compute dtype. audios[idx] = (wav if wav.dtype == torch.float32 else wav.to(torch.float32)).reshape(-1) for req_id, finished, segment_finished in zip( @@ -560,7 +560,7 @@ class Qwen3TTSCode2Wav(nn.Module): ) device = self.vllm_config.device_config.device - self.decoder.to(device=device, dtype=torch.float32) + self.decoder.to(device=device, dtype=self.vllm_config.model_config.dtype) # Precompute SnakeBeta exp caches (benefits both Triton and eager paths) if hasattr(self.decoder, "precompute_snake_caches"): diff --git a/vllm_omni/platforms/npu/models/qwen3_tts.py b/vllm_omni/platforms/npu/models/qwen3_tts.py index a91d2cf80..124e60624 100644 --- a/vllm_omni/platforms/npu/models/qwen3_tts.py +++ b/vllm_omni/platforms/npu/models/qwen3_tts.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Patch Qwen3-TTS NPU runtime setup and weights. @@ -68,11 +68,16 @@ def _patched_load_weights(self, weights): assert _original_load_weights is not None loaded = _original_load_weights(self, weights) device = self.vllm_config.device_config.device - runtime_dtype = getattr(self, "_npu_decoder_runtime_dtype", lambda _: torch.float32)(device) - self.decoder.to(device=device, dtype=runtime_dtype) + platform_dtype_override = getattr(self, "_npu_decoder_runtime_dtype", None) + if platform_dtype_override is not None: + # 310P requires FP16 and deliberately overrides the stage dtype. + runtime_dtype = platform_dtype_override(device) + self.decoder.to(device=device, dtype=runtime_dtype) + if hasattr(self.decoder, "precompute_snake_caches"): + self.decoder.precompute_snake_caches() _prepare_npu_decoder_weights(self.decoder) - if runtime_dtype != torch.float32 and hasattr(self.decoder, "precompute_snake_caches"): - self.decoder.precompute_snake_caches() + actual_dtype = next(self.decoder.parameters()).dtype + logger.info("Prepared NPU Code2Wav decoder with runtime dtype %s", actual_dtype) return loaded @@ -108,8 +113,8 @@ def _apply_a5_prompt_embeds_builder_patch() -> None: # 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 + setattr(prompt_embeds_builder, "Qwen3TTSPromptEmbedsBuilder", _Qwen3TTSPromptEmbedsBuilderA5) + setattr(qwen3_tts_talker, "Qwen3TTSPromptEmbedsBuilder", _Qwen3TTSPromptEmbedsBuilderA5) _A5_PATCHED = True logger.debug("Applied A5 prompt-embeds builder patch for Qwen3-TTS")