From ab0da662c069618e0cd35b309d8f41310c997d03 Mon Sep 17 00:00:00 2001 From: R0CKSTAR Date: Wed, 29 Jul 2026 10:46:27 +0800 Subject: [PATCH] [Hardware] Enable CosyVoice3 on Moore Threads MUSA (#5164) Signed-off-by: Xiaodong Ye --- .../text_to_speech/cosyvoice3/end2end.py | 7 +++++-- tests/engine/test_output_processor.py | 4 ++++ vllm_omni/deploy/cosyvoice3.yaml | 2 -- .../models/cosyvoice3_audio/cosyvoice3_dit.py | 4 ++-- .../models/cosyvoice3/cosyvoice3.py | 16 +++++++++++++++- vllm_omni/outputs/output_modality.py | 17 ++++++++++++++++- 6 files changed, 42 insertions(+), 8 deletions(-) diff --git a/examples/offline_inference/text_to_speech/cosyvoice3/end2end.py b/examples/offline_inference/text_to_speech/cosyvoice3/end2end.py index e5c266f7a..ccfb1aaef 100644 --- a/examples/offline_inference/text_to_speech/cosyvoice3/end2end.py +++ b/examples/offline_inference/text_to_speech/cosyvoice3/end2end.py @@ -182,10 +182,13 @@ def run_e2e(): print(f"Multimodal output keys: {mm.keys()}") if "audio" in mm: audio_out = mm["audio"] + # The model emits its own rate (24 kHz for CosyVoice3); take + # it from the payload so the file is not mislabeled. + sample_rate = int(mm["sr"]) if "sr" in mm.keys() else 24000 print(f"Generated Audio Shape: {audio_out.shape}") out_path = f"output_{i}.wav" - sf.write(out_path, audio_out.cpu().numpy().squeeze(), 22050) - print(f"Saved audio to {out_path}") + sf.write(out_path, audio_out.cpu().numpy().squeeze(), sample_rate) + print(f"Saved audio to {out_path} (sr={sample_rate})") else: print("No multimodal output found.") except Exception as e: diff --git a/tests/engine/test_output_processor.py b/tests/engine/test_output_processor.py index df0911aa0..81f0c8303 100644 --- a/tests/engine/test_output_processor.py +++ b/tests/engine/test_output_processor.py @@ -59,6 +59,10 @@ def _make_state(output_kind: RequestOutputKind): return OmniRequestState(**_DEFAULT_STATE_KWARGS, output_kind=output_kind) +def test_output_modality_name_uses_string_value(): + assert str(OutputModalityNames.AUDIO) == "audio" + + def test_init_empty_dict(): """Ensure mm_accumulated is initially empty.""" assert _make_state(RequestOutputKind.CUMULATIVE).mm_accumulated == {} diff --git a/vllm_omni/deploy/cosyvoice3.yaml b/vllm_omni/deploy/cosyvoice3.yaml index d81cf4a9b..551b42af6 100644 --- a/vllm_omni/deploy/cosyvoice3.yaml +++ b/vllm_omni/deploy/cosyvoice3.yaml @@ -24,8 +24,6 @@ # first. Raising connector_get_sleep_s also relaxes the consumer's polling # cadence, but trades against TTFP — keep it as a last resort. async_chunk: true -dtype: float32 - connectors: connector_of_shared_memory: name: SharedMemoryConnector diff --git a/vllm_omni/diffusion/models/cosyvoice3_audio/cosyvoice3_dit.py b/vllm_omni/diffusion/models/cosyvoice3_audio/cosyvoice3_dit.py index 6333819b0..5ae75b256 100644 --- a/vllm_omni/diffusion/models/cosyvoice3_audio/cosyvoice3_dit.py +++ b/vllm_omni/diffusion/models/cosyvoice3_audio/cosyvoice3_dit.py @@ -135,8 +135,8 @@ class DiTAttention(nn.Module): # The diffusion Attention layer expects (batch, seq, heads, head_dim) out = self.attn(query, key, value, attn_metadata=None) - # Reshape back: (batch, seq, dim) - out = out.view(batch_size, seq_len, self.inner_dim) + # Some attention backends return a non-contiguous layout. + out = out.reshape(batch_size, seq_len, self.inner_dim) out = out.to(query.dtype) # Output projection diff --git a/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3.py b/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3.py index b3e178565..f2f9a6a0f 100644 --- a/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3.py +++ b/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3.py @@ -73,6 +73,19 @@ def _cosyvoice3_trt_enabled() -> bool: return os.environ.get("COSYVOICE3_TRT", "1") not in ("0", "false", "False", "") +def _campplus_onnx_providers() -> list[str]: + """ONNX-Runtime providers for the campplus speaker-embedding session. + + Prefer ``MUSAExecutionProvider`` (from the onnxruntime-musa build) with a + CPU fallback when it is available; otherwise CPU only. The MUSA kernels + differ from the CPU reference by a few percent on the embedding, which + conditions voice cloning -- verify voice similarity if that matters. + """ + if "MUSAExecutionProvider" in onnxruntime.get_available_providers(): + return ["MUSAExecutionProvider", "CPUExecutionProvider"] + return ["CPUExecutionProvider"] + + class CosyVoice3MultiModalProcessingInfo(BaseProcessingInfo): def get_hf_config(self): """If the config is not already present pass it @@ -163,7 +176,7 @@ class CosyVoice3MultiModalProcessor(BaseMultiModalProcessor[CosyVoice3MultiModal campplus_session = onnxruntime.InferenceSession( campplus_onnx_path, sess_options=option, - providers=["CPUExecutionProvider"], + providers=_campplus_onnx_providers(), ) return { @@ -922,6 +935,7 @@ class CosyVoice3Model( if getattr(self, "_code2wav_trt_done", False): return self._code2wav_trt_done = True + if not (_cosyvoice3_trt_enabled() and torch.cuda.is_available()): return onnx_path = self._resolve_flow_estimator_onnx() diff --git a/vllm_omni/outputs/output_modality.py b/vllm_omni/outputs/output_modality.py index 2536a8ade..743669454 100644 --- a/vllm_omni/outputs/output_modality.py +++ b/vllm_omni/outputs/output_modality.py @@ -8,7 +8,22 @@ for type-safe multimodal output routing and tensor merging. from __future__ import annotations import re -from enum import Enum, Flag, StrEnum, auto +import sys +from enum import Enum, Flag, auto + +if sys.version_info >= (3, 11): + from enum import StrEnum +else: + # TODO: Remove this Python 3.10 compatibility shim when Python 3.10 support is dropped. + class StrEnum(str, Enum): + """``enum.StrEnum`` for Python 3.10, which predates it. + + ``__str__`` is taken from ``str`` so ``str(member)`` yields the member + value rather than ``"ClassName.MEMBER"``, matching 3.11 semantics. + """ + + __str__ = str.__str__ + _MODALITY_ALIASES: dict[str, str] = { "speech": "audio",