Merge branch 'main' into fix/pr4709-review-followup

vllm_omni/model_executor/models/bagel/bagel.py

@@ -587,9 +587,9 @@ class OmniBagelForConditionalGeneration(BagelForConditionalGeneration):
         input_ids: torch.Tensor | None,
         positions: torch.Tensor | None,
         inputs_embeds: torch.Tensor | None,
-        req_ids: list[str],
-        num_computed_tokens: list[int],
-        num_scheduled_tokens: list[int],
+        req_ids: Sequence[str],
+        num_computed_tokens: Sequence[int],
+        num_scheduled_tokens: Sequence[int],
         input_ids_buffer: torch.Tensor | None = None,
     ) -> tuple[torch.Tensor | None, torch.Tensor | None]:
         """Restore input_ids so _adjust_positions_for_img2img can locate
@@ -599,7 +599,7 @@ class OmniBagelForConditionalGeneration(BagelForConditionalGeneration):
             input_ids = input_ids_buffer
         return input_ids, positions
 
-    def flush_pending_metadata(self, req_ids: list[str]) -> None:
+    def flush_pending_metadata(self, req_ids: Sequence[str]) -> None:
         """Map pending metadata (batch order) to req_ids after forward().
 
         Guard: if a request already has metadata with ``image_shape``

vllm_omni/model_executor/models/common/ming/audio_vae.py

@@ -138,8 +138,8 @@ class Encoder(nn.Module):
         self.norm = nn.LayerNorm(config.hidden_size)
         self.patch_size = patch_size
         if patch_size != -1:
-            aggregator_config = _qwen2_config({**encoder_args, "num_hidden_layers": 4})
-            self.aggregator = Qwen2Model(aggregator_config)
+            config.num_hidden_layers = 4
+            self.aggregator = Qwen2Model(config)
             self.cls_embed = nn.Parameter(torch.rand(1, 1, config.hidden_size))
             self.cls_embed.data.normal_(0, 0.02)
 
@@ -346,8 +346,9 @@ class AudioVAE(PreTrainedModel):
         h, _ = self.encoder(waveform)
         h = h.transpose(1, 2)
 
-        mean = torch.chunk(h, 2, dim=1)[0]
-        latent = mean.transpose(1, 2)
+        mean, scale = torch.chunk(h, 2, dim=1)
+        std = F.softplus(scale) + 1e-4
+        latent = (mean + std * torch.randn_like(mean)).transpose(1, 2)
         return latent, frame_num
 
     @torch.inference_mode()

vllm_omni/model_executor/models/common/ming/dit.py

@@ -6,6 +6,19 @@ from torch import nn
 from x_transformers.x_transformers import apply_rotary_pos_emb
 
 
+def _apply_rope_cached(t, cos, sin):
+    """Bit-exact RoPE for ``scale==1`` / interleaved-half.
+    Uses precomputed ``cos``/``sin`` per seq_len to skip per-step trig and cat."""
+    rot_dim = cos.shape[-1]
+    tr = t[..., :rot_dim]
+    x = tr.reshape(*tr.shape[:-1], -1, 2)
+    rot_half = torch.stack((-x[..., 1], x[..., 0]), dim=-1).reshape_as(tr)
+    out = tr * cos + rot_half * sin
+    if rot_dim < t.shape[-1]:  # partial rotary (unused by ming heads); keep general
+        out = torch.cat((out, t[..., rot_dim:]), dim=-1)
+    return out.type(t.dtype)
+
+
 class RMSNorm(nn.Module):
     def __init__(self, dim, eps=1e-6):
         super().__init__()
@@ -66,12 +79,33 @@ class Attention(nn.Module):
         self.to_out.append(nn.Dropout(dropout))
         self.pe_attn_head = pe_attn_head
         self.attn_mask_enabled = attn_mask_enabled
+        # cos/sin cache keyed by seq_len (freqs are deterministic in seq_len).
+        self._rope_cache: dict[int, tuple[torch.Tensor, torch.Tensor]] = {}
+        # TODO: materialize fused weights eagerly outside any torch.compile
+        # scope (e.g. post-weight-load);
+        self._qkv_w: torch.Tensor | None = None
+        self._qkv_b: torch.Tensor | None = None
+
+    def _qkv(self, x):
+        if self._qkv_w is None:
+            self._qkv_w = torch.cat([self.to_q.weight, self.to_k.weight, self.to_v.weight], dim=0)
+            if self.to_q.bias is not None:
+                self._qkv_b = torch.cat([self.to_q.bias, self.to_k.bias, self.to_v.bias], dim=0)
+        qkv = F.linear(x, self._qkv_w, self._qkv_b)
+        return qkv.split(self.inner_dim, dim=-1)
+
+    def _rope_cos_sin(self, freqs):
+        seq_len = freqs.shape[-2]
+        cached = self._rope_cache.get(seq_len)
+        if cached is None:
+            f = freqs.unsqueeze(1) if freqs.ndim == 3 else freqs  # b n d -> b 1 n d for 4D q/k
+            cached = (f.cos(), f.sin())
+            self._rope_cache[seq_len] = cached
+        return cached
 
     def forward(self, x, mask=None, rope=None):
         batch_size = x.shape[0]
-        query = self.to_q(x)
-        key = self.to_k(x)
-        value = self.to_v(x)
+        query, key, value = self._qkv(x)
 
         inner_dim = key.shape[-1]
         head_dim = inner_dim // self.heads
@@ -86,14 +120,21 @@ class Attention(nn.Module):
 
         if rope is not None:
             freqs, xpos_scale = rope
-            q_xpos_scale, k_xpos_scale = (xpos_scale, xpos_scale**-1.0) if xpos_scale is not None else (1.0, 1.0)
-            if self.pe_attn_head is not None:
-                on = self.pe_attn_head
-                query[:, :on, :, :] = apply_rotary_pos_emb(query[:, :on, :, :], freqs, q_xpos_scale)
-                key[:, :on, :, :] = apply_rotary_pos_emb(key[:, :on, :, :], freqs, k_xpos_scale)
+            scale_is_one = xpos_scale is None or (not torch.is_tensor(xpos_scale) and xpos_scale == 1.0)
+            if scale_is_one and self.pe_attn_head is None:
+                # Fast path: cached, cat-free RoPE (bit-exact for scale==1, full rotary).
+                cos, sin = self._rope_cos_sin(freqs)
+                query = _apply_rope_cached(query, cos, sin)
+                key = _apply_rope_cached(key, cos, sin)
             else:
-                query = apply_rotary_pos_emb(query, freqs, q_xpos_scale)
-                key = apply_rotary_pos_emb(key, freqs, k_xpos_scale)
+                q_xpos_scale, k_xpos_scale = (xpos_scale, xpos_scale**-1.0) if xpos_scale is not None else (1.0, 1.0)
+                if self.pe_attn_head is not None:
+                    on = self.pe_attn_head
+                    query[:, :on, :, :] = apply_rotary_pos_emb(query[:, :on, :, :], freqs, q_xpos_scale)
+                    key[:, :on, :, :] = apply_rotary_pos_emb(key[:, :on, :, :], freqs, k_xpos_scale)
+                else:
+                    query = apply_rotary_pos_emb(query, freqs, q_xpos_scale)
+                    key = apply_rotary_pos_emb(key, freqs, k_xpos_scale)
 
         if self.attn_mask_enabled and mask is not None:
             valid_sample_indices = mask.any(dim=1)

vllm_omni/model_executor/models/common/qwen3_code_predictor.py

@@ -3,7 +3,7 @@
 Shared by Qwen3-Omni and Qwen3-TTS talker models.
 
 * SDPA attention (F.scaled_dot_product_attention) with native GQA support
-* HF-compatible numerics (float32 RMSNorm, float32 RoPE, separate linear layers)
+* HF-compatible CPU/CUDA numerics with NPU-only fused norm/RoPE fast paths
 * Per-call embedding buffer to avoid cross-request aliasing
 * Pre-allocated position_ids (read-only, safe to persist)
 * torch.compile (epilogue_fusion=False) on inner transformer by default
@@ -14,7 +14,7 @@ Shared by Qwen3-Omni and Qwen3-TTS talker models.
 from __future__ import annotations
 
 import dataclasses
-from collections.abc import Iterable
+from collections.abc import Iterable, Sequence
 
 import torch
 import torch.nn as nn
@@ -24,37 +24,56 @@ from vllm.logger import init_logger
 from vllm.model_executor.layers.vocab_parallel_embedding import VocabParallelEmbedding
 from vllm.model_executor.model_loader.weight_utils import default_weight_loader
 
+from vllm_omni.diffusion.layers.custom_op import CustomOp
 from vllm_omni.platforms import current_omni_platform
 
 logger = init_logger(__name__)
 
+_GeneratorLike = torch.Generator | Sequence[torch.Generator | None] | None
+_UNIFORM_EPS = 1e-20
+
+if current_omni_platform.is_npu():
+    import torch_npu
+
 
 # ===================================================================
-# HF-numerics-compatible layers for code predictor
+# Portable layers for code predictor
 # ===================================================================
 #
 # These use plain PyTorch ops (nn.Linear, manual RMSNorm in float32,
 # rotate_half RoPE) to produce outputs numerically identical to the
-# HuggingFace reference. vLLM's fused kernels (RMSNorm, QKVParallel,
-# get_rope) introduce small precision differences that compound across
-# the autoregressive steps of the code predictor, causing severe
-# audio quality degradation.
+# HuggingFace reference on CPU/CUDA. vLLM's fused kernels (RMSNorm,
+# QKVParallel, get_rope) introduce small precision differences that compound
+# across the autoregressive steps of the code predictor, causing severe audio
+# quality degradation. The Ascend fused norm/RoPE kernels below are dispatched
+# by the current device platform.
 #
 # See: https://github.com/vllm-project/vllm-omni/issues/2274
 
 
-class _RMSNorm(nn.Module):
-    """RMSNorm matching HuggingFace's implementation exactly.
-
-    Computes variance in float32 to avoid bfloat16 precision loss.
-    """
+class _RMSNorm(CustomOp):
+    """RMSNorm with HuggingFace-compatible CPU/CUDA math and an NPU fast path."""
 
     def __init__(self, hidden_size: int, eps: float = 1e-6) -> None:
         super().__init__()
         self.weight = nn.Parameter(torch.ones(hidden_size))
         self.variance_epsilon = eps
 
-    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+    def forward_npu(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        hidden_states, _ = torch_npu.npu_rms_norm(
+            hidden_states,
+            self.weight,
+            self.variance_epsilon,
+        )
+        return hidden_states
+
+    def forward_cuda(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        return self.forward_native(hidden_states)
+
+    def forward_xpu(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        return self.forward_native(hidden_states)
+
+    def forward_native(self, hidden_states: torch.Tensor) -> torch.Tensor:
         input_dtype = hidden_states.dtype
         hidden_states = hidden_states.to(torch.float32)
         variance = hidden_states.pow(2).mean(-1, keepdim=True)
@@ -69,11 +88,8 @@ def _rotate_half(x: torch.Tensor) -> torch.Tensor:
     return torch.cat((-x2, x1), dim=-1)
 
 
-class _RotaryEmbedding(nn.Module):
-    """RoPE matching HuggingFace's implementation exactly.
-
-    Forces float32 computation for cos/sin, matching HF's torch.autocast(enabled=False).
-    """
+class _RotaryEmbedding(CustomOp):
+    """RoPE with HuggingFace-compatible CPU/CUDA math and cached NPU tables."""
 
     def __init__(self, config) -> None:
         super().__init__()
@@ -85,8 +101,24 @@ class _RotaryEmbedding(nn.Module):
         rope_theta = getattr(config, "rope_theta", 10000.0)
         inv_freq = 1.0 / (rope_theta ** (torch.arange(0, head_dim, 2, dtype=torch.float32) / head_dim))
         self.register_buffer("inv_freq", inv_freq, persistent=False)
+        if current_omni_platform.is_npu():
+            max_seq = int(getattr(config, "num_code_groups", 0) or 0) + 1
+            positions = torch.arange(max_seq, dtype=torch.float32)
+            freqs = torch.outer(positions, inv_freq)
+            emb = torch.cat((freqs, freqs), dim=-1)
+            self.register_buffer("cos_cached", emb.cos(), persistent=False)
+            self.register_buffer("sin_cached", emb.sin(), persistent=False)
 
-    def forward(self, x: torch.Tensor, position_ids: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+    def forward_npu(self, x: torch.Tensor, position_ids: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        return self.cos_cached[position_ids].to(dtype=x.dtype), self.sin_cached[position_ids].to(dtype=x.dtype)
+
+    def forward_cuda(self, x: torch.Tensor, position_ids: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        return self.forward_native(x, position_ids)
+
+    def forward_xpu(self, x: torch.Tensor, position_ids: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        return self.forward_native(x, position_ids)
+
+    def forward_native(self, x: torch.Tensor, position_ids: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
         # position_ids: [batch, seq_len]
         inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1)
         position_ids_expanded = position_ids[:, None, :].float()
@@ -142,7 +174,6 @@ class CodePredictorAttention(nn.Module):
         self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=False)
         self.q_norm = _RMSNorm(self.head_dim, eps=config.rms_norm_eps)
         self.k_norm = _RMSNorm(self.head_dim, eps=config.rms_norm_eps)
-
         if current_omni_platform.is_npu():
             if self.max_seq > 2048:
                 raise ValueError(
@@ -165,8 +196,6 @@ class CodePredictorAttention(nn.Module):
         bsz: int,
         seq_len: int,
     ) -> torch.Tensor:
-        import torch_npu
-
         q_f, k_f, v_f = q, k, v
         if self.is_gqa:
             k_f = (
@@ -227,10 +256,13 @@ class CodePredictorAttention(nn.Module):
         # cos/sin are [batch, seq_len, head_dim], need unsqueeze at dim=1 for heads
         cos = cos.unsqueeze(1)  # [batch, 1, seq_len, head_dim]
         sin = sin.unsqueeze(1)
-        q = (q * cos) + (_rotate_half(q) * sin)
-        k = (k * cos) + (_rotate_half(k) * sin)
-
-        if not current_omni_platform.is_npu():
+        if current_omni_platform.is_npu():
+            q = torch_npu.npu_rotary_mul(q, cos, sin)
+            k = torch_npu.npu_rotary_mul(k, cos, sin)
+            attn_out = self._forward_npu_attention(q, k, v, bsz, seq_len)
+        else:
+            q = (q * cos) + (_rotate_half(q) * sin)
+            k = (k * cos) + (_rotate_half(k) * sin)
             attn_out = F.scaled_dot_product_attention(
                 q,
                 k,
@@ -239,8 +271,6 @@ class CodePredictorAttention(nn.Module):
                 is_causal=True,
                 enable_gqa=self.is_gqa,
             )
-        else:
-            attn_out = self._forward_npu_attention(q, k, v, bsz, seq_len)
 
         attn_out = attn_out.transpose(1, 2).reshape(bsz, seq_len, -1)
         return self.o_proj(attn_out)
@@ -287,10 +317,17 @@ class CodePredictorDecoderLayer(nn.Module):
         residual = hidden_states
         hidden_states = self.input_layernorm(hidden_states)
         hidden_states = self.self_attn(hidden_states, position_embeddings)
-        hidden_states = residual + hidden_states
-
-        residual = hidden_states
-        hidden_states = self.post_attention_layernorm(hidden_states)
+        if current_omni_platform.is_npu():
+            hidden_states, _, residual = torch_npu.npu_add_rms_norm(
+                hidden_states,
+                residual,
+                self.post_attention_layernorm.weight,
+                self.post_attention_layernorm.variance_epsilon,
+            )
+        else:
+            hidden_states = residual + hidden_states
+            residual = hidden_states
+            hidden_states = self.post_attention_layernorm(hidden_states)
         hidden_states = self.mlp(hidden_states)
         hidden_states = residual + hidden_states
         return hidden_states
@@ -605,6 +642,30 @@ class CodePredictorWrapper(nn.Module):
                 values.add(parsed)
         return values
 
+    @staticmethod
+    def _normalize_generators(
+        generator: _GeneratorLike, batch_size: int
+    ) -> torch.Generator | list[torch.Generator | None] | None:
+        if generator is None or isinstance(generator, torch.Generator):
+            return generator
+
+        row_generators = list(generator)
+        if len(row_generators) != batch_size:
+            raise ValueError(f"Expected {batch_size} per-row generators, but got {len(row_generators)}.")
+        return row_generators
+
+    @classmethod
+    def _sample_codes_gumbel(cls, logits: torch.Tensor, generator: _GeneratorLike = None) -> torch.Tensor:
+        """Sample ``logits`` via Gumbel-max with optional per-row generators."""
+        row_generators = cls._normalize_generators(generator, int(logits.shape[0]))
+        u = torch.empty_like(logits, dtype=torch.float32)
+        if isinstance(row_generators, list):
+            for row, row_generator in enumerate(row_generators):
+                u[row : row + 1].uniform_(_UNIFORM_EPS, 1.0 - _UNIFORM_EPS, generator=row_generator)
+        else:
+            u.uniform_(_UNIFORM_EPS, 1.0 - _UNIFORM_EPS, generator=row_generators)
+        return (logits.float() - torch.log(-torch.log(u))).argmax(dim=-1, keepdim=True)
+
     def _prefix_seq_lens(self, max_seq: int) -> list[int]:
         all_seq_lens = list(range(2, max_seq))
         if not self._prefix_graph_seq_lens:
@@ -742,6 +803,27 @@ class CodePredictorWrapper(nn.Module):
     #  Forward -- re-prefill + inline sampling
     # ------------------------------------------------------------------
 
+    @staticmethod
+    def _multinomial(
+        probs: torch.Tensor,
+        generator: torch.Generator | None,
+        generators: Sequence[torch.Generator | None] | None,
+    ) -> torch.Tensor:
+        """Sample one code per row, optionally with per-row generators.
+
+        Per-row generators keep explicitly-seeded requests deterministic in a
+        multi-row batch: each row consumes draws only from its own generator,
+        so the transformer forward can stay batched (#4883).
+        """
+        if generators is None:
+            return torch.multinomial(probs, num_samples=1, generator=generator)
+        return torch.cat(
+            [
+                torch.multinomial(probs[row : row + 1], num_samples=1, generator=row_generator)
+                for row, row_generator in enumerate(generators)
+            ]
+        )
+
     @torch.inference_mode()
     def forward(
         self,
@@ -753,9 +835,13 @@ class CodePredictorWrapper(nn.Module):
         top_k: int = 50,
         top_p: float = 1.0,
         generator: torch.Generator | None = None,
+        generators: Sequence[torch.Generator | None] | None = None,
     ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
         """Predict residual codebooks 1..G-1 autoregressively via re-prefill."""
         bsz = int(layer0_code.shape[0])
+        if generators is not None and len(generators) != bsz:
+            raise ValueError(f"generators must have one entry per row: got {len(generators)} for batch {bsz}")
+        sample_generator: _GeneratorLike = generators if generators is not None else generator
         num_groups = self._num_groups
         device = layer0_code.device
 
@@ -823,7 +909,6 @@ class CodePredictorWrapper(nn.Module):
             # Use captured device graph if available, otherwise call compiled fn.
             device_graph_entry = self._device_graphs.get(graph_key)
 
-            # Run transformer (device graph replay or compiled forward)
             if device_graph_entry is not None:
                 device_graph_entry[0].replay()
                 hidden_out = device_graph_entry[1]
@@ -832,9 +917,20 @@ class CodePredictorWrapper(nn.Module):
 
             logits = lm_heads[step - 1](hidden_out[:bsz, step, :])
 
-            # Sample next code
+            # Sample next code via Gumbel-max.
+            #
+            # ``argmax_i(logits_i + Gumbel_i)`` with
+            # ``Gumbel_i = -log(-log(u_i)), u_i ~ Uniform(0, 1)`` is
+            # distributionally identical to sampling from ``softmax(logits)``.
+            # In this file the motivations are practical rather than graph
+            # related: it is measurably cheaper than ``softmax + multinomial``
+            # on the B x 2048 shapes used here, it stays well-defined for
+            # degenerate masked rows with a surviving finite entry (and is more
+            # defensive than ``multinomial`` around fully-masked/NaN inputs),
+            # and the helper below can honor either one batch generator or one
+            # generator per seeded row.
             if stored_mode:
-                # "stored" mode: top-k -> top-p -> softmax -> multinomial
+                # "stored" mode: top-k -> top-p -> Gumbel-max
                 if s_top_k > 0:
                     topk_vals, _ = logits.topk(s_top_k, dim=-1)
                     logits = logits.masked_fill(logits < topk_vals[:, -1:], float("-inf"))
@@ -845,17 +941,15 @@ class CodePredictorWrapper(nn.Module):
                     remove_mask = (cumulative_probs - sorted_probs) >= s_top_p
                     sorted_logits[remove_mask] = float("-inf")
                     logits = sorted_logits.scatter(1, sorted_idx, sorted_logits)
-                probs = F.softmax(logits, dim=-1, dtype=torch.float32)
-                code = torch.multinomial(probs, num_samples=1, generator=generator)
+                code = self._sample_codes_gumbel(logits, generator=sample_generator)
             else:
-                # "per_call" mode: temperature-scaled + top-k
+                # "per_call" mode: temperature-scaled + top-k -> Gumbel-max
                 if use_sampling:
                     scaled = logits * inv_temperature
                     if top_k > 0:
                         topk_vals, _ = scaled.topk(top_k, dim=-1)
                         scaled = scaled.masked_fill(scaled < topk_vals[:, -1:], float("-inf"))
-                    probs = F.softmax(scaled, dim=-1, dtype=torch.float32)
-                    code = torch.multinomial(probs, num_samples=1, generator=generator)
+                    code = self._sample_codes_gumbel(scaled, generator=sample_generator)
                 else:
                     code = logits.argmax(dim=-1, keepdim=True)
 
@@ -878,6 +972,18 @@ class CodePredictorWrapper(nn.Module):
     #  Weight loading
     # ------------------------------------------------------------------
 
+    def _prepare_npu_weights(self) -> None:
+        from vllm_ascend.utils import maybe_trans_nz
+
+        linear_count = 0
+        with torch.no_grad():
+            # Pack linear weights once for NPU matmul.
+            for module in self.modules():
+                if isinstance(module, nn.Linear):
+                    module.weight.data = maybe_trans_nz(module.weight.data)
+                    linear_count += 1
+        logger.info("Prepared NPU code predictor weights: linear=%d", linear_count)
+
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
         """Load weights directly (no fused projection remapping needed)."""
         loaded: set[str] = set()
@@ -904,4 +1010,7 @@ class CodePredictorWrapper(nn.Module):
             weight_loader(param, w)
             loaded.add(name)
 
+        if current_omni_platform.is_npu():
+            self._prepare_npu_weights()
+
         return loaded

vllm_omni/model_executor/models/cosyvoice3/pipeline.py

@@ -5,9 +5,8 @@
 Stage 0: Talker   — text prompt → speech tokens (LLM autoregressive).
 Stage 1: Code2Wav — flow-matching decoder → acoustic features → waveform.
   * ``sync_process_input_func`` (``text2flow_token_only``) runs when
-    ``deploy.async_chunk=false``: stage 1 builds full-sequence flow input
-    from the emitted speech tokens. ``text2flow`` is the
-    ``custom_process_input_func`` variant.
+    ``deploy.async_chunk=false``: stage 1 allocates placeholder slots while
+    bulk tensors arrive via ``text2flow_full_payload`` on the connector.
   * ``async_chunk_process_next_stage_input_func`` runs when
     ``deploy.async_chunk=true``: stage 0 streams codec chunks to stage 1
     through the shared-memory connector.
@@ -47,7 +46,6 @@ COSYVOICE3_PIPELINE = PipelineConfig(
             final_output=True,
             final_output_type="audio",
             engine_output_type="latent",
-            custom_process_input_func=f"{_PROC}.text2flow",
             sync_process_input_func=f"{_PROC}.text2flow_token_only",
         ),
     ),

vllm_omni/model_executor/models/higgs_audio_v2/higgs_audio_v2_tokenizer.py

@@ -22,6 +22,7 @@ generation.
 
 from __future__ import annotations
 
+import json
 import os
 import re
 from typing import Any
@@ -170,14 +171,25 @@ _AUDIO_TOKENIZER_PATH_ENVS = (
 )
 
 
+def _is_higgs_audio_tokenizer_config(config_path: str) -> bool:
+    try:
+        with open(config_path, encoding="utf-8") as f:
+            config = json.load(f)
+    except (OSError, json.JSONDecodeError):
+        return False
+    return config.get("model_type") == "higgs_audio_v2_tokenizer"
+
+
 def _normalize_audio_tokenizer_dir(path: str) -> str | None:
     if not path:
         return None
     expanded = os.path.abspath(os.path.expanduser(path))
-    if os.path.isfile(os.path.join(expanded, "config.json")):
+    config_path = os.path.join(expanded, "config.json")
+    if os.path.isfile(config_path) and _is_higgs_audio_tokenizer_config(config_path):
         return expanded
     nested = os.path.join(expanded, _K2_OMNIVOICE_SUBDIR)
-    if os.path.isfile(os.path.join(nested, "config.json")):
+    nested_config_path = os.path.join(nested, "config.json")
+    if os.path.isfile(nested_config_path) and _is_higgs_audio_tokenizer_config(nested_config_path):
         return nested
     return None
 
@@ -201,7 +213,9 @@ def _resolve_audio_tokenizer_dir() -> str | None:
         filename=f"{_K2_OMNIVOICE_SUBDIR}/config.json",
     )
     if isinstance(cached_config, str) and os.path.isfile(cached_config):
-        return os.path.dirname(cached_config)
+        candidate = _normalize_audio_tokenizer_dir(os.path.dirname(cached_config))
+        if candidate is not None:
+            return candidate
 
     from huggingface_hub.constants import HF_HUB_CACHE
 
@@ -240,7 +254,9 @@ def _load_audio_tokenizer():
             _K2_OMNIVOICE_REPO,
             allow_patterns=[f"{_K2_OMNIVOICE_SUBDIR}/*"],
         )
-        audio_tokenizer_dir = os.path.join(repo_path, _K2_OMNIVOICE_SUBDIR)
+        audio_tokenizer_dir = _normalize_audio_tokenizer_dir(repo_path)
+        if audio_tokenizer_dir is None:
+            raise RuntimeError(f"Downloaded {_K2_OMNIVOICE_REPO} does not contain a valid Higgs audio tokenizer")
     device = current_omni_platform.get_torch_device()
     model = HiggsAudioV2TokenizerModel.from_pretrained(audio_tokenizer_dir).to(device)
     return model.eval()

vllm_omni/model_executor/models/hunyuan_image3/hunyuan_image3.py

@@ -222,14 +222,8 @@ class HunyuanModel(HunYuanModel):
             if self.config.tie_word_embeddings and "lm_head.weight" in name:
                 continue
 
-            if self.quant_config is not None and (scale_name := self.quant_config.get_cache_scale(name)):
-                # Loading kv cache scales for compressed-tensors quantization
-                param = params_dict[scale_name]
-                weight_loader = getattr(param, "weight_loader", default_weight_loader)
-                loaded_weight = loaded_weight[0]
-                weight_loader(param, loaded_weight)
-                continue
-
+            # KV-cache scales are renamed to .attn.{k,v,q}_scale by the
+            # quant_config mapper in the outer AutoWeightsLoader.
             is_found = False
             for param_name, weight_name, shard_id in stacked_params_mapping:
                 if weight_name not in name:

vllm_omni/model_executor/models/mammoth_moda2/mammoth_moda2.py

@@ -483,19 +483,15 @@ class MammothModa2Qwen2ForCausalLM(nn.Module, SupportsPP):
         ]
         params_dict = dict(self.named_parameters(remove_duplicate=False))
         loaded_params: set[str] = set()
+        if self.quant_config is not None:
+            cache_scale_mapper = self.quant_config.get_cache_scale_mapper()
+            if cache_scale_mapper is not None:
+                weights = cache_scale_mapper.apply(weights)
 
         for name, loaded_weight in weights:
             if "rotary_emb.inv_freq" in name:
                 continue
 
-            if self.quant_config is not None and (scale_name := self.quant_config.get_cache_scale(name)):
-                param = params_dict[scale_name]
-                weight_loader = getattr(param, "weight_loader", default_weight_loader)
-                loaded_weight = loaded_weight if loaded_weight.dim() == 0 else loaded_weight[0]
-                weight_loader(param, loaded_weight)
-                loaded_params.add(scale_name)
-                continue
-
             for param_name, weight_name, shard_id in stacked_params_mapping:
                 if weight_name not in name:
                     continue

vllm_omni/model_executor/models/mammoth_moda2/pipeline.py

@@ -0,0 +1,66 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""MammothModa2 pipeline topology (frozen).
+
+Stage 0: AR  — multimodal understanding + latent generation
+Stage 1: DiT — latent → image
+
+For text/image understanding tasks (text output only), use MAMMOTH_MODA2_AR_PIPELINE.
+"""
+
+from vllm_omni.config.stage_config import (
+    PipelineConfig,
+    StageExecutionType,
+    StagePipelineConfig,
+)
+
+_PROC = "vllm_omni.model_executor.stage_input_processors.mammoth_moda2"
+
+MAMMOTH_MODA2_PIPELINE = PipelineConfig(
+    model_type="mammoth_moda2",
+    model_arch="MammothModa2ForConditionalGeneration",
+    hf_architectures=("Mammothmoda2Model", "MammothModa2ForConditionalGeneration"),
+    stages=(
+        StagePipelineConfig(
+            stage_id=0,
+            model_stage="ar",
+            execution_type=StageExecutionType.LLM_AR,
+            input_sources=(),
+            final_output=False,
+            owns_tokenizer=True,
+            requires_multimodal_data=True,
+            engine_output_type="latent",
+        ),
+        StagePipelineConfig(
+            stage_id=1,
+            model_stage="dit",
+            execution_type=StageExecutionType.LLM_GENERATION,
+            input_sources=(0,),
+            final_output=True,
+            final_output_type="image",
+            owns_tokenizer=False,
+            requires_multimodal_data=False,
+            engine_output_type="image",
+            custom_process_input_func=f"{_PROC}.ar2dit",
+        ),
+    ),
+)
+
+# Single-stage AR variant for understanding / summarization tasks
+MAMMOTH_MODA2_AR_PIPELINE = PipelineConfig(
+    model_type="mammoth_moda2_ar",
+    model_arch="MammothModa2ForConditionalGeneration",
+    stages=(
+        StagePipelineConfig(
+            stage_id=0,
+            model_stage="ar",
+            execution_type=StageExecutionType.LLM_AR,
+            input_sources=(),
+            final_output=True,
+            final_output_type="text",
+            owns_tokenizer=True,
+            requires_multimodal_data=True,
+            engine_output_type="text",
+        ),
+    ),
+)

vllm_omni/model_executor/models/mimo_audio/mimo_audio_llm.py

@@ -1153,17 +1153,14 @@ class MiMoAudioLLMForConditionalGeneration(nn.Module, SupportsMultiModal, Suppor
         params_dict = dict(self.named_parameters(remove_duplicate=False))
 
         loaded_params: set[str] = set()
+        if self.quant_config is not None:
+            cache_scale_mapper = self.quant_config.get_cache_scale_mapper()
+            if cache_scale_mapper is not None:
+                weights = cache_scale_mapper.apply(weights)
+
         for name, loaded_weight in weights:
             if name.startswith("model."):
                 name = "model." + name
-            if self.quant_config is not None and (scale_name := self.quant_config.get_cache_scale(name)):
-                # Loading kv cache quantization scales
-                param = params_dict[scale_name]
-                weight_loader = getattr(param, "weight_loader", default_weight_loader)
-                loaded_weight = loaded_weight if loaded_weight.dim() == 0 else loaded_weight[0]
-                weight_loader(param, loaded_weight)
-                loaded_params.add(scale_name)
-                continue
 
             for param_name, weight_name, shard_id in stacked_params_mapping:
                 if name.startswith("input_local_transformer.") or name.startswith("local_transformer."):

vllm_omni/model_executor/models/mimo_audio/pipeline.py

@@ -56,7 +56,6 @@ MIMO_AUDIO_PIPELINE = PipelineConfig(
             final_output=True,
             final_output_type="audio",
             engine_output_type="audio",
-            custom_process_input_func=f"{_PROC}.llm2code2wav",
             sync_process_input_func=f"{_PROC}.llm2code2wav_token_only",
             sampling_constraints={"detokenize": False},
         ),

vllm_omni/model_executor/models/ming_tts/config_ming_tts.py

@@ -22,6 +22,7 @@ from .constants import (
     DEFAULT_SIGMA,
     DEFAULT_TEMPERATURE,
     HISTORY_PATCH_SIZE,
+    INITIAL_LATENT_CHUNK_SIZE,
     KEY_CFG,
     KEY_CHUNK_ID,
     KEY_DECODE_STEP,
@@ -239,6 +240,7 @@ class MingTTSConfig:
     max_decode_steps: int = MAX_DECODE_STEPS
 
     latent_chunk_size: int = LATENT_CHUNK_SIZE
+    initial_latent_chunk_size: int = INITIAL_LATENT_CHUNK_SIZE
     latent_left_context: int = LATENT_LEFT_CONTEXT
 
     text_eos_token_id: int = TEXT_EOS_TOKEN_ID

vllm_omni/model_executor/models/ming_tts/constants.py

@@ -58,6 +58,7 @@ DEFAULT_TEMPERATURE = 0.0
 
 # Connector / Stage-2 streaming defaults (runtime tuning)
 LATENT_CHUNK_SIZE = 25
+INITIAL_LATENT_CHUNK_SIZE = 4
 LATENT_LEFT_CONTEXT = 0
 MAX_DECODE_STEPS = 200
 

vllm_omni/model_executor/models/ming_tts/ming_tts_llm.py

@@ -326,6 +326,28 @@ class MingLLMModel(nn.Module):
         try:
             from .fm.cfm_cudagraph import CFMGraphExecutor, CFMSampler
 
+            # TODO(perf):
+            #   (1) compile per-block (dit_model.blocks) instead of the
+            #   whole forward for tighter fusion / fewer graph breaks;
+            #   (2) move the compile + a synthetic warmup forward to a post-weight-load hook
+            dit_model = self.flowloss.cfm.model
+            if not getattr(dit_model, "_ming_compiled", False):
+                try:
+                    # Scoped options (avoid mutating global inductor config).
+                    dit_model.forward = torch.compile(
+                        dit_model.forward,
+                        fullgraph=False,
+                        dynamic=False,
+                        options={
+                            "triton.cudagraphs": False,
+                            "triton.cudagraph_trees": False,
+                        },
+                    )
+                    dit_model._ming_compiled = True
+                    logger.info("Ming CFM DiT torch.compile enabled (inductor cudagraphs off).")
+                except Exception as exc:
+                    logger.warning("Ming CFM torch.compile failed (%s); using uncompiled DiT.", exc)
+
             sampler = CFMSampler(self.flowloss.cfm.model, steps=_CFM_STEPS)
             self._cfm_graph = CFMGraphExecutor(
                 sampler,

vllm_omni/model_executor/models/ming_tts/patch_emission.py

@@ -127,7 +127,7 @@ def _validate_ming_decode_window(
             continue
         max_steps = _resolve_runtime_int(info, KEY_MAX_DECODE_STEPS, default_max_decode_steps)
         min_steps = _resolve_optional_runtime_int(info, KEY_MIN_DECODE_STEPS, 0)
-        min_required = max(min_stop_step + 1, min_steps)
+        min_required = max(min_stop_step + 2, min_steps)
         if max_steps < min_required:
             req_id = info.get(KEY_REQUEST_ID, f"idx={i}")
             raise ValueError(
@@ -148,7 +148,7 @@ def _resolve_ming_stop_decision(
     audio_dummy_token_id: int,
     text_eos_token_id: int,
 ) -> tuple[str, bool, bool, int, int]:
-    min_required_decode_steps = max(min_stop_step + 1, min_decode_steps)
+    min_required_decode_steps = max(min_stop_step + 2, min_decode_steps)
     if max_decode_steps < min_required_decode_steps:
         raise RuntimeError(
             "Invalid Ming decode window: "

vllm_omni/model_executor/models/ming_tts/validation.py

@@ -172,6 +172,9 @@ def validate_ming_tts_config(cfg: Any) -> None:
 
     if cfg.latent_chunk_size <= 0:
         raise ValueError(f"latent_chunk_size must be > 0, got {cfg.latent_chunk_size}.")
+    initial_latent_chunk_size = getattr(cfg, "initial_latent_chunk_size", 0)
+    if initial_latent_chunk_size < 0:
+        raise ValueError(f"initial_latent_chunk_size must be >= 0, got {initial_latent_chunk_size}.")
     if cfg.latent_left_context < 0:
         raise ValueError(f"latent_left_context must be >= 0, got {cfg.latent_left_context}.")
     if cfg.max_decode_steps <= 0:

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/cuda_graph_streaming_decoder_wrapper.py

@@ -0,0 +1,164 @@
+# Copyright 2026 OpenMOSS and the vLLM-Omni team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License").
+"""CUDA graph wrapper for MOSS-TTS Local streaming codec decode."""
+
+from __future__ import annotations
+
+import time
+from collections.abc import Callable
+from dataclasses import dataclass
+
+import torch
+import torch.nn as nn
+from torch.cuda import CUDAGraph
+from vllm.logger import init_logger
+from vllm.platforms import current_platform
+
+logger = init_logger(__name__)
+
+
+@dataclass
+class _CapturedStreamingDecodeGraph:
+    graph: CUDAGraph
+    static_codes: torch.Tensor
+    static_lengths: torch.Tensor
+    static_audio: torch.Tensor
+    static_audio_lengths: torch.Tensor
+
+
+class CUDAGraphStreamingDecoderWrapper:
+    """Replay fixed-size streaming ``_decode_frame`` steps with CUDA graphs.
+
+    The wrapped codec must already be inside ``codec.streaming(batch_size)``.
+    Graphs are keyed by the streaming step frame count ``T``. The batch width is
+    fixed to the codec stream slot count.
+    """
+
+    def __init__(
+        self,
+        codec: nn.Module,
+        *,
+        batch_size: int,
+        num_quantizers: int,
+        reset_streaming_state: Callable[[], None] | None = None,
+    ) -> None:
+        self.codec = codec
+        self.batch_size = int(batch_size)
+        self.num_quantizers = int(num_quantizers)
+        self.reset_streaming_state = reset_streaming_state
+        self.graphs: dict[int, _CapturedStreamingDecodeGraph] = {}
+        self._pool = None
+        self._warmed_up = False
+
+    @property
+    def is_ready(self) -> bool:
+        return bool(self.graphs)
+
+    @property
+    def capture_sizes(self) -> list[int]:
+        return sorted(self.graphs)
+
+    @torch.no_grad()
+    def warmup(self, device: torch.device, capture_sizes: list[int]) -> None:
+        if self._warmed_up:
+            return
+        self._warmed_up = True
+        if device.type != "cuda":
+            return
+        sizes = sorted({int(size) for size in capture_sizes if int(size) > 0}, reverse=True)
+        if not sizes:
+            return
+
+        logger.info(
+            "MOSS-TTS streaming decoder CUDA graph warmup: batch_size=%d n_vq=%d capture_sizes=%s",
+            self.batch_size,
+            self.num_quantizers,
+            list(reversed(sizes)),
+        )
+        start_s = time.perf_counter()
+        with torch.cuda.device(device):
+            for size in sizes:
+                try:
+                    self._capture(size, device)
+                    logger.info("  Captured MOSS-TTS streaming decoder CUDA graph for size=%d", size)
+                except Exception:
+                    self.graphs.pop(size, None)
+                    logger.warning(
+                        "  Failed to capture MOSS-TTS streaming decoder CUDA graph for size=%d; "
+                        "this size will use eager decode",
+                        size,
+                        exc_info=True,
+                    )
+        logger.info(
+            "MOSS-TTS streaming decoder CUDA graph warmup complete: %d/%d captured in %.1f ms",
+            len(self.graphs),
+            len(sizes),
+            (time.perf_counter() - start_s) * 1000.0,
+        )
+
+    @torch.no_grad()
+    def _capture(self, size: int, device: torch.device) -> None:
+        codes = torch.zeros(
+            self.num_quantizers,
+            self.batch_size,
+            size,
+            dtype=torch.long,
+            device=device,
+        )
+        lengths = torch.full((self.batch_size,), size, dtype=torch.long, device=device)
+        exec_mask = torch.ones(self.batch_size, dtype=torch.bool, device=device)
+
+        self.codec._set_streaming_exec_mask(exec_mask)
+        stream = torch.cuda.Stream()
+        stream.wait_stream(torch.cuda.current_stream())
+        with torch.cuda.stream(stream):
+            _ = self.codec._decode_frame(codes, lengths)
+        torch.cuda.current_stream().wait_stream(stream)
+        torch.accelerator.synchronize(device)
+
+        if self.reset_streaming_state is not None:
+            self.reset_streaming_state()
+        self.codec._set_streaming_exec_mask(exec_mask)
+        if self._pool is None:
+            self._pool = current_platform.get_global_graph_pool()
+        graph = CUDAGraph()
+        with torch.cuda.graph(graph, pool=self._pool, capture_error_mode="thread_local"):
+            output = self.codec._decode_frame(codes, lengths)
+
+        if output.audio is None or output.audio_lengths is None:
+            raise RuntimeError("MOSS-TTS streaming decoder CUDA graph capture produced empty audio.")
+
+        self.graphs[size] = _CapturedStreamingDecodeGraph(
+            graph=graph,
+            static_codes=codes,
+            static_lengths=lengths,
+            static_audio=output.audio,
+            static_audio_lengths=output.audio_lengths,
+        )
+
+    @torch.no_grad()
+    def decode(
+        self,
+        codes_step: torch.Tensor,
+        exec_mask: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor] | None:
+        if not codes_step.is_cuda or torch.cuda.is_current_stream_capturing():
+            return None
+        n_vq, batch_size, step_t = codes_step.shape
+        if int(n_vq) != self.num_quantizers or int(batch_size) != self.batch_size:
+            return None
+        # The streaming codec advances its internal KV offsets by tensor T, so
+        # bucket-padding a smaller step into a larger captured graph would write
+        # padded frames into state. Capture exact step sizes instead.
+        entry = self.graphs.get(int(step_t))
+        if entry is None:
+            return None
+
+        self.codec._set_streaming_exec_mask(exec_mask)
+        entry.static_codes.copy_(codes_step)
+        entry.graph.replay()
+        return entry.static_audio, entry.static_audio_lengths
+
+
+__all__ = ["CUDAGraphStreamingDecoderWrapper"]

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,15 @@ 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.cuda_graph_streaming_decoder_wrapper import (
+    CUDAGraphStreamingDecoderWrapper,
+)
 from vllm_omni.model_executor.models.moss_tts.moss_codec_cudagraph import (
     MossTTSCUDAGraphCodecWrapper,
 )
@@ -27,6 +37,121 @@ 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,
+        cudagraph_capture_sizes: list[int] | None = None,
+    ) -> 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))
+        self._cudagraph_wrapper: CUDAGraphStreamingDecoderWrapper | None = None
+        capture_sizes = sorted({int(size) for size in (cudagraph_capture_sizes or []) if int(size) > 0})
+        if capture_sizes and self._device.type == "cuda":
+            self._cudagraph_wrapper = CUDAGraphStreamingDecoderWrapper(
+                codec,
+                batch_size=self._batch_size,
+                num_quantizers=self._n_vq,
+                reset_streaming_state=lambda: self.reset_slots(list(range(self._batch_size))),
+            )
+            self._cudagraph_wrapper.warmup(self._device, capture_sizes)
+            self.reset_slots(list(range(self._batch_size)))
+            if not self._cudagraph_wrapper.is_ready:
+                self._cudagraph_wrapper = None
+
+    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
+
+        graph_output: tuple[torch.Tensor, torch.Tensor] | None = None
+        if self._cudagraph_wrapper is not None:
+            graph_output = self._cudagraph_wrapper.decode(codes_step, exec_mask)
+
+        used_cudagraph = graph_output is not None
+        if used_cudagraph:
+            audio_tensor, lengths_tensor = graph_output
+        else:
+            self._codec._set_streaming_exec_mask(exec_mask)
+            result = self._codec._decode_frame(codes_step, codes_lengths)
+            if result.audio is None:
+                return {}
+            audio_tensor = result.audio
+            lengths_tensor = result.audio_lengths
+
+        audio = audio_tensor.detach().to("cpu", torch.float32)
+        lengths = lengths_tensor.detach().to("cpu") if lengths_tensor 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 +200,19 @@ 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_chunk_frames: int = self._connector_int("codec_chunk_frames", default=0)
+        self._stream_max_step_frames: int = self._stream_chunk_frames or 100
+        self._stream_req_slots: dict[str, int] = {}
+        self._stream_pending_codes: dict[str, list[torch.Tensor]] = {}
+        self._stream_starved_reqs: set[str] = set()
+        self._codec_streaming: bool = self._connector_bool("codec_streaming", default=False)
+        default_graph_capture_max = self._stream_chunk_frames or min(self._stream_max_step_frames, 32)
+        default_graph_capture_sizes = list(range(1, max(1, default_graph_capture_max) + 1))
+        self._streaming_cudagraph_capture_sizes = self._streaming_cudagraph_capture_sizes_from_compilation_config(
+            default_graph_capture_sizes
+        )
 
     # ------------------------------------------------------------------
     # vLLM-Omni stubs (codec has no AR loop)
@@ -112,9 +250,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 +261,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:
@@ -137,27 +275,59 @@ class MossTTSCodecDecoder(nn.Module):
                 },
             )
 
+        if self._codec_streaming and runtime_additional_information is None:
+            return OmniOutput(
+                text_hidden_states=None,
+                multimodal_outputs={
+                    "model_outputs": [empty] * num_req,
+                    "sr": [sr_tensor] * num_req,
+                },
+            )
+
         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 +335,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", self._codec_streaming))
+            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 +359,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,21 +399,273 @@ 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", self._codec_streaming)):
+                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,
+            cudagraph_capture_sizes=self._streaming_cudagraph_capture_sizes,
+        )
+        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
 
+    def _connector_bool(self, name: str, default: bool = False) -> bool:
+        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 bool(extra_cfg[name])
+        return default
+
+    def _streaming_cudagraph_capture_sizes_from_compilation_config(
+        self,
+        default: list[int],
+    ) -> list[int]:
+        if getattr(self.vllm_config.model_config, "enforce_eager", True):
+            return []
+        compilation_config = getattr(self.vllm_config, "compilation_config", None)
+        capture_sizes = getattr(compilation_config, "cudagraph_capture_sizes", None)
+        if capture_sizes:
+            max_step_frames = max(1, int(self._stream_max_step_frames))
+            return sorted({int(size) for size in capture_sizes if 0 < int(size) <= max_step_frames})
+        return list(default)
+
     # ------------------------------------------------------------------
     # Weight loading
     # ------------------------------------------------------------------
@@ -272,14 +707,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 +725,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 +767,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(
@@ -326,7 +785,9 @@ class MossTTSCodecDecoder(nn.Module):
             self._n_channels,
         )
 
-        self._maybe_enable_decoder_cudagraph(device)
+        self._configure_decoder_cudagraph(device)
+        if self._codec_streaming and self._streaming_cudagraph_capture_sizes:
+            self._ensure_stream_session()
 
         # vLLM's track_weights_loading() compares the returned set against
         # ``self.named_parameters()``. After ``self._codec = codec`` above,
@@ -336,28 +797,43 @@ 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)
         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."""
+    def _configure_decoder_cudagraph(self, device: torch.device) -> None:
+        """Select the codec CUDA Graph path.
+
+        ``enforce_eager`` is the single graph on/off switch. If graphing is
+        enabled, ``codec_streaming`` decides whether decode uses the persistent
+        streaming-state wrapper or the offline full-chunk wrapper.
+        """
         if getattr(self.vllm_config.model_config, "enforce_eager", True):
+            self._streaming_cudagraph_capture_sizes = []
             return
         if self._codec is None:
             return
+        if self._codec_streaming:
+            logger.info(
+                "MOSS-TTS codec CUDA Graph selected streaming wrapper: capture_sizes=%s",
+                self._streaming_cudagraph_capture_sizes,
+            )
+            return
+
+        self._enable_non_streaming_decoder_cudagraph(device)
+
+    def _enable_non_streaming_decoder_cudagraph(self, device: torch.device) -> None:
+        if self._codec is None:
+            return
 
         # Read capture sizes from the connector's extra config (same convention
         # as Qwen3-TTS), falling back to a sensible default covering common

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/moss_codec_cudagraph.py

@@ -1,6 +1,6 @@
 """CUDA Graph acceleration for the MOSS Audio Tokenizer codec decoder.
 
-Captures MossAudioTokenizerModel._decode for a set of fixed frame-count
+Captures MossAudioTokenizerModel._decode_frame for a set of fixed frame-count
 bucket sizes, then replays the captured graph at inference time to eliminate
 kernel-launch overhead.  Inputs that exceed all captured sizes fall back to
 eager execution transparently.
@@ -24,7 +24,7 @@ logger = init_logger(__name__)
 
 
 class MossTTSCUDAGraphCodecWrapper:
-    """CUDA Graph wrapper for MossAudioTokenizerModel._decode.
+    """CUDA Graph wrapper for MossAudioTokenizerModel._decode_frame.
 
     Graphs are keyed by padded_T (int).  On each call the actual T is
     bucket-matched to the smallest pre-captured size >= T.  The static code
@@ -99,7 +99,7 @@ class MossTTSCUDAGraphCodecWrapper:
             dummy_codes = torch.zeros(nq, 1, size, dtype=torch.long, device=device)
             dummy_lengths = torch.tensor([size], dtype=torch.long, device=device)
             with torch.no_grad():
-                _ = self.model._decode(dummy_codes, dummy_lengths)
+                _ = self.model._decode_frame(dummy_codes, dummy_lengths)
 
         torch.accelerator.synchronize(device)
 
@@ -132,13 +132,13 @@ class MossTTSCUDAGraphCodecWrapper:
 
         # Extra eager warmup inside capture to ensure all kernels are compiled.
         with torch.no_grad():
-            _ = self.model._decode(static_codes, static_lengths)
+            _ = self.model._decode_frame(static_codes, static_lengths)
         torch.accelerator.synchronize(device)
 
         graph = CUDAGraph()
         with torch.no_grad():
             with torch.cuda.graph(graph, pool=current_platform.get_global_graph_pool()):
-                static_out = self.model._decode(static_codes, static_lengths)
+                static_out = self.model._decode_frame(static_codes, static_lengths)
 
         self.graphs[size] = graph
         self.static_codes[size] = static_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/models/omnivoice/pipeline.py

@@ -0,0 +1,31 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""OmniVoice pipeline topology (frozen).
+
+Single-stage diffusion model for text-to-speech.
+Stage 0: DiT — text → iterative unmasking → codebook tokens → DAC decode → 24kHz audio
+"""
+
+from vllm_omni.config.stage_config import (
+    PipelineConfig,
+    StageExecutionType,
+    StagePipelineConfig,
+)
+
+OMNIVOICE_PIPELINE = PipelineConfig(
+    model_type="omnivoice",
+    model_arch="OmniVoicePipeline",
+    stages=(
+        StagePipelineConfig(
+            stage_id=0,
+            model_stage="dit",
+            execution_type=StageExecutionType.DIFFUSION,
+            input_sources=(),
+            final_output=True,
+            final_output_type="audio",
+            owns_tokenizer=False,
+            requires_multimodal_data=False,
+            engine_output_type="audio",
+        ),
+    ),
+)

vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni.py

@@ -17,7 +17,6 @@ from vllm.logger import init_logger
 from vllm.model_executor.models.interfaces import SupportsMRoPE, SupportsMultiModal, SupportsPP
 from vllm.model_executor.models.qwen2_5_omni_thinker import (
     Qwen2_5OmniConditionalGenerationMixin,
-    Qwen2_5OmniThinkerProcessingInfo,
 )
 from vllm.model_executor.models.utils import init_vllm_registered_model, maybe_prefix
 from vllm.model_executor.models.vision import (
@@ -39,6 +38,7 @@ from vllm_omni.model_executor.models.output_templates import OmniOutput
 from vllm_omni.model_executor.models.qwen2_5_omni.qwen2_5_omni_thinker import (
     Qwen2_5OmniThinkerDummyInputsBuilder,
     Qwen2_5OmniThinkerMultiModalProcessor,
+    Qwen2_5OmniThinkerProcessingInfo,
 )
 from vllm_omni.model_executor.models.utils import add_prefix_to_loaded_weights, split_list_into_ranges
 from vllm_omni.platforms import current_omni_platform

vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_thinker.py

@@ -4,6 +4,7 @@ from collections.abc import Iterable, Iterator, Mapping
 from typing import Any
 
 import numpy as np
+import PIL.Image as PILImage
 import torch
 from torch import nn
 from transformers.models.qwen2_5_omni.configuration_qwen2_5_omni import (
@@ -26,16 +27,21 @@ from vllm.model_executor.models.module_mapping import MultiModelKeys
 from vllm.model_executor.models.qwen2_5_omni_thinker import (
     Qwen2_5OmniAudioFeatureInputs,
     Qwen2_5OmniThinkerDummyInputsBuilder,
-    Qwen2_5OmniThinkerProcessingInfo,
     check_interleaved_audio_video,
     merge_interleaved_embeddings,
 )
 from vllm.model_executor.models.qwen2_5_omni_thinker import (
     Qwen2_5OmniConditionalGenerationMixin as Qwen2_5OmniConditionalGenerationMixinBase,
 )
+from vllm.model_executor.models.qwen2_5_omni_thinker import (
+    Qwen2_5OmniThinkerMultiModalDataParser as _Qwen2_5OmniThinkerMultiModalDataParserBase,
+)
 from vllm.model_executor.models.qwen2_5_omni_thinker import (
     Qwen2_5OmniThinkerMultiModalProcessor as _Qwen2_5OmniThinkerMultiModalProcessorBase,
 )
+from vllm.model_executor.models.qwen2_5_omni_thinker import (
+    Qwen2_5OmniThinkerProcessingInfo as _Qwen2_5OmniThinkerProcessingInfoBase,
+)
 from vllm.model_executor.models.qwen2_5_vl import (
     Qwen2_5_VisionTransformer,
     Qwen2_5_VLImageEmbeddingInputs,
@@ -56,12 +62,13 @@ from vllm.multimodal.inputs import (
     MultiModalFeatureSpec,
     MultiModalKwargsItems,
 )
-from vllm.multimodal.parse import MultiModalDataItems
+from vllm.multimodal.parse import MultiModalDataItems, VideoProcessorItems
 from vllm.multimodal.processing.processor import (
     MultiModalPromptUpdates,
     PlaceholderFeaturesInfo,
 )
 from vllm.sequence import IntermediateTensors
+from vllm.utils.collection_utils import is_list_of
 
 from vllm_omni.quantization.component_config import (
     resolve_encoder_quant_config,
@@ -74,11 +81,143 @@ except (ImportError, ModuleNotFoundError):
 logger = init_logger(__name__)
 
 
+def _presampled_videos_hf_kwargs(
+    mm_data: Mapping[str, object],
+    mm_kwargs: Mapping[str, object],
+) -> Mapping[str, object]:
+    """Adjust HF video kwargs for videos pre-sampled by vLLM's video loader.
+
+    When ``video_metadata`` is present (emitted by
+    ``Qwen2_5OmniVideoProcessorItems``), the frames were already sampled
+    according to ``media_io_kwargs``, so the HF processor must not re-sample
+    them, and it needs the sampled fps (instead of its default) to compute
+    ``video_second_per_grid`` correctly. Otherwise the audio/video temporal
+    alignment breaks under ``use_audio_in_video=True``.
+    """
+    video_metadata = mm_data.get("video_metadata")
+    if not video_metadata:
+        return mm_kwargs
+
+    mm_kwargs = dict(mm_kwargs)
+    videos_kwargs = dict(mm_kwargs.get("videos_kwargs") or {})
+    videos_kwargs["do_sample_frames"] = False
+
+    def _compute_sampled_video_fps(metadata) -> float | None:
+        duration = getattr(metadata, "duration", None)
+        indices = getattr(metadata, "frames_indices", None)
+        if not duration or not indices or duration <= 0:
+            return None
+        return len(indices) / float(duration)
+
+    # An explicit user-provided fps takes precedence.
+    if "fps" not in videos_kwargs:
+        fps_values = [_compute_sampled_video_fps(m) for m in video_metadata]
+        # HF accepts a single fps per call and uses it for every video's
+        # video_second_per_grid.
+        known_fps = [fps for fps in fps_values if fps is not None]
+        unique_fps = set(known_fps)
+        if len(unique_fps) == 1:
+            videos_kwargs["fps"] = known_fps[0]
+        elif len(unique_fps) > 1:
+            logger.warning(
+                f"Mixed sampled FPS {sorted(unique_fps)} in one request; HF accepts a single fps, using {known_fps[0]}."
+            )
+            videos_kwargs["fps"] = known_fps[0]
+
+    mm_kwargs["videos_kwargs"] = videos_kwargs
+    return mm_kwargs
+
+
+class Qwen2_5OmniVideoProcessorItems(VideoProcessorItems):
+    """Video items that carry the loader's ``(frames, metadata)`` tuples.
+
+    The tuples are kept in ``data`` so that mm hashing and cache-miss
+    re-parsing retain the metadata; ``get_processor_data`` unpacks them into
+    the ``videos`` + ``video_metadata`` arguments of the HF processor.
+    """
+
+    def get_processor_data(self) -> Mapping[str, object]:
+        from transformers.video_utils import VideoMetadata
+
+        videos = [frames for frames, _ in self.data]
+        video_metadata = [
+            VideoMetadata(**{k: v for k, v in metadata.items() if k != "do_sample_frames"}) for _, metadata in self.data
+        ]
+        return {"videos": videos, "video_metadata": video_metadata}
+
+
+class Qwen2_5OmniThinkerMultiModalDataParser(_Qwen2_5OmniThinkerMultiModalDataParserBase):
+    def _parse_video_data(self, data):
+        if data is None or isinstance(data, dict) or self.is_embeddings(data):
+            return super()._parse_video_data(data)
+
+        # Normalize to a list of per-video items (same as the base parser).
+        if (is_list_of(data, PILImage.Image) and len(data) > 0) or (
+            isinstance(data, (np.ndarray, torch.Tensor)) and data.ndim == 4
+        ):
+            data_items = [data]
+        elif isinstance(data, (np.ndarray, torch.Tensor)):
+            data_items = [elem for elem in data]
+        elif isinstance(data, tuple) and len(data) == 2:
+            data_items = [data]
+        else:
+            data_items = data
+
+        videos_with_metadata = [self._get_video_with_metadata(item) for item in data_items]
+
+        # Metadata is optional: videos fetched by vLLM's media connector
+        # arrive as (frames, metadata) tuples, while plain arrays (e.g.
+        # offline `multi_modal_data` or dummy profiling inputs) carry no
+        # metadata and parse as before.
+        if any(metadata is None for _, metadata in videos_with_metadata):
+            return super()._parse_video_data(data)
+
+        return Qwen2_5OmniVideoProcessorItems(
+            videos_with_metadata,
+            metadata=[metadata for _, metadata in videos_with_metadata],
+        )
+
+
+class Qwen2_5OmniThinkerProcessingInfo(_Qwen2_5OmniThinkerProcessingInfoBase):
+    def get_data_parser(self):
+        feature_extractor = self.get_feature_extractor()
+
+        return Qwen2_5OmniThinkerMultiModalDataParser(
+            spatial_merge_size=self.get_hf_config().vision_config.spatial_merge_size,
+            target_sr=feature_extractor.sampling_rate,
+            target_channels=self.get_target_channels(),
+            expected_hidden_size=self._get_expected_hidden_size(),
+        )
+
+
 class Qwen2_5OmniThinkerMultiModalProcessor(
     _Qwen2_5OmniThinkerMultiModalProcessorBase,
 ):
     """Override to fix use_audio_in_video detection when mm cache returns None."""
 
+    def _cached_apply_hf_processor(self, inputs, timing_ctx):
+        # If use_audio_in_video, process video and audio together; otherwise,
+        # skip cache to avoid errors.
+        # Prevents partial cache hits from causing processing failures.
+        if inputs.hf_processor_mm_kwargs.get("use_audio_in_video"):
+            return self._apply_hf_processor(inputs, timing_ctx)
+        return super()._cached_apply_hf_processor(inputs, timing_ctx)
+
+    def _call_hf_processor(
+        self,
+        prompt: str,
+        mm_data: Mapping[str, object],
+        mm_kwargs: Mapping[str, object],
+        tok_kwargs: Mapping[str, object],
+    ):
+        mm_kwargs = _presampled_videos_hf_kwargs(mm_data, mm_kwargs)
+        return super()._call_hf_processor(
+            prompt=prompt,
+            mm_data=mm_data,
+            mm_kwargs=mm_kwargs,
+            tok_kwargs=tok_kwargs,
+        )
+
     def _maybe_apply_prompt_updates(
         self,
         mm_items: MultiModalDataItems,

vllm_omni/model_executor/models/qwen2_5_omni/qwen2_old.py

@@ -335,19 +335,10 @@ class Qwen2Model(nn.Module):
             ("gate_up_proj", "up_proj", 1),
         ]
         params_dict = dict(self.named_parameters(remove_duplicate=False))
-        get_cache_scale = getattr(self.quant_config, "get_cache_scale", None)
         loaded_params: set[str] = set()
         for name, loaded_weight in weights:
             if "rotary_emb.inv_freq" in name:
                 continue
-            if get_cache_scale is not None and (scale_name := get_cache_scale(name)):
-                # Loading kv cache quantization scales
-                param = params_dict[scale_name]
-                weight_loader = getattr(param, "weight_loader", default_weight_loader)
-                loaded_weight = loaded_weight if loaded_weight.dim() == 0 else loaded_weight[0]
-                weight_loader(param, loaded_weight)
-                loaded_params.add(scale_name)
-                continue
             for param_name, weight_name, shard_id in stacked_params_mapping:
                 if weight_name not in name:
                     continue

vllm_omni/model_executor/models/qwen3_omni/pipeline.py

@@ -9,6 +9,7 @@ Stage 2: Code2Wav — RVQ codes → audio waveform
 
 from transformers import Qwen3OmniMoeConfig
 
+from vllm_omni.config.endpoint_policy import EndpointRestriction, OmniServingCapability
 from vllm_omni.config.stage_config import (
     PipelineConfig,
     StageExecutionType,
@@ -21,6 +22,12 @@ _PROC = "vllm_omni.model_executor.stage_input_processors.qwen3_omni"
 QWEN3_OMNI_PIPELINE = PipelineConfig(
     model_type="qwen3_omni_moe",
     model_arch="Qwen3OmniMoeForConditionalGeneration",
+    endpoint_restrictions=(
+        EndpointRestriction(
+            OmniServingCapability.COMPLETIONS,
+            "Qwen3-Omni requires chat template structure for thinker-talker handoff. Use /v1/chat/completions instead.",
+        ),
+    ),
     stages=(
         StagePipelineConfig(
             stage_id=0,
@@ -44,7 +51,6 @@ QWEN3_OMNI_PIPELINE = PipelineConfig(
             input_sources=(0,),
             hf_config_name="talker_config",
             engine_output_type="latent",
-            custom_process_input_func=f"{_PROC}.thinker2talker",
             sync_process_input_func=f"{_PROC}.thinker2talker_token_only",
             custom_process_next_stage_input_func=(f"{_PROC}.talker2code2wav_full_payload"),
             async_chunk_process_next_stage_input_func=(f"{_PROC}.talker2code2wav_async_chunk"),
@@ -62,7 +68,6 @@ QWEN3_OMNI_PIPELINE = PipelineConfig(
             final_output_type="audio",
             hf_config_name="thinker_config",
             engine_output_type="audio",
-            custom_process_input_func=f"{_PROC}.talker2code2wav",
             sampling_constraints={"detokenize": True},
         ),
     ),

vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_thinker.py

@@ -74,7 +74,6 @@ from vllm.model_executor.models.module_mapping import MultiModelKeys
 from vllm.model_executor.models.qwen2_5_omni_thinker import (
     Qwen2_5OmniAudioFeatureInputs,
     Qwen2_5OmniThinkerDummyInputsBuilder,
-    Qwen2_5OmniThinkerMultiModalProcessor,
     check_interleaved_audio_video,
     merge_interleaved_embeddings,
 )
@@ -120,6 +119,9 @@ from vllm.transformers_utils.processor import cached_processor_from_config
 
 from vllm_omni.model_executor.models.qwen2_5_omni.qwen2_5_omni_thinker import (
     Qwen2_5OmniConditionalGenerationMixin,
+    Qwen2_5OmniThinkerMultiModalDataParser,
+    Qwen2_5OmniThinkerMultiModalProcessor,
+    _presampled_videos_hf_kwargs,
 )
 from vllm_omni.quantization.component_config import (
     PRE_QUANTIZED_METHODS,
@@ -621,6 +623,19 @@ class Qwen3OmniMoeThinkerProcessingInfo(Qwen2AudioProcessingInfo, Qwen2_5_VLProc
         assert isinstance(feature_extractor, WhisperFeatureExtractor)
         return feature_extractor
 
+    def get_data_parser(self):
+        feature_extractor = self.get_feature_extractor()
+
+        # Install the Omni parser, which keeps video metadata from vLLM's
+        # video loader (see Qwen2_5OmniVideoProcessorItems); the parser
+        # inherited from Qwen2AudioProcessingInfo would drop it.
+        return Qwen2_5OmniThinkerMultiModalDataParser(
+            spatial_merge_size=self.get_hf_config().vision_config.spatial_merge_size,
+            target_sr=feature_extractor.sampling_rate,
+            target_channels=1,
+            expected_hidden_size=self._get_expected_hidden_size(),
+        )
+
     def get_supported_mm_limits(self) -> Mapping[str, int | None]:
         return {"audio": None, "image": None, "video": None}
 
@@ -671,8 +686,11 @@ class Qwen3OmniMoeThinkerMultiModalProcessor(
         tok_kwargs: Mapping[str, object],
     ) -> BatchFeature:
         mm_data = dict(mm_data)
+        mm_kwargs = dict(mm_kwargs)
         audios = mm_data.pop("audios", [])
 
+        mm_kwargs = dict(_presampled_videos_hf_kwargs(mm_data, mm_kwargs))
+
         def pad_to_hop_length(x: np.ndarray, hop_length: int) -> np.ndarray:
             length = x.shape[-1]
             if length % hop_length != 0:
@@ -927,9 +945,13 @@ class Qwen3OmniMoeThinkerMultiModalProcessor(
 
             audio_in_video_item_idx += 1
 
-            second_per_grid_ts = hf_processor_mm_kwargs.get("second_per_grid_ts", None)
-            if second_per_grid_ts:
-                video_second_per_grid_t = second_per_grid_ts[item_idx]
+            # Prefer the HF-computed value (temporal_patch_size / sampled fps)
+            # over request kwargs;
+            second_per_grid_ts = out_mm_data.get("second_per_grid_ts")
+            if second_per_grid_ts is None:
+                second_per_grid_ts = hf_processor_mm_kwargs.get("second_per_grid_ts", None)
+            if second_per_grid_ts is not None:
+                video_second_per_grid_t = float(second_per_grid_ts[item_idx])
             else:
                 video_second_per_grid_t = 2.0
 

vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_code_predictor_vllm.py

@@ -13,6 +13,7 @@ from vllm_omni.model_executor.models.common.qwen3_code_predictor import (
     CodePredictorWrapper,
     CodePredictorWrapperConfig,
 )
+from vllm_omni.platforms import current_omni_platform
 
 from .configuration_qwen3_tts import Qwen3TTSTalkerCodePredictorConfig, Qwen3TTSTalkerConfig
 
@@ -21,7 +22,7 @@ Qwen3TTSTalkerCodePredictorModelVLLM = CodePredictorBaseModel
 
 
 class Qwen3TTSTalkerCodePredictorForConditionalGenerationVLLM(CodePredictorWrapper):
-    """Qwen3-TTS code predictor (CUDA graphs, per-call sampling, projection)."""
+    """Qwen3-TTS code predictor (per-call sampling, projection)."""
 
     def __init__(
         self,
@@ -35,7 +36,7 @@ class Qwen3TTSTalkerCodePredictorForConditionalGenerationVLLM(CodePredictorWrapp
             vllm_config=vllm_config,
             cp_config=config,
             wrapper_config=CodePredictorWrapperConfig(
-                use_cuda_graphs=True,
+                use_cuda_graphs=current_omni_platform.is_npu(),
                 use_parallel_embedding=False,
                 use_projection=(config.hidden_size != talker_config.hidden_size),
                 return_proj_buf=False,

vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_talker.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 import copy
-from collections.abc import Iterable, Mapping
+from collections.abc import Iterable, Mapping, Sequence
 from typing import Any
 
 import numpy as np
@@ -314,6 +314,12 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module):
         # tensor read; postprocess receives the tail-only slice instead, which
         # avoids ~18 ms merge + ~6 ms write per step (Sy0307 profile, #3665).
         self.requires_full_prefix_cached_hidden_states = False
+        # Stage 1 (code2wav) consumes runtime audio codes, not hidden states,
+        # so the inter-stage pooler payload never needs a CPU hidden-states
+        # view. Opting out lets the runner skip the per-step blocking
+        # ``hidden_states[:n].to("cpu")`` (measured at 24% of stage-0 on-CPU
+        # samples at c=64 on H20).
+        self.omni_pooler_payload_include_hidden = False
         # ``codes.audio`` is only needed for future prefix-hit reconstruction
         # after a request has produced codec rows. Keep per-step rows on GPU and
         # materialize the CPU OmniTensorPrefixCache entry once at completion.
@@ -323,7 +329,33 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module):
         # OmniGPUModelRunner will store talker_mtp output under this key in
         # per-request additional_information.
         self.talker_mtp_output_key = ("codes", "audio")
-
+        self.talker_mtp_graph_safe = True
+        # talker_mtp samples with per-row generators, so explicitly-seeded
+        # requests stay batched instead of one scalar forward per row (#4883).
+        # This is only valid while talker_mtp receives the *unpadded* active
+        # batch, i.e. when it is NOT graph-wrapped. The runner wraps talker_mtp
+        # (padded batches + a single captured RNG stream) whenever full
+        # cudagraphs are enabled, so disable per-row generators in that case.
+        # Consequence: under full cudagraphs, per-request ``tts_local_seed`` is
+        # not reproducible (matching Qwen3-Omni, which has no per-row seeding);
+        # eager mode keeps the per-row generator fast-path.
+        # TODO(#4923): the model should not read ``cudagraph_mode`` — it is an
+        # upstream (runner) concern, and the runner already re-derives the same
+        # wrap decision in ``_init_talker_mtp``. Once all TTS models declare
+        # ``talker_mtp_graph_safe`` we should move this gating into the model
+        # runner (models only expose the static capability flag, and users can
+        # then configure ``cudagraph_mode`` freely). This also means mtp seeding
+        # is silently dropped when the decode batch > 1 under full cudagraphs;
+        # the follow-up should refactor the per-row generator in the runner to
+        # be cudagraph-safe so seeded batches stay reproducible.
+        cudagraph_mode = getattr(vllm_config.compilation_config, "cudagraph_mode", None)
+        talker_mtp_graph_wrapped = (
+            self.talker_mtp_graph_safe and cudagraph_mode is not None and cudagraph_mode.has_full_cudagraphs()
+        )
+        self.talker_mtp_accepts_per_row_generators = not talker_mtp_graph_wrapped
+        self.use_async_omni_output = True
+        self.eager_omni_postprocess_before_async_output = True
+        self.omni_pooler_payload_include_hidden = False
         self.model = Qwen3Model(vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model"))
 
         if get_pp_group().is_last_rank:
@@ -1060,6 +1092,7 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module):
         top_k: int | None = None,
         top_p: float | None = None,
         generator: torch.Generator | None = None,
+        generators: Sequence[torch.Generator | None] | None = None,
         **kwargs: Any,
     ) -> tuple[torch.Tensor, torch.Tensor]:
         """GPU fast-path used by OmniGPUModelRunner to predict residual codebooks (1..Q-1).
@@ -1099,6 +1132,7 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module):
             top_k=top_k,
             top_p=top_p,
             generator=generator,
+            generators=generators,
         )  # [B, Q]
 
         # Map invalid layer-0 ids (e.g. EOS) to PAD=0 so SpeechTokenizer sees only real codes.

vllm_omni/model_executor/models/voxcpm2/pipeline.py

@@ -26,6 +26,7 @@ VOXCPM2_PIPELINE = PipelineConfig(
             final_output_type="audio",
             owns_tokenizer=True,
             engine_output_type="audio",
+            scheduler_cls="vllm_omni.model_executor.models.voxcpm2.scheduler.VoxCPM2OmniARAsyncScheduler",
             sampling_constraints={
                 "detokenize": False,
                 "stop_token_ids": [1],

vllm_omni/model_executor/models/voxcpm2/runtime_config.py

@@ -0,0 +1,123 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from __future__ import annotations
+
+import dataclasses
+from typing import Any
+
+from vllm.logger import init_logger
+
+logger = init_logger(__name__)
+
+
+@dataclasses.dataclass(frozen=True)
+class _VoxCPM2RuntimeConfig:
+    enable_profiling: bool = False
+    enable_nvtx_profile: bool = False
+    enable_loc_dit_layer_nvtx: bool = False
+    enable_loc_dit_fused_qkv: bool = True
+    enable_loc_dit_fused_mlp: bool = True
+    enable_loc_dit_zero_dt_cache: bool = True
+    enable_loc_dit_fast_rope: bool = False
+    enable_loc_dit_skip_qkv_contig: bool = True
+    enable_loc_dit_reduce_overhead_no_cg: bool = False
+    enable_loc_dit_fullgraph_no_cg: bool = False
+    cfg_cutoff_ratio: float = 1.0
+    decode_graph_capture_policy: str = "all"
+    enable_vae_cuda_graph: bool = False
+    enable_cfm_cuda_graph: bool = False
+    enable_cfm_prealloc_output: bool = False
+    enable_batched_cfm: bool = True
+    deterministic_cfm_noise: bool = False
+    deterministic_cfm_seed: int = 20260601
+    audio_emit_every: int = 1
+    vae_decode_every: int = 3
+    enable_delayed_audio_copy: bool = False
+    delayed_audio_copy_use_events: bool = False
+    coalesce_audio_d2h: bool = True
+    enable_batched_vae_decode: bool = True
+    enable_batched_fsq_fusion: bool = False
+    batched_fsq_fusion_max_batch: int = 32
+    enable_batched_prefill_tail: bool = False
+    enable_unified_decode_graph: bool = True
+    unified_decode_graph_max_batch_size: int = 64
+
+    @classmethod
+    def from_vllm_config(cls, vllm_config: Any) -> _VoxCPM2RuntimeConfig:
+        model_config = vllm_config.model_config
+        raw = getattr(model_config, "voxcpm2_runtime_config", None)
+        if raw is None:
+            raw = getattr(getattr(model_config, "hf_config", None), "voxcpm2_runtime_config", None)
+        if raw is None:
+            return cls()
+        if hasattr(raw, "to_dict"):
+            raw = raw.to_dict()
+        elif not isinstance(raw, dict) and hasattr(raw, "__dict__"):
+            raw = vars(raw)
+        if not isinstance(raw, dict):
+            logger.warning("Ignoring invalid voxcpm2_runtime_config=%r; expected a dict.", raw)
+            return cls()
+
+        fields = {field.name: field for field in dataclasses.fields(cls)}
+        values: dict[str, Any] = {}
+        for key, value in raw.items():
+            if key not in fields:
+                logger.warning("Ignoring unknown VoxCPM2 runtime config key: %s", key)
+                continue
+            default = getattr(cls(), key)
+            values[key] = cls._coerce_value(key, value, default)
+        return cls(**values)._normalized()
+
+    def _normalized(self) -> _VoxCPM2RuntimeConfig:
+        cfg = self
+        if cfg.enable_batched_cfm and cfg.enable_cfm_cuda_graph:
+            logger.warning(
+                "VoxCPM2 batched CFM and CFM CUDA Graph are mutually exclusive; "
+                "disabling CFM CUDA Graph and keeping batched CFM."
+            )
+            cfg = dataclasses.replace(cfg, enable_cfm_cuda_graph=False)
+        if cfg.enable_cfm_prealloc_output and not cfg.enable_cfm_cuda_graph:
+            logger.warning("VoxCPM2 CFM preallocated output requires CFM CUDA Graph; disabling it.")
+            cfg = dataclasses.replace(cfg, enable_cfm_prealloc_output=False)
+        if cfg.delayed_audio_copy_use_events and not cfg.enable_delayed_audio_copy:
+            logger.warning("VoxCPM2 delayed audio copy events require delayed audio copy; disabling event polling.")
+            cfg = dataclasses.replace(cfg, delayed_audio_copy_use_events=False)
+        if cfg.enable_batched_vae_decode and (
+            cfg.enable_delayed_audio_copy or cfg.audio_emit_every > 1 or cfg.enable_vae_cuda_graph
+        ):
+            logger.warning(
+                "VoxCPM2 batched VAE decode is incompatible with delayed audio copy, "
+                "audio_emit_every > 1, and VAE CUDA Graph; disabling batched VAE decode."
+            )
+            cfg = dataclasses.replace(cfg, enable_batched_vae_decode=False)
+        if cfg.enable_unified_decode_graph:
+            if cfg.enable_cfm_cuda_graph:
+                logger.info("VoxCPM2 unified decode graph supersedes independent CFM CUDA Graph; disabling it.")
+                cfg = dataclasses.replace(cfg, enable_cfm_cuda_graph=False, enable_cfm_prealloc_output=False)
+            if not cfg.enable_batched_cfm:
+                logger.info("VoxCPM2 unified decode graph requires batched CFM; enabling it.")
+                cfg = dataclasses.replace(cfg, enable_batched_cfm=True)
+        return cfg
+
+    def unified_decode_graph_available(self, *, use_cuda_graph: bool) -> bool:
+        return bool(use_cuda_graph and self.enable_unified_decode_graph and not self.deterministic_cfm_noise)
+
+    @staticmethod
+    def _coerce_value(key: str, value: Any, default: Any) -> Any:
+        if isinstance(default, bool):
+            if isinstance(value, str):
+                return value.strip().lower() in ("1", "true", "yes", "on")
+            return bool(value)
+        if isinstance(default, int) and not isinstance(default, bool):
+            value = int(value)
+            if key in {"audio_emit_every", "vae_decode_every", "batched_fsq_fusion_max_batch"}:
+                return max(1, value)
+            return value
+        if isinstance(default, float):
+            if key == "cfg_cutoff_ratio":
+                return min(1.0, max(0.0, float(value)))
+            return float(value)
+        if isinstance(default, str):
+            return str(value)
+        return value

vllm_omni/model_executor/models/voxcpm2/scheduler.py

@@ -0,0 +1,42 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from __future__ import annotations
+
+from vllm.v1.request import RequestStatus
+
+from vllm_omni.core.sched.omni_ar_scheduler import OmniARAsyncScheduler
+from vllm_omni.platforms import current_omni_platform
+
+from .runtime_config import _VoxCPM2RuntimeConfig
+
+
+class VoxCPM2OmniARAsyncScheduler(OmniARAsyncScheduler):
+    """VoxCPM2 scheduler variant for full unified decode graph serving.
+
+    VoxCPM2's full unified decode graph only applies to pure decode batches.
+    When a decode-ready request is already running, this scheduler defers new
+    waiting admissions for the current tick so the decode batch can stay on the
+    VoxCPM2 graph path. This is a model-local serving policy, not a generic AR
+    scheduler rule.
+    """
+
+    def _unified_decode_graph_enabled(self) -> bool:
+        runtime_config = _VoxCPM2RuntimeConfig.from_vllm_config(self.vllm_config)
+        return runtime_config.unified_decode_graph_available(use_cuda_graph=current_omni_platform.is_cuda())
+
+    def _should_defer_waiting_for_unified_decode_graph(self) -> bool:
+        if not self._unified_decode_graph_enabled():
+            return False
+        if not self.waiting or not self.running:
+            return False
+
+        for request in self.running:
+            if getattr(request, "status", None) != RequestStatus.RUNNING or request.is_finished():
+                continue
+            if self._get_confirmed_num_computed_tokens(request) >= request.num_prompt_tokens:
+                return True
+        return False
+
+    def _should_defer_waiting_admission(self) -> bool:
+        return self._should_defer_waiting_for_unified_decode_graph()

vllm_omni/model_executor/models/voxtral_tts/voxtral_tts.py

@@ -209,7 +209,12 @@ class VoxtralTTSForConditionalGeneration(
 
     def tts_preprocess(self, input_ids: torch.Tensor, input_embeds: torch.Tensor, **info_dict: dict | None):
         self.post_process_idx = 0
-        audio_tokens = info_dict.pop("audio", None)
+        codes = info_dict.get("codes")
+        audio_tokens = codes.get("audio") if isinstance(codes, Mapping) else None
+        if audio_tokens is None:
+            # Backward compatible with request state written before #4527 moved
+            # Voxtral TTS feedback frames under ``codes.audio``.
+            audio_tokens = info_dict.pop("audio", None)
         if audio_tokens is not None:
             kwargs = {"audio_tokens": audio_tokens.to(input_ids.device)}
             multimodal_embeddings = self.model.embed_multimodal(**kwargs)