[Model] Optimize MOSS-TTS Local batched execution and streaming codec

.buildkite/cuda/test-ready.yml

@@ -93,7 +93,7 @@ steps:
     commands:
       - |
         timeout 15m bash -c "
-                  pytest -sv tests/engine/ tests/model_executor/ -m 'core_model and cuda' --run-level 'core_model'
+                  pytest -sv tests/engine/ tests/model_executor/ tests/worker/test_batched_omni_output.py -m 'core_model and cuda' --run-level 'core_model'
         "
     mirror_hardwares: l4_1
 

examples/online_serving/text_to_speech/README.md

@@ -259,21 +259,25 @@ python examples/online_serving/text_to_speech/indextts2/speech_client.py \
 0.6B DualAR TTS at 44.1 kHz, 11 languages, zero-shot voice cloning.
 
 ### Prerequisites
+
 None beyond the base install: the neural audio codec is implemented in tree and
 its weights (`codec.pth`) ship with the checkpoint.
 
 ### Launch
+
 ```bash
 vllm serve Audio8/Audio8-TTS-Preview-0.6b --omni --port 8092
 # or:
 ./audio8_tts/run_server.sh
 ```
+
 The deploy config auto-loads from `vllm_omni/deploy/audio8_tts.yaml` (HF
 `model_type` is `arktts`). Do **not** pass `--trust-remote-code`: vllm-omni
 registers its own `arktts` config, and transformers would otherwise prefer the
 checkpoint's remote code and bypass it.
 
 ### Text-only synthesis
+
 ```bash
 curl -X POST http://localhost:8092/v1/audio/speech \
     -H "Content-Type: application/json" \
@@ -284,6 +288,7 @@ curl -X POST http://localhost:8092/v1/audio/speech \
 ```
 
 ### Voice cloning
+
 ```bash
 curl -X POST http://localhost:8092/v1/audio/speech \
     -H "Content-Type: application/json" \
@@ -293,18 +298,21 @@ curl -X POST http://localhost:8092/v1/audio/speech \
         "ref_text": "The exact transcript of the reference recording."
     }' --output cloned.wav
 ```
+
 `ref_audio` also accepts a `data:audio/wav;base64,...` URL. `ref_text` is
 mandatory whenever `ref_audio` is present and must match the recording.
 Uploading a voice via `POST /v1/audio/voices` lets the server reuse the encoded
 reference codes across requests (`voice: "<name>"`).
 
 ### Streaming
+
 ```bash
 python audio8_tts/speech_client.py --text "Welcome to Audio8 TTS." --stream --output out.pcm
 ffplay -f s16le -ar 44100 -ac 1 out.pcm
 ```
 
 ### Gradio demo
+
 ```bash
 ./audio8_tts/run_gradio_demo.sh          # server + demo
 python audio8_tts/gradio_demo.py --api-base http://localhost:8092   # demo only
@@ -328,7 +336,7 @@ these as a lower bound. 20 prompts from `seed_tts_smoke/en`, mean audio 4.5 s,
 deploy defaults:
 
 | | HF reference (batch=1) | vllm-omni c=1 | c=4 | c=8 |
-|---|---|---|---|---|
+| --- | --- | --- | --- | --- |
 | RTF (mean) | 1.008 | **0.19** | 0.30 | 0.54 |
 | Throughput (req/s) | 0.23 | 1.18 | 2.84 | **2.93** |
 | E2E latency (mean, ms) | 4375 | 850 | 1349 | 2361 |
@@ -347,7 +355,7 @@ job holding ~80% SM, no repeats), so read them as directional, not as measured
 speedups. Same setup, varying stage 0's `max_num_seqs` (stage 1 stays at 1):
 
 | stage-0 `max_num_seqs` | client concurrency | req/s | RTF | TTFP mean / p99 (ms) | underrun p99 (s) |
-|---|---|---|---|---|---|
+| --- | --- | --- | --- | --- | --- |
 | 4 (default) | 1 | 1.18 | 0.19 | 63 / 67 | 0.00 |
 | 4 | 4 | 2.84 | 0.30 | 105 / 148 | 0.00 |
 | 4 | 8 | 2.93 | 0.54 | 1124 / 1642 | 0.00 |
@@ -376,6 +384,7 @@ underrun. Raise it to 8 if first-packet latency under load matters more than a
 rare sub-300 ms gap.
 
 ### Notes
+
 - Output: 44.1 kHz mono; ~21.5 codec frames per second.
 - No built-in speaker presets. Omit `voice` for a random timbre, or clone one.
 - `max_new_tokens` caps generated codec frames (1 frame ~= 46 ms).
@@ -553,6 +562,46 @@ python examples/online_serving/text_to_speech/ming_flash_omni_tts/speech_client.
 
 ---
 
+## MOSS-TTS Local Transformer v1.5
+
+For a single H200, the optional
+[`moss_tts_local_h200.yaml`](../../../vllm_omni/deploy/moss_tts_local_h200.yaml)
+deployment places both the talker and codec on logical GPU 0:
+
+```bash
+CUDA_VISIBLE_DEVICES=0 VLLM_OMNI_EVENT_DRIVEN_ORCH=1 \
+    vllm serve OpenMOSS-Team/MOSS-TTS-Local-Transformer-v1.5 \
+    --omni --trust-remote-code \
+    --deploy-config vllm_omni/deploy/moss_tts_local_h200.yaml \
+    --disable-log-stats
+```
+
+Run from the repository root and select an available physical GPU with
+`CUDA_VISIBLE_DEVICES`. The command enables event-driven orchestration and
+disables detailed per-request statistics logging to reduce CPU overhead at
+high concurrency. Remove `--disable-log-stats` when those statistics are needed;
+keep these settings identical when comparing performance.
+This preset configures 256 request slots per stage,
+a 32 GiB talker KV cache, talker CUDA Graph buckets through 512 scheduled tokens,
+and codec batch buckets through 256. It requires more than 80 GiB of GPU memory;
+request capacity also depends on input and generated lengths. The default
+[`moss_tts_local.yaml`](../../../vllm_omni/deploy/moss_tts_local.yaml) remains
+available for smaller deployments.
+
+The preset enables the optional codec backend with
+`hf_overrides.codec_attention_backend: triton` on stage 1. It preserves the
+streaming ring-cache mask and uses BF16 attention with 64-dimensional heads;
+other attention shapes use PyTorch SDPA. Set the backend to `sdpa` to use the
+default attention implementation. Codec terminal tails share execution only
+when they already map to the same padded CUDA Graph; returned audio retains
+each request's actual length.
+
+Voice cloning requests use `ref_audio` and `ref_text`. For streaming output,
+set `stream: true`, `stream_format: "audio"`, and `response_format: "pcm"`;
+the native PCM format is 48 kHz, stereo, signed 16-bit little-endian. Compute
+audio throughput as `PCM bytes / (48000 * 2 * 2) / elapsed seconds`, and compare
+the same requests, concurrency, and reference-cache state.
+
 ## MOSS-TTS-Nano
 
 Single-stage 0.1B AR LM + MOSS-Audio-Tokenizer-Nano codec at 48 kHz mono. Every request must include `ref_audio`; there are no built-in speaker presets.

recipes/OpenMOSS/MOSS-TTS.md

@@ -16,17 +16,25 @@ Use this recipe for 24 kHz multilingual TTS with voice cloning (20 languages inc
 Chinese and English). Choose a variant based on your latency and quality requirements:
 
 | Model | Params | Use case |
-|---|---|---|
+| --- | --- | --- |
 | MOSS-TTS | 8B | General TTS, highest quality |
 | MOSS-TTS-v1.5 | 8B | General TTS upgrade of 1.0: 31 languages, steadier cloning, `[pause Xs]` markers (set `language` for best results); same `MossTTSDelay` API |
 | MOSS-TTS-Realtime | 1.7B | Lowest latency (TTFB ~180 ms), streaming-first |
 | MOSS-TTSD-v1.0 | 8B | Multi-turn dialogue TTS |
 | MOSS-SoundEffect | 8B | Sound effect synthesis from text description |
 | MOSS-VoiceGenerator | 1.7B | Zero-shot voice design |
 
-All variants share the same codec (`OpenMOSS-Team/MOSS-Audio-Tokenizer`, ~7 GB) and
+The variants above share the same codec (`OpenMOSS-Team/MOSS-Audio-Tokenizer`, ~7 GB) and
 output 24 kHz mono audio.
 
+MOSS-TTS-Local-Transformer-v1.5 uses MOSS-Audio-Tokenizer-v2 and outputs 48 kHz
+stereo audio. For Local voice cloning through `/v1/audio/speech`, provide an
+accurate `ref_text` transcript alongside `ref_audio`. The adapter uses the
+reference transcript and audio as a continuation prefix, then generates only
+the requested `input` speech. Without a nonblank `ref_text`, Local uses
+audio-reference generation. Reference transcripts must match the reference
+audio; they are not style instructions.
+
 ## References
 
 - Offline inference example: [`examples/offline_inference/text_to_speech/moss_tts/`](../../examples/offline_inference/text_to_speech/moss_tts/)

tests/benchmarks/test_seed_tts_dataset_variants.py

@@ -26,6 +26,7 @@
 
 if _MODULE_NAME not in sys.modules:
     _spec = importlib.util.spec_from_file_location(_MODULE_NAME, _MODULE_PATH)
+    assert _spec is not None and _spec.loader is not None
     _mod = importlib.util.module_from_spec(_spec)
     sys.modules[_MODULE_NAME] = _mod
     _spec.loader.exec_module(_mod)
@@ -237,7 +238,8 @@ def test_seed_tts_design_dataset_rejects_missing_description(seed_tts_design_roo
         assert req.seed_tts_utterance_id == "ok"
 
 
-def test_seed_tts_whisper_transcribe_passes_attention_mask(monkeypatch):
+@pytest.mark.parametrize("audio_seconds", [0.1, 61])
+def test_seed_tts_whisper_transcribe_passes_attention_mask(monkeypatch, audio_seconds):
     from vllm_omni.benchmarks.data_modules import seed_tts_eval
 
     calls = {}
@@ -252,8 +254,9 @@ def to(self, device):
             return self
 
     class FakeProcessor:
-        def __call__(self, wav, *, sampling_rate, return_tensors, return_attention_mask=False):
+        def __call__(self, wav, *, sampling_rate, return_tensors, return_attention_mask=False, **kwargs):
             calls["return_attention_mask"] = return_attention_mask
+            calls["processor_kwargs"] = kwargs
             assert sampling_rate == 16000
             assert return_tensors == "pt"
             assert len(wav) > 0
@@ -262,9 +265,10 @@ def __call__(self, wav, *, sampling_rate, return_tensors, return_attention_mask=
                 attention_mask=FakeTensor("mask") if return_attention_mask else None,
             )
 
-        def get_decoder_prompt_ids(self, *, language, task):
+        def get_decoder_prompt_ids(self, *, language, task, no_timestamps=True):
             assert language == "english"
             assert task == "transcribe"
+            calls["no_timestamps"] = no_timestamps
             return [(1, 2)]
 
         def batch_decode(self, predicted_ids, *, skip_special_tokens):
@@ -283,10 +287,18 @@ def generate(self, input_features, **kwargs):
     monkeypatch.setattr(seed_tts_eval, "_en_model", FakeModel())
     monkeypatch.setattr(seed_tts_eval, "_device", "cuda:1")
 
-    text = seed_tts_eval._transcribe_en_f32_16k(np.ones(1600, dtype=np.float32))
+    text = seed_tts_eval._transcribe_en_f32_16k(np.ones(int(audio_seconds * 16000), dtype=np.float32))
 
     assert text == "hello"
     assert calls["return_attention_mask"] is True
     assert calls["input_features"].device == "cuda:1"
     assert calls["generate_kwargs"]["attention_mask"].device == "cuda:1"
     assert calls["generate_kwargs"]["forced_decoder_ids"] == [(1, 2)]
+    if audio_seconds > 30:
+        assert calls["processor_kwargs"] == {"truncation": False, "padding": "longest"}
+        assert calls["generate_kwargs"]["return_timestamps"] is True
+        assert calls["no_timestamps"] is False
+    else:
+        assert calls["processor_kwargs"] == {}
+        assert "return_timestamps" not in calls["generate_kwargs"]
+        assert calls["no_timestamps"] is True

tests/entrypoints/openai_api/test_tts_adapter.py

@@ -166,6 +166,56 @@ def test_moss_tts_seed_falls_back_to_stage_default(adapter_cls, mocker):
     assert prepared.tts_params["seed"] == [42]
 
 
+@pytest.mark.parametrize(
+    "variant,ref_text,mode",
+    [
+        ("local", " Reference. ", "continuation"),
+        ("local", None, "generation"),
+        ("local", "  ", "generation"),
+        ("tts", "Reference.", "generation"),
+    ],
+)
+def test_moss_reference_transcript_mode(variant, ref_text, mode, mocker):
+    import torch
+
+    reference = [torch.ones((3, 12), dtype=torch.int64)]
+    unified = torch.arange(52, dtype=torch.int64).reshape(1, 4, 13)
+    processor = mocker.Mock(return_value={"input_ids": unified})
+    server = mocker.Mock(_moss_variant="local", uploaded_speakers={"speaker": {}})
+    server._voice_created_at.return_value = 123
+    request = OpenAICreateSpeechRequest(
+        input="Target.", ref_text=ref_text, language="English", voice="Speaker", seed=0, max_new_tokens=2048
+    )
+    adapter = MossTTSAdapter(SpeechServingContext(server=server))
+
+    adapter._moss_variant = variant
+    mocker.patch.object(adapter, "_get_moss_processor", return_value=processor)
+    mocker.patch.object(
+        adapter, "_encode_moss_references", new=mocker.AsyncMock(return_value=(reference, {0: "reference-key"}))
+    )
+
+    prepared = asyncio.run(adapter.build(request, [], has_inline_ref_audio=False))
+
+    if mode == "continuation":
+        processor.build_user_message.assert_called_once_with(text="Reference. Target.", language="English")
+        processor.build_assistant_message.assert_called_once_with(audio_codes_list=reference)
+        conversation = [processor.build_user_message.return_value, processor.build_assistant_message.return_value]
+    else:
+        processor.build_user_message.assert_called_once_with(text="Target.", language="English", reference=reference)
+        processor.build_assistant_message.assert_not_called()
+        conversation = [processor.build_user_message.return_value]
+    processor.assert_called_once_with(conversations=[conversation], mode=mode)
+    assert prepared.prompt["prompt_token_ids"] == unified[0, :, 0].tolist()
+    assert torch.equal(prepared.tts_params["codes"]["ref"], unified[0, :, 1:])
+    assert prepared.tts_params["max_new_frames"] == [2048]
+    assert prepared.tts_params["ref_audio_cache_key"] == "reference-key"
+    assert prepared.tts_params["voice_name"] == ["speaker"]
+    assert prepared.tts_params["voice_created_at"] == [123]
+    assert prepared.tts_params["seed"] == [0]
+    assert prepared.prompt["cache_salt"]
+    assert request.input == "Target."
+
+
 def test_qwen3_tts_metadata():
     assert Qwen3TTSAdapter.backend == "ar"
     assert issubclass(Qwen3TTSAdapter, ARTTSAdapter)

tests/entrypoints/test_omni_entrypoints.py

@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 from __future__ import annotations
 
 import queue
@@ -97,9 +100,11 @@ def __init__(
     ) -> None:
         self.model = model
         self.config_path = None
-        self.stage_configs: list[Any] = []
         self.stage_metadata = stage_metadata or [THREE_STAGE_META[-1]]
         self.num_stages = len(self.stage_metadata)
+        self.stage_configs = [
+            StageConfig(stage_id=i, model_stage="dummy-model").to_omegaconf() for i in range(self.num_stages)
+        ]
         self.default_sampling_params_list = default_sampling_params_list or [
             SamplingParams(max_tokens=8) for _ in range(self.num_stages)
         ]
@@ -222,6 +227,25 @@ def test_resolve_sampling_params_list_preserves_stage_constraints():
     assert caller_params.stop_token_ids == [7]
 
 
+@pytest.mark.parametrize("use_defaults", [False, True])
+def test_moss_local_output_policy_preserves_codec_streaming(use_defaults):
+    from vllm_omni.model_executor.models.moss_tts.pipeline import MOSS_TTS_LOCAL_PIPELINE
+
+    base = _make_base()
+    base.engine.num_stages = 2
+    base.sampling_constraints_list = [stage.sampling_constraints for stage in MOSS_TTS_LOCAL_PIPELINE.stages]
+    base.default_sampling_params_list = [
+        base._apply_sampling_constraints(SamplingParams(), constraints)
+        for constraints in base.sampling_constraints_list
+    ]
+    caller = [SamplingParams(output_kind=RequestOutputKind.DELTA) for _ in range(2)]
+    result = base.resolve_sampling_params_list(None if use_defaults else caller, allow_delta_coercion=True)
+
+    assert [params.output_kind for params in result] == [RequestOutputKind.FINAL_ONLY, RequestOutputKind.DELTA]
+    assert all(params.output_kind == RequestOutputKind.DELTA for params in caller)
+    assert base.default_sampling_params_list[0].output_kind == RequestOutputKind.FINAL_ONLY
+
+
 def _stage_spec(
     stage_id: int,
     *,

tests/model_executor/models/moss_tts/test_local_depth_rope_cache.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
 """CPU parity tests for MOSS-TTS Local Depth RoPE caching."""
 
 import pytest
@@ -42,3 +42,16 @@ def test_rope_cache_matches_reference_and_reuses_storage() -> None:
     # Smaller slices retain the fixed allocation used by CUDA Graph.
     attn.prepare_rope_cache(6, device, torch.float32)
     assert attn._rope_cos_cache.data_ptr() == cache_ptr
+
+
+def test_frame_local_kv_matches_full_prefix_and_reuses_slots() -> None:
+    model = MossTTSLocalDepthTransformer(_test_config()).eval()
+    model.h[0].attn.prepare_rope_cache(12, torch.device("cpu"), torch.float32)
+    cache = (torch.empty(2, 4, 12, 8), torch.empty(2, 4, 12, 8))
+    with torch.inference_mode():
+        for _ in range(2):
+            inputs = torch.randn(2, 12, 32)
+            for position in range(12):
+                expected = model._forward_prefix(inputs[:, : position + 1])[:, -1:]
+                actual = model._forward_prefix(inputs[:, position : position + 1], cache, position)
+                torch.testing.assert_close(actual, expected, atol=1e-6, rtol=1e-5)