[Model] Add Ming-omni-tts MoE 16.8B-A3B + CFM CUDAGraph (#4341)

Signed-off-by: LHXuuu <xulianhao.xlh@antgroup.com>
Signed-off-by: Yuanheng Zhao <jonathan.zhaoyh@gmail.com>
Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
Co-authored-by: Yuanheng Zhao <jonathan.zhaoyh@gmail.com>
Co-authored-by: Yuanheng Zhao <54058983+yuanheng-zhao@users.noreply.github.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>

tests/model_executor/models/ming_tts/test_ming_shared_modules.py

@@ -137,8 +137,12 @@ def test_ming_compute_logits_falls_back_to_dummy_token_id():
 
 def test_ming_forward_non_decode_return_clears_cached_forced_next_token_ids():
     class FakeBackbone:
-        def __call__(self, **kwargs):
-            return kwargs["inputs_embeds"]
+        # MingLLMModel.forward calls the backbone positionally
+        # (input_ids, positions, intermediate_tensors, inputs_embeds) so the
+        # same signature works for both Qwen2Model (``positions``) and
+        # BailingMoeModel (``position_ids``).
+        def __call__(self, input_ids, positions, intermediate_tensors, inputs_embeds, **kwargs):
+            return inputs_embeds
 
     model = _make_ming_logits_model()
     model.model = FakeBackbone()

vllm_omni/config/pipeline_registry.py

@@ -63,7 +63,10 @@ from vllm_omni.model_executor.models.ming_flash_omni.pipeline import (
     MING_FLASH_OMNI_THINKER_ONLY_PIPELINE,
     MING_FLASH_OMNI_TTS_PIPELINE,
 )
-from vllm_omni.model_executor.models.ming_tts.pipeline import MING_TTS_PIPELINE
+from vllm_omni.model_executor.models.ming_tts.pipeline import (
+    MING_TTS_MOE_PIPELINE,
+    MING_TTS_PIPELINE,
+)
 from vllm_omni.model_executor.models.minicpmo_4_5.pipeline import MINICPMO_4_5_PIPELINE
 from vllm_omni.model_executor.models.moss_tts.pipeline import (
     MOSS_TTS_LOCAL_PIPELINE,
@@ -106,6 +109,7 @@ OMNI_PIPELINES: dict[str, PipelineConfig | PipelineResolverFunc] = {
     "cosyvoice3": COSYVOICE3_PIPELINE,
     "mimo_audio": MIMO_AUDIO_PIPELINE,
     "ming_tts": MING_TTS_PIPELINE,
+    "ming_tts_moe": MING_TTS_MOE_PIPELINE,
     "voxtral_tts": VOXTRAL_TTS_PIPELINE,
     "glm_tts": GLM_TTS_PIPELINE,
     "fish_qwen3_omni": FISH_SPEECH_PIPELINE,

vllm_omni/deploy/ming_tts_moe.yaml

@@ -0,0 +1,66 @@
+# Ming-omni-tts-16.8B-A3B (MoE) deploy: stage-0 bailing_moe LLM+flow -> stage-1 audio VAE.
+# Same two-stage pipeline as the dense ming_tts.yaml; the Stage-0 backbone is the
+# 16.8B-A3B bailing_moe MoE (~31GB bf16).
+#
+# pipeline: HF config reports model_type="bailingmm", which is not a registry key,
+# already declared by MING_TTS_MOE_PIPELINE for auto-detection.
+#
+# NOTE: ~31GB bf16 + the tiny audio VAE fit one 80GB+ GPU, so both stages colocate
+# on GPU0 (Stage-0 0.60 util ≈ 58GB + Stage-1 0.25 util ≈ 24GB ≈ 0.85, fits 97GB
+# H20). For smaller cards raise Stage-0 tensor_parallel_size or split the stages
+# onto separate devices. Multi-request batching is not yet validated (max_num_seqs: 1).
+pipeline: ming_tts_moe
+async_chunk: true
+# The MoE ships a custom tokenizer (tokenization_bailing.BailingTokenizer) that
+# must be executed to load. The model config still resolves to the registered
+# MingMoeConfig (vLLM _CONFIG_REGISTRY precedence over the repo auto_map).
+trust_remote_code: true
+dtype: bfloat16
+enable_prefix_caching: false
+
+connectors:
+  connector_of_shared_memory:
+    name: SharedMemoryConnector
+    extra:
+      latent_chunk_size: 25
+      latent_left_context: 0
+
+stages:
+  - stage_id: 0
+    max_num_seqs: 1
+    tensor_parallel_size: 1
+    gpu_memory_utilization: 0.60
+    enforce_eager: true
+    enable_prefix_caching: false
+    async_scheduling: false
+    max_num_batched_tokens: 8192
+    max_model_len: 8192
+    # Pin the unquantized MoE kernel (wer diff between vllm 0.22 vs 0.23)
+    moe_backend: triton
+    devices: "0"
+    output_connectors:
+      to_stage_1: connector_of_shared_memory
+    default_sampling_params:
+      temperature: 0.0
+      top_p: 1.0
+      top_k: -1
+      max_tokens: 512
+      detokenize: true
+
+  - stage_id: 1
+    max_num_seqs: 1
+    gpu_memory_utilization: 0.25
+    enforce_eager: true
+    enable_prefix_caching: false
+    async_scheduling: false
+    max_num_batched_tokens: 8192
+    max_model_len: 8192
+    devices: "0"
+    input_connectors:
+      from_stage_0: connector_of_shared_memory
+    default_sampling_params:
+      temperature: 0.0
+      top_p: 1.0
+      top_k: -1
+      max_tokens: 1
+      detokenize: false

vllm_omni/engine/arg_utils.py

@@ -40,7 +40,10 @@ def _register_omni_hf_configs() -> None:
         from vllm_omni.model_executor.models.indextts2.configuration_indextts2 import (
             IndexTTS2Config,
         )
-        from vllm_omni.model_executor.models.ming_tts.config_ming_tts import MingDenseConfig
+        from vllm_omni.model_executor.models.ming_tts.config_ming_tts import (
+            MingDenseConfig,
+            MingMoeConfig,
+        )
         from vllm_omni.model_executor.models.moss_tts.configuration_moss_tts import (
             MossTTSLocalConfig,
             MossTTSRealtimeConfig,
@@ -66,6 +69,7 @@ def _register_omni_hf_configs() -> None:
 
     for model_type, config_cls in [
         ("dense", MingDenseConfig),
+        ("bailingmm", MingMoeConfig),
         ("indextts2", IndexTTS2Config),
         ("moss_tts_local", MossTTSLocalConfig),
         ("moss_tts_realtime", MossTTSRealtimeConfig),

vllm_omni/entrypoints/openai/serving_speech.py

@@ -2596,7 +2596,8 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
 
         if self._tts_tokenizer is None:
             model_name = self.engine_client.model_config.model
-            self._tts_tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=False)
+            trust_remote_code = bool(getattr(self.engine_client.model_config, "trust_remote_code", False))
+            self._tts_tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=trust_remote_code)
 
         ref_text = request.ref_text
         prompt_waveform = self._build_ming_prompt_waveform(ref_audio_data) if ref_text is not None else None
@@ -3585,10 +3586,17 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
         elif self._tts_model_type == "ming_tts" and sampling_params_list:
             import copy
 
-            from vllm_omni.model_executor.models.ming_tts.config_ming_tts import TEXT_EOS_TOKEN_ID
+            from vllm_omni.model_executor.models.ming_tts.config_ming_tts import (
+                MOE_TEXT_EOS_TOKEN_ID,
+                TEXT_EOS_TOKEN_ID,
+            )
+
+            hf_config = self.engine_client.model_config.hf_config
+            is_moe = getattr(hf_config, "model_type", "") == "bailingmm"
+            stop_token_id = MOE_TEXT_EOS_TOKEN_ID if is_moe else TEXT_EOS_TOKEN_ID
 
             sampling_params_list = copy.deepcopy(sampling_params_list)
-            sampling_params_list[0].stop_token_ids = [int(TEXT_EOS_TOKEN_ID)]
+            sampling_params_list[0].stop_token_ids = [int(stop_token_id)]
             if request.max_new_tokens is not None:
                 # Ming emits TEXT_EOS after the latent decode budget is exhausted, so
                 # Stage-0 needs one extra token beyond ming_max_decode_steps.

vllm_omni/model_executor/models/common/ming/audio_vae.py

@@ -346,11 +346,8 @@ class AudioVAE(PreTrainedModel):
         h, _ = self.encoder(waveform)
         h = h.transpose(1, 2)
 
-        mean, logvar = torch.chunk(h, 2, dim=1)
-        logvar = torch.clamp(logvar, -30.0, 20.0)
-        std = torch.exp(0.5 * logvar)
-        latent = mean + std * torch.randn_like(mean)
-        latent = latent.transpose(1, 2)
+        mean = torch.chunk(h, 2, dim=1)[0]
+        latent = mean.transpose(1, 2)
         return latent, frame_num
 
     @torch.inference_mode()

vllm_omni/model_executor/models/ming_tts/__init__.py

@@ -1,6 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
 # NOTE: Do not import model classes in this file. Importing any
 # submodule in this package triggers __init__.py execution, and
 # both the model registry and pipeline registry import submodules

vllm_omni/model_executor/models/ming_tts/audio_prep.py

@@ -14,7 +14,6 @@ from .config_ming_tts import (
     PATCH_SIZE,
     SAMPLE_RATE,
     VAE_PATCH_SIZE,
-    VISION_START_TOKEN_ID,
     MingTTSConfig,
 )
 
@@ -232,14 +231,17 @@ def _find_audio_placeholder_positions(input_ids: torch.Tensor, cfg: MingTTSConfi
     return filtered if filtered.numel() > 0 else dummy_pos
 
 
-def _find_speaker_placeholder_positions(input_ids: torch.Tensor, hf_config: Any) -> list[int]:
-    vision_start_token_id = getattr(hf_config, "vision_start_token_id", VISION_START_TOKEN_ID)
-    vision_start_pos = (input_ids == int(vision_start_token_id)).nonzero(as_tuple=True)[0]
-    if vision_start_pos.numel() == 0:
+def _find_speaker_placeholder_positions(input_ids: torch.Tensor, cfg: MingTTSConfig) -> list[int]:
+    # The speaker embedding is written at the slot AFTER the placeholder marker:
+    # dense ``<|vision_start|>``+1 (the ``<|vision_pad|>`` slot), moe ``<spk>``+1
+    # (the ``<audioPatch>`` slot). See upstream prepare_input_embed.
+    placeholder_token_id = int(cfg.speaker_placeholder_token_id)
+    placeholder_pos = (input_ids == placeholder_token_id).nonzero(as_tuple=True)[0]
+    if placeholder_pos.numel() == 0:
         return []
 
     slots = []
-    for pos in vision_start_pos:
+    for pos in placeholder_pos:
         slot = int(pos.item()) + 1
         if slot < int(input_ids.shape[0]):
             slots.append(slot)

vllm_omni/model_executor/models/ming_tts/config_ming_tts.py

@@ -42,6 +42,12 @@ from .constants import (
     LLM_HIDDEN_SIZE,
     LLM_VOCAB_SIZE,
     MAX_DECODE_STEPS,
+    MOE_AUDIO_DUMMY_TOKEN_ID,
+    MOE_AUDIO_END_TOKEN_ID,
+    MOE_AUDIO_EOS_TOKEN_ID,
+    MOE_AUDIO_START_TOKEN_ID,
+    MOE_SPK_TOKEN_ID,
+    MOE_TEXT_EOS_TOKEN_ID,
     PATCH_SIZE,
     SAMPLE_RATE,
     STOP_HEAD_MIN_STEPS,
@@ -89,10 +95,126 @@ class MingDenseConfig(PretrainedConfig):
         return self.llm_config
 
 
+class BailingMoeConfig(PretrainedConfig):
+    """Stage-0 LLM config for the Ming MoE variant.
+
+    Ported from upstream inclusionAI ``configuration_bailing_moe.py``; the
+    community vLLM ``BailingMoeModel`` is duck-typed over these attributes.
+    Only the fields the backbone reads are kept; unknown keys flow through
+    ``**kwargs`` to ``PretrainedConfig``.
+    """
+
+    model_type = "bailing_moe"
+
+    def __init__(
+        self,
+        vocab_size: int = 126464,
+        hidden_size: int = 2048,
+        intermediate_size: int | None = None,
+        num_hidden_layers: int = 28,
+        num_attention_heads: int = 16,
+        num_key_value_heads: int = 0,
+        hidden_act: str = "silu",
+        use_qkv_bias: bool = False,
+        use_bias: bool = True,
+        rms_norm_eps: float = 1e-05,
+        norm_head: bool = False,
+        tie_word_embeddings: bool = False,
+        max_position_embeddings: int = 32768,
+        rope_theta: float = 10000.0,
+        use_cache: bool = True,
+        rope_scaling: Any = None,
+        pad_token_id: int = 126081,
+        num_experts: int = 64,
+        num_shared_experts: int = 2,
+        num_experts_per_tok: int = 6,
+        norm_topk_prob: bool = True,
+        moe_intermediate_size: int | None = None,
+        first_k_dense_replace: int = 0,
+        head_dim: int | None = None,
+        output_router_logits: bool = False,
+        multi_gate: bool = False,
+        image_patch_token: int = 126346,
+        use_grouped_gemm: bool = False,
+        **kwargs: Any,
+    ) -> None:
+        self.vocab_size = vocab_size
+        self.hidden_size = hidden_size
+        self.intermediate_size = intermediate_size
+        self.num_hidden_layers = num_hidden_layers
+        self.num_attention_heads = num_attention_heads
+        self.num_key_value_heads = num_key_value_heads
+        self.hidden_act = hidden_act
+        self.use_qkv_bias = use_qkv_bias
+        self.use_bias = use_bias
+        self.norm_head = norm_head
+        self.rms_norm_eps = rms_norm_eps
+        self.max_position_embeddings = max_position_embeddings
+        self.rope_theta = rope_theta
+        self.use_cache = use_cache
+        self.head_dim = head_dim or self.hidden_size // self.num_attention_heads
+        self.rope_scaling = rope_scaling
+        self.num_experts = num_experts
+        self.num_shared_experts = num_shared_experts
+        self.num_experts_per_tok = num_experts_per_tok
+        self.norm_topk_prob = norm_topk_prob
+        self.moe_intermediate_size = moe_intermediate_size
+        self.first_k_dense_replace = first_k_dense_replace
+        self.output_router_logits = output_router_logits
+        self.multi_gate = multi_gate
+        self.image_patch_token = image_patch_token
+        self.use_grouped_gemm = use_grouped_gemm
+        super().__init__(pad_token_id=pad_token_id, tie_word_embeddings=tie_word_embeddings, **kwargs)
+
+
+def _coerce_bailing_moe_config(value: Any) -> BailingMoeConfig:
+    if isinstance(value, BailingMoeConfig):
+        cfg = value
+    elif isinstance(value, PretrainedConfig):
+        cfg = BailingMoeConfig.from_dict(value.to_dict())
+    elif isinstance(value, dict):
+        cfg = BailingMoeConfig.from_dict(dict(value))
+    else:
+        raise TypeError(f"Unsupported llm_config type for Ming MoE config: {type(value)!r}")
+    # The multimodal bailing_moe declares 3D mrope (rope_scaling type "3D").
+    # Ming-TTS Stage-0 carries no vision/video tokens, so mrope degenerates to
+    # standard 1D RoPE (all three position dims are equal). Strip rope_scaling
+    # so vLLM's get_rope uses plain RoPE — it has no "3D" scaling type.
+    cfg.rope_scaling = None
+    cfg.rope_parameters = {"rope_type": "default", "rope_theta": cfg.rope_theta}
+    return cfg
+
+
+class MingMoeConfig(PretrainedConfig):
+    # Upstream inclusionAI Ming-omni-tts-16.8B-A3B HF config reports model_type="bailingmm".
+    model_type = "bailingmm"
+
+    def __init__(
+        self,
+        llm_config: BailingMoeConfig | dict[str, Any] | None = None,
+        ditar_config: dict[str, Any] | None = None,
+        aggregator_config: dict[str, Any] | None = None,
+        audio_tokenizer_config: AudioVAEConfig | dict[str, Any] | None = None,
+        architectures: list[str] | None = None,
+        **kwargs: Any,
+    ) -> None:
+        self.llm_config = _coerce_bailing_moe_config(llm_config or {})
+        self.ditar_config = dict(ditar_config or {})
+        self.aggregator_config = dict(aggregator_config or {})
+        self.audio_tokenizer_config = _coerce_audio_vae_config(audio_tokenizer_config)
+        super().__init__(architectures=architectures, **kwargs)
+
+    def get_text_config(self, decoder: bool = False, **kwargs: Any) -> BailingMoeConfig:
+        del decoder, kwargs
+        return self.llm_config
+
+
 @dataclass
 class MingTTSConfig:
     """Flat config object shared by Stage-1 and Stage-2. Build via from_hf_config()."""
 
+    model_variant: str = "dense"  # "dense" (qwen2 backbone) | "moe" (bailing_moe backbone)
+
     llm_hidden_size: int = LLM_HIDDEN_SIZE
     llm_vocab_size: int = LLM_VOCAB_SIZE
     llm_config: dict[str, Any] = field(default_factory=dict)
@@ -124,6 +246,7 @@ class MingTTSConfig:
     audio_start_token_id: int = AUDIO_START_TOKEN_ID
     audio_end_token_id: int = AUDIO_END_TOKEN_ID
     audio_eos_token_id: int = AUDIO_EOS_TOKEN_ID
+    speaker_placeholder_token_id: int = VISION_START_TOKEN_ID  # dense <|vision_start|>; moe <spk>
 
     @classmethod
     def from_hf_config(cls, hf_config: PretrainedConfig) -> MingTTSConfig:
@@ -144,7 +267,12 @@ class MingTTSConfig:
 
         enc_hop_size = _nested_get(atc, "enc_kwargs", "hop_size", default=AUDIO_FRAME_HOP)
 
+        # Variant is decided by the Stage-0 LLM backbone family, not the wrapper
+        # model_type: dense -> qwen2, MoE -> bailing_moe.
+        model_variant = "moe" if str(llm_dict.get("model_type", "")) == "bailing_moe" else "dense"
+
         cfg = cls(
+            model_variant=model_variant,
             llm_hidden_size=llm_dict.get("hidden_size", LLM_HIDDEN_SIZE),
             llm_vocab_size=llm_dict.get("vocab_size", LLM_VOCAB_SIZE),
             llm_config=llm_dict,
@@ -158,6 +286,15 @@ class MingTTSConfig:
             sample_rate=atc_sample_rate,
             audio_frame_hop=enc_hop_size if enc_hop_size is not None else AUDIO_FRAME_HOP,
         )
+        if model_variant == "moe":
+            # The bailing tokenizer uses a different vocab; override the dense
+            # (Qwen2) special-token defaults with the bailing token IDs.
+            cfg.text_eos_token_id = MOE_TEXT_EOS_TOKEN_ID
+            cfg.audio_dummy_token_id = MOE_AUDIO_DUMMY_TOKEN_ID
+            cfg.audio_start_token_id = MOE_AUDIO_START_TOKEN_ID
+            cfg.audio_end_token_id = MOE_AUDIO_END_TOKEN_ID
+            cfg.audio_eos_token_id = MOE_AUDIO_EOS_TOKEN_ID
+            cfg.speaker_placeholder_token_id = MOE_SPK_TOKEN_ID
         return cfg
 
     def validate(self) -> None:
@@ -195,7 +332,9 @@ __all__ = [
     "LLM_HIDDEN_SIZE",
     "LLM_VOCAB_SIZE",
     "MAX_DECODE_STEPS",
+    "BailingMoeConfig",
     "MingDenseConfig",
+    "MingMoeConfig",
     "MingTTSConfig",
     "PATCH_SIZE",
     "SAMPLE_RATE",

vllm_omni/model_executor/models/ming_tts/constants.py

@@ -15,6 +15,22 @@ VISION_START_TOKEN_ID = 151652  # <|vision_start|>
 TEXT_EOS_TOKEN_ID = 151669  # <text_eos>
 
 
+# ---------------------------------------------------------------------------
+# MoE (bailing_moe / Ming-omni-tts-16.8B-A3B) token IDs
+# (confirmed from the bailing tokenizer_config.json — different vocab than the
+# dense Qwen2 tokenizer above). The bailing tokenizer has NO <text_eos>; the
+# audio AR loop terminates purely on the stop-head (see upstream
+# modeling_bailingmm.sample), so <end_of_audio> doubles as the AR stop token.
+# ---------------------------------------------------------------------------
+
+MOE_AUDIO_DUMMY_TOKEN_ID = 126357  # <audioPatch>
+MOE_AUDIO_START_TOKEN_ID = 126358  # <audio>
+MOE_AUDIO_END_TOKEN_ID = 126359  # </audio>
+MOE_AUDIO_EOS_TOKEN_ID = 126356  # <end_of_audio>
+MOE_TEXT_EOS_TOKEN_ID = 126356  # no <text_eos> in bailing; reuse <end_of_audio> as AR stop
+MOE_SPK_TOKEN_ID = 126368  # <spk> — speaker-embedding placeholder (dense uses <|vision_start|>)
+
+
 # ---------------------------------------------------------------------------
 # Architectural constants (confirmed from original config.json)
 # ---------------------------------------------------------------------------

vllm_omni/model_executor/models/ming_tts/fm/cfm_cudagraph.py

@@ -0,0 +1,154 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""CUDAGraph-accelerated CFM (flow-matching) diffusion head for Ming-TTS.
+
+The per-AR-step flow head runs ``steps`` ODE iterations, each a tiny DiT forward
+(history+patch tokens, CFG cond+null batched). In eager mode these are dozens of
+micro kernels per step and the wall time is dominated by kernel-launch / Python
+overhead, not GPU math. Capturing the whole sampling loop + Aggregator + stop
+head into one CUDA graph removes that overhead.
+
+Adapted from adaspeech/ming_flash ``CFMGraphExecutor``. Captured for batch=1
+(Ming Stage-0 runs ``max_num_seqs: 1``); other batch sizes fall back to eager.
+"""
+
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+
+from vllm_omni.model_executor.models.common.ming.fm import build_timesteps
+
+
+class CFMSampler(nn.Module):
+    """Unrolled CFM ODE solver, graph-capturable (no ``torch.randn``/``.item()``).
+
+    All randomness (``y0``, ``sde_rnd``) and the timestep schedule ``t`` are
+    passed in as float32 tensors so the loop is deterministic during replay and
+    the integration matches the eager float32 path. Equivalent to eager
+    ``Solver.integrate`` for ``temperature==0`` (the TTS default).
+    """
+
+    def __init__(self, dit_model: nn.Module, steps: int = 10):
+        super().__init__()
+        self.model = dit_model
+        self.steps = steps
+
+    @torch.no_grad()
+    def forward(self, llm_cond, lat_cond, y0, t, sde_args, sde_rnd):
+        # y0 / t / sde_rnd / sde_args are float32; integrate in float32 and cast
+        # only the DiT inputs to the model dtype, mirroring eager FlowLoss.sample.
+        model_dtype = llm_cond.dtype
+        for step in range(self.steps):
+            dt = t[step + 1] - t[step]
+            pred_cfg = self.model.forward_with_cfg(
+                x=y0.to(model_dtype),
+                t=t[step].to(model_dtype),
+                c=llm_cond,
+                latent_history=lat_cond,
+                cfg_scale=sde_args[0],
+                patch_size=y0.shape[1],
+            )
+            pred, null_pred = torch.chunk(pred_cfg, 2, dim=0)
+            flow = (pred + (pred - null_pred) * sde_args[0]).float()
+            y0 = y0 + flow * dt
+            if step < self.steps - 1:
+                y0 = y0 + sde_args[1] * (sde_args[2] ** 0.5) * (dt.abs() ** 0.5) * sde_rnd[step]
+        return y0
+
+
+class CFMGraphExecutor:
+    """Capture CFM sampling + Aggregator + stop head into one CUDA graph.
+
+    First ``execute`` captures the graph; later calls copy inputs into the
+    static placeholders and ``replay()``. Runtime controls (cfg/sigma/temp) are
+    fed via the ``sde_args`` placeholder, so they are tunable without recapture.
+    """
+
+    def __init__(
+        self,
+        cfm_sampler: CFMSampler,
+        aggregator: nn.Module,
+        stop_head: nn.Module,
+        *,
+        patch_size: int,
+        latent_dim: int,
+        steps: int = 10,
+    ):
+        self.cfm_sampler = cfm_sampler
+        self.aggregator = aggregator
+        self.stop_head = stop_head
+        self.patch_size = int(patch_size)
+        self.latent_dim = int(latent_dim)
+        self.steps = int(steps)
+        self.initialized = False
+        self.graph: torch.cuda.CUDAGraph | None = None
+
+    def _timesteps(self, device):
+        # Raw EPSS schedule + sway in float32, identical to eager build_timesteps.
+        return build_timesteps(self.steps, device=device, dtype=torch.float32, use_epss=True, sway_sampling_coef=-1.0)
+
+    @torch.no_grad()
+    def execute(self, llm_cond, his_lat, cfg_strength, sigma, temperature):
+        bsz = llm_cond.shape[0]
+        device = llm_cond.device
+        # ODE state / noise / timesteps are float32 to match the eager path.
+        y0 = torch.randn((bsz, self.patch_size, self.latent_dim), device=device, dtype=torch.float32)
+        t = self._timesteps(device)
+        sde_rnd = torch.randn((self.steps, bsz, self.patch_size, self.latent_dim), device=device, dtype=torch.float32)
+
+        if not self.initialized:
+            self._initialize_graph(llm_cond, his_lat, y0, t, sde_rnd)
+
+        self._llm_cond_ph.copy_(llm_cond)
+        self._his_lat_ph.copy_(his_lat)
+        self._y0_ph.copy_(y0)
+        self._t_ph.copy_(t)
+        self._sde_args_ph[0] = float(cfg_strength)
+        self._sde_args_ph[1] = float(sigma)
+        self._sde_args_ph[2] = float(temperature)
+        self._sde_rnd_ph.copy_(sde_rnd)
+
+        self.graph.replay()
+
+        return self._gen_lat_ph.clone(), self._embeds_ph.clone(), self._stop_ph.clone()
+
+    def _initialize_graph(self, llm_cond, his_lat, y0, t, sde_rnd):
+        model_dtype = llm_cond.dtype
+        self._llm_cond_ph = torch.empty_like(llm_cond)
+        self._his_lat_ph = torch.empty_like(his_lat)
+        self._y0_ph = torch.empty_like(y0)  # float32
+        self._t_ph = torch.empty_like(t)  # float32
+        self._sde_args_ph = torch.empty(3, device=llm_cond.device, dtype=torch.float32)
+        self._sde_rnd_ph = torch.empty_like(sde_rnd)  # float32
+
+        self._llm_cond_ph.copy_(llm_cond)
+        self._his_lat_ph.copy_(his_lat)
+        self._y0_ph.copy_(y0)
+        self._t_ph.copy_(t)
+        self._sde_args_ph[0], self._sde_args_ph[1], self._sde_args_ph[2] = 2.0, 0.25, 0.0
+        self._sde_rnd_ph.copy_(sde_rnd)
+
+        def _run():
+            # Solver runs in float32; cast the final latent to the model dtype
+            # for the Aggregator / stop head (matches eager ``.to(dtype)``).
+            gen_lat = self.cfm_sampler(
+                self._llm_cond_ph, self._his_lat_ph, self._y0_ph, self._t_ph, self._sde_args_ph, self._sde_rnd_ph
+            ).to(model_dtype)
+            embeds = self.aggregator(gen_lat)
+            stop = self.stop_head(self._llm_cond_ph[:, -1, :]).softmax(dim=-1)
+            return gen_lat, embeds, stop
+
+        # Warm up on a side stream to stabilize allocations before capture.
+        warmup = torch.cuda.Stream()
+        warmup.wait_stream(torch.cuda.current_stream())
+        with torch.cuda.stream(warmup):
+            for _ in range(3):
+                _run()
+        torch.cuda.current_stream().wait_stream(warmup)
+
+        self.graph = torch.cuda.CUDAGraph()
+        with torch.cuda.graph(self.graph):
+            self._gen_lat_ph, self._embeds_ph, self._stop_ph = _run()
+
+        self.initialized = True

vllm_omni/model_executor/models/ming_tts/fm/dit.py

@@ -87,12 +87,17 @@ class DiT(nn.Module):
         return x
 
     def forward_with_cfg(self, x, t, c, cfg_scale, latent_history, patch_size):
-        if not cfg_scale == 1:
-            x = torch.cat([x, x], dim=0)
-            latent_history = torch.cat([latent_history, latent_history], dim=0)
-            fake_latent = torch.zeros_like(c)
-            c = torch.cat([c, fake_latent], dim=0)
+        # Always run the classifier-free-guidance path (cond + uncond). TTS
+        # inference always uses cfg>1, and branching on ``cfg_scale`` is avoided
+        # because under CUDA-graph capture it is a CUDA tensor — ``bool(cfg_scale)``
+        # would force a device sync and break capture. ``expand`` (not ``repeat``)
+        # avoids a fresh allocation each step.
+        x = torch.cat([x, x], dim=0)
+        latent_history = torch.cat([latent_history, latent_history], dim=0)
+        fake_latent = torch.zeros_like(c)
+        c = torch.cat([c, fake_latent], dim=0)
         if t.ndim == 0:
-            t = t.repeat(x.shape[0])
+            t = t.reshape(1)
+        t = t.expand(x.shape[0])
         model_out = self.forward(x, t, c, latent_history)
         return model_out[:, -patch_size:, :]

vllm_omni/model_executor/models/ming_tts/ming_tts.py

@@ -182,7 +182,7 @@ class MingTTSForConditionalGeneration(nn.Module, SupportsPP, CustomProcessMixin)
                 use_zero_spk_emb=bool(info_dict.get("use_zero_spk_emb", False)),
             )
         if speaker_embeddings is not None and len(speaker_embeddings) > 0:
-            speaker_slots = _find_speaker_placeholder_positions(input_ids, self.vllm_config.model_config.hf_config)
+            speaker_slots = _find_speaker_placeholder_positions(input_ids, self.ming_config)
             if len(speaker_slots) < len(speaker_embeddings):
                 raise RuntimeError(
                     "Could not locate enough speaker placeholder slots: "

vllm_omni/model_executor/models/ming_tts/ming_tts_llm.py

@@ -3,6 +3,7 @@
 # Adopted from https://github.com/inclusionAI/Ming-omni-tts/blob/main/modeling_bailingmm.py
 from __future__ import annotations
 
+import os
 from collections.abc import Iterable
 from typing import Any
 
@@ -51,6 +52,9 @@ from .patch_emission import (
 
 logger = init_logger(__name__)
 
+# CFM ODE integration steps; must match CFM.sample(steps=...) default (10).
+_CFM_STEPS = 10
+
 
 class MingLLMModel(nn.Module):
     hf_to_vllm_mapper = WeightsMapper(
@@ -66,7 +70,16 @@ class MingLLMModel(nn.Module):
         self.vllm_config = vllm_config
         self.prefix = prefix
         self.fm_dtype = _resolve_ming_runtime_dtype(vllm_config)
-        self.model = Qwen2Model(vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model"))
+        if self.ming_config.model_variant == "moe":
+            from vllm.model_executor.models.bailing_moe import BailingMoeModel
+
+            # BailingMoeModel reads ``vllm_config.model_config.hf_config`` directly
+            # (no get_text_config()), so re-wrap with the nested bailing_moe config.
+            llm_config = vllm_config.model_config.hf_config.get_text_config()
+            llm_vllm_config = vllm_config.with_hf_config(llm_config)
+            self.model = BailingMoeModel(vllm_config=llm_vllm_config, prefix=maybe_prefix(prefix, "model"))
+        else:
+            self.model = Qwen2Model(vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model"))
         self.linear_proj_audio = Aggregator(
             in_channels=self.ming_config.latent_dim,
             llm_input_dim=self.ming_config.llm_hidden_size,
@@ -86,6 +99,11 @@ class MingLLMModel(nn.Module):
         self._pending_postprocess_updates: dict[str, dict[str, Any]] = {}
         self._last_ming_next_token_ids = None
         self._last_text_mode = False
+        # CUDAGraph for the flow-matching (CFM) diffusion head — the TTFP/latency
+        # hot spot. Built lazily on the first batch=1 CUDA decode; set
+        # MING_CFM_CUDAGRAPH=0 to force the eager flow head.
+        self._cfm_graph = None
+        self._cfm_graph_enabled = os.environ.get("MING_CFM_CUDAGRAPH", "1") == "1"
 
     def embed_input_ids(
         self, input_ids: torch.Tensor, inputs_embeds: torch.Tensor | None = None, **_: Any
@@ -118,11 +136,14 @@ class MingLLMModel(nn.Module):
             min_stop_step=int(self.ming_config.stop_head_min_steps),
             default_max_decode_steps=self.ming_config.max_decode_steps,
         )
+        # Positional call: Qwen2Model names the 2nd arg ``positions`` while
+        # BailingMoeModel names it ``position_ids`` — both share the signature
+        # (input_ids, positions, intermediate_tensors, inputs_embeds).
         backbone_out = self.model(
-            input_ids=input_ids,
-            positions=positions,
-            intermediate_tensors=intermediate_tensors,
-            inputs_embeds=inputs_embeds,
+            input_ids,
+            positions,
+            intermediate_tensors,
+            inputs_embeds,
         )
         if isinstance(backbone_out, IntermediateTensors):
             self._last_ming_next_token_ids = None
@@ -294,6 +315,33 @@ class MingLLMModel(nn.Module):
             logprobs_tensors=None,
         )
 
+    def _maybe_build_cfm_graph(self, z_diff_cond: torch.Tensor):
+        """Lazily build the CFM CUDAGraph executor (batch=1, CUDA only)."""
+        if not self._cfm_graph_enabled:
+            return None
+        if z_diff_cond.shape[0] != 1 or not z_diff_cond.is_cuda:
+            return None
+        if self._cfm_graph is not None:
+            return self._cfm_graph
+        try:
+            from .fm.cfm_cudagraph import CFMGraphExecutor, CFMSampler
+
+            sampler = CFMSampler(self.flowloss.cfm.model, steps=_CFM_STEPS)
+            self._cfm_graph = CFMGraphExecutor(
+                sampler,
+                self.linear_proj_audio,
+                self.stop_head,
+                patch_size=self.ming_config.patch_size,
+                latent_dim=self.ming_config.latent_dim,
+                steps=_CFM_STEPS,
+            )
+            logger.info("Ming CFM CUDAGraph enabled (steps=%d, patch=%d).", _CFM_STEPS, self.ming_config.patch_size)
+        except Exception as exc:
+            logger.warning("Ming CFM CUDAGraph unavailable (%s); using eager flow head.", exc)
+            self._cfm_graph_enabled = False
+            return None
+        return self._cfm_graph
+
     def _decode_one_step(
         self,
         *,
@@ -317,23 +365,45 @@ class MingLLMModel(nn.Module):
         z_diff_cond = hidden_states.to(dtype=self.fm_dtype).unsqueeze(1)
         if not torch.isfinite(z_diff_cond).all():
             raise RuntimeError("Non-finite z_diff_cond before FlowLoss.sample().")
-        sampled_token_latent = self.flowloss.sample(
-            z=z_diff_cond,
-            latent_history=latent_history,
-            cfg=cfg_scale,
-            patch_size=self.ming_config.patch_size,
-            sigma=sigma,
-            temperature=temperature,
-        )
+
+        # Fast path: CUDAGraph-captured flow head (CFM sampling + Aggregator +
+        # stop head). Falls back permanently to the eager path on any failure.
+        sampled_token_latent = next_embeds = stop_probs = None
+        # cfg < 1e-5 disables CFG; the eager flow head handles this with a
+        # dedicated unconditional branch (c=zeros) that the captured graph does
+        # not replicate, so fall back to eager there to match FlowLoss.sample.
+        graph_exec = self._maybe_build_cfm_graph(z_diff_cond) if cfg_scale >= 1e-5 else None
+        if graph_exec is not None:
+            try:
+                sampled_token_latent, next_embeds, stop_full = graph_exec.execute(
+                    z_diff_cond, latent_history, cfg_scale, sigma, temperature
+                )
+                stop_probs = stop_full[:, 1]
+            except Exception as exc:
+                logger.warning("Ming CFM CUDAGraph failed (%s); falling back to eager flow head.", exc)
+                self._cfm_graph = None
+                self._cfm_graph_enabled = False
+                sampled_token_latent = next_embeds = stop_probs = None
+
+        if sampled_token_latent is None:
+            sampled_token_latent = self.flowloss.sample(
+                z=z_diff_cond,
+                latent_history=latent_history,
+                cfg=cfg_scale,
+                patch_size=self.ming_config.patch_size,
+                sigma=sigma,
+                temperature=temperature,
+            )
+            # Aggregator expects [Batch, Time, Dimension] = [B, 4, 64] and returns [B, 1, H].
+            next_embeds = self.linear_proj_audio(sampled_token_latent)
+            stop_probs = self.stop_head(hidden_states.to(dtype=self.fm_dtype)).softmax(dim=-1)[:, 1]
+
         expected_shape = (hidden_states.shape[0], self.ming_config.patch_size, self.ming_config.latent_dim)
         if tuple(sampled_token_latent.shape) != expected_shape:
             raise RuntimeError(
                 f"FlowLoss output shape mismatch: got {tuple(sampled_token_latent.shape)}, expected {expected_shape}"
             )
         new_history = torch.cat([latent_history[:, self.ming_config.patch_size :, :], sampled_token_latent], dim=1)
-        # Aggregator expects [Batch, Time, Dimension] = [B, 4, 64] and returns [B, 1, H].
-        next_embeds = self.linear_proj_audio(sampled_token_latent)
-        stop_probs = self.stop_head(hidden_states.to(dtype=self.fm_dtype)).softmax(dim=-1)[:, 1]
         if not torch.isfinite(sampled_token_latent).all():
             raise RuntimeError("Non-finite sampled_token_latent in Ming decode step.")
         if not torch.isfinite(next_embeds).all():

vllm_omni/model_executor/models/ming_tts/pipeline.py

@@ -13,7 +13,10 @@ _PROC = "vllm_omni.model_executor.stage_input_processors.ming_tts"
 MING_TTS_PIPELINE = PipelineConfig(
     model_type="ming_tts",
     model_arch="MingTTSForConditionalGeneration",
-    hf_architectures=("MingTTSForConditionalGeneration",),
+    hf_architectures=("MingTTSForConditionalGeneration", "BailingMMNativeForConditionalGeneration"),
+    # The dense (0.5B) and MoE (16.8B) share architectures=["BailingMMNativeForConditionalGeneration"]
+    # Here we disambiguate by the model_type the upstream HF config reports
+    hf_config_predicate=lambda c: str(getattr(c, "model_type", "")) == "dense",
     stages=(
         StagePipelineConfig(
             stage_id=0,
@@ -52,3 +55,17 @@ MING_TTS_PIPELINE = PipelineConfig(
         ),
     ),
 )
+
+
+# MoE variant (inclusionAI/Ming-omni-tts-16.8B-A3B)
+# Same two-stage topology and model class as the above,
+# but with a distinct model_type reported by the upstream HF config.
+# Keep this for auto-detection of deploy config yaml.
+MING_TTS_MOE_PIPELINE = PipelineConfig(
+    model_type="ming_tts_moe",
+    model_arch="MingTTSForConditionalGeneration",
+    hf_architectures=("MingTTSForConditionalGeneration", "BailingMMNativeForConditionalGeneration"),
+    # Disambiguate from the dense-0.5B ming_tts pipeline
+    hf_config_predicate=lambda c: str(getattr(c, "model_type", "")) == "bailingmm",
+    stages=MING_TTS_PIPELINE.stages,
+)

vllm_omni/model_executor/models/ming_tts/prompt_assembly.py

@@ -176,6 +176,84 @@ def build_dense_prompt_token_ids(
     )
 
 
+def build_moe_prompt_token_ids(
+    tokenizer: Any,
+    *,
+    prompt: str,
+    text: str,
+    instruction: str | None = None,
+    prompt_text: str | None = None,
+    speaker_count: int = 0,
+    prompt_patch_count: int = 0,
+) -> list[int]:
+    """Stage-0 prompt for the Ming MoE (bailing_moe) variant.
+
+    Mirrors the ``model_type != 'dense'`` branch of upstream
+    ``BailingMMNativeForConditionalGeneration.prepare_input_embed``: no system
+    message, ``<role>HUMAN/ASSISTANT</role>`` turns, and a ``<spk><audioPatch></spk>``
+    speaker placeholder (vs. the dense ``<|vision_*|>`` placeholder).
+    """
+    speaker_prompt = []
+    for idx in range(int(speaker_count)):
+        speaker_prompt.extend(
+            tokenizer.encode(f"  speaker_{idx + 1}:")
+            + tokenizer.encode("<spk>")
+            + tokenizer.encode("<audioPatch>")
+            + tokenizer.encode("</spk>\n")
+        )
+    instruction_prompt = (
+        tokenizer.encode(instruction) + tokenizer.encode("<|endoftext|>") if instruction is not None else []
+    )
+    prompt_text_tokens = (
+        tokenizer.encode(prompt_text) if int(prompt_patch_count) > 0 and prompt_text is not None else []
+    )
+    audio_patch_token_id = tokenizer.convert_tokens_to_ids("<audioPatch>")
+    if audio_patch_token_id == tokenizer.unk_token_id:
+        raise ValueError("Ming tokenizer is missing required <audioPatch> token.")
+    prompt_latent_tokens = [audio_patch_token_id] * int(prompt_patch_count)
+    text_input_prefix = (
+        []
+        if all(token in text for token in ("Genre: ", "Mood: ", "Instrument: ", "Theme: ", "Duration: "))
+        else tokenizer.encode(" Text input:\n")
+    )
+    return (
+        tokenizer.encode("<role>HUMAN</role>")
+        + tokenizer.encode(prompt)
+        + speaker_prompt
+        + text_input_prefix
+        + prompt_text_tokens
+        + tokenizer.encode(text)
+        + tokenizer.encode("<role>ASSISTANT</role>")
+        + instruction_prompt
+        + tokenizer.encode("<audio>")
+        + prompt_latent_tokens
+    )
+
+
+def build_prompt_token_ids(model_variant: str, tokenizer: Any, **kwargs: Any) -> list[int]:
+    """Dispatch to the dense or MoE Stage-0 prompt builder by model variant."""
+    if model_variant == "moe":
+        return build_moe_prompt_token_ids(tokenizer, **kwargs)
+    return build_dense_prompt_token_ids(tokenizer, **kwargs)
+
+
+def detect_model_variant(tokenizer: Any) -> str:
+    """Infer the Ming variant from the tokenizer's special tokens.
+
+    The bailing (MoE) tokenizer carries ``<role>`` chat markers; the dense
+    Qwen2 tokenizer does not. Used to pick the prompt format when the caller
+    does not pass an explicit variant.
+    """
+    try:
+        role_id = tokenizer.convert_tokens_to_ids("<role>")
+    except Exception:
+        return "dense"
+    unk = getattr(tokenizer, "unk_token_id", None)
+    if role_id is None or (unk is not None and role_id == unk):
+        return "dense"
+    return "moe"
+
+
 def build_ming_dense_prompt(
     tokenizer: Any,
     *,
@@ -189,7 +267,10 @@ def build_ming_dense_prompt(
     speaker_embedding: Any = None,
     use_zero_spk_emb: bool = False,
     request_id: str | None = None,
+    model_variant: str | None = None,
 ) -> dict[str, Any]:
+    if model_variant is None:
+        model_variant = detect_model_variant(tokenizer)
     instruction_text = create_instruction(instruction)
     speaker_embeddings = coerce_speaker_embeddings(speaker_embedding, use_zero_spk_emb=use_zero_spk_emb)
     effective_runtime_controls = resolve_effective_runtime_controls(text=text, runtime_controls=runtime_controls)
@@ -217,7 +298,8 @@ def build_ming_dense_prompt(
             prompt_latent_value, patch_size=PATCH_SIZE, latent_dim=LATENT_DIM
         )
 
-    prompt_token_ids = build_dense_prompt_token_ids(
+    prompt_token_ids = build_prompt_token_ids(
+        model_variant,
         tokenizer,
         prompt=prompt,
         text=text,
@@ -266,7 +348,10 @@ __all__ = [
     "DEFAULT_PROMPT",
     "build_dense_prompt_token_ids",
     "build_ming_dense_prompt",
+    "build_moe_prompt_token_ids",
+    "build_prompt_token_ids",
     "create_instruction",
+    "detect_model_variant",
     "estimate_decode_step_window_for_duration",
     "estimate_decode_steps_for_duration",
     "parse_duration_seconds",

vllm_omni/model_executor/models/ming_tts/validation.py

@@ -14,6 +14,9 @@ from .constants import (
     LATENT_DIM,
     LLM_HIDDEN_SIZE,
     LLM_VOCAB_SIZE,
+    MOE_AUDIO_DUMMY_TOKEN_ID,
+    MOE_AUDIO_EOS_TOKEN_ID,
+    MOE_TEXT_EOS_TOKEN_ID,
     PATCH_SIZE,
     SAMPLE_RATE,
 )
@@ -70,19 +73,23 @@ def _nested_get(obj: Any, *keys: str, default: Any = None) -> Any:
 
 def validate_ming_tts_config(cfg: Any) -> None:
     """Run before GPU allocation/weight loading. Raises ValueError on mismatches."""
-    if cfg.audio_dummy_token_id != 151705:
+    is_moe = getattr(cfg, "model_variant", "dense") == "moe"
+    exp_audio_dummy = MOE_AUDIO_DUMMY_TOKEN_ID if is_moe else 151705
+    exp_audio_eos = MOE_AUDIO_EOS_TOKEN_ID if is_moe else 151704
+    exp_text_eos = MOE_TEXT_EOS_TOKEN_ID if is_moe else 151669
+    if cfg.audio_dummy_token_id != exp_audio_dummy:
         raise ValueError(
-            f"audio_dummy_token_id={cfg.audio_dummy_token_id}, expected 151705 (<audioPatch>). "
+            f"audio_dummy_token_id={cfg.audio_dummy_token_id}, expected {exp_audio_dummy} (<audioPatch>). "
             "Wrong tokenizer/checkpoint?"
         )
-    if cfg.audio_eos_token_id != 151704:
+    if cfg.audio_eos_token_id != exp_audio_eos:
         raise ValueError(
-            f"audio_eos_token_id={cfg.audio_eos_token_id}, expected 151704 (<end_of_audio>). "
+            f"audio_eos_token_id={cfg.audio_eos_token_id}, expected {exp_audio_eos} (<end_of_audio>). "
             "Wrong tokenizer/checkpoint?"
         )
-    if cfg.text_eos_token_id != 151669:
+    if cfg.text_eos_token_id != exp_text_eos:
         raise ValueError(
-            f"text_eos_token_id={cfg.text_eos_token_id}, expected 151669 (<text_eos>). Wrong tokenizer/checkpoint?"
+            f"text_eos_token_id={cfg.text_eos_token_id}, expected {exp_text_eos}. Wrong tokenizer/checkpoint?"
         )
 
     if cfg.audio_tokenizer_config is None:
@@ -102,13 +109,16 @@ def validate_ming_tts_config(cfg: Any) -> None:
             f"history_patch_size mismatch: got {cfg.history_patch_size}, expected {HISTORY_PATCH_SIZE}. "
             "Check ditar_config.history_patch_size."
         )
-    if cfg.llm_hidden_size != LLM_HIDDEN_SIZE:
-        raise ValueError(
-            f"llm_hidden_size mismatch: got {cfg.llm_hidden_size}, expected {LLM_HIDDEN_SIZE}. "
-            "Check llm_config.hidden_size."
-        )
-    if cfg.llm_vocab_size != LLM_VOCAB_SIZE:
-        raise ValueError(f"llm_vocab_size mismatch: got {cfg.llm_vocab_size}, expected {LLM_VOCAB_SIZE}.")
+    # Absolute hidden/vocab pins are dense-only; the MoE backbone carries its
+    # own (larger) dims, cross-checked against llm_config.hidden_size below.
+    if not is_moe:
+        if cfg.llm_hidden_size != LLM_HIDDEN_SIZE:
+            raise ValueError(
+                f"llm_hidden_size mismatch: got {cfg.llm_hidden_size}, expected {LLM_HIDDEN_SIZE}. "
+                "Check llm_config.hidden_size."
+            )
+        if cfg.llm_vocab_size != LLM_VOCAB_SIZE:
+            raise ValueError(f"llm_vocab_size mismatch: got {cfg.llm_vocab_size}, expected {LLM_VOCAB_SIZE}.")
     if cfg.sample_rate != SAMPLE_RATE:
         raise ValueError(f"sample_rate mismatch: got {cfg.sample_rate}, expected {SAMPLE_RATE}.")