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

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

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]