[Feature][Qwen3-TTS] Make Code2Wav honor the stage dtype (#6059)

Signed-off-by: XIN GAO <1037396230@qq.com>

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"))
 

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]

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

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

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

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"):

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")