[Hardware] Enable CosyVoice3 on Moore Threads MUSA (#5164)

Signed-off-by: Xiaodong Ye <xiaodong.ye@mthreads.com>

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:

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 == {}

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

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

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

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