[Model] Add Ming-omni-tts dense 0.5B pipeline (#2906)

Signed-off-by: akshatvishu <akshatnayak197@gmail.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>

tests/e2e/offline_inference/test_ming_flash_omni_expansion.py

@@ -51,7 +51,7 @@ test_params = [(m, c) for m in models for c in [get_eager_config()]]
 
 @pytest.mark.core_model
 @pytest.mark.omni
-@hardware_test(res={"cuda": "H100"}, num_cards=4)
+@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=4)
 @pytest.mark.parametrize("omni_runner", test_params_thinker, indirect=True)
 def test_text_to_text(omni_runner, omni_runner_handler) -> None:
     """
@@ -67,7 +67,7 @@ def test_text_to_text(omni_runner, omni_runner_handler) -> None:
 
 @pytest.mark.core_model
 @pytest.mark.omni
-@hardware_test(res={"cuda": "H100"}, num_cards=4)
+@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=4)
 @pytest.mark.parametrize("omni_runner", test_params_thinker, indirect=True)
 def test_image_to_text(omni_runner, omni_runner_handler) -> None:
     """
@@ -84,7 +84,7 @@ def test_image_to_text(omni_runner, omni_runner_handler) -> None:
 
 @pytest.mark.core_model
 @pytest.mark.omni
-@hardware_test(res={"cuda": "H100"}, num_cards=4)
+@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=4)
 @pytest.mark.parametrize("omni_runner", test_params_thinker, indirect=True)
 def test_audio_to_text(omni_runner, omni_runner_handler) -> None:
     """
@@ -103,7 +103,7 @@ def test_audio_to_text(omni_runner, omni_runner_handler) -> None:
 
 @pytest.mark.core_model
 @pytest.mark.omni
-@hardware_test(res={"cuda": "H100"}, num_cards=4)
+@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=4)
 @pytest.mark.parametrize("omni_runner", test_params_thinker, indirect=True)
 def test_video_to_text(omni_runner, omni_runner_handler) -> None:
     """
@@ -120,7 +120,7 @@ def test_video_to_text(omni_runner, omni_runner_handler) -> None:
 
 @pytest.mark.core_model
 @pytest.mark.omni
-@hardware_test(res={"cuda": "H100"}, num_cards=4)
+@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=4)
 @pytest.mark.parametrize("omni_runner", test_params_thinker, indirect=True)
 def test_mixed_to_text(omni_runner, omni_runner_handler) -> None:
     """
@@ -140,7 +140,7 @@ def test_mixed_to_text(omni_runner, omni_runner_handler) -> None:
 
 @pytest.mark.advanced_model
 @pytest.mark.omni
-@hardware_test(res={"cuda": "H100"}, num_cards=4)
+@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=4)
 @pytest.mark.parametrize("omni_runner", test_params, indirect=True)
 def test_text_to_audio(omni_runner, omni_runner_handler) -> None:
     """
@@ -156,7 +156,7 @@ def test_text_to_audio(omni_runner, omni_runner_handler) -> None:
 
 @pytest.mark.advanced_model
 @pytest.mark.omni
-@hardware_test(res={"cuda": "H100"}, num_cards=4)
+@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=4)
 @pytest.mark.parametrize("omni_runner", test_params, indirect=True)
 def test_image_to_audio(omni_runner, omni_runner_handler) -> None:
     """

tests/e2e/offline_inference/test_ming_tts.py

@@ -0,0 +1,295 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""End-to-end offline inference tests for Ming-omni-tts."""
+
+import asyncio
+import uuid
+
+import numpy as np
+import pytest
+import torch
+from transformers import AutoTokenizer
+from vllm import SamplingParams
+from vllm.inputs import tokens_input
+
+from tests.helpers.mark import hardware_test
+from tests.helpers.runtime import OmniRunner
+from tests.helpers.stage_config import get_deploy_config_path
+from vllm_omni import AsyncOmni
+from vllm_omni.model_executor.models.ming_tts.config_ming_tts import (
+    KEY_MAX_DECODE_STEPS,
+    SAMPLE_RATE,
+    TEXT_EOS_TOKEN_ID,
+)
+from vllm_omni.model_executor.models.ming_tts.prompt_assembly import DEFAULT_PROMPT, build_ming_dense_prompt
+
+MODEL = "inclusionAI/Ming-omni-tts-0.5B"
+DEPLOY_CONFIG = get_deploy_config_path("ming_tts.yaml")
+TEST_TEXT = "我会一直在这里陪着你，直到你慢慢地沉入那个最温柔的梦里。"
+TEST_INSTRUCTION = "轻柔的ASMR耳语，慢速，贴近麦克风"
+MIN_AUDIO_SAMPLES = 1000
+
+
+@pytest.fixture(scope="module")
+def ming_tokenizer():
+    return AutoTokenizer.from_pretrained(MODEL, trust_remote_code=False)
+
+
+@pytest.fixture
+def ming_engine():
+    with OmniRunner(
+        MODEL,
+        deploy_config=DEPLOY_CONFIG,
+        stage_init_timeout=300,
+        enforce_eager=True,
+    ) as runner:
+        yield runner.omni
+
+
+@pytest.fixture
+def async_omni_engine():
+    engine = AsyncOmni(
+        model=MODEL,
+        deploy_config=DEPLOY_CONFIG,
+        stage_init_timeout=300,
+        enforce_eager=True,
+    )
+    yield engine
+    engine.shutdown()
+
+
+def _build_prompt(
+    tokenizer,
+    *,
+    text: str = TEST_TEXT,
+    instruction=TEST_INSTRUCTION,
+    use_zero_spk_emb: bool = True,
+) -> dict:
+    prompt_dict = build_ming_dense_prompt(
+        tokenizer,
+        prompt=DEFAULT_PROMPT,
+        text=text,
+        instruction=instruction,
+        runtime_controls={KEY_MAX_DECODE_STEPS: 200},
+        use_zero_spk_emb=use_zero_spk_emb,
+    )
+    prompt = tokens_input(prompt_token_ids=prompt_dict["prompt_token_ids"])
+    prompt["prompt"] = prompt_dict["prompt"]
+    prompt["text"] = prompt_dict["text"]
+    prompt["additional_information"] = prompt_dict["additional_information"]
+    return prompt
+
+
+def _sampling_params_list() -> list[SamplingParams]:
+    return [
+        SamplingParams(
+            temperature=0.0,
+            max_tokens=201,
+            stop_token_ids=[int(TEXT_EOS_TOKEN_ID)],
+        ),
+        SamplingParams(temperature=0.0, max_tokens=1),
+    ]
+
+
+def _flatten_audio(audio) -> torch.Tensor:
+    if isinstance(audio, list):
+        parts = [torch.as_tensor(item, dtype=torch.float32).reshape(-1).cpu() for item in audio]
+        parts = [item for item in parts if item.numel() > 0]
+        if not parts:
+            return torch.zeros((0,), dtype=torch.float32)
+        return torch.cat(parts, dim=0)
+    return torch.as_tensor(audio, dtype=torch.float32).reshape(-1).cpu()
+
+
+def _extract_audio(multimodal_output: dict) -> torch.Tensor:
+    assert isinstance(multimodal_output, dict), f"Expected dict, got {type(multimodal_output)}"
+    audio = multimodal_output.get("audio", multimodal_output.get("model_outputs"))
+    assert audio is not None, f"No audio output found, keys={list(multimodal_output.keys())}"
+    waveform = _flatten_audio(audio)
+    if waveform.numel() == 0:
+        raise RuntimeError("Generated audio waveform is empty")
+    return waveform
+
+
+def _extract_sample_rate(multimodal_output: dict) -> int:
+    sample_rate = multimodal_output.get("sr")
+    if sample_rate is None:
+        raise RuntimeError("Expected multimodal_output['sr']")
+    if isinstance(sample_rate, list):
+        sample_rate = sample_rate[-1]
+    if hasattr(sample_rate, "item"):
+        sample_rate = sample_rate.item()
+    return int(sample_rate)
+
+
+def _extract_final_audio_outputs(outputs):
+    final_outputs = []
+    for item in outputs:
+        if getattr(item, "final_output_type", None) == "audio":
+            final_outputs.append(item)
+            continue
+        request_output = getattr(item, "request_output", None)
+        if request_output is None:
+            continue
+        multimodal_output = getattr(request_output, "multimodal_output", None)
+        if isinstance(multimodal_output, dict):
+            final_outputs.append(item)
+            continue
+        completions = getattr(request_output, "outputs", None) or []
+        if any(isinstance(getattr(completion, "multimodal_output", None), dict) for completion in completions):
+            final_outputs.append(item)
+    return final_outputs
+
+
+def _extract_multimodal_output(output) -> dict:
+    multimodal_output = getattr(output, "multimodal_output", None)
+    if isinstance(multimodal_output, dict):
+        return multimodal_output
+
+    request_output = getattr(output, "request_output", None)
+    if request_output is not None:
+        multimodal_output = getattr(request_output, "multimodal_output", None)
+        if isinstance(multimodal_output, dict):
+            return multimodal_output
+        completions = getattr(request_output, "outputs", None) or []
+        for completion in completions:
+            multimodal_output = getattr(completion, "multimodal_output", None)
+            if isinstance(multimodal_output, dict):
+                return multimodal_output
+
+    raise AssertionError("No multimodal audio output found in Ming generate results")
+
+
+@pytest.mark.advanced_model
+@pytest.mark.omni
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+def test_ming_tts_offline_basic(ming_engine, ming_tokenizer) -> None:
+    """Test blocking Ming generation through Omni."""
+    outputs = ming_engine.generate(
+        prompts=[_build_prompt(ming_tokenizer)],
+        sampling_params_list=_sampling_params_list(),
+        py_generator=False,
+    )
+    final_outputs = _extract_final_audio_outputs(outputs)
+    assert len(final_outputs) == 1, f"Expected one final audio output, got {len(final_outputs)}"
+    multimodal_output = _extract_multimodal_output(final_outputs[0])
+    waveform = _extract_audio(multimodal_output)
+    sample_rate = _extract_sample_rate(multimodal_output)
+    assert waveform.ndim == 1
+    assert waveform.shape[0] == waveform.numel()
+    assert waveform.numel() > MIN_AUDIO_SAMPLES
+    assert np.max(np.abs(waveform.numpy())) > 0.01, "Audio appears silent"
+    assert sample_rate == SAMPLE_RATE, f"Expected Ming output sample rate {SAMPLE_RATE}, got {sample_rate}"
+
+
+@pytest.mark.advanced_model
+@pytest.mark.omni
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+def test_ming_tts_speaker_conditioning_differs(ming_engine, ming_tokenizer) -> None:
+    """Test that different Ming speaker controls produce different waveform outputs."""
+    style_outputs = ming_engine.generate(
+        prompts=[_build_prompt(ming_tokenizer)],
+        sampling_params_list=_sampling_params_list(),
+        py_generator=False,
+    )
+    ip_outputs = ming_engine.generate(
+        prompts=[_build_prompt(ming_tokenizer, text=TEST_TEXT, instruction={"IP": "灵小甄"}, use_zero_spk_emb=True)],
+        sampling_params_list=_sampling_params_list(),
+        py_generator=False,
+    )
+
+    style_final_outputs = _extract_final_audio_outputs(style_outputs)
+    ip_final_outputs = _extract_final_audio_outputs(ip_outputs)
+    assert len(style_final_outputs) == 1, "No style audio output produced"
+    assert len(ip_final_outputs) == 1, "No IP audio output produced"
+
+    style_waveform = _extract_audio(_extract_multimodal_output(style_final_outputs[0]))
+    ip_waveform = _extract_audio(_extract_multimodal_output(ip_final_outputs[0]))
+    assert style_waveform.numel() > MIN_AUDIO_SAMPLES
+    assert ip_waveform.numel() > MIN_AUDIO_SAMPLES
+    assert np.max(np.abs(style_waveform.numpy())) > 0.01, "Style audio appears silent"
+    assert np.max(np.abs(ip_waveform.numpy())) > 0.01, "IP audio appears silent"
+
+    overlap = min(int(style_waveform.numel()), int(ip_waveform.numel()))
+    mean_abs_diff = torch.mean(torch.abs(style_waveform[:overlap] - ip_waveform[:overlap])).item()
+    assert style_waveform.shape != ip_waveform.shape or mean_abs_diff > 1e-4, (
+        "Speaker-conditioned outputs should differ, but style and IP waveforms were effectively identical"
+    )
+
+
+@pytest.mark.advanced_model
+@pytest.mark.omni
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+def test_ming_tts_multiple_prompts_queued(ming_engine, ming_tokenizer) -> None:
+    """Regression: supported max_num_seqs=1 config must still drain queued prompts."""
+    prompts = [
+        _build_prompt(
+            ming_tokenizer,
+            text="第一条语音用于验证 Ming dense 队列中的长请求可以完成。",
+            instruction="平静自然的旁白，语速中等",
+        ),
+        _build_prompt(
+            ming_tokenizer,
+            text="第二条短语音也必须生成完成。",
+            instruction="清晰明亮的女声",
+        ),
+    ]
+    outputs = ming_engine.generate(
+        prompts=prompts,
+        sampling_params_list=_sampling_params_list(),
+        py_generator=False,
+    )
+    final_outputs = _extract_final_audio_outputs(outputs)
+    assert len(final_outputs) == len(prompts), f"Expected {len(prompts)} audio outputs, got {len(final_outputs)}"
+    for i, output in enumerate(final_outputs):
+        waveform = _extract_audio(_extract_multimodal_output(output))
+        duration_s = waveform.numel() / SAMPLE_RATE
+        assert 0.1 < duration_s < 30.0, f"Request {i} audio duration out of range: {duration_s:.2f}s"
+        assert np.max(np.abs(waveform.numpy())) > 0.01, f"Request {i} audio appears silent"
+
+
+@pytest.mark.advanced_model
+@pytest.mark.omni
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+def test_ming_tts_offline_streaming(async_omni_engine, ming_tokenizer) -> None:
+    """Test async_chunk streaming Ming generation through AsyncOmni."""
+
+    async def _run() -> None:
+        all_audio_chunks = []
+        accumulated_samples = 0
+        chunk_idx = 0
+        sample_rate = None
+        async for stage_output in async_omni_engine.generate(
+            prompt=_build_prompt(ming_tokenizer),
+            request_id=str(uuid.uuid4()),
+            sampling_params_list=_sampling_params_list(),
+        ):
+            multimodal_output = stage_output.multimodal_output or {}
+            audio = multimodal_output.get("audio")
+            if "sr" in multimodal_output:
+                sample_rate = _extract_sample_rate(multimodal_output)
+            if audio is None:
+                continue
+            finished = stage_output.finished
+            if isinstance(audio, torch.Tensor):
+                if finished:
+                    audio_chunk = audio[accumulated_samples:].float().detach().cpu()
+                else:
+                    audio_chunk = audio.float().detach().cpu()
+            elif isinstance(audio, list):
+                audio_chunk = _flatten_audio(audio)
+            else:
+                audio_chunk = torch.as_tensor(audio, dtype=torch.float32).reshape(-1).cpu()
+            accumulated_samples += int(audio_chunk.numel())
+            chunk_idx += 1
+            if audio_chunk.numel() > 0:
+                all_audio_chunks.append(audio_chunk)
+        assert all_audio_chunks, "No streaming audio chunks received"
+        waveform = torch.cat(all_audio_chunks, dim=0)
+        assert waveform.numel() > MIN_AUDIO_SAMPLES
+        assert np.max(np.abs(waveform.numpy())) > 0.01, "Audio appears silent"
+        assert sample_rate is not None, "Streaming path did not return a sample rate"
+        assert sample_rate == SAMPLE_RATE, f"Expected Ming output sample rate {SAMPLE_RATE}, got {sample_rate}"
+
+    asyncio.run(_run())

tests/e2e/online_serving/test_ming_tts.py

@@ -0,0 +1,92 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+E2E online serving tests for Ming-omni-tts (dense) model.
+Tests text-to-audio via /v1/audio/speech endpoint.
+"""
+
+import os
+
+os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
+
+import pytest
+
+from tests.helpers.mark import hardware_test
+from tests.helpers.runtime import OmniServerParams
+from tests.helpers.stage_config import get_deploy_config_path
+
+pytestmark = [pytest.mark.advanced_model, pytest.mark.omni]
+
+MODEL = "inclusionAI/Ming-omni-tts-0.5B"
+DEPLOY_CONFIG = get_deploy_config_path("ming_tts.yaml")
+
+no_async_chunk_params = [
+    pytest.param(
+        OmniServerParams(
+            model=MODEL,
+            stage_config_path=DEPLOY_CONFIG,
+            server_args=["--enforce-eager", "--no-async-chunk"],
+        ),
+        id="no_async_chunk",
+    )
+]
+
+async_chunk_params = [
+    pytest.param(
+        OmniServerParams(
+            model=MODEL,
+            stage_config_path=DEPLOY_CONFIG,
+            server_args=["--enforce-eager"],
+        ),
+        id="async_chunk",
+    )
+]
+
+
+def get_prompt(prompt_type="zh"):
+    prompts = {
+        "zh": "今天天气真不错，适合出去散散步。",
+        "zh_short": "这款产品的名字，叫变态坑爹牛肉丸。",
+    }
+    return prompts.get(prompt_type, prompts["zh"])
+
+
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+@pytest.mark.parametrize("omni_server", no_async_chunk_params, indirect=True)
+def test_text_to_audio_non_streaming_001(omni_server, openai_client) -> None:
+    """
+    Deploy Setting: ming_tts.yaml with --no-async-chunk
+    Input Modal: text
+    Output Modal: audio
+    Input Setting: stream=False
+    Datasets: two concurrent requests
+    """
+    request_config = {
+        "model": omni_server.model,
+        "input": get_prompt("zh"),
+        "stream": False,
+        "response_format": "wav",
+        "timeout": 300.0,
+    }
+    openai_client.send_audio_speech_request(request_config, request_num=2)
+
+
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+@pytest.mark.parametrize("omni_server", async_chunk_params, indirect=True)
+def test_text_to_audio_streaming_001(omni_server, openai_client) -> None:
+    """
+    Deploy Setting: ming_tts.yaml (async_chunk=true)
+    Input Modal: text + voice
+    Output Modal: audio (streamed)
+    Input Setting: stream=True
+    Datasets: single request
+    """
+    request_config = {
+        "model": omni_server.model,
+        "input": get_prompt("zh_short"),
+        "voice": "灵小甄",
+        "stream": True,
+        "response_format": "wav",
+        "timeout": 300.0,
+    }
+    openai_client.send_audio_speech_request(request_config)

tests/model_executor/models/ming_tts/test_ming_shared_modules.py

@@ -0,0 +1,185 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from vllm_omni.entrypoints.openai.serving_speech import OmniOpenAIServingSpeech
+from vllm_omni.model_executor.models.common.ming.audio_vae import AudioVAEConfig
+from vllm_omni.model_executor.models.common.ming.fm import Solver
+from vllm_omni.model_executor.models.ming_tts.constants import (
+    AGGREGATOR_HIDDEN_SIZE,
+    HISTORY_PATCH_SIZE,
+    LATENT_DIM,
+    LLM_HIDDEN_SIZE,
+    LLM_VOCAB_SIZE,
+    PATCH_SIZE,
+    SAMPLE_RATE,
+    VAE_PATCH_SIZE,
+)
+from vllm_omni.model_executor.models.ming_tts.flowloss_head import FlowLoss
+from vllm_omni.model_executor.models.ming_tts.ming_tts_llm import MingLLMModel
+from vllm_omni.model_executor.models.ming_tts.validation import validate_ming_tts_config
+from vllm_omni.model_executor.models.output_templates import OmniOutput
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.tts]
+
+
+def test_ming_tts_audio_vae_uses_common_config():
+    """AudioVAEConfig is shared by Ming dense and Ming flash modules."""
+    cfg = AudioVAEConfig(sample_rate=16000, patch_size=-1)
+
+    assert cfg.sample_rate == 16000
+    assert cfg.patch_size == -1
+
+
+def test_ming_tts_cfm_solver_uses_common_implementation():
+    """Ming dense imports the shared solver implementation directly."""
+    assert Solver.__module__ == "vllm_omni.model_executor.models.common.ming.fm"
+
+
+def test_ming_tts_flowloss_preserves_checkpoint_prefix():
+    flowloss = FlowLoss(z_channels=4, llm_cond_dim=8, hidden_size=16, depth=1, num_heads=2)
+
+    assert any(name.startswith("cfm.model.") for name in flowloss.state_dict())
+
+
+def test_ming_dense_validation_rejects_semantic_audio_vae_config():
+    """Dense 0.5B validation rejects semantic AudioVAE configs."""
+    cfg = SimpleNamespace(
+        audio_dummy_token_id=151705,
+        audio_eos_token_id=151704,
+        text_eos_token_id=151669,
+        audio_tokenizer_config=AudioVAEConfig(
+            sample_rate=SAMPLE_RATE,
+            patch_size=VAE_PATCH_SIZE,
+            semantic_module_kwargs={"whisper_encoder": {}},
+            enc_kwargs={"latent_dim": LATENT_DIM, "input_dim": 882, "hop_size": 882},
+            dec_kwargs={"latent_dim": LATENT_DIM, "output_dim": 882},
+        ),
+        latent_dim=LATENT_DIM,
+        patch_size=PATCH_SIZE,
+        history_patch_size=HISTORY_PATCH_SIZE,
+        llm_hidden_size=LLM_HIDDEN_SIZE,
+        llm_vocab_size=LLM_VOCAB_SIZE,
+        sample_rate=SAMPLE_RATE,
+        vae_patch_size=VAE_PATCH_SIZE,
+        llm_config={"hidden_size": LLM_HIDDEN_SIZE},
+        aggregator_config={"hidden_size": AGGREGATOR_HIDDEN_SIZE},
+        ditar_config={"hidden_size": AGGREGATOR_HIDDEN_SIZE},
+        latent_chunk_size=1,
+        latent_left_context=0,
+        max_decode_steps=1,
+        stop_head_threshold=0.5,
+        stop_head_min_steps=0,
+    )
+
+    with pytest.raises(ValueError, match="semantic_module_kwargs"):
+        validate_ming_tts_config(cfg)
+
+
+def test_ming_instruction_parser_preserves_dense_and_flash_defaults():
+    """Ming dense and Ming flash keep distinct instruction defaults."""
+    serving = object.__new__(OmniOpenAIServingSpeech)
+    serving.uploaded_speakers = {"uploaded": {}}
+
+    dense_plain = serving._parse_ming_instruction(SimpleNamespace(instructions="calm", language=None, voice=None))
+    assert dense_plain == "calm"
+
+    dense_with_fields = serving._parse_ming_instruction(
+        SimpleNamespace(instructions="calm", language="Auto", voice="灵小甄")
+    )
+    assert dense_with_fields == {"IP": "灵小甄", "风格": "calm"}
+
+    flash_fields = serving._parse_ming_instruction_fields(
+        SimpleNamespace(instructions="calm", language="粤语", voice="灵小甄")
+    )
+    assert flash_fields == {"风格": "calm"}
+
+
+def _make_ming_logits_model(vocab_size=8):
+    model = object.__new__(MingLLMModel)
+    model.ming_config = SimpleNamespace(
+        llm_vocab_size=vocab_size,
+        max_decode_steps=1,
+        stop_head_min_steps=0,
+        text_eos_token_id=7,
+    )
+    model._last_text_mode = False
+    model._last_ming_next_token_ids = None
+    return model
+
+
+def test_ming_compute_logits_uses_cached_forced_next_token_ids():
+    model = _make_ming_logits_model()
+    model._last_ming_next_token_ids = [2, 5]
+
+    logits = MingLLMModel.compute_logits(model, torch.zeros((2, 4)), SimpleNamespace())
+
+    assert logits.shape == (2, 8)
+    assert logits[0, 2].item() == 0.0
+    assert logits[1, 5].item() == 0.0
+    assert torch.isneginf(logits[0, [0, 1, 3, 4, 5, 6, 7]]).all()
+    assert torch.isneginf(logits[1, [0, 1, 2, 3, 4, 6, 7]]).all()
+    assert model._last_ming_next_token_ids is None
+
+
+def test_ming_compute_logits_falls_back_to_dummy_token_id():
+    model = _make_ming_logits_model()
+
+    logits = MingLLMModel.compute_logits(model, torch.zeros((1, 4)), SimpleNamespace())
+    assert logits.shape == (1, 8)
+    assert logits[0, 7].item() == 0.0
+    assert torch.isneginf(logits[0, [0, 1, 2, 3, 4, 5, 6]]).all()
+
+
+def test_ming_forward_non_decode_return_clears_cached_forced_next_token_ids():
+    class FakeBackbone:
+        def __call__(self, **kwargs):
+            return kwargs["inputs_embeds"]
+
+    model = _make_ming_logits_model()
+    model.model = FakeBackbone()
+    model._last_ming_next_token_ids = [2]
+
+    output = MingLLMModel.forward(
+        model,
+        input_ids=torch.tensor([1]),
+        positions=torch.tensor([0]),
+        inputs_embeds=torch.zeros((1, 4)),
+        model_intermediate_buffer=[],
+    )
+
+    assert isinstance(output, OmniOutput)
+    assert output.multimodal_outputs is None
+    assert model._last_ming_next_token_ids is None
+
+
+def test_ming_compute_logits_rejects_forced_token_batch_mismatch():
+    model = _make_ming_logits_model()
+    model._last_ming_next_token_ids = [2]
+
+    with pytest.raises(RuntimeError, match="batch mismatch"):
+        MingLLMModel.compute_logits(model, torch.zeros((2, 4)), SimpleNamespace())
+
+
+def test_ming_compute_logits_text_mode_delegates_to_backbone():
+    class FakeBackbone:
+        def __init__(self):
+            self.hidden_states = None
+
+        def compute_logits(self, hidden_states):
+            self.hidden_states = hidden_states
+            return torch.ones((hidden_states.shape[0], 3))
+
+    model = _make_ming_logits_model(vocab_size=3)
+    model._last_text_mode = True
+    model.model = FakeBackbone()
+    hidden_states = torch.zeros((2, 4))
+
+    logits = MingLLMModel.compute_logits(model, hidden_states, SimpleNamespace())
+
+    assert torch.equal(model.model.hidden_states, hidden_states)
+    assert torch.equal(logits, torch.ones((2, 3)))

tests/model_executor/stage_input_processors/test_qwen3_omni_streaming_helpers.py

@@ -921,6 +921,7 @@ def test_cosyvoice3_full_payload_replace_keys_present() -> None:
 def test_ming_flash_omni_thinker2talker_token_only_smoke() -> None:
     """Smoke: ming_flash_omni token-only carries voice metadata."""
     from vllm_omni.model_executor.stage_input_processors.ming_flash_omni import (
+        thinker2talker,
         thinker2talker_token_only,
     )
 
@@ -938,14 +939,15 @@ def test_ming_flash_omni_thinker2talker_token_only_smoke() -> None:
 
     src = [_Wrap("hello world")]
     prompt = _Prompt({"voice_name": "ZH_FEMALE", "prompt_text": "ref text"})
-    out = thinker2talker_token_only(src, prompt=prompt)
-    assert len(out) == 1
-    assert out[0]["prompt_token_ids"] == [0]  # talker self-tokenizes; dummy id
-    info = out[0]["additional_information"]
-    assert info["text"] == "hello world"
-    assert info["voice_name"] == "ZH_FEMALE"
-    assert info["prompt_text"] == "ref text"
-    assert info["ming_task"] == "omni"
+    for func in (thinker2talker, thinker2talker_token_only):
+        out = func(src, prompt=prompt)
+        assert len(out) == 1
+        assert out[0]["prompt_token_ids"] == [0]
+        info = out[0]["additional_information"]
+        assert info["text"] == "hello world"
+        assert info["voice_name"] == "ZH_FEMALE"
+        assert info["prompt_text"] == "ref text"
+        assert info["ming_task"] == "omni"
 
 
 def test_qwen2_5_omni_thinker2talker_token_only_smoke() -> None: