diff --git a/tests/model_executor/models/ming_tts/test_ming_shared_modules.py b/tests/model_executor/models/ming_tts/test_ming_shared_modules.py index 7aab6133b..c5966dd60 100644 --- a/tests/model_executor/models/ming_tts/test_ming_shared_modules.py +++ b/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() diff --git a/vllm_omni/config/pipeline_registry.py b/vllm_omni/config/pipeline_registry.py index 6feded0e9..7bf79d44c 100644 --- a/vllm_omni/config/pipeline_registry.py +++ b/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, diff --git a/vllm_omni/deploy/ming_tts_moe.yaml b/vllm_omni/deploy/ming_tts_moe.yaml new file mode 100644 index 000000000..fe7806b1d --- /dev/null +++ b/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 diff --git a/vllm_omni/engine/arg_utils.py b/vllm_omni/engine/arg_utils.py index d5e21d419..d97fb26e5 100644 --- a/vllm_omni/engine/arg_utils.py +++ b/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), diff --git a/vllm_omni/entrypoints/openai/serving_speech.py b/vllm_omni/entrypoints/openai/serving_speech.py index da1418a09..1a3f13391 100644 --- a/vllm_omni/entrypoints/openai/serving_speech.py +++ b/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. diff --git a/vllm_omni/model_executor/models/common/ming/audio_vae.py b/vllm_omni/model_executor/models/common/ming/audio_vae.py index 6924ae548..89520f515 100644 --- a/vllm_omni/model_executor/models/common/ming/audio_vae.py +++ b/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() diff --git a/vllm_omni/model_executor/models/ming_tts/__init__.py b/vllm_omni/model_executor/models/ming_tts/__init__.py index fb6738a29..da4076940 100644 --- a/vllm_omni/model_executor/models/ming_tts/__init__.py +++ b/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 diff --git a/vllm_omni/model_executor/models/ming_tts/audio_prep.py b/vllm_omni/model_executor/models/ming_tts/audio_prep.py index b63e70cc8..378cd222d 100644 --- a/vllm_omni/model_executor/models/ming_tts/audio_prep.py +++ b/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 ````+1 + # (the ```` 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) diff --git a/vllm_omni/model_executor/models/ming_tts/config_ming_tts.py b/vllm_omni/model_executor/models/ming_tts/config_ming_tts.py index 8ef8a3a0e..a83989a03 100644 --- a/vllm_omni/model_executor/models/ming_tts/config_ming_tts.py +++ b/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 @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", diff --git a/vllm_omni/model_executor/models/ming_tts/constants.py b/vllm_omni/model_executor/models/ming_tts/constants.py index 4b3f527da..a189e4285 100644 --- a/vllm_omni/model_executor/models/ming_tts/constants.py +++ b/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 # +# --------------------------------------------------------------------------- +# 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 ; the +# audio AR loop terminates purely on the stop-head (see upstream +# modeling_bailingmm.sample), so doubles as the AR stop token. +# --------------------------------------------------------------------------- + +MOE_AUDIO_DUMMY_TOKEN_ID = 126357 # +MOE_AUDIO_START_TOKEN_ID = 126358 # +MOE_AUDIO_EOS_TOKEN_ID = 126356 # +MOE_TEXT_EOS_TOKEN_ID = 126356 # no in bailing; reuse as AR stop +MOE_SPK_TOKEN_ID = 126368 # — speaker-embedding placeholder (dense uses <|vision_start|>) + + # --------------------------------------------------------------------------- # Architectural constants (confirmed from original config.json) # --------------------------------------------------------------------------- diff --git a/vllm_omni/model_executor/models/ming_tts/fm/cfm_cudagraph.py b/vllm_omni/model_executor/models/ming_tts/fm/cfm_cudagraph.py new file mode 100644 index 000000000..7a7280034 --- /dev/null +++ b/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 diff --git a/vllm_omni/model_executor/models/ming_tts/fm/dit.py b/vllm_omni/model_executor/models/ming_tts/fm/dit.py index 9aa05d4a2..3eb13e987 100644 --- a/vllm_omni/model_executor/models/ming_tts/fm/dit.py +++ b/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:, :] diff --git a/vllm_omni/model_executor/models/ming_tts/ming_tts.py b/vllm_omni/model_executor/models/ming_tts/ming_tts.py index 4140b403b..70a414186 100644 --- a/vllm_omni/model_executor/models/ming_tts/ming_tts.py +++ b/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: " diff --git a/vllm_omni/model_executor/models/ming_tts/ming_tts_llm.py b/vllm_omni/model_executor/models/ming_tts/ming_tts_llm.py index e7e59631f..ae4d61f68 100644 --- a/vllm_omni/model_executor/models/ming_tts/ming_tts_llm.py +++ b/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(): diff --git a/vllm_omni/model_executor/models/ming_tts/pipeline.py b/vllm_omni/model_executor/models/ming_tts/pipeline.py index 31ec98e5a..762dcf6fd 100644 --- a/vllm_omni/model_executor/models/ming_tts/pipeline.py +++ b/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, +) diff --git a/vllm_omni/model_executor/models/ming_tts/prompt_assembly.py b/vllm_omni/model_executor/models/ming_tts/prompt_assembly.py index c732e07a0..36817022a 100644 --- a/vllm_omni/model_executor/models/ming_tts/prompt_assembly.py +++ b/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, ``HUMAN/ASSISTANT`` turns, and a ```` + 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("") + + tokenizer.encode("") + + tokenizer.encode("\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("") + if audio_patch_token_id == tokenizer.unk_token_id: + raise ValueError("Ming tokenizer is missing required 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("HUMAN") + + tokenizer.encode(prompt) + + speaker_prompt + + text_input_prefix + + prompt_text_tokens + + tokenizer.encode(text) + + tokenizer.encode("ASSISTANT") + + instruction_prompt + + tokenizer.encode("