[tts][bugfix] Fix CosyVoice3 sampling and stage handoff (#6424)

Signed-off-by: akshatvishu <akshatnayak197@gmail.com>

tests/model_executor/models/cosyvoice3/test_cosyvoice3_components.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 CosyVoice3 components."""
 
 from types import SimpleNamespace
@@ -9,10 +9,41 @@ import torch
 import torch.nn as nn
 
 from tests.helpers.mark import hardware_test
+from vllm_omni.model_executor.models.cosyvoice3.code2wav_core.hifigan import (
+    CausalHiFTGenerator,
+)
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
 
+@pytest.fixture
+def causal_hift():
+    return CausalHiFTGenerator(
+        base_channels=32,
+        upsample_rates=[2, 2],
+        upsample_kernel_sizes=[4, 4],
+        source_resblock_kernel_sizes=[3, 3],
+        source_resblock_dilation_sizes=[[1, 3, 5], [1, 3, 5]],
+    )
+
+
+def test_causal_hift_moves_stft_window_with_model(causal_hift):
+    assert causal_hift.get_buffer("stft_window") is causal_hift.stft_window
+    assert "stft_window" not in causal_hift.state_dict()
+    causal_hift.to(dtype=torch.float64)
+    assert causal_hift.stft_window.dtype == torch.float64
+
+
+def test_causal_hift_stft_moves_window_to_input_device(causal_hift):
+    waveform = torch.empty((1, 64), device="meta")
+
+    real, imag = causal_hift._stft(waveform)
+
+    assert real.device == waveform.device
+    assert imag.device == waveform.device
+    assert causal_hift.stft_window.device == waveform.device
+
+
 class TestPreLookaheadLayer:
     """Tests for PreLookaheadLayer."""
 

tests/model_executor/models/cosyvoice3/test_cosyvoice3_config.py

@@ -0,0 +1,18 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+import pytest
+
+from vllm_omni.config.model import OmniModelArchConfigConvertor
+from vllm_omni.transformers_utils.configs.cosyvoice3 import CosyVoice3Config
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def test_attention_geometry_matches_qwen_backbone() -> None:
+    config = CosyVoice3Config()
+    model_arch_config = OmniModelArchConfigConvertor(config, config.get_text_config()).convert()
+
+    assert model_arch_config.total_num_attention_heads == 14
+    assert model_arch_config.total_num_kv_heads == 2
+    assert model_arch_config.head_size == 64

tests/model_executor/models/cosyvoice3/test_cosyvoice3_model_helpers.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
 
 from __future__ import annotations
 
@@ -401,6 +401,52 @@ def test_sample_tolerates_padded_rows_without_history():
     assert out.sampled_token_ids.shape == (2, 1)
 
 
+def test_sample_excludes_non_finite_logits():
+    model = _make_talker_model()
+    metadata = _make_sampling_metadata(output_token_ids=[[]])
+    metadata.temperature.fill_(0.5)
+    logits = torch.tensor([[float("nan"), 1.0, float("inf"), float("-inf")]], dtype=torch.bfloat16)
+
+    out = model.sample(logits, metadata)
+
+    assert out is not None
+    assert out.sampled_token_ids.tolist() == [[1]]
+
+
+def test_sample_preserves_allowed_token_mask_with_invalid_logits():
+    model = _make_talker_model()
+    metadata = _make_sampling_metadata(output_token_ids=[[]])
+    metadata.allowed_token_ids_mask = torch.tensor([[True, False, True]])
+    logits = torch.tensor([[float("nan"), 1.0, float("inf")]], dtype=torch.float32)
+
+    out = model.sample(logits, metadata)
+
+    assert out is not None
+    assert out.sampled_token_ids.tolist() == [[1]]
+
+
+def test_sample_rejects_rows_without_finite_logits():
+    model = _make_talker_model()
+    metadata = _make_sampling_metadata(output_token_ids=[[]])
+    metadata.allowed_token_ids_mask = torch.tensor([[True, False, True]])
+    logits = torch.tensor([[0.0, float("nan"), 0.0]], dtype=torch.float32)
+
+    with pytest.raises(ValueError, match="no finite logits"):
+        model.sample(logits, metadata)
+
+
+def test_sample_keeps_only_finite_token_after_ras_rejection():
+    model = _make_talker_model()
+    metadata = _make_sampling_metadata(output_token_ids=[[1] * 10])
+    metadata.allowed_token_ids_mask = torch.tensor([[True, False, True]])
+    logits = torch.tensor([[0.0, 1.0, 0.0]], dtype=torch.float32)
+
+    out = model.sample(logits, metadata)
+
+    assert out is not None
+    assert out.sampled_token_ids.tolist() == [[1]]
+
+
 def test_gpu_ar_model_runner_prefers_model_sampler_when_opted_in():
     metadata = _make_sampling_metadata(output_token_ids=[[1, 2, 3]])
     expected = SamplerOutput(
@@ -423,9 +469,14 @@ def test_gpu_ar_model_runner_prefers_model_sampler_when_opted_in():
     _, GPUARModelRunner = _cosyvoice3_model_and_runner()
     runner = object.__new__(GPUARModelRunner)
     runner.input_batch = _DummyInputBatch()
+
+    def model_sample(logits, sampling_metadata):
+        calls.append(logits.clone())
+        return expected
+
     runner.model = SimpleNamespace(
         prefer_model_sampler=True,
-        sample=lambda logits, sampling_metadata: calls.append(logits.clone()) or expected,
+        sample=model_sample,
     )
     runner.sampler = lambda **_: (_ for _ in ()).throw(AssertionError("fallback sampler should not be used"))
 
@@ -459,12 +510,14 @@ def test_gpu_ar_model_runner_supplies_req_output_history_to_model_sampler():
     _, GPUARModelRunner = _cosyvoice3_model_and_runner()
     runner = object.__new__(GPUARModelRunner)
     runner.input_batch = _DummyInputBatch()
+
+    def model_sample(logits, sampling_metadata):
+        seen_histories.append([list(x) for x in sampling_metadata.output_token_ids])
+        return SamplerOutput(sampled_token_ids=torch.tensor([[7]], dtype=torch.int32), logprobs_tensors=None)
+
     runner.model = SimpleNamespace(
         prefer_model_sampler=True,
-        sample=lambda logits, sampling_metadata: (
-            seen_histories.append([list(x) for x in sampling_metadata.output_token_ids])
-            or SamplerOutput(sampled_token_ids=torch.tensor([[7]], dtype=torch.int32), logprobs_tensors=None)
-        ),
+        sample=model_sample,
     )
     runner.sampler = lambda **_: (_ for _ in ()).throw(AssertionError("fallback sampler should not be used"))
 
@@ -504,12 +557,14 @@ def test_gpu_ar_model_runner_repairs_async_placeholders_for_model_sampler():
     _, GPUARModelRunner = _cosyvoice3_model_and_runner()
     runner = object.__new__(GPUARModelRunner)
     runner.input_batch = _DummyInputBatch()
+
+    def model_sample(logits, sampling_metadata):
+        seen_histories.append([list(x) for x in sampling_metadata.output_token_ids])
+        return SamplerOutput(sampled_token_ids=torch.tensor([[7]], dtype=torch.int32), logprobs_tensors=None)
+
     runner.model = SimpleNamespace(
         prefer_model_sampler=True,
-        sample=lambda logits, sampling_metadata: (
-            seen_histories.append([list(x) for x in sampling_metadata.output_token_ids])
-            or SamplerOutput(sampled_token_ids=torch.tensor([[7]], dtype=torch.int32), logprobs_tensors=None)
-        ),
+        sample=model_sample,
     )
     runner.sampler = lambda **_: (_ for _ in ()).throw(AssertionError("fallback sampler should not be used"))
 

tests/model_executor/stage_input_processors/test_cosyvoice3_stage_input_processors.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
 
 from collections import defaultdict
 from types import SimpleNamespace
@@ -7,6 +7,7 @@ from types import SimpleNamespace
 import pytest
 import torch
 
+from vllm_omni.data_entry_keys import serialize_payload
 from vllm_omni.model_executor.stage_input_processors.cosyvoice3 import (
     talker2code2wav_async_chunk,
     text2flow_full_payload,
@@ -153,6 +154,36 @@ def test_talker2code2wav_async_chunk_final_payload_uses_absolute_token_offset():
     assert payload.embed.embedding is not None
 
 
+def test_talker2code2wav_async_chunk_decodes_bfloat16_conditioning():
+    speech_feat = torch.tensor([[[0.1, 0.2], [0.3, 0.4]]], dtype=torch.bfloat16)
+    additional_information = serialize_payload(
+        {
+            "embed": {
+                "speech_token": torch.tensor([[11]], dtype=torch.long),
+                "speech_feat": speech_feat,
+                "embedding": torch.tensor([[0.5, 0.6]], dtype=torch.bfloat16),
+            }
+        }
+    )
+    request = SimpleNamespace(
+        external_req_id="rid-bfloat16",
+        output_token_ids=[1],
+        additional_information=additional_information,
+        is_finished=lambda: True,
+    )
+
+    payload = talker2code2wav_async_chunk(
+        transfer_manager=_transfer_manager(),
+        multimodal_output=None,
+        request=request,
+        is_finished=True,
+    )
+
+    assert payload is not None
+    assert payload.embed.speech_feat.dtype == torch.bfloat16
+    assert torch.equal(payload.embed.speech_feat, speech_feat)
+
+
 def test_talker2code2wav_async_chunk_emits_eof_when_finished_without_valid_codes():
     transfer_manager = _transfer_manager(chunk_frames=25)
     request = SimpleNamespace(