Skip to content

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

Merged
hsliuustc0106 merged 16 commits into
vllm-project:mainfrom
LHXuuu:ming_tts_moe
Jun 25, 2026
Merged

hsliuustc0106 merged 16 commits into
vllm-project:mainfrom
LHXuuu:ming_tts_moe

Conversation

@LHXuuu

@LHXuuu LHXuuu commented Jun 11, 2026

Copy link
Copy Markdown
Contributor

Purpose

This PR adds support for Ming-omni-tts-16.8B-A3B (the MoE variant) to vLLM-Omni, building on the dense 0.5B pipeline merged in #2906.

Model adaptation

  • bailing_moe backbone selection via vllm_config.with_hf_config(...), with new MingMoeConfig / BailingMoeConfig and config registration for model_type="bailingmm" / "bailing_moe". The 3D mrope rope_scaling is stripped — TTS carries no vision tokens, so it degenerates to standard 1D RoPE.
  • MoE prompt builder (<role> chat format, <spk> speaker placeholder) with tokenizer-based variant dispatch; variant-aware token IDs / validation / speaker slots.
  • deploy/ming_tts_moe.yaml — single-GPU deploy (both stages colocate on one 80GB+ card).

Performance optimization

  • CFM CUDAGraph for the flow-matching diffusion head — the TTFP/RTF hot spot. The per-step CFM sampling (10-step ODE) + Aggregator + stop head is captured into a single CUDA graph, removing the per-step kernel-launch overhead of dozens of tiny DiT forwards. Enabled by default with an automatic eager fallback (MING_CFM_CUDAGRAPH=0 to disable).

Resolves: #1461

Test Plan

Test Result


Essential Elements of an Effective PR Description Checklist
  • The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
  • The test plan. Please provide the test scripts & test commands. Please state the reasons if your codes don't require additional test scripts. For test file guidelines, please check the test style doc
  • The test results. Please paste the results comparison before and after, or the e2e results.
  • (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model. Please run mkdocs serve to sync the documentation editions to ./docs.
  • (Optional) Release notes update. If your change is user-facing, please update the release notes draft.

BEFORE SUBMITTING, please read https://github.com/vllm-project/vllm-omni/blob/main/CONTRIBUTING.md
(anything written below this line will be removed by GitHub Actions)

Adapt the bailing_moe (16.8B-A3B) variant on top of the dense Ming-omni-tts
pipeline:
- MingMoeConfig / BailingMoeConfig + variant detection; register
  bailingmm / bailing_moe configs; strip 3D mrope rope_scaling for the
  no-vision TTS path
- select BailingMoeModel backbone via vllm_config.with_hf_config; positional
  backbone forward call (positions vs position_ids)
- MoE prompt builder (<role> chat, <spk> speaker placeholder) + tokenizer-
  based variant dispatch; variant-aware token ids / validation / speaker slot
- deploy/ming_tts_moe.yaml
- CFM CUDAGraph for the flow-matching head (fm/cfm_cudagraph.py + graph-safe
  forward_with_cfg); ~2.1x TTFP/RTF, lazy build with eager fallback via
  MING_CFM_CUDAGRAPH

Signed-off-by: LHXuuu <xulianhao.xlh@antgroup.com>
@chatgpt-codex-connector

Copy link
Copy Markdown

Codex usage limits have been reached for code reviews. Please check with the admins of this repo to increase the limits by adding credits.
Credits must be used to enable repository wide code reviews.

@hsliuustc0106 hsliuustc0106 left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Graph-focused review comments posted inline.

def execute(self, llm_cond, his_lat, cfg_strength, sigma, temperature):
bsz = llm_cond.shape[0]
device, dtype = llm_cond.device, llm_cond.dtype
y0 = torch.randn((bsz, self.patch_size, self.latent_dim), device=device, dtype=dtype)

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The graph path changes the numeric behavior of CFM sampling here. Eager FlowLoss.sample() creates the initial noise in fp32 and builds fp32 timesteps, then casts only the DiT inputs/final result; this uses llm_cond.dtype, which is bf16 under the new deploy. Because the graph is enabled by default, this can change the latent trajectory/audio quality for both MoE and existing dense Ming TTS. Can we keep the solver state/timesteps/noise/sde args in fp32 and only cast x/t for the model call to match eager?

@yuanheng-zhao yuanheng-zhao Jun 15, 2026 •

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Same concern at here #4341 (comment) @LHXuuu

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Fixed in dbb70be. The CUDAGraph path now matches the eager FlowLoss.sample dtype boundary: the whole ODE solve (y0, timesteps, SDE noise, Euler step) runs in fp32, only the DiT forward is cast to the model dtype (bf16), and the final latent is cast back at the end.

Comment thread vllm_omni/model_executor/models/ming_tts/fm/cfm_cudagraph.py
Comment thread vllm_omni/model_executor/models/ming_tts/fm/cfm_cudagraph.py
Comment thread vllm_omni/model_executor/models/ming_tts/ming_tts_llm.py Outdated
@Sy0307

Sy0307 commented Jun 12, 2026

Copy link
Copy Markdown
Collaborator

The overall MoE adaptation looks directionally right to me: the prompt format / special tokens are aligned with the upstream behavior, treating the TTS-only path as regular 1D RoPE is a reasonable choice, and the CFM CUDAGraph work is a useful optimization direction.

However, I found two startup-path blockers in real inference. The common root cause is that this patch adapts the model forward path, but it does not fully account for what the vLLM engine consumes during initialization: both the normalized config and the tokenizer can be touched before the request ever reaches the model forward.

1. Stripping 3D mRoPE leaves rope_parameters unusable for vLLM's Bailing backbone

The code currently strips the 3D mRoPE config with:

cfg.rope_scaling = None

I agree with the intent here: the TTS path has no vision/video tokens, so the 3D mRoPE path should degenerate to regular 1D RoPE. The problem is that, in the vLLM runtime I tested, the adapted config ends up with:

rope_scaling None
rope_parameters None

But vLLM's BailingAttention initializes RoPE by mutating this dict:

config.rope_parameters["partial_rotary_factor"] = rotary_dim / self.head_dim

So the model can fail during backbone initialization before weight loading. If this PR intentionally downgrades 3D mRoPE to regular 1D RoPE, it should also explicitly rebuild the regular RoPE parameters, for example:

cfg.rope_scaling = None
cfg.rope_parameters = {"rope_type": "default", "rope_theta": cfg.rope_theta}

With that change, the backbone can proceed past initialization into weight loading.

2. The tokenizer dependency should not be on the engine startup critical path

The MoE deploy enables:

trust_remote_code: true

and the serving path loads the tokenizer through that route. The tokenizer implementation imports:

from chat_format import Chat

This may be an environment/dependency issue rather than a tokenizer logic bug, but the current deploy path makes that dependency fatal during engine startup. In the runtime environment I tested, the dependency is not available, so tokenizer loading fails with:

ImportError: No module named 'chat_format'

This is not only a prompt-builder issue. vLLM core also initializes a tokenizer through StructuredOutputManager, so StageEngineCore can fail before reaching READY.

For this TTS path, the engine receives prompt_token_ids, so it should not need the custom tokenizer for detokenization. I think the MoE deploy should either declare/install this dependency explicitly, or avoid putting it on the engine startup path by setting:

engine_args:
  skip_tokenizer_init: true

The prompt-building side can still fall back to the local fast tokenizer file when the custom tokenizer import fails:

PreTrainedTokenizerFast(tokenizer_file="tokenizer.json")

After applying the RoPE fix, skipping engine-side tokenizer initialization, and using the fast-tokenizer fallback for prompt construction, Stage 0 and Stage 1 can initialize and produce an audio tensor. Without these changes, the deploy config looks complete, but real inference can fail during engine initialization rather than during model generation.

@yuanheng-zhao yuanheng-zhao left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

For now, there exist extra steps to serve the model inclusionAI/Ming-omni-tts-16.8B-A3B - we have to:

  1. download the model repo to local
  2. modify tokenizer_config.json
    • remove the auto_map
    • replace tokenizer class as "tokenizer_class": "PreTrainedTokenizerFast"

We shall find out more elegant workarounds, as these manual resolutions make it inconvenient to serve the model.

Some generated audio outputs for references:

ming_tts_moe_basic.wav
ming_tts_moe_podcast_style.wav
ming_tts_moe_zero_shot.wav

@hsliuustc0106 hsliuustc0106 added the tts code related to tts models label Jun 14, 2026
@LHXuuu

LHXuuu commented Jun 15, 2026

Copy link
Copy Markdown
Contributor Author

For now, there exist extra steps to serve the model inclusionAI/Ming-omni-tts-16.8B-A3B - we have to:

  1. download the model repo to local

  2. modify tokenizer_config.json

    • remove the auto_map
    • replace tokenizer class as "tokenizer_class": "PreTrainedTokenizerFast"

We shall find out more elegant workarounds, as these manual resolutions make it inconvenient to serve the model.

Some generated audio outputs for references:

ming_tts_moe_basic.wav ming_tts_moe_podcast_style.wav ming_tts_moe_zero_shot.wav

@yuanheng-zhao @Sy0307 I've already reported this to the model team, and they've fixed the HF weights/tokenizer upstream.

@yuanheng-zhao

Copy link
Copy Markdown
Collaborator

For now, there exist extra steps to serve the model inclusionAI/Ming-omni-tts-16.8B-A3B - we have to:

  1. download the model repo to local

  2. modify tokenizer_config.json

    • remove the auto_map
    • replace tokenizer class as "tokenizer_class": "PreTrainedTokenizerFast"

We shall find out more elegant workarounds, as these manual resolutions make it inconvenient to serve the model.
Some generated audio outputs for references:
ming_tts_moe_basic.wav ming_tts_moe_podcast_style.wav ming_tts_moe_zero_shot.wav

@yuanheng-zhao @Sy0307 I've already reported this to the model team, and they've fixed the HF weights/tokenizer upstream.

Now it works with model ID.

vllm-omni serve inclusionAI/Ming-omni-tts-16.8B-A3B \
	--deploy-config vllm_omni/deploy/ming_tts_moe.yaml \
	--omni 

Signed-off-by: Yuanheng Zhao <jonathan.zhaoyh@gmail.com>
@yuanheng-zhao

yuanheng-zhao commented Jun 15, 2026 •

Copy link
Copy Markdown
Collaborator

Tests on 354c23f, now default deploy config yaml works:

inclusionAI/Ming-omni-tts-16.8B-A3B

vllm-omni serve inclusionAI/Ming-omni-tts-16.8B-A3B \
	--omni \
	--port 8091

And using the example requests in recipe recipes/inclusionAI/Ming-omni-tts.md, got

ming_tts_basic-moe16.8B.wav
ming_tts_podcast_style-moe16.8B.wav
ming_tts_zero_shot-moe16.8B.wav

inclusionAI/Ming-omni-tts-0.5B

vllm-omni serve inclusionAI/Ming-omni-tts-0.5B  --omni  --port 8091

ming_tts_basic-dense0.5B.wav
ming_tts_podcast_style-dense0.5B.wav
ming_tts_zero_shot-dense0.5B.wav

Comment thread vllm_omni/engine/arg_utils.py Outdated
def execute(self, llm_cond, his_lat, cfg_strength, sigma, temperature):
bsz = llm_cond.shape[0]
device, dtype = llm_cond.device, llm_cond.dtype
y0 = torch.randn((bsz, self.patch_size, self.latent_dim), device=device, dtype=dtype)

@yuanheng-zhao yuanheng-zhao Jun 15, 2026 •

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Same concern at here #4341 (comment) @LHXuuu

@hsliuustc0106

Copy link
Copy Markdown
Collaborator

any progress?

… mrope

Signed-off-by: LHXuuu <xulianhao.xlh@antgroup.com>
@LHXuuu

LHXuuu commented Jun 17, 2026

Copy link
Copy Markdown
Contributor Author

The overall MoE adaptation looks directionally right to me: the prompt format / special tokens are aligned with the upstream behavior, treating the TTS-only path as regular 1D RoPE is a reasonable choice, and the CFM CUDAGraph work is a useful optimization direction.

However, I found two startup-path blockers in real inference. The common root cause is that this patch adapts the model forward path, but it does not fully account for what the vLLM engine consumes during initialization: both the normalized config and the tokenizer can be touched before the request ever reaches the model forward.

1. Stripping 3D mRoPE leaves rope_parameters unusable for vLLM's Bailing backbone

The code currently strips the 3D mRoPE config with:

cfg.rope_scaling = None

I agree with the intent here: the TTS path has no vision/video tokens, so the 3D mRoPE path should degenerate to regular 1D RoPE. The problem is that, in the vLLM runtime I tested, the adapted config ends up with:

rope_scaling None
rope_parameters None

But vLLM's BailingAttention initializes RoPE by mutating this dict:

config.rope_parameters["partial_rotary_factor"] = rotary_dim / self.head_dim

So the model can fail during backbone initialization before weight loading. If this PR intentionally downgrades 3D mRoPE to regular 1D RoPE, it should also explicitly rebuild the regular RoPE parameters, for example:

cfg.rope_scaling = None
cfg.rope_parameters = {"rope_type": "default", "rope_theta": cfg.rope_theta}

With that change, the backbone can proceed past initialization into weight loading.

2. The tokenizer dependency should not be on the engine startup critical path

The MoE deploy enables:

trust_remote_code: true

and the serving path loads the tokenizer through that route. The tokenizer implementation imports:

from chat_format import Chat

This may be an environment/dependency issue rather than a tokenizer logic bug, but the current deploy path makes that dependency fatal during engine startup. In the runtime environment I tested, the dependency is not available, so tokenizer loading fails with:

ImportError: No module named 'chat_format'

This is not only a prompt-builder issue. vLLM core also initializes a tokenizer through StructuredOutputManager, so StageEngineCore can fail before reaching READY.

For this TTS path, the engine receives prompt_token_ids, so it should not need the custom tokenizer for detokenization. I think the MoE deploy should either declare/install this dependency explicitly, or avoid putting it on the engine startup path by setting:

engine_args:
  skip_tokenizer_init: true

The prompt-building side can still fall back to the local fast tokenizer file when the custom tokenizer import fails:

PreTrainedTokenizerFast(tokenizer_file="tokenizer.json")

After applying the RoPE fix, skipping engine-side tokenizer initialization, and using the fast-tokenizer fallback for prompt construction, Stage 0 and Stage 1 can initialize and produce an audio tensor. Without these changes, the deploy config looks complete, but real inference can fail during engine initialization rather than during model generation.

@Sy0307 Good catch, thanks — fixed.

Signed-off-by: LHXuuu <xulianhao.xlh@antgroup.com>
@yuanheng-zhao yuanheng-zhao added this to the v0.24.0 milestone Jun 18, 2026
LHXuuu added 3 commits June 21, 2026 23:30
# Conflicts:
#	vllm_omni/model_executor/models/ming_tts/ming_tts_llm.py
…ow head

Signed-off-by: LHXuuu <xulianhao.xlh@antgroup.com>
Signed-off-by: LHXuuu <xulianhao.xlh@antgroup.com>
@Gaohan123

Copy link
Copy Markdown
Collaborator

@yuanheng-zhao @Sy0307 PTAL again

Signed-off-by: Yuanheng Zhao <jonathan.zhaoyh@gmail.com>
@LHXuuu LHXuuu changed the title [WIP][Model] Add Ming-omni-tts MoE 16.8B-A3B + CFM CUDAGraph [Model] Add Ming-omni-tts MoE 16.8B-A3B + CFM CUDAGraph Jun 24, 2026
Signed-off-by: LHXuuu <xulianhao.xlh@antgroup.com>
atc = cfg.audio_tokenizer_config
semantic_module_kwargs = _nested_get(atc, "semantic_module_kwargs", default=None)
if semantic_module_kwargs is not None:
raise ValueError("Ming dense 0.5B expects audio_tokenizer_config.semantic_module_kwargs to be null.")

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Better distinguish model type or remove "dense 0.5B", as we support the MoE-16.8B now

)

# Fast path: CUDAGraph-captured flow head (CFM sampling + Aggregator +
# stop head). Falls back permanently to the eager path on any failure.

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Let's consider if we want to make it fallback permanently to eager path on any failure, or re-try in future

Comment thread vllm_omni/model_executor/models/ming_tts/fm/cfm_cudagraph.py
LHXuuu added 2 commits June 25, 2026 00:18
…lver.integrate

Signed-off-by: LHXuuu <xulianhao.xlh@antgroup.com>
… 1e-5)

Signed-off-by: LHXuuu <xulianhao.xlh@antgroup.com>
@yuanheng-zhao
yuanheng-zhao enabled auto-merge (squash) June 24, 2026 16:35
@yuanheng-zhao yuanheng-zhao added the ready label to trigger buildkite CI label Jun 24, 2026
Signed-off-by: LHXuuu <xulianhao.xlh@antgroup.com>
auto-merge was automatically disabled June 25, 2026 03:31

Head branch was pushed to by a user without write access

@LHXuuu
LHXuuu requested a review from yenuo26 as a code owner June 25, 2026 03:31
@princepride princepride added ready label to trigger buildkite CI and removed ready label to trigger buildkite CI labels Jun 25, 2026
@yuanheng-zhao

Copy link
Copy Markdown
Collaborator

The CI errors encountered seem to be unrelated to this PR. PTAL @Gaohan123 , @yenuo26 , @tjtanaa

@hsliuustc0106 hsliuustc0106 left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

lgtm

hsliuustc0106 and others added 3 commits June 25, 2026 17:06
Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
…voice-clone collapse

Signed-off-by: LHXuuu <xulianhao.xlh@antgroup.com>
… merge

Signed-off-by: LHXuuu <xulianhao.xlh@antgroup.com>
@hsliuustc0106
hsliuustc0106 merged commit 0b04f2a into vllm-project:main Jun 25, 2026
6 checks passed
@LHXuuu
LHXuuu deleted the ming_tts_moe branch July 10, 2026 07:32
Moore-Z added a commit to Moore-Z/vllm-omni-RJ that referenced this pull request Jul 14, 2026
… ODE integration, seeded FM noise, light package init

Review fixes from PR vllm-project#4765 (linyueqian):

* High-1: prefill branch now pushes a (req_id, None) placeholder onto
  _results_queue so compute_logits' positional pairing holds for mixed
  prefill+decode batches (voxcpm2 pattern).  Verified: unit tests on the
  pairing contract + 3-request e2e, all pass.
* High-2: Euler ODE integration state (z, times, CFG blend) now fp32;
  DiT matmuls stay bf16 under autocast, result cast back to sequence
  dtype (same fix as Ming CFM, PR vllm-project#4341).
* Medium-1: FM noise drawn from a per-request per-step blake2b-seeded
  Generator instead of the global CUDA RNG (voxcpm2 parity).  Verified:
  in-session replay is bit-identical.
* Medium-3: drop eager talker import from package __init__ (registry
  imports submodules directly; voxcpm2 convention).

Signed-off-by: Moore-Z <zhumoore9661996@gmail.com>
Moore-Z added a commit to Moore-Z/vllm-omni-RJ that referenced this pull request Jul 15, 2026
… ODE integration, seeded FM noise, light package init

Review fixes from PR vllm-project#4765 (linyueqian):

* High-1: prefill branch now pushes a (req_id, None) placeholder onto
  _results_queue so compute_logits' positional pairing holds for mixed
  prefill+decode batches (voxcpm2 pattern).  Verified: unit tests on the
  pairing contract + 3-request e2e, all pass.
* High-2: Euler ODE integration state (z, times, CFG blend) now fp32;
  DiT matmuls stay bf16 under autocast, result cast back to sequence
  dtype (same fix as Ming CFM, PR vllm-project#4341).
* Medium-1: FM noise drawn from a per-request per-step blake2b-seeded
  Generator instead of the global CUDA RNG (voxcpm2 parity).  Verified:
  in-session replay is bit-identical.
* Medium-3: drop eager talker import from package __init__ (registry
  imports submodules directly; voxcpm2 convention).

Signed-off-by: Moore-Z <zhumoore9661996@gmail.com>
Moore-Z added a commit to Moore-Z/vllm-omni-RJ that referenced this pull request Aug 11, 2026
… ODE integration, seeded FM noise, light package init

Review fixes from PR vllm-project#4765 (linyueqian):

* High-1: prefill branch now pushes a (req_id, None) placeholder onto
  _results_queue so compute_logits' positional pairing holds for mixed
  prefill+decode batches (voxcpm2 pattern).  Verified: unit tests on the
  pairing contract + 3-request e2e, all pass.
* High-2: Euler ODE integration state (z, times, CFG blend) now fp32;
  DiT matmuls stay bf16 under autocast, result cast back to sequence
  dtype (same fix as Ming CFM, PR vllm-project#4341).
* Medium-1: FM noise drawn from a per-request per-step blake2b-seeded
  Generator instead of the global CUDA RNG (voxcpm2 parity).  Verified:
  in-session replay is bit-identical.
* Medium-3: drop eager talker import from package __init__ (registry
  imports submodules directly; voxcpm2 convention).

Signed-off-by: Moore-Z <zhumoore9661996@gmail.com>
khairulkabir1661 pushed a commit to khairulkabir1661/vllm-omni that referenced this pull request Sep 25, 2026
…#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>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

ready label to trigger buildkite CI tts code related to tts models

Projects

None yet

Development

Successfully merging this pull request may close these issues.

6 participants