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

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

tests/e2e/online_serving/test_cosyvoice3_tts_expansion.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
 """
 E2E Online tests for CosyVoice3 TTS model with voice cloning.
 
@@ -17,12 +17,16 @@ from tests.helpers.mark import hardware_test
 from tests.helpers.media import load_test_audio_data_url
 from tests.helpers.runtime import OmniServerParams
 from tests.helpers.stage_config import get_deploy_config_path
+from vllm_omni.platforms import current_omni_platform
 
-_SKIP_ISSUE_6416 = pytest.mark.skip(
-    reason="https://github.com/vllm-project/vllm-omni/issues/6416",
-)
+if current_omni_platform.is_rocm():
+    os.environ.setdefault("COSYVOICE3_TRT", "0")
 
-pytestmark = [pytest.mark.slow, pytest.mark.tts, _SKIP_ISSUE_6416]
+pytestmark = [
+    pytest.mark.slow,
+    pytest.mark.tts,
+    pytest.mark.core_model,
+]
 
 MODEL = "FunAudioLLM/Fun-CosyVoice3-0.5B-2512"
 
@@ -46,6 +50,13 @@ def get_prompt(prompt_type="zh"):
     return prompts.get(prompt_type, prompts["zh"])
 
 
+def _set_rocm_request_timeout(request_config: dict) -> None:
+    if current_omni_platform.is_rocm():
+        # A cold AITER MHA kernel build took 301 seconds on MI300X.
+        # Allow 600 seconds for the first request; later requests reuse the cache.
+        request_config["timeout"] = 600.0
+
+
 tts_server_params = [
     pytest.param(
         OmniServerParams(
@@ -69,7 +80,7 @@ tts_async_chunk_server_params = [
 ]
 
 
-@hardware_test(res={"cuda": "H100"}, num_cards=1)
+@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=1)
 @pytest.mark.parametrize("omni_server", tts_server_params, indirect=True)
 def test_voice_clone_zh_001(omni_server, openai_client) -> None:
     """
@@ -88,10 +99,11 @@ def test_voice_clone_zh_001(omni_server, openai_client) -> None:
         "ref_audio": REF_AUDIO_URL,
         "ref_text": REF_TEXT,
     }
+    _set_rocm_request_timeout(request_config)
     openai_client.send_audio_speech_request(request_config)
 
 
-@hardware_test(res={"cuda": "H100"}, num_cards=1)
+@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=1)
 @pytest.mark.parametrize("omni_server", tts_async_chunk_server_params, indirect=True)
 def test_voice_clone_zh_002(omni_server, openai_client) -> None:
     """
@@ -111,10 +123,11 @@ def test_voice_clone_zh_002(omni_server, openai_client) -> None:
         "ref_audio": REF_AUDIO_URL,
         "ref_text": REF_TEXT,
     }
+    _set_rocm_request_timeout(request_config)
     openai_client.send_audio_speech_request(request_config)
 
 
-@hardware_test(res={"cuda": "H100"}, num_cards=1)
+@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=1)
 @pytest.mark.parametrize("omni_server", tts_server_params, indirect=True)
 def test_voice_clone_en_001(omni_server, openai_client) -> None:
     """
@@ -133,4 +146,5 @@ def test_voice_clone_en_001(omni_server, openai_client) -> None:
         "ref_audio": REF_AUDIO_URL,
         "ref_text": REF_TEXT,
     }
+    _set_rocm_request_timeout(request_config)
     openai_client.send_audio_speech_request(request_config)

tests/engine/test_data_entry_keys.py

@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 """Tests for data_entry_keys."""
 
 import msgspec
@@ -372,12 +375,12 @@ class TestSerializeDeserializePayload:
         assert torch.equal(restored["hidden_states"]["layers"][24], torch.tensor([3.0]))
 
     def test_tensor_dtype_preserved(self):
-        # bfloat16 excluded: numpy() doesn't support it; callers must cast before serializing.
-        for dtype in [torch.float16, torch.float32, torch.int64, torch.int32, torch.bool]:
+        for dtype in [torch.float16, torch.bfloat16, torch.float32, torch.int64, torch.int32, torch.bool]:
             original: OmniPayload = {"codes": {"audio": torch.tensor([1], dtype=dtype)}}
             wire = serialize_payload(original)
             restored = deserialize_payload(wire)
             assert restored["codes"]["audio"].dtype == dtype, f"dtype mismatch for {dtype}"
+            assert torch.equal(restored["codes"]["audio"], original["codes"]["audio"])
 
     def test_tensor_shape_preserved(self):
         t = torch.randn(3, 4, 5)

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(

vllm_omni/data_entry_keys.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
 """Structured payload types for inter-stage communication.
 
 Categories under ``OmniPayload``:
@@ -12,6 +12,7 @@ Categories under ``OmniPayload``:
 
 from __future__ import annotations
 
+from collections.abc import Mapping
 from typing import TYPE_CHECKING, Any, TypedDict
 
 import msgspec
@@ -117,7 +118,9 @@ class OmniPayload(TypedDict, total=False):
 # ── msgspec.Struct mirror of the TypedDicts (runtime-validated) ──
 
 
-class _StructBase(msgspec.Struct, omit_defaults=True, kw_only=True, forbid_unknown_fields=True):
+class _StructBase(  # type: ignore[call-arg]
+    msgspec.Struct, omit_defaults=True, kw_only=True, forbid_unknown_fields=True
+):
     pass
 
 
@@ -310,7 +313,7 @@ def _dtype_to_name(dtype: torch.dtype) -> str:
 _NESTED_KEYS = frozenset({"hidden_states", "embed", "ids", "codes", "meta"})
 
 
-def flatten_payload(payload: dict[str, Any]) -> dict[str, Any]:
+def flatten_payload(payload: Mapping[str, Any]) -> dict[str, Any]:
     """Flatten a nested ``OmniPayload`` to dotted keys.
 
     Nested sub-dicts under ``_NESTED_KEYS`` are expanded:
@@ -361,17 +364,25 @@ def _serialize_tensor(t: torch.Tensor) -> AdditionalInformationEntry:
 
     t_cpu = t.detach().to("cpu").contiguous()
     return AdditionalInformationEntry(
-        tensor_data=t_cpu.numpy().tobytes(),
+        tensor_data=t_cpu.flatten().view(torch.uint8).numpy().tobytes(),
         tensor_shape=list(t_cpu.shape),
         tensor_dtype=_dtype_to_name(t_cpu.dtype),
     )
 
 
 def _deserialize_tensor(entry: AdditionalInformationEntry) -> torch.Tensor:
-    dt = np.dtype(entry.tensor_dtype or "float32")
-    arr = np.frombuffer(entry.tensor_data, dtype=dt)  # type: ignore[arg-type]
-    arr = arr.reshape(entry.tensor_shape)
-    return torch.from_numpy(arr.copy())
+    dtype_name = entry.tensor_dtype or "float32"
+    dtype = getattr(torch, dtype_name, None)
+    if not isinstance(dtype, torch.dtype):
+        raise ValueError(f"Unsupported tensor dtype: {dtype_name}")
+
+    if entry.tensor_shape is None:
+        raise ValueError("Tensor shape is required")
+    if not entry.tensor_data:
+        return torch.empty(0, dtype=dtype).reshape(entry.tensor_shape)
+
+    data = torch.frombuffer(bytearray(entry.tensor_data), dtype=torch.uint8)
+    return data.view(dtype).reshape(entry.tensor_shape)
 
 
 def serialize_payload(

vllm_omni/model_executor/models/cosyvoice3/code2wav_core/hifigan.py

@@ -494,7 +494,7 @@ class HiFTGenerator(nn.Module):
             self.istft_params["n_fft"],
             self.istft_params["hop_len"],
             self.istft_params["n_fft"],
-            window=self.stft_window,
+            window=self._get_stft_window(x),
             return_complex=True,
         )
         spec = torch.view_as_real(spec)  # [B, F, TT, 2]
@@ -509,10 +509,15 @@ class HiFTGenerator(nn.Module):
             self.istft_params["n_fft"],
             self.istft_params["hop_len"],
             self.istft_params["n_fft"],
-            window=self.stft_window,
+            window=self._get_stft_window(magnitude),
         )
         return inverse_transform
 
+    def _get_stft_window(self, tensor: torch.Tensor) -> torch.Tensor:
+        if self.stft_window.device != tensor.device:
+            self.stft_window = self.stft_window.to(tensor.device)
+        return self.stft_window
+
     def _decode_pre_istft(
         self,
         x: torch.Tensor,

vllm_omni/model_executor/models/cosyvoice3/cosyvoice3.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
 import os
 from collections.abc import Iterable, Mapping, Sequence
 from functools import partial
@@ -35,6 +35,7 @@ from vllm.multimodal.processing import (
 from vllm.sequence import IntermediateTensors
 from vllm.v1.outputs import SamplerOutput
 from vllm.v1.sample.metadata import SamplingMetadata
+from vllm.v1.sample.ops.topk_topp_sampler import random_sample
 from vllm.v1.sample.sampler import Sampler
 
 from vllm_omni.data_entry_keys import EmbeddingsStruct, OmniPayloadStruct, to_dict, to_struct
@@ -566,8 +567,8 @@ class CosyVoice3Model(
         return float(value)
 
     @staticmethod
-    def _multinomial_sample(probs: torch.Tensor, generator: torch.Generator | None = None) -> torch.Tensor:
-        return torch.multinomial(probs, 1, replacement=True, generator=generator).reshape(())
+    def _random_sample_one(probs: torch.Tensor, generator: torch.Generator | None = None) -> torch.Tensor:
+        return random_sample(probs.unsqueeze(0), {} if generator is None else {0: generator}).reshape(())
 
     @classmethod
     def _nucleus_sample_one(
@@ -602,7 +603,7 @@ class CosyVoice3Model(
         # First token always passes (cum_before[0] = 0 < top_p for any top_p > 0),
         # so ``weights`` is guaranteed to have at least one nonzero entry. The
         # final ``.item()`` is the ONLY D2H sync per call.
-        sample_idx = torch.multinomial(weights, 1, replacement=True, generator=generator)
+        sample_idx = cls._random_sample_one(weights, generator=generator)
         return int(sorted_idx[sample_idx].item())
 
     @classmethod
@@ -632,9 +633,14 @@ class CosyVoice3Model(
             rep_num = int((recent == top_id).sum().item())
             if rep_num >= win_size * tau_r:
                 weighted_scores = weighted_scores.clone()
+                original_score = weighted_scores[top_id].clone()
                 weighted_scores[top_id] = float("-inf")
-                fallback_probs = weighted_scores.softmax(dim=0)
-                top_id = int(cls._multinomial_sample(fallback_probs, generator=generator).item())
+                weighted_scores[top_id] = torch.where(
+                    torch.isfinite(weighted_scores).any(),
+                    weighted_scores[top_id],
+                    original_score,
+                )
+                top_id = int(cls._random_sample_one(weighted_scores.softmax(dim=0), generator=generator).item())
         return top_id
 
     def _cosyvoice3_ras_enabled(self, sampling_metadata: SamplingMetadata) -> bool:
@@ -677,6 +683,10 @@ class CosyVoice3Model(
             logits.masked_fill_(sampling_metadata.allowed_token_ids_mask, float("-inf"))
         for processor in sampling_metadata.logitsprocs.non_argmax_invariant:
             logits = processor.apply(logits)
+        finite_logits = torch.isfinite(logits)
+        if not finite_logits.any(dim=-1).all().item():
+            raise ValueError("CosyVoice3 sampling received a row with no finite logits")
+        logits.masked_fill_(~finite_logits, float("-inf"))
 
         sampling_cfg = dict(self.config.llm.get("sampling", {}))
         default_top_p = float(sampling_cfg.get("top_p", 0.8))

vllm_omni/model_executor/stage_input_processors/cosyvoice3.py

@@ -1,11 +1,10 @@
 # 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 collections.abc import Mapping
 from contextlib import nullcontext
 from typing import Any
 
-import numpy as np
 import torch
 from vllm.inputs import TextPrompt
 from vllm.logger import init_logger
@@ -16,6 +15,7 @@ from vllm_omni.data_entry_keys import (
     MetaStruct,
     OmniPayloadStruct,
 )
+from vllm_omni.engine.serialization import deserialize_additional_information
 from vllm_omni.inputs.data import OmniTokensPrompt
 from vllm_omni.model_executor.models.cosyvoice3.utils import unpad_prompt_conditioning
 
@@ -93,32 +93,6 @@ def _to_cpu_tensor(x: Any) -> torch.Tensor | None:
     return None
 
 
-def _decode_additional_information(raw_info: Any) -> dict[str, Any]:
-    if raw_info is None:
-        return {}
-    if isinstance(raw_info, dict):
-        return raw_info
-
-    entries = getattr(raw_info, "entries", None)
-    if not isinstance(entries, dict):
-        return {}
-
-    decoded: dict[str, Any] = {}
-    for key, entry in entries.items():
-        tensor_data = getattr(entry, "tensor_data", None)
-        if tensor_data is not None:
-            dtype_name = getattr(entry, "tensor_dtype", "float32")
-            tensor_shape = getattr(entry, "tensor_shape", None)
-            if tensor_shape is None:
-                continue
-            dt = np.dtype(dtype_name)
-            arr = np.frombuffer(tensor_data, dtype=dt).reshape(tensor_shape)
-            decoded[key] = torch.from_numpy(arr.copy())
-        else:
-            decoded[key] = getattr(entry, "list_data", None)
-    return decoded
-
-
 def talker2code2wav_async_chunk(
     transfer_manager: Any,
     multimodal_output: dict[str, Any] | None,
@@ -149,7 +123,7 @@ def talker2code2wav_async_chunk(
         request_state = transfer_manager.request_payload.get(request_id)
         if not isinstance(request_state, dict) or "_cosyvoice3_async_state" not in request_state:
             with nullcontext():
-                info = _decode_additional_information(getattr(request, "additional_information", None))
+                info = deserialize_additional_information(getattr(request, "additional_information", None))
                 info_embed = info.get("embed", {}) if isinstance(info, dict) else {}
                 prompt_payload = {}
                 cond_keys = ("speech_token", "speech_feat", "embedding", "speech_token_len")

vllm_omni/transformers_utils/configs/cosyvoice3.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 transformers import AutoConfig
 from transformers.configuration_utils import PretrainedConfig
 
@@ -18,6 +18,7 @@ class CosyVoice3Config(PretrainedConfig):
         self.llm_output_size = 896
         self.hidden_size = self.llm_output_size
         self.num_attention_heads = 14
+        self.num_key_value_heads = 2
         self.num_hidden_layers = 24
         self.spk_embed_dim = 192
         self.token_frame_rate = 25