[BugFix][TTS] Fix MOSS Realtime serving and generation (#5661)

Signed-off-by: akshatvishu <akshatnayak197@gmail.com>

tests/e2e/offline_inference/test_moss_tts_expansion.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 """E2E offline inference tests for MOSS-VoiceGenerator (MossTTSDelayModel, 1.7B).
 
 Weekly expansion coverage (``pytest.mark.slow``) for the delay-model path.
@@ -74,7 +74,6 @@ def _get_test_config() -> str:
 
 pytestmark = [
     pytest.mark.slow,
-    pytest.mark.skip(reason="https://github.com/vllm-project/vllm-omni/issues/4700"),
     pytest.mark.tts,
     pytest.mark.parametrize(
         "omni_runner",
@@ -137,7 +136,7 @@ def _collect_audio(omni_runner: OmniRunner, request: dict) -> tuple[torch.Tensor
     return torch.cat(chunks, dim=0), sr_final
 
 
-@hardware_test(res={"cuda": "L4"}, num_cards=1)
+@hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards=1)
 def test_moss_tts_delay_english(omni_runner: OmniRunner) -> None:
     """VoiceGenerator: English instruction produces non-empty 24 kHz audio."""
     req = _build_request(
@@ -151,7 +150,7 @@ def test_moss_tts_delay_english(omni_runner: OmniRunner) -> None:
     assert not torch.all(audio == 0), "Audio is silence"
 
 
-@hardware_test(res={"cuda": "L4"}, num_cards=1)
+@hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards=1)
 def test_moss_tts_delay_chinese(omni_runner: OmniRunner) -> None:
     """VoiceGenerator: Chinese input produces non-empty audio."""
     req = _build_request(
@@ -165,7 +164,7 @@ def test_moss_tts_delay_chinese(omni_runner: OmniRunner) -> None:
     assert not torch.all(audio == 0)
 
 
-@hardware_test(res={"cuda": "L4"}, num_cards=1)
+@hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards=1)
 def test_moss_tts_delay_batch(omni_runner: OmniRunner) -> None:
     """VoiceGenerator: batch of two requests each returns non-empty audio."""
     requests = [

tests/e2e/offline_inference/test_moss_tts_realtime.py

@@ -95,7 +95,6 @@ def _get_test_config() -> str:
 # ---------------------------------------------------------------------------
 
 pytestmark = [
-    pytest.mark.skip(reason="https://github.com/vllm-project/vllm-omni/issues/4700"),
     pytest.mark.slow,
     pytest.mark.tts,
     pytest.mark.parametrize(
@@ -272,7 +271,7 @@ def _collect_audio(omni_runner: OmniRunner, request: dict) -> tuple[torch.Tensor
 # ---------------------------------------------------------------------------
 
 
-@hardware_test(res={"cuda": "L4"}, num_cards=1)
+@hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards=1)
 def test_moss_tts_realtime_english(omni_runner: OmniRunner, ref_audio_path: str) -> None:
     """MossTTSRealtime: English voice_clone produces non-empty 24 kHz audio."""
     req = _build_request(ref_audio_path, "This is a real-time TTS streaming test.")
@@ -283,7 +282,7 @@ def test_moss_tts_realtime_english(omni_runner: OmniRunner, ref_audio_path: str)
     assert not torch.all(audio == 0), "Audio is silence"
 
 
-@hardware_test(res={"cuda": "L4"}, num_cards=1)
+@hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards=1)
 def test_moss_tts_realtime_chinese(omni_runner: OmniRunner, ref_audio_path: str) -> None:
     """MossTTSRealtime: Chinese input produces non-empty audio."""
     req = _build_request(ref_audio_path, "你好，这是语音合成测试。")

tests/e2e/offline_inference/test_moss_tts_realtime_expansion.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 """E2E offline inference tests for MOSS-TTS-Realtime (MossTTSRealtime, 1.7B).
 
 Exercises the local-transformer generation path on L4. Delay-model coverage is in
@@ -33,7 +33,6 @@ _OMNI_RUNNER_PARAM = (
 pytestmark = [
     pytest.mark.slow,
     pytest.mark.tts,
-    pytest.mark.skip(reason="https://github.com/vllm-project/vllm-omni/issues/4700"),
     pytest.mark.parametrize("omni_runner", [_OMNI_RUNNER_PARAM], indirect=True),
 ]
 
@@ -118,7 +117,7 @@ def _collect_audio(omni_runner: OmniRunner, request: dict) -> tuple[torch.Tensor
     raise AssertionError("No stage outputs received")
 
 
-@hardware_test(res={"cuda": "L4"})
+@hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards=1)
 def test_moss_tts_realtime_english(omni_runner: OmniRunner, ref_audio_path) -> None:
     """MossTTSRealtime: English voice_clone produces non-empty 24 kHz audio."""
     req = _build_request("This is a real-time TTS streaming test.", ref_audio_path)

tests/e2e/online_serving/test_moss_tts_expansion.py

@@ -30,20 +30,16 @@ from tests.helpers.mark import hardware_test
 from tests.helpers.runtime import OmniServerParams
 from tests.helpers.stage_config import get_deploy_config_path
 
-_SKIP_ISSUE_6417 = pytest.mark.skip(
-    reason="https://github.com/vllm-project/vllm-omni/issues/6417",
-)
-
 pytestmark = [
+    pytest.mark.advanced_model,
     pytest.mark.slow,
     pytest.mark.tts,
-    _SKIP_ISSUE_6417,
 ]
 
 MODEL = "OpenMOSS-Team/MOSS-TTS-Realtime"
 REF_AUDIO_URL = "https://raw.githubusercontent.com/OpenMOSS/MOSS-TTS/HEAD/assets/audio/reference_zh_1.wav"
-# Voice-clone output is not reliably transcribed by the Whisper check used at
-# full_model run_level; assert non-trivial WAV payload size instead.
+# Reject a response that contains only a header or a very short audio payload.
+# Advanced and full runs also validate the transcript.
 _MIN_AUDIO_BYTES = 10_000
 
 
@@ -77,7 +73,9 @@ def get_prompt(prompt_type: str = "text") -> str:
     """
     prompts = {
         "text": "Hello, this is a real-time voice cloning demo for testing.",
-        "chinese": "你好，这是一段实时语音合成测试。",
+        # Match the offline test prompt. Whisper-small can transcribe "一段"
+        # as "一本", which makes this short sample fail the similarity check.
+        "chinese": "你好，这是语音合成测试。",
     }
     return prompts.get(prompt_type, prompts["text"])
 
@@ -94,7 +92,7 @@ tts_server_params = [
 ]
 
 
-@hardware_test(res={"cuda": "L4"}, num_cards=1)
+@hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards=1)
 @pytest.mark.parametrize("omni_server", tts_server_params, indirect=True)
 def test_text_to_audio_001(omni_server, online_client, ref_audio_data_url) -> None:
     """
@@ -105,9 +103,8 @@ def test_text_to_audio_001(omni_server, online_client, ref_audio_data_url) -> No
     Input Setting: stream=False
     Datasets: single request
 
-    NOTE: ``min_audio_bytes`` skips Whisper transcript similarity — cloned
-    voice output is often empty or mismatched under ASR without indicating a
-    real serving regression.
+    ``min_audio_bytes`` adds a payload-size check. It does not disable the
+    transcript check used by advanced and full runs.
     """
     request_config = {
         "model": omni_server.model,
@@ -121,7 +118,7 @@ def test_text_to_audio_001(omni_server, online_client, ref_audio_data_url) -> No
     online_client.send_audio_speech_request(request_config)
 
 
-@hardware_test(res={"cuda": "L4"}, num_cards=1)
+@hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards=1)
 @pytest.mark.parametrize("omni_server", tts_server_params, indirect=True)
 def test_text_to_audio_002_streaming(omni_server, online_client, ref_audio_data_url) -> None:
     """
@@ -150,7 +147,7 @@ def test_text_to_audio_002_streaming(omni_server, online_client, ref_audio_data_
     online_client.send_audio_speech_request(request_config)
 
 
-@hardware_test(res={"cuda": "L4"}, num_cards=1)
+@hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards=1)
 @pytest.mark.parametrize("omni_server", tts_server_params, indirect=True)
 def test_text_to_audio_003_chinese(omni_server, online_client, ref_audio_data_url) -> None:
     """
@@ -161,7 +158,8 @@ def test_text_to_audio_003_chinese(omni_server, online_client, ref_audio_data_ur
     Input Setting: stream=False
     Datasets: single request
 
-    NOTE: same ``min_audio_bytes`` rationale as test_text_to_audio_001.
+    ``min_audio_bytes`` adds a payload-size check. It does not disable the
+    transcript check used by advanced and full runs.
     """
     request_config = {
         "model": omni_server.model,

tests/entrypoints/openai_api/test_moss_tts_realtime_serving.py

@@ -0,0 +1,302 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+import asyncio
+import threading
+from types import SimpleNamespace
+from unittest.mock import AsyncMock
+
+import pytest
+import torch
+
+from vllm_omni.entrypoints.openai.protocol.audio import OpenAICreateSpeechRequest
+from vllm_omni.entrypoints.openai.tts_adapters import moss_tts as adapter_module
+from vllm_omni.entrypoints.openai.tts_adapters.base import SpeechServingContext
+from vllm_omni.entrypoints.openai.tts_adapters.moss_tts import MossTTSAdapter
+from vllm_omni.model_executor.models.moss_tts import reference_encoder
+from vllm_omni.utils.speaker_cache import SpeakerEmbeddingCache
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+@pytest.mark.parametrize(
+    ("hf_config", "expected_codec_path"),
+    [
+        (SimpleNamespace(), "OpenMOSS-Team/MOSS-Audio-Tokenizer"),
+        (
+            SimpleNamespace(codec_model_name_or_path="/models/custom-moss-codec"),
+            "/models/custom-moss-codec",
+        ),
+    ],
+    ids=["default-codec", "configured-codec"],
+)
+def test_realtime_components_use_the_realtime_model_and_codec(
+    monkeypatch: pytest.MonkeyPatch,
+    hf_config: SimpleNamespace,
+    expected_codec_path: str,
+) -> None:
+    engine_client = SimpleNamespace(
+        model_config=SimpleNamespace(
+            model="OpenMOSS-Team/MOSS-TTS-Realtime",
+            hf_config=hf_config,
+        )
+    )
+
+    server = MossTTSAdapter(SpeechServingContext(server=object(), engine_client=engine_client))
+
+    tokenizer = object()
+    codec = type("Codec", (), {"to": lambda self, device: self, "eval": lambda self: self})()
+    processor_calls = []
+
+    class Processor:
+        def __init__(self, *, tokenizer):
+            processor_calls.append(tokenizer)
+
+    class_calls = []
+    tokenizer_calls = []
+    codec_calls = []
+
+    def fake_get_class(class_reference, model_id):
+        class_calls.append((class_reference, model_id))
+        return Processor
+
+    def fake_load_tokenizer(model_id, *, trust_remote_code):
+        tokenizer_calls.append((model_id, trust_remote_code))
+        return tokenizer
+
+    def fake_load_codec(model_id, *, trust_remote_code):
+        codec_calls.append((model_id, trust_remote_code))
+        return codec
+
+    monkeypatch.setattr(adapter_module, "get_class_from_dynamic_module", fake_get_class)
+    monkeypatch.setattr(adapter_module.AutoTokenizer, "from_pretrained", fake_load_tokenizer)
+    monkeypatch.setattr(adapter_module.AutoModel, "from_pretrained", fake_load_codec)
+
+    components = server._get_moss_realtime_components()
+
+    assert components[0] is tokenizer
+    assert components[2] is codec
+    assert server._get_moss_realtime_components() is components
+    assert class_calls == [
+        (
+            "processing_mossttsrealtime.MossTTSRealtimeProcessor",
+            "OpenMOSS-Team/MOSS-TTS-Realtime",
+        )
+    ]
+    assert tokenizer_calls == [("OpenMOSS-Team/MOSS-TTS-Realtime", True)]
+    assert codec_calls == [(expected_codec_path, True)]
+    assert processor_calls == [tokenizer]
+
+
+def test_realtime_serving_builds_the_talker_prompt(monkeypatch: pytest.MonkeyPatch) -> None:
+    server = MossTTSAdapter(SpeechServingContext(server=SimpleNamespace(uploaded_speakers={})))
+    server._moss_variant = "realtime"
+    tokenizer = object()
+    processor = object()
+    codec = object()
+    request_thread = threading.get_ident()
+
+    def load_components():
+        assert threading.get_ident() != request_thread
+        return tokenizer, processor, codec
+
+    server._get_moss_realtime_components = load_components
+
+    reference_codes = torch.arange(64, dtype=torch.int64).reshape(4, 16)
+    encode = AsyncMock(return_value=([reference_codes], {0: "reference-cache-key"}))
+    monkeypatch.setattr(server, "_encode_moss_references", encode)
+
+    build_call = None
+
+    def fake_build(actual_tokenizer, actual_processor, text, actual_codes):
+        nonlocal build_call
+        build_call = (actual_tokenizer, actual_processor, text, actual_codes)
+        return {
+            "prompt_token_ids": [10, 11],
+            "codes": {"ref": actual_codes},
+            "ids": {"all": [12]},
+        }
+
+    monkeypatch.setattr(adapter_module, "build_realtime_prompt", fake_build)
+
+    params = asyncio.run(
+        server._build_moss_tts_params(
+            OpenAICreateSpeechRequest(
+                input="speak this text",
+                ref_audio="data:audio/wav;base64,AAAA",
+                max_new_tokens=50,
+            )
+        )
+    )
+
+    encode.assert_awaited_once()
+    assert encode.call_args.kwargs == {"has_inline_ref_audio": False, "two_speaker": False}
+    assert encode.call_args.args[0].ref_audio == "data:audio/wav;base64,AAAA"
+    assert build_call == (tokenizer, processor, "speak this text", reference_codes)
+    assert params == {
+        "prompt_token_ids": [10, 11],
+        "codes": {"ref": reference_codes},
+        "ids": {"all": [12]},
+        "max_new_frames": [50],
+        "ref_audio_cache_key": "reference-cache-key",
+    }
+    assert "prompt_audio_array" not in params
+
+
+def test_cancelled_component_waiter_does_not_duplicate_codec_load(monkeypatch):
+    adapter = MossTTSAdapter(
+        SpeechServingContext(
+            server=object(),
+            engine_client=SimpleNamespace(
+                model_config=SimpleNamespace(model="OpenMOSS-Team/MOSS-TTS-Realtime", hf_config=SimpleNamespace())
+            ),
+        )
+    )
+    started = threading.Event()
+    release = threading.Event()
+    second_started = threading.Event()
+    loads = []
+    codec = type("Codec", (), {"to": lambda self, device: self, "eval": lambda self: self})()
+
+    def load_codec(*args, **kwargs):
+        loads.append(args)
+        started.set()
+        assert release.wait(5), "codec loader was not released"
+        return codec
+
+    monkeypatch.setattr(adapter_module, "get_class_from_dynamic_module", lambda *args: lambda **kwargs: object())
+    monkeypatch.setattr(adapter_module.AutoTokenizer, "from_pretrained", lambda *args, **kwargs: object())
+    monkeypatch.setattr(adapter_module.AutoModel, "from_pretrained", load_codec)
+
+    def second_load():
+        second_started.set()
+        return adapter._get_moss_realtime_components()
+
+    async def run():
+        first = asyncio.create_task(asyncio.to_thread(adapter._get_moss_realtime_components))
+        try:
+            assert await asyncio.to_thread(started.wait, 5)
+            first.cancel()
+            with pytest.raises(asyncio.CancelledError):
+                await first
+            second = asyncio.create_task(asyncio.to_thread(second_load))
+            try:
+                assert await asyncio.to_thread(second_started.wait, 5)
+            finally:
+                release.set()
+            components = await asyncio.wait_for(second, 5)
+            assert components[2] is codec
+            assert adapter._get_moss_realtime_components() is components
+            assert len(loads) == 1
+        finally:
+            release.set()
+
+    asyncio.run(run())
+
+
+@pytest.mark.parametrize("voice_name", [None, "speaker"])
+def test_realtime_reference_cache_reuses_codes(voice_name):
+    cache = SpeakerEmbeddingCache()
+    resolve = AsyncMock(return_value=([0.0], 24000, "audio-key"))
+    request_thread = threading.get_ident()
+    encode_calls = []
+    expected = torch.arange(64).reshape(4, 16)
+
+    def batch_encode(wavs, *, num_quantizers):
+        assert threading.get_ident() != request_thread
+        assert num_quantizers == 16
+        encode_calls.append(wavs)
+        return SimpleNamespace(audio_codes=expected.T.unsqueeze(1), audio_codes_lengths=torch.tensor([4]))
+
+    encoder = reference_encoder.build_reference_encoder(
+        SimpleNamespace(batch_encode=batch_encode), variant="realtime", speaker_cache=cache
+    )
+
+    async def run():
+        kwargs = dict(
+            resolve_ref_audio=resolve,
+            get_artifact_key=lambda key: "content-key",
+            voice_name=voice_name,
+            voice_created_at=1,
+        )
+        try:
+            first, first_key = await encoder.encode("reference", **kwargs)
+            second, second_key = await encoder.encode("reference", **kwargs)
+            assert torch.equal(first, expected)
+            assert torch.equal(second, expected)
+            assert first.dtype == torch.int64
+            assert first_key == "audio-key"
+            assert second_key == (None if voice_name else "audio-key")
+            second.zero_()
+            third, _ = await encoder.encode("another-reference", **kwargs)
+            assert torch.equal(third, expected)
+            assert len(encode_calls) == 1
+            assert resolve.await_count == (1 if voice_name else 3)
+            key = cache.make_cache_key(
+                voice_name or "ref:content-key", "moss_tts_realtime_nq16", 1 if voice_name else 0
+            )
+            assert cache.get(key)["codes"].dtype == torch.int32
+            if voice_name:
+                kwargs["voice_created_at"] = 2
+                await encoder.encode("reference", **kwargs)
+                assert len(encode_calls) == 2
+                assert resolve.await_count == 2
+        finally:
+            await encoder.aclose()
+
+    asyncio.run(run())
+
+
+def test_realtime_reference_batch_trims_each_item():
+    cache = SpeakerEmbeddingCache()
+    codes = torch.arange(16 * 2 * 5).reshape(16, 2, 5)
+    calls = []
+
+    def batch_encode(wavs, *, num_quantizers):
+        calls.append(wavs)
+        assert num_quantizers == 16
+        assert [wav.shape for wav in wavs] == [torch.Size([8]), torch.Size([12])]
+        return SimpleNamespace(audio_codes=codes, audio_codes_lengths=torch.tensor([3, 5]))
+
+    encoder = reference_encoder.build_reference_encoder(
+        SimpleNamespace(batch_encode=batch_encode), variant="realtime", speaker_cache=cache
+    )
+
+    async def resolve(ref):
+        return [0.0] * (8 if ref == "short" else 12), 24000, ref
+
+    async def run():
+        try:
+            short, long = await asyncio.gather(
+                encoder.encode("short", resolve_ref_audio=resolve, get_artifact_key=lambda key: key),
+                encoder.encode("long", resolve_ref_audio=resolve, get_artifact_key=lambda key: key),
+            )
+            torch.testing.assert_close(short[0], codes[:, 0, :3].T)
+            torch.testing.assert_close(long[0], codes[:, 1, :5].T)
+            assert len(calls) == 1
+        finally:
+            await encoder.aclose()
+
+    asyncio.run(run())
+
+
+def test_failed_realtime_reference_encoding_does_not_populate_cache():
+    cache = SpeakerEmbeddingCache()
+    resolve = AsyncMock(return_value=([0.0], 24000, "audio-key"))
+
+    def fail(*args, **kwargs):
+        raise ValueError("invalid reference")
+
+    encoder = reference_encoder.build_reference_encoder(
+        SimpleNamespace(batch_encode=fail), variant="realtime", speaker_cache=cache
+    )
+
+    async def run():
+        try:
+            with pytest.raises(ValueError, match="invalid reference"):
+                await encoder.encode("reference", resolve_ref_audio=resolve, get_artifact_key=lambda key: None)
+            assert cache.memory_bytes() == 0
+        finally:
+            await encoder.aclose()
+
+    asyncio.run(run())

tests/model_executor/models/moss_tts/test_audio_tokenizer_v2.py

@@ -0,0 +1,54 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Regression tests for MOSS audio tokenizer v2 projection modules."""
+
+from __future__ import annotations
+
+import pytest
+import torch
+import torch.nn as nn
+
+from vllm_omni.model_executor.models.moss_tts.audio_tokenizer_v2 import (
+    MossAudioTokenizerProjectedTransformer,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def test_projected_transformer_keeps_learned_projections_when_dimensions_match() -> None:
+    module = MossAudioTokenizerProjectedTransformer(
+        input_dimension=8,
+        output_dimension=8,
+        d_model=8,
+        module_type="transformer",
+        num_heads=1,
+        num_layers=0,
+        positional_embedding="rope",
+    )
+
+    assert isinstance(module.input_proj, nn.Linear)
+    assert isinstance(module.output_proj, nn.Linear)
+    assert module.input_proj.weight.shape == (8, 8)
+    assert module.output_proj.weight.shape == (8, 8)
+
+    x = torch.randn(2, 8, 5)
+    lengths = torch.tensor([5, 3], dtype=torch.long)
+    output, output_lengths = module(x, lengths)
+
+    assert output.shape == x.shape
+    torch.testing.assert_close(output_lengths, lengths)
+
+
+def test_projected_transformer_keeps_linear_projections_when_dimensions_differ() -> None:
+    module = MossAudioTokenizerProjectedTransformer(
+        input_dimension=8,
+        output_dimension=6,
+        d_model=10,
+        module_type="transformer",
+        num_heads=1,
+        num_layers=0,
+        positional_embedding="rope",
+    )
+
+    assert isinstance(module.input_proj, nn.Linear)
+    assert isinstance(module.output_proj, nn.Linear)

tests/model_executor/models/moss_tts/test_realtime_chunked_prefill.py

@@ -0,0 +1,86 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from types import SimpleNamespace
+from typing import Any
+
+import pytest
+import torch
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def test_realtime_talker_preserves_state_across_chunked_prefill() -> None:
+    from vllm_omni.model_executor.models.moss_tts.modeling_moss_tts_talker import (
+        MossTTSRealtimeTalkerForGeneration,
+    )
+
+    talker = MossTTSRealtimeTalkerForGeneration.__new__(MossTTSRealtimeTalkerForGeneration)
+    talker.n_vq = 2
+    talker.audio_pad_token = 1024
+    talker.text_vocab_size = 16
+    talker.text_pad_id = 2
+    talker._batch_state = None
+    talker._batch_state_spans = None
+
+    embedded_audio: list[torch.Tensor | None] = []
+
+    def build_input_embeds(input_ids, audio_codes):
+        embedded_audio.append(audio_codes.clone() if audio_codes is not None else None)
+        return torch.zeros((input_ids.shape[0], 1))
+
+    generated: list[None] = []
+
+    def generate_frame(*args, **kwargs):
+        generated.append(None)
+        return torch.ones((1, talker.n_vq), dtype=torch.long)
+
+    talker._build_input_embeds = build_input_embeds
+    talker.local_transformer = SimpleNamespace(generate_frame=generate_frame)
+    talker.local_lm_heads = []
+
+    ref_codes = torch.arange(10, dtype=torch.long).reshape(5, 2)
+    info: dict[str, Any] = {
+        "codes": {"ref": ref_codes},
+        "ids": {"all": [7, 8]},
+        "max_new_frames": [3],
+    }
+    chunks = [(0, 2, False), (2, 2, False), (4, 1, True)]
+
+    state: dict[str, Any] | None = None
+    for computed, length, eligible in chunks:
+        _, _, update = talker.preprocess(
+            torch.ones(length, dtype=torch.long),
+            None,
+            _omni_is_prefill=True,
+            _omni_num_computed_tokens=computed,
+            _omni_prompt_len=5,
+            **info,
+        )
+        info.update(update)
+        if state is None:
+            state = info["audio_state"]
+        else:
+            assert info["audio_state"] is state
+
+        output = talker.make_omni_output(
+            torch.ones((length, 1)),
+            model_intermediate_buffer=[info],
+            request_token_spans=[(0, length)],
+            request_sample_eligible=[eligible],
+        )
+        if not eligible:
+            assert output.multimodal_outputs == {}
+        logits = talker.compute_logits(torch.ones((1, 1)))
+        assert logits.argmax().item() == (7 if eligible else talker.text_pad_id)
+
+    assert state is not None
+    assert info["ref_offset"] == 5
+    assert state["step"] == 1
+    assert state["text_cursor"] == 1
+    assert state["remaining_text"] == [7, 8]
+    assert len(generated) == 1
+    assert output.multimodal_outputs["codes"]["audio"][0].shape == (1, 2)
+    for actual, expected in zip(embedded_audio, ref_codes.split((2, 2, 1)), strict=True):
+        assert actual is not None
+        torch.testing.assert_close(actual, expected)

tests/model_executor/models/moss_tts/test_realtime_generation_limit.py

@@ -0,0 +1,53 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def test_realtime_talker_stops_at_max_new_frames() -> None:
+    from vllm_omni.model_executor.models.moss_tts.modeling_moss_tts_talker import (
+        MossTTSRealtimeTalkerForGeneration,
+    )
+
+    talker = MossTTSRealtimeTalkerForGeneration.__new__(MossTTSRealtimeTalkerForGeneration)
+    talker.n_vq = 4
+    talker.audio_pad_token = 1024
+    talker.text_vocab_size = 8
+    talker.text_pad_id = 2
+    talker.audio_eos_id = 3
+    talker._batch_state = None
+    talker._build_input_embeds = lambda input_ids, audio_codes: torch.zeros((input_ids.shape[0], 1))
+
+    generated: list[None] = []
+
+    def generate_frame(*args, **kwargs):
+        generated.append(None)
+        return torch.ones((1, talker.n_vq), dtype=torch.long)
+
+    talker.local_transformer = SimpleNamespace(generate_frame=generate_frame)
+    talker.local_lm_heads = []
+
+    _, _, info = talker.preprocess(
+        torch.tensor([1, 2]),
+        None,
+        max_new_frames=[2],
+    )
+    assert info["audio_state"]["max_new_frames"] == 2
+
+    for _ in range(3):
+        output = talker.make_omni_output(
+            torch.ones((1, 1)),
+            runtime_additional_information=[info],
+            request_token_spans=[(0, 1)],
+        )
+
+    assert len(generated) == 2
+    assert info["audio_state"]["step"] == 2
+    assert info["audio_state"]["is_stopping"] is True
+    assert output.multimodal_outputs == {}
+    assert talker.compute_logits(torch.ones((1, 1))).argmax().item() == talker.audio_eos_id

tests/model_executor/models/moss_tts/test_realtime_prompt.py

@@ -0,0 +1,78 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+import numpy as np
+import pytest
+import torch
+
+from vllm_omni.model_executor.models.moss_tts.realtime_prompt import build_realtime_prompt
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+class _Tokenizer:
+    def encode(self, text: str, add_special_tokens: bool = False) -> list[int]:
+        del add_special_tokens
+        if text == "<|im_start|>assistant\n":
+            return [20, 21]
+        return list(range(100, 100 + len(text.split())))
+
+
+class _Processor:
+    audio_channel_pad = 1024
+    audio_bos_token = 1025
+
+    def __init__(self) -> None:
+        self.reference_codes = None
+
+    def make_ensemble(self, prompt_audio_tokens):
+        self.reference_codes = prompt_audio_tokens
+        return np.full((3, 17), self.audio_channel_pad, dtype=np.int64)
+
+
+def test_build_realtime_prompt_aligns_reference_and_prefill_text() -> None:
+    processor = _Processor()
+    reference_codes = torch.arange(80, dtype=torch.int64).reshape(5, 16)
+
+    params = build_realtime_prompt(
+        _Tokenizer(),
+        processor,
+        "one two three",
+        reference_codes,
+    )
+
+    assert np.array_equal(processor.reference_codes, reference_codes.numpy())
+    assert params["prompt_token_ids"] == [1024, 1024, 1024, 20, 21, 100, 101, 102]
+    audio_grid = params["codes"]["ref"]
+    assert audio_grid.shape == (len(params["prompt_token_ids"]), 16)
+    assert audio_grid[-1, 0].item() == processor.audio_bos_token
+    assert "ids" not in params
+
+
+def test_build_realtime_prompt_streams_text_after_first_twelve_tokens() -> None:
+    params = build_realtime_prompt(
+        _Tokenizer(),
+        _Processor(),
+        " ".join(f"word{i}" for i in range(15)),
+        torch.zeros((2, 16), dtype=torch.int64),
+    )
+
+    assert params["prompt_token_ids"][-12:] == list(range(100, 112))
+    assert params["ids"] == {"all": [112, 113, 114]}
+
+
+def test_build_realtime_prompt_rejects_wrong_processor_layout(monkeypatch: pytest.MonkeyPatch) -> None:
+    processor = _Processor()
+    monkeypatch.setattr(
+        processor,
+        "make_ensemble",
+        lambda prompt_audio_tokens: np.zeros((3, 16), dtype=np.int64),
+    )
+
+    with pytest.raises(ValueError, match=r"expected \(sequence_length, 17\)"):
+        build_realtime_prompt(
+            _Tokenizer(),
+            processor,
+            "one",
+            torch.zeros((2, 16), dtype=torch.int64),
+        )

vllm_omni/entrypoints/openai/tts_adapters/moss_tts.py

@@ -6,8 +6,12 @@ Both variants share the same build/validate flow (``_build_moss_tts_params``
 handles each); they are registered under distinct model-type names.
 """
 
+import asyncio
+import threading
 from typing import TYPE_CHECKING, Any, cast
 
+from transformers import AutoModel, AutoTokenizer
+from transformers.dynamic_module_utils import get_class_from_dynamic_module
 from vllm.inputs import tokens_input
 
 from vllm_omni.entrypoints.openai.tts_adapters import register_tts_adapter
@@ -18,6 +22,7 @@ from vllm_omni.entrypoints.openai.tts_adapters.base import (
     apply_max_new_tokens,
     conditioning_cache_salt,
 )
+from vllm_omni.model_executor.models.moss_tts.realtime_prompt import build_realtime_prompt
 
 if TYPE_CHECKING:
     from vllm_omni.entrypoints.openai.protocol.audio import OpenAICreateSpeechRequest
@@ -30,6 +35,7 @@ class _MossTTSAdapterBase(ARTTSAdapter):
         super().__init__(ctx)
         self._moss_variant = None if self.name == "moss_tts_nano" else self._detect_moss_variant()
         self._moss_processor_cache = None
+        self._moss_realtime_components_lock = threading.Lock()
 
     @property
     def engine_client(self):
@@ -63,7 +69,7 @@ class _MossTTSAdapterBase(ARTTSAdapter):
         # inside the model package; this layer only supplies the processor and
         # the process-wide speaker cache.
         encoder = build_reference_encoder(
-            self._get_moss_processor(),
+            self._get_moss_realtime_components()[2] if self._moss_variant == "realtime" else self._get_moss_processor(),
             variant=cast(str, self._moss_variant),
             speaker_cache=self._speaker_cache,
         )
@@ -157,6 +163,45 @@ class _MossTTSAdapterBase(ARTTSAdapter):
         self._moss_processor_cache = proc
         return proc
 
+    def _get_moss_realtime_components(self):
+        with self._moss_realtime_components_lock:
+            cached = getattr(self, "_moss_realtime_components", None)
+            if cached is not None:
+                return cached
+
+            model_id = self.engine_client.model_config.model
+            processor_cls = get_class_from_dynamic_module(
+                "processing_mossttsrealtime.MossTTSRealtimeProcessor",
+                model_id,
+            )
+            tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
+            processor = processor_cls(tokenizer=tokenizer)
+            hf_config = self.engine_client.model_config.hf_config
+            codec_path = str(
+                getattr(
+                    hf_config,
+                    "codec_model_name_or_path",
+                    getattr(
+                        hf_config,
+                        "audio_tokenizer_name_or_path",
+                        "OpenMOSS-Team/MOSS-Audio-Tokenizer",
+                    ),
+                )
+            )
+            codec = (
+                AutoModel.from_pretrained(
+                    codec_path,
+                    trust_remote_code=True,
+                )
+                # Reference encoding runs in the API process. Keep the codec on CPU
+                # so it does not reserve accelerator memory outside the model workers.
+                .to("cpu")
+                .eval()
+            )
+            cached = (tokenizer, processor, codec)
+            self._moss_realtime_components = cached
+            return cached
+
     async def _build_moss_tts_params(
         self,
         request: "OpenAICreateSpeechRequest",
@@ -196,24 +241,18 @@ class _MossTTSAdapterBase(ARTTSAdapter):
             params["ref_audio_cache_key"] = cache_key
             return params
 
-        # ---- MOSS-TTS-Realtime: keep the old prompt_audio_array path ----
-        # ``AutoProcessor.from_pretrained`` doesn't auto-discover
-        # ``MossTTSRealtimeProcessor`` (no ``processor_config.json`` in the
-        # snapshot), and Realtime's prompt format diverges from MossTTSDelay
-        # (16-channel grid, separate per-step text feed). The
-        # ``prompt_audio_array`` shape lines up well enough with what the
-        # talker reads for short prompts; full Realtime support needs a
-        # separate processor.from_module path which we don't wire here.
         if v == "realtime":
-            params = {
-                "text": [request.input or ""],
-                "mode": ["voice_clone"],
-            }
+            tokenizer, processor, _ = await asyncio.to_thread(self._get_moss_realtime_components)
+            references, realtime_resolve_keys = await self._encode_moss_references(
+                request,
+                has_inline_ref_audio=has_inline_ref_audio,
+                two_speaker=False,
+            )
+            params = build_realtime_prompt(tokenizer, processor, request.input or "", references[0])
             if request.max_new_tokens is not None:
                 params["max_new_frames"] = [request.max_new_tokens]
-            wav_list, sr, cache_key = await self._resolve_ref_audio(cast(str, request.ref_audio))
-            params["prompt_audio_array"] = [[wav_list, sr]]
-            params["ref_audio_cache_key"] = cache_key
+            if 0 in realtime_resolve_keys:
+                params["ref_audio_cache_key"] = realtime_resolve_keys[0]
             return params
 
         # ---- MossTTSDelay family (tts/ttsd/sound_effect/voice_generator)

vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_talker.py

@@ -44,10 +44,10 @@ def _maybe_prefix(prefix: str, name: str) -> str:
 
 
 def _iter_state_row_spans(
-    states: Sequence[dict[str, Any]],
+    states: Sequence[dict[str, Any] | None],
     spans: Sequence[tuple[int, int]] | None,
     num_rows: int,
-) -> Iterable[tuple[dict[str, Any], int, int]]:
+) -> Iterable[tuple[dict[str, Any] | None, int, int]]:
     """Yield state/logit-row ranges without assuming equal request row counts."""
     if num_rows <= 0 or not states:
         return
@@ -798,6 +798,9 @@ class MossTTSRealtimeTalkerForGeneration(nn.Module):
     have_multimodal_outputs: bool = True
     has_preprocess: bool = True
     has_postprocess: bool = True
+    # Intermediate prefill chunks reach make_omni_output but must not produce
+    # audio until the scheduler has completed the prompt.
+    requires_request_sample_eligibility: bool = True
 
     AUDIO_BOS = 1025
     AUDIO_EOS = 1026
@@ -858,7 +861,7 @@ class MossTTSRealtimeTalkerForGeneration(nn.Module):
         # We still need a logits_processor so vLLM downstream stays happy with
         # something callable; a minimal pass-through suffices.
         self.logits_processor = LogitsProcessor(self.text_vocab_size)
-        self._batch_state: list[dict[str, Any]] | None = None
+        self._batch_state: list[dict[str, Any] | None] | None = None
 
         # Stacked weight cache built after load_weights().
         # shape: (n_vq, audio_vocab_size, hidden_size)
@@ -972,13 +975,14 @@ class MossTTSRealtimeTalkerForGeneration(nn.Module):
         span_len = int(input_ids.shape[0])
         audio_state = info_dict.get("audio_state")
         is_first_call = not isinstance(audio_state, dict)
+        is_prefill = bool(info_dict.get("_omni_is_prefill", span_len > 1))
 
-        if span_len > 1 or is_first_call:
+        if is_prefill or is_first_call:
             # Prefill: read the full reference-audio code grid (channels x T)
             # from ``info_dict["codes"]["ref"]``; slice the chunk that aligns
             # with this prefill window.
             ref_codes = (info_dict.get("codes", {}) or {}).get("ref")
-            ref_offset = int(info_dict.get("ref_offset", 0))
+            ref_offset = int(info_dict.get("_omni_num_computed_tokens", info_dict.get("ref_offset", 0)))
             chunk_audio = None
             if isinstance(ref_codes, torch.Tensor) and ref_codes.numel() > 0:
                 if ref_codes.dim() == 1 and ref_codes.numel() % self.n_vq == 0:
@@ -994,22 +998,24 @@ class MossTTSRealtimeTalkerForGeneration(nn.Module):
                 if chunk_audio is not None
                 else torch.full((self.n_vq,), self.audio_pad_token, dtype=torch.long, device=device)
             )
-            # Capture the streaming-text list from the request (set by the
-            # realtime end2end builder; empty for the delay variants which
-            # don't take this code path anyway).
-            remaining = (info_dict.get("ids", {}) or {}).get("all")
-            if not isinstance(remaining, list):
-                remaining = []
             info_update: dict[str, Any] = {
-                "audio_state": {
+                "audio_codes": {"current": current_codes},
+                "ref_offset": ref_offset + span_len,
+            }
+            if is_first_call:
+                remaining = (info_dict.get("ids", {}) or {}).get("all")
+                if not isinstance(remaining, list):
+                    remaining = []
+                max_new_frames = info_dict.get("max_new_frames")
+                if isinstance(max_new_frames, (list, tuple)):
+                    max_new_frames = max_new_frames[0] if max_new_frames else None
+                info_update["audio_state"] = {
                     "is_stopping": False,
                     "step": 0,
                     "text_cursor": 0,
                     "remaining_text": list(remaining),
-                },
-                "audio_codes": {"current": current_codes},
-                "ref_offset": ref_offset + span_len,
-            }
+                    "max_new_frames": int(max_new_frames) if max_new_frames is not None else -1,
+                }
             return input_ids, embeds, info_update
 
         # Decode step: text_token from the just-sampled vLLM logit, plus the
@@ -1046,13 +1052,21 @@ class MossTTSRealtimeTalkerForGeneration(nn.Module):
         info_dicts: list[dict[str, Any]] = (
             kwargs.get("model_intermediate_buffer") or kwargs.get("runtime_additional_information") or []
         )
+        sample_eligible = kwargs.get("request_sample_eligible")
+        eligible = [
+            sample_eligible is None or i >= len(sample_eligible) or bool(sample_eligible[i])
+            for i in range(len(info_dicts))
+        ]
 
         # Defensive state init.
         for info in info_dicts:
             if isinstance(info, dict) and not isinstance(info.get("audio_state"), dict):
                 info["audio_state"] = {"is_stopping": False, "step": 0}
 
-        self._batch_state = [(info["audio_state"] if isinstance(info, dict) else {}) for info in info_dicts]
+        self._batch_state = [
+            (info["audio_state"] if isinstance(info, dict) and eligible[i] else None)
+            for i, info in enumerate(info_dicts)
+        ]
         self._batch_state_spans = kwargs.get("request_token_spans")
 
         # See delay talker: real per-request row spans from the runner are
@@ -1074,6 +1088,8 @@ class MossTTSRealtimeTalkerForGeneration(nn.Module):
             for i, info in enumerate(info_dicts):
                 if not isinstance(info, dict):
                     continue
+                if not eligible[i]:
+                    continue
                 row_start, row_end = spans[i]
                 row_end = min(int(row_end), num_rows)
                 if row_start >= row_end:
@@ -1115,10 +1131,13 @@ class MossTTSRealtimeTalkerForGeneration(nn.Module):
                     )
 
                 ch0 = int(new_codes[0].item())
+                state["step"] = int(state.get("step", 0)) + 1
+                max_new_frames = int(state.get("max_new_frames", -1))
+                if max_new_frames > 0 and state["step"] >= max_new_frames:
+                    state["is_stopping"] = True
                 # Stop condition mirrors upstream: codebook 0 == eos_audio_id.
                 if ch0 == self.AUDIO_EOS:
                     state["is_stopping"] = True
-                    state["step"] = int(state.get("step", 0)) + 1
                     info["audio_codes"] = {
                         "current": new_codes,
                         "accumulated": (info.get("audio_codes", {}) or {}).get("accumulated"),
@@ -1127,7 +1146,6 @@ class MossTTSRealtimeTalkerForGeneration(nn.Module):
 
                 if ch0 in (self.AUDIO_BOS, self.audio_pad_token):
                     # Skip the bos / pad frames — they don't decode to real audio.
-                    state["step"] = int(state.get("step", 0)) + 1
                     info["audio_codes"] = {
                         "current": new_codes,
                         "accumulated": (info.get("audio_codes", {}) or {}).get("accumulated"),
@@ -1141,7 +1159,6 @@ class MossTTSRealtimeTalkerForGeneration(nn.Module):
                     updated_acc = new_codes.unsqueeze(0)
 
                 info["audio_codes"] = {"current": new_codes, "accumulated": updated_acc}
-                state["step"] = int(state.get("step", 0)) + 1
                 per_req_codes[i] = new_codes.unsqueeze(0)
                 have_codes = True
 

vllm_omni/model_executor/models/moss_tts/realtime_prompt.py

@@ -0,0 +1,60 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""Prompt construction for MOSS-TTS-Realtime."""
+
+from typing import Any
+
+import numpy as np
+import torch
+
+REALTIME_AUDIO_CHANNELS = 16
+REALTIME_PREFILL_TEXT_TOKENS = 12
+
+
+def build_realtime_prompt(
+    tokenizer: Any,
+    processor: Any,
+    text: str,
+    reference_codes: torch.Tensor,
+) -> dict[str, Any]:
+    """Build the mixed text/audio prompt consumed by the Realtime talker."""
+    audio_tokens = reference_codes.detach().cpu().numpy()
+    system_grid = np.asarray(processor.make_ensemble(prompt_audio_tokens=audio_tokens), dtype=np.int64)
+    expected_width = REALTIME_AUDIO_CHANNELS + 1
+    if system_grid.ndim != 2 or system_grid.shape[1] != expected_width:
+        raise ValueError(
+            f"MOSS-TTS-Realtime processor returned shape {system_grid.shape}; "
+            f"expected (sequence_length, {expected_width})"
+        )
+
+    assistant_ids = tokenizer.encode("<|im_start|>assistant\n", add_special_tokens=False)
+    assistant_grid = np.full(
+        (len(assistant_ids), expected_width),
+        processor.audio_channel_pad,
+        dtype=np.int64,
+    )
+    assistant_grid[:, 0] = assistant_ids
+
+    text_ids = tokenizer.encode(text, add_special_tokens=False)
+    if not text_ids:
+        raise ValueError("MOSS-TTS-Realtime input is empty after tokenization")
+
+    prefill_len = min(len(text_ids), REALTIME_PREFILL_TEXT_TOKENS)
+    text_grid = np.full(
+        (prefill_len, expected_width),
+        processor.audio_channel_pad,
+        dtype=np.int64,
+    )
+    text_grid[:, 0] = text_ids[:prefill_len]
+    text_grid[-1, 1] = processor.audio_bos_token
+
+    prompt_grid = np.concatenate((system_grid, assistant_grid, text_grid), axis=0)
+    params: dict[str, Any] = {
+        "prompt_token_ids": prompt_grid[:, 0].tolist(),
+        "codes": {"ref": torch.from_numpy(prompt_grid[:, 1:].copy())},
+    }
+    remaining_text_ids = text_ids[prefill_len:]
+    if remaining_text_ids:
+        params["ids"] = {"all": remaining_text_ids}
+    return params

vllm_omni/model_executor/models/moss_tts/reference_encoder.py

@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 """Reference-audio encoding + speaker cache for the MOSS-TTS-family talker.
 
 This lives in the model package (not the shared serving layer) so all
@@ -360,6 +363,18 @@ class MossReferenceEncoder:
         await self._batcher.aclose()
 
 
+class MossRealtimeReferenceEncoder(MossReferenceEncoder):
+    """Use the realtime codec with the shared reference cache and batcher."""
+
+    def _encode_prepared(self, prepared: list[torch.Tensor]) -> list[torch.Tensor]:
+        with torch.no_grad():
+            encoded = self._processor.batch_encode([wav.squeeze(0) for wav in prepared], num_quantizers=self._n_vq)
+        return [
+            encoded.audio_codes[:, i, : int(length.item())].transpose(0, 1).contiguous()
+            for i, length in enumerate(encoded.audio_codes_lengths)
+        ]
+
+
 def build_reference_encoder(
     processor: Any,
     *,
@@ -369,9 +384,17 @@ def build_reference_encoder(
     """Build the per-server encoder for a MOSS-TTS ``variant``.
 
     Derives the encode geometry (``n_vq`` and the working sample rate) from the
-    upstream processor's ``model_config`` so the variant knowledge stays in the
-    model package rather than in ``serving_speech.py``.
+    upstream processor's ``model_config``; Realtime takes its codec directly
+    and uses 16 codebooks at 24 kHz. Variant knowledge stays in the model package.
     """
+    if variant == "realtime":
+        return MossRealtimeReferenceEncoder(
+            processor,
+            variant=variant,
+            n_vq=16,
+            sr_target=24000,
+            speaker_cache=speaker_cache,
+        )
     n_vq = int(getattr(processor.model_config, "n_vq", 32))
     # Local-v1.5 encodes reference audio at a fixed 24 kHz working rate
     # regardless of its 48 kHz stereo *output* codec -- mirrors the offline