[Model] Support MOSS-TTS-local-v1.5 (#4664)

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

benchmarks/tts/model_configs.yaml

@@ -80,6 +80,17 @@ models:
     task_extra_body:
       voice_clone: {}
 
+  # MOSS-TTS-Local-Transformer-v1.5: Qwen3 backbone + 1-layer GPT2-style local
+  # depth transformer (n_vq=12), paired with the 48 kHz stereo
+  # MOSS-Audio-Tokenizer-v2 codec. Same voice-clone request shape as the
+  # other MOSS-TTS variants.
+  OpenMOSS-Team/MOSS-TTS-Local-Transformer-v1.5:
+    supported_tasks: [voice_clone]
+    backend: openai-audio-speech
+    endpoint: /v1/audio/speech
+    task_extra_body:
+      voice_clone: {}
+
   # MOSS-TTSD synthesizes [S1]/[S2] dialogue. The seed-tts voice_clone dataset
   # has single-speaker rows; bench piggybacks on it (both speakers share the
   # same ref_audio). For real dialogue evaluation use the ttsd dataset.

examples/offline_inference/text_to_speech/moss_tts/end2end.py

@@ -1,11 +1,12 @@
 """Offline inference example for MOSS-TTS variants.
 
-Covers all five models:
-  - OpenMOSS-Team/MOSS-TTS           (8B, general TTS)
-  - OpenMOSS-Team/MOSS-TTSD-v1.0     (8B, dialogue TTS)
-  - OpenMOSS-Team/MOSS-SoundEffect   (8B, sound effect synthesis)
-  - OpenMOSS-Team/MOSS-TTS-Realtime  (1.7B, low-latency streaming)
-  - OpenMOSS-Team/MOSS-VoiceGenerator(1.7B, voice design)
+Covers six models:
+  - OpenMOSS-Team/MOSS-TTS                       (8B, general TTS)
+  - OpenMOSS-Team/MOSS-TTSD-v1.0                 (8B, dialogue TTS)
+  - OpenMOSS-Team/MOSS-SoundEffect               (8B, sound effect synthesis)
+  - OpenMOSS-Team/MOSS-TTS-Realtime              (1.7B, low-latency streaming)
+  - OpenMOSS-Team/MOSS-VoiceGenerator             (1.7B, voice design)
+  - OpenMOSS-Team/MOSS-TTS-Local-Transformer-v1.5 (n_vq=12, 48 kHz stereo)
 
 Usage
 -----
@@ -210,12 +211,18 @@ def run_tts(args: argparse.Namespace) -> None:
         "OpenMOSS-Team/MOSS-SoundEffect": "moss_sound_effect",
         "OpenMOSS-Team/MOSS-TTS-Realtime": "moss_tts_realtime",
         "OpenMOSS-Team/MOSS-VoiceGenerator": "moss_voice_generator",
+        "OpenMOSS-Team/MOSS-TTS-Local-Transformer-v1.5": "moss_tts_local",
     }
     deploy_config = deploy_map.get(args.model, "moss_tts")
 
     is_sound_effect = "SoundEffect" in args.model
     is_voice_gen = "VoiceGenerator" in args.model
     is_realtime = "Realtime" in args.model
+    # Local-v1.5 ships its own AutoProcessor (processor_config.json +
+    # processing_moss_tts.py), so it reuses the same _build_unified_codes
+    # path as the Delay variants below — just a different n_vq/codec.
+    is_local = "Local" in args.model
+    output_sr = 48000 if is_local else 24000
 
     ref_audio = None
     builder_kwargs: dict = {}
@@ -314,27 +321,33 @@ def run_tts(args: argparse.Namespace) -> None:
             wav = mm.get("model_outputs")
         if wav is None:
             continue
+        # Local-v1.5's codec emits stereo chunks shaped (C, T); keep the
+        # channel axis intact (don't flatten) so concatenation below stays
+        # along the time axis. Mono variants stay (T,) as before.
         if isinstance(wav, list):
-            non_empty = [w.reshape(-1) for w in wav if isinstance(w, torch.Tensor) and w.numel() > 0]
+            non_empty = [w for w in wav if isinstance(w, torch.Tensor) and w.numel() > 0]
             if non_empty and t_first_audio is None:
                 t_first_audio = _time.monotonic()
             chunks.extend(non_empty)
         elif isinstance(wav, torch.Tensor) and wav.numel() > 0:
             if t_first_audio is None:
                 t_first_audio = _time.monotonic()
-            chunks.append(wav.reshape(-1))
+            chunks.append(wav)
 
     t_gen_end = _time.monotonic()
     if not chunks:
         print("No audio in output — check model loading logs.")
         return
 
-    wav_t = torch.cat(chunks).float().cpu()
-    sf.write(args.output, wav_t.numpy(), 24000)
-    print(f"Saved {len(wav_t) / 24000:.2f}s of audio to {args.output}")
+    cat_dim = -1 if chunks[0].dim() == 2 else 0
+    wav_t = torch.cat(chunks, dim=cat_dim).float().cpu()
+    n_samples = wav_t.shape[-1]
+    # soundfile wants (T,) mono or (T, C) multi-channel.
+    sf.write(args.output, (wav_t.transpose(0, 1) if wav_t.dim() == 2 else wav_t).numpy(), output_sr)
+    print(f"Saved {n_samples / output_sr:.2f}s of audio to {args.output}")
     ttfa_ms = (t_first_audio - t_gen_start) * 1000.0 if t_first_audio is not None else float("nan")
     gen_total_ms = (t_gen_end - t_gen_start) * 1000.0
-    print(f"TTFA_MS={ttfa_ms:.1f} GEN_TOTAL_MS={gen_total_ms:.1f} AUDIO_S={len(wav_t) / 24000:.2f}")
+    print(f"TTFA_MS={ttfa_ms:.1f} GEN_TOTAL_MS={gen_total_ms:.1f} AUDIO_S={n_samples / output_sr:.2f}")
 
 
 def main() -> None:

vllm_omni/benchmarks/patch/patch.py

@@ -52,6 +52,8 @@ from vllm_omni.benchmarks.data_modules.ttsd_dataset import TTSDDataset
 from vllm_omni.metrics import definitions as defs
 
 _AUDIO_CONTINUITY_THRESHOLD_ENV = "VLLM_OMNI_BENCH_AUDIO_CONTINUITY_THRESHOLD_S"
+_AUDIO_SAMPLE_RATE_ENV = "VLLM_OMNI_BENCH_AUDIO_SAMPLE_RATE"
+_AUDIO_CHANNELS_ENV = "VLLM_OMNI_BENCH_AUDIO_CHANNELS"
 RETURN_STAGE_METRICS_FIELD = "return_stage_metrics"
 _IMAGE_STAGE_METRICS_BACKENDS = frozenset({"openai-image-edits-omni"})
 _PRINT_STAGE = False
@@ -109,6 +111,32 @@ def _audio_continuity_threshold_s() -> float:
     return max(value, 0.0)
 
 
+def _audio_pcm_format() -> tuple[int, int]:
+    """Return ``(sample_rate, channels)`` of the streamed PCM response.
+
+    Defaults to 24 kHz mono (Qwen3-TTS / VoxCPM2 / most MOSS-TTS variants).
+    Override for models with a different codec output format (e.g.
+    MOSS-TTS-Local-Transformer-v1.5, which streams 48 kHz stereo PCM) via
+    ``VLLM_OMNI_BENCH_AUDIO_SAMPLE_RATE`` / ``VLLM_OMNI_BENCH_AUDIO_CHANNELS`` -
+    getting this wrong silently skews audio_duration/audio_rtf/underrun stats.
+    """
+    sample_rate = 24000
+    channels = 1
+    raw_sr = os.environ.get(_AUDIO_SAMPLE_RATE_ENV)
+    if raw_sr:
+        try:
+            sample_rate = int(raw_sr)
+        except ValueError:
+            logger.warning("Invalid %s=%r; using default %d", _AUDIO_SAMPLE_RATE_ENV, raw_sr, sample_rate)
+    raw_ch = os.environ.get(_AUDIO_CHANNELS_ENV)
+    if raw_ch:
+        try:
+            channels = int(raw_ch)
+        except ValueError:
+            logger.warning("Invalid %s=%r; using default %d", _AUDIO_CHANNELS_ENV, raw_ch, channels)
+    return sample_rate, channels
+
+
 get_samples_old = datasets.get_samples
 
 _DEFAULT_DAILY_OMNI_REPO = "liarliar/Daily-Omni"
@@ -1133,10 +1161,10 @@ async def async_request_openai_audio_speech(
     output = MixRequestFuncOutput()
     output.prompt_len = request_func_input.prompt_len
 
-    # PCM format: 16-bit signed, 24 kHz, mono
-    sample_rate = 24000
+    # PCM format: 16-bit signed; sample_rate/channels are model-dependent
+    # (see _audio_pcm_format docstring).
+    sample_rate, channels = _audio_pcm_format()
     sample_width = 2  # 16-bit = 2 bytes
-    channels = 1
 
     st = time.perf_counter()
     output.start_time = st

vllm_omni/config/pipeline_registry.py

@@ -65,7 +65,11 @@ from vllm_omni.model_executor.models.ming_flash_omni.pipeline import (
 )
 from vllm_omni.model_executor.models.ming_tts.pipeline import MING_TTS_PIPELINE
 from vllm_omni.model_executor.models.minicpmo_4_5.pipeline import MINICPMO_4_5_PIPELINE
-from vllm_omni.model_executor.models.moss_tts.pipeline import MOSS_TTS_PIPELINE, MOSS_TTS_REALTIME_PIPELINE
+from vllm_omni.model_executor.models.moss_tts.pipeline import (
+    MOSS_TTS_LOCAL_PIPELINE,
+    MOSS_TTS_PIPELINE,
+    MOSS_TTS_REALTIME_PIPELINE,
+)
 from vllm_omni.model_executor.models.moss_tts_nano.pipeline import MOSS_TTS_NANO_PIPELINE
 from vllm_omni.model_executor.models.qwen2_5_omni.pipeline import (
     QWEN2_5_OMNI_PIPELINE,
@@ -112,6 +116,7 @@ OMNI_PIPELINES: dict[str, PipelineConfig | PipelineResolverFunc] = {
     "moss_tts_nano": MOSS_TTS_NANO_PIPELINE,
     "moss_tts_delay": MOSS_TTS_PIPELINE,
     "moss_tts_realtime": MOSS_TTS_REALTIME_PIPELINE,
+    "moss_tts_local": MOSS_TTS_LOCAL_PIPELINE,
     "minicpmo_4_5": MINICPMO_4_5_PIPELINE,
     "higgs_audio_v2": HIGGS_AUDIO_V2_PIPELINE,
     "higgs_multimodal_qwen3": HIGGS_AUDIO_V3_PIPELINE,

vllm_omni/deploy/moss_tts_local.yaml

@@ -0,0 +1,56 @@
+# MOSS-TTS-Local-Transformer-v1.5 (MossTTSLocalModel, n_vq=12, 48 kHz stereo)
+# HF: OpenMOSS-Team/MOSS-TTS-Local-Transformer-v1.5
+# Codec: OpenMOSS-Team/MOSS-Audio-Tokenizer-v2 (48 kHz, stereo, channel-interleaved RVQ)
+# Uses a 1-layer GPT2-style local depth transformer (no delay warm-up), reset every
+# frame — same per-frame timing (80 ms/frame, 12.5 fps) as MOSS-TTS-Realtime.
+async_chunk: true
+
+connectors:
+  shm:
+    name: SharedMemoryConnector
+    extra:
+      shm_threshold_bytes: 65536
+      codec_streaming: true
+      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
+      codec_left_context_frames: 0
+
+stages:
+  - stage_id: 0
+    max_num_seqs: 8
+    gpu_memory_utilization: 0.60
+    enable_prefix_caching: false
+    async_scheduling: true
+    max_num_batched_tokens: 4096
+    max_model_len: 4096
+    devices: "0"
+    output_connectors:
+      to_stage_1: shm
+    default_sampling_params:
+      temperature: 1.0
+      top_p: 0.95
+      top_k: 50
+      max_tokens: 4096
+      seed: 42
+      repetition_penalty: 1.1
+
+  - stage_id: 1
+    max_num_seqs: 1
+    gpu_memory_utilization: 0.15
+    enforce_eager: true
+    enable_prefix_caching: false
+    async_scheduling: true
+    max_num_batched_tokens: 65536
+    max_model_len: 65536
+    devices: "0"
+    input_connectors:
+      from_stage_0: shm
+    default_sampling_params:
+      temperature: 0.0
+      top_p: 1.0
+      top_k: -1
+      max_tokens: 65536
+      seed: 42
+      repetition_penalty: 1.0

vllm_omni/engine/arg_utils.py

@@ -41,6 +41,10 @@ def _register_omni_hf_configs() -> None:
             IndexTTS2Config,
         )
         from vllm_omni.model_executor.models.ming_tts.config_ming_tts import MingDenseConfig
+        from vllm_omni.model_executor.models.moss_tts.configuration_moss_tts import (
+            MossTTSLocalConfig,
+            MossTTSRealtimeConfig,
+        )
         from vllm_omni.model_executor.models.qwen3_tts.configuration_qwen3_tts import (
             Qwen3TTSConfig,
         )
@@ -63,6 +67,8 @@ def _register_omni_hf_configs() -> None:
     for model_type, config_cls in [
         ("dense", MingDenseConfig),
         ("indextts2", IndexTTS2Config),
+        ("moss_tts_local", MossTTSLocalConfig),
+        ("moss_tts_realtime", MossTTSRealtimeConfig),
         ("qwen3_tts", Qwen3TTSConfig),
         ("cosyvoice3", CosyVoice3Config),
         ("glm_tts", GLMTTSConfig),

vllm_omni/entrypoints/openai/serving_speech.py

@@ -81,6 +81,7 @@ _VOXCPM2_TTS_MODEL_STAGES = {"latent_generator"}
 _MING_TTS_MODEL_STAGES = {"ming_tts"}
 _MOSS_TTS_MODEL_STAGES = {"moss_tts_nano"}
 _MOSS_TTS_FULL_MODEL_STAGES = {"moss_tts", "moss_tts_codec"}
+_MOSS_TTS_LOCAL_MODEL_STAGES = {"moss_tts_local", "moss_tts_local_codec"}
 _HIGGS_AUDIO_V2_TTS_MODEL_STAGES = {"higgs_audio_v2"}
 _HIGGS_V3_TTS_MODEL_STAGES = {"higgs_audio_v3"}
 _GLM_TTS_MODEL_STAGES = {"glm_tts"}
@@ -99,6 +100,7 @@ _TTS_MODEL_STAGES: set[str] = (
     | _MING_TTS_MODEL_STAGES
     | _MOSS_TTS_MODEL_STAGES
     | _MOSS_TTS_FULL_MODEL_STAGES
+    | _MOSS_TTS_LOCAL_MODEL_STAGES
     | _GLM_TTS_MODEL_STAGES
     | _STEP_AUDIO2_TTS_MODEL_STAGES
     | _INDEXTTS2_TTS_MODEL_STAGES
@@ -182,6 +184,18 @@ def _create_wav_header(sample_rate: int, num_channels: int = 1, bits_per_sample:
     return header
 
 
+def _infer_audio_num_channels(audio: np.ndarray) -> int:
+    """Infer channel count before streaming PCM bytes are wrapped as WAV."""
+    if audio.ndim == 3 and audio.shape[0] == 1:
+        audio = audio[0]
+    if audio.ndim == 2:
+        if audio.shape[0] in (1, 2):
+            return int(audio.shape[0])
+        if audio.shape[1] in (1, 2):
+            return int(audio.shape[1])
+    return 1
+
+
 def _sanitize_filename(filename: str) -> str:
     """Sanitize filename to prevent path traversal attacks.
 
@@ -688,6 +702,8 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
             return "moss_tts_nano"
         if model_stage in _MOSS_TTS_FULL_MODEL_STAGES:
             return "moss_tts"
+        if model_stage in _MOSS_TTS_LOCAL_MODEL_STAGES:
+            return "moss_tts"
         if model_stage in _HIGGS_AUDIO_V2_TTS_MODEL_STAGES:
             return "higgs_audio_v2"
         if model_stage in _HIGGS_V3_TTS_MODEL_STAGES:
@@ -1675,6 +1691,8 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
             name = ""
         if "realtime" in name:
             return "realtime"
+        if "local" in name:
+            return "local"
         if "ttsd" in name:
             return "ttsd"
         if "soundeffect" in name:
@@ -1701,9 +1719,17 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
                 return "Input text cannot be empty"
 
         v = self._moss_variant
-        if v in (None, "tts", "realtime"):
+        if v in (None, "tts", "realtime", "local"):
             if request.ref_audio is None:
-                label = "MOSS-TTS-Nano" if v is None else ("MOSS-TTS-Realtime" if v == "realtime" else "MOSS-TTS")
+                label = (
+                    "MOSS-TTS-Nano"
+                    if v is None
+                    else (
+                        "MOSS-TTS-Realtime"
+                        if v == "realtime"
+                        else ("MOSS-TTS-Local-Transformer" if v == "local" else "MOSS-TTS")
+                    )
+                )
                 return f"{label} requires 'ref_audio' (reference audio for voice cloning)."
             return self._validate_ref_audio_format(request.ref_audio)
 
@@ -1806,11 +1832,26 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
             params["prompt_audio_array"] = [[wav_list, sr]]
             return params
 
-        # ---- MossTTSDelay family (tts/ttsd/sound_effect/voice_generator):
-        # call the upstream processor server-side to produce unified codes ----
+        # ---- MossTTSDelay family (tts/ttsd/sound_effect/voice_generator)
+        # and MOSS-TTS-Local-Transformer-v1.5: call the upstream processor
+        # server-side to produce unified codes. Local-v1.5 ships its own
+        # AutoProcessor (processor_config.json + processing_moss_tts.py) and
+        # reuses this exact build_user_message/encode_audios_from_wav path in
+        # the offline example (examples/.../moss_tts/end2end.py:
+        # _build_unified_codes) -- it is NOT in the same boat as Realtime
+        # (no processor_config.json there), so it must not fall back to the
+        # prompt_audio_array path above (which the talker's preprocess()
+        # never reads -- info_dict["codes"]["ref"] is the only thing it
+        # consumes, so skipping this path silently drops all voice-clone
+        # conditioning and produces unconditioned/garbage audio online). ----
         proc = self._get_moss_processor()
         n_vq = int(getattr(proc.model_config, "n_vq", 32))
-        sr_target = int(getattr(proc.model_config, "sampling_rate", 24000))
+        # Local-v1.5 encodes reference audio at a fixed 24 kHz working rate
+        # regardless of its 48 kHz stereo *output* codec -- mirrors the
+        # offline example's hardcoded encode_audios_from_wav(sampling_rate=24000)
+        # for this variant; proc.model_config.sampling_rate there is the
+        # output rate (48000), the wrong value to resample the reference into.
+        sr_target = 24000 if v == "local" else int(getattr(proc.model_config, "sampling_rate", 24000))
 
         # Reference-audio encoding + speaker caching lives in the model package
         # (moss_tts.reference_encoder), mirroring Fish Speech / CosyVoice3 /
@@ -1837,7 +1878,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
             )
 
         user_kwargs: dict[str, Any] = {"text": request.input or ""}
-        if v in ("tts", "realtime"):
+        if v in ("tts", "local"):
             user_kwargs["reference"] = [await _encode_ref(request.ref_audio)]
         elif v == "ttsd":
             refs = [await _encode_ref(request.ref_audio)]
@@ -2654,7 +2695,12 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
                         assert sr_raw is not None, (
                             "First audio chunk must include sample rate metadata for WAV streaming"
                         )
-                        wav_header = _create_wav_header(sample_rate=sample_rate_val, num_channels=1, bits_per_sample=16)
+                        num_channels = _infer_audio_num_channels(np.asarray(chunk_np))
+                        wav_header = _create_wav_header(
+                            sample_rate=sample_rate_val,
+                            num_channels=num_channels,
+                            bits_per_sample=16,
+                        )
                         yield wav_header
                         first_chunk = False
 

vllm_omni/entrypoints/openai/tts_adapters/moss_tts.py

@@ -54,5 +54,5 @@ class MossTTSNanoAdapter(_MossTTSAdapterBase):
 
 @register_tts_adapter
 class MossTTSAdapter(_MossTTSAdapterBase):
-    stage_keys = frozenset({"moss_tts", "moss_tts_codec"})
+    stage_keys = frozenset({"moss_tts", "moss_tts_codec", "moss_tts_local", "moss_tts_local_codec"})
     name = "moss_tts"

vllm_omni/model_executor/models/moss_tts/audio_tokenizer.py

@@ -45,6 +45,8 @@ class MossAudioTokenizerConfig(PretrainedConfig):
         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 = 1,
+        enable_channel_interleave: bool = True,
         quantizer_type: str = "rlfq",
         quantizer_kwargs: dict[str, Any] | None = None,
         **kwargs: Any,
@@ -54,6 +56,8 @@ class MossAudioTokenizerConfig(PretrainedConfig):
         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.encoder_kwargs = encoder_kwargs or _default_encoder_kwargs()
         self.decoder_kwargs = decoder_kwargs or _default_decoder_kwargs()
         if quantizer_kwargs is None:
@@ -201,13 +205,21 @@ class _Attention(nn.Module):
     """Causal multi-head self-attention with RoPE, no streaming KV cache."""
 
     def __init__(
-        self, embed_dim: int, num_heads: int, causal: bool, max_period: float, device=None, dtype=None
+        self,
+        embed_dim: int,
+        num_heads: int,
+        causal: bool,
+        max_period: float,
+        context: int | None = None,
+        device=None,
+        dtype=None,
     ) -> None:
         super().__init__()
         self.embed_dim = embed_dim
         self.num_heads = num_heads
         self.causal = causal
         self.max_period = max_period
+        self.context = context
         kw = {"device": device, "dtype": dtype}
         self.in_proj = nn.Linear(embed_dim, 3 * embed_dim, bias=False, **kw)
         self.out_proj = nn.Linear(embed_dim, embed_dim, bias=False, **kw)
@@ -234,7 +246,38 @@ class _Attention(nn.Module):
         qkv = self.in_proj(x).reshape(B, T, 3, H, D).permute(2, 0, 3, 1, 4)
         q, k, v = qkv[0], qkv[1], qkv[2]  # each (B, H, T, D)
         q, k = _apply_rope(q, k, self.max_period)
-        out = F.scaled_dot_product_attention(q, k, v, is_causal=self.causal)
+        if self.context is not None and self.context < T:
+            # Local-windowed causal attention: query i may only see keys in
+            # [i - context + 1, i]. Matches upstream's per-stage receptive
+            # field — wider context windows are not numerically equivalent.
+            #
+            # A dense (T, T) mask is O(T^2) memory (e.g. 4 GiB at T=65536) —
+            # for long utterances decoded in a single non-streaming call this
+            # OOMs. Process queries in blocks against only their reachable
+            # key range so peak mask memory is O(block * (block + context))
+            # regardless of T.
+            block = 4096
+            if T <= block:
+                positions = torch.arange(T, device=x.device)
+                delta = positions.view(-1, 1) - positions.view(1, -1)
+                mask = (delta >= 0) & (delta < self.context)
+                out = F.scaled_dot_product_attention(q, k, v, attn_mask=mask)
+            else:
+                outs = []
+                for start in range(0, T, block):
+                    end = min(start + block, T)
+                    k_start = max(0, start - self.context + 1)
+                    q_pos = torch.arange(start, end, device=x.device)
+                    k_pos = torch.arange(k_start, end, device=x.device)
+                    delta = q_pos.view(-1, 1) - k_pos.view(1, -1)
+                    blk_mask = (delta >= 0) & (delta < self.context)
+                    out_blk = F.scaled_dot_product_attention(
+                        q[:, :, start:end], k[:, :, k_start:end], v[:, :, k_start:end], attn_mask=blk_mask
+                    )
+                    outs.append(out_blk)
+                out = torch.cat(outs, dim=2)
+        else:
+            out = F.scaled_dot_product_attention(q, k, v, is_causal=self.causal)
         out = out.transpose(1, 2).reshape(B, T, self.embed_dim)
         return self.out_proj(out)
 
@@ -248,6 +291,7 @@ class _TransformerLayer(nn.Module):
         causal: bool,
         max_period: float,
         layer_scale: float | None,
+        context: int | None = None,
         device=None,
         dtype=None,
     ) -> None:
@@ -255,7 +299,7 @@ class _TransformerLayer(nn.Module):
         kw = {"device": device, "dtype": dtype}
         self.norm1 = nn.LayerNorm(d_model, eps=1e-5, **kw)
         self.norm2 = nn.LayerNorm(d_model, eps=1e-5, **kw)
-        self.attn = _Attention(d_model, num_heads, causal, max_period, device, dtype)
+        self.attn = _Attention(d_model, num_heads, causal, max_period, context, device, dtype)
         self.ff1 = nn.Linear(d_model, dim_feedforward, bias=False, **kw)
         self.ff2 = nn.Linear(dim_feedforward, d_model, bias=False, **kw)
         ls_kw = cast(dict[str, object], kw)
@@ -282,13 +326,16 @@ class _Transformer(nn.Module):
         causal: bool,
         max_period: float,
         layer_scale: float | None,
+        context: int | None = None,
         device=None,
         dtype=None,
     ) -> None:
         super().__init__()
         self.layers = nn.ModuleList(
             [
-                _TransformerLayer(d_model, num_heads, dim_feedforward, causal, max_period, layer_scale, device, dtype)
+                _TransformerLayer(
+                    d_model, num_heads, dim_feedforward, causal, max_period, layer_scale, context, device, dtype
+                )
                 for _ in range(num_layers)
             ]
         )
@@ -307,11 +354,12 @@ class _ProjectedTransformer(nn.Module):
     ) -> None:
         super().__init__()
         self.downsample_ratio: int = 1
-        self.in_proj = nn.Linear(input_dimension, d_model, bias=False) if d_model != input_dimension else nn.Identity()
+        # Upstream always materializes a learned Linear here regardless of
+        # whether input_dimension == d_model — substituting Identity when
+        # the dims happen to match silently drops a trained projection.
+        self.in_proj = nn.Linear(input_dimension, d_model, bias=False)
         self.transformer = _Transformer(d_model=d_model, **kwargs)
-        self.out_proj = (
-            nn.Linear(d_model, output_dimension, bias=False) if d_model != output_dimension else nn.Identity()
-        )
+        self.out_proj = nn.Linear(d_model, output_dimension, bias=False)
 
     def forward(self, x: torch.Tensor, lengths: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
         x = self.in_proj(x.transpose(1, 2))  # (B, D, T) → (B, T, d_model)
@@ -456,10 +504,21 @@ class _ResidualQ(nn.Module):
 
 
 def _build_modules(
-    specs: list[dict[str, Any]], is_downsample: bool, context_duration: float, sampling_rate: float
-) -> nn.ModuleList:
+    specs: list[dict[str, Any]], is_downsample: bool, default_context_duration: float, frame_rate: float
+) -> tuple[nn.ModuleList, float]:
+    """Build encoder/decoder stage modules.
+
+    ``frame_rate`` tracks the running frames-per-second at the current stage
+    (updated after each module by its downsample/upsample ratio, mirroring
+    upstream's ``current_frame_rate``). Each ``Transformer`` module's local
+    attention window (``context``, in frames) is derived from its own
+    ``context_duration`` (seconds) at the rate *at that point in the stack* —
+    not a global constant — so it must be computed here, not by the caller.
+    Returns the final rate too, since the decoder's starting rate is the
+    encoder's bottleneck rate, not the raw sampling rate.
+    """
     modules: list[nn.Module] = []
-    rate = float(sampling_rate)
+    rate = float(frame_rate)
     for spec in specs:
         spec = dict(spec)
         if spec["module_type"] == "PatchedPretransform":
@@ -471,16 +530,18 @@ def _build_modules(
             spec.pop("positional_embedding", None)
             spec.pop("norm", None)
             spec.pop("causal", None)
+            context_duration = float(spec.pop("context_duration", default_context_duration))
             m = _ProjectedTransformer(
                 module_type="Transformer",
                 causal=True,
                 max_period=spec.pop("max_period", 10000),
                 layer_scale=spec.pop("layer_scale", None),
+                context=int(round(rate * context_duration)),
                 **spec,
             )
         modules.append(m)
         rate = rate / m.downsample_ratio if is_downsample else rate * m.downsample_ratio
-    return nn.ModuleList(modules)
+    return nn.ModuleList(modules), rate
 
 
 class MossAudioTokenizerModel(PreTrainedModel):
@@ -494,12 +555,27 @@ class MossAudioTokenizerModel(PreTrainedModel):
         super().__init__(config)
         self.sampling_rate = config.sampling_rate
         self.downsample_rate = config.downsample_rate
+        # Real v1 checkpoints store these as an explicit JSON `null` rather
+        # than omitting the key, so `getattr(..., default)` doesn't apply —
+        # normalize None to the same defaults here.
+        self.number_channels = getattr(config, "number_channels", 1) or 1
+        self.enable_channel_interleave = getattr(config, "enable_channel_interleave", True)
+        if self.enable_channel_interleave is None:
+            self.enable_channel_interleave = True
 
         ctx = config.causal_transformer_context_duration
-        sr = float(config.sampling_rate)
-        self.encoder = _build_modules(config.encoder_kwargs, is_downsample=True, context_duration=ctx, sampling_rate=sr)
-        self.decoder = _build_modules(
-            copy.deepcopy(config.decoder_kwargs), is_downsample=False, context_duration=ctx, sampling_rate=sr
+        channel_interleave_factor = (
+            self.number_channels if (self.enable_channel_interleave and self.number_channels > 1) else 1
+        )
+        frame_rate = float(config.sampling_rate) * channel_interleave_factor
+        self.encoder, bottleneck_rate = _build_modules(
+            config.encoder_kwargs, is_downsample=True, default_context_duration=ctx, frame_rate=frame_rate
+        )
+        self.decoder, _ = _build_modules(
+            copy.deepcopy(config.decoder_kwargs),
+            is_downsample=False,
+            default_context_duration=ctx,
+            frame_rate=bottleneck_rate,
         )
 
         kw = dict(config.quantizer_kwargs)
@@ -512,15 +588,29 @@ class MossAudioTokenizerModel(PreTrainedModel):
         wav_list: list[torch.Tensor],
         num_quantizers: int | None = None,
     ) -> MossAudioTokenizerEncoderOutput:
-        """Encode a list of 1-D waveform tensors → RVQ codes (NQ, B, T)."""
+        """Encode a list of waveform tensors → RVQ codes (NQ, B, T).
+
+        Each element of ``wav_list`` is ``(T,)``/``(1, T)`` for a mono codec
+        (``number_channels == 1``) or ``(number_channels, T)`` otherwise.
+        """
         device = wav_list[0].device
         B = len(wav_list)
-        max_len = max(w.shape[-1] for w in wav_list)
-        x = torch.zeros(B, 1, max_len, device=device)
+        C = self.number_channels
+        normalized: list[torch.Tensor] = []
         lengths = torch.zeros(B, device=device, dtype=torch.long)
         for i, w in enumerate(wav_list):
-            x[i, 0, : w.shape[-1]] = w
-            lengths[i] = w.shape[-1]
+            if C == 1:
+                w_i = w.unsqueeze(0) if w.dim() == 1 else w
+            else:
+                if w.dim() != 2 or w.shape[0] != C:
+                    raise ValueError(f"Expected wav_list[{i}] to have shape ({C}, T), got {tuple(w.shape)}.")
+                w_i = w
+            normalized.append(w_i)
+            lengths[i] = w_i.shape[-1]
+        max_len = int(lengths.max().item())
+        x = torch.zeros(B, C, max_len, device=device, dtype=wav_list[0].dtype)
+        for i, w_i in enumerate(normalized):
+            x[i, :, : w_i.shape[-1]] = w_i
         return self._encode(x, lengths, num_quantizers)
 
     @torch.no_grad()
@@ -541,6 +631,45 @@ class MossAudioTokenizerModel(PreTrainedModel):
             lengths[i] = c.shape[-1]
         return self._decode(codes, lengths)
 
+    def _flatten_channels_for_codec(
+        self,
+        input_values: torch.Tensor,
+        input_lengths: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """Pad to a ``downsample_rate`` multiple, then interleave channels into one stream.
+
+        Mirrors upstream's ``_flatten_channels_for_codec``: for a stereo
+        codec, ``(B, C, T)`` becomes ``(B, 1, T*C)`` with samples interleaved
+        frame-by-frame across channels, so the mono encoder/decoder stack can
+        process a multi-channel waveform unchanged.
+        """
+        if input_values.shape[-1] % self.downsample_rate != 0:
+            pad_length = self.downsample_rate - (input_values.shape[-1] % self.downsample_rate)
+            input_values = F.pad(input_values, (0, pad_length))
+        if self.number_channels > 1 and self.enable_channel_interleave:
+            input_values = input_values.transpose(1, 2).contiguous().view(input_values.shape[0], 1, -1)
+            input_lengths = input_lengths * self.number_channels
+        return input_values, input_lengths
+
+    def _restore_channels_from_codec(
+        self,
+        output_values: torch.Tensor,
+        output_lengths: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """Inverse of ``_flatten_channels_for_codec``: de-interleave back to ``(B, C, T)``."""
+        if self.number_channels == 1 or not self.enable_channel_interleave:
+            return output_values.float(), output_lengths
+        output_values = (
+            output_values.squeeze(1)
+            .contiguous()
+            .view(output_values.shape[0], -1, self.number_channels)
+            .transpose(1, 2)
+            .contiguous()
+            .float()
+        )
+        output_lengths = torch.div(output_lengths, self.number_channels, rounding_mode="floor")
+        return output_values, output_lengths
+
     def _encode(
         self,
         x: torch.Tensor,
@@ -549,9 +678,7 @@ class MossAudioTokenizerModel(PreTrainedModel):
     ) -> MossAudioTokenizerEncoderOutput:
         if x.dim() == 2:
             x = x.unsqueeze(1)
-        T = x.shape[-1]
-        if T % self.downsample_rate:
-            x = F.pad(x, (0, self.downsample_rate - T % self.downsample_rate))
+        x, lengths = self._flatten_channels_for_codec(x, lengths)
         e, e_len = x, lengths
         for m in self.encoder:
             e, e_len = m(e, e_len)
@@ -567,6 +694,7 @@ class MossAudioTokenizerModel(PreTrainedModel):
         d, d_len = z, lengths
         for m in self.decoder:
             d, d_len = m(d, d_len)
+        d, d_len = self._restore_channels_from_codec(d, d_len)
         return MossAudioTokenizerDecoderOutput(audio=d, audio_lengths=d_len)
 
 

vllm_omni/model_executor/models/moss_tts/configuration_moss_tts.py

@@ -17,6 +17,7 @@ from __future__ import annotations
 
 from transformers import AutoConfig
 from transformers.configuration_utils import PretrainedConfig
+from transformers.models.gpt2.configuration_gpt2 import GPT2Config
 from transformers.models.qwen3 import Qwen3Config
 
 
@@ -235,7 +236,80 @@ class MossTTSRealtimeConfig(PretrainedConfig):
         return self.language_config
 
 
+class MossTTSLocalConfig(PretrainedConfig):
+    """Config for MossTTSLocalModel (MOSS-TTS-Local-Transformer-v1.5).
+
+    Like ``MossTTSDelayConfig``, the HF checkpoint stores a nested Qwen3 config
+    (``qwen3_config``) which we unwrap via ``get_text_config()``. It also stores
+    a nested GPT2-style config (``gpt2_config``) describing the single-layer
+    local depth transformer that decodes the ``n_vq`` audio codebooks per frame
+    (see ``modeling_moss_tts_local_depth.py``).
+    """
+
+    model_type = "moss_tts_local"
+
+    def __init__(
+        self,
+        qwen3_config: dict | None = None,
+        gpt2_config: dict | None = None,
+        n_vq: int = 12,
+        audio_vocab_size: int = 1024,
+        audio_pad_token_id: int = 1024,
+        pad_token_id: int = 151643,
+        im_start_token_id: int = 151644,
+        im_end_token_id: int = 151645,
+        audio_start_token_id: int = 151669,
+        audio_end_token_id: int = 151670,
+        audio_user_slot_token_id: int = 151654,
+        audio_assistant_slot_token_id: int = 151656,
+        sampling_rate: int = 48000,
+        audio_tokenizer_name_or_path: str = "OpenMOSS-Team/MOSS-Audio-Tokenizer-v2",
+        local_text_head_mode: str = "binary",
+        **kwargs: object,
+    ) -> None:
+        if qwen3_config is None:
+            qwen3_config = {}
+        if isinstance(qwen3_config, dict):
+            qwen3_config = dict(qwen3_config)
+            qwen3_config.pop("model_type", None)
+            self.qwen3_config = Qwen3Config(**qwen3_config)
+        else:
+            self.qwen3_config = qwen3_config
+
+        if isinstance(gpt2_config, dict):
+            gpt2_config = dict(gpt2_config)
+            gpt2_config.pop("model_type", None)
+            self.gpt2_config = GPT2Config(**gpt2_config)
+        else:
+            self.gpt2_config = gpt2_config
+
+        super().__init__(pad_token_id=pad_token_id, **kwargs)
+
+        self.n_vq = n_vq
+        self.audio_vocab_size = audio_vocab_size
+        self.audio_pad_token_id = audio_pad_token_id
+        self.audio_pad_code = audio_pad_token_id
+        self.im_start_token_id = im_start_token_id
+        self.im_end_token_id = im_end_token_id
+        self.audio_start_token_id = audio_start_token_id
+        self.audio_end_token_id = audio_end_token_id
+        self.audio_user_slot_token_id = audio_user_slot_token_id
+        self.audio_assistant_slot_token_id = audio_assistant_slot_token_id
+        self.sampling_rate = sampling_rate
+        self.audio_tokenizer_name_or_path = audio_tokenizer_name_or_path
+        self.local_text_head_mode = local_text_head_mode
+
+    def get_text_config(self, **_: object) -> Qwen3Config:
+        return self.qwen3_config
+
+
 AutoConfig.register("moss_tts_delay", MossTTSDelayConfig)
 AutoConfig.register("moss_tts_realtime", MossTTSRealtimeConfig)
+AutoConfig.register("moss_tts_local", MossTTSLocalConfig)
 
-__all__ = ["MossTTSDelayConfig", "MossTTSRealtimeConfig", "MossTTSLocalTransformerConfig"]
+__all__ = [
+    "MossTTSDelayConfig",
+    "MossTTSRealtimeConfig",
+    "MossTTSLocalTransformerConfig",
+    "MossTTSLocalConfig",
+]

vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_codec.py

@@ -5,7 +5,6 @@
 
 from __future__ import annotations
 
-import os
 from collections.abc import Iterable
 from typing import Any
 
@@ -42,10 +41,8 @@ class MossTTSCodecDecoder(nn.Module):
     ``n_vq=32`` (MOSS-TTS) or ``n_vq=16`` (all other variants) without
     swapping weights.
 
-    The codec checkpoint path defaults to the value stored in
-    ``vllm_config.model_config.hf_config.codec_model_name_or_path`` but can
-    be overridden by setting the environment variable
-    ``MOSS_TTS_CODEC_PATH``.
+    The codec checkpoint path comes from
+    ``vllm_config.model_config.hf_config.codec_model_name_or_path``.
     """
 
     input_modalities = "audio"
@@ -64,11 +61,19 @@ class MossTTSCodecDecoder(nn.Module):
 
         cfg = vllm_config.model_config.hf_config
         self._n_vq: int = int(getattr(cfg, "n_vq", getattr(cfg, "rvq", 16)))
-        self._codec_path: str = str(getattr(cfg, "codec_model_name_or_path", "OpenMOSS-Team/MOSS-Audio-Tokenizer"))
+        self._codec_path: str = str(
+            getattr(
+                cfg,
+                "codec_model_name_or_path",
+                getattr(cfg, "audio_tokenizer_name_or_path", "OpenMOSS-Team/MOSS-Audio-Tokenizer"),
+            )
+        )
 
-        # Resolved at load_weights() time
+        # Resolved at load_weights() time, once the codec checkpoint's own
+        # config (sampling rate, channel count) is known.
         self._codec: MossAudioTokenizerModel | None = None
         self._cuda_graph_wrapper: MossTTSCUDAGraphCodecWrapper | None = None
+        self._n_channels: int = 1
         self._sr_tensor = torch.tensor(self._OUTPUT_SAMPLE_RATE, dtype=torch.int32)
 
     # ------------------------------------------------------------------
@@ -119,7 +124,7 @@ class MossTTSCodecDecoder(nn.Module):
         """
         sr_tensor = self._sr_tensor
         empty = torch.zeros((0,), dtype=torch.float32)
-        info_list: list[dict[str, Any]] = runtime_additional_information or []
+        info_list: list[dict[str, Any]] = runtime_additional_information or [{}]
         num_req = max(len(info_list), 1)
 
         if self._codec is None:
@@ -147,7 +152,7 @@ class MossTTSCodecDecoder(nn.Module):
         # kwargs if available, else assume one request.
         ids_flat = input_ids.reshape(-1).to(dtype=torch.long)
         num_scheduled_tokens = kwargs.get("num_scheduled_tokens")
-        if isinstance(num_scheduled_tokens, list) and len(num_scheduled_tokens) == num_req:
+        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))
@@ -176,12 +181,8 @@ class MossTTSCodecDecoder(nn.Module):
             codebook_size = self._codec.config.codebook_size
             codes_nq_t = codes_nq_t.clamp_(0, int(codebook_size) - 1)
 
-            # left_context_size: number of left-context frames prepended to
-            # this chunk by the stage input processor (overlap-add pattern,
-            # same as Qwen3-TTS code2wav).  We decode them together with the
-            # real chunk for smooth boundaries, then trim them from the output.
             meta = (info.get("meta", {}) if isinstance(info, dict) else {}) or {}
-            left_ctx = meta.get("left_context_size", 0)
+            left_ctx = self._runtime_value(info, meta, "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):
@@ -196,28 +197,40 @@ class MossTTSCodecDecoder(nn.Module):
             if out.audio is None:
                 continue
 
-            wav = out.audio.reshape(-1).to(dtype=torch.float32).cpu()
+            # ``out.audio`` is ``(1, C, T)``; keep the channel axis for
+            # stereo codecs (Local-v1.5) and flatten to ``(T,)`` for mono
+            # ones (Delay/Realtime) to preserve their existing output shape.
+            wav = out.audio[0].to(dtype=torch.float32).cpu()
             if out.audio_lengths is not None:
-                wav = wav[: int(out.audio_lengths[0].item())]
+                wav = wav[..., : int(out.audio_lengths[0].item())]
 
-            # Trim left-context samples.
+            # Trim left-context samples (per-channel sample axis, so the
+            # trim amount is identical for mono and interleaved-stereo).
             if left_ctx > 0:
-                trim = min(left_ctx * self._codec.downsample_rate, wav.shape[0])
+                trim = min(left_ctx * self._codec.downsample_rate, wav.shape[-1])
                 if trim < left_ctx * self._codec.downsample_rate:
                     logger.warning(
                         "left_ctx trim (%d samples) exceeds wav length (%d); returning empty audio.",
                         left_ctx * self._codec.downsample_rate,
-                        wav.shape[0],
+                        wav.shape[-1],
                     )
-                wav = wav[trim:]
+                wav = wav[..., trim:]
 
-            audios[i] = wav
+            audios[i] = wav.reshape(-1) if self._n_channels == 1 else wav
 
         return OmniOutput(
             text_hidden_states=None,
             multimodal_outputs={"model_outputs": audios, "sr": srs},
         )
 
+    @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]
+        return default
+
     # ------------------------------------------------------------------
     # Weight loading
     # ------------------------------------------------------------------
@@ -234,11 +247,10 @@ class MossTTSCodecDecoder(nn.Module):
         for _ in weights:
             pass
 
-        codec_path = os.environ.get("MOSS_TTS_CODEC_PATH", self._codec_path)
+        codec_path = self._codec_path
         logger.info("Loading MOSS Audio Tokenizer from %s", codec_path)
 
-        codec_cfg = MossAudioTokenizerConfig.from_pretrained(codec_path)
-        codec = MossAudioTokenizerModel(codec_cfg)
+        codec_cfg, codec = self._build_codec(codec_path)
 
         model_loader = DefaultModelLoader(self.vllm_config.load_config)
         source = DefaultModelLoader.Source(
@@ -255,10 +267,19 @@ class MossTTSCodecDecoder(nn.Module):
         # and the rest stay at their random init, which produces noise that
         # sounds correct in duration but is structurally garbage.
         _SUFFIX_REMAP: list[tuple[str, str]] = [
+            # v1 (MOSS-Audio-Tokenizer) naming.
             (".self_attn.in_projs.0.", ".attn.in_proj."),
             (".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."),
             (".layer_scale_1.", ".ls1."),
             (".layer_scale_2.", ".ls2."),
             (".input_proj.", ".in_proj."),
@@ -295,11 +316,14 @@ 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)
+        self._sr_tensor = torch.tensor(int(codec_cfg.sampling_rate), dtype=torch.int32)
 
         logger.info(
-            "MOSS Audio Tokenizer loaded: sampling_rate=%d, n_vq=%d",
+            "MOSS Audio Tokenizer loaded: sampling_rate=%d, n_vq=%d, n_channels=%d",
             codec_cfg.sampling_rate,
             codec_cfg.num_quantizers,
+            self._n_channels,
         )
 
         self._maybe_enable_decoder_cudagraph(device)
@@ -310,6 +334,24 @@ class MossTTSCodecDecoder(nn.Module):
         # mirror that here.
         return {f"_codec.{name}" for name, _ in codec.named_parameters()}
 
+    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)
+            return codec_cfg, codec
+        except Exception:
+            logger.exception(
+                "Failed to instantiate official MOSS Audio Tokenizer via HF remote code; "
+                "falling back to vendored codec."
+            )
+
+        codec_cfg = MossAudioTokenizerConfig.from_pretrained(codec_path)
+        codec = MossAudioTokenizerModel(codec_cfg)
+        return codec_cfg, codec
+
     def _maybe_enable_decoder_cudagraph(self, device: torch.device) -> None:
         """Capture CUDA Graphs for the codec decoder if enforce_eager is False."""
         if getattr(self.vllm_config.model_config, "enforce_eager", True):

vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_local_depth.py

@@ -0,0 +1,223 @@
+# Copyright 2026 OpenMOSS and the vLLM-Omni team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License").
+"""Per-frame depth transformer for MossTTSLocalModel (MOSS-TTS-Local-Transformer-v1.5).
+
+A 1-layer GPT2-style block that decodes the ``n_vq`` audio codebook codes for
+one audio frame, run inside the talker's ``make_omni_output`` independent of
+vLLM's main scheduler -- mirrors ``MossTTSRealtimeLocalTransformer``
+(``modeling_moss_tts_local.py``) in role, but the algorithm and numerics are
+faithful to the official ``gpt2_decoder.py`` / ``modeling_moss_tts.py``
+(GPT2-style LayerNorm + bias, SiLU MLP, **interleaved/GPT-J-style RoPE** --
+not vLLM's neox-style concat-half rotation) rather than Realtime's
+Qwen3-style ``CodePredictorBaseModel``.
+
+Its KV cache and positions reset to 0 every audio frame: there is no
+cross-frame state. Submodule names (``h.0.*`` / ``ln_f``) match the
+checkpoint 1:1 so ``load_weights()`` needs no remapping.
+"""
+
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from vllm_omni.model_executor.models.moss_tts.modeling_moss_tts_local import (
+    _sample_token,
+)
+
+
+class _MossTTSLocalAttention(nn.Module):
+    """GPT2-style fused-QKV self-attention with interleaved RoPE.
+
+    Faithful to the official ``MossTTSNanoGPT2Attention``: ``rotate_half``
+    operates on even/odd index pairs (GPT-J style), and ``cos``/``sin`` are
+    built via ``repeat_interleave(2, dim=-1)`` rather than the neox-style
+    concat-half construction vLLM uses elsewhere.
+    """
+
+    def __init__(self, hidden_size: int, n_head: int, rope_base: float) -> None:
+        super().__init__()
+        if hidden_size % n_head != 0:
+            raise ValueError(f"hidden_size={hidden_size} must be divisible by n_head={n_head}")
+        self.n_head = n_head
+        self.head_dim = hidden_size // n_head
+        self.embed_dim = hidden_size
+        self.c_attn = nn.Linear(hidden_size, 3 * hidden_size, bias=True)
+        self.c_proj = nn.Linear(hidden_size, hidden_size, bias=True)
+        inv_freq = 1.0 / (rope_base ** (torch.arange(0, self.head_dim, 2, dtype=torch.float32) / self.head_dim))
+        self.register_buffer("inv_freq", inv_freq, persistent=False)
+
+    def _rope_cos_sin(
+        self, seq_len: int, device: torch.device, dtype: torch.dtype
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        position_ids = torch.arange(seq_len, device=device, dtype=torch.float32)
+        freqs = torch.einsum("s,d->sd", position_ids, self.inv_freq.to(device=device))
+        cos = freqs.cos().repeat_interleave(2, dim=-1).to(dtype)
+        sin = freqs.sin().repeat_interleave(2, dim=-1).to(dtype)
+        return cos.view(1, seq_len, 1, self.head_dim), sin.view(1, seq_len, 1, self.head_dim)
+
+    @staticmethod
+    def _rotate_half(x: torch.Tensor) -> torch.Tensor:
+        even = x[..., ::2]
+        odd = x[..., 1::2]
+        return torch.stack((-odd, even), dim=-1).reshape_as(x)
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        """``hidden_states``: ``(B, S, H)``. Re-prefills with a fresh causal mask over ``[0, S)`` every call."""
+        batch_size, seq_len, _ = hidden_states.shape
+        qkv = self.c_attn(hidden_states)
+        query, key, value = qkv.split(self.embed_dim, dim=-1)
+        query = query.view(batch_size, seq_len, self.n_head, self.head_dim)
+        key = key.view(batch_size, seq_len, self.n_head, self.head_dim)
+        value = value.view(batch_size, seq_len, self.n_head, self.head_dim)
+
+        cos, sin = self._rope_cos_sin(seq_len, hidden_states.device, hidden_states.dtype)
+        query = query * cos + self._rotate_half(query) * sin
+        key = key * cos + self._rotate_half(key) * sin
+
+        query = query.transpose(1, 2)
+        key = key.transpose(1, 2)
+        value = value.transpose(1, 2)
+        attn_output = F.scaled_dot_product_attention(query, key, value, is_causal=True)
+        attn_output = attn_output.transpose(1, 2).reshape(batch_size, seq_len, self.embed_dim)
+        return self.c_proj(attn_output)
+
+
+class _MossTTSLocalMLP(nn.Module):
+    def __init__(self, hidden_size: int, inner_size: int) -> None:
+        super().__init__()
+        self.fc_in = nn.Linear(hidden_size, inner_size, bias=True)
+        self.fc_out = nn.Linear(inner_size, hidden_size, bias=True)
+        self.act = nn.SiLU()
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        return self.fc_out(self.act(self.fc_in(hidden_states)))
+
+
+class _MossTTSLocalBlock(nn.Module):
+    def __init__(self, hidden_size: int, n_head: int, inner_size: int, rope_base: float, eps: float) -> None:
+        super().__init__()
+        self.ln_1 = nn.LayerNorm(hidden_size, eps=eps)
+        self.attn = _MossTTSLocalAttention(hidden_size, n_head, rope_base)
+        self.ln_2 = nn.LayerNorm(hidden_size, eps=eps)
+        self.mlp = _MossTTSLocalMLP(hidden_size, inner_size)
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        hidden_states = hidden_states + self.attn(self.ln_1(hidden_states))
+        hidden_states = hidden_states + self.mlp(self.ln_2(hidden_states))
+        return hidden_states
+
+
+class MossTTSLocalDepthTransformer(nn.Module):
+    """Per-frame depth transformer for MOSS-TTS-Local-Transformer-v1.5.
+
+    Per frame:
+      - position 0's input is the backbone's last hidden state; its output
+        feeds BOTH the binary continue/stop head (``local_text_lm_head``)
+        and codebook-0's head (``audio_lm_heads[0]``) simultaneously.
+      - codebooks 1..n_vq-1 are sampled sequentially: each sampled code is
+        re-embedded (``audio_embeddings[c]``) and appended as the next
+        position, re-prefilling the block over the growing (<=n_vq) sequence
+        with a fresh causal mask each call -- mathematically identical to
+        incremental KV-cache decoding since attention is strictly causal and
+        only the last position is ever read, but avoids any cache plumbing
+        given the trivially short sequence length.
+    """
+
+    def __init__(self, gpt2_config, hidden_size: int | None = None) -> None:
+        super().__init__()
+        self.hidden_size = int(hidden_size if hidden_size is not None else gpt2_config.n_embd)
+        n_head = int(gpt2_config.n_head)
+        inner_size = int(gpt2_config.n_inner)
+        eps = float(getattr(gpt2_config, "layer_norm_epsilon", 1e-5))
+        rope_base = float(getattr(gpt2_config, "rope_base", 1_000_000.0))
+        self.h = nn.ModuleList([_MossTTSLocalBlock(self.hidden_size, n_head, inner_size, rope_base, eps)])
+        self.ln_f = nn.LayerNorm(self.hidden_size, eps=eps)
+
+    def _forward_prefix(self, seq_embeds: torch.Tensor) -> torch.Tensor:
+        hidden_states = seq_embeds
+        for block in self.h:
+            hidden_states = block(hidden_states)
+        return self.ln_f(hidden_states)
+
+    @torch.no_grad()
+    def generate_frame(
+        self,
+        backbone_last_hidden: torch.Tensor,  # (B, H)
+        audio_lm_heads: nn.ModuleList,  # n_vq x Linear(H -> audio_vocab_size)
+        audio_embeddings: nn.ModuleList,  # n_vq x Embedding(audio_vocab_size, H)
+        local_text_lm_head: nn.Module,  # Linear(H -> 2): [continue, stop]
+        *,
+        n_vq: int,
+        do_sample: bool = True,
+        temperature: float = 1.0,
+        top_k: int = 50,
+        top_p: float = 1.0,
+        text_temperature: float = 1.0,
+        text_top_k: int = 50,
+        text_top_p: float = 1.0,
+        repetition_penalty: float = 1.0,
+        history_per_codebook: list[list[int]] | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """Generate one audio frame for batch B.
+
+        Returns ``(should_continue, codes)``: ``should_continue`` is a
+        ``(B,)`` bool tensor (``True`` iff the binary head picked the
+        "continue" candidate, i.e. logits index 0); ``codes`` is a
+        ``(B, n_vq)`` LongTensor of sampled codebook indices.
+
+        ``history_per_codebook[c]`` is a list of recently-emitted token ids
+        for codebook ``c``; when ``repetition_penalty != 1.0`` those tokens'
+        logits get scaled down before sampling (mirrors upstream's
+        ``_apply_repetition_penalty``).
+        """
+        batch_size = backbone_last_hidden.shape[0]
+        dtype = self.ln_f.weight.dtype
+
+        embeds = backbone_last_hidden.new_zeros((batch_size, n_vq, self.hidden_size), dtype=dtype)
+        embeds[:, 0, :] = backbone_last_hidden.to(dtype)
+
+        hidden = self._forward_prefix(embeds[:, :1, :])
+        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)
+        should_continue = binary_choice.eq(0)
+        import os as _os
+
+        if _os.environ.get("MOSS_TTS_DEBUG_STOP"):
+            import logging as _logging
+
+            _logging.getLogger("moss_tts_debug").warning(
+                "binary_logits=%s choice=%s", binary_logits.tolist(), binary_choice.tolist()
+            )
+
+        codes = backbone_last_hidden.new_zeros((batch_size, n_vq), dtype=torch.long)
+        for channel_index in range(n_vq):
+            channel_logits = audio_lm_heads[channel_index](local_hidden).float()
+            if (
+                repetition_penalty != 1.0
+                and history_per_codebook is not None
+                and channel_index < len(history_per_codebook)
+            ):
+                hist = history_per_codebook[channel_index]
+                if hist:
+                    hist_t = torch.tensor(hist, dtype=torch.long, device=channel_logits.device)
+                    sel = channel_logits.index_select(-1, hist_t)
+                    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)
+            codes[:, channel_index] = channel_token
+
+            if channel_index + 1 < n_vq:
+                embeds[:, channel_index + 1, :] = audio_embeddings[channel_index](channel_token).to(dtype)
+                hidden = self._forward_prefix(embeds[:, : channel_index + 2, :])
+                local_hidden = hidden[:, channel_index + 1, :]
+
+        return should_continue, codes
+
+
+__all__ = ["MossTTSLocalDepthTransformer"]

vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_talker.py

@@ -6,7 +6,7 @@
 from __future__ import annotations
 
 import copy
-from collections.abc import Iterable
+from collections.abc import Iterable, Sequence
 from typing import Any
 
 import torch
@@ -22,11 +22,15 @@ from vllm.v1.sample.metadata import SamplingMetadata
 
 from vllm_omni.model_executor.models.moss_tts.configuration_moss_tts import (
     MossTTSDelayConfig,
+    MossTTSLocalConfig,
     MossTTSRealtimeConfig,
 )
 from vllm_omni.model_executor.models.moss_tts.modeling_moss_tts_local import (
     MossTTSRealtimeLocalTransformer,
 )
+from vllm_omni.model_executor.models.moss_tts.modeling_moss_tts_local_depth import (
+    MossTTSLocalDepthTransformer,
+)
 from vllm_omni.model_executor.models.output_templates import OmniOutput
 
 logger = init_logger(__name__)
@@ -36,6 +40,40 @@ def _maybe_prefix(prefix: str, name: str) -> str:
     return f"{prefix}.{name}" if prefix else name
 
 
+def _iter_state_row_spans(
+    states: Sequence[dict[str, Any]],
+    spans: Sequence[tuple[int, int]] | None,
+    num_rows: int,
+) -> Iterable[tuple[dict[str, Any], int, int]]:
+    """Yield state/logit-row ranges without assuming equal request row counts."""
+    if num_rows <= 0 or not states:
+        return
+
+    if spans is not None:
+        if len(spans) != len(states):
+            raise RuntimeError(f"request_token_spans has {len(spans)} entries for {len(states)} request states")
+        if all(row_end <= num_rows for _, row_end in spans):
+            for state, (row_start, row_end) in zip(states, spans, strict=False):
+                if row_start < row_end:
+                    yield state, int(row_start), int(row_end)
+            return
+        if num_rows != len(states):
+            raise RuntimeError(
+                "request_token_spans describe full hidden rows, but current logits "
+                f"have {num_rows} rows for {len(states)} states and are not one-row-per-request"
+            )
+
+    if num_rows == len(states):
+        for i, state in enumerate(states):
+            yield state, i, i + 1
+        return
+
+    raise RuntimeError(
+        "MOSS-TTS compute_logits requires request_token_spans or one-row-per-request logits; "
+        f"got {num_rows} rows for {len(states)} request states"
+    )
+
+
 # ---------------------------------------------------------------------------
 # MossTTSDelayTalkerForGeneration
 # ---------------------------------------------------------------------------
@@ -209,14 +247,13 @@ class MossTTSDelayTalkerForGeneration(nn.Module):
         n_vq = self.n_vq
         device = logits.device
         vocab_size = logits.shape[-1]
-        rows_per_state = max(1, logits.shape[0] // max(1, len(states)))
 
-        for i, state in enumerate(states):
-            if not isinstance(state, dict):
-                continue
-            row_start = i * rows_per_state
-            row_end = min(row_start + rows_per_state, logits.shape[0])
-            if row_start >= row_end:
+        for state, row_start, row_end in _iter_state_row_spans(
+            states,
+            getattr(self, "_batch_state_spans", None),
+            logits.shape[0],
+        ):
+            if state is None:
                 continue
             row = logits[row_start:row_end]
 
@@ -569,6 +606,7 @@ class MossTTSDelayTalkerForGeneration(nn.Module):
         """
         if isinstance(model_outputs, OmniOutput):
             self._batch_state = None
+            self._batch_state_spans = None
             return model_outputs
 
         hidden = model_outputs  # (S, H)
@@ -590,11 +628,11 @@ class MossTTSDelayTalkerForGeneration(nn.Module):
         # Stash the per-row state for compute_logits to apply masks. logits
         # rows align with hidden rows, which align with info_dicts in order.
         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")
 
         # Per-request (start, end) hidden-row spans from the runner. In mixed
         # prefill+decode steps the per-request token counts differ, so an equal
-        # rows-per-request split would sample codes from the wrong rows; fall
-        # back to that split only when spans are unavailable.
+        # rows-per-request split would sample codes from the wrong rows.
         spans = kwargs.get("request_token_spans")
 
         # One accumulated-codes tensor per request, in batch order. The generic
@@ -607,18 +645,18 @@ class MossTTSDelayTalkerForGeneration(nn.Module):
 
         if hidden.numel() > 0:
             num_rows = hidden.shape[0]
-            rows_per_req = max(1, num_rows // max(1, len(info_dicts) or 1))
+            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):
                     per_req_codes.append(hidden.new_empty((0, self.n_vq), dtype=torch.long))
                     continue
-                if spans is not None and i < len(spans):
-                    row_start, row_end = spans[i]
-                    row_end = min(int(row_end), num_rows)
-                else:
-                    row_start = i * rows_per_req
-                    row_end = min(row_start + rows_per_req, num_rows)
+                row_start, row_end = spans[i]
+                row_end = min(int(row_end), num_rows)
                 if row_start >= row_end:
                     per_req_codes.append(hidden.new_empty((0, self.n_vq), dtype=torch.long))
                     continue
@@ -864,13 +902,12 @@ class MossTTSRealtimeTalkerForGeneration(nn.Module):
         logits = hidden_states.new_full((B, V), float("-inf"))
 
         states = self._batch_state or []
-        rows_per_state = max(1, B // max(1, len(states) or 1))
-        for i, state in enumerate(states):
-            r0 = i * rows_per_state
-            r1 = min(r0 + rows_per_state, B)
-            if r0 >= r1:
-                continue
-            if not isinstance(state, dict):
+        for state, r0, r1 in _iter_state_row_spans(
+            states,
+            getattr(self, "_batch_state_spans", None),
+            B,
+        ):
+            if state is None:
                 logits[r0:r1, self.text_pad_id] = 0.0
                 continue
             if state.get("is_stopping"):
@@ -991,6 +1028,7 @@ class MossTTSRealtimeTalkerForGeneration(nn.Module):
     ) -> OmniOutput:
         if isinstance(model_outputs, OmniOutput):
             self._batch_state = None
+            self._batch_state_spans = None
             return model_outputs
 
         hidden = model_outputs  # (S, H)
@@ -1004,9 +1042,10 @@ class MossTTSRealtimeTalkerForGeneration(nn.Module):
                 info["audio_state"] = {"is_stopping": False, "step": 0}
 
         self._batch_state = [(info["audio_state"] if isinstance(info, dict) else {}) for info in info_dicts]
+        self._batch_state_spans = kwargs.get("request_token_spans")
 
-        # See delay talker: real per-request row spans from the runner, with the
-        # equal-split as fallback when unavailable.
+        # See delay talker: real per-request row spans from the runner are
+        # required because mixed prefill+decode steps have unequal row counts.
         spans = kwargs.get("request_token_spans")
         # Per-request accumulated codes in batch order, pre-filled with empty
         # placeholders so every skip path (stopped / eos / bos / pad) keeps
@@ -1016,16 +1055,16 @@ class MossTTSRealtimeTalkerForGeneration(nn.Module):
         have_codes = False
         if hidden.numel() > 0 and info_dicts:
             num_rows = hidden.shape[0]
-            rows_per_req = max(1, num_rows // max(1, len(info_dicts) or 1))
+            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
-                if spans is not None and i < len(spans):
-                    row_start, row_end = spans[i]
-                    row_end = min(int(row_end), num_rows)
-                else:
-                    row_start = i * rows_per_req
-                    row_end = min(row_start + rows_per_req, num_rows)
+                row_start, row_end = spans[i]
+                row_end = min(int(row_end), num_rows)
                 if row_start >= row_end:
                     continue
 
@@ -1189,7 +1228,430 @@ class MossTTSRealtimeTalkerForGeneration(nn.Module):
         return loaded
 
 
+# ---------------------------------------------------------------------------
+# MossTTSLocalTalkerForGeneration
+# ---------------------------------------------------------------------------
+
+
+class MossTTSLocalTalkerForGeneration(nn.Module):
+    """Stage-0 talker for MossTTSLocalModel (MOSS-TTS-Local-Transformer-v1.5).
+
+    Architecture differs from both existing variants:
+
+    * Backbone (Qwen3; checkpoint prefix ``transformer.*``, unlike Delay/
+      Realtime's ``language_model.*``) consumes additive text+audio fusion
+      the same way Delay does (real text embeddings + masked per-codebook
+      audio embeddings from ``info["codes"]["ref"]`` / ``info["audio_codes"]
+      ["current"]``), but there is no delay-pattern countdown: the prompt
+      already ends right at the audio_start boundary, and every step from
+      the first decode onward drives exactly one audio frame.
+    * The checkpoint's own ``text_lm_head`` is loaded (trained for a
+      teacher-forced text loss) but never consulted at inference: the
+      official reference's own ``generate()`` hardcodes the next
+      text-channel input to ``audio_assistant_slot_token_id`` regardless of
+      what ``text_lm_head`` would predict, because
+      ``local_text_head_mode == "binary"`` for this checkpoint.
+      ``compute_logits`` therefore synthesises a one-hot logit row
+      (Realtime-style), forced to ``audio_assistant_slot_token_id`` while
+      continuing or ``im_end_token_id`` once the local transformer's binary
+      head decides to stop.
+    * 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
+      condition: the local transformer's binary head selects the
+      "audio_end" candidate (logits index 1).
+    """
+
+    have_multimodal_outputs: bool = True
+    has_preprocess: bool = True
+    has_postprocess: bool = True
+
+    def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
+        super().__init__()
+        self.vllm_config = vllm_config
+        self.config: MossTTSLocalConfig = vllm_config.model_config.hf_config
+
+        qwen3_cfg = self.config.qwen3_config
+        hidden_size = int(qwen3_cfg.hidden_size)
+        self.hidden_size = hidden_size
+        self.n_vq: int = int(self.config.n_vq)
+        self.audio_vocab_size: int = int(self.config.audio_vocab_size)
+        self.audio_pad_token_id: int = int(self.config.audio_pad_token_id)
+        self.text_vocab_size: int = int(qwen3_cfg.vocab_size)
+
+        self.audio_assistant_slot_token_id: int = int(self.config.audio_assistant_slot_token_id)
+        self.im_end_token_id: int = int(self.config.im_end_token_id)
+
+        # Qwen3 backbone. get_text_config() is sufficient for vLLM to size
+        # KV cache / heads correctly, as long as AutoConfig resolves to the
+        # local MossTTSLocalConfig rather than the HF remote-code class.
+        self.model = Qwen3Model(
+            vllm_config=vllm_config,
+            prefix=_maybe_prefix(prefix, "model"),
+        )
+
+        # Real text LM head -- loaded from the checkpoint but never consulted
+        # at inference (see class docstring); kept purely as a destination
+        # for ``text_lm_head.weight`` so load_weights() loads every tensor.
+        self.text_lm_head = ParallelLMHead(
+            self.text_vocab_size,
+            hidden_size,
+            bias=False,
+            prefix=_maybe_prefix(prefix, "text_lm_head"),
+        )
+        self.logits_processor = LogitsProcessor(self.text_vocab_size)
+
+        # Per-codebook audio embeddings/heads. audio_pad_token_id (==
+        # audio_vocab_size) is outside the embedding table -- masked to zero
+        # contribution rather than stored as a 1025th row (mirrors upstream
+        # ``_build_inputs_embeds``'s ``valid_mask`` masking).
+        self.audio_embeddings = nn.ModuleList(
+            [nn.Embedding(self.audio_vocab_size, hidden_size) for _ in range(self.n_vq)]
+        )
+        self.audio_lm_heads = nn.ModuleList(
+            [nn.Linear(hidden_size, self.audio_vocab_size, bias=False) for _ in range(self.n_vq)]
+        )
+        self.local_text_lm_head = nn.Linear(hidden_size, 2, bias=False)
+        self.local_transformer = MossTTSLocalDepthTransformer(self.config.gpt2_config, hidden_size=hidden_size)
+
+        self._batch_state: list[dict[str, Any]] | None = None
+        # 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.gpu_resident_buffer_keys: set[tuple[str, str]] = {
+            ("audio_codes", "current"),
+            ("audio_codes", "accumulated"),
+            ("hidden_states", "last"),
+        }
+
+    # ------------------------------------------------------------------
+    # vLLM hooks
+    # ------------------------------------------------------------------
+
+    def embed_input_ids(self, input_ids: torch.Tensor, **_: Any) -> torch.Tensor:
+        return self.model.embed_tokens(input_ids)
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        positions: torch.Tensor,
+        intermediate_tensors: IntermediateTensors | None = None,
+        inputs_embeds: torch.Tensor | None = None,
+        **_: Any,
+    ) -> torch.Tensor | IntermediateTensors:
+        return self.model(
+            input_ids=input_ids,
+            positions=positions,
+            intermediate_tensors=intermediate_tensors,
+            inputs_embeds=inputs_embeds,
+        )
+
+    def compute_logits(
+        self,
+        hidden_states: torch.Tensor | OmniOutput,
+        sampling_metadata: SamplingMetadata | None = None,
+    ) -> torch.Tensor | None:
+        """Synthesise a one-hot text logit row per request.
+
+        The text channel is never freely generated here: the upstream
+        reference hardcodes it to ``audio_assistant_slot_token_id`` while
+        audio frames keep coming, and only emits ``im_end_token_id`` once the
+        local transformer's binary head decides to stop (see class
+        docstring for why the real ``text_lm_head`` is bypassed).
+        """
+        if isinstance(hidden_states, OmniOutput):
+            hidden_states = hidden_states.text_hidden_states
+        if hidden_states is None or hidden_states.numel() == 0:
+            return None
+
+        num_rows = hidden_states.shape[0]
+        logits = hidden_states.new_full((num_rows, self.text_vocab_size), float("-inf"))
+
+        states = self._batch_state or []
+        for state, row_start, row_end in _iter_state_row_spans(
+            states,
+            getattr(self, "_batch_state_spans", None),
+            num_rows,
+        ):
+            if state is not None and state.get("is_stopping"):
+                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
+        if not states:
+            logits[:, self.audio_assistant_slot_token_id] = 0.0
+        return logits
+
+    # ------------------------------------------------------------------
+    # Embedding (text + audio codebooks, additive)
+    # ------------------------------------------------------------------
+
+    def _audio_embed(self, codes: torch.Tensor) -> torch.Tensor:
+        """``codes``: ``(T, n_vq)`` long, possibly containing ``audio_pad_token_id``
+        (outside the valid embedding range). Returns the ``(T, H)`` additive
+        audio embedding, masking out pad positions per codebook.
+        """
+        device = codes.device
+        valid_mask = codes.ne(self.audio_pad_token_id)  # (T, n_vq)
+        safe_codes = codes.masked_fill(~valid_mask, 0).clamp(0, self.audio_vocab_size - 1)
+        if self._stacked_audio_emb_w is not None:
+            n_vq_idx = torch.arange(self.n_vq, device=device)
+            codes_t = safe_codes.t()  # (n_vq, T)
+            gathered = self._stacked_audio_emb_w[n_vq_idx.unsqueeze(1), codes_t]  # (n_vq, T, H)
+            gathered = gathered * valid_mask.t().unsqueeze(-1)
+            return gathered.sum(0)
+        embeds: torch.Tensor | None = None
+        for i, emb_layer in enumerate(self.audio_embeddings):
+            col = emb_layer(safe_codes[:, i]) * valid_mask[:, i].unsqueeze(-1)
+            embeds = col if embeds is None else embeds + col
+        return embeds
+
+    def preprocess(
+        self,
+        input_ids: torch.Tensor,
+        input_embeds: torch.Tensor | None,
+        **info_dict: Any,
+    ) -> tuple[torch.Tensor, torch.Tensor, dict[str, Any]]:
+        """Build per-step input embeddings (text + audio additive fusion).
+
+        Prefill: initialise the per-request ``{is_stopping, step}`` state.
+        Decode: combine the forced text token (from ``compute_logits``) with
+        the audio codes the local transformer produced last step.
+        """
+        device = input_ids.device
+        span_len = int(input_ids.shape[0])
+        audio_state = info_dict.get("audio_state")
+        is_first_call = not isinstance(audio_state, dict)
+
+        if span_len > 1 or is_first_call:
+            embeds = self.model.embed_tokens(input_ids)
+
+            ref_codes = (info_dict.get("codes", {}) or {}).get("ref")
+            ref_offset = int(info_dict.get("ref_offset", 0))
+            last_ref_row = None
+            if isinstance(ref_codes, torch.Tensor) and ref_codes.numel() > 0:
+                if ref_codes.dim() == 1 and ref_codes.numel() % self.n_vq == 0:
+                    ref_codes = ref_codes.view(-1, self.n_vq)
+                if isinstance(ref_codes, torch.Tensor) and ref_codes.dim() == 2:
+                    end_off = ref_offset + span_len
+                    chunk = ref_codes[ref_offset:end_off]
+                    if chunk.numel() > 0 and chunk.shape[0] == span_len:
+                        codes = chunk.to(device=device, dtype=torch.long)
+                        embeds = embeds + self._audio_embed(codes)
+                        last_ref_row = codes[-1].detach()
+
+            current_codes = (
+                last_ref_row
+                if last_ref_row is not None
+                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
+
+            info_update: dict[str, Any] = {
+                "audio_state": {"is_stopping": False, "step": 0, "max_new_frames": max_new_frames},
+                "audio_codes": {"current": current_codes},
+                "ref_offset": ref_offset + span_len,
+            }
+            return input_ids, embeds, info_update
+
+        # Decode step: input_ids is the forced text token from compute_logits.
+        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)
+        else:
+            combined = text_embed
+        return input_ids, combined, {}
+
+    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()}}
+
+    # ------------------------------------------------------------------
+    # Per-step audio generation via local transformer
+    # ------------------------------------------------------------------
+
+    def make_omni_output(
+        self,
+        model_outputs: torch.Tensor | OmniOutput,
+        **kwargs: Any,
+    ) -> OmniOutput:
+        if isinstance(model_outputs, OmniOutput):
+            self._batch_state = None
+            self._batch_state_spans = None
+            return model_outputs
+
+        hidden = model_outputs  # (S, H)
+        info_dicts: list[dict[str, Any]] = (
+            kwargs.get("model_intermediate_buffer") or kwargs.get("runtime_additional_information") or []
+        )
+
+        for info in info_dicts:
+            if isinstance(info, dict) and not isinstance(info.get("audio_state"), dict):
+                info["audio_state"] = {"is_stopping": False, "step": 0, "max_new_frames": -1}
+
+        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]
+        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
+                have_codes = True
+
+        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
+        return OmniOutput(
+            text_hidden_states=hidden,
+            multimodal_outputs={
+                "codes": {"audio": per_req_codes},
+                "meta": {"finished": torch.tensor([True], dtype=torch.bool, device=device)},
+            },
+        )
+
+    # ------------------------------------------------------------------
+    # Weight loading
+    # ------------------------------------------------------------------
+
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        """Map checkpoint weight names to vLLM-Omni module names.
+
+        Checkpoint layout (prefix ``transformer.*`` for the backbone --
+        unlike Delay/Realtime's ``language_model.*``):
+          transformer.*               → model.*            (Qwen3 backbone)
+          text_lm_head.weight         → text_lm_head.weight
+          audio_embeddings.{i}.weight → audio_embeddings.{i}.weight
+          audio_lm_heads.{i}.weight   → audio_lm_heads.{i}.weight
+          local_text_lm_head.weight   → local_text_lm_head.weight
+          local_transformer.h.0.*     → local_transformer.h.0.*
+          local_transformer.ln_f.*    → local_transformer.ln_f.*
+
+        Every "tied" pair in the upstream config is stored as a full,
+        separate tensor in this checkpoint (confirmed via the raw
+        safetensors header: 438 keys, no dedup), so aside from the backbone
+        prefix this is pure 1:1 name-based loading.
+        """
+        loaded: set[str] = set()
+        params_dict = dict(self.named_parameters())
+
+        backbone_weights: list[tuple[str, torch.Tensor]] = []
+        skipped: list[str] = []
+
+        for name, tensor in weights:
+            if name.startswith("transformer."):
+                backbone_weights.append((name[len("transformer.") :], tensor))
+                continue
+            if name in params_dict:
+                default_weight_loader(params_dict[name], tensor)
+                loaded.add(name)
+            else:
+                skipped.append(name)
+
+        backbone_loaded = self.model.load_weights(iter(backbone_weights))
+        for n in backbone_loaded:
+            loaded.add(f"model.{n}")
+
+        self._stacked_audio_emb_w = torch.stack(
+            [e.weight.detach() for e in self.audio_embeddings], dim=0
+        )  # (n_vq, audio_vocab_size, hidden_size)
+
+        logger.info(
+            "[MossTTSLocal] loaded %d/%d params; skipped=%d (first 5: %s)",
+            len(loaded),
+            len(params_dict),
+            len(skipped),
+            skipped[:5],
+        )
+        not_loaded = [n for n in params_dict if n not in loaded]
+        if not_loaded:
+            logger.warning("[MossTTSLocal] %d params NOT loaded (first 5: %s)", len(not_loaded), not_loaded[:5])
+        return loaded
+
+
 __all__ = [
     "MossTTSDelayTalkerForGeneration",
     "MossTTSRealtimeTalkerForGeneration",
+    "MossTTSLocalTalkerForGeneration",
 ]

vllm_omni/model_executor/models/moss_tts/pipeline.py

@@ -80,9 +80,39 @@ MOSS_TTS_REALTIME_PIPELINE = PipelineConfig(
     ),
 )
 
-# The pipeline config is otherwise the same for all 5 variants; the per-variant
-# differences (n_vq, backbone size, generation strategy) are encoded in the
-# HF config.json and the deploy YAML. Realtime is split out because it has a
-# different talker architecture (MossTTSRealtime vs MossTTSDelayModel).
+MOSS_TTS_LOCAL_PIPELINE = PipelineConfig(
+    model_type="moss_tts_local",
+    model_arch="MossTTSLocalModel",  # different talker class: GPT2-style local depth transformer
+    stages=(
+        StagePipelineConfig(
+            stage_id=0,
+            model_stage="moss_tts_local",
+            execution_type=StageExecutionType.LLM_AR,
+            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},
+        ),
+        StagePipelineConfig(
+            stage_id=1,
+            model_stage="moss_tts_local_codec",
+            execution_type=StageExecutionType.LLM_GENERATION,
+            input_sources=(0,),
+            final_output=True,
+            final_output_type="audio",
+            engine_output_type="audio",
+            model_arch="MossTTSCodecDecoder",
+            sync_process_input_func=f"{_PROC}.talker2codec",
+            sampling_constraints={"detokenize": True},
+        ),
+    ),
+)
 
-__all__ = ["MOSS_TTS_PIPELINE", "MOSS_TTS_REALTIME_PIPELINE"]
+# The pipeline config is otherwise the same for all variants; the per-variant
+# differences (n_vq, backbone size, generation strategy) are encoded in the
+# HF config.json and the deploy YAML. Realtime and Local are split out because
+# they have different talker architectures from the delay variant
+# (MossTTSRealtime / MossTTSLocalModel vs MossTTSDelayModel).
+
+__all__ = ["MOSS_TTS_PIPELINE", "MOSS_TTS_REALTIME_PIPELINE", "MOSS_TTS_LOCAL_PIPELINE"]

vllm_omni/model_executor/models/registry.py

@@ -294,6 +294,12 @@ _OMNI_MODELS = {
         "modeling_moss_tts_talker",
         "MossTTSRealtimeTalkerForGeneration",
     ),
+    # MossTTSLocalModel: MOSS-TTS-Local-Transformer-v1.5
+    "MossTTSLocalModel": (
+        "moss_tts",
+        "modeling_moss_tts_talker",
+        "MossTTSLocalTalkerForGeneration",
+    ),
     # Stage-1 codec decoder (shared by all 5 variants)
     "MossTTSCodecDecoder": (
         "moss_tts",