[Feature] Support MOSS-TTS-Local-v1.5 Streaming (#4804)

Signed-off-by: gcanlin <canlinguosdu@gmail.com>

tests/worker/test_omni_gpu_model_runner.py

@@ -275,7 +275,7 @@ def test_talker_mtp_forward_passes_qwen3_tts_subtalker_sampling_params_to_talker
     runner = _make_runner(req_ids=("r1",), hidden_size=4)
     runner.requests["r1"].sampling_params = SimpleNamespace(
         seed=42,
-        extra_args={"qwen3_tts_request_seed": 42},
+        extra_args={"tts_local_seed": 42},
     )
     runner.talker_mtp = CaptureTalkerMTP()
     runner.vllm_config = SimpleNamespace(
@@ -319,11 +319,11 @@ def test_talker_mtp_forward_keeps_explicit_seeded_requests_scalar(monkeypatch):
     runner = _make_runner(req_ids=("r1", "r2"), hidden_size=4)
     runner.requests["r1"].sampling_params = SimpleNamespace(
         seed=11,
-        extra_args={"qwen3_tts_request_seed": 11},
+        extra_args={"tts_local_seed": 11},
     )
     runner.requests["r2"].sampling_params = SimpleNamespace(
         seed=22,
-        extra_args={"qwen3_tts_request_seed": 22},
+        extra_args={"tts_local_seed": 22},
     )
     runner.talker_mtp = CaptureTalkerMTP()
     runner.vllm_config = SimpleNamespace(model_config=SimpleNamespace(subtalker_sampling_params={}))

vllm_omni/benchmarks/data_modules/seed_tts_eval.py

@@ -36,7 +36,9 @@ Enable with ``SEED_TTS_WER_EVAL=1`` or ``--seed-tts-wer-eval``. Install optional
     pip install 'vllm-omni[dev]'
 
 Env: ``SEED_TTS_EVAL_DEVICE`` (e.g. ``cuda:0``, ``cpu``); ``SEED_TTS_HF_WHISPER_MODEL``
-defaults to ``openai/whisper-large-v3`` (override for debugging only).
+defaults to ``openai/whisper-large-v3`` (override for debugging only). Set
+``SEED_TTS_WER_SAVE_AUDIO_DIR`` to save the captured 24 kHz mono WAV used by
+WER evaluation for each synthesized utterance.
 """
 
 from __future__ import annotations
@@ -50,6 +52,7 @@ import string
 import tempfile
 import threading
 import wave
+from pathlib import Path
 from typing import Any
 
 import numpy as np
@@ -87,6 +90,30 @@ def pcm_s16le_mono_to_wav_bytes(pcm: bytes, *, sample_rate: int = 24000) -> byte
     return buf.getvalue()
 
 
+def _safe_filename_part(value: Any, *, default: str = "item", max_len: int = 96) -> str:
+    text = str(value or "").strip() or default
+    safe = "".join(ch if ch.isalnum() or ch in ("-", "_", ".") else "_" for ch in text)
+    safe = safe.strip("._") or default
+    return safe[:max_len]
+
+
+def _save_seed_tts_eval_audio(
+    pcm: bytes,
+    *,
+    output_dir: Path | None,
+    index: int,
+    utterance_id: Any,
+    locale: str,
+) -> str | None:
+    if output_dir is None:
+        return None
+    output_dir.mkdir(parents=True, exist_ok=True)
+    stem = f"{index:05d}_{_safe_filename_part(utterance_id)}_{_safe_filename_part(locale)}"
+    path = output_dir / f"{stem}.wav"
+    path.write_bytes(pcm_s16le_mono_to_wav_bytes(pcm, sample_rate=24000))
+    return str(path)
+
+
 def _get_eval_device() -> str:
     explicit = os.environ.get("SEED_TTS_EVAL_DEVICE", "").strip()
     if explicit:
@@ -512,13 +539,18 @@ def compute_seed_tts_wer_metrics(
     sim_skipped_no_ref = 0
     utmos_failed = 0
     utmos_on = _eval_submetric_enabled("SEED_TTS_UTMOS_EVAL", default=False)
+    save_audio_raw = os.environ.get("SEED_TTS_WER_SAVE_AUDIO_DIR", "").strip()
+    save_audio_dir = Path(save_audio_raw).expanduser() if save_audio_raw else None
+    saved_audio = 0
+    save_audio_failed = 0
 
-    for req, out in zip(input_requests, outputs, strict=True):
+    for index, (req, out) in enumerate(zip(input_requests, outputs, strict=True)):
         assert isinstance(req, SeedTTSSampleRequest)
         ref = req.prompt
         locale = req.seed_tts_locale or "en"
         row_lang = "zh" if locale.lower().startswith("zh") else "en"
         utmos_v: float | None = None
+        audio_path: str | None = None
 
         if not out.success:
             request_failed += 1
@@ -545,6 +577,23 @@ def compute_seed_tts_wer_metrics(
                     }
                 )
             continue
+        try:
+            audio_path = _save_seed_tts_eval_audio(
+                pcm,
+                output_dir=save_audio_dir,
+                index=index,
+                utterance_id=req.seed_tts_utterance_id,
+                locale=locale,
+            )
+            if audio_path:
+                saved_audio += 1
+        except OSError as e:
+            save_audio_failed += 1
+            logger.warning(
+                "Seed-TTS WER audio save failed for utterance=%s: %s",
+                req.seed_tts_utterance_id,
+                e,
+            )
 
         wav_16k = _pcm_s16le_to_f32_16k(pcm)
         if len(wav_16k) == 0:
@@ -555,6 +604,7 @@ def compute_seed_tts_wer_metrics(
                         "utterance_id": req.seed_tts_utterance_id,
                         "locale": locale,
                         "error": "empty_audio",
+                        "audio_path": audio_path,
                     }
                 )
             continue
@@ -609,6 +659,7 @@ def compute_seed_tts_wer_metrics(
                         "locale": locale,
                         "error": "asr_exception",
                         "detail": str(e)[:500],
+                        "audio_path": audio_path,
                     }
                 )
             continue
@@ -621,6 +672,7 @@ def compute_seed_tts_wer_metrics(
                         "utterance_id": req.seed_tts_utterance_id,
                         "locale": locale,
                         "error": "empty_asr",
+                        "audio_path": audio_path,
                     }
                 )
             continue
@@ -637,6 +689,7 @@ def compute_seed_tts_wer_metrics(
                         "locale": locale,
                         "error": "wer_compute_failed",
                         "detail": str(e)[:500],
+                        "audio_path": audio_path,
                     }
                 )
             continue
@@ -673,6 +726,8 @@ def compute_seed_tts_wer_metrics(
                 "reference_raw": raw_truth,
                 "asr_raw": raw_hypo,
             }
+            if audio_path:
+                row["audio_path"] = audio_path
             if sim_v is not None:
                 row["sim"] = sim_v
             if utmos_v is not None:
@@ -697,7 +752,11 @@ def compute_seed_tts_wer_metrics(
         "seed_tts_utmos_mean": statistics.fmean(utmos_values) if utmos_values else None,
         "seed_tts_utmos_median": statistics.median(utmos_values) if utmos_values else None,
         "seed_tts_utmos_failed": utmos_failed,
+        "seed_tts_saved_audio": saved_audio,
+        "seed_tts_save_audio_failed": save_audio_failed,
     }
+    if save_audio_dir is not None:
+        result["seed_tts_save_audio_dir"] = str(save_audio_dir)
     if include_per_item:
         result["seed_tts_wer_eval_items"] = items
     return result
@@ -729,6 +788,13 @@ def print_seed_tts_wer_summary(metrics: dict[str, Any]) -> None:
     print("{:<40} {:<10}".format("Request failed:", metrics.get("seed_tts_request_failed", 0)))
     print("{:<40} {:<10}".format("No PCM captured:", metrics.get("seed_tts_no_pcm", 0)))
     print("{:<40} {:<10}".format("ASR / WER failed:", metrics.get("seed_tts_asr_failed", 0)))
+    save_dir = metrics.get("seed_tts_save_audio_dir")
+    if save_dir:
+        print("{:<40} {:<10}".format("Saved eval WAVs:", metrics.get("seed_tts_saved_audio", 0)))
+        print("{:<40} {}".format("Saved eval WAV dir:", save_dir))
+        failed = int(metrics.get("seed_tts_save_audio_failed", 0) or 0)
+        if failed:
+            print("{:<40} {:<10}".format("Save eval WAV errors:", failed))
     if sim_ev or metrics.get("seed_tts_sim_skipped_no_ref") or metrics.get("seed_tts_sim_failed"):
         print("{:<40} {:<10}".format("SIM evaluated (higher ~ closer):", sim_ev))
         sm = metrics.get("seed_tts_sim_mean")

vllm_omni/benchmarks/patch/patch.py

@@ -137,6 +137,30 @@ def _audio_pcm_format() -> tuple[int, int]:
     return sample_rate, channels
 
 
+def _pcm_s16le_to_seed_tts_wer_bytes(
+    pcm_bytes: bytes,
+    *,
+    sample_rate: int,
+    channels: int,
+) -> bytes:
+    """Normalize streamed raw PCM to the 24 kHz mono PCM used by Seed-TTS WER."""
+    if not pcm_bytes:
+        return b""
+    channels = max(1, int(channels))
+    pcm = np.frombuffer(pcm_bytes, dtype=np.int16)
+    if channels > 1:
+        usable = (pcm.size // channels) * channels
+        pcm = pcm[:usable].reshape(-1, channels).mean(axis=1)
+    pcm_f32 = pcm.astype(np.float32) / 32767.0
+    if int(sample_rate) != 24000 and pcm_f32.size:
+        from vllm.multimodal.audio import AudioResampler
+
+        resampler = AudioResampler(target_sr=24000)
+        pcm_f32 = resampler.resample(pcm_f32, orig_sr=int(sample_rate))
+    pcm_f32 = np.clip(pcm_f32, -1.0, 1.0)
+    return (pcm_f32 * 32767).astype(np.int16).tobytes()
+
+
 get_samples_old = datasets.get_samples
 
 _DEFAULT_DAILY_OMNI_REPO = "liarliar/Daily-Omni"
@@ -1216,7 +1240,15 @@ async def async_request_openai_audio_speech(
                 output.audio_continuity_ok = continuity.is_continuous
                 output.audio_underrun_event_count = continuity.underrun_event_count
                 if pcm_capture is not None and pcm_capture:
-                    output.tts_output_pcm_bytes = bytes(pcm_capture)
+                    try:
+                        output.tts_output_pcm_bytes = _pcm_s16le_to_seed_tts_wer_bytes(
+                            bytes(pcm_capture),
+                            sample_rate=sample_rate,
+                            channels=channels,
+                        )
+                    except Exception as ex:
+                        logger.warning("Seed-TTS WER PCM normalization failed: %s", ex)
+                        output.tts_output_pcm_bytes = bytes(pcm_capture)
                 elif capture_wer_pcm:
                     ct = response.headers.get("Content-Type", "")
                     logger.warning(

vllm_omni/deploy/moss_tts_local.yaml

@@ -14,8 +14,10 @@ connectors:
       connector_get_sleep_s: 0.005
       connector_get_max_wait_first_chunk: 1000
       connector_get_max_wait: 300
-      codec_chunk_frames: 15          # 15 frames * 80 ms/frame = 1.2 s chunks
+      initial_codec_chunk_frames: 1   # 1 frame * 80 ms/frame = 80 ms first audio chunk
+      codec_chunk_frames: 15          # 15 frames * 80 ms/frame = 1.2 s steady chunks
       codec_left_context_frames: 0
+      codec_stream_slots: 8
 
 stages:
   - stage_id: 0
@@ -37,7 +39,7 @@ stages:
       repetition_penalty: 1.1
 
   - stage_id: 1
-    max_num_seqs: 1
+    max_num_seqs: 8
     gpu_memory_utilization: 0.15
     enforce_eager: true
     enable_prefix_caching: false

vllm_omni/entrypoints/openai/serving_speech.py

@@ -3705,15 +3705,14 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
                 sampling_params_list[0].max_tokens = int(request.max_new_tokens) + 1
 
         if request.seed is not None and sampling_params_list:
-            if sampling_params_list is self.engine_client.default_sampling_params_list:
-                import copy
+            import copy
 
-                sampling_params_list = copy.deepcopy(sampling_params_list)
-            sampling_params_list[0].seed = request.seed
-            if self._tts_model_type == "qwen3_tts":
-                if sampling_params_list[0].extra_args is None:
-                    sampling_params_list[0].extra_args = {}
-                sampling_params_list[0].extra_args["qwen3_tts_request_seed"] = request.seed
+            sampling_params_list = copy.deepcopy(sampling_params_list)
+            stage0_params = sampling_params_list[0]
+            stage0_params.seed = request.seed
+            if stage0_params.extra_args is None:
+                stage0_params.extra_args = {}
+            stage0_params.extra_args["tts_local_seed"] = request.seed
 
         generator = self.engine_client.generate(
             prompt=prompt,

vllm_omni/model_executor/models/moss_tts/configuration_moss_audio_tokenizer_v2.py

@@ -0,0 +1,359 @@
+# Copyright 2026 OpenMOSS and the HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""MossAudioTokenizer model configuration"""
+
+from typing import Any
+
+from transformers.configuration_utils import PreTrainedConfig
+from transformers.utils import logging
+
+logger = logging.get_logger(__name__)
+
+
+class MossAudioTokenizerConfig(PreTrainedConfig):
+    r"""
+    This is the configuration class to store the configuration of a [`MossAudioTokenizerModel`]. It is used to
+    instantiate a MossAudioTokenizer model according to the specified arguments, defining the model architecture.
+
+    Instantiating a configuration with the defaults will yield a similar configuration to that of the
+    [VoiceAgentGroup/moss_audio_tokenizer](https://huggingface.co/VoiceAgentGroup/moss_audio_tokenizer) architecture.
+
+    Configuration objects inherit from [`PreTrainedConfig`] and can be used to control the model outputs. Read the
+    documentation from [`PreTrainedConfig`] for more information.
+
+    Args:
+        sampling_rate (`int`, *optional*, defaults to 24000):
+            The sampling rate at which the audio waveform should be digitalized expressed in hertz (Hz).
+        downsample_rate (`int`, *optional*, defaults to 1920):
+            Total downsampling rate from waveform to tokens.
+        causal_transformer_context_duration (`float`, *optional*, defaults to 10.0):
+            Context duration in seconds for causal transformer.
+        encoder_kwargs (`list[dict]`, *optional*):
+            List of encoder module configurations. Each dict specifies a module type and its parameters.
+        decoder_kwargs (`list[dict]`, *optional*):
+            List of decoder module configurations in execution order.
+        quantizer_type (`str`, *optional*, defaults to `"rvq"`):
+            Quantizer type. Options include `"rvq"`, `"spec_rvq"`, `"rlfq"`, `"random_prefix_rlfq"`.
+        quantizer_kwargs (`dict`, *optional*):
+            Configuration for the quantizer including `input_dim`, `rvq_dim`, `output_dim`, `num_quantizers`,
+            `codebook_size`, and `codebook_dim`.
+
+    Example:
+
+    ```python
+    >>> from transformers import MossAudioTokenizerModel, MossAudioTokenizerConfig
+
+    >>> # Initializing a MossAudioTokenizer style configuration
+    >>> configuration = MossAudioTokenizerConfig()
+
+    >>> # Initializing a model (with random weights) from the configuration
+    >>> model = MossAudioTokenizerModel(configuration)
+
+    >>> # Accessing the model configuration
+    >>> configuration = model.config
+    ```
+    """
+
+    model_type = "moss-audio-tokenizer"
+
+    # Backward-compatible alias used by some checkpoints.
+    attribute_map = {"sample_rate": "sampling_rate"}
+
+    sampling_rate: int
+    downsample_rate: int
+    causal_transformer_context_duration: float
+    encoder_kwargs: list[dict[str, Any]]
+    decoder_kwargs: list[dict[str, Any]]
+    number_channels: int
+    enable_channel_interleave: bool
+    attention_implementation: str
+    compute_dtype: str
+    codec_weight_dtype: str
+    quantizer_type: str
+    quantizer_kwargs: dict[str, Any]
+
+    def __init__(
+        self,
+        version: str | None = None,
+        sampling_rate: int = 48000,
+        downsample_rate: int = 3840,
+        causal_transformer_context_duration: float = 10.0,
+        encoder_kwargs: list[dict[str, Any]] | None = None,
+        decoder_kwargs: list[dict[str, Any]] | None = None,
+        number_channels: int = 2,
+        enable_channel_interleave: bool = True,
+        attention_implementation: str = "sdpa",
+        compute_dtype: str = "fp32",
+        codec_weight_dtype: str = "fp32",
+        quantizer_type: str = "rlfq",
+        quantizer_kwargs: dict[str, Any] | None = None,
+        **kwargs,
+    ):
+        # Some checkpoints might include an incorrect/legacy `model_type` (e.g. "speech_tokenizer").
+        # We drop it to avoid overriding the class-level `model_type`.
+        kwargs.pop("model_type", None)
+        if "channels_numbers" in kwargs:
+            number_channels = kwargs.pop("channels_numbers")
+        if "enable_channel_interleave" in kwargs:
+            enable_channel_interleave = kwargs.pop("enable_channel_interleave")
+        if "attention_backend" in kwargs and attention_implementation == "sdpa":
+            attention_implementation = kwargs.pop("attention_backend")
+        if "codec_compute_dtype" in kwargs and compute_dtype == "fp32":
+            compute_dtype = kwargs.pop("codec_compute_dtype")
+        if "codec_load_dtype" in kwargs and codec_weight_dtype == "fp32":
+            codec_weight_dtype = kwargs.pop("codec_load_dtype")
+
+        # `version` is accepted for compatibility but not used in modeling.
+        self.version = version
+        self.sampling_rate = sampling_rate
+        self.downsample_rate = downsample_rate
+        self.causal_transformer_context_duration = causal_transformer_context_duration
+        self.number_channels = number_channels
+        self.enable_channel_interleave = enable_channel_interleave
+        self.attention_implementation = attention_implementation
+        self.compute_dtype = compute_dtype
+        self.codec_weight_dtype = codec_weight_dtype
+        # Default encoder configuration
+        if encoder_kwargs is None:
+            encoder_kwargs = [
+                {
+                    "module_type": "PatchedPretransform",
+                    "patch_size": 240,
+                },
+                {
+                    "module_type": "Transformer",
+                    "input_dimension": 240,
+                    "output_dimension": 384,
+                    "d_model": 768,
+                    "num_heads": 12,
+                    "num_layers": 12,
+                    "dim_feedforward": 3072,
+                    "causal": True,
+                    "norm": "layer_norm",
+                    "positional_embedding": "rope",
+                    "max_period": 10000,
+                    "gating": "none",
+                    "layer_scale": 0.01,
+                    "conv_layout": True,
+                },
+                {
+                    "module_type": "PatchedPretransform",
+                    "patch_size": 2,
+                },
+                {
+                    "module_type": "Transformer",
+                    "input_dimension": 768,
+                    "output_dimension": 384,
+                    "d_model": 768,
+                    "num_heads": 12,
+                    "num_layers": 12,
+                    "dim_feedforward": 3072,
+                    "causal": True,
+                    "norm": "layer_norm",
+                    "positional_embedding": "rope",
+                    "max_period": 10000,
+                    "gating": "none",
+                    "layer_scale": 0.01,
+                    "conv_layout": True,
+                },
+                {
+                    "module_type": "PatchedPretransform",
+                    "patch_size": 2,
+                },
+                {
+                    "module_type": "Transformer",
+                    "input_dimension": 768,
+                    "output_dimension": 640,
+                    "d_model": 768,
+                    "num_heads": 12,
+                    "num_layers": 12,
+                    "dim_feedforward": 3072,
+                    "causal": True,
+                    "norm": "layer_norm",
+                    "positional_embedding": "rope",
+                    "max_period": 10000,
+                    "gating": "none",
+                    "layer_scale": 0.01,
+                    "conv_layout": True,
+                },
+                {
+                    "module_type": "PatchedPretransform",
+                    "patch_size": 2,
+                },
+                {
+                    "module_type": "Transformer",
+                    "input_dimension": 1280,
+                    "output_dimension": 768,
+                    "d_model": 1280,
+                    "num_heads": 20,
+                    "num_layers": 32,
+                    "dim_feedforward": 5120,
+                    "causal": True,
+                    "norm": "layer_norm",
+                    "positional_embedding": "rope",
+                    "max_period": 10000,
+                    "gating": "none",
+                    "layer_scale": 0.01,
+                    "conv_layout": True,
+                },
+            ]
+        self.encoder_kwargs = encoder_kwargs
+
+        # Default decoder configuration (execution order)
+        if decoder_kwargs is None:
+            decoder_kwargs = [
+                {
+                    "module_type": "Transformer",
+                    "input_dimension": 768,
+                    "output_dimension": 1280,
+                    "d_model": 1280,
+                    "num_heads": 20,
+                    "num_layers": 32,
+                    "dim_feedforward": 5120,
+                    "causal": True,
+                    "norm": "layer_norm",
+                    "positional_embedding": "rope",
+                    "max_period": 10000,
+                    "gating": "none",
+                    "layer_scale": 0.01,
+                    "conv_layout": True,
+                },
+                {
+                    "module_type": "PatchedPretransform",
+                    "patch_size": 2,
+                },
+                {
+                    "module_type": "Transformer",
+                    "input_dimension": 640,
+                    "output_dimension": 768,
+                    "d_model": 768,
+                    "num_heads": 12,
+                    "num_layers": 12,
+                    "dim_feedforward": 3072,
+                    "causal": True,
+                    "norm": "layer_norm",
+                    "positional_embedding": "rope",
+                    "max_period": 10000,
+                    "gating": "none",
+                    "layer_scale": 0.01,
+                    "conv_layout": True,
+                },
+                {
+                    "module_type": "PatchedPretransform",
+                    "patch_size": 2,
+                },
+                {
+                    "module_type": "Transformer",
+                    "input_dimension": 384,
+                    "output_dimension": 768,
+                    "d_model": 768,
+                    "num_heads": 12,
+                    "num_layers": 12,
+                    "dim_feedforward": 3072,
+                    "causal": True,
+                    "norm": "layer_norm",
+                    "positional_embedding": "rope",
+                    "max_period": 10000,
+                    "gating": "none",
+                    "layer_scale": 0.01,
+                    "conv_layout": True,
+                },
+                {
+                    "module_type": "PatchedPretransform",
+                    "patch_size": 2,
+                },
+                {
+                    "module_type": "Transformer",
+                    "input_dimension": 384,
+                    "output_dimension": 768,
+                    "d_model": 768,
+                    "num_heads": 12,
+                    "num_layers": 12,
+                    "dim_feedforward": 3072,
+                    "causal": True,
+                    "norm": "layer_norm",
+                    "positional_embedding": "rope",
+                    "max_period": 10000,
+                    "gating": "none",
+                    "layer_scale": 0.01,
+                    "conv_layout": True,
+                },
+                {
+                    "module_type": "PatchedPretransform",
+                    "patch_size": 2,
+                },
+                {
+                    "module_type": "Transformer",
+                    "input_dimension": 384,
+                    "output_dimension": 240,
+                    "d_model": 768,
+                    "num_heads": 12,
+                    "num_layers": 12,
+                    "dim_feedforward": 3072,
+                    "causal": True,
+                    "norm": "layer_norm",
+                    "positional_embedding": "rope",
+                    "max_period": 10000,
+                    "gating": "none",
+                    "layer_scale": 0.01,
+                    "conv_layout": True,
+                },
+                {
+                    "module_type": "PatchedPretransform",
+                    "patch_size": 240,
+                },
+            ]
+        self.decoder_kwargs = decoder_kwargs
+
+        # Default quantizer configuration
+        if quantizer_kwargs is None:
+            quantizer_kwargs = {
+                "input_dim": 768,
+                "rvq_dim": 512,
+                "output_dim": 768,
+                "num_quantizers": 32,
+                "codebook_size": 1024,
+                "codebook_dim": 8,
+                "quantizer_type": "rlfq",
+            }
+
+        # Handle quantizer_type from kwargs or config
+        kw_qtype = quantizer_kwargs.get("quantizer_type", None)
+        if kw_qtype is not None:
+            self.quantizer_type = kw_qtype
+        else:
+            self.quantizer_type = quantizer_type
+            quantizer_kwargs["quantizer_type"] = quantizer_type
+
+        self.quantizer_kwargs = quantizer_kwargs
+
+        super().__init__(**kwargs)
+
+    @property
+    def num_quantizers(self) -> int:
+        """Return the number of quantizers from quantizer_kwargs."""
+        return self.quantizer_kwargs.get("num_quantizers", 32)
+
+    @property
+    def codebook_size(self) -> int:
+        """Return the codebook size from quantizer_kwargs."""
+        return self.quantizer_kwargs.get("codebook_size", 4096)
+
+    @property
+    def frame_rate(self) -> float:
+        """Return the frame rate (tokens per second)."""
+        return self.sampling_rate / self.downsample_rate
+
+
+__all__ = ["MossAudioTokenizerConfig"]

vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_codec.py

@@ -5,6 +5,7 @@
 
 from __future__ import annotations
 
+import contextlib
 from collections.abc import Iterable
 from typing import Any
 
@@ -19,6 +20,12 @@ from vllm_omni.model_executor.models.moss_tts.audio_tokenizer import (
     MossAudioTokenizerConfig,
     MossAudioTokenizerModel,
 )
+from vllm_omni.model_executor.models.moss_tts.audio_tokenizer_v2 import (
+    MossAudioTokenizerModel as MossAudioTokenizerV2Model,
+)
+from vllm_omni.model_executor.models.moss_tts.configuration_moss_audio_tokenizer_v2 import (
+    MossAudioTokenizerConfig as MossAudioTokenizerV2Config,
+)
 from vllm_omni.model_executor.models.moss_tts.moss_codec_cudagraph import (
     MossTTSCUDAGraphCodecWrapper,
 )
@@ -27,6 +34,96 @@ from vllm_omni.model_executor.models.output_templates import OmniOutput
 logger = init_logger(__name__)
 
 
+class _MossCodecStreamSession:
+    """Persistent streaming decode session for vendored MOSS-Audio-Tokenizer-v2."""
+
+    def __init__(
+        self,
+        codec: nn.Module,
+        *,
+        stream_slots: int,
+        n_vq: int,
+    ) -> None:
+        self._codec = codec
+        self._stream_slots = int(stream_slots)
+        self._batch_size = self._stream_slots
+        self._n_vq = int(n_vq)
+        self._device = next(codec.parameters()).device
+        self._free_stream_slots = list(range(self._stream_slots))
+        self._exit_stack = contextlib.ExitStack()
+        self._closed = False
+        with torch.no_grad():
+            self._exit_stack.enter_context(codec.streaming(self._batch_size))
+
+    def acquire(self) -> int | None:
+        if not self._free_stream_slots:
+            return None
+        return self._free_stream_slots.pop()
+
+    def release(self, slot: int) -> None:
+        if self._closed:
+            return
+        self.reset_slots([slot])
+        self._free_stream_slots.append(slot)
+
+    def reset_slots(self, slots: list[int]) -> None:
+        if not slots:
+            return
+        reset_mask = torch.zeros(self._batch_size, dtype=torch.bool, device=self._device)
+        reset_mask[slots] = True
+
+        def _reset(module: nn.Module) -> None:
+            state = getattr(module, "_streaming_state", None)
+            if state is not None:
+                state.reset(reset_mask.to(state.device))
+
+        with torch.no_grad():
+            self._codec.apply(_reset)
+
+    def close(self) -> None:
+        if self._closed:
+            return
+        with torch.no_grad():
+            self._exit_stack.close()
+        self._closed = True
+
+    @torch.no_grad()
+    def step(self, slot_codes: dict[int, torch.Tensor]) -> dict[int, torch.Tensor]:
+        if not slot_codes:
+            return {}
+        step_lengths = {int(codes.shape[1]) for codes in slot_codes.values()}
+        if len(step_lengths) != 1:
+            raise ValueError(f"MOSS codec streaming step needs uniform T, got {sorted(step_lengths)}")
+        (step_t,) = step_lengths
+        codes_step = torch.zeros(
+            self._n_vq,
+            self._batch_size,
+            step_t,
+            dtype=torch.long,
+            device=self._device,
+        )
+        codes_lengths = torch.zeros(self._batch_size, dtype=torch.long, device=self._device)
+        exec_mask = torch.zeros(self._batch_size, dtype=torch.bool, device=self._device)
+        for slot, codes in slot_codes.items():
+            codes_step[:, slot, :] = codes.to(device=self._device, dtype=torch.long)
+            codes_lengths[slot] = int(codes.shape[1])
+            exec_mask[slot] = True
+
+        self._codec._set_streaming_exec_mask(exec_mask)
+        result = self._codec._decode_frame(codes_step, codes_lengths)
+        if result.audio is None:
+            return {}
+        audio = result.audio.detach().to("cpu", torch.float32)
+        lengths = result.audio_lengths.detach().to("cpu") if result.audio_lengths is not None else None
+        out: dict[int, torch.Tensor] = {}
+        for slot in slot_codes:
+            wav = audio[slot]
+            if lengths is not None:
+                wav = wav[..., : int(lengths[slot].item())]
+            out[slot] = wav.contiguous()
+        return out
+
+
 class MossTTSCodecDecoder(nn.Module):
     """Stage-1 decoder for all MOSS-TTS variants.
 
@@ -75,6 +172,12 @@ class MossTTSCodecDecoder(nn.Module):
         self._cuda_graph_wrapper: MossTTSCUDAGraphCodecWrapper | None = None
         self._n_channels: int = 1
         self._sr_tensor = torch.tensor(self._OUTPUT_SAMPLE_RATE, dtype=torch.int32)
+        self._stream_session: _MossCodecStreamSession | None = None
+        self._stream_slots: int = self._connector_int("codec_stream_slots", default=0)
+        self._stream_max_step_frames: int = self._connector_int("codec_max_step_frames", default=100)
+        self._stream_req_slots: dict[str, int] = {}
+        self._stream_pending_codes: dict[str, list[torch.Tensor]] = {}
+        self._stream_starved_reqs: set[str] = set()
 
     # ------------------------------------------------------------------
     # vLLM-Omni stubs (codec has no AR loop)
@@ -112,9 +215,9 @@ class MossTTSCodecDecoder(nn.Module):
         chunk transfer adapter assigns those to ``request.prompt_token_ids``,
         which arrives here as ``input_ids`` concatenated across all requests.
         Per-request slice boundaries are computed from
-        ``kwargs["num_scheduled_tokens"]`` (a list of token counts, one per
-        request).  ``runtime_additional_information`` carries per-request
-        metadata such as ``left_context_size``.
+        ``kwargs["seq_token_counts"]`` (token counts, one per request).
+        ``runtime_additional_information`` carries per-request metadata such as
+        ``left_context_size``.
 
         Returns
         -------
@@ -123,8 +226,8 @@ class MossTTSCodecDecoder(nn.Module):
           multimodal_outputs["sr"]            — list of scalar int32 tensors
         """
         sr_tensor = self._sr_tensor
-        empty = torch.zeros((0,), dtype=torch.float32)
-        info_list: list[dict[str, Any]] = runtime_additional_information or [{}]
+        empty = self._empty_audio()
+        info_list: list[dict[str, Any]] = list(runtime_additional_information or [{}])
         num_req = max(len(info_list), 1)
 
         if self._codec is None:
@@ -140,24 +243,47 @@ class MossTTSCodecDecoder(nn.Module):
         audios: list[torch.Tensor] = [empty] * num_req
         srs: list[torch.Tensor] = [sr_tensor] * num_req
         device = next(self._codec.parameters()).device
+        streaming_work: list[tuple[int, str, torch.Tensor, bool]] = []
 
         if input_ids is None or input_ids.numel() == 0:
+            for i, wav in self._finish_empty_streaming_requests(info_list).items():
+                audios[i] = wav.reshape(-1) if wav.ndim == 1 or int(wav.shape[0]) == 1 else wav
             return OmniOutput(
                 text_hidden_states=None,
                 multimodal_outputs={"model_outputs": audios, "sr": srs},
             )
 
-        # ``input_ids`` is concatenated across all requests; split by
-        # query_start_loc-style offsets from ``num_scheduled_tokens`` in
-        # kwargs if available, else assume one request.
+        # ``input_ids`` is concatenated across all requests. vLLM-Omni runners
+        # pass the per-request lengths via the shared code2wav contract
+        # ``seq_token_counts``.
         ids_flat = input_ids.reshape(-1).to(dtype=torch.long)
-        num_scheduled_tokens = kwargs.get("num_scheduled_tokens")
-        if isinstance(num_scheduled_tokens, (list, tuple)) and len(num_scheduled_tokens) == num_req:
-            offsets = [0]
-            for n in num_scheduled_tokens:
-                offsets.append(offsets[-1] + int(n))
-        else:
-            offsets = [0, int(ids_flat.shape[0])]
+        token_counts = self._normalize_seq_token_counts(kwargs.get("seq_token_counts"))
+        if token_counts is None:
+            raise RuntimeError(
+                "MossTTS codec requires seq_token_counts; otherwise concatenated "
+                "codec tokens cannot be split per request."
+            )
+        if sum(token_counts) != int(ids_flat.shape[0]):
+            raise RuntimeError(
+                "MossTTS codec seq_token_counts mismatch: "
+                f"counts={token_counts}, sum={sum(token_counts)}, input_tokens={int(ids_flat.shape[0])}."
+            )
+
+        num_req = len(token_counts)
+        if len(info_list) < num_req:
+            info_list.extend({} for _ in range(num_req - len(info_list)))
+        elif len(info_list) > num_req:
+            info_list = info_list[:num_req]
+        if len(audios) < num_req:
+            audios.extend(empty for _ in range(num_req - len(audios)))
+            srs.extend(sr_tensor for _ in range(num_req - len(srs)))
+        elif len(audios) > num_req:
+            audios = audios[:num_req]
+            srs = srs[:num_req]
+
+        offsets = [0]
+        for n in token_counts:
+            offsets.append(offsets[-1] + int(n))
 
         for i, info in enumerate(info_list):
             if i + 1 >= len(offsets):
@@ -165,6 +291,14 @@ class MossTTSCodecDecoder(nn.Module):
             seg = ids_flat[offsets[i] : offsets[i + 1]]
             if seg.numel() == 0:
                 continue
+            meta = (info.get("meta", {}) if isinstance(info, dict) else {}) or {}
+            finished = bool(meta.get("stream_finished", meta.get("finished", False)))
+            streaming_enabled = bool(meta.get("codec_streaming", False))
+            code_flat_numel = meta.get("code_flat_numel")
+            if streaming_enabled and finished and code_flat_numel is not None and int(code_flat_numel) == 0:
+                for _, wav in self._finish_empty_streaming_requests([info]).items():
+                    audios[i] = wav.reshape(-1) if wav.ndim == 1 or int(wav.shape[0]) == 1 else wav
+                continue
             if seg.numel() % self._n_vq != 0:
                 logger.warning(
                     "MossTTS codec input length %d not divisible by n_vq %d; skipping.",
@@ -181,14 +315,19 @@ class MossTTSCodecDecoder(nn.Module):
             codebook_size = self._codec.config.codebook_size
             codes_nq_t = codes_nq_t.clamp_(0, int(codebook_size) - 1)
 
-            meta = (info.get("meta", {}) if isinstance(info, dict) else {}) or {}
-            left_ctx = self._runtime_value(info, meta, "left_context_size", 0)
+            left_ctx = meta.get("left_context_size", 0)
             if isinstance(left_ctx, (list, tuple)):
                 left_ctx = int(left_ctx[0]) if left_ctx else 0
             elif isinstance(left_ctx, torch.Tensor):
                 left_ctx = int(left_ctx.reshape(-1)[0].item()) if left_ctx.numel() else 0
             left_ctx = int(left_ctx)
 
+            req_key = self._runtime_request_key(info, meta, i)
+
+            if streaming_enabled:
+                streaming_work.append((i, req_key, codes_nq_t, finished))
+                continue
+
             if self._cuda_graph_wrapper is not None:
                 out = self._cuda_graph_wrapper.decode(codes_nq_t)
             else:
@@ -216,19 +355,246 @@ class MossTTSCodecDecoder(nn.Module):
                     )
                 wav = wav[..., trim:]
 
-            audios[i] = wav.reshape(-1) if self._n_channels == 1 else wav
+            audios[i] = wav.reshape(-1) if wav.ndim == 1 or int(wav.shape[0]) == 1 else wav
+
+        if streaming_work:
+            for i, wav in self._decode_streaming_batch(streaming_work).items():
+                audios[i] = wav.reshape(-1) if wav.ndim == 1 or int(wav.shape[0]) == 1 else wav
 
         return OmniOutput(
             text_hidden_states=None,
             multimodal_outputs={"model_outputs": audios, "sr": srs},
         )
 
+    def _finish_empty_streaming_requests(self, info_list: list[dict[str, Any]]) -> dict[int, torch.Tensor]:
+        """Release codec stream state for empty finish sentinels.
+
+        Stage-0 can finish on a step that emits no new audio frame. The stage
+        input processor forwards that as an empty payload with finished=true.
+        If a request never acquired a stream slot, do not offline-decode its
+        buffered codes here: streaming requests must only emit client deltas.
+        """
+        session = self._stream_session
+        outputs: dict[int, torch.Tensor] = {}
+        if session is None:
+            return outputs
+        for i, info in enumerate(info_list):
+            if not isinstance(info, dict):
+                continue
+            meta = (info.get("meta", {}) or {}) if isinstance(info.get("meta", {}), dict) else {}
+            if not bool(meta.get("codec_streaming", False)):
+                continue
+            finished = bool(meta.get("stream_finished", meta.get("finished", False)))
+            if not finished:
+                continue
+            req_key = self._runtime_request_key(info, meta, i)
+            slot = self._stream_req_slots.get(req_key)
+            pending = req_key in self._stream_pending_codes
+            if slot is not None or pending:
+                if pending and slot is None:
+                    logger.warning(
+                        "MOSS codec stream request %s finished before a codec stream slot became available; "
+                        "dropping buffered codes instead of offline-decoding a non-delta waveform.",
+                        req_key,
+                    )
+                self._finish_stream_request(req_key, session, slot)
+        return outputs
+
     @staticmethod
-    def _runtime_value(info: Any, meta: dict[str, Any], name: str, default: Any = None) -> Any:
-        if name in meta:
-            return meta[name]
-        if isinstance(info, dict) and name in info:
-            return info[name]
+    def _normalize_seq_token_counts(value: Any) -> list[int] | None:
+        if value is None:
+            return None
+        if not isinstance(value, (list, tuple)):
+            raise TypeError(
+                "MossTTS codec expects seq_token_counts to be a list/tuple of per-request token counts, "
+                f"got {type(value).__name__}."
+            )
+        counts = [int(item) for item in value]
+        if not counts:
+            return None
+        for count in counts:
+            if count < 0:
+                raise ValueError(f"MossTTS codec seq_token_counts must be non-negative, got {counts}.")
+        return counts
+
+    def _runtime_request_key(self, info: Any, meta: dict[str, Any], index: int) -> str:
+        for value in (
+            meta.get("req_id"),
+            info.get("request_id") if isinstance(info, dict) else None,
+        ):
+            if isinstance(value, (list, tuple)):
+                value = value[0] if value else None
+            if value is not None:
+                return str(value)
+        return f"moss-codec-stream-{index}"
+
+    def _empty_audio(self) -> torch.Tensor:
+        if self._n_channels > 1:
+            return torch.zeros((self._n_channels, 0), dtype=torch.float32)
+        return torch.zeros((0,), dtype=torch.float32)
+
+    def _ensure_stream_session(self) -> _MossCodecStreamSession | None:
+        if self._codec is None:
+            return None
+        if self._stream_session is not None:
+            return self._stream_session
+        slots = self._stream_slots
+        if slots <= 0:
+            scheduler_cfg = getattr(self.vllm_config, "scheduler_config", None)
+            slots = int(getattr(scheduler_cfg, "max_num_seqs", 1) or 1)
+        self._stream_session = _MossCodecStreamSession(
+            self._codec,
+            stream_slots=max(1, slots),
+            n_vq=self._n_vq,
+        )
+        return self._stream_session
+
+    def _decode_streaming_batch(
+        self,
+        items: list[tuple[int, str, torch.Tensor, bool]],
+    ) -> dict[int, torch.Tensor]:
+        session = self._ensure_stream_session()
+        if session is None:
+            return {}
+
+        outputs: dict[int, torch.Tensor] = {}
+        grouped: dict[int, list[tuple[int, str, int, torch.Tensor, bool]]] = {}
+        max_step_frames = max(1, int(self._stream_max_step_frames))
+
+        for output_index, request_id, codes_nq_t, finished in items:
+            pending = self._stream_pending_codes.get(request_id)
+            slot = self._stream_req_slots.get(request_id)
+            if slot is None:
+                slot = session.acquire()
+                if slot is None:
+                    self._append_stream_pending(request_id, codes_nq_t)
+                    if request_id not in self._stream_starved_reqs:
+                        logger.warning(
+                            "MOSS codec streaming slots exhausted; buffering %s until a stream slot is available.",
+                            request_id,
+                        )
+                        self._stream_starved_reqs.add(request_id)
+                    if finished:
+                        logger.warning(
+                            "MOSS codec stream request %s finished before a codec stream slot became available; "
+                            "dropping buffered codes instead of offline-decoding a non-delta waveform.",
+                            request_id,
+                        )
+                        self._finish_stream_request(request_id, session, None)
+                    continue
+                self._stream_req_slots[request_id] = slot
+
+            if pending:
+                self._append_stream_pending(request_id, codes_nq_t)
+                replay_codes = self._pop_stream_pending(request_id)
+                wav = self._decode_stream_slot_sequence(session, slot, replay_codes)
+                if wav is not None:
+                    outputs[output_index] = wav
+                if finished:
+                    self._finish_stream_request(request_id, session, slot)
+                continue
+
+            if int(codes_nq_t.shape[1]) > max_step_frames:
+                wav = self._decode_stream_slot_sequence(session, slot, codes_nq_t)
+                if wav is not None:
+                    outputs[output_index] = wav
+                if finished:
+                    self._finish_stream_request(request_id, session, slot)
+                continue
+
+            grouped.setdefault(int(codes_nq_t.shape[1]), []).append(
+                (output_index, request_id, slot, codes_nq_t, finished)
+            )
+
+        for group in grouped.values():
+            plan = {slot: codes_nq_t for _, _, slot, codes_nq_t, _ in group}
+            decoded = session.step(plan)
+            for output_index, request_id, slot, _, finished in group:
+                wav = decoded.get(slot)
+                if wav is not None:
+                    outputs[output_index] = wav
+                if finished:
+                    self._finish_stream_request(request_id, session, slot)
+
+        return outputs
+
+    def _append_stream_pending(self, request_id: str, codes_nq_t: torch.Tensor) -> None:
+        self._stream_pending_codes.setdefault(request_id, []).append(
+            codes_nq_t.detach().to("cpu", torch.long).contiguous()
+        )
+
+    def _pop_stream_pending(self, request_id: str) -> torch.Tensor:
+        pending = self._stream_pending_codes.pop(request_id, [])
+        if not pending:
+            return torch.empty((self._n_vq, 0), dtype=torch.long)
+        return torch.cat(pending, dim=1).contiguous()
+
+    def _decode_stream_slot_sequence(
+        self,
+        session: _MossCodecStreamSession,
+        slot: int,
+        codes_nq_t: torch.Tensor,
+    ) -> torch.Tensor | None:
+        if codes_nq_t.numel() == 0:
+            return None
+        max_step_frames = max(1, int(self._stream_max_step_frames))
+        parts: list[torch.Tensor] = []
+        for start in range(0, int(codes_nq_t.shape[1]), max_step_frames):
+            chunk = codes_nq_t[:, start : start + max_step_frames]
+            decoded = session.step({slot: chunk})
+            wav = decoded.get(slot)
+            if wav is not None:
+                parts.append(wav)
+        if not parts:
+            return None
+        return torch.cat(parts, dim=-1) if len(parts) > 1 else parts[0]
+
+    def _finish_stream_request(
+        self,
+        request_id: str,
+        session: _MossCodecStreamSession,
+        slot: int | None,
+    ) -> None:
+        if slot is not None:
+            session.release(slot)
+        self._stream_req_slots.pop(request_id, None)
+        self._stream_pending_codes.pop(request_id, None)
+        self._stream_starved_reqs.discard(request_id)
+
+    def on_requests_finished(self, finished_req_ids: set[str] | list[str]) -> None:
+        """Release codec streaming slots when requests finish outside payload flow.
+
+        Normal streaming completion releases slots from ``_decode_streaming_batch``
+        when the Stage-0 payload carries ``finished=True``. Client disconnects
+        and engine-side aborts can finish a request without delivering that
+        terminal payload, so the runner calls this hook from its finished-request
+        path to avoid leaking stream slots and buffered codes.
+        """
+        session = self._stream_session
+        for req_id in finished_req_ids:
+            request_id = str(req_id)
+            slot = self._stream_req_slots.get(request_id)
+            has_state = (
+                slot is not None or request_id in self._stream_pending_codes or request_id in self._stream_starved_reqs
+            )
+            if not has_state:
+                continue
+            if session is not None:
+                self._finish_stream_request(request_id, session, slot)
+            else:
+                self._stream_req_slots.pop(request_id, None)
+                self._stream_pending_codes.pop(request_id, None)
+                self._stream_starved_reqs.discard(request_id)
+
+    def _connector_int(self, name: str, default: int = 0) -> int:
+        model_cfg = getattr(self.vllm_config, "model_config", None)
+        connector_cfg = getattr(model_cfg, "stage_connector_config", None)
+        if isinstance(connector_cfg, dict):
+            extra_cfg: dict | None = connector_cfg.get("extra", connector_cfg)
+        else:
+            extra_cfg = getattr(connector_cfg, "extra", None)
+        if isinstance(extra_cfg, dict) and name in extra_cfg:
+            return int(extra_cfg[name])
         return default
 
     # ------------------------------------------------------------------
@@ -272,14 +638,12 @@ class MossTTSCodecDecoder(nn.Module):
             (".self_attn.out_projs.0.", ".attn.out_proj."),
             (".linear1.", ".ff1."),
             (".linear2.", ".ff2."),
-            # v2 (MOSS-Audio-Tokenizer-v2) uses different submodule names for
-            # the same attention/FFN sublayers — no trailing "s"/index on
-            # in_proj/out_proj, and an `ffn.{0,2}` Sequential instead of
-            # separate linear1/linear2 attributes.
-            (".self_attn.in_proj.", ".attn.in_proj."),
-            (".self_attn.out_proj.", ".attn.out_proj."),
-            (".ffn.0.", ".ff1."),
-            (".ffn.2.", ".ff2."),
+            # v2 checkpoint names use singular in_proj/out_proj and ffn.{0,2};
+            # the vendored module keeps the original MOSS layer names.
+            (".self_attn.in_proj.", ".self_attn.in_projs.0."),
+            (".self_attn.out_proj.", ".self_attn.out_projs.0."),
+            (".ffn.0.", ".linear1."),
+            (".ffn.2.", ".linear2."),
             (".layer_scale_1.", ".ls1."),
             (".layer_scale_2.", ".ls2."),
             (".input_proj.", ".in_proj."),
@@ -292,21 +656,39 @@ class MossTTSCodecDecoder(nn.Module):
                     return name.replace(src, dst)
             return name
 
-        loaded_total = 0
+        loaded_names: set[str] = set()
         skipped: list[str] = []
+        shape_mismatches: list[tuple[str, str, tuple[int, ...], tuple[int, ...]]] = []
         for name, tensor in codec_weights:
             # Try direct name first (e.g. ``quantizer.input_proj.*`` exists
             # under the same name in both layouts), then the remap (transformer
             # submodules need ``.linear1.``→``.ff1.`` etc.).
             tgt = name if name in params_dict else _remap(name)
             if tgt in params_dict:
+                expected_shape = tuple(params_dict[tgt].shape)
+                actual_shape = tuple(tensor.shape)
+                if expected_shape != actual_shape:
+                    shape_mismatches.append((name, tgt, actual_shape, expected_shape))
+                    continue
                 default_weight_loader(params_dict[tgt], tensor)
-                loaded_total += 1
+                loaded_names.add(tgt)
             else:
                 skipped.append(name)
+
+        missing = sorted(set(params_dict) - loaded_names)
+        if missing or skipped or shape_mismatches:
+            raise RuntimeError(
+                "MOSS Audio Tokenizer weights were not fully loaded: "
+                f"loaded={len(loaded_names)}/{len(params_dict)} "
+                f"missing={len(missing)} skipped={len(skipped)} "
+                f"shape_mismatches={len(shape_mismatches)}; "
+                f"first_missing={missing[:5]} "
+                f"first_skipped={skipped[:5]} "
+                f"first_shape_mismatches={shape_mismatches[:3]}"
+            )
         logger.info(
             "MOSS Audio Tokenizer weights: loaded=%d/%d skipped=%d (first skipped: %s)",
-            loaded_total,
+            len(loaded_names),
             len(params_dict),
             len(skipped),
             skipped[:3] if skipped else "none",
@@ -316,7 +698,15 @@ class MossTTSCodecDecoder(nn.Module):
         codec.to(device=device, dtype=torch.float32)
         codec.eval()
         self._codec = codec
-        self._n_channels = int(getattr(codec_cfg, "number_channels", 1) or 1)
+        inferred_channels = 2 if "v2" in codec_path.lower() else 1
+        self._n_channels = int(
+            getattr(
+                codec_cfg,
+                "number_channels",
+                getattr(codec_cfg, "num_channels", inferred_channels),
+            )
+            or inferred_channels
+        )
         self._sr_tensor = torch.tensor(int(codec_cfg.sampling_rate), dtype=torch.int32)
 
         logger.info(
@@ -336,16 +726,13 @@ class MossTTSCodecDecoder(nn.Module):
 
     def _build_codec(self, codec_path: str) -> tuple[Any, nn.Module]:
         try:
-            from transformers import AutoConfig, AutoModel
-
-            codec_cfg = AutoConfig.from_pretrained(codec_path, trust_remote_code=True)
-            codec = AutoModel.from_config(codec_cfg, trust_remote_code=True)
-            logger.info("Using MOSS Audio Tokenizer remote-code classes from %s", codec_path)
+            codec_cfg = MossAudioTokenizerV2Config.from_pretrained(codec_path)
+            codec = MossAudioTokenizerV2Model(codec_cfg)
+            logger.info("Using vendored MOSS Audio Tokenizer v2 classes from %s", codec_path)
             return codec_cfg, codec
         except Exception:
             logger.exception(
-                "Failed to instantiate official MOSS Audio Tokenizer via HF remote code; "
-                "falling back to vendored codec."
+                "Failed to instantiate vendored MOSS Audio Tokenizer v2; falling back to legacy vendored codec."
             )
 
         codec_cfg = MossAudioTokenizerConfig.from_pretrained(codec_path)

vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_local.py

@@ -126,6 +126,7 @@ def _sample_token(
     top_k: int,
     top_p: float,
     do_sample: bool,
+    generator: torch.Generator | None = None,
 ) -> torch.Tensor:
     """Top-k + top-p sampling (matches upstream's ``sample_token`` for the
     inference branch).
@@ -152,7 +153,7 @@ def _sample_token(
 
     probs = F.softmax(logits, dim=-1)
     flat = probs.reshape(-1, probs.shape[-1])
-    sampled = torch.multinomial(flat, num_samples=1).reshape(probs.shape[:-1])
+    sampled = torch.multinomial(flat, num_samples=1, generator=generator).reshape(probs.shape[:-1])
     return sampled
 
 

vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_local_depth.py

@@ -160,6 +160,7 @@ class MossTTSLocalDepthTransformer(nn.Module):
         text_top_p: float = 1.0,
         repetition_penalty: float = 1.0,
         history_per_codebook: list[list[int]] | None = None,
+        generator: torch.Generator | None = None,
     ) -> tuple[torch.Tensor, torch.Tensor]:
         """Generate one audio frame for batch B.
 
@@ -183,7 +184,16 @@ class MossTTSLocalDepthTransformer(nn.Module):
         local_hidden = hidden[:, 0, :]
 
         binary_logits = local_text_lm_head(local_hidden).float()
-        binary_choice = _sample_token(binary_logits, text_temperature, text_top_k, text_top_p, do_sample)
+        # This is a binary continue/stop gate. The checkpoint expects sampling
+        # here; greedy argmax is biased toward "continue" and may never stop.
+        binary_choice = _sample_token(
+            binary_logits,
+            text_temperature,
+            text_top_k,
+            text_top_p,
+            do_sample,
+            generator=generator,
+        )
         should_continue = binary_choice.eq(0)
         import os as _os
 
@@ -209,7 +219,14 @@ class MossTTSLocalDepthTransformer(nn.Module):
                     pos = sel > 0
                     sel = torch.where(pos, sel / repetition_penalty, sel * repetition_penalty)
                     channel_logits.index_copy_(-1, hist_t, sel)
-            channel_token = _sample_token(channel_logits, temperature, top_k, top_p, do_sample)
+            channel_token = _sample_token(
+                channel_logits,
+                temperature,
+                top_k,
+                top_p,
+                do_sample,
+                generator=generator,
+            )
             codes[:, channel_index] = channel_token
 
             if channel_index + 1 < n_vq:

vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_talker.py

@@ -1139,26 +1139,19 @@ class MossTTSRealtimeTalkerForGeneration(nn.Module):
 
                 info["audio_codes"] = {"current": new_codes, "accumulated": updated_acc}
                 state["step"] = int(state.get("step", 0)) + 1
-                per_req_codes[i] = updated_acc
+                per_req_codes[i] = new_codes.unsqueeze(0)
                 have_codes = True
 
         if not have_codes:
             return OmniOutput(text_hidden_states=hidden, multimodal_outputs={})
 
-        # Emit one accumulated-codes tensor per request as a list; the splitter
-        # routes ``element[idx]`` to request idx, so each Stage-1 request decodes
-        # its own codes instead of request 0's.
-        # The realtime variant emits raw codes (no delay pattern), so we signal
-        # that to the chunk processor via a 1-D bool tensor in ``meta.finished``
-        # (already a tensor field on the schema). The processor reads its truth
-        # value as "skip de-delay" only for realtime; the delay path doesn't
-        # populate ``meta`` at all.
-        device = hidden.device
+        # Emit one current-frame tensor per request as a list; the splitter
+        # routes ``element[idx]`` to request idx, so each Stage-1 request receives
+        # only its own newly generated raw codec row.
         return OmniOutput(
             text_hidden_states=hidden,
             multimodal_outputs={
                 "codes": {"audio": per_req_codes},
-                "meta": {"finished": torch.tensor([True], dtype=torch.bool, device=device)},
             },
         )
 
@@ -1266,7 +1259,7 @@ class MossTTSLocalTalkerForGeneration(nn.Module):
     * Per-step audio generation runs ``MossTTSLocalDepthTransformer``
       (1-layer GPT2-style local transformer, interleaved RoPE, ``n_vq=12``
       codebooks -- structurally different from Realtime's Qwen3-style
-      ``CodePredictorBaseModel``) inside ``make_omni_output``. Stop
+      ``CodePredictorBaseModel``) in the runner-owned ``talker_mtp`` path. Stop
       condition: the local transformer's binary head selects the
       "audio_end" candidate (logits index 1).
     """
@@ -1327,6 +1320,8 @@ class MossTTSLocalTalkerForGeneration(nn.Module):
         # Stacked embedding cache built after load_weights() for vectorised
         # additive-fusion gathers (avoids an n_vq-iteration Python loop).
         self._stacked_audio_emb_w: torch.Tensor | None = None
+        self.mtp_hidden_size = hidden_size
+        self.talker_mtp_accepts_req_infos = True
 
         self.gpu_resident_buffer_keys: set[tuple[str, str]] = {
             ("audio_codes", "current"),
@@ -1384,6 +1379,15 @@ class MossTTSLocalTalkerForGeneration(nn.Module):
             num_rows,
         ):
             if state is not None and state.get("is_stopping"):
+                forced_im_end_count = int(state.get("_forced_im_end_count", 0))
+                if forced_im_end_count == 1:
+                    logger.warning(
+                        "MOSS-TTS Local request is still running after "
+                        "im_end_token_id=%d was forced once; verify stage-0 "
+                        "stop_token_ids/eos includes this id.",
+                        self.im_end_token_id,
+                    )
+                state["_forced_im_end_count"] = forced_im_end_count + 1
                 logits[row_start:row_end, self.im_end_token_id] = 0.0
             else:
                 logits[row_start:row_end, self.audio_assistant_slot_token_id] = 0.0
@@ -1455,38 +1459,159 @@ class MossTTSLocalTalkerForGeneration(nn.Module):
                 else torch.full((self.n_vq,), self.audio_pad_token_id, dtype=torch.long, device=device)
             )
 
-            max_new_frames_req = info_dict.get("max_new_frames")
-            if isinstance(max_new_frames_req, (list, tuple)) and max_new_frames_req:
-                max_new_frames_req = max_new_frames_req[0]
-            try:
-                max_new_frames = int(max_new_frames_req) if max_new_frames_req is not None else -1
-            except (TypeError, ValueError):
-                max_new_frames = -1
+            max_new_frames = info_dict.get("max_new_frames", [-1])[0]
 
             info_update: dict[str, Any] = {
-                "audio_state": {"is_stopping": False, "step": 0, "max_new_frames": max_new_frames},
-                "audio_codes": {"current": current_codes},
+                "audio_state": {
+                    "is_stopping": False,
+                    "step": 0,
+                    "max_new_frames": max_new_frames,
+                },
+                "audio_codes": {"current": current_codes, "emit": False},
                 "ref_offset": ref_offset + span_len,
             }
             return input_ids, embeds, info_update
 
-        # Decode step: input_ids is the forced text token from compute_logits.
+        # Decode step: return the plain forced text-token embedding. The local
+        # depth transformer runs in OmniGPUModelRunner.talker_mtp before the
+        # backbone forward; it adds the freshly sampled audio-code embedding
+        # and stores the generated code frame in model_intermediate_buffer.
         text_embed = self.model.embed_tokens(input_ids.reshape(1))  # (1, H)
-        prev_codes = (info_dict.get("audio_codes", {}) or {}).get("current")
-        if isinstance(prev_codes, torch.Tensor) and prev_codes.numel() == self.n_vq:
-            codes = prev_codes.to(device=device, dtype=torch.long).unsqueeze(0)  # (1, n_vq)
-            combined = text_embed + self._audio_embed(codes)
+        hs = info_dict.get("hidden_states", {}) or {}
+        last_hidden = hs.get("last")
+        if isinstance(last_hidden, torch.Tensor) and last_hidden.numel() > 0:
+            mtp_hidden = last_hidden.to(device=device, dtype=text_embed.dtype).reshape(1, -1)
         else:
-            combined = text_embed
-        return input_ids, combined, {}
+            mtp_hidden = torch.zeros((1, self.hidden_size), device=device, dtype=text_embed.dtype)
+        return input_ids, text_embed, {"mtp_inputs": (mtp_hidden, torch.zeros_like(mtp_hidden))}
 
     def postprocess(self, hidden_states: torch.Tensor, **_: Any) -> dict[str, Any]:
         if hidden_states.numel() == 0:
             return {}
         return {"hidden_states": {"last": hidden_states[-1].detach()}}
 
+    @torch.inference_mode()
+    def talker_mtp(
+        self,
+        input_ids: torch.Tensor,
+        input_embeds: torch.Tensor,
+        last_talker_hidden: torch.Tensor,
+        text_step: torch.Tensor,
+        do_sample: bool | None = None,
+        temperature: float | None = None,
+        top_k: int | None = None,
+        top_p: float | None = None,
+        generator: torch.Generator | None = None,
+        req_infos: list[dict[str, Any]] | None = None,
+        **_: Any,
+    ) -> tuple[torch.Tensor, None]:
+        bsz = int(input_embeds.shape[0])
+        dev = input_embeds.device
+        input_embeds_out = input_embeds.reshape(bsz, -1).clone()
+        req_infos = req_infos or [{} for _ in range(bsz)]
+
+        def write_update(info: dict[str, Any], update: dict[str, Any]) -> None:
+            for key, value in update.items():
+                if isinstance(value, dict):
+                    target = info.setdefault(key, {})
+                    target.update(value)
+                else:
+                    info[key] = value
+
+        # Match MOSS-TTS Local v1.5's model-card/SGLang defaults. The generic
+        # stage SamplingParams are for the backbone token loop and should not
+        # silently retune the local depth transformer's audio-channel sampler.
+        audio_temperature = 1.7
+        audio_top_k = 25
+        audio_top_p = 0.8
+        do_sample_val = True if do_sample is None else bool(do_sample)
+
+        for i in range(bsz):
+            info = req_infos[i] if i < len(req_infos) and isinstance(req_infos[i], dict) else {}
+            state = dict(info.get("audio_state", {}) or {})
+            step = int(state.get("step", 0))
+            max_new_frames = int(state.get("max_new_frames", -1))
+            acc = (info.get("audio_codes", {}) or {}).get("accumulated")
+
+            if state.get("is_stopping"):
+                write_update(
+                    info,
+                    {
+                        "audio_state": state,
+                        "audio_codes": {
+                            "current": input_embeds.new_empty((0, self.n_vq), dtype=torch.long),
+                            "emit": False,
+                        },
+                    },
+                )
+                continue
+
+            rep_window = 50
+            if isinstance(acc, torch.Tensor) and acc.numel() > 0:
+                tail = acc[-rep_window:].long().cpu().tolist()
+                hist_per_cb = [[row[cb] for row in tail] for cb in range(self.n_vq)]
+            else:
+                hist_per_cb = [[] for _ in range(self.n_vq)]
+
+            should_continue_t, new_codes_b = self.local_transformer.generate_frame(
+                last_talker_hidden[i : i + 1],
+                self.audio_lm_heads,
+                self.audio_embeddings,
+                self.local_text_lm_head,
+                n_vq=self.n_vq,
+                do_sample=do_sample_val,
+                temperature=audio_temperature,
+                top_k=audio_top_k,
+                top_p=audio_top_p,
+                repetition_penalty=1.0,
+                history_per_codebook=hist_per_cb,
+                generator=generator,
+            )
+            should_continue = bool(should_continue_t[0].item())
+            new_codes = new_codes_b.squeeze(0).to(device=dev, dtype=torch.long)
+
+            if not should_continue or (0 <= max_new_frames <= step):
+                state["is_stopping"] = True
+                state["step"] = step + 1
+                write_update(
+                    info,
+                    {
+                        "audio_state": state,
+                        "audio_codes": {
+                            "current": input_embeds.new_empty((0, self.n_vq), dtype=torch.long),
+                            "accumulated": acc,
+                            "emit": False,
+                        },
+                    },
+                )
+                continue
+
+            if isinstance(acc, torch.Tensor) and acc.numel() > 0:
+                updated_acc = torch.cat([acc.to(dev), new_codes.unsqueeze(0)], dim=0)
+            else:
+                updated_acc = new_codes.unsqueeze(0)
+
+            state["step"] = step + 1
+            current = new_codes.unsqueeze(0)
+            input_embeds_out[i : i + 1] = input_embeds_out[i : i + 1] + self._audio_embed(current).to(
+                dtype=input_embeds_out.dtype
+            )
+            write_update(
+                info,
+                {
+                    "audio_state": state,
+                    "audio_codes": {
+                        "current": current,
+                        "accumulated": updated_acc,
+                        "emit": True,
+                    },
+                },
+            )
+
+        return input_embeds_out, None
+
     # ------------------------------------------------------------------
-    # Per-step audio generation via local transformer
+    # Package runner-generated audio frames
     # ------------------------------------------------------------------
 
     def make_omni_output(
@@ -1511,91 +1636,41 @@ class MossTTSLocalTalkerForGeneration(nn.Module):
         self._batch_state = [(info["audio_state"] if isinstance(info, dict) else {}) for info in info_dicts]
         self._batch_state_spans = kwargs.get("request_token_spans")
 
-        spans = kwargs.get("request_token_spans")
-        per_req_codes: list[torch.Tensor] = [hidden.new_empty((0, self.n_vq), dtype=torch.long) for _ in info_dicts]
+        per_req_codes: list[torch.Tensor] = []
         have_codes = False
-
-        if hidden.numel() > 0 and info_dicts:
-            num_rows = hidden.shape[0]
-            if spans is None or len(spans) != len(info_dicts):
-                raise RuntimeError(
-                    "MOSS-TTS make_omni_output requires request_token_spans "
-                    f"for {len(info_dicts)} request infos, got {0 if spans is None else len(spans)}"
-                )
-            for i, info in enumerate(info_dicts):
-                if not isinstance(info, dict):
-                    continue
-                row_start, row_end = spans[i]
-                row_end = min(int(row_end), num_rows)
-                if row_start >= row_end:
-                    continue
-
-                state = info.get("audio_state", {}) or {}
-                if state.get("is_stopping"):
-                    continue  # already stopped -- no more audio frames
-
-                last_h = hidden[row_end - 1].unsqueeze(0)  # (1, H)
-
-                # Sampling parameters mirror Realtime's defaults; repetition
-                # penalty applies over a 50-frame rolling window per codebook.
-                rep_window = 50
-                acc_for_hist = (info.get("audio_codes", {}) or {}).get("accumulated")
-                if isinstance(acc_for_hist, torch.Tensor) and acc_for_hist.numel() > 0:
-                    tail = acc_for_hist[-rep_window:].long().cpu().tolist()
-                    hist_per_cb = [[row[cb] for row in tail] for cb in range(self.n_vq)]
-                else:
-                    hist_per_cb = [[] for _ in range(self.n_vq)]
-
-                should_continue_t, new_codes = self.local_transformer.generate_frame(
-                    last_h,
-                    self.audio_lm_heads,
-                    self.audio_embeddings,
-                    self.local_text_lm_head,
-                    n_vq=self.n_vq,
-                    do_sample=True,
-                    temperature=0.8,
-                    top_k=30,
-                    top_p=0.6,
-                    repetition_penalty=1.1,
-                    history_per_codebook=hist_per_cb,
-                )
-                should_continue = bool(should_continue_t[0].item())
-                new_codes = new_codes.squeeze(0)  # (n_vq,)
-
-                step = int(state.get("step", 0))
-                max_new_frames = int(state.get("max_new_frames", -1))
-                if not should_continue or (0 <= max_new_frames <= step):
-                    state["is_stopping"] = True
-                    state["step"] = step + 1
-                    info["audio_codes"] = {
-                        "current": new_codes,
-                        "accumulated": (info.get("audio_codes", {}) or {}).get("accumulated"),
-                    }
-                    continue  # don't append the stop frame to accumulated
-
-                acc = (info.get("audio_codes", {}) or {}).get("accumulated")
-                if isinstance(acc, torch.Tensor) and acc.numel() > 0:
-                    updated_acc = torch.cat([acc.to(new_codes.device), new_codes.unsqueeze(0)], dim=0)
-                else:
-                    updated_acc = new_codes.unsqueeze(0)
-
-                info["audio_codes"] = {"current": new_codes, "accumulated": updated_acc}
-                state["step"] = step + 1
-                per_req_codes[i] = updated_acc
+        for info in info_dicts:
+            current = (info.get("audio_codes", {}) or {}).get("current") if isinstance(info, dict) else None
+            emit = bool((info.get("audio_codes", {}) or {}).get("emit")) if isinstance(info, dict) else False
+            if emit and isinstance(current, torch.Tensor) and current.numel() > 0:
+                codes = current.to(device=hidden.device, dtype=torch.long)
+                if codes.dim() == 1:
+                    codes = codes.unsqueeze(0)
+                per_req_codes.append(codes)
                 have_codes = True
+            else:
+                per_req_codes.append(hidden.new_empty((0, self.n_vq), dtype=torch.long))
 
         if not have_codes:
             return OmniOutput(text_hidden_states=hidden, multimodal_outputs={})
 
-        # Local-v1.5 emits raw codes (no delay pattern), so signal that to the
-        # chunk processor the same way Realtime does, via ``meta.finished``
-        # (read as "skip de-delay" by ``talker2codec_async_chunk``).
-        device = hidden.device
+        # Local-v1.5 emits the current raw code frame (no delay pattern); the
+        # pipeline routes those rows through ``talker2codec_raw_async_chunk``.
+        per_req_ref_codes: list[torch.Tensor] = []
+        for info in info_dicts:
+            ref_codes = (info.get("codes", {}) or {}).get("ref") if isinstance(info, dict) else None
+            if isinstance(ref_codes, torch.Tensor) and ref_codes.numel() > 0:
+                ref_codes = ref_codes.detach()
+                if ref_codes.dim() == 1 and ref_codes.numel() % self.n_vq == 0:
+                    ref_codes = ref_codes.view(-1, self.n_vq)
+                if ref_codes.dim() == 2 and ref_codes.shape[1] == self.n_vq:
+                    per_req_ref_codes.append(ref_codes.to(device=hidden.device, dtype=torch.long))
+                    continue
+            per_req_ref_codes.append(hidden.new_empty((0, self.n_vq), dtype=torch.long))
+
         return OmniOutput(
             text_hidden_states=hidden,
             multimodal_outputs={
-                "codes": {"audio": per_req_codes},
-                "meta": {"finished": torch.tensor([True], dtype=torch.bool, device=device)},
+                "codes": {"audio": per_req_codes, "ref": per_req_ref_codes},
             },
         )
 

vllm_omni/model_executor/models/moss_tts/pipeline.py

@@ -31,7 +31,7 @@ MOSS_TTS_PIPELINE = PipelineConfig(
             input_sources=(),
             owns_tokenizer=True,
             engine_output_type="latent",
-            async_chunk_process_next_stage_input_func=(f"{_PROC}.talker2codec_async_chunk"),
+            async_chunk_process_next_stage_input_func=(f"{_PROC}.talker2codec_delay_async_chunk"),
             sampling_constraints={
                 "detokenize": False,
             },
@@ -62,7 +62,7 @@ MOSS_TTS_REALTIME_PIPELINE = PipelineConfig(
             input_sources=(),
             owns_tokenizer=True,
             engine_output_type="latent",
-            async_chunk_process_next_stage_input_func=(f"{_PROC}.talker2codec_async_chunk"),
+            async_chunk_process_next_stage_input_func=(f"{_PROC}.talker2codec_raw_async_chunk"),
             sampling_constraints={"detokenize": False},
         ),
         StagePipelineConfig(
@@ -91,8 +91,11 @@ MOSS_TTS_LOCAL_PIPELINE = PipelineConfig(
             input_sources=(),
             owns_tokenizer=True,
             engine_output_type="latent",
-            async_chunk_process_next_stage_input_func=(f"{_PROC}.talker2codec_async_chunk"),
-            sampling_constraints={"detokenize": False},
+            async_chunk_process_next_stage_input_func=(f"{_PROC}.talker2codec_raw_async_chunk"),
+            sampling_constraints={
+                "detokenize": False,
+                "stop_token_ids": [151645],
+            },
         ),
         StagePipelineConfig(
             stage_id=1,

vllm_omni/model_executor/stage_input_processors/moss_tts.py

@@ -5,6 +5,8 @@
 
 from __future__ import annotations
 
+from collections import defaultdict
+from collections.abc import Mapping
 from typing import Any
 
 import torch
@@ -15,6 +17,8 @@ from vllm_omni.data_entry_keys import CodesStruct, MetaStruct, OmniPayloadStruct
 
 logger = init_logger(__name__)
 
+_MOSS_AUDIO_PAD_CODE = 1024
+
 
 # ---------------------------------------------------------------------------
 # Helpers
@@ -90,13 +94,13 @@ def talker2codec(
 # ---------------------------------------------------------------------------
 
 
-def talker2codec_async_chunk(
+def talker2codec_delay_async_chunk(
     transfer_manager: Any,
     multimodal_output: dict[str, Any] | None,
     request: Any,
     is_finished: bool = False,
 ) -> OmniPayloadStruct | None:
-    """Emit accumulated audio codes to Stage 1 as they arrive from Stage 0.
+    """Emit delay-patterned accumulated audio codes to Stage 1.
 
     State is maintained in ``transfer_manager`` keyed by request ID.
     A chunk is forwarded to Stage 1 when either:
@@ -138,14 +142,6 @@ def talker2codec_async_chunk(
                     req_state["accumulated"] = new_rows
                 else:
                     req_state["accumulated"] = torch.cat([req_state["accumulated"], new_rows], dim=0)
-        # Realtime variant emits raw (un-delay-patterned) codes; record that
-        # so the emit step below skips the apply_de_delay_pattern transform.
-        # Realtime sets ``meta.finished`` to a 1-D bool tensor; the delay
-        # variant leaves ``meta`` empty.
-        meta_in = pooling_output.get("meta", {}) or {}
-        flag = meta_in.get("finished")
-        if isinstance(flag, torch.Tensor) and flag.numel() >= 1 and bool(flag.reshape(-1)[0].item()):
-            req_state["skip_dedelay"] = True
 
     acc = req_state["accumulated"]
     if acc is None or acc.numel() == 0:
@@ -181,15 +177,12 @@ def talker2codec_async_chunk(
     #   1. de-delay   ``(T+nq-1, nq)`` → ``(T, nq)``
     #   2. drop rows that are entirely pad (separators between audio segments,
     #      and the leading text-mode rows that precede the first audio_start).
-    # The realtime talker emits raw codes (no delay), so steps 1 is skipped.
     chunk_codes_long = chunk_codes.to(torch.long).cpu().contiguous()  # (T_chunk, NQ)
     nq = int(chunk_codes_long.shape[1])
     t_chunk = int(chunk_codes_long.shape[0])
     audio_pad_code = 1024  # MOSS-TTS audio_pad_code; same value across variants.
 
-    if req_state.get("skip_dedelay"):
-        de_delayed = chunk_codes_long
-    elif t_chunk > nq:
+    if t_chunk > nq:
         de_delayed = chunk_codes_long.new_zeros((t_chunk - nq + 1, nq))
         for i in range(nq):
             de_delayed[:, i] = chunk_codes_long[i : i + de_delayed.shape[0], i]
@@ -224,9 +217,133 @@ def talker2codec_async_chunk(
         codes=CodesStruct(audio=codec_flat),
         meta=MetaStruct(
             left_context_size=left_context,
-            finished=torch.tensor(is_finished, dtype=torch.bool),
+            finished=torch.tensor(bool(is_finished), dtype=torch.bool),
         ),
     )
 
 
-__all__ = ["talker2codec", "talker2codec_async_chunk"]
+def talker2codec_raw_async_chunk(
+    transfer_manager: Any,
+    multimodal_output: dict[str, Any] | None,
+    request: Any,
+    is_finished: bool = False,
+) -> OmniPayloadStruct | None:
+    """Async processor for MOSS-TTS Local/Realtime raw codec rows.
+
+    Stage 0 emits newly generated raw codec rows shaped ``[T, n_vq]`` (normally
+    ``[1, n_vq]`` per decode step). This processor buffers those new rows until
+    a codec chunk is ready, then forwards the chunk to Stage 1. No delay-pattern
+    de-delay is applied on this path.
+    """
+    external_req_id = getattr(request, "external_req_id", None)
+    req_id = str(external_req_id if external_req_id is not None else getattr(request, "request_id", id(request)))
+
+    if not hasattr(transfer_manager, "code_prompt_token_ids"):
+        transfer_manager.code_prompt_token_ids = defaultdict(list)
+    if not hasattr(transfer_manager, "request_payload"):
+        transfer_manager.request_payload = {}
+    if not hasattr(transfer_manager, "put_req_chunk"):
+        transfer_manager.put_req_chunk = defaultdict(int)
+
+    pending_frames = transfer_manager.code_prompt_token_ids[req_id]
+
+    if isinstance(multimodal_output, Mapping):
+        codes_dict = multimodal_output.get("codes", {}) or {}
+        new_frames = codes_dict.get("audio")
+        if isinstance(new_frames, torch.Tensor) and new_frames.numel() > 0:
+            frames_cpu = new_frames.detach().to("cpu", torch.long).contiguous()
+            if frames_cpu.ndim == 1:
+                frames_cpu = frames_cpu.reshape(1, -1)
+            if frames_cpu.ndim != 2:
+                raise ValueError(f"MOSS raw codec frames must be 2-D, got {tuple(frames_cpu.shape)}")
+            valid_rows = frames_cpu.ne(_MOSS_AUDIO_PAD_CODE).any(dim=1)
+            for frame in frames_cpu[valid_rows]:
+                pending_frames.append(frame.clone())
+        # Raw/local streaming should mirror the non-streaming path: the codec
+        # decodes only generated audio rows. Reference audio conditions the
+        # talker, but feeding its codes into the codec streaming state adds a
+        # long first-packet prime step and changes the decoder state relative
+        # to non-streaming output.
+
+    connector = getattr(transfer_manager, "connector", None)
+    raw_cfg = getattr(connector, "config", {}) or {}
+    cfg = raw_cfg.get("extra", raw_cfg) if isinstance(raw_cfg, dict) else {}
+    cfg = cfg if isinstance(cfg, dict) else {}
+    chunk_frames = int(cfg.get("codec_chunk_frames", 15) or 15)
+    initial_chunk_frames = int(cfg.get("initial_codec_chunk_frames") or 0)
+    if chunk_frames <= 0:
+        raise ValueError(f"codec_chunk_frames must be positive for MOSS raw streaming, got {chunk_frames}")
+    if initial_chunk_frames < 0:
+        raise ValueError(f"initial_codec_chunk_frames must be non-negative, got {initial_chunk_frames}")
+    if initial_chunk_frames > chunk_frames:
+        logger.warning(
+            "initial_codec_chunk_frames=%d > codec_chunk_frames=%d, clamping.",
+            initial_chunk_frames,
+            chunk_frames,
+        )
+        initial_chunk_frames = chunk_frames
+
+    pending = len(pending_frames)
+    emitted_any = int(transfer_manager.put_req_chunk.get(req_id, 0)) > 0
+    threshold = initial_chunk_frames if initial_chunk_frames > 0 and not emitted_any else chunk_frames
+    if pending <= 0:
+        if is_finished:
+            transfer_manager.code_prompt_token_ids.pop(req_id, None)
+            transfer_manager.request_payload.pop(req_id, None)
+            return OmniPayloadStruct(
+                # A non-empty sentinel is required so Stage-1 is scheduled.
+                # ``code_flat_numel=0`` tells the codec this is a control-only
+                # finish packet, not an audio code.
+                codes=CodesStruct(audio=torch.tensor([0], dtype=torch.long)),
+                meta=MetaStruct(
+                    req_id=[req_id],
+                    left_context_size=0,
+                    codec_streaming=True,
+                    codec_chunk_frames=0,
+                    codec_left_context_frames=0,
+                    code_flat_numel=0,
+                    stream_finished=torch.tensor(True, dtype=torch.bool),
+                    finished=torch.tensor(True, dtype=torch.bool),
+                ),
+                request_id=req_id,
+            )
+        return None
+    if not is_finished and pending < threshold:
+        return None
+
+    emit_frames = pending if is_finished else threshold
+    chunk_rows = pending_frames[:emit_frames]
+    del pending_frames[:emit_frames]
+    chunk_codes = torch.stack(
+        [row.to(torch.long).cpu() for row in chunk_rows],
+        dim=0,
+    ).contiguous()
+    finished = bool(is_finished and len(pending_frames) == 0)
+
+    codec_flat = chunk_codes.transpose(0, 1).contiguous().reshape(-1).to(torch.long)
+
+    if finished:
+        transfer_manager.code_prompt_token_ids.pop(req_id, None)
+        transfer_manager.request_payload.pop(req_id, None)
+
+    return OmniPayloadStruct(
+        codes=CodesStruct(audio=codec_flat),
+        meta=MetaStruct(
+            req_id=[req_id],
+            left_context_size=0,
+            codec_streaming=True,
+            codec_chunk_frames=int(chunk_codes.shape[0]),
+            codec_left_context_frames=0,
+            code_flat_numel=int(codec_flat.numel()),
+            stream_finished=torch.tensor(finished, dtype=torch.bool),
+            finished=torch.tensor(finished, dtype=torch.bool),
+        ),
+        request_id=req_id,
+    )
+
+
+__all__ = [
+    "talker2codec",
+    "talker2codec_delay_async_chunk",
+    "talker2codec_raw_async_chunk",
+]

vllm_omni/worker/gpu_generation_model_runner.py

@@ -76,6 +76,10 @@ class GPUGenerationModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin
 
     def _update_request_states(self, scheduler_output: SchedulerOutput):
         # remove requests
+        # Some stateful vocoder model may need to clean the state
+        # to avoid the leak of slots when the requests have been aborted.
+        if scheduler_output.finished_req_ids and hasattr(self.model, "on_requests_finished"):
+            self.model.on_requests_finished(scheduler_output.finished_req_ids)
         for req_id in scheduler_output.finished_req_ids:
             self.input_batch.remove_request(req_id)
         scheduled_req_ids = scheduler_output.num_scheduled_tokens.keys()
@@ -789,6 +793,10 @@ class GPUGenerationModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin
             if hasattr(self.model, "get_dummy_runtime_additional_information"):
                 runtime_addi = self.model.get_dummy_runtime_additional_information(num_reqs)
                 model_kwargs["runtime_additional_information"] = runtime_addi
+            # Generation-stage code2wav models split flat input_ids using this
+            # shared contract. Dummy/profile runs use padded synthetic inputs,
+            # so expose one synthetic segment covering the actual tensor length.
+            model_kwargs["seq_token_counts"] = [int(num_tokens_padded)]
 
             if self.uses_mrope:
                 positions = self.mrope_positions.gpu[:, :num_tokens_padded]

vllm_omni/worker/gpu_model_runner.py

@@ -1821,7 +1821,9 @@ class OmniGPUModelRunner(GPUModelRunner):
         def _explicit_talker_seed(req_id: str) -> int | None:
             sampling_params = getattr(self.requests[req_id], "sampling_params", None)
             extra_args = getattr(sampling_params, "extra_args", None) if sampling_params is not None else None
-            seed = extra_args.get("qwen3_tts_request_seed") if isinstance(extra_args, dict) else None
+            seed = None
+            if isinstance(extra_args, dict):
+                seed = extra_args.get("tts_local_seed")
             return int(seed) if seed is not None else None
 
         if decode_batch_size > 1 and any(_explicit_talker_seed(req_id) is not None for req_id in decode_req_ids):
@@ -1869,6 +1871,11 @@ class OmniGPUModelRunner(GPUModelRunner):
         }
         if generator is not None:
             talker_kwargs["generator"] = generator
+        if getattr(self.model, "talker_mtp_accepts_req_infos", False):
+            talker_kwargs["req_ids"] = decode_req_ids
+            talker_kwargs["req_infos"] = [
+                self.model_intermediate_buffer.setdefault(req_id, {}) for req_id in decode_req_ids
+            ]
         with current_omni_platform.set_forward_context(
             None, self.vllm_config, cudagraph_runtime_mode=_cudagraph_mode, batch_descriptor=batch_desc
         ):
@@ -1888,8 +1895,9 @@ class OmniGPUModelRunner(GPUModelRunner):
             start_offsets = [int(self.query_start_loc.cpu[id_to_index[req_id]]) for req_id in decode_req_ids]
         for idx, (req_id, start_offset) in enumerate(zip(decode_req_ids, start_offsets, strict=True)):
             inputs_embeds[start_offset : start_offset + 1] = req_embeds[idx : idx + 1]
-            update_dict = {out_key[0]: {out_key[1]: code_predictor_codes[idx : idx + 1]}}
-            self._merge_additional_information_update(req_id, update_dict)
+            if code_predictor_codes is not None:
+                update_dict = {out_key[0]: {out_key[1]: code_predictor_codes[idx : idx + 1]}}
+                self._merge_additional_information_update(req_id, update_dict)
 
     def _model_forward(
         self,