Merge main into dev/vllm-align for vLLM 0.31 release

Signed-off-by: tzhouam <tzhouam@connect.ust.hk>

vllm_omni/model_executor/models/breeze_tts_2/depth_decoder.py

@@ -16,6 +16,14 @@ from vllm.platforms import current_platform
 from vllm_omni.platforms import current_omni_platform
 
 
+def _depth_layer_compile_options() -> dict[str, bool]:
+    """Keep ROCm away from the unstable Triton GEMM autotune path."""
+    return {
+        "epilogue_fusion": False,
+        "max_autotune": not current_platform.is_rocm(),
+    }
+
+
 def sample_logits(
     logits: torch.Tensor,
     temperature: float,
@@ -283,7 +291,7 @@ class BreezeDepthDecoder(nn.Module):
                     BreezeDepthLayer.forward,
                     fullgraph=True,
                     dynamic=True,
-                    options={"epilogue_fusion": False, "max_autotune": True},
+                    options=_depth_layer_compile_options(),
                 )
                 self._compiled_sampler = torch.compile(sample_graph_logits, fullgraph=True, dynamic=False)
             entry = BreezeDepthGraph(

vllm_omni/model_executor/models/common/fused_sampling.py

@@ -0,0 +1,97 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Fused Gumbel selection after native nucleus probability reductions.
+
+Keep ATen's top-k threshold, sort, softmax and cumulative sum, including
+boundary ties. Random numbers remain indexed by original token id. Only the
+final mask, Gumbel transform and argmax are fused; no vocabulary scatter is
+needed. This avoids changing which tokens survive a nucleus boundary through
+different reduction orders.
+"""
+
+import torch
+from vllm.triton_utils import tl, triton
+
+
+@triton.jit
+def _sorted_nucleus_gumbel(
+    logits,
+    indices,
+    probabilities,
+    cumulative,
+    uniforms,
+    output,
+    vocab: tl.constexpr,
+    logits_stride: tl.constexpr,
+    token_stride: tl.constexpr,
+    uniform_stride: tl.constexpr,
+    top_p: tl.constexpr,
+    block: tl.constexpr,
+):
+    row = tl.program_id(0)
+    offsets = tl.arange(0, block)
+    valid = offsets < vocab
+    values = tl.load(logits + row * logits_stride + offsets * token_stride, valid, other=-float("inf")).to(tl.float32)
+    if top_p < 1.0:
+        ids = tl.load(indices + row * vocab + offsets, valid, other=0)
+        probs = tl.load(probabilities + row * vocab + offsets, valid, other=0)
+        cdf = tl.load(cumulative + row * vocab + offsets, valid, other=0)
+        values = tl.where((cdf - probs) >= top_p, -float("inf"), values)
+    else:
+        ids = offsets
+    noise = tl.load(uniforms + row * uniform_stride + ids, valid, other=0.5)
+    score = values - tl.extra.cuda.libdevice.log(-tl.extra.cuda.libdevice.log(noise))
+    # torch.argmax returns the first original token for a tie or NaN.
+    first_nan = tl.min(tl.where(valid & (score != score), ids, 2147483647), 0)
+    maximum = tl.max(tl.where(valid, score, -float("inf")), 0)
+    winner = tl.min(tl.where(valid & (score == maximum), ids, 2147483647), 0)
+    tl.store(output + row, tl.where(first_nan < 2147483647, first_nan, winner))
+
+
+def sample_top_k_top_p_gumbel(
+    logits: torch.Tensor, uniforms: torch.Tensor, *, top_k: int, top_p: float
+) -> torch.Tensor:
+    """Sample CUDA rows without scattering sorted logits back to the vocabulary.
+
+    Uniforms are caller-owned FP32 values with the same shape as logits; their
+    generation and generator advancement are deliberately outside this kernel.
+    Softmax/scan use the reference ATen operations. The Gumbel logarithms
+    still use libdevice, so this is not a universal bitwise-equivalence claim.
+    """
+    if logits.ndim != 2 or not logits.is_cuda or uniforms.device != logits.device or torch.version.hip is not None:
+        raise ValueError("fused sampling requires two-dimensional CUDA logits and uniforms")
+    if uniforms.shape != logits.shape or uniforms.dtype != torch.float32 or uniforms.stride(1) != 1:
+        raise ValueError("uniforms must be FP32 with the logits shape and contiguous rows")
+    rows, vocab = logits.shape
+    # The stored-mode caller disables top-k for both 0 and negative values
+    # (vLLM commonly uses -1). Keep that behavior when fusion is enabled.
+    top_k = max(top_k, 0)
+    if vocab == 0 or top_k > vocab or not 0 < top_p <= 1:
+        raise ValueError("invalid top-k/top-p sampling parameters")
+    if top_k:
+        threshold = logits.topk(top_k, dim=-1).values[:, -1:]
+        logits = logits.masked_fill(logits < threshold, -float("inf"))
+    indices = probabilities = cumulative = None
+    if top_p < 1:
+        logits, indices = logits.sort(dim=-1, descending=True)
+        probabilities = logits.softmax(dim=-1, dtype=torch.float32)
+        cumulative = probabilities.cumsum(-1)
+    result = torch.empty((rows, 1), dtype=torch.long, device=logits.device)
+    if rows:
+        _sorted_nucleus_gumbel[(rows,)](
+            logits,
+            indices,
+            probabilities,
+            cumulative,
+            uniforms,
+            result,
+            vocab,
+            logits.stride(0),
+            logits.stride(1),
+            uniforms.stride(0),
+            top_p,
+            triton.next_power_of_2(vocab),
+            num_warps=4,
+            enable_fp_fusion=False,
+        )
+    return result

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

@@ -219,7 +219,8 @@ class ISTFTHead(FourierHead):
         # phase = torch.atan2(y, x)
         # S = mag * torch.exp(phase * 1j)
         # better directly produce the complex value
-        S = mag * (x + 1j * y)
+        # float32 operands: bf16 promotes to bcomplex32, which torch.fft cannot handle
+        S = mag.float() * (x.float() + 1j * y.float())
         audio, audio_buffer, window_buffer = self.istft(
             S, audio_buffer=audio_buffer, window_buffer=window_buffer, streaming=streaming, last_chunk=last_chunk
         )

vllm_omni/model_executor/models/common/qwen3_code_predictor.py

@@ -18,6 +18,7 @@ Shared by Qwen3-Omni and Qwen3-TTS talker models.
 from __future__ import annotations
 
 import dataclasses
+import time
 from collections.abc import Iterable, Sequence
 
 import torch
@@ -158,8 +159,8 @@ class CodePredictorAttention(nn.Module):
     """Multi-head self-attention for code predictor.
 
     Uses ``F.scaled_dot_product_attention`` with HF-compatible RoPE and RMSNorm.
-    No KV cache -- the code predictor always re-prefills the full (short)
-    sequence each AR step.
+    The default path re-prefills the short sequence; the opt-in cached path
+    appends only new positions.
 
     Input : [B, seq_len, hidden_size]
     Output: [B, seq_len, hidden_size]
@@ -316,6 +317,53 @@ class CodePredictorAttention(nn.Module):
         attn_out = attn_out.transpose(1, 2).reshape(bsz, seq_len, -1)
         return self.o_proj(attn_out)
 
+    def forward_cached(
+        self,
+        hidden_states: torch.Tensor,
+        position_embeddings: tuple[torch.Tensor, torch.Tensor],
+        k_cache: torch.Tensor,
+        v_cache: torch.Tensor,
+        write_positions: torch.Tensor,
+        key_mask: torch.Tensor,
+    ) -> torch.Tensor:
+        """Attend only the new positions, reusing earlier positions' keys and values.
+
+        ``k_cache``/``v_cache`` are ``[B, num_kv_heads, max_seq, head_dim]``;
+        the new positions' keys and values are written at ``write_positions``
+        (``[n]``). ``key_mask`` (``[1, 1, n, max_seq]``, bool) keeps key
+        positions up to each query's, so slots not yet written in this call are
+        never read and every single-position step has the same shapes.
+        """
+        bsz, seq_len, _ = hidden_states.shape
+        qkv = self.qkv_proj(hidden_states)
+        q_raw, k_raw, v_raw = self._split_qkv(qkv)
+        q = self.q_norm(q_raw.view(bsz, seq_len, self.num_heads, self.head_dim)).transpose(1, 2)
+        k = self.k_norm(k_raw.view(bsz, seq_len, self.num_kv_heads, self.head_dim)).transpose(1, 2)
+        v = v_raw.view(bsz, seq_len, self.num_kv_heads, self.head_dim).transpose(1, 2)
+
+        cos, sin = position_embeddings
+        cos = cos.unsqueeze(1)
+        sin = sin.unsqueeze(1)
+        q = (q * cos) + (_rotate_half(q) * sin)
+        k = (k * cos) + (_rotate_half(k) * sin)
+        k_cache.index_copy_(2, write_positions, k.to(k_cache.dtype))
+        v_cache.index_copy_(2, write_positions, v.to(v_cache.dtype))
+        if getattr(self, "use_short_kv_attention", False):
+            from .short_kv_attention import short_kv_attention
+
+            attn_out = short_kv_attention(q, k_cache, v_cache, write_positions, self.scaling)
+        else:
+            attn_out = F.scaled_dot_product_attention(
+                q,
+                k_cache,
+                v_cache,
+                attn_mask=key_mask,
+                scale=self.scaling,
+                enable_gqa=self.is_gqa,
+            )
+        attn_out = attn_out.transpose(1, 2).reshape(bsz, seq_len, -1)
+        return self.o_proj(attn_out)
+
 
 # ===================================================================
 #  MLP
@@ -351,7 +399,7 @@ class CodePredictorMLP(nn.Module):
 
 
 class CodePredictorDecoderLayer(nn.Module):
-    """Transformer decoder layer (SDPA, no KV cache)."""
+    """Transformer decoder layer with default re-prefill and opt-in KV reuse."""
 
     def __init__(self, config, *, prefix: str = "") -> None:
         super().__init__()
@@ -383,9 +431,29 @@ class CodePredictorDecoderLayer(nn.Module):
         hidden_states = residual + hidden_states
         return hidden_states
 
+    def forward_cached(
+        self,
+        hidden_states: torch.Tensor,
+        position_embeddings: tuple[torch.Tensor, torch.Tensor],
+        k_cache: torch.Tensor,
+        v_cache: torch.Tensor,
+        write_positions: torch.Tensor,
+        key_mask: torch.Tensor,
+    ) -> torch.Tensor:
+        residual = hidden_states
+        hidden_states = self.input_layernorm(hidden_states)
+        hidden_states = self.self_attn.forward_cached(
+            hidden_states, position_embeddings, k_cache, v_cache, write_positions, key_mask
+        )
+        hidden_states = residual + hidden_states
+        residual = hidden_states
+        hidden_states = self.post_attention_layernorm(hidden_states)
+        hidden_states = self.mlp(hidden_states)
+        return residual + hidden_states
+
 
 # ===================================================================
-#  Base Transformer Model (re-prefill, no KV cache)
+#  Base Transformer Model (re-prefill, or incremental with a K/V buffer)
 # ===================================================================
 
 
@@ -454,6 +522,29 @@ class CodePredictorBaseModel(nn.Module):
             hidden_states = self.norm(hidden_states)
         return hidden_states.to(input_dtype)
 
+    def forward_cached(
+        self,
+        inputs_embeds: torch.Tensor,
+        position_ids: torch.Tensor,
+        kv_caches: Sequence[tuple[torch.Tensor, torch.Tensor]],
+        write_positions: torch.Tensor,
+        key_mask: torch.Tensor,
+    ) -> torch.Tensor:
+        """``forward`` for the new positions only (``position_ids`` ``[B, n]``), reusing earlier K/V."""
+        input_dtype = inputs_embeds.dtype
+        use_fp32 = input_dtype == torch.float16 and inputs_embeds.device.type not in ("cpu", "npu")
+        if use_fp32:
+            inputs_embeds = inputs_embeds.float()
+        hidden_states = inputs_embeds
+        with torch.amp.autocast(inputs_embeds.device.type, enabled=use_fp32, dtype=torch.float32):
+            position_embeddings = self.rotary_emb(hidden_states, position_ids)
+            for layer, (k_cache, v_cache) in zip(self.layers, kv_caches):
+                hidden_states = layer.forward_cached(
+                    hidden_states, position_embeddings, k_cache, v_cache, write_positions, key_mask
+                )
+            hidden_states = self.norm(hidden_states)
+        return hidden_states.to(input_dtype)
+
     # Mapping from HF checkpoint shard name to the (fused parameter,
     # slice index) pair used to assemble the fused tensor.  See
     # ``load_weights`` below.
@@ -626,8 +717,14 @@ class CodePredictorWrapper(nn.Module):
       4. torch.compile on inner transformer.
       5. Inline sampling (top-k + top-p) -- no custom op overhead.
       6. Optional manual CUDA graph capture per batch-size bucket.
+
+    ``code_predictor_kv_cache`` (CUDA, opt-in) switches to incremental K/V
+    decoding: at large batches the re-prefill FLOPs are not negligible.
     """
 
+    _kv_cache_enabled: bool = False
+    _fused_sampling_enabled: bool = False
+
     def __init__(
         self,
         *,
@@ -672,6 +769,7 @@ class CodePredictorWrapper(nn.Module):
         # Sampling defaults for "stored" mode
         self._top_k: int = 50
         self._top_p: float = 0.8
+        self._do_sample: bool = True
 
         # Lazily initialised state
         self._proj_buf: torch.Tensor | None = None
@@ -686,6 +784,7 @@ class CodePredictorWrapper(nn.Module):
         # ``None`` keeps the legacy power-of-two bucket derivation.
         self._execution_batch_buckets: list[int] | None = None
         prefix_graph_cfg = self._stage_connector_extra_config(vllm_config)
+        self._fused_sampling_enabled = self._parse_bool_config(prefix_graph_cfg.get("code_predictor_fused_sampling"))
         prefix_graphs_requested = self._parse_bool_config(prefix_graph_cfg.get("code_predictor_prefix_graphs"))
         is_npu = current_omni_platform.is_npu()
         # MRv2 captures the whole Talker MTP call. It can therefore capture
@@ -705,15 +804,28 @@ class CodePredictorWrapper(nn.Module):
             prefix_graph_cfg.get("code_predictor_prefix_graph_seq_lens")
         )
         self._prefix_reprefill_seq_lens = tuple(self._prefix_seq_lens(self._num_groups + 1))
+        # Incremental decoding (opt-in, CUDA): each AR step runs only its new
+        # position against a static K/V buffer instead of re-prefilling the
+        # whole prefix (at batch 64 the re-prefill runs 17x the positions).
+        self._kv_cache_enabled = (
+            self._parse_bool_config(prefix_graph_cfg.get("code_predictor_kv_cache")) and current_omni_platform.is_cuda()
+        )
+        self._short_kv_requested = self._parse_bool_config(prefix_graph_cfg.get("code_predictor_short_kv"))
+        self._short_kv_enabled = False
+        self._kv_buf: torch.Tensor | None = None
+        self._compiled_kv_fwd = None
+        # Per AR step: (new positions [n], key mask [1, 1, n, max_seq]).
+        self._kv_steps: dict[int, tuple[torch.Tensor, torch.Tensor]] = {}
 
     def get_input_embeddings(self) -> nn.ModuleList:
         return self.model.get_input_embeddings()
 
-    def set_sampling_params(self, top_k: int = 50, top_p: float = 0.8) -> None:
-        """Configure sampling parameters to maintain consistency with previous implementation."""
+    def set_sampling_params(self, top_k: int = 50, top_p: float = 0.8, do_sample: bool = True) -> None:
+        """Configure "stored"-mode sampling; ``do_sample=False`` selects greedy argmax."""
         self._top_k = top_k
         self._top_p = top_p
-        logger.debug("Sampling parameters updated: top_k=%d, top_p=%.2f", top_k, top_p)
+        self._do_sample = do_sample
+        logger.debug("Sampling parameters updated: top_k=%d, top_p=%.2f, do_sample=%s", top_k, top_p, do_sample)
 
     # ------------------------------------------------------------------
     #  Lazy-init helpers
@@ -730,6 +842,35 @@ class CodePredictorWrapper(nn.Module):
         ):
             return
         self._proj_buf = torch.zeros(bsz, max_seq, self._cp_hidden, dtype=dtype, device=device)
+        if self._kv_cache_enabled:
+            cfg = self.config
+            head_dim = int(getattr(cfg, "head_dim", None) or cfg.hidden_size // cfg.num_attention_heads)
+            self._kv_buf = torch.zeros(
+                int(cfg.num_hidden_layers),
+                2,
+                bsz,
+                int(cfg.num_key_value_heads),
+                max_seq,
+                head_dim,
+                dtype=dtype,
+                device=device,
+            )
+
+    def _kv_views(self, bsz: int) -> list[tuple[torch.Tensor, torch.Tensor]]:
+        kv_buf = self._kv_buf
+        assert kv_buf is not None
+        return [(kv_buf[i, 0, :bsz], kv_buf[i, 1, :bsz]) for i in range(kv_buf.shape[0])]
+
+    def _ensure_kv_steps(self, device: torch.device) -> None:
+        if self._kv_steps:
+            return
+        max_seq = self._num_groups + 1
+        keys = torch.arange(max_seq, device=device)
+        for step in range(1, self._num_groups):
+            # Step 1 runs positions 0 (talker hidden) and 1 (CB0); step s > 1 runs position s.
+            positions = torch.arange(0, 2, device=device) if step == 1 else torch.tensor([step], device=device)
+            key_mask = (keys[None, :] <= positions[:, None]).view(1, 1, positions.numel(), max_seq)
+            self._kv_steps[step] = (positions, key_mask)
 
     def _setup_compile(self) -> None:
         """Lazily set up torch.compile with optional device graph capture."""
@@ -742,6 +883,9 @@ class CodePredictorWrapper(nn.Module):
         self._model_dtype = next(self.model.parameters()).dtype
         self._lm_heads_list = list(self.lm_head)
         self._codec_embeds_list = list(self.model.codec_embedding)
+        if self._model_dtype == torch.float16 or not current_omni_platform.supports_torch_inductor():
+            # The fp16 body runs in fp32 under autocast; keep re-prefill there.
+            self._kv_cache_enabled = False
 
         # Torch 2.13 XPU Dynamo can double-register built-in handlers when
         # spawned workers compile this predictor. Keep this narrow path eager
@@ -768,10 +912,28 @@ class CodePredictorWrapper(nn.Module):
             dynamic=False,
             options={"epilogue_fusion": False},
         )
+        self._short_kv_enabled = (
+            self._kv_cache_enabled
+            and self._short_kv_requested
+            and current_omni_platform.is_cuda()
+            and self._model_dtype == torch.bfloat16
+        )
+        for layer in self.model.layers:
+            layer.self_attn.use_short_kv_attention = self._short_kv_enabled
+        if self._kv_cache_enabled:
+            self._compiled_kv_fwd = torch.compile(
+                self.model.forward_cached,
+                dynamic=self._short_kv_enabled,
+                options={"epilogue_fusion": False},
+            )
         with torch._dynamo.config.patch(cache_size_limit=self._compile_cache_size_limit()):
             self._warmup_buckets()
 
-        if self._wrapper_config.use_cuda_graphs:
+        if self._kv_cache_enabled:
+            # Fifteen short steps per call: they rely on the runner capturing
+            # the whole MTP call (MRv2) rather than on per-step inner graphs.
+            logger.info("code_predictor: torch.compile + incremental K/V decoding")
+        elif self._wrapper_config.use_cuda_graphs:
             self._capture_cuda_graphs()
             logger.info("code_predictor: torch.compile (no epilogue fusion) + CUDA graphs")
         else:
@@ -792,7 +954,13 @@ class CodePredictorWrapper(nn.Module):
             extra_cfg = connector_cfg.get("extra", connector_cfg)
         else:
             extra_cfg = getattr(connector_cfg, "extra", None)
-        return extra_cfg if isinstance(extra_cfg, dict) else {}
+        extra_cfg = extra_cfg if isinstance(extra_cfg, dict) else {}
+        # A single-stage deployment has no connector edge to carry model
+        # options; it sets them in the stage's ``additional_config``.
+        additional = getattr(vllm_config, "additional_config", None)
+        if isinstance(additional, dict) and additional:
+            return {**additional, **extra_cfg}
+        return extra_cfg
 
     @staticmethod
     def _parse_bool_config(value: object) -> bool:
@@ -840,13 +1008,13 @@ class CodePredictorWrapper(nn.Module):
         return row_generators
 
     @classmethod
-    def _sample_codes_gumbel(
+    def _sampling_uniforms(
         cls,
         logits: torch.Tensor,
         generator: _GeneratorLike = None,
         uniforms: torch.Tensor | None = None,
     ) -> torch.Tensor:
-        """Sample ``logits`` via Gumbel-max with optional precomputed noise."""
+        """Draw one FP32 uniform per original token, preserving generator state."""
         if uniforms is None:
             row_generators = cls._normalize_generators(generator, int(logits.shape[0]))
             u = torch.empty_like(logits, dtype=torch.float32)
@@ -862,6 +1030,17 @@ class CodePredictorWrapper(nn.Module):
                     f"logits={tuple(logits.shape)}"
                 )
             u = uniforms
+        return u
+
+    @classmethod
+    def _sample_codes_gumbel(
+        cls,
+        logits: torch.Tensor,
+        generator: _GeneratorLike = None,
+        uniforms: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        """Sample ``logits`` via Gumbel-max with optional precomputed noise."""
+        u = cls._sampling_uniforms(logits, generator, uniforms)
         return (logits.float() - torch.log(-torch.log(u))).argmax(dim=-1, keepdim=True)
 
     def _prefix_seq_lens(self, max_seq: int) -> list[int]:
@@ -924,6 +1103,8 @@ class CodePredictorWrapper(nn.Module):
             if needs_full_graph:
                 full_graph_entries += len(prefix_buckets)
             required_entries = full_graph_entries + len(prefix_buckets) * len(prefix_seq_lens)
+        if self._kv_cache_enabled:
+            required_entries = 4  # {first step, later steps} x {batch 1, dynamic batch}
         return max(torch._dynamo.config.cache_size_limit, required_entries)
 
     def _warmup_buckets(self) -> None:
@@ -938,7 +1119,29 @@ class CodePredictorWrapper(nn.Module):
         self._ensure_buffers(device, self._model_dtype, max(self._bucket_sizes))
         proj_buf = self._proj_buf
 
-        if self._prefix_reprefill_enabled:
+        if self._kv_cache_enabled:
+            self._ensure_kv_steps(device)
+            started = time.perf_counter()
+            for bsz in self._bucket_sizes:
+                caches = self._kv_views(bsz)
+                for step in (1, 2):
+                    positions, key_mask = self._kv_steps[step]
+                    lo = 0 if step == 1 else step
+                    hidden = proj_buf[:bsz, lo : step + 1, :]
+                    pos_ids = positions.unsqueeze(0).expand(bsz, -1)
+                    if bsz > 1:
+                        # One graph for every batch >= 2 per step shape (4 compiles, not 2 per bucket):
+                        # a cold per-bucket compile outlasted the stage handshake timeout.
+                        for tensor in (hidden, pos_ids, *(t for kv in caches for t in kv)):
+                            torch._dynamo.mark_dynamic(tensor, 0)
+                    for _ in range(2):
+                        self._compiled_kv_fwd(hidden, pos_ids, caches, positions, key_mask)
+            logger.info(
+                "code_predictor: incremental K/V decoding warmed for buckets %s in %.1f s",
+                self._bucket_sizes,
+                time.perf_counter() - started,
+            )
+        elif self._prefix_reprefill_enabled:
             prefix_seq_lens = self._prefix_reprefill_seq_lens
             needs_full_graph = set(prefix_seq_lens) != set(range(2, max_seq))
             for bsz in self._bucket_sizes:
@@ -1135,6 +1338,21 @@ class CodePredictorWrapper(nn.Module):
             scaled = scaled.masked_fill(scaled < topk_vals[:, -1:], float("-inf"))
         return self._sample_codes_gumbel(scaled, generator=generator, uniforms=uniforms)
 
+    def _validate_sampling_inputs(
+        self,
+        bsz: int,
+        generators: Sequence[torch.Generator | None] | None,
+        sample_uniforms: torch.Tensor | None,
+    ) -> None:
+        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}")
+        if sample_uniforms is not None:
+            expected_shape = (bsz, self._num_groups - 1, int(self.config.vocab_size))
+            if tuple(sample_uniforms.shape) != expected_shape:
+                raise ValueError(
+                    f"sample_uniforms must have shape {expected_shape}, got {tuple(sample_uniforms.shape)}"
+                )
+
     @torch.inference_mode()
     def forward(
         self,
@@ -1152,14 +1370,7 @@ class CodePredictorWrapper(nn.Module):
         """Predict residual codebooks 1..G-1 autoregressively via re-prefill."""
         bsz = int(layer0_code.shape[0])
         num_groups = self._num_groups
-        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}")
-        if sample_uniforms is not None:
-            expected_shape = (bsz, num_groups - 1, int(self.config.vocab_size))
-            if tuple(sample_uniforms.shape) != expected_shape:
-                raise ValueError(
-                    f"sample_uniforms must have shape {expected_shape}, got {tuple(sample_uniforms.shape)}"
-                )
+        self._validate_sampling_inputs(bsz, generators, sample_uniforms)
         sample_generator: _GeneratorLike = generators if generators is not None else generator
         device = layer0_code.device
 
@@ -1206,9 +1417,29 @@ class CodePredictorWrapper(nn.Module):
             all_codes = torch.empty(bsz, num_groups, dtype=torch.long, device=device)
             all_codes[:, 0] = layer0_code.reshape(bsz)
 
+        kv_caches = None
+        if self._kv_cache_enabled:
+            # Masked slots still enter the attention matmul with zero weight, so
+            # clear them: a non-finite value left by an earlier call would leak.
+            if not self._short_kv_enabled:
+                self._kv_buf[:, :, :padded_bsz].zero_()
+            kv_caches = self._kv_views(padded_bsz)
         # Autoregressive loop: predict layers 1..G-1
         for step in range(1, num_groups):
-            logits = self._predict_step_logits(proj_buf, bsz, padded_bsz, step, is_npu_capturing)
+            if kv_caches is not None:
+                # Step 1 runs positions 0-1; later steps append one position.
+                positions, key_mask = self._kv_steps[step]
+                lo = 0 if step == 1 else step
+                new_hidden = self._compiled_kv_fwd(
+                    proj_buf[:padded_bsz, lo : step + 1, :],
+                    positions.unsqueeze(0).expand(padded_bsz, -1),
+                    kv_caches,
+                    positions,
+                    key_mask,
+                )
+                logits = self._lm_heads_list[step - 1](new_hidden[:bsz, -1, :])
+            else:
+                logits = self._predict_step_logits(proj_buf, bsz, padded_bsz, step, is_npu_capturing)
 
             # Sample next code via Gumbel-max.
             #
@@ -1222,7 +1453,15 @@ class CodePredictorWrapper(nn.Module):
             # 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:
+            if stored_mode and not self._do_sample:
+                code = logits.argmax(dim=-1, keepdim=True)
+            elif stored_mode and self._fused_sampling_enabled and logits.is_cuda and torch.version.hip is None:
+                from .fused_sampling import sample_top_k_top_p_gumbel
+
+                step_uniforms = sample_uniforms[:, step - 1, :] if sample_uniforms is not None else None
+                uniforms = self._sampling_uniforms(logits, sample_generator, step_uniforms)
+                code = sample_top_k_top_p_gumbel(logits, uniforms, top_k=s_top_k, top_p=s_top_p)
+            elif stored_mode:
                 # "stored" mode: top-k -> top-p -> Gumbel-max
                 if s_top_k > 0:
                     topk_vals, _ = logits.topk(s_top_k, dim=-1)

vllm_omni/model_executor/models/common/talker_first_audio.py

@@ -0,0 +1,29 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Shared eligibility; each model separately controls its first-audio opt-in."""
+
+from typing import Any
+
+import torch
+
+from vllm_omni.platforms import current_omni_platform
+
+
+def supports_talker_first_audio(vllm_config: Any) -> bool:
+    """Require CUDA MRv2 streaming, local single-rank execution and no prefix replay.
+
+    Other configurations retain the regular codec path. These conditions do
+    not enable the feature: model-specific options must also accept it.
+    """
+    model = vllm_config.model_config
+    if not (getattr(model, "use_v2_model_runner", False) and getattr(model, "async_chunk", False)):
+        return False
+    parallel = vllm_config.parallel_config
+    return (
+        current_omni_platform.is_cuda()
+        and torch.device(vllm_config.device_config.device).type == "cuda"
+        and parallel.tensor_parallel_size == 1
+        and parallel.pipeline_parallel_size == 1
+        and parallel.distributed_executor_backend in (None, "uni")
+        and not vllm_config.cache_config.enable_prefix_caching
+    )

vllm_omni/model_executor/models/cosyvoice3/code2wav_core/hifigan.py

@@ -966,7 +966,20 @@ class CausalHiFTGenerator(HiFTGenerator):
         self.conv_pre_look_right = conv_pre_look_right
         self.f0_predictor = f0_predictor
 
+    def enable_decode_graphs(self) -> None:
+        """Opt into exact-shape CUDA replay after frozen weights are loaded."""
+        from .hift_graph import HiFTDecodeGraphs
+
+        self.enable_cached_istft()
+        self._decode_graphs = HiFTDecodeGraphs(self._decode_eager)
+
     def decode(self, x: torch.Tensor, s: torch.Tensor = torch.zeros(1, 1, 0), finalize: bool = True) -> torch.Tensor:
+        runner = getattr(self, "_decode_graphs", None)
+        if runner is not None:
+            return runner.run(x, s, finalize)
+        return self._decode_eager(x, s, finalize=finalize)
+
+    def _decode_eager(self, x: torch.Tensor, s: torch.Tensor, finalize: bool = True) -> torch.Tensor:
         s_stft_real, s_stft_imag = self._stft(s.squeeze(1))
         if finalize is True:
             x = self.conv_pre(x)

vllm_omni/model_executor/models/cosyvoice3/code2wav_core/hift_graph.py

@@ -0,0 +1,85 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Bounded CUDA graph replay for HiFT's deterministic waveform decoder.
+
+F0, sine/noise generation, phase carry, and per-request streaming state remain
+outside the graph. Each replay consumes current mel and source tensors and
+returns an owned waveform, because multiple requests can reuse one shape.
+"""
+
+from collections import OrderedDict
+from collections.abc import Callable
+from dataclasses import dataclass
+
+import torch
+
+
+@dataclass
+class _DecodeGraph:
+    graph: torch.cuda.CUDAGraph
+    mel: torch.Tensor
+    source: torch.Tensor
+    output: torch.Tensor
+
+
+class HiFTDecodeGraphs:
+    """Capture recurrent exact shapes, with at most 32 resident graphs.
+
+    ``decode`` must have warmed its ISTFT envelope before capture and remain
+    deterministic for its mel, source and finalize arguments. Uncommon shapes,
+    CPU execution and nested CUDA capture use that same eager implementation.
+    """
+
+    def __init__(self, decode: Callable[..., torch.Tensor], max_graphs: int = 32) -> None:
+        self.decode = decode
+        self.max_graphs = max_graphs
+        self.graphs: dict[tuple, _DecodeGraph] = {}
+        self.seen: OrderedDict[tuple, int] = OrderedDict()
+        self.pool = None
+
+    @torch.inference_mode()
+    def run(self, mel: torch.Tensor, source: torch.Tensor, finalize: bool) -> torch.Tensor:
+        """Replay a recurring layout without aliasing another request's output."""
+        if not mel.is_cuda or not source.is_cuda or torch.cuda.is_current_stream_capturing():
+            return self.decode(mel, source, finalize=finalize)
+        autocast = torch.is_autocast_enabled("cuda")
+        key = (
+            mel.shape,
+            source.shape,
+            mel.device,
+            mel.dtype,
+            source.dtype,
+            finalize,
+            torch.backends.cudnn.allow_tf32,
+            torch.backends.cuda.matmul.allow_tf32,
+            autocast,
+            torch.get_autocast_dtype("cuda") if autocast else None,
+        )
+        state = self.graphs.get(key)
+        if state is None:
+            if len(self.graphs) >= self.max_graphs:
+                return self.decode(mel, source, finalize=finalize)
+            count = self.seen.pop(key, 0) + 1
+            self.seen[key] = count
+            if len(self.seen) > 128:
+                self.seen.popitem(last=False)
+            if count < 3:
+                return self.decode(mel, source, finalize=finalize)
+            static_mel, static_source = mel.clone(), source.clone()
+            stream = torch.cuda.Stream(device=mel.device)
+            stream.wait_stream(torch.cuda.current_stream(mel.device))
+            with torch.cuda.stream(stream):
+                self.decode(static_mel, static_source, finalize=finalize)
+            torch.cuda.current_stream(mel.device).wait_stream(stream)
+            graph = torch.cuda.CUDAGraph()
+            with torch.cuda.graph(graph, pool=self.pool):
+                output = self.decode(static_mel, static_source, finalize=finalize)
+            if self.pool is None:
+                self.pool = graph.pool()
+            state = _DecodeGraph(graph, static_mel, static_source, output)
+            self.graphs[key] = state
+        state.mel.copy_(mel)
+        state.source.copy_(source)
+        state.graph.replay()
+        # A later request may replay this shape before the caller consumes it.
+        return state.output.clone()

vllm_omni/model_executor/models/cosyvoice3/code2wav_core/packed_dit.py

@@ -170,7 +170,7 @@ class PackedDiT:
             t,
             rows,
             attention,
-            max_seqlen_q=attention.page_table.shape[1],
+            max_seqlen_q=attention.max_seqlen_q,
             block_modulation=block_modulation,
             final_modulation=final_modulation,
             rope=rope,
@@ -771,7 +771,9 @@ class FlowGraphRunner:
         attention.heads, attention.head_dim = self.heads, self.head_dim
         attention.cache_seqlens, attention.cu_seqlens_q = state.cache_seqlens, state.cu_seqlens_q
         attention.page_table = state.page_table
-        attention.max_seqlen_q = state.page_table.shape[1]
+        # Streaming queries cover one training chunk, even when their KV
+        # prefix and the packed graph allocation contain thousands of frames.
+        attention.max_seqlen_q = estimator.dit.static_chunk_size if key[0] else state.page_table.shape[1]
         mark_packed_compile_metadata(rows, attention)
         mu_cfg = torch.cat((state.mu, torch.zeros_like(state.mu)), dim=1)
         cond_cfg = torch.cat((state.cond, torch.zeros_like(state.cond)), dim=1)

vllm_omni/model_executor/models/cosyvoice3/cosyvoice3_code2wav.py

@@ -769,4 +769,11 @@ class CosyVoice3Code2Wav(nn.Module):
         if f0_folded == 0:
             logger.warning("HiFT F0 predictor had no weight-norm layers to fold; check config drift")
         self.hift.eval()
+        if device.type == "cuda" and cosyvoice3_packed_streaming_enabled():
+            if os.getenv("COSYVOICE3_CACHED_ISTFT", "0") == "1":
+                self.hift.enable_cached_istft()
+                logger.info("CosyVoice3 streaming HiFT: cached ISTFT enabled")
+            if os.getenv("COSYVOICE3_HIFT_GRAPH", "0") == "1":
+                self.hift.enable_decode_graphs()
+                logger.info("CosyVoice3 streaming HiFT: decode CUDA graphs enabled")
         logger.info(f"Loaded hift weights from {hift_path}")

vllm_omni/model_executor/models/moss_tts/audio_tokenizer_v2.py

@@ -539,11 +539,15 @@ class RingKVCache:
             end_offset = self.end_offset.index_select(0, slots)
             row_cache = self.cache.index_select(1, slots)
 
-            indexes = torch.arange(T, device=end_offset.device, dtype=end_offset.dtype)
+            # Upsampling can make a chunk longer than the ring. Only the last
+            # capacity tokens survive; writing all T tokens gives scatter_
+            # duplicate destinations and nondeterministic cache contents.
+            write_length = min(T, self.capacity)
+            indexes = torch.arange(T - write_length, T, device=end_offset.device, dtype=end_offset.dtype)
             indexes = (indexes + end_offset.view(-1, 1)) % self.capacity
-            scatter_indexes = indexes.view(B, 1, T, 1).expand(-1, H, T, D)
-            row_cache[0].scatter_(2, scatter_indexes, k)
-            row_cache[1].scatter_(2, scatter_indexes, v)
+            scatter_indexes = indexes.view(B, 1, write_length, 1).expand(-1, H, write_length, D)
+            row_cache[0].scatter_(2, scatter_indexes, k[:, :, -write_length:])
+            row_cache[1].scatter_(2, scatter_indexes, v[:, :, -write_length:])
             # Live and graph-padding rows always map to distinct slots. The
             # latter map only to scratch state, so this write cannot corrupt a
             # request even though dense graph operators still execute it.
@@ -729,7 +733,13 @@ class MossAudioTokenizerMultiheadAttention(StreamingModule):
         if self.context is None:
             capacity = self.weights_per_step if self.weights_per_step else 1024
         else:
-            capacity = self.context
+            # RingKVCache.complete writes a whole chunk before attending, so a
+            # ring of exactly ``context`` entries lets a T-token chunk evict up
+            # to T tokens of history that the chunk's own queries should still
+            # see (and, for T > context, evict part of the chunk itself).
+            # ``_ring_headroom`` (max chunk tokens at this layer) keeps every
+            # in-window key resident; 0 preserves the legacy truncation.
+            capacity = self.context + int(getattr(self, "_ring_headroom", 0))
 
         kv_cache = RingKVCache(
             batch_size,
@@ -791,29 +801,63 @@ class MossAudioTokenizerMultiheadAttention(StreamingModule):
         if self.rope:
             q, k = self.rope(q, k, offset, time_before_heads=False)
 
-        k, v, pos_k = self._complete_kv(k, v, execution_context)
-        pos_k = pos_k[:, None]
-
-        if self.causal:
-            pos_q = offset.view(-1, 1, 1) + torch.arange(T, device=q.device, dtype=torch.long).view(-1, 1)
-            delta = pos_q - pos_k
-            attn_bias = (pos_k >= 0) & (delta >= 0)
-            if self.context is not None:
-                attn_bias = attn_bias & (delta < self.context)
-            attn_bias = attn_bias[:, None]
+        slot_attention = getattr(self, "_slot_attention", None)
+        slot_attention_rows = getattr(self, "_slot_attention_rows", None)
+        advance_kv_offset = False
+        slot_path = state is not None and execution_context is not None and self.causal
+        if slot_path and slot_attention_rows is not None:
+            # One launch: the gathered MHA offset equals the ring's end_offset
+            # for every slot, and both advance together below.
+            valid_lengths = execution_context.valid_rows.to(dtype=torch.int32) * T
+            x = slot_attention_rows(
+                q,
+                k,
+                v,
+                state.kv_cache.cache,
+                offset,
+                execution_context.state_slot_ids,
+                valid_lengths,
+                self.context if self.context is not None else -1,
+            )
+            advance_kv_offset = True
+        elif slot_attention is not None and slot_path:
+            # Current streaming batches use one exact T. Padding rows own
+            # scratch slots and must not advance persistent attention state.
+            valid_lengths = execution_context.valid_rows.to(dtype=torch.int32) * T
+            x = slot_attention(
+                q,
+                k,
+                v,
+                state.kv_cache.cache,
+                state.kv_cache.end_offset,
+                execution_context.state_slot_ids,
+                valid_lengths,
+                self.context if self.context is not None else -1,
+            )
         else:
-            attn_bias = None
+            k, v, pos_k = self._complete_kv(k, v, execution_context)
+            pos_k = pos_k[:, None]
 
-        streaming_attention = getattr(self, "_streaming_attention", None)
-        if (
-            streaming_attention is not None
-            and attn_bias is not None
-            and q.dtype == torch.bfloat16
-            and q.shape[-1] == 64
-        ):
-            x = streaming_attention(q, k, v, attn_bias)
-        else:
-            x = F.scaled_dot_product_attention(q, k, v, attn_bias, dropout_p=0.0)
+            if self.causal:
+                pos_q = offset.view(-1, 1, 1) + torch.arange(T, device=q.device, dtype=torch.long).view(-1, 1)
+                delta = pos_q - pos_k
+                attn_bias = (pos_k >= 0) & (delta >= 0)
+                if self.context is not None:
+                    attn_bias = attn_bias & (delta < self.context)
+                attn_bias = attn_bias[:, None]
+            else:
+                attn_bias = None
+
+            streaming_attention = getattr(self, "_streaming_attention", None)
+            if (
+                streaming_attention is not None
+                and attn_bias is not None
+                and q.dtype == torch.bfloat16
+                and q.shape[-1] == 64
+            ):
+                x = streaming_attention(q, k, v, attn_bias)
+            else:
+                x = F.scaled_dot_product_attention(q, k, v, attn_bias, dropout_p=0.0)
         x = x.transpose(1, 2).reshape(B, T, self.embed_dim)
         x = apply_weights_per_step(self.out_projs, self.weights_per_step_schedule, x, offset_cpu)
 
@@ -824,6 +868,8 @@ class MossAudioTokenizerMultiheadAttention(StreamingModule):
             else:
                 next_offset = torch.where(execution_context.valid_rows, offset + T, offset)
                 state.offset.index_copy_(0, execution_context.state_slot_ids, next_offset)
+                if advance_kv_offset:
+                    state.kv_cache.end_offset.index_copy_(0, execution_context.state_slot_ids, next_offset)
         return x
 
 
@@ -1624,6 +1670,7 @@ class MossAudioTokenizerModel(MossAudioTokenizerPreTrainedModel):
         decoder_kwargs_list = copy.deepcopy(config.decoder_kwargs)
         self.decoder = nn.ModuleList()
 
+        decoder_input_frame_rate = current_frame_rate
         for decoder_kwargs_i in decoder_kwargs_list:
             decoder_kwargs_i = dict(decoder_kwargs_i)
             if decoder_kwargs_i["module_type"] == "PatchedPretransform":
@@ -1638,6 +1685,11 @@ class MossAudioTokenizerModel(MossAudioTokenizerPreTrainedModel):
                         context=int(round(current_frame_rate * context_duration)),
                     )
                 )
+                # Tokens this transformer sees per input code frame. Streaming
+                # ring headroom is sized from it (see initialize_decoder_state_pool).
+                self.decoder[-1].tokens_per_input_frame = max(
+                    1, int(round(current_frame_rate / decoder_input_frame_rate))
+                )
             current_frame_rate *= self.decoder[-1].downsample_ratio
 
         expected_output_frame_rate = float(self.sampling_rate * channel_interleave_factor)
@@ -1731,10 +1783,29 @@ class MossAudioTokenizerModel(MossAudioTokenizerPreTrainedModel):
         self._decoder_state_capacity = 0
         self._decoder_slot_offsets = None
 
-    def initialize_decoder_state_pool(self, state_capacity: int, scratch_capacity: int = 0) -> None:
-        """Allocate persistent decoder state independently of execution B."""
-        if state_capacity <= 0 or scratch_capacity < 0:
-            raise ValueError(f"Invalid decoder state capacities: state={state_capacity}, scratch={scratch_capacity}.")
+    def initialize_decoder_state_pool(
+        self, state_capacity: int, scratch_capacity: int = 0, chunk_frames: int = 0
+    ) -> None:
+        """Allocate persistent decoder state independently of execution B.
+
+        ``chunk_frames`` > 0 sizes every decoder attention ring as
+        ``context + chunk_frames * tokens_per_input_frame`` so that a chunk of
+        up to ``chunk_frames`` code frames never evicts keys inside its own
+        causal window; chunked streaming then matches whole-sequence decoding.
+        ``chunk_frames`` = 0 keeps the legacy ``context``-sized rings.
+        """
+        if state_capacity <= 0 or scratch_capacity < 0 or chunk_frames < 0:
+            raise ValueError(
+                "Invalid decoder state capacities: "
+                f"state={state_capacity}, scratch={scratch_capacity}, chunk_frames={chunk_frames}."
+            )
+        for module in self.decoder:
+            if not isinstance(module, MossAudioTokenizerProjectedTransformer):
+                continue
+            tokens_per_frame = int(getattr(module, "tokens_per_input_frame", 1))
+            for attention in module.modules():
+                if isinstance(attention, MossAudioTokenizerMultiheadAttention):
+                    attention._ring_headroom = chunk_frames * tokens_per_frame
         self._start_streaming(state_capacity + scratch_capacity, decoder_only=True)
         self._decoder_state_capacity = state_capacity + scratch_capacity
 

vllm_omni/model_executor/models/moss_tts/cuda_graph_streaming_decoder_wrapper.py

@@ -81,15 +81,31 @@ class CUDAGraphStreamingDecoderWrapper:
         frame_sizes: list[int],
         num_quantizers: int,
         vllm_config: VllmConfig,
+        scratch_base: int | None = None,
+        private_pool: bool = False,
+        compiled_decode: nn.Module | None = None,
     ) -> None:
         self.codec = codec
         self.state_capacity = int(state_capacity)
+        # Graph padding rows use scratch slots [scratch_base, scratch_base+B).
+        # A wrapper replaying concurrently with another needs its own scratch
+        # range and graph memory pool.
+        self.scratch_base = self.state_capacity if scratch_base is None else int(scratch_base)
+        self._private_pool = bool(private_pool)
+        # Graphs bake in the cuBLAS workspace of the stream they were captured
+        # on, and torch.cuda.graph shares one default capture stream. Split-K
+        # GEMMs synchronize through that workspace, so two graphs replaying
+        # concurrently from one capture stream can spin on each other forever.
+        self._capture_stream: torch.cuda.Stream | None = None
         self.batch_sizes = sorted({int(size) for size in batch_sizes if 0 < int(size) <= state_capacity})
         self.frame_sizes = sorted({int(size) for size in frame_sizes if int(size) > 0})
         self.num_quantizers = int(num_quantizers)
         self.graphs: dict[tuple[int, int], _CapturedStreamingDecodeGraph] = {}
         self._pool = None
         self._warmed_up = False
+        if compiled_decode is not None:
+            self._compiled_decode: nn.Module | None = compiled_decode
+            return
         # vLLM owns Inductor compilation; this wrapper remains the sole owner
         # of CUDA Graph capture/replay because it understands persistent codec
         # state slots. Disable vLLM's CUDA Graph layer to avoid nested graphs.
@@ -98,7 +114,7 @@ class CUDAGraphStreamingDecoderWrapper:
         compile_config.compilation_config.cudagraph_mode = CUDAGraphMode.NONE
         compile_config.compilation_config.static_forward_context = {}
         with set_current_vllm_config(compile_config):
-            self._compiled_decode: nn.Module | None = _MossStreamingDecodeCompileAdapter(
+            self._compiled_decode = _MossStreamingDecodeCompileAdapter(
                 codec,
                 vllm_config=compile_config,
             )
@@ -175,7 +191,7 @@ class CUDAGraphStreamingDecoderWrapper:
                     frame_size,
                     exc_info=True,
                 )
-                scratch_slots = self.state_capacity + torch.arange(
+                scratch_slots = self.scratch_base + torch.arange(
                     batch_size,
                     dtype=torch.long,
                     device=device,
@@ -206,10 +222,13 @@ class CUDAGraphStreamingDecoderWrapper:
             device=device,
         )
         lengths = torch.zeros(batch_size, dtype=torch.long, device=device)
-        scratch_slots = self.state_capacity + torch.arange(batch_size, dtype=torch.long, device=device)
+        scratch_slots = self.scratch_base + torch.arange(batch_size, dtype=torch.long, device=device)
         valid_rows = torch.zeros(batch_size, dtype=torch.bool, device=device)
 
-        stream = torch.cuda.Stream()
+        if self._private_pool and self._capture_stream is None:
+            self._capture_stream = torch.cuda.Stream(device=device)
+        # Warm up on the capture stream so its workspace exists before capture.
+        stream = self._capture_stream or torch.cuda.Stream()
         stream.wait_stream(torch.cuda.current_stream())
         with torch.cuda.stream(stream):
             for _ in range(3):
@@ -219,9 +238,11 @@ class CUDAGraphStreamingDecoderWrapper:
         self.codec.reset_decoder_state_slots(scratch_slots)
 
         if self._pool is None:
-            self._pool = current_platform.get_global_graph_pool()
+            self._pool = (
+                torch.cuda.graph_pool_handle() if self._private_pool else current_platform.get_global_graph_pool()
+            )
         graph = CUDAGraph()
-        with torch.cuda.graph(graph, pool=self._pool, capture_error_mode="thread_local"):
+        with torch.cuda.graph(graph, pool=self._pool, stream=self._capture_stream, capture_error_mode="thread_local"):
             audio, audio_lengths = decode(codes, lengths, scratch_slots, valid_rows)
 
         self.graphs[(batch_size, frame_size)] = _CapturedStreamingDecodeGraph(
@@ -272,7 +293,7 @@ class CUDAGraphStreamingDecoderWrapper:
         entry.static_lengths.zero_()
         entry.static_lengths[:actual_batch_size].fill_(int(frame_size))
         entry.static_state_slot_ids.copy_(
-            self.state_capacity + torch.arange(batch_size, dtype=torch.long, device=entry.static_state_slot_ids.device)
+            self.scratch_base + torch.arange(batch_size, dtype=torch.long, device=entry.static_state_slot_ids.device)
         )
         entry.static_state_slot_ids[:actual_batch_size].copy_(state_slot_ids, non_blocking=True)
         entry.static_valid_rows.zero_()

vllm_omni/model_executor/models/moss_tts/first_chunk_fast_path.py

@@ -0,0 +1,333 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Decode a stream's first codec chunk as soon as Stage 1 receives it.
+
+The Stage-1 engine loop only sees a new chunk between steps, and a step waits
+for its codec graph to finish. A first chunk that arrives mid-step therefore
+waits for the rest of that step, then for its own replay. This path takes the
+first chunk from the receive thread, replays a small dedicated graph on its
+own stream while the main step runs, and emits the audio straight to the
+engine output queue. The main path later receives the same chunk, finds it
+already decoded, and continues the stream from the state written here.
+"""
+
+from __future__ import annotations
+
+import math
+import queue
+import threading
+import time
+from dataclasses import dataclass
+from typing import Any
+
+import torch
+from vllm.logger import init_logger
+
+from vllm_omni.platforms import current_omni_platform
+from vllm_omni.worker_v2.first_audio_sender import FirstAudioSink, _PreparedDelivery
+
+logger = init_logger(__name__)
+
+
+class _SlotHandoff:
+    """Orders the main path's first use of a slot after the fast decode."""
+
+    __slots__ = ("done", "event", "error")
+
+    def __init__(self) -> None:
+        self.done = threading.Event()
+        self.event: torch.cuda.Event | None = None
+        self.error: Exception | None = None
+
+
+@dataclass
+class _Job:
+    request_id: str
+    request_key: str
+    delivery: _PreparedDelivery
+    slot: int
+    codes: torch.Tensor  # (n_vq, T) int64 on host
+    handoff: _SlotHandoff
+
+
+class MossFirstChunkFastPath:
+    def __init__(
+        self,
+        session: Any,
+        wrapper: Any,
+        *,
+        n_vq: int,
+        frames: int,
+        codebook_size: int,
+        samples_per_frame: int,
+        n_channels: int,
+        sample_rate: torch.Tensor,
+        device: torch.device,
+        gate_main: bool = False,
+        handoff_timeout_s: float = 30.0,
+    ) -> None:
+        if not math.isfinite(handoff_timeout_s) or handoff_timeout_s <= 0:
+            raise ValueError("handoff_timeout_s must be positive and finite")
+        self._handoff_timeout_s = float(handoff_timeout_s)
+        self._session = session
+        # When set, the main codec stream waits for an in-flight fast decode
+        # before its next replay, so the small first-chunk graph runs mostly
+        # uncontended instead of interleaving with a large regular replay.
+        self._gate_main = bool(gate_main)
+        self._inflight: torch.cuda.Event | None = None
+        self._wrapper = wrapper
+        self._n_vq = int(n_vq)
+        self.frames = int(frames)
+        self._codebook_size = int(codebook_size)
+        self._samples = int(frames) * int(samples_per_frame)
+        self._n_channels = int(n_channels)
+        self._sr = sample_rate
+        if device.index is None:
+            device = torch.device(device.type, torch.accelerator.current_device_index())
+        self._device = device
+        self._max_batch = max(wrapper.batch_sizes)
+        self._stream = torch.cuda.Stream(device=device, priority=-1)
+        self._host_codes = torch.zeros((self._n_vq, self._max_batch, self.frames), dtype=torch.long).pin_memory()
+        self._host_slots = torch.zeros(self._max_batch, dtype=torch.long).pin_memory()
+        audio_shape = (self._max_batch, self._n_channels, self._samples)
+        self._host_audio = torch.empty(audio_shape, dtype=torch.float32).pin_memory()
+        self._copied = torch.cuda.Event()
+        self._sink: FirstAudioSink | None = None
+        # Request key -> slot for chunks decoded here and not yet seen by the main path.
+        self._decoded: dict[str, int] = {}
+        # A failed output-queue handoff must be replayed from the already
+        # decoded samples. Re-decoding would advance the streaming state twice.
+        self._unsent_audio: dict[str, torch.Tensor] = {}
+        self._handoffs: dict[int, _SlotHandoff] = {}
+        self._lock = threading.Lock()
+        self._jobs: queue.SimpleQueue[_Job | None] = queue.SimpleQueue()
+        self._thread: threading.Thread | None = None
+        self._closed = False
+        self._worker_error: Exception | None = None
+        self._emitted = 0
+        self._batches = 0
+        self._decode_s = 0.0
+
+    def bind(self, sink: FirstAudioSink) -> None:
+        if not callable(getattr(sink, "prepare", None)):
+            raise TypeError("The codec first-chunk sink requires prepare(request_ids)")
+        with self._lock:
+            if self._closed:
+                raise RuntimeError("The codec first-chunk decoder is closed")
+            self._sink = sink
+            if self._thread is None:
+                self._thread = threading.Thread(target=self._run, name="moss-codec-first-chunk", daemon=True)
+                self._thread.start()
+
+    def close(self) -> None:
+        with self._lock:
+            if not self._closed:
+                self._closed = True
+                self._jobs.put(None)
+            thread = self._thread
+        if thread is not None:
+            thread.join(timeout=5.0)
+            # A stalled decode may still own GPU state: keep its thread and
+            # handoff, rather than permitting a restart or slot reuse.
+            if thread.is_alive():
+                logger.warning("Codec first-chunk decoder is still stopping; slot waits remain bounded")
+
+    # ------------------------------------------------------------------
+    # Receive thread
+    # ------------------------------------------------------------------
+    def submit(
+        self,
+        request_id: str,
+        request_key: str,
+        codes_flat: Any,
+        req_slots: dict[str, int],
+    ) -> bool:
+        """Claim a first chunk for fast decode; False leaves it to the main path."""
+        # Own the codes: the job outlives this call and the payload it came from.
+        codes = torch.as_tensor(codes_flat, dtype=torch.long).reshape(-1).clone()
+        if codes.numel() != self._n_vq * self.frames:
+            return False
+        with self._lock:
+            if (
+                self._closed
+                or self._sink is None
+                or self._thread is None
+                or not self._thread.is_alive()
+                or request_key in req_slots
+            ):
+                return False
+            delivery = self._sink.prepare([request_id])
+            if request_id not in delivery.routes:
+                return False
+            slot = self._session.acquire(oldest=True)
+            if slot is None:
+                return False
+            handoff = _SlotHandoff()
+            self._handoffs[slot] = handoff
+            self._decoded[request_key] = slot
+            req_slots[request_key] = slot
+            codes = codes.reshape(self._n_vq, self.frames)
+            # Enqueue under the admission lock: close's sentinel cannot
+            # overtake an accepted job.
+            self._jobs.put(_Job(request_id, request_key, delivery, slot, codes, handoff))
+        return True
+
+    # ------------------------------------------------------------------
+    # Main (engine) thread
+    # ------------------------------------------------------------------
+    def get_request_slot(self, request_key: str, req_slots: dict[str, int]) -> int | None:
+        """Observe a completed admission before deciding whether cleanup is needed.
+
+        The scheduler removes cancelled routes before the runner's finish hook.
+        A submit that already froze its route must finish leasing its slot
+        before that hook checks it. Waiting for decode happens outside this
+        lock, so the worker can still complete the slot handoff.
+        """
+        with self._lock:
+            return req_slots.get(request_key)
+
+    def take_decoded(self, request_key: str) -> bool:
+        """True once per request whose first chunk this path already decoded."""
+        with self._lock:
+            return self._decoded.pop(request_key, None) is not None
+
+    def take_unsent_audio(self, request_key: str) -> torch.Tensor | None:
+        """Return first-chunk samples when the out-of-band sink rejected them."""
+        with self._lock:
+            return self._unsent_audio.pop(request_key, None)
+
+    def forget(self, request_key: str) -> None:
+        with self._lock:
+            self._decoded.pop(request_key, None)
+            self._unsent_audio.pop(request_key, None)
+
+    def gate(self) -> None:
+        """Queue the current stream behind the latest fast decode, if still running."""
+        event = self._inflight
+        if self._gate_main and event is not None and not event.query():
+            torch.cuda.current_stream(self._device).wait_event(event)
+
+    def order_after(self, slot: int) -> None:
+        """Make the current stream's next use of ``slot`` follow the fast decode."""
+        with self._lock:
+            handoff = self._handoffs.get(slot)
+        if handoff is None:
+            return
+        if not handoff.done.wait(timeout=self._handoff_timeout_s):
+            # Keep the handoff. The worker may still be writing the slot, so
+            # neither cancellation nor a retry may release/re-decode it.
+            raise TimeoutError(f"Codec first-chunk slot {slot} did not complete within {self._handoff_timeout_s:g}s")
+        if handoff.error is not None:
+            raise RuntimeError(f"Codec first-chunk decode failed for slot {slot}") from handoff.error
+        if handoff.event is not None:
+            torch.cuda.current_stream(self._device).wait_event(handoff.event)
+        with self._lock:
+            self._handoffs.pop(slot, None)
+
+    # ------------------------------------------------------------------
+    # Decode thread
+    # ------------------------------------------------------------------
+    def _run(self) -> None:
+        try:
+            current_omni_platform.set_device(self._device)
+            self._run_jobs()
+        except Exception as error:
+            logger.exception("MOSS codec first-chunk worker failed")
+            self._worker_error = error
+        finally:
+            with self._lock:
+                self._closed = True
+                for handoff in self._handoffs.values():
+                    if not handoff.done.is_set():
+                        handoff.error = self._worker_error or RuntimeError("Codec first-chunk worker stopped")
+                        handoff.done.set()
+            # Accepted, unprocessed routes must not be left waiting for PCM.
+            while True:
+                try:
+                    job = self._jobs.get_nowait()
+                except queue.Empty:
+                    break
+                if job is not None:
+                    self._fail_delivery(job)
+
+    @staticmethod
+    def _fail_delivery(job: _Job) -> None:
+        try:
+            job.delivery.fail([job.request_id])
+        except Exception:
+            logger.exception("Codec first-chunk error delivery failed for %s", job.request_id)
+
+    def _run_jobs(self) -> None:
+        while True:
+            job = self._jobs.get()
+            if job is None:
+                return
+            jobs = [job]
+            while len(jobs) < self._max_batch:
+                try:
+                    extra = self._jobs.get_nowait()
+                except queue.Empty:
+                    break
+                if extra is None:
+                    self._jobs.put(None)
+                    break
+                jobs.append(extra)
+            try:
+                start = time.perf_counter()
+                self._decode(jobs)
+                self._decode_s += time.perf_counter() - start
+                self._batches += 1
+                before, self._emitted = self._emitted, self._emitted + len(jobs)
+                if self._emitted // 1000 != before // 1000:
+                    logger.info(
+                        "MOSS codec first-chunk fast path: %d chunks in %d batches, %.2f ms/batch",
+                        self._emitted,
+                        self._batches,
+                        1e3 * self._decode_s / self._batches,
+                    )
+            except Exception as error:
+                logger.exception("MOSS codec first-chunk fast decode failed for %s", [j.request_id for j in jobs])
+                # A partial graph replay may have advanced the slot. Surface
+                # the failure instead of re-decoding and advancing it twice.
+                for item in jobs:
+                    item.handoff.error = error
+                    item.handoff.done.set()
+                    self._fail_delivery(item)
+                raise
+
+    @torch.no_grad()
+    def _decode(self, jobs: list[_Job]) -> None:
+        n = len(jobs)
+        for row, item in enumerate(jobs):
+            self._host_codes[:, row].copy_(item.codes)
+            self._host_slots[row] = item.slot
+        stream = self._stream
+        with torch.cuda.stream(stream):
+            for item in jobs:
+                self._session.order_after_reset(item.slot, stream)
+            codes = self._host_codes[:, :n].to(self._device, non_blocking=True).clamp_(0, self._codebook_size - 1)
+            slots = self._host_slots[:n].to(self._device, non_blocking=True)
+            result = self._wrapper.decode(codes, slots)
+            if result is None:
+                raise RuntimeError(f"no first-chunk codec graph for B={n}, T={self.frames}")
+            audio = result[0][:n, ..., : self._samples].float()
+            state_written = torch.cuda.Event()
+            state_written.record(stream)
+            self._inflight = state_written
+            self._host_audio[:n].view_as(audio).copy_(audio, non_blocking=True)
+            self._copied.record(stream)
+        for item in jobs:
+            item.handoff.event = state_written
+        self._copied.synchronize()
+        for row, item in enumerate(jobs):
+            wav = self._host_audio[row].clone()
+            if self._n_channels == 1:
+                wav = wav.reshape(-1)
+            try:
+                item.delivery([item.request_id], [wav], self._sr)
+            except Exception:
+                logger.exception("MOSS codec first-chunk output failed for %s", item.request_id)
+                with self._lock:
+                    self._unsent_audio[item.request_key] = wav
+            finally:
+                item.handoff.done.set()

vllm_omni/model_executor/models/moss_tts/local_model_state.py

@@ -0,0 +1,447 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""MOSS Local request-slot state, independent of each step's batch order.
+
+State writes, gathers and MTP replay run on the runner stream. Output gathers
+own their storage; the runner retains its normal asynchronous output snapshot.
+No request dictionary holds a view into a mutable slot or CUDA graph buffer.
+"""
+
+import numpy as np
+import torch
+from vllm.config import CUDAGraphMode
+from vllm.forward_context import set_forward_context
+from vllm.logger import init_logger
+from vllm.utils.torch_utils import async_tensor_h2d
+from vllm.v1.worker.gpu.input_batch import get_num_sampled_and_rejected
+from vllm.v1.worker.gpu.sample.output import SamplerOutput
+from vllm.v1.worker.gpu.sample.sampler import Sampler
+
+from vllm_omni.model_executor.output_snapshot import PackedOutputSnapshot
+from vllm_omni.worker_v2.model_states.omni_model_state import OmniModelState
+
+logger = init_logger(__name__)
+
+
+def _metadata_to_device(data: np.ndarray, device: torch.device) -> torch.Tensor:
+    # CPU state tests need no CUDA allocator. The serving CUDA path retains
+    # vLLM's asynchronous pinned upload and its ownership/lifetime handling.
+    if device.type == "cpu":
+        return torch.from_numpy(data)
+    return async_tensor_h2d(data, device=device)
+
+
+class _CodeRowsSnapshot(PackedOutputSnapshot):
+    """Owned per-request code rows backed by one gathered device tensor.
+
+    The generic snapshot flattens every per-request leaf, concatenates them
+    into a slab and slices each back, then slices each again after the D2H.
+    Here the gathered ``[decode_rows, n_vq]`` tensor is already owned, so the
+    copy is a single D2H and the per-request views come from one ``split``.
+    """
+
+    def __init__(self, codes: torch.Tensor, decode_rows: list[int], num_reqs: int) -> None:
+        dict.__init__(self, {"codes": {"audio": self._rows(codes, decode_rows, num_reqs)}})
+        self._codes = codes
+        self._decode_rows = decode_rows
+        self._num_reqs = num_reqs
+        self.producer_event = None
+
+    @staticmethod
+    def _rows(codes: torch.Tensor, decode_rows: list[int], num_reqs: int) -> list[torch.Tensor]:
+        outputs = [codes.new_empty((0, codes.shape[-1]))] * num_reqs
+        for batch_row, row in zip(decode_rows, codes.split(1), strict=True):
+            outputs[batch_row] = row
+        return outputs
+
+    def copy_to_cpu(self, copy_tensor):
+        cpu = copy_tensor(self._codes)
+        return {"codes": {"audio": self._rows(cpu, self._decode_rows, self._num_reqs)}}
+
+
+class MossLocalModelState(OmniModelState):
+    def __init__(self, vllm_config, model, encoder_cache, device):
+        super().__init__(vllm_config, model, encoder_cache, device)
+        self._init_slot_buffers(self.scheduler_config.max_num_seqs, model.hidden_size, device, self.dtype)
+        self._batch_prefill = bool(getattr(model.config, "mrv2_batch_prefill", False))
+        self._direct_tokens = bool(getattr(model.config, "mrv2_direct_tokens", False))
+        self._local_eager_mtp = bool(getattr(model.config, "mrv2_eager_mtp", False))
+        logger.info(
+            "MOSS Local MRV2 batch prefill=%s, determined tokens=%s, eager MTP=%s",
+            self._batch_prefill,
+            self._direct_tokens,
+            self._local_eager_mtp,
+        )
+        logger.info("MOSS Local MRV2 GPU slot state enabled: capacity=%d", self.scheduler_config.max_num_seqs)
+
+    def _init_slot_buffers(self, capacity, hidden_size, device, dtype):
+        self._hidden_pool = torch.zeros((capacity, hidden_size), dtype=dtype, device=device)
+        self._codes_pool = torch.full(
+            (capacity, self.model.n_vq), self.model.audio_pad_token_id, dtype=torch.long, device=device
+        )
+        # Execution eligibility mirrors audio_state.is_stopping. Binary MTP
+        # continuation is a per-step output, not a sticky request-state update.
+        self._active_pool = torch.zeros(capacity, dtype=torch.bool, device=device)
+        self._active_host = [False] * capacity
+        self._decode_rows = []
+        self._prefill_staging = []
+        self._batch_prefill = False
+        self._direct_tokens = False
+        # Eager MTP: a frame is sampled right after the forward that produced
+        # its hidden state, one step earlier than the canonical preprocess
+        # path. Its input-embedding contribution (audio embedding times emit
+        # mask) is added to the next step's text embedding; ``_keep_pool``
+        # chains emission so a stopped stream cannot emit again.
+        self._local_eager_mtp = False
+        self._eager_emb_pool = torch.zeros((capacity, hidden_size), dtype=dtype, device=device)
+        self._keep_pool = torch.zeros(capacity, dtype=torch.bool, device=device)
+        self._completing_rows = []
+        self._local_eager_rows = []
+        # A slot takes the eager path only after its final prefill chunk ran
+        # an eager MTP in this state; others keep the canonical path, so a
+        # reused slot never reads a previous request's frame state.
+        self._eager_host = [False] * capacity
+        self._mtp_dispatcher = None
+
+    def sample_determined_tokens(self, batch, sampler):
+        """Skip full-vocabulary sampling only for unconstrained Local tokens.
+
+        The audio/binary MTP sampler has already made the random decision. The
+        outer text distribution has exactly one finite logit. Keep the original
+        path for processors that can mask it or request distribution metadata.
+        """
+        if not self._direct_tokens or type(sampler) is not Sampler:
+            return None
+        rows = batch.idx_mapping_np
+        if (
+            sampler.compute_nans
+            or sampler.return_sampling_mask
+            or sampler.trace_replay_state is not None
+            or sampler.get_logprobs_dims(rows) is not None
+            or np.any(sampler.logit_bias_state.use_logit_bias[rows])
+            or np.any(sampler.penalties_state.use_penalty[rows])
+            or np.any(sampler.bad_words_state.num_bad_words.np[rows])
+            or (
+                sampler.thinking_budget_state.enabled
+                and np.any(sampler.thinking_budget_state.use_thinking_budget[rows])
+            )
+        ):
+            return None
+        keep = self.model._batch_should_continue
+        if keep is None or keep.numel() != batch.num_reqs:
+            return None
+        # Owned storage: subsequent forwards can replace the continue mask
+        # while these token IDs are being copied or delivered asynchronously.
+        tokens = torch.where(keep, self.model.audio_assistant_slot_token_id, self.model.im_end_token_id)
+        num_sampled, num_rejected = get_num_sampled_and_rejected(
+            batch.seq_lens.new_ones(batch.num_reqs),
+            batch.seq_lens,
+            batch.cu_num_logits,
+            batch.idx_mapping,
+            sampler.req_states.prefill_len.gpu,
+        )
+        return SamplerOutput(tokens.reshape(-1, 1), None, None, num_sampled, num_rejected)
+
+    def _apply_reference_batch(self, embeds, references):
+        """One owned pinned upload for reference codes and their token positions.
+
+        Host storage may be reused only after its H2D event completes. Device
+        storage is consumed and overwritten in runner-stream order. If all
+        cached slots are busy, allocate a transient snapshot instead of waiting.
+        """
+        if not references:
+            return
+        count = sum(codes.shape[0] for _, codes in references)
+        capacity = max(count, self._static_inputs_embeds.shape[0])
+        staging = None
+        for item in self._prefill_staging:
+            if item[0].shape[0] >= count and (item[2] is None or item[2].query()):
+                staging = item
+                break
+        if staging is None:
+            host = torch.empty((capacity, self.model.n_vq + 1), dtype=torch.long, pin_memory=embeds.is_cuda)
+            device = torch.empty_like(host, device=embeds.device) if embeds.is_cuda else host
+            staging = [host, device, torch.cuda.Event() if embeds.is_cuda else None]
+            if len(self._prefill_staging) < 3:
+                self._prefill_staging.append(staging)
+        host, device, event = staging
+        offset = 0
+        for start, codes in references:
+            end = offset + codes.shape[0]
+            host[offset:end, 0].numpy()[:] = np.arange(start, start + codes.shape[0])
+            host[offset:end, 1:].copy_(codes)
+            offset = end
+        if embeds.is_cuda:
+            device[:count].copy_(host[:count], non_blocking=True)
+            event.record(torch.cuda.current_stream(embeds.device))
+        packed = device[:count]
+        positions = packed[:, 0]
+        audio = self.model._audio_embed(packed[:, 1:])
+        embeds.index_copy_(0, positions, embeds.index_select(0, positions) + audio)
+
+    def add_request(self, req_index, new_req_data):
+        super().add_request(req_index, new_req_data)
+        self._hidden_pool[req_index].zero_()
+        self._codes_pool[req_index].fill_(self.model.audio_pad_token_id)
+        self._active_pool[req_index].fill_(True)
+        self._active_host[req_index] = True
+        self._eager_host[req_index] = False
+        self._keep_pool[req_index].fill_(False)
+        self._eager_emb_pool[req_index].zero_()
+
+    def remove_request(self, req_index_or_id):
+        # Freeing is CPU bookkeeping. Reinitialization happens on admission,
+        # in stream order, so pending output snapshots remain independent.
+        super().remove_request(req_index_or_id)
+
+    def _select_rows(self, tensor, rows):
+        if rows == list(range(len(rows))):
+            return tensor[: len(rows)]
+        indices = _metadata_to_device(np.asarray(rows, dtype=np.int64), tensor.device)
+        return tensor.index_select(0, indices)
+
+    def run_preprocess(self, input_batch, model_inputs, req_states=None, mtp_batch_descriptor_dispatcher=None):
+        input_ids = model_inputs.get("input_ids")
+        if input_ids is None:
+            input_ids = input_batch.input_ids
+        embeds = model_inputs.get("inputs_embeds")
+        if embeds is None:
+            embeds = self.model.embed_input_ids(input_ids[: input_batch.num_tokens])
+            model_inputs["inputs_embeds"] = embeds
+        elif self._static_inputs_embeds is not None and embeds.data_ptr() == self._static_inputs_embeds.data_ptr():
+            embeds[: input_batch.num_tokens].copy_(self.model.embed_input_ids(input_ids[: input_batch.num_tokens]))
+
+        decode_rows = []
+        completing_rows = []
+        references = []
+        for row in range(input_batch.num_reqs):
+            slot = int(input_batch.idx_mapping_np[row])
+            buf = self.intermediate_buffer.buffers[slot]
+            if not buf or str(buf.get("req_id", "")).startswith("_warmup_"):
+                continue
+            start = int(input_batch.query_start_loc_np[row])
+            count = int(input_batch.num_scheduled_tokens[row])
+            prompt_len = (
+                None if req_states is None else self._get_req_state_value(getattr(req_states, "prompt_len", None), slot)
+            )
+            computed = (
+                None
+                if req_states is None
+                else self._get_req_state_value(getattr(req_states, "num_computed_tokens", None), slot)
+            )
+            if computed is None:
+                computed = self._get_input_batch_num_computed(input_batch, slot, row)
+            prefill = computed < prompt_len if computed is not None and prompt_len is not None else count > 1
+            if prefill and computed is not None and prompt_len is not None and computed + count >= prompt_len:
+                completing_rows.append(row)
+            if count == 1 and not prefill and isinstance(buf.get("audio_state"), dict):
+                active = not bool(buf["audio_state"].get("is_stopping"))
+                if active != self._active_host[slot]:
+                    self._active_pool[slot].fill_(active)
+                    self._active_host[slot] = active
+                decode_rows.append(row)
+                continue
+            if self._batch_prefill:
+                ref = (buf.get("codes", {}) or {}).get("ref")
+                ref_offset = int(buf.get("ref_offset", 0))
+                # Device-side references keep the canonical path; do not add
+                # an implicit D2H just to place them in a host staging buffer.
+                if not isinstance(ref, torch.Tensor) or ref.device.type == "cpu":
+                    if isinstance(ref, torch.Tensor) and ref.numel():
+                        if ref.dim() == 1 and ref.numel() % self.model.n_vq == 0:
+                            ref = ref.view(-1, self.model.n_vq)
+                        if ref.dim() == 2:
+                            chunk = ref[ref_offset : ref_offset + count]
+                            if chunk.numel() and chunk.shape[0] == count:
+                                references.append((start, chunk))
+                    # Text embeddings have already been computed for all input
+                    # tokens. Match the canonical prefill state transition.
+                    self.intermediate_buffer.update(
+                        slot,
+                        {
+                            "audio_state": {"is_stopping": False},
+                            "ref_offset": ref_offset + count,
+                        },
+                        self.model.gpu_resident_buffer_keys,
+                    )
+                    if not self._active_host[slot]:
+                        self._active_pool[slot].fill_(True)
+                        self._active_host[slot] = True
+                    continue
+            # Reference conditioning and chunked-prefill coordinates retain
+            # the model's canonical scalar implementation.
+            info = dict(buf)
+            info["_omni_is_prefill"] = prefill
+            if prompt_len is not None:
+                info["_omni_prompt_len"] = prompt_len
+            if computed is not None:
+                info["_omni_num_computed_tokens"] = computed
+            ids = input_ids[start : start + count]
+            previous = embeds[start : start + count]
+            new_ids, new_embeds, updates = self.model.preprocess(ids, previous, **info)
+            if "mtp_inputs" in updates:
+                raise RuntimeError("MOSS slot state requires an admitted prefill before decode")
+            if self._preprocess_result_needs_writeback(ids, new_ids):
+                ids.copy_(new_ids)
+            if self._preprocess_result_needs_writeback(previous, new_embeds):
+                previous.copy_(new_embeds)
+            self.intermediate_buffer.update(slot, updates, self.model.gpu_resident_buffer_keys)
+            active = not bool(updates.get("audio_state", {}).get("is_stopping", False))
+            if active != self._active_host[slot]:
+                self._active_pool[slot].fill_(active)
+                self._active_host[slot] = active
+
+        self._apply_reference_batch(embeds, references)
+        self._decode_rows = decode_rows
+        if self._local_eager_mtp:
+            self._completing_rows = completing_rows
+            self._mtp_dispatcher = mtp_batch_descriptor_dispatcher
+            eager_rows = [row for row in decode_rows if self._eager_host[int(input_batch.idx_mapping_np[row])]]
+            canonical_rows = [row for row in decode_rows if not self._eager_host[int(input_batch.idx_mapping_np[row])]]
+            self._local_eager_rows = eager_rows
+            if eager_rows:
+                # Their frame was sampled after the previous forward.
+                slots = self._select_rows(input_batch.idx_mapping, eager_rows)
+                offsets = self._mtp_offsets[: len(eager_rows)]
+                offsets.copy_(self._select_rows(input_batch.query_start_loc, eager_rows))
+                embeds.index_copy_(
+                    0, offsets, embeds.index_select(0, offsets) + self._eager_emb_pool.index_select(0, slots)
+                )
+            if canonical_rows:
+                self._run_slot_mtp(input_batch, input_ids, embeds, canonical_rows, mtp_batch_descriptor_dispatcher)
+        elif decode_rows:
+            self._run_slot_mtp(input_batch, input_ids, embeds, decode_rows, mtp_batch_descriptor_dispatcher)
+
+    def _run_slot_mtp(self, batch, input_ids, embeds, rows, dispatcher):
+        bsz = len(rows)
+        slots = self._select_rows(batch.idx_mapping, rows)
+        # Upstream query_start_loc is int32; index_copy_ requires int64.
+        offsets = self._mtp_offsets[:bsz]
+        offsets.copy_(self._select_rows(batch.query_start_loc, rows))
+        ids, emb = self._mtp_input_ids[:bsz], self._mtp_input_embeds[:bsz]
+        hidden, step = self._mtp_hidden[:bsz], self._mtp_text_step[:bsz]
+        ids.copy_(input_ids.index_select(0, offsets))  # MRV2 token inputs are int32.
+        torch.index_select(embeds, 0, offsets, out=emb)
+        torch.index_select(self._hidden_pool, 0, slots, out=hidden)
+        step[:, 0].copy_(self._active_pool.index_select(0, slots))
+
+        buffers = [self.intermediate_buffer.buffers[int(batch.idx_mapping_np[row])] for row in rows]
+        req_ids = [str(buf["req_id"]) for buf in buffers]
+        generators = [
+            self._get_mtp_generator(req_id, buf.get("sampling_params"), ids.device)
+            for req_id, buf in zip(req_ids, buffers, strict=True)
+        ]
+        can_graph = self._is_mtp_graph_runner() and not any(g is not None for g in generators)
+        desc = dispatcher(bsz) if can_graph and dispatcher is not None else None
+        mode = getattr(desc, "cg_mode", CUDAGraphMode.FULL) if can_graph else CUDAGraphMode.NONE
+        num_tokens = int(desc.num_tokens) if desc is not None else bsz
+        with set_forward_context(
+            None, self.vllm_config, num_tokens=num_tokens, cudagraph_runtime_mode=mode, batch_descriptor=desc
+        ):
+            new_emb, codes = self._call_mtp_with_sampling(
+                self._mtp_input_ids[:num_tokens],
+                self._mtp_input_embeds[:num_tokens],
+                self._mtp_hidden[:num_tokens],
+                self._mtp_text_step[:num_tokens],
+                buffers=buffers,
+                req_ids=req_ids,
+                generators=generators,
+            )
+        embeds.index_copy_(0, offsets, new_emb[:bsz])
+        self._codes_pool.index_copy_(0, slots, codes[:bsz])
+
+    def _run_local_eager_mtp(self, batch, rows, completing):
+        """Sample the next frame from this forward's hidden state.
+
+        Same MTP inputs as the canonical next-step preprocess (hidden state
+        and active flag); the input embedding is zero, so the returned
+        embedding is exactly the audio contribution the next step adds to its
+        text embedding. A decode row emits only if its previous frame did.
+        """
+        bsz = len(rows)
+        slots = self._select_rows(batch.idx_mapping, rows)
+        ids, emb = self._mtp_input_ids[:bsz], self._mtp_input_embeds[:bsz]
+        hidden, step = self._mtp_hidden[:bsz], self._mtp_text_step[:bsz]
+        ids.zero_()  # the MTP does not read token ids
+        emb.zero_()
+        torch.index_select(self._hidden_pool, 0, slots, out=hidden)
+        gate = self._keep_pool.index_select(0, slots)
+        if completing:
+            first = _metadata_to_device(
+                np.asarray([row in completing for row in rows], dtype=np.bool_),
+                gate.device,
+            )
+            gate = gate | first
+        step[:, 0].copy_(self._active_pool.index_select(0, slots) & gate)
+
+        buffers = [self.intermediate_buffer.buffers[int(batch.idx_mapping_np[row])] for row in rows]
+        req_ids = [str(buf["req_id"]) for buf in buffers]
+        generators = [
+            self._get_mtp_generator(req_id, buf.get("sampling_params"), ids.device)
+            for req_id, buf in zip(req_ids, buffers, strict=True)
+        ]
+        dispatcher = self._mtp_dispatcher
+        can_graph = self._is_mtp_graph_runner() and not any(g is not None for g in generators)
+        desc = dispatcher(bsz) if can_graph and dispatcher is not None else None
+        mode = getattr(desc, "cg_mode", CUDAGraphMode.FULL) if can_graph else CUDAGraphMode.NONE
+        num_tokens = int(desc.num_tokens) if desc is not None else bsz
+        with set_forward_context(
+            None, self.vllm_config, num_tokens=num_tokens, cudagraph_runtime_mode=mode, batch_descriptor=desc
+        ):
+            contribution, codes = self._call_mtp_with_sampling(
+                self._mtp_input_ids[:num_tokens],
+                self._mtp_input_embeds[:num_tokens],
+                self._mtp_hidden[:num_tokens],
+                self._mtp_text_step[:num_tokens],
+                buffers=buffers,
+                req_ids=req_ids,
+                generators=generators,
+            )
+        codes = codes[:bsz]
+        self._eager_emb_pool.index_copy_(0, slots, contribution[:bsz])
+        self._codes_pool.index_copy_(0, slots, codes)
+        self._keep_pool.index_copy_(0, slots, codes.ne(self.model.audio_pad_token_id).any(dim=-1))
+
+    def run_postprocess(self, hidden_states, input_batch):
+        # query_start_loc is already device-resident and authoritative for
+        # mixed/chunked-prefill batches. index_copy snapshots graph outputs.
+        last = input_batch.query_start_loc[1 : input_batch.num_reqs + 1] - 1
+        self._hidden_pool.index_copy_(
+            0, input_batch.idx_mapping[: input_batch.num_reqs], hidden_states.index_select(0, last)
+        )
+        if self._local_eager_mtp:
+            completing = self._completing_rows
+            rows = sorted(self._local_eager_rows + completing)
+            # Canonical decode rows emit their preprocess frame; eager rows
+            # and completed prefills emit the frame sampled below.
+            self._decode_rows = sorted(set(self._decode_rows) | set(completing))
+            if rows:
+                self._run_local_eager_mtp(input_batch, rows, set(completing))
+                for row in completing:
+                    self._eager_host[int(input_batch.idx_mapping_np[row])] = True
+
+    def postprocess_model_output(self, model_output, input_batch, req_states):
+        # Local backbone returns a tensor; keep the standard runner's output
+        # snapshot/materializer contract and omit hidden payload transport.
+        if not isinstance(model_output, torch.Tensor):
+            raise TypeError("MOSS Local slot state requires a tensor backbone output")
+        model = self.model
+        model._batch_state = None
+        model._batch_state_spans = None
+        model._batch_should_continue = torch.ones(input_batch.num_reqs, device=model_output.device, dtype=torch.bool)
+        if not self._decode_rows:
+            return model_output, {}
+        # This gather is an owned snapshot. Later slot writes cannot change
+        # data already handed to the runner, even before its D2H completes.
+        slots = self._select_rows(input_batch.idx_mapping, self._decode_rows)
+        codes = self._codes_pool.index_select(0, slots)
+        keep = codes.ne(model.audio_pad_token_id).any(dim=-1)
+        if self._decode_rows == list(range(input_batch.num_reqs)):
+            model._batch_should_continue = keep
+        else:
+            rows = _metadata_to_device(
+                np.asarray(self._decode_rows, dtype=np.int64),
+                codes.device,
+            )
+            model._batch_should_continue.index_copy_(0, rows, keep)
+        return model_output, _CodeRowsSnapshot(codes, list(self._decode_rows), input_batch.num_reqs)

vllm_omni/model_executor/models/moss_tts/local_short_attention.py

@@ -0,0 +1,100 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Single-query attention for the frame-local MOSS depth transformer.
+
+The caller has already written this token's K/V and provides only its valid
+prefix. There is no persistent state here. BF16 probability rounding matches
+the FlashAttention P/V multiply, but reduction order is not bit-identical.
+"""
+
+import torch
+from vllm.triton_utils import tl, triton
+
+
+@triton.jit
+def _short_attention(
+    q_ptr,
+    k_ptr,
+    v_ptr,
+    out_ptr,
+    qs0: tl.constexpr,
+    qs1: tl.constexpr,
+    ks0: tl.constexpr,
+    ks1: tl.constexpr,
+    ks2: tl.constexpr,
+    vs0: tl.constexpr,
+    vs1: tl.constexpr,
+    vs2: tl.constexpr,
+    heads: tl.constexpr,
+    dim: tl.constexpr,
+    tokens: tl.constexpr,
+    block_d: tl.constexpr,
+    block_t: tl.constexpr,
+):
+    row = tl.program_id(0)
+    batch, head = row // heads, row % heads
+    d, t = tl.arange(0, block_d), tl.arange(0, block_t)
+    query = tl.load(q_ptr + batch * qs0 + head * qs1 + d, d < dim, 0).to(tl.float32)
+    keys = tl.load(
+        k_ptr + batch * ks0 + head * ks1 + t[:, None] * ks2 + d[None, :],
+        (t[:, None] < tokens) & (d[None, :] < dim),
+        0,
+    ).to(tl.float32)
+    logits = tl.sum(keys * query[None, :], axis=1) * (dim**-0.5 * 1.4426950408889634)
+    logits = tl.where(t < tokens, logits, float("-inf"))
+    probs = tl.exp2(logits - tl.max(logits, axis=0))
+    denominator = tl.sum(probs, axis=0)
+    values = tl.load(
+        v_ptr + batch * vs0 + head * vs1 + t[:, None] * vs2 + d[None, :],
+        (t[:, None] < tokens) & (d[None, :] < dim),
+        0,
+    )
+    probs = probs.to(values.dtype).to(tl.float32)
+    output = tl.sum(values.to(tl.float32) * probs[:, None], axis=0) / denominator
+    tl.store(out_ptr + row * dim + d, output, d < dim)
+
+
+@torch.library.custom_op("vllm_omni::moss_local_short_attention", mutates_args=())
+def local_short_attention(query: torch.Tensor, key: torch.Tensor, value: torch.Tensor) -> torch.Tensor:
+    batch, heads, seq_len, dim = query.shape
+    if not (
+        query.is_cuda
+        and key.device == query.device
+        and value.device == query.device
+        and query.dtype == key.dtype == value.dtype == torch.bfloat16
+        and seq_len == 1
+        and 0 < dim <= 128
+        and key.shape == value.shape
+        and key.shape[:2] == query.shape[:2]
+        and key.shape[-1] == dim
+        and 0 < key.shape[2] <= 16
+        and query.stride(-1) == key.stride(-1) == value.stride(-1) == 1
+    ):
+        raise ValueError("MOSS short attention requires CUDA BF16 single-query inputs and 1..16 valid keys")
+    output = query.new_empty((batch, heads, 1, dim))
+    _short_attention[(batch * heads,)](
+        query,
+        key,
+        value,
+        output,
+        query.stride(0),
+        query.stride(1),
+        key.stride(0),
+        key.stride(1),
+        key.stride(2),
+        value.stride(0),
+        value.stride(1),
+        value.stride(2),
+        heads,
+        dim,
+        key.shape[2],
+        triton.next_power_of_2(dim),
+        triton.next_power_of_2(key.shape[2]),
+        num_warps=1,
+    )
+    return output
+
+
+@local_short_attention.register_fake
+def _fake_local_short_attention(query: torch.Tensor, key: torch.Tensor, value: torch.Tensor) -> torch.Tensor:
+    return query.new_empty(query.shape)

vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_codec.py

@@ -8,7 +8,10 @@
 
 from __future__ import annotations
 
+import threading
 from collections.abc import Iterable
+from contextlib import nullcontext
+from functools import partial
 from typing import Any
 
 import torch
@@ -32,11 +35,16 @@ from vllm_omni.model_executor.models.moss_tts.configuration_moss_audio_tokenizer
 from vllm_omni.model_executor.models.moss_tts.cuda_graph_streaming_decoder_wrapper import (
     CUDAGraphStreamingDecoderWrapper,
 )
+from vllm_omni.model_executor.models.moss_tts.first_chunk_fast_path import (
+    MossFirstChunkFastPath,
+)
 from vllm_omni.model_executor.models.moss_tts.moss_codec_cudagraph import (
     MossTTSCUDAGraphCodecWrapper,
 )
 from vllm_omni.model_executor.models.output_templates import OmniOutput
+from vllm_omni.model_executor.output_snapshot import PackedOutputSnapshot, pack_output_snapshot
 from vllm_omni.model_executor.stage_input_processors.chunk_size_utils import parse_chunk_ramp
+from vllm_omni.worker_v2.first_audio_sender import FirstAudioSink
 
 logger = init_logger(__name__)
 
@@ -53,18 +61,46 @@ class _MossCodecStreamSession:
         vllm_config: VllmConfig,
         graph_batch_sizes: list[int] | None = None,
         graph_frame_sizes: list[int] | None = None,
+        gpu_output: bool = False,
+        chunk_frames: int = 0,
+        ring_headroom: bool = False,
+        fast_batch_sizes: list[int] | None = None,
+        fast_frames: int = 0,
     ) -> None:
         self._codec = codec
         self._state_capacity = int(state_capacity)
+        self._chunk_frames = int(chunk_frames)
+        self._ring_headroom = bool(ring_headroom)
         self._n_vq = int(n_vq)
         self._device = next(codec.parameters()).device
+        self._gpu_output = gpu_output and self._device.type == "cuda"
         self._free_stream_slots = list(reversed(range(self._state_capacity)))
         self._leased_slots: set[int] = set()
+        # The first-chunk fast path leases slots from its decode thread.
+        self._slot_lock = threading.Lock()
+        self._slot_reset_events: dict[int, torch.cuda.Event] = {}
         self._closed = False
+        self._metadata_index = 0
+        self._metadata_ring = []
+        if self._gpu_output:
+            # The metadata DMA is tiny, but a short ring forces the producer
+            # CPU thread to synchronize with the codec stream while it waits
+            # for a host slab to become reusable.  Keep enough entries for a
+            # scheduler burst so staging remains enqueue-only in practice.
+            for _ in range(64):
+                host = torch.empty(2 * self._state_capacity, dtype=torch.long, pin_memory=True)
+                gpu = torch.empty_like(host, device=self._device)
+                self._metadata_ring.append((host, host.numpy(), gpu, torch.cuda.Event(), False))
         self._cudagraph_wrapper: CUDAGraphStreamingDecoderWrapper | None = None
+        self._fast_wrapper: CUDAGraphStreamingDecoderWrapper | None = None
         batch_sizes = sorted({int(size) for size in (graph_batch_sizes or []) if 0 < int(size) <= self._state_capacity})
         frame_sizes = sorted({int(size) for size in (graph_frame_sizes or []) if int(size) > 0})
         scratch_capacity = max(batch_sizes, default=0) if self._device.type in ("cuda", "npu") else 0
+        fast_batch_sizes = sorted({int(size) for size in (fast_batch_sizes or []) if 0 < int(size)})
+        if not (fast_frames > 0 and fast_batch_sizes and batch_sizes and frame_sizes and self._device.type == "cuda"):
+            fast_batch_sizes = []
+        fast_scratch_base = self._state_capacity + scratch_capacity
+        scratch_capacity += max(fast_batch_sizes, default=0)
         self._total_state_capacity = self._state_capacity + scratch_capacity
         self._state_slot_ids = torch.arange(
             self._total_state_capacity,
@@ -76,7 +112,14 @@ class _MossCodecStreamSession:
         if not callable(initialize_state_pool):
             raise RuntimeError("The streaming codec does not implement a decoder state pool.")
         with torch.no_grad():
-            initialize_state_pool(self._state_capacity, scratch_capacity)
+            if self._ring_headroom:
+                if self._chunk_frames <= 0:
+                    raise ValueError("codec_ring_headroom requires a positive codec chunk frame count.")
+                # Size every attention ring for context + one full chunk so a
+                # chunk write never evicts keys inside its own causal window.
+                initialize_state_pool(self._state_capacity, scratch_capacity, chunk_frames=self._chunk_frames)
+            else:
+                initialize_state_pool(self._state_capacity, scratch_capacity)
         if batch_sizes and frame_sizes and self._device.type in ("cuda", "npu"):
             if self._device.type == "npu":
                 from vllm_omni.platforms.npu.models.moss_tts_streaming_decode_wrapper import (
@@ -101,16 +144,34 @@ class _MossCodecStreamSession:
                     vllm_config=vllm_config,
                 )
             self._cudagraph_wrapper.warmup(self._device)
+            if fast_batch_sizes and self._cudagraph_wrapper.is_ready:
+                main_compiled = self._cudagraph_wrapper._compiled_decode
+                self._fast_wrapper = CUDAGraphStreamingDecoderWrapper(
+                    codec,
+                    state_capacity=self._state_capacity,
+                    batch_sizes=fast_batch_sizes,
+                    frame_sizes=[int(fast_frames)],
+                    num_quantizers=self._n_vq,
+                    vllm_config=vllm_config,
+                    scratch_base=fast_scratch_base,
+                    private_pool=True,
+                    compiled_decode=main_compiled if main_compiled is not None else codec.decode_streaming_tensors,
+                )
+                self._fast_wrapper.warmup(self._device)
+                if not self._fast_wrapper.is_ready:
+                    self._fast_wrapper = None
             self.reset_slots(list(range(self._state_capacity + scratch_capacity)))
             if not self._cudagraph_wrapper.is_ready:
                 self._cudagraph_wrapper = None
 
-    def acquire(self) -> int | None:
-        if not self._free_stream_slots:
-            return None
-        slot = self._free_stream_slots.pop()
-        self._leased_slots.add(slot)
-        return slot
+    def acquire(self, *, oldest: bool = False) -> int | None:
+        """Lease a slot; ``oldest`` takes the one whose reset was queued first."""
+        with self._slot_lock:
+            if not self._free_stream_slots:
+                return None
+            slot = self._free_stream_slots.pop(0 if oldest else -1)
+            self._leased_slots.add(slot)
+            return slot
 
     def release(self, slot: int, *, state_already_reset: bool = False) -> None:
         if self._closed:
@@ -119,8 +180,23 @@ class _MossCodecStreamSession:
             return
         if not state_already_reset:
             self.reset_slots([slot])
-        self._leased_slots.remove(slot)
-        self._free_stream_slots.append(slot)
+        reset_event = None
+        if self._fast_wrapper is not None:
+            # The fast path decodes on another stream; it must not read this
+            # slot before the reset queued on the current stream has run.
+            reset_event = torch.cuda.Event()
+            reset_event.record(torch.cuda.current_stream(self._device))
+        with self._slot_lock:
+            if reset_event is not None:
+                self._slot_reset_events[slot] = reset_event
+            self._leased_slots.remove(slot)
+            self._free_stream_slots.append(slot)
+
+    def order_after_reset(self, slot: int, stream: torch.cuda.Stream) -> None:
+        with self._slot_lock:
+            event = self._slot_reset_events.pop(slot, None)
+        if event is not None and not event.query():
+            stream.wait_event(event)
 
     def _reset_slot_ids(self, slot_ids: torch.Tensor) -> None:
         reset_streaming_slots = getattr(self._codec, "reset_decoder_state_slots", None)
@@ -139,9 +215,29 @@ class _MossCodecStreamSession:
             raise ValueError(f"Decoder state slots must be contiguous, got {slots}")
         self._reset_slot_ids(self._state_slot_ids[start : start + len(slots)])
 
+    def _stage_slot_ids(self, slots: list[int], terminal_slots: set[int]) -> tuple[torch.Tensor, torch.Tensor]:
+        # Host slabs cannot be overwritten until their DMA completes. Device
+        # slabs are reused in producer-stream order, after all previous reads.
+        index = self._metadata_index
+        host, view, gpu, event, recorded = self._metadata_ring[index]
+        if recorded and not event.query():
+            event.synchronize()
+        terminal_ids = [slot for slot in slots if slot in terminal_slots]
+        n, end = len(slots), len(slots) + len(terminal_ids)
+        view[:n] = slots
+        view[n:end] = terminal_ids
+        gpu[:end].copy_(host[:end], non_blocking=True)
+        event.record(torch.cuda.current_stream(self._device))
+        self._metadata_ring[index] = (host, view, gpu, event, True)
+        self._metadata_index = (index + 1) % len(self._metadata_ring)
+        return gpu[:n], gpu[n:end]
+
     def close(self) -> None:
         if self._closed:
             return
+        for _, _, _, event, recorded in self._metadata_ring:
+            if recorded:
+                event.synchronize()
         close_state_pool = getattr(self._codec, "close_decoder_state_pool", None)
         with torch.no_grad():
             if callable(close_state_pool):
@@ -154,6 +250,7 @@ class _MossCodecStreamSession:
         slot_codes: dict[int, torch.Tensor],
         *,
         terminal_slots: set[int] | None = None,
+        output_buffers: dict[int, torch.Tensor] | None = None,
     ) -> dict[int, torch.Tensor]:
         if not slot_codes:
             return {}
@@ -178,17 +275,22 @@ class _MossCodecStreamSession:
             ],
             dim=1,
         )
-        state_slot_ids = torch.tensor(slots, device=self._device, dtype=torch.long)
+        staged_terminal_ids = None
+        if self._metadata_ring:
+            state_slot_ids, staged_terminal_ids = self._stage_slot_ids(slots, terminal_slots)
+        else:
+            state_slot_ids = torch.tensor(slots, device=self._device, dtype=torch.long)
 
         graph_output: tuple[torch.Tensor, torch.Tensor, int] | None = None
         if self._cudagraph_wrapper is not None:
-            # A terminal tail can safely pad T to a larger graph: decoder
-            # attention is causal, audio_lengths trims the padding, and every
-            # affected state slot is reset immediately after this step.
+            # An all-terminal group may borrow the next larger frame graph.
+            # The decoder is causal, so cropping the padded tail on the time
+            # axis recovers the exact-T waveform up to attention-ring
+            # eviction (see codec_ring_headroom); the state is reset below.
             graph_output = self._cudagraph_wrapper.decode(
                 codes_step,
                 state_slot_ids,
-                allow_frame_padding=len(terminal_slots) == len(slots),
+                allow_frame_padding=bool(terminal_slots) and len(terminal_slots) == len(slots),
             )
 
         used_cudagraph = graph_output is not None
@@ -214,11 +316,69 @@ class _MossCodecStreamSession:
             audio_tensor = result.audio
 
         if terminal_slots:
-            terminal_rows = [row for row, slot in enumerate(slots) if slot in terminal_slots]
-            terminal_slot_ids = state_slot_ids if len(terminal_rows) == len(slots) else state_slot_ids[terminal_rows]
+            if staged_terminal_ids is not None:
+                terminal_slot_ids = staged_terminal_ids
+            else:
+                terminal_rows = [row for row, slot in enumerate(slots) if slot in terminal_slots]
+                terminal_slot_ids = (
+                    state_slot_ids if len(terminal_rows) == len(slots) else state_slot_ids[terminal_rows]
+                )
             self._reset_slot_ids(terminal_slot_ids)
 
-        audio = audio_tensor.detach().to("cpu", torch.float32)
+        if self._gpu_output and output_buffers is not None:
+            # Write directly into the caller-owned per-forward slab.  This
+            # preserves graph output before the next replay without creating
+            # a temporary waveform tensor that would be packed again later.
+            out: dict[int, torch.Tensor] = {}
+            row = 0
+            while row < len(slots):
+                destination = output_buffers.get(slots[row])
+                if destination is None:
+                    row += 1
+                    continue
+                end = row + 1
+                if end < len(slots) and destination.is_contiguous():
+                    storage = destination.untyped_storage().data_ptr()
+                    offset = destination.storage_offset()
+                    row_elements = destination.numel()
+                    while end < len(slots):
+                        following = output_buffers.get(slots[end])
+                        if not (
+                            following is not None
+                            and following.shape == destination.shape
+                            and following.dtype == destination.dtype
+                            and following.device == destination.device
+                            and following.is_contiguous()
+                            and following.untyped_storage().data_ptr() == storage
+                            and following.storage_offset() == offset + (end - row) * row_elements
+                        ):
+                            break
+                        end += 1
+                if end > row + 1:
+                    # A mixed-frame forward can leave gaps between equal-sized
+                    # rows in its slab. Copy each adjacent run as one batch,
+                    # without overwriting the intervening requests' output.
+                    destination = destination.as_strided(
+                        (end - row, *destination.shape), (destination.numel(), *destination.stride())
+                    )
+                    source = audio_tensor[row:end]
+                else:
+                    source = audio_tensor[row]
+                # Crop time before reshaping: flattening first mixes stereo
+                # channels whenever a row is shorter than the decode batch.
+                destination.copy_(source[..., : destination.shape[-1]].view_as(destination))
+                out.update((slot, output_buffers[slot]) for slot in slots[row:end])
+                row = end
+            return out
+
+        if self._gpu_output:
+            # A graph's static output is overwritten on every replay, including
+            # multiple replays within one forward. Snapshot on the producer
+            # stream before returning; record_stream alone cannot prevent this
+            # overwrite. MRV2 owns the subsequent pinned, asynchronous D2H.
+            audio = audio_tensor.detach().to(dtype=torch.float32, copy=True)
+        else:
+            audio = audio_tensor.detach().to("cpu", torch.float32)
         out: dict[int, torch.Tensor] = {}
         for row, slot in enumerate(slots):
             audio_length = int(slot_codes[slot].shape[1]) * self._samples_per_frame
@@ -270,6 +430,7 @@ class MossTTSCodecDecoder(nn.Module):
     requires_exact_input_shape: bool = True
 
     _OUTPUT_SAMPLE_RATE: int = 24_000
+    _first_chunk_fast_path: MossFirstChunkFastPath | None = None
 
     def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
         super().__init__()
@@ -292,20 +453,38 @@ class MossTTSCodecDecoder(nn.Module):
         self._n_channels: int = 1
         self._sr_tensor = torch.tensor(self._OUTPUT_SAMPLE_RATE, dtype=torch.int32)
         self._stream_session: _MossCodecStreamSession | None = None
+        self._codec_stream: torch.cuda.Stream | None = None
         scheduler_cfg = getattr(self.vllm_config, "scheduler_config", None)
         self._stream_state_capacity = max(1, int(getattr(scheduler_cfg, "max_num_seqs", 1) or 1))
         self._initial_stream_chunk_frames: int = self._connector_int("initial_codec_chunk_frames", 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
+        # Opt-in: size codec attention rings as context + one chunk so chunked
+        # streaming reproduces whole-sequence decoding (costs ring memory).
+        self._stream_ring_headroom: bool = bool(self._connector_int("codec_ring_headroom", default=0))
         self._stream_req_slots: dict[str, int] = {}
+        # Opt-in: decode each stream's first chunk from the receive thread.
+        self._first_chunk_fast = bool(self._connector_int("codec_first_chunk_fast_path", default=0))
+        self._first_chunk_fast_path: MossFirstChunkFastPath | None = None
         self._async_chunk = bool(getattr(self.vllm_config.model_config, "async_chunk", False))
+        self._gpu_stream_output = (
+            self._async_chunk
+            and bool(getattr(vllm_config.model_config, "use_v2_model_runner", False))
+            and getattr(cfg, "model_type", None) == "moss_tts_local"
+            # Diagnostic switch: 0 copies each decode group's batch to the host synchronously.
+            and bool(self._connector_int("codec_gpu_stream_output", default=1))
+        )
         self._streaming_graph_batch_sizes = self._streaming_graph_batch_sizes_from_compilation_config()
-        # codec_chunk_ramp ladder sizes join the streaming graph frame-size
-        # set so every ladder (B, T) combination is captured during warmup.
+        ramp = parse_chunk_ramp(self._connector_extra(), steady=self._stream_chunk_frames)
+        if ramp:
+            # The sender gives the ramp precedence over initial chunk size.
+            # The dedicated first-chunk graph must use that same first shape.
+            self._initial_stream_chunk_frames = ramp[0]
+        self._stream_max_step_frames = max(
+            self._stream_max_step_frames, self._initial_stream_chunk_frames, *(ramp or [])
+        )
         self._streaming_graph_frame_sizes = _resolve_streaming_graph_frame_sizes(
-            self._initial_stream_chunk_frames,
-            self._stream_chunk_frames,
-            self._connector_extra_cfg(),
+            self._initial_stream_chunk_frames, self._stream_chunk_frames, self._connector_extra()
         )
 
     # ------------------------------------------------------------------
@@ -337,6 +516,41 @@ class MossTTSCodecDecoder(nn.Module):
         inputs_embeds: torch.Tensor | None = None,
         runtime_additional_information: list[dict[str, Any]] | None = None,
         **kwargs: Any,
+    ) -> OmniOutput:
+        if self._gpu_stream_output and input_ids is not None and input_ids.is_cuda:
+            if self._codec_stream is None:
+                self._codec_stream = torch.cuda.Stream(device=input_ids.device)
+            main_stream = torch.cuda.current_stream(input_ids.device)
+            # Runner input storage is reused by the next batch on main_stream.
+            # Own it before allowing that preparation to overlap codec decode.
+            owned_input_ids = input_ids.clone()
+            with torch.cuda.stream(self._codec_stream):
+                self._codec_stream.wait_stream(main_stream)
+                owned_input_ids.record_stream(self._codec_stream)
+                output = self._forward(
+                    owned_input_ids,
+                    positions,
+                    intermediate_tensors,
+                    inputs_embeds,
+                    runtime_additional_information,
+                    **kwargs,
+                )
+                if isinstance(output.multimodal_outputs, PackedOutputSnapshot):
+                    output.multimodal_outputs.record_producer_event(self._codec_stream)
+            return output
+        return self._forward(
+            input_ids, positions, intermediate_tensors, inputs_embeds, runtime_additional_information, **kwargs
+        )
+
+    @torch.no_grad()
+    def _forward(
+        self,
+        input_ids: torch.Tensor | None = None,
+        positions: torch.Tensor | None = None,
+        intermediate_tensors: Any = None,
+        inputs_embeds: torch.Tensor | None = None,
+        runtime_additional_information: list[dict[str, Any]] | None = None,
+        **kwargs: Any,
     ) -> OmniOutput:
         """Decode audio VQ codes to waveform.
 
@@ -493,13 +707,37 @@ class MossTTSCodecDecoder(nn.Module):
             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():
+            stream_output_slab = None
+            stream_output_buffers: dict[int, torch.Tensor] = {}
+            if self._gpu_stream_output and device.type == "cuda":
+                total_elements = 0
+                for output_index, _, codes_nq_t, _ in streaming_work:
+                    samples = int(codes_nq_t.shape[1]) * self._codec.downsample_rate
+                    total_elements += samples * (self._n_channels if self._n_channels > 1 else 1)
+                if total_elements:
+                    stream_output_slab = torch.empty(total_elements, dtype=torch.float32, device=device)
+                    offset = 0
+                    for output_index, _, codes_nq_t, _ in streaming_work:
+                        samples = int(codes_nq_t.shape[1]) * self._codec.downsample_rate
+                        elements = samples * (self._n_channels if self._n_channels > 1 else 1)
+                        shape = (self._n_channels, samples) if self._n_channels > 1 else (samples,)
+                        stream_output_buffers[output_index] = stream_output_slab.narrow(0, offset, elements).view(shape)
+                        offset += elements
+            for i, wav in self._decode_streaming_batch(
+                streaming_work,
+                output_buffers=stream_output_buffers or None,
+            ).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},
-        )
+        payload = {"model_outputs": audios, "sr": srs}
+        if self._gpu_stream_output and device.type == "cuda":
+            # One transfer per dtype/device preserves ragged request lengths
+            # without issuing one D2H per waveform. A fresh slot owns each
+            # output until the runner completes its copy; no reusable ring.
+            packed = pack_output_snapshot(payload, {}, max_buckets=0, reuse_existing_storage=True)
+            if packed is not None:
+                payload = packed
+        return OmniOutput(text_hidden_states=None, multimodal_outputs=payload)
 
     @staticmethod
     def _normalize_seq_token_counts(value: Any) -> list[int] | None:
@@ -546,25 +784,104 @@ class MossTTSCodecDecoder(nn.Module):
             vllm_config=self.vllm_config,
             graph_batch_sizes=self._streaming_graph_batch_sizes,
             graph_frame_sizes=self._streaming_graph_frame_sizes,
+            gpu_output=self._gpu_stream_output,
+            chunk_frames=self._stream_max_step_frames,
+            ring_headroom=self._stream_ring_headroom,
+            fast_batch_sizes=[1, 2, 4, 8] if self._first_chunk_fast else None,
+            fast_frames=self._initial_stream_chunk_frames if self._first_chunk_fast else 0,
         )
+        if self._stream_session._fast_wrapper is not None:
+            self._first_chunk_fast_path = MossFirstChunkFastPath(
+                self._stream_session,
+                self._stream_session._fast_wrapper,
+                n_vq=self._n_vq,
+                frames=self._initial_stream_chunk_frames,
+                codebook_size=int(self._codec.config.codebook_size),
+                samples_per_frame=int(self._codec.downsample_rate),
+                n_channels=self._n_channels,
+                sample_rate=self._sr_tensor,
+                device=next(self._codec.parameters()).device,
+                gate_main=bool(self._connector_int("codec_first_chunk_gate", default=0)),
+            )
+            logger.info(
+                "MOSS-TTS codec first-chunk fast path ready: B=%s T=%d",
+                self._stream_session._fast_wrapper.capture_sizes,
+                self._initial_stream_chunk_frames,
+            )
+        elif self._first_chunk_fast:
+            logger.warning("MOSS-TTS codec first-chunk fast path requested but unavailable; using the main path.")
+        if self._stream_ring_headroom:
+            logger.info(
+                "MOSS-TTS codec attention rings sized with headroom for %d-frame chunks",
+                self._stream_max_step_frames,
+            )
         return self._stream_session
 
+    def bind_first_chunk_fast_path(self, sink: FirstAudioSink) -> Any:
+        """Start the fast path with an output sink; returns the receive hook."""
+        fast = self._first_chunk_fast_path
+        if fast is None:
+            return None
+        fast.bind(sink)
+        return self._submit_first_chunk
+
+    def _submit_first_chunk(self, request_id: str, handle: Any, payload: Any) -> bool:
+        fast = self._first_chunk_fast_path
+        if fast is None or not isinstance(payload, dict):
+            return False
+        codes = payload.get("codes")
+        codes = codes.get("audio") if isinstance(codes, dict) else None
+        if codes is None:
+            return False
+        meta = payload.get("meta") if isinstance(payload.get("meta"), dict) else {}
+        key = meta.get("req_id")
+        if isinstance(key, (list, tuple)):
+            key = key[0] if key else None
+        if key is None:
+            key = getattr(handle, "external_req_id", None) or request_id
+        return fast.submit(request_id, str(key), codes, self._stream_req_slots)
+
     def _decode_streaming_batch(
         self,
         items: list[tuple[int, str, torch.Tensor, bool]],
+        *,
+        output_buffers: dict[int, torch.Tensor] | None = None,
     ) -> dict[int, torch.Tensor]:
         session = self._ensure_stream_session()
         if session is None:
             return {}
 
         outputs: dict[int, torch.Tensor] = {}
+        unsent_first_audio: 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))
         graph = session._cudagraph_wrapper
+        # Terminal tails join the regular padded frame bucket when every
+        # execution batch width has a graph, so a step stays one replay
+        # instead of one eager or plain-graph decode per distinct tail length.
         coalesce_tails = graph is not None and max(graph.batch_sizes, default=0) >= self._stream_state_capacity
 
+        fast = self._first_chunk_fast_path
         for output_index, request_id, codes_nq_t, finished in items:
             slot = self._stream_req_slots.get(request_id)
+            if fast is not None and slot is not None:
+                fast.order_after(slot)
+                if fast.take_decoded(request_id):
+                    unsent = fast.take_unsent_audio(request_id)
+                    if unsent is not None:
+                        unsent_first_audio[output_index] = unsent
+                    # The fast path already emitted this chunk's audio and
+                    # advanced the slot state past it.
+                    codes_nq_t = codes_nq_t[:, fast.frames :]
+                    if codes_nq_t.shape[1] == 0:
+                        if unsent is not None:
+                            outputs[output_index] = unsent
+                        if finished:
+                            self._finish_stream_request(request_id, session, slot)
+                        continue
+                    if output_buffers is not None:
+                        # Sized for the untrimmed chunk; decode into a fresh tensor.
+                        output_buffers.pop(output_index, None)
             if slot is None:
                 slot = session.acquire()
                 if slot is None:
@@ -580,9 +897,11 @@ class MossTTSCodecDecoder(nn.Module):
                     slot,
                     codes_nq_t,
                     terminal=finished,
+                    output_buffer=output_buffers.get(output_index) if output_buffers is not None else None,
                 )
                 if wav is not None:
-                    outputs[output_index] = wav
+                    prefix = unsent_first_audio.pop(output_index, None)
+                    outputs[output_index] = torch.cat((prefix, wav), dim=-1) if prefix is not None else wav
                 if finished:
                     self._finish_stream_request(
                         request_id,
@@ -593,21 +912,33 @@ class MossTTSCodecDecoder(nn.Module):
                 continue
 
             frame_count = int(codes_nq_t.shape[1])
-            # Only combine tails that already used the same padded graph.
-            # Preserve exact-size graphs (including T=1) and the eager path.
             group_frames = frame_count
             if coalesce_tails and finished:
+                # Only terminal rows may be padded: their state is reset after
+                # the step, and the output is cropped to the exact frame count.
                 group_frames = graph._select_frame_size(frame_count, allow_padding=True) or frame_count
             grouped.setdefault(group_frames, []).append((output_index, request_id, slot, codes_nq_t, finished))
 
+        if fast is not None and grouped:
+            fast.gate()
         for group in grouped.values():
             plan = {slot: codes_nq_t for _, _, slot, codes_nq_t, _ in group}
             terminal_slots = {slot for _, _, slot, _, finished in group if finished}
-            decoded = session.step(plan, terminal_slots=terminal_slots)
+            step_buffers = (
+                {
+                    slot: output_buffers[output_index]
+                    for output_index, _, slot, _, _ in group
+                    if output_buffers is not None and output_index in output_buffers
+                }
+                if output_buffers is not None
+                else None
+            )
+            decoded = session.step(plan, terminal_slots=terminal_slots, output_buffers=step_buffers)
             for output_index, request_id, slot, _, finished in group:
                 wav = decoded.get(slot)
                 if wav is not None:
-                    outputs[output_index] = wav
+                    prefix = unsent_first_audio.pop(output_index, None)
+                    outputs[output_index] = torch.cat((prefix, wav), dim=-1) if prefix is not None else wav
                 if finished:
                     self._finish_stream_request(
                         request_id,
@@ -625,6 +956,7 @@ class MossTTSCodecDecoder(nn.Module):
         codes_nq_t: torch.Tensor,
         *,
         terminal: bool = False,
+        output_buffer: torch.Tensor | None = None,
     ) -> torch.Tensor | None:
         if codes_nq_t.numel() == 0:
             return None
@@ -633,13 +965,24 @@ class MossTTSCodecDecoder(nn.Module):
         for start in range(0, int(codes_nq_t.shape[1]), max_step_frames):
             chunk = codes_nq_t[:, start : start + max_step_frames]
             is_terminal_chunk = terminal and start + max_step_frames >= int(codes_nq_t.shape[1])
+            chunk_samples = int(chunk.shape[1]) * self._codec.downsample_rate
+            chunk_buffer = (
+                output_buffer[
+                    ..., start * self._codec.downsample_rate : start * self._codec.downsample_rate + chunk_samples
+                ]
+                if output_buffer is not None
+                else None
+            )
             decoded = session.step(
                 {slot: chunk},
                 terminal_slots={slot} if is_terminal_chunk else None,
+                output_buffers={slot: chunk_buffer} if chunk_buffer is not None else None,
             )
             wav = decoded.get(slot)
             if wav is not None:
                 parts.append(wav)
+        if output_buffer is not None:
+            return output_buffer
         if not parts:
             return None
         return torch.cat(parts, dim=-1) if len(parts) > 1 else parts[0]
@@ -665,29 +1008,30 @@ class MossTTSCodecDecoder(nn.Module):
         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)
-            if slot is None:
-                continue
-            if session is not None:
-                self._finish_stream_request(request_id, session, slot)
-            else:
-                self._stream_req_slots.pop(request_id, None)
+        # Cancellation and terminal cleanup share the decoder's stream order;
+        # a reused slot must not be reset ahead of an outstanding decode.
+        context = torch.cuda.stream(self._codec_stream) if self._codec_stream is not None else nullcontext()
+        with context:
+            session = self._stream_session
+            for req_id in finished_req_ids:
+                request_id = str(req_id)
+                fast = self._first_chunk_fast_path
+                slot = (
+                    fast.get_request_slot(request_id, self._stream_req_slots)
+                    if fast is not None
+                    else self._stream_req_slots.get(request_id)
+                )
+                if slot is None:
+                    continue
+                if fast is not None:
+                    fast.order_after(slot)
+                    fast.forget(request_id)
+                if session is not None:
+                    self._finish_stream_request(request_id, session, slot)
+                else:
+                    self._stream_req_slots.pop(request_id, None)
 
-    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_extra_cfg(self) -> dict:
+    def _connector_extra(self) -> dict[str, Any]:
         model_cfg = getattr(self.vllm_config, "model_config", None)
         connector_cfg = getattr(model_cfg, "stage_connector_config", None)
         if isinstance(connector_cfg, dict):
@@ -696,6 +1040,9 @@ class MossTTSCodecDecoder(nn.Module):
             extra_cfg = getattr(connector_cfg, "extra", None)
         return extra_cfg if isinstance(extra_cfg, dict) else {}
 
+    def _connector_int(self, name: str, default: int = 0) -> int:
+        return int(self._connector_extra().get(name, default))
+
     def _streaming_graph_batch_sizes_from_compilation_config(self) -> list[int]:
         if getattr(self.vllm_config.model_config, "enforce_eager", True):
             return []
@@ -816,15 +1163,33 @@ class MossTTSCodecDecoder(nn.Module):
             codec.decoder.to(dtype=torch.bfloat16)
         attention_backend = getattr(self.vllm_config.model_config.hf_config, "codec_attention_backend", "sdpa")
         if attention_backend != "sdpa":
-            if attention_backend != "triton" or device.type != "cuda":
+            if attention_backend not in {"triton", "triton_slot"} or device.type != "cuda":
                 raise ValueError(f"Unsupported codec attention backend/device: {attention_backend}/{device.type}")
             from vllm_omni.model_executor.models.moss_tts.audio_tokenizer_v2 import MossAudioTokenizerMultiheadAttention
             from vllm_omni.model_executor.models.moss_tts.streaming_attention import masked_attention
 
+            if attention_backend == "triton_slot":
+                from vllm_omni.model_executor.models.moss_tts.slot_attention import (
+                    slot_ring_attention,
+                    slot_ring_attention_rows,
+                )
+
+                skip_empty_tiles = bool(self._connector_int("codec_skip_empty_attention_tiles", default=0))
+                fused_slots = bool(self._connector_int("codec_fused_slot_attention", default=0))
+                if skip_empty_tiles and fused_slots:
+                    raise ValueError("codec_skip_empty_attention_tiles requires unfused slot attention")
+                if skip_empty_tiles:
+                    slot_ring_attention = partial(slot_ring_attention, skip_empty_tiles=True)
+                    logger.info("MOSS codec slot attention: skipping empty tiles for T <= 32")
+
             for module in codec.decoder.modules():
                 if isinstance(module, MossAudioTokenizerMultiheadAttention):
                     module._streaming_attention = masked_attention
-            logger.info("Enabled Triton masked attention for the streaming codec decoder")
+                    if attention_backend == "triton_slot":
+                        module._slot_attention = slot_ring_attention
+                        if self._connector_int("codec_fused_slot_attention", default=0):
+                            module._slot_attention_rows = slot_ring_attention_rows
+            logger.info("Enabled codec attention backend=%s", attention_backend)
         build_decode_lut = getattr(codec.quantizer, "build_decode_lut", None)
         if callable(build_decode_lut):
             lut_dtype = torch.bfloat16 if device.type == "cuda" else torch.float32

vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_local_depth.py

@@ -21,7 +21,8 @@ checkpoint 1:1 so ``load_weights()`` needs no remapping.
 
 from __future__ import annotations
 
-from collections.abc import Sequence
+import os
+from collections.abc import Callable, Sequence
 
 import torch
 import torch.nn as nn
@@ -52,6 +53,7 @@ class _MossTTSLocalAttention(nn.Module):
         self.n_head = n_head
         self.head_dim = hidden_size // n_head
         self.embed_dim = hidden_size
+        self._short_attention = os.environ.get("VLLM_OMNI_MOSS_LOCAL_SHORT_ATTN", "0") == "1"
         self.c_attn = nn.Linear(hidden_size, 3 * hidden_size, bias=True)
         self.c_proj = nn.Linear(hidden_size, hidden_size, bias=True)
         inv_freq = 1.0 / (rope_base ** (torch.arange(0, self.head_dim, 2, dtype=torch.float32) / self.head_dim))
@@ -134,7 +136,19 @@ class _MossTTSLocalAttention(nn.Module):
             v_cache[:, :, position : position + 1].copy_(value)
             key = k_cache[:, :, : position + 1]
             value = v_cache[:, :, : position + 1]
-        attn_output = F.scaled_dot_product_attention(query, key, value, is_causal=kv_cache is None)
+        if (
+            self._short_attention
+            and kv_cache is not None
+            and query.is_cuda
+            and query.dtype == torch.bfloat16
+            and self.head_dim <= 128
+            and key.shape[2] <= 16
+        ):
+            from vllm_omni.model_executor.models.moss_tts.local_short_attention import local_short_attention
+
+            attn_output = local_short_attention(query, key, value)
+        else:
+            attn_output = F.scaled_dot_product_attention(query, key, value, is_causal=kv_cache is None)
         attn_output = attn_output.transpose(1, 2).reshape(batch_size, seq_len, self.embed_dim)
         return self.c_proj(attn_output)
 
@@ -193,6 +207,7 @@ class MossTTSLocalDepthTransformer(nn.Module):
         self.h = nn.ModuleList([_MossTTSLocalBlock(self.hidden_size, n_head, inner_size, rope_base, eps)])
         self.ln_f = nn.LayerNorm(self.hidden_size, eps=eps)
         self._compiled_forward_prefix = None
+        self._compiled_audio_sampler = None
 
     def _forward_prefix(
         self,
@@ -216,6 +231,14 @@ class MossTTSLocalDepthTransformer(nn.Module):
             dynamic=True,
             options={"epilogue_fusion": False},
         )
+        if os.environ.get("VLLM_OMNI_MOSS_LOCAL_COMPILE_AUDIO_SAMPLER", "0") == "1":
+            # Keep the existing top-k/top-p algorithm and torch RNG. Explicit
+            # per-request generators use the original helper below. The
+            # binary continue/stop head is deliberately unchanged.
+            self._compiled_audio_sampler = torch.compile(
+                _sample_token, fullgraph=True, dynamic=True, options={"fallback_random": True}
+            )
+            logger.info("MOSS-TTS Local compiled audio-channel sampler enabled")
         logger.info("MOSS-TTS local depth frame-local KV execution enabled with torch.compile")
 
     def _run_prefix(
@@ -239,6 +262,7 @@ class MossTTSLocalDepthTransformer(nn.Module):
         do_sample: bool,
         generator: torch.Generator | None,
         generators: Sequence[torch.Generator | None] | None = None,
+        sample_token: Callable[..., torch.Tensor] | None = None,
     ) -> torch.Tensor:
         """Compute channel logits, apply repetition penalty, sample, store."""
         channel_logits = audio_lm_heads[channel_index](local_hidden).float()
@@ -250,7 +274,7 @@ 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_token = (sample_token or _sample_token)(
             channel_logits,
             temperature,
             top_k,
@@ -351,6 +375,7 @@ class MossTTSLocalDepthTransformer(nn.Module):
                 do_sample=do_sample,
                 generator=generator,
                 generators=generators,
+                sample_token=self._compiled_audio_sampler if generator is None and generators is None else None,
             )
 
             if channel_index + 1 < n_vq:

vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_talker.py

@@ -37,6 +37,7 @@ from vllm_omni.model_executor.models.moss_tts.modeling_moss_tts_local_depth impo
 )
 from vllm_omni.model_executor.models.output_templates import OmniOutput
 from vllm_omni.platforms import current_omni_platform
+from vllm_omni.worker.sampling_utils import get_tts_local_seed
 
 logger = init_logger(__name__)
 
@@ -1675,6 +1676,31 @@ class MossTTSLocalTalkerForGeneration(nn.Module):
         return input_embeds_out, output_codes
 
     # ------------------------------------------------------------------
+    # MRV2 reads model-owned capabilities; retain the V1 hook and platform
+    # graph-safety setting as the canonical implementation.
+    mtp = talker_mtp
+    get_mtp_seed = staticmethod(get_tts_local_seed)
+
+    def create_mrv2_model_state(self, vllm_config, encoder_cache, device):
+        from vllm_omni.worker_v2.model_states.omni_model_state import OmniModelState
+
+        state_cls = OmniModelState
+        if getattr(self.config, "mrv2_gpu_slot_state", False):
+            from .local_model_state import MossLocalModelState
+
+            state_cls = MossLocalModelState
+        return state_cls(vllm_config, self, encoder_cache, device)
+
+    @property
+    def mtp_output_key(self) -> tuple[str, str]:
+        """Return the request-buffer destination for one local audio frame."""
+        return self.talker_mtp_output_key
+
+    @property
+    def mtp_graph_safe(self) -> bool:
+        """Honor the same platform graph-safety override as the V1 runner."""
+        return self.talker_mtp_graph_safe
+
     # Package runner-generated audio frames
     # ------------------------------------------------------------------
 

vllm_omni/model_executor/models/moss_tts/pipeline.py

@@ -97,6 +97,7 @@ MOSS_TTS_LOCAL_PIPELINE = PipelineConfig(
         StagePipelineConfig(
             stage_id=0,
             model_stage="moss_tts_local",
+            supports_native_mrv2_data_plane=True,
             execution_type=StageExecutionType.LLM_AR,
             input_sources=(),
             owns_tokenizer=True,
@@ -114,6 +115,7 @@ MOSS_TTS_LOCAL_PIPELINE = PipelineConfig(
         StagePipelineConfig(
             stage_id=1,
             model_stage="moss_tts_local_codec",
+            supports_native_mrv2_data_plane=True,
             execution_type=StageExecutionType.LLM_GENERATION,
             input_sources=(0,),
             final_output=True,

vllm_omni/model_executor/models/moss_tts/reference_encoder.py

@@ -25,7 +25,10 @@ so importing it from the API-server process does not pull the talker/codec.
 from __future__ import annotations
 
 import asyncio
+import contextlib
 import hashlib
+import os
+from collections import OrderedDict
 from collections.abc import Awaitable, Callable
 from functools import lru_cache
 from typing import Any
@@ -35,22 +38,93 @@ import torch
 from numpy.typing import NDArray
 from vllm.logger import init_logger
 
+from vllm_omni.model_executor.models.moss_tts.reference_encoder_graphs import (
+    FINE_BATCH_SIZES,
+    FINE_BUCKET_SECONDS,
+    MossReferenceEncoderGraphs,
+    install_windowed_attention,
+)
+from vllm_omni.model_executor.models.moss_tts.shared_reference_encoder import (
+    SharedReferenceEncoderClient,
+    SharedReferenceEncoderHost,
+    SharedReferenceEncoderStartupError,
+    elect_host,
+)
+
 logger = init_logger(__name__)
 
 # Coalescing defaults. Kept as module constants (not env/CLI) to match the
 # ``_REF_AUDIO_RESOLVE_CACHE_MAX_*`` convention in serving_speech.py.
 _REF_ENCODE_BATCH_WINDOW_MS = 10.0
-_REF_ENCODE_MAX_BATCH = 8
+_REF_ENCODE_MAX_BATCH = 32
+# Experiment switches for cold-reference latency (see the A/B notes):
+# batch window in ms; batches in flight per API process with a shared encoder;
+# "high" puts reference encodes on the highest-priority CUDA stream.
+_REF_ENCODE_WINDOW_ENV = "VLLM_OMNI_MOSS_REF_BATCH_WINDOW_MS"
+_REF_ENCODE_INFLIGHT_ENV = "VLLM_OMNI_MOSS_REF_INFLIGHT"
+_REF_ENCODE_PRIORITY_ENV = "VLLM_OMNI_MOSS_REF_STREAM_PRIORITY"
+# "1" resamples and loudness-normalizes clips on the encoder's GPU instead of
+# the API process's CPU (not bit-identical to the CPU path).
+_REF_ENCODE_GPU_PREP_ENV = "VLLM_OMNI_MOSS_REF_GPU_PREP"
 
 _INT32_MAX = 2**31
+# Inline (data: URI) references already resolved and encoded, keyed by the
+# serving resolve key. Codes are small; the URI itself is not retained.
+_INLINE_INDEX_MAX_ENTRIES = 4096
+# Optional directory (e.g. under /dev/shm) that API processes of one server
+# share, so an inline reference encoded by any frontend is reused by all.
+_SHARED_CODES_DIR_ENV = "VLLM_OMNI_MOSS_REF_CODES_SHARED_DIR"
+# "0" keeps one reference encoder per API process even with a shared directory.
+_SHARED_ENCODER_ENV = "VLLM_OMNI_MOSS_REF_SHARED_ENCODER"
+_shared_role: str | None = None
+# Concurrent encode workers in the shared host, each with its own CUDA graphs.
+_SHARED_WORKERS_ENV = "VLLM_OMNI_MOSS_REF_ENCODER_WORKERS"
+_SHARED_WORKER_BATCH_SECONDS = 96.0
+
+
+def shared_encoder_role() -> str:
+    """``"host"``, ``"client"`` or ``""`` (no shared encoder) for this process.
+
+    Decided once, before the processor places its tokenizer on the GPU, so
+    client processes can leave it on the CPU.
+    """
+    global _shared_role
+    if _shared_role is None:
+        shared_dir = os.environ.get(_SHARED_CODES_DIR_ENV)
+        if not shared_dir or os.environ.get(_SHARED_ENCODER_ENV, "1") == "0":
+            _shared_role = ""
+        else:
+            os.makedirs(shared_dir, exist_ok=True)
+            _shared_role = "host" if elect_host(shared_dir) else "client"
+    return _shared_role
+
+
+# Opt in to reference graphs; existing deployments keep the eager encoder.
+_REF_ENCODE_GRAPHS_ENV = "VLLM_OMNI_MOSS_REF_GRAPHS"
+
+
+def reference_graphs_enabled() -> bool:
+    return os.environ.get(_REF_ENCODE_GRAPHS_ENV, "0") == "1"
+
+
+def reference_stream_priority() -> int:
+    """CUDA stream priority for reference encodes (lower is more urgent)."""
+    if os.environ.get(_REF_ENCODE_PRIORITY_ENV, "") == "high" and torch.cuda.is_available():
+        return torch.cuda.Stream.priority_range()[1]
+    return 0
+
+
+# "sdpa" keeps the tokenizer's dense masked SDPA attention. By default the
+# reference encoder attends through a local-window flash kernel, the backend
+# the tokenizer config selects (flash_attention_2) when flash-attn is present.
+_REF_ENCODE_ATTN_ENV = "VLLM_OMNI_MOSS_REF_ATTN"
+# "0" captures the encoder's eager kernels. By default the captured encoder is
+# compiled first, which fuses its many small elementwise kernels.
+_REF_ENCODE_COMPILE_ENV = "VLLM_OMNI_MOSS_REF_COMPILE"
 
 _Waveform = NDArray[np.float32] | list[float]
 
 
-def _sha1(s: str) -> str:
-    return hashlib.sha1((s or "").encode("utf-8")).hexdigest()
-
-
 @lru_cache(maxsize=16)
 def _reference_resampler(sr: int, sr_target: int) -> torch.nn.Module:
     import torchaudio
@@ -72,6 +146,13 @@ def _prep_wav_sync(waveform: _Waveform, sr: int, sr_target: int) -> torch.Tensor
     return wav
 
 
+def _inline_reference_key(ref_str: str) -> str | None:
+    """SHA-1 of a ``data:`` URI, matching the serving resolve key for it."""
+    if not isinstance(ref_str, str) or ref_str[:5].lower() != "data:":
+        return None
+    return hashlib.sha1(ref_str.encode("utf-8")).hexdigest()
+
+
 def _to_compact_codes(codes: torch.Tensor) -> torch.Tensor:
     """Downcast RVQ codes to int32 for compact caching."""
     codes = codes.detach().cpu().contiguous()
@@ -119,18 +200,26 @@ class _RefEncodeBatcher:
         *,
         window_ms: float,
         max_batch: int,
+        max_inflight: int = 1,
     ):
         # encode_batch_fn: sync, takes [(waveform, sr), ...], returns a list of
         # (codes_tensor | Exception) aligned to the input order.
         self._encode_batch_fn = encode_batch_fn
         self._window_s = max(0.0, float(window_ms) / 1000.0)
         self._max_batch = max(1, int(max_batch))
+        # Batches encoded at once. Clips that arrive while every slot is busy
+        # are queued and form the next batch.
+        self._max_inflight = max(1, int(max_inflight))
         self._queue: asyncio.Queue | None = None
+        self._slots: asyncio.Semaphore | None = None
+        self._inflight: set[asyncio.Task] = set()
         self._drainer: asyncio.Task | None = None
 
     def _ensure_started(self) -> None:
         if self._queue is None:
             self._queue = asyncio.Queue()
+        if self._slots is None:
+            self._slots = asyncio.Semaphore(self._max_inflight)
         if self._drainer is None or self._drainer.done():
             self._drainer = asyncio.create_task(self._drain_loop())
 
@@ -141,16 +230,27 @@ class _RefEncodeBatcher:
         return await fut
 
     async def _drain_loop(self) -> None:
-        assert self._queue is not None
+        assert self._queue is not None and self._slots is not None
         while True:
-            first = await self._queue.get()
-            jobs = None
+            await self._slots.acquire()
             try:
-                jobs = await self._coalesce(first)
-                await self._run_batch(jobs)
-            finally:
-                # Do not retain the completed batch while waiting for new work.
-                del first, jobs
+                jobs = await self._coalesce(await self._queue.get())
+            except BaseException:
+                self._slots.release()
+                raise
+            task = asyncio.create_task(self._run_and_release(jobs))
+            self._inflight.add(task)
+            task.add_done_callback(self._inflight.discard)
+            # Do not retain the batch here while waiting for new work.
+            del jobs
+
+    async def _run_and_release(self, jobs: list[tuple[_Waveform, int, asyncio.Future]]) -> None:
+        try:
+            await self._run_batch(jobs)
+        finally:
+            del jobs
+            assert self._slots is not None
+            self._slots.release()
 
     async def _coalesce(self, first: tuple) -> list[tuple]:
         """Group ``first`` with jobs arriving within the batch window."""
@@ -200,6 +300,8 @@ class _RefEncodeBatcher:
             except (asyncio.CancelledError, Exception):  # noqa: BLE001
                 pass
             self._drainer = None
+        for task in list(self._inflight):
+            task.cancel()
 
 
 class MossReferenceEncoder:
@@ -213,7 +315,7 @@ class MossReferenceEncoder:
         n_vq: int,
         sr_target: int,
         speaker_cache: Any,
-        window_ms: float = _REF_ENCODE_BATCH_WINDOW_MS,
+        window_ms: float | None = None,
         max_batch: int = _REF_ENCODE_MAX_BATCH,
     ):
         self._processor = processor
@@ -224,8 +326,33 @@ class MossReferenceEncoder:
         # model_type namespaces the whole family so a moss_tts server never
         # collides with another model's speaker-cache entries.
         self._model_type = f"moss_tts_{variant}_nq{int(n_vq)}"
-        self._inflight: dict[str, asyncio.Task] = {}
-        self._batcher = _RefEncodeBatcher(self._encode_batch_sync, window_ms=window_ms, max_batch=max_batch)
+        self._inflight: dict[str | tuple[str, str], asyncio.Task] = {}
+        self._inline_index: OrderedDict[str, torch.Tensor] = OrderedDict()
+        self._shared_codes_dir = os.environ.get(_SHARED_CODES_DIR_ENV) or None
+        if self._shared_codes_dir:
+            os.makedirs(self._shared_codes_dir, exist_ok=True)
+        # Realtime encodes with its own codec and never shares the encoder.
+        self._shares_encoder = bool(self._shared_codes_dir) and type(self)._encode_prepared is (
+            MossReferenceEncoder._encode_prepared
+        )
+        self._shared_host: SharedReferenceEncoderHost | None = None
+        self._shared_client: SharedReferenceEncoderClient | None = None
+        self._graphs: MossReferenceEncoderGraphs | None = None
+        self._attention_installed = False
+        self._graphs_enabled = reference_graphs_enabled()
+        if window_ms is None:
+            window_ms = float(os.environ.get(_REF_ENCODE_WINDOW_ENV, _REF_ENCODE_BATCH_WINDOW_MS))
+        # A shared encoder serves several batches at once (one per host
+        # worker); a process's own graphs replay one batch at a time.
+        inflight = 1
+        if self._shares_encoder and os.environ.get(_SHARED_ENCODER_ENV, "1") != "0":
+            inflight = max(1, int(os.environ.get(_REF_ENCODE_INFLIGHT_ENV, "1")))
+        self._local_stream: torch.cuda.Stream | None = None
+        self._gpu_prep = os.environ.get(_REF_ENCODE_GPU_PREP_ENV, "0") == "1"
+        self._gpu_resamplers: dict[torch.device, torch.nn.Module] = {}
+        self._batcher = _RefEncodeBatcher(
+            self._encode_batch_sync, window_ms=window_ms, max_batch=max_batch, max_inflight=inflight
+        )
 
     def _make_cache_key(self, name: str, created_at: int) -> tuple:
         return self._speaker_cache.make_cache_key(name, model_type=self._model_type, created_at=int(created_at))
@@ -251,6 +378,7 @@ class MossReferenceEncoder:
         caller salts those requests with ``voice_created_at`` instead).
         """
         voice_name, created_at = _registered_voice(voice_name, voice_created_at)
+        inline_key: str | None = None
 
         if voice_name:
             # A named voice has a stable key that does not depend on the
@@ -260,6 +388,16 @@ class MossReferenceEncoder:
             if cached is not None:
                 return _clone_out(cached["codes"]), None
         else:
+            inline_key = _inline_reference_key(ref_str)
+            if inline_key is not None:
+                codes = self._inline_index.get(inline_key)
+                if codes is None:
+                    codes = self._load_shared_codes(inline_key)
+                if codes is not None:
+                    # A data: URI is immutable content, and its serving
+                    # resolve key is this same digest: skip the resolve hop.
+                    self._inline_index.move_to_end(inline_key)
+                    return _clone_out(codes), inline_key
             # Anonymous refs have no pre-resolve hot path: the content key must
             # come from the resolve itself (mtime/size from a single stat) so an
             # on-disk edit invalidates the cached codes; the flight body
@@ -267,17 +405,55 @@ class MossReferenceEncoder:
             # when the resolve cache is warm. The flight key is the request-side
             # reference (not the content hash), so concurrent requests for the
             # same ref_str also share the resolve/download, not just the encode.
-            flight_key = "ref:" + _sha1(ref_str)
+            # Key by the locator itself: equality is exact, and hashing a
+            # large inline data URI once per request object is much cheaper
+            # than a separate SHA-1 of its UTF-8 copy on the event loop.
+            flight_key = ("ref", ref_str)
 
         codes, resolve_key = await self._single_flight(
             flight_key,
             lambda: self._resolve_and_encode(ref_str, resolve_ref_audio, get_artifact_key, voice_name, created_at),
         )
+        if not voice_name and inline_key is not None and resolve_key == inline_key:
+            # Only index when the resolver's key is exactly this digest.
+            self._remember_inline(inline_key, codes)
+            self._store_shared_codes(inline_key, codes)
         return _clone_out(codes), resolve_key
 
+    def _remember_inline(self, key: str, codes: torch.Tensor) -> None:
+        self._inline_index[key] = codes
+        self._inline_index.move_to_end(key)
+        while len(self._inline_index) > _INLINE_INDEX_MAX_ENTRIES:
+            self._inline_index.popitem(last=False)
+
+    def _load_shared_codes(self, key: str) -> torch.Tensor | None:
+        if not self._shared_codes_dir:
+            return None
+        try:
+            array = np.load(os.path.join(self._shared_codes_dir, key + ".npy"), allow_pickle=False)
+        except (OSError, ValueError):
+            return None
+        if array.ndim != 2 or array.shape[1] != self._n_vq or array.dtype not in (np.int32, np.int64):
+            return None
+        codes = torch.from_numpy(array)
+        self._remember_inline(key, codes)
+        return codes
+
+    def _store_shared_codes(self, key: str, codes: torch.Tensor) -> None:
+        if not self._shared_codes_dir:
+            return
+        path = os.path.join(self._shared_codes_dir, key + ".npy")
+        tmp = f"{path}.{os.getpid()}.tmp"
+        try:
+            with open(tmp, "wb") as handle:
+                np.save(handle, codes.detach().cpu().numpy(), allow_pickle=False)
+            os.replace(tmp, path)
+        except OSError:
+            logger.warning("MOSS ref codes: could not write shared entry %s", path, exc_info=True)
+
     async def _single_flight(
         self,
-        flight_key: str,
+        flight_key: str | tuple[str, str],
         start_flight: Callable[[], Awaitable[tuple[torch.Tensor, str]]],
     ) -> tuple[torch.Tensor, str]:
         """Join the in-flight encode for ``flight_key``, starting one if absent.
@@ -300,7 +476,7 @@ class MossReferenceEncoder:
         # a duplicate resolve/encode instead of joining this one. Identity
         # guard: only drop the slot if it still holds *our* task (a later
         # request may have replaced it after ours completed).
-        def _retire(t: asyncio.Task, key: str = flight_key) -> None:
+        def _retire(t: asyncio.Task, key: str | tuple[str, str] = flight_key) -> None:
             if self._inflight.get(key) is t:
                 self._inflight.pop(key, None)
 
@@ -350,29 +526,228 @@ class MossReferenceEncoder:
         prepared_idx: list[int] = []
         for i, (waveform, sr) in enumerate(payload):
             try:
-                prepared.append(_prep_wav_sync(waveform, sr, self._sr_target))
+                prepared.append(self._encoder_input(_prep_wav_sync(waveform, sr, self._sr_target)))
                 prepared_idx.append(i)
             except Exception as exc:  # noqa: BLE001 — isolate this clip's failure
                 results[i] = exc
         if not prepared:
             return results
 
-        try:
-            codes_list = self._encode_prepared(prepared)
-            for local_i, orig_i in enumerate(prepared_idx):
-                results[orig_i] = codes_list[local_i]
-        except Exception:  # noqa: BLE001 — batch forward failed; retry item-by-item
-            logger.warning("MOSS ref batch encode (n=%d) failed; falling back to per-item", len(prepared))
-            for local_i, orig_i in enumerate(prepared_idx):
-                try:
-                    results[orig_i] = self._encode_prepared([prepared[local_i]])[0]
-                except Exception as exc:  # noqa: BLE001 — isolate this clip's failure
-                    results[orig_i] = exc
+        for orig_i, codes in zip(prepared_idx, self._encode_shared_or_local(prepared)):
+            results[orig_i] = codes
         return results
 
-    def _encode_prepared(self, prepared: list[torch.Tensor]) -> list[torch.Tensor]:
-        with torch.no_grad():
-            return self._processor.encode_audios_from_wav(prepared, sampling_rate=self._sr_target, n_vq=self._n_vq)
+    def prepare(self) -> None:
+        """Capture the CUDA graphs (or start the shared encoder) before the first request."""
+        role = shared_encoder_role() if self._shares_encoder else ""
+        if role == "host":
+            self._host().wait_until_ready()
+        elif role == "client":
+            if self._shared_client is None:
+                self._shared_client = SharedReferenceEncoderClient(self._shared_codes_dir)
+            self._shared_client.wait_until_ready()
+        elif role == "":
+            self._reference_graphs()
+
+    def _host(self) -> SharedReferenceEncoderHost:
+        if self._shared_host is None:
+            # Graph capture runs in the background; clips queue until it ends.
+            self._shared_host = SharedReferenceEncoderHost(
+                self._shared_codes_dir,
+                self._host_workers(),
+                stream_priority=reference_stream_priority(),
+                batch_window_ms=float(os.environ.get("VLLM_OMNI_MOSS_REF_HOST_WINDOW_MS", "0")),
+            )
+        return self._shared_host
+
+    def _encode_shared_or_local(self, prepared: list[torch.Tensor]) -> list:
+        role = shared_encoder_role() if self._shares_encoder else ""
+        if role == "host":
+            return self._host().encode(prepared)
+        if role == "client":
+            if self._shared_client is None:
+                self._shared_client = SharedReferenceEncoderClient(self._shared_codes_dir)
+            try:
+                return self._shared_client.encode(prepared)
+            except SharedReferenceEncoderStartupError:
+                raise
+            except Exception:  # noqa: BLE001 — the host is unreachable: encode here
+                logger.warning("MOSS shared reference encoder unavailable; encoding locally", exc_info=True)
+                self._place_tokenizer_for_local_encode()
+        with self._local_stream_context():
+            return self._encode_local(prepared)
+
+    def _local_stream_context(self):
+        tokenizer = getattr(self._processor, "audio_tokenizer", None)
+        params = next(tokenizer.parameters(), None) if isinstance(tokenizer, torch.nn.Module) else None
+        if params is None or params.device.type != "cuda":
+            return contextlib.nullcontext()
+        if self._local_stream is None:
+            self._local_stream = torch.cuda.Stream(device=params.device, priority=reference_stream_priority())
+        return torch.cuda.stream(self._local_stream)
+
+    def _host_workers(self) -> list:
+        """Encode workers of the shared host, each with its own graphs and stream."""
+        count = max(1, int(os.environ.get(_SHARED_WORKERS_ENV, "4")))
+        if not self._graph_capable():
+            return [(self._encode_local, None)] * count
+        workers = []
+        for _ in range(count):
+            graphs = MossReferenceEncoderGraphs(
+                self._processor.audio_tokenizer,
+                n_vq=self._n_vq,
+                batch_sizes=FINE_BATCH_SIZES,
+                bucket_seconds=FINE_BUCKET_SECONDS,
+                max_batch_seconds=_SHARED_WORKER_BATCH_SECONDS,
+                compile_core=self._compile_graphs(),
+                batched_transfer=os.environ.get("VLLM_OMNI_MOSS_REF_BATCHED_TRANSFER", "0") == "1",
+                singleton_bucket_seconds=self._singleton_bucket_seconds(),
+            )
+            workers.append(((lambda wavs, g=graphs: self._encode_local(wavs, g)), graphs.capture))
+        return workers
+
+    def _encode_local(self, prepared: list[torch.Tensor], graphs: MossReferenceEncoderGraphs | None = None) -> list:
+        """One batched forward, retried per clip so one bad clip fails alone."""
+        try:
+            return list(self._encode_prepared(prepared, graphs))
+        except Exception:  # noqa: BLE001 — batch forward failed; retry item-by-item
+            logger.warning("MOSS ref batch encode (n=%d) failed; falling back to per-item", len(prepared))
+        results: list = []
+        for wav in prepared:
+            try:
+                results.append(self._encode_prepared([wav], graphs)[0])
+            except Exception as exc:  # noqa: BLE001 — isolate this clip's failure
+                results.append(exc)
+        return results
+
+    def _place_tokenizer_for_local_encode(self) -> None:
+        """A client process keeps the tokenizer off the GPU until it must encode."""
+        tokenizer = getattr(self._processor, "audio_tokenizer", None)
+        device = getattr(self._processor, "_vllm_omni_ref_encoder_device", None)
+        if tokenizer is not None and device is not None and next(tokenizer.parameters()).device != device:
+            self._processor.audio_tokenizer = tokenizer.to(device)
+
+    def _encode_prepared(
+        self, prepared: list[torch.Tensor], graphs: MossReferenceEncoderGraphs | None = None
+    ) -> list[torch.Tensor]:
+        # ``prepared`` clips come from _encoder_input, at ``_clip_rate()``.
+        graphs = graphs or self._reference_graphs()
+        if graphs is None:
+            with torch.no_grad():
+                return self._processor.encode_audios_from_wav(
+                    prepared, sampling_rate=self._clip_rate(), n_vq=self._n_vq
+                )
+        codes = graphs.encode([self._finish_prepare(wav) for wav in prepared])
+        missing = [i for i, c in enumerate(codes) if c is None]
+        if missing:
+            with torch.no_grad():
+                eager = self._processor.encode_audios_from_wav(
+                    [prepared[i] for i in missing], sampling_rate=self._clip_rate(), n_vq=self._n_vq
+                )
+            for i, c in zip(missing, eager):
+                codes[i] = c
+        return codes
+
+    def _graph_capable(self) -> bool:
+        """True when the tokenizer can run graphs; installs windowed attention once."""
+        if not self._graphs_enabled:
+            return False
+        tokenizer = getattr(self._processor, "audio_tokenizer", None)
+        device = next(tokenizer.parameters()).device if tokenizer is not None else None
+        if device is None or device.type != "cuda" or not hasattr(tokenizer, "_codec_inference_autocast"):
+            self._graphs_enabled = False
+            return False
+        if not self._attention_installed and os.environ.get(_REF_ENCODE_ATTN_ENV, "flash") == "flash":
+            logger.info(
+                "MOSS reference encoder: %d attention layers use local-window flash attention",
+                install_windowed_attention(
+                    tokenizer,
+                    fa_version=int(os.environ.get("VLLM_OMNI_MOSS_REF_FA_VERSION", "2")),
+                    cache_rope=os.environ.get("VLLM_OMNI_MOSS_REF_CACHE_ROPE", "0") == "1",
+                    skip_padded_query_mask=os.environ.get("VLLM_OMNI_MOSS_REF_SKIP_PADDED_QUERY_MASK", "0") == "1",
+                ),
+            )
+        self._attention_installed = True
+        return True
+
+    @staticmethod
+    def _compile_graphs() -> bool:
+        return os.environ.get(_REF_ENCODE_COMPILE_ENV, "1") != "0"
+
+    @staticmethod
+    def _singleton_bucket_seconds() -> tuple[float, ...] | None:
+        step_ms = int(os.environ.get("VLLM_OMNI_MOSS_REF_SINGLETON_BUCKET_MS", "0"))
+        if step_ms == 0:
+            return None
+        if not 80 <= step_ms <= 1000:
+            raise ValueError("MOSS reference singleton bucket step must be 80..1000 ms, or 0 to disable")
+        return tuple(ms / 1000 for ms in range(3000, 10001, step_ms))
+
+    def _reference_graphs(self) -> MossReferenceEncoderGraphs | None:
+        """Capture on first use, on the encode worker thread that replays them."""
+        if self._graphs is None:
+            if not self._graph_capable():
+                return None
+            graphs = MossReferenceEncoderGraphs(
+                self._processor.audio_tokenizer,
+                n_vq=self._n_vq,
+                compile_core=self._compile_graphs(),
+                batched_transfer=os.environ.get("VLLM_OMNI_MOSS_REF_BATCHED_TRANSFER", "0") == "1",
+                singleton_bucket_seconds=self._singleton_bucket_seconds(),
+            )
+            graphs.capture()
+            self._graphs = graphs
+        return self._graphs
+
+    def _tokenizer_rate(self) -> int:
+        return int(getattr(getattr(self._processor, "model_config", None), "sampling_rate", self._sr_target))
+
+    def _encoder_input(self, wav: torch.Tensor) -> torch.Tensor:
+        """A prepped clip at the tokenizer rate, one or two channels.
+
+        This is the costly half of the processor's per-clip preparation; it runs
+        in the process that received the request, even with a shared encoder.
+        """
+        import torchaudio
+
+        if wav.ndim == 1:
+            wav = wav.unsqueeze(0)
+        if wav.shape[0] > 2:
+            wav = wav[:2]
+        if self._clip_rate() != self._sr_target:
+            # Channels resample independently, so a mono clip is resampled
+            # once before it is duplicated; the result is the same.
+            wav = torchaudio.functional.resample(wav, self._sr_target, self._tokenizer_rate())
+        return wav
+
+    def _clip_rate(self) -> int:
+        """Sample rate of the clips _encoder_input hands to the encoder."""
+        return self._sr_target if self._gpu_prep else self._tokenizer_rate()
+
+    def _finish_prepare(self, wav: torch.Tensor) -> torch.Tensor:
+        """The rest of the processor's preparation: two channels, normalized loudness."""
+        if self._gpu_prep:
+            device = next(self._processor.audio_tokenizer.parameters()).device
+            wav = wav.to(device)
+            if self._clip_rate() != self._tokenizer_rate():
+                wav = self._gpu_resampler(device)(wav)
+        if wav.shape[0] == 1:
+            wav = wav.repeat(2, 1)
+        return self._processor.loudness_normalize(wav)
+
+    def _gpu_resampler(self, device: torch.device) -> torch.nn.Module:
+        resampler = self._gpu_resamplers.get(device)
+        if resampler is None:
+            import torchaudio
+
+            # The same windowed-sinc filter as torchaudio.functional.resample's defaults.
+            resampler = torchaudio.transforms.Resample(self._sr_target, self._tokenizer_rate()).to(device)
+            self._gpu_resamplers[device] = resampler
+        return resampler
+
+    def _processor_prepare(self, wav: torch.Tensor) -> torch.Tensor:
+        """Per-clip preparation of the processor's ``encode_audios_from_wav``."""
+        return self._finish_prepare(self._encoder_input(wav))
 
     async def aclose(self) -> None:
         """Release the batcher drainer + any in-flight encodes (tests/shutdown)."""
@@ -385,7 +760,10 @@ class MossReferenceEncoder:
 class MossRealtimeReferenceEncoder(MossReferenceEncoder):
     """Use the realtime codec with the shared reference cache and batcher."""
 
-    def _encode_prepared(self, prepared: list[torch.Tensor]) -> list[torch.Tensor]:
+    def _encoder_input(self, wav: torch.Tensor) -> torch.Tensor:
+        return wav  # the realtime codec takes the clip at the working rate
+
+    def _encode_prepared(self, prepared: list[torch.Tensor], graphs: Any = None) -> list[torch.Tensor]:
         with torch.no_grad():
             encoded = self._processor.batch_encode([wav.squeeze(0) for wav in prepared], num_quantizers=self._n_vq)
         return [

vllm_omni/model_executor/models/moss_tts/reference_encoder_graphs.py

@@ -0,0 +1,399 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""CUDA graphs for batched MOSS reference-audio encoding.
+
+Reference clips are encoded by the upstream MOSS audio tokenizer that the
+processor loads. Its eager forward is dominated by kernel launches for the
+small batches that serving produces. This runner replays graphs of the same
+encoder and quantizer modules, captured per ``(batch, length)`` bucket. For a
+given padded input the codes are bit-identical to the tokenizer's own
+``_encode_frame``. Only the per-clip crop, which needs host values, runs
+outside the graph.
+
+The encoder is causal, so padding a clip to its length bucket leaves its valid
+frames unchanged up to kernel selection for the padded shape, the same effect
+that batching already has when the tokenizer pads a batch to its longest clip.
+"""
+
+from __future__ import annotations
+
+import math
+import types
+from dataclasses import dataclass
+
+import torch
+from vllm.logger import init_logger
+
+logger = init_logger(__name__)
+
+DEFAULT_BATCH_SIZES = (1, 2, 4, 8, 16, 32)
+DEFAULT_BUCKET_SECONDS = (4.0, 6.0, 8.0, 10.0, 12.0, 16.0, 20.0)
+# Fine steps keep padding small: a padded row or padded time is encoded in
+# full. Worth the extra graphs where several encoders share one GPU.
+FINE_BATCH_SIZES = (1, 2, 3, 4, 5, 6, 8)
+FINE_BUCKET_SECONDS = (3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 10.0, 12.0, 14.0, 16.0, 20.0)
+# Graphs are captured only while batch * bucket stays within this much audio,
+# which bounds the shared graph memory pool by its largest batch.
+DEFAULT_MAX_BATCH_SECONDS = 256.0
+
+
+@torch.compile(dynamic=False)
+def _reference_rope_factors(seqlen: int, head_dim: int, max_period: float, device: torch.device):
+    # Use the same compiled FP32 arithmetic as the encoder. Eager libdevice
+    # sin/cos and compiled sin/cos can round differently near a VQ boundary.
+    ds = torch.arange(head_dim // 2, device=device, dtype=torch.float32)
+    freqs = torch.exp(ds * (-math.log(max_period) * 2 / head_dim))
+    angles = torch.arange(seqlen, device=device, dtype=torch.float32)[:, None] * freqs[None, :]
+    return torch.cos(angles), torch.sin(angles)
+
+
+def _cached_reference_rope(self, q: torch.Tensor, k: torch.Tensor, cache: dict):
+    if self.rope is None:
+        return q, k
+    batch, heads, seqlen, dim = q.shape
+    key = (dim, self.rope.max_period, q.device)
+    if key not in cache or cache[key][0].shape[0] < seqlen:
+        cache[key] = _reference_rope_factors(seqlen, dim, self.rope.max_period, q.device)
+    cos, sin = cache[key]
+    cos, sin = cos[:seqlen], sin[:seqlen]
+    qr, qi = q.view(batch, heads, seqlen, dim // 2, 2).unbind(-1)
+    kr, ki = k.view(batch, heads, seqlen, dim // 2, 2).unbind(-1)
+    qr, qi, kr, ki = qr.float(), qi.float(), kr.float(), ki.float()
+    # Preserve the original arithmetic, casts, and B/H/T/D output layout.
+    qo = torch.stack(((qr * cos - qi * sin).to(q.dtype), (qr * sin + qi * cos).to(q.dtype)), dim=-1)
+    ko = torch.stack(((kr * cos - ki * sin).to(k.dtype), (kr * sin + ki * cos).to(k.dtype)), dim=-1)
+    return qo.view(batch, heads, seqlen, dim), ko.view(batch, heads, seqlen, dim)
+
+
+def _windowed_attention(
+    self,
+    x: torch.Tensor,
+    input_lengths: torch.Tensor,
+    *,
+    sdpa,
+    fa_version: int = 2,
+    rope_cache: dict | None = None,
+    skip_padded_query_mask: bool = False,
+) -> torch.Tensor:
+    """The tokenizer's non-streaming attention with a local-window flash kernel.
+
+    Same projection, RoPE and padded-row zeroing as the tokenizer's SDPA path.
+    That path masks each query to ``context`` keys (distance 0 .. context-1)
+    inside a dense block; the flash kernel visits only that window.
+    """
+    batch, seqlen, _ = x.shape
+    q, k, v = self._project_qkv(x)
+    q, k = self._apply_dense_rope(q, k) if rope_cache is None else _cached_reference_rope(self, q, k, rope_cache)
+    if q.dtype not in (torch.bfloat16, torch.float16):
+        return sdpa(x, input_lengths)
+    from vllm.vllm_flash_attn import flash_attn_varlen_func
+
+    heads, dim = q.shape[1], q.shape[-1]
+
+    def packed(t: torch.Tensor) -> torch.Tensor:
+        return t.transpose(1, 2).reshape(batch * seqlen, heads, dim)
+
+    starts = torch.arange(0, (batch + 1) * seqlen, seqlen, device=x.device, dtype=torch.int32)
+    window = [self.context - 1, 0] if self.causal and self.context is not None else [-1, -1]
+    out = flash_attn_varlen_func(
+        packed(q),
+        packed(k),
+        packed(v),
+        max_seqlen_q=seqlen,
+        cu_seqlens_q=starts,
+        max_seqlen_k=seqlen,
+        cu_seqlens_k=starts,
+        causal=self.causal,
+        window_size=window,
+        fa_version=fa_version,
+    ).view(batch, seqlen, heads, dim)
+    if not skip_padded_query_mask:
+        valid = (torch.arange(seqlen, device=x.device).view(1, seqlen) < input_lengths.view(-1, 1)).view(
+            batch, seqlen, 1, 1
+        )
+        out = torch.where(valid, out, torch.zeros((), device=out.device, dtype=out.dtype))
+    return out.reshape(batch, seqlen, self.embed_dim)
+
+
+def _validate_unmasked_reference_encoder(tokenizer: torch.nn.Module) -> None:
+    """Restrict the optimization to the audited HF encoder's valid-prefix contract.
+
+    Causal attention and pointwise transformer operations cannot read future
+    padded rows. Its patch-downsample floors lengths, so every valid output
+    group consists entirely of valid input rows. The quantizer is pointwise
+    and the caller crops codes by those lengths. Padded *outputs* may differ.
+    This argument does not apply to the streaming decoder or to upsampling.
+    """
+    transformers = 0
+    for module in tokenizer.encoder:
+        name = type(module).__name__
+        if name == "MossAudioTokenizerPatchedPretransform" and module.is_downsample:
+            continue
+        if name == "MossAudioTokenizerProjectedTransformer":
+            layers = module.transformer.layers
+            if layers and all(layer.self_attn.causal for layer in layers):
+                transformers += 1
+                continue
+        raise ValueError("Unmasked reference queries require the causal MOSS patch-downsample encoder")
+    if not transformers:
+        raise ValueError("Unmasked reference queries require at least one causal MOSS transformer")
+    if type(tokenizer.quantizer).__name__ not in {
+        "MossAudioTokenizerResidualVQ",
+        "MossAudioTokenizerResidualLFQ",
+    }:
+        raise ValueError("Unmasked reference queries require the pointwise MOSS quantizer")
+    for module in tokenizer.quantizer.modules():
+        if isinstance(module, torch.nn.Conv1d) and module.kernel_size != (1,):
+            raise ValueError("Unmasked reference queries do not support temporal quantizer convolutions")
+
+
+def install_windowed_attention(
+    tokenizer: torch.nn.Module,
+    *,
+    fa_version: int = 2,
+    cache_rope: bool = False,
+    skip_padded_query_mask: bool = False,
+) -> int:
+    """Route the encoder's non-streaming attention through a local-window flash kernel."""
+    if skip_padded_query_mask:
+        _validate_unmasked_reference_encoder(tokenizer)
+        logger.info("MOSS reference encoder: skipping padded query masks (valid-prefix outputs only)")
+    count = 0
+    rope_cache = {} if cache_rope else None
+    if cache_rope:
+        tokenizer._moss_reference_rope_cache = rope_cache
+        tokenizer._moss_reference_rope_prepared_samples = 0
+    for module in tokenizer.encoder.modules():
+        sdpa = getattr(module, "_forward_non_streaming_sdpa", None)
+        if sdpa is None or not hasattr(module, "_project_qkv") or not getattr(module, "causal", False):
+            continue
+
+        def attention(self, x, input_lengths, _sdpa=sdpa):
+            return _windowed_attention(
+                self,
+                x,
+                input_lengths,
+                sdpa=_sdpa,
+                fa_version=fa_version,
+                rope_cache=rope_cache,
+                skip_padded_query_mask=skip_padded_query_mask,
+            )
+
+        module._forward_non_streaming_sdpa = types.MethodType(attention, module)
+        count += 1
+    return count
+
+
+@dataclass
+class _Graph:
+    graph: torch.cuda.CUDAGraph
+    audio: torch.Tensor  # (B, C, T) float32 input
+    lengths: torch.Tensor  # (B,) int64 input, in samples per channel
+    codes: torch.Tensor  # (n_vq, B, frames) output
+    code_lengths: torch.Tensor  # (B,) output
+
+
+class MossReferenceEncoderGraphs:
+    """Encode prepared reference waveforms with captured tokenizer graphs."""
+
+    def __init__(
+        self,
+        tokenizer: torch.nn.Module,
+        *,
+        n_vq: int,
+        batch_sizes: tuple[int, ...] = DEFAULT_BATCH_SIZES,
+        bucket_seconds: tuple[float, ...] = DEFAULT_BUCKET_SECONDS,
+        max_batch_seconds: float = DEFAULT_MAX_BATCH_SECONDS,
+        compile_core: bool = False,
+        batched_transfer: bool = False,
+        singleton_bucket_seconds: tuple[float, ...] | None = None,
+    ) -> None:
+        self._tokenizer = tokenizer
+        self._compile_core = compile_core
+        self._batched_transfer = batched_transfer
+        self._n_vq = int(n_vq)
+        self._device = next(tokenizer.parameters()).device
+        self._channels = int(tokenizer.number_channels)
+        frame = int(tokenizer.downsample_rate)
+        sampling_rate = int(tokenizer.sampling_rate)
+        # Whole frames, so the tokenizer's own frame padding is a no-op.
+        self._buckets = sorted({int(math.ceil(s * sampling_rate / frame)) * frame for s in bucket_seconds})
+        self._singleton_buckets = sorted(
+            set(self._buckets)
+            | {int(math.ceil(s * sampling_rate / frame)) * frame for s in (singleton_bucket_seconds or ())}
+        )
+        self._batch_sizes = sorted({int(b) for b in batch_sizes if int(b) > 0})
+        self._max_batch_samples = float(max_batch_seconds) * sampling_rate
+        self._graphs: dict[tuple[int, int], _Graph] = {}
+        self._pool = None
+        self._transfer_buffers = None
+
+    @property
+    def max_batch(self) -> int:
+        return self._batch_sizes[-1]
+
+    @property
+    def captured(self) -> list[tuple[int, int]]:
+        return sorted(self._graphs)
+
+    def _core(self, audio: torch.Tensor, lengths: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        # Mirrors the tokenizer's _encode_frame without its host-side crop.
+        tokenizer = self._tokenizer
+        values, value_lengths = tokenizer._flatten_channels_for_codec(audio, lengths)
+        with tokenizer._codec_inference_autocast():
+            hidden, hidden_lengths = values, value_lengths
+            for module in tokenizer.encoder:
+                hidden, hidden_lengths = module(hidden, hidden_lengths)
+        _, codes, code_lengths = tokenizer.quantizer(hidden.float(), hidden_lengths, self._n_vq)
+        return codes, code_lengths
+
+    @torch.no_grad()
+    def capture(self) -> None:
+        if self._device.type != "cuda" or self._graphs:
+            return
+        self._pool = torch.cuda.graph_pool_handle()
+        stream = torch.cuda.Stream(device=self._device)
+        prepared = getattr(self._tokenizer, "_moss_reference_rope_prepared_samples", None)
+        max_samples = max(self._singleton_buckets)
+        if prepared is not None and prepared < max_samples:
+            # Prime immutable FP32 factors at the largest length before any
+            # capture or compilation. Smaller lengths use a prefix of the
+            # same table; the dictionary stays fixed throughout capture and
+            # workers share tables without introducing dynamic-shape guards
+            # for a separate dictionary key at each sequence length.
+            audio = torch.zeros(1, self._channels, max_samples, device=self._device)
+            lengths = torch.full((1,), max_samples, dtype=torch.long, device=self._device)
+            self._core(audio, lengths)
+            self._tokenizer._moss_reference_rope_prepared_samples = max_samples
+        core = self._core
+        if self._compile_core:
+            # Fuses the encoder's many small elementwise kernels (casts, RoPE,
+            # residuals). Only capture runs the compiled code; replays do not.
+            core = torch.compile(self._core, dynamic=True)
+        # Each encoder stage is its own module shape: allow a compile per stage.
+        with torch._dynamo.config.patch(recompile_limit=64):
+            self._capture_all(core, stream)
+        if self._batched_transfer and self._graphs:
+            # One reusable slab per worker, rather than one per graph bucket.
+            entries = list(self._graphs.values())
+            self._transfer_buffers = (
+                torch.empty(max(e.audio.numel() for e in entries), dtype=entries[0].audio.dtype, pin_memory=True),
+                torch.empty(max(e.lengths.numel() for e in entries), dtype=entries[0].lengths.dtype, pin_memory=True),
+                torch.empty(max(e.codes.numel() for e in entries), dtype=entries[0].codes.dtype, pin_memory=True),
+                torch.empty(
+                    max(e.code_lengths.numel() for e in entries), dtype=entries[0].code_lengths.dtype, pin_memory=True
+                ),
+                torch.cuda.Event(),
+            )
+        logger.info(
+            "MOSS reference encoder CUDA graphs: batch=%s buckets(s)=%s compiled=%s captured=%d",
+            self._batch_sizes,
+            [round(s / int(self._tokenizer.sampling_rate), 2) for s in self._buckets],
+            self._compile_core,
+            len(self._graphs),
+        )
+        if self._singleton_buckets != self._buckets:
+            logger.info(
+                "MOSS reference encoder singleton buckets(s): %s",
+                [round(s / int(self._tokenizer.sampling_rate), 2) for s in self._singleton_buckets],
+            )
+
+    def _capture_all(self, core, stream: torch.cuda.Stream) -> None:
+        for batch in reversed(self._batch_sizes):
+            buckets = self._singleton_buckets if batch == 1 else self._buckets
+            for samples in reversed(buckets):
+                if batch > 1 and batch * samples > self._max_batch_samples:
+                    continue
+                audio = torch.zeros(batch, self._channels, samples, device=self._device)
+                lengths = torch.full((batch,), samples, dtype=torch.long, device=self._device)
+                try:
+                    stream.wait_stream(torch.cuda.current_stream(self._device))
+                    with torch.cuda.stream(stream):
+                        # Lazy library setup must happen before capture; one run does it.
+                        core(audio, lengths)
+                    torch.cuda.current_stream(self._device).wait_stream(stream)
+                    graph = torch.cuda.CUDAGraph()
+                    with torch.cuda.graph(graph, pool=self._pool, stream=stream):
+                        codes, code_lengths = core(audio, lengths)
+                except Exception:
+                    logger.warning(
+                        "MOSS reference encoder graph (B=%d, samples=%d) not captured; using eager",
+                        batch,
+                        samples,
+                        exc_info=True,
+                    )
+                    continue
+                self._graphs[(batch, samples)] = _Graph(graph, audio, lengths, codes, code_lengths)
+
+    @torch.no_grad()
+    def encode(self, wavs: list[torch.Tensor]) -> list[torch.Tensor | None]:
+        """Codes ``(frames, n_vq)`` int64 on CPU per clip; None where no graph fits.
+
+        ``wavs`` are prepared the way the processor prepares them: ``(C, T)``
+        float32 after channel handling, resampling and loudness normalization.
+        Clips are grouped longest first, so each replay pads to the bucket of
+        similar-length clips.
+        """
+        out: list[torch.Tensor | None] = [None] * len(wavs)
+        order = sorted(range(len(wavs)), key=lambda i: -int(wavs[i].shape[-1]))
+        pos = 0
+        while pos < len(order):
+            longest = int(wavs[order[pos]].shape[-1])
+            # Fine singleton buckets reduce padding without fragmenting an
+            # already assembled batch into separate graph replays.
+            buckets = self._singleton_buckets if len(order) - pos == 1 else self._buckets
+            samples = next((s for s in buckets if s >= longest and (1, s) in self._graphs), None)
+            if samples is None:
+                pos += 1  # longer than every bucket: left to the caller
+                continue
+            widest = max(b for b in self._batch_sizes if (b, samples) in self._graphs)
+            rows = order[pos : pos + widest]
+            batch = next(b for b in self._batch_sizes if b >= len(rows) and (b, samples) in self._graphs)
+            for row, codes in zip(rows, self._replay(self._graphs[(batch, samples)], [wavs[i] for i in rows])):
+                out[row] = codes
+            pos += len(rows)
+        return out
+
+    def _replay(self, entry: _Graph, wavs: list[torch.Tensor]) -> list[torch.Tensor]:
+        if self._transfer_buffers is not None and all(w.device.type == "cpu" for w in wavs):
+            return self._replay_batched_transfer(entry, wavs)
+        count, batch, samples = len(wavs), entry.audio.shape[0], entry.audio.shape[-1]
+        # Unused rows encode full-length silence: a fully masked row would be
+        # all-NaN inside attention, and rows never mix.
+        entry.audio.zero_()
+        lengths = [int(w.shape[-1]) for w in wavs] + [samples] * (batch - count)
+        entry.lengths.copy_(torch.tensor(lengths, dtype=torch.long))
+        for row, wav in enumerate(wavs):
+            entry.audio[row, :, : wav.shape[-1]].copy_(wav)
+        entry.graph.replay()
+        codes = entry.codes[:, :count].cpu()
+        code_lengths = entry.code_lengths[:count].tolist()
+        return [codes[:, row, : int(code_lengths[row])].transpose(0, 1).contiguous().long() for row in range(count)]
+
+    def _replay_batched_transfer(self, entry: _Graph, wavs: list[torch.Tensor]) -> list[torch.Tensor]:
+        """Stage one batch on pinned host memory, then wait once for its outputs."""
+        audio_slab, lengths_slab, codes_slab, code_lengths_slab, copied = self._transfer_buffers
+        count, batch, samples = len(wavs), entry.audio.shape[0], entry.audio.shape[-1]
+        audio = audio_slab[: entry.audio.numel()].view_as(entry.audio)
+        lengths = lengths_slab[:batch]
+        audio.zero_()
+        lengths.fill_(samples)
+        for row, wav in enumerate(wavs):
+            audio[row, :, : wav.shape[-1]].copy_(wav)
+            lengths[row] = wav.shape[-1]
+        entry.audio.copy_(audio, non_blocking=True)
+        entry.lengths.copy_(lengths, non_blocking=True)
+        entry.graph.replay()
+        result = entry.codes[:, :count]
+        codes = codes_slab[: result.numel()].view_as(result)
+        code_lengths = code_lengths_slab[:count]
+        codes.copy_(result, non_blocking=True)
+        code_lengths.copy_(entry.code_lengths[:count], non_blocking=True)
+        copied.record(torch.cuda.current_stream(self._device))
+        copied.synchronize()
+        # Returned codes must outlive this worker's next replay, including
+        # one-frame clips whose transpose is already contiguous.
+        return [
+            codes[:, row, : int(code_lengths[row])].transpose(0, 1).to(dtype=torch.long, copy=True).contiguous()
+            for row in range(count)
+        ]

vllm_omni/model_executor/models/moss_tts/shared_reference_encoder.py

@@ -0,0 +1,334 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""One reference-audio encoder shared by the API processes of a server.
+
+Each API process otherwise loads its own copy of the audio tokenizer on the
+GPU and encodes only the clips it receives. When the processes share a
+directory (``VLLM_OMNI_MOSS_REF_CODES_SHARED_DIR``), the first one to take the lock in it
+hosts the encoder and serves the others over a Unix socket in that directory.
+Requests from all processes are merged into batches, and only the host keeps
+the tokenizer and its CUDA graphs on the GPU.
+
+Messages carry raw arrays and a JSON header; nothing is unpickled. A client
+that cannot reach the host encodes locally instead.
+"""
+
+from __future__ import annotations
+
+import collections
+import fcntl
+import json
+import os
+import threading
+import time
+from collections.abc import Callable
+from contextlib import nullcontext
+from multiprocessing.connection import Client, Connection, Listener
+
+import numpy as np
+import torch
+from vllm.logger import init_logger
+
+logger = init_logger(__name__)
+
+_LOCK_NAME = "ref-encoder.lock"
+_SOCKET_NAME = "ref-encoder.sock"
+# Clips one worker merges into an encode, across all API processes. Several
+# workers encode concurrently on their own streams: one serial queue of large
+# batches makes a burst of new references wait behind each other. A free
+# worker takes what is queued without waiting for more: API processes already
+# gather their clips before sending them.
+_MAX_BATCH = 8
+# A client gives up on the host after this long and encodes locally.
+_CONNECT_TIMEOUT_S = 120.0
+_REQUEST_TIMEOUT_S = 120.0
+# First-time compilation can outlast a normal encode request. Readiness has a
+# separate deadline and must finish before the API starts accepting requests.
+_STARTUP_TIMEOUT_S = 1200.0
+
+EncodeFn = Callable[[list[torch.Tensor]], list]
+
+# Open for the life of the process: closing it would release the host lock.
+_host_lock_handle = None
+
+
+class SharedReferenceEncoderStartupError(RuntimeError):
+    """The shared encoder cannot safely accept requests yet."""
+
+
+def elect_host(shared_dir: str) -> bool:
+    """True in the one process that holds the encoder lock for ``shared_dir``.
+
+    The lock is held for the life of the process, so a restarted server
+    elects a new host even if the old socket file is still present.
+    """
+    global _host_lock_handle
+    path = os.path.join(shared_dir, _LOCK_NAME)
+    handle = open(path, "a+")  # noqa: SIM115 - kept open to hold the lock
+    try:
+        fcntl.flock(handle, fcntl.LOCK_EX | fcntl.LOCK_NB)
+    except OSError:
+        handle.close()
+        return False
+    _host_lock_handle = handle
+    return True
+
+
+def _send(conn: Connection, header: dict, arrays: list[np.ndarray]) -> None:
+    conn.send_bytes(json.dumps(header).encode())
+    for array in arrays:
+        conn.send_bytes(np.ascontiguousarray(array).tobytes())
+
+
+def _recv(conn: Connection) -> tuple[dict, list[bytes]]:
+    header = json.loads(conn.recv_bytes().decode())
+    return header, [conn.recv_bytes() for _ in range(int(header["count"]))]
+
+
+class _Job:
+    __slots__ = ("wav", "done", "result", "queued_at")
+
+    def __init__(self, wav: torch.Tensor) -> None:
+        self.wav = wav
+        self.done = threading.Event()
+        self.result: torch.Tensor | BaseException | None = None
+        self.queued_at = time.monotonic()
+
+
+class SharedReferenceEncoderHost:
+    """Serves reference encodes for every API process sharing ``shared_dir``.
+
+    ``workers`` pairs each worker's encode function with an optional warmup
+    (for example, capturing that worker's CUDA graphs).
+    """
+
+    def __init__(
+        self,
+        shared_dir: str,
+        workers: list[tuple[EncodeFn, Callable[[], object] | None]],
+        *,
+        stream_priority: int = 0,
+        batch_window_ms: float = 0.0,
+    ) -> None:
+        self._stream_priority = stream_priority
+        self._batch_window_s = max(0.0, float(batch_window_ms)) / 1000
+        self._pending: collections.deque[_Job] = collections.deque()
+        # Includes jobs already taken by a worker, so shutdown can release
+        # every caller without waiting for an in-flight GPU operation.
+        self._outstanding: set[_Job] = set()
+        self._ready = threading.Condition()
+        self._path = os.path.join(shared_dir, _SOCKET_NAME)
+        if os.path.exists(self._path):
+            os.unlink(self._path)  # left by a previous host; we hold the lock
+        self._listener = Listener(self._path, family="AF_UNIX")
+        os.chmod(self._path, 0o600)
+        self._closed = False
+        self._started = threading.Event()
+        self._startup_error: Exception | None = None
+        threading.Thread(
+            target=self._start_workers, args=(workers,), name="moss-ref-encoder-start", daemon=True
+        ).start()
+        threading.Thread(target=self._accept_loop, name="moss-ref-encoder-accept", daemon=True).start()
+        logger.info("MOSS shared reference encoder serving at %s with %d workers", self._path, len(workers))
+
+    def encode(self, wavs: list[torch.Tensor]) -> list:
+        """Encode on behalf of the host process itself."""
+        self.wait_until_ready()
+        jobs = [_Job(wav) for wav in wavs]
+        with self._ready:
+            if self._closed:
+                raise RuntimeError("shared reference encoder is closed")
+            self._outstanding.update(jobs)
+            self._pending.extend(jobs)
+            self._ready.notify(len(jobs))
+        for job in jobs:
+            job.done.wait()
+        return [job.result for job in jobs]
+
+    def close(self) -> None:
+        with self._ready:
+            self._closed = True
+            self._started.set()
+            for job in self._outstanding:
+                job.result = RuntimeError("shared reference encoder closed before encoding completed")
+                job.done.set()
+            self._outstanding.clear()
+            self._pending.clear()
+            self._ready.notify_all()
+        try:
+            self._listener.close()
+        except OSError:
+            pass
+
+    def wait_until_ready(self) -> None:
+        if not self._started.wait(_STARTUP_TIMEOUT_S):
+            raise SharedReferenceEncoderStartupError("shared reference encoder warmup timed out")
+        if self._closed or self._startup_error is not None:
+            raise SharedReferenceEncoderStartupError("shared reference encoder startup failed") from self._startup_error
+
+    def _accept_loop(self) -> None:
+        while not self._closed:
+            try:
+                conn = self._listener.accept()
+            except OSError:
+                return
+            threading.Thread(target=self._serve, args=(conn,), name="moss-ref-encoder-conn", daemon=True).start()
+
+    def _serve(self, conn: Connection) -> None:
+        with conn:
+            while True:
+                try:
+                    header, payloads = _recv(conn)
+                except (EOFError, OSError):
+                    return
+                if header.get("op") == "ready":
+                    try:
+                        self.wait_until_ready()
+                        error = None
+                    except SharedReferenceEncoderStartupError as exc:
+                        error = str(exc)
+                    try:
+                        _send(conn, {"count": 0, "ready": error is None, "error": error}, [])
+                    except OSError:
+                        return
+                    continue
+                wavs = [
+                    torch.from_numpy(np.frombuffer(data, dtype=np.float32).reshape(shape).copy())
+                    for data, shape in zip(payloads, header["shapes"])
+                ]
+                results = self.encode(wavs)
+                errors = [None if isinstance(r, torch.Tensor) else repr(r) for r in results]
+                codes = [r.to(torch.int32).numpy() for r in results if isinstance(r, torch.Tensor)]
+                reply = {"count": len(codes), "shapes": [list(c.shape) for c in codes], "errors": errors}
+                try:
+                    _send(conn, reply, codes)
+                except OSError:
+                    return
+
+    def _start_workers(self, workers: list[tuple[EncodeFn, Callable[[], object] | None]]) -> None:
+        # Every warmup (graph capture) finishes before any worker encodes: CUDA
+        # work in another thread would invalidate a capture in progress.
+        try:
+            for _, warmup in workers:
+                if warmup is not None:
+                    warmup()
+            if self._closed:
+                return
+            for index, (encode_fn, _) in enumerate(workers):
+                threading.Thread(
+                    target=self._encode_loop, args=(encode_fn,), name=f"moss-ref-encoder-{index}", daemon=True
+                ).start()
+            logger.info("MOSS shared reference encoder ready: %d workers", len(workers))
+        except Exception as exc:  # noqa: BLE001 - wake readiness waiters with the failure
+            self._startup_error = exc
+            logger.exception("MOSS shared reference encoder startup failed")
+        finally:
+            self._started.set()
+
+    def _encode_loop(self, encode_fn: EncodeFn) -> None:
+        stream = torch.cuda.Stream(priority=self._stream_priority) if torch.cuda.is_available() else None
+        with torch.cuda.stream(stream) if stream is not None else nullcontext():
+            self._serve_jobs(encode_fn)
+
+    def _serve_jobs(self, encode_fn: EncodeFn) -> None:
+        while True:
+            with self._ready:
+                while True:
+                    if self._closed:
+                        return
+                    if not self._pending:
+                        self._ready.wait()
+                        continue
+                    remaining = self._pending[0].queued_at + self._batch_window_s - time.monotonic()
+                    if remaining > 0 and len(self._pending) < _MAX_BATCH:
+                        # Keep clips in the common queue while collecting;
+                        # workers must not each reserve a singleton first.
+                        self._ready.wait(remaining)
+                        continue
+                    break
+                jobs = [self._pending.popleft() for _ in range(min(_MAX_BATCH, len(self._pending)))]
+            try:
+                results = encode_fn([job.wav for job in jobs])
+                if len(results) != len(jobs):
+                    raise RuntimeError(
+                        f"shared reference encoder returned {len(results)} results for {len(jobs)} clips"
+                    )
+            except Exception as error:  # noqa: BLE001 - every waiter gets the failure
+                results = [error] * len(jobs)
+            with self._ready:
+                for job, result in zip(jobs, results, strict=True):
+                    # close() may already have completed this caller with an
+                    # error. A late GPU result must not replace that outcome.
+                    if not job.done.is_set():
+                        job.result = result
+                        job.done.set()
+                    self._outstanding.discard(job)
+
+
+class SharedReferenceEncoderClient:
+    """Sends reference clips to the host process of ``shared_dir``."""
+
+    def __init__(self, shared_dir: str) -> None:
+        self._path = os.path.join(shared_dir, _SOCKET_NAME)
+        # One connection per request in flight; idle ones are reused.
+        self._idle: list[Connection] = []
+        self._lock = threading.Lock()
+        self._connected = False
+
+    def _connect(self) -> Connection:
+        # The host may still be starting until the first connection succeeds;
+        # later reconnects are quick.
+        deadline = time.monotonic() + (1.0 if self._connected else _CONNECT_TIMEOUT_S)
+        while True:
+            try:
+                conn = Client(self._path, family="AF_UNIX")
+                self._connected = True
+                break
+            except (FileNotFoundError, ConnectionRefusedError):
+                if time.monotonic() >= deadline:
+                    raise
+                time.sleep(0.2)
+        try:
+            _send(conn, {"op": "ready", "count": 0}, [])
+            if not conn.poll(_STARTUP_TIMEOUT_S):
+                raise SharedReferenceEncoderStartupError("shared reference encoder readiness timed out")
+            header, _ = _recv(conn)
+            if not header.get("ready"):
+                raise SharedReferenceEncoderStartupError(header.get("error", "shared reference encoder not ready"))
+            return conn
+        except BaseException:
+            conn.close()
+            raise
+
+    def wait_until_ready(self) -> None:
+        """Startup handshake, without encoding or populating the reference cache."""
+        try:
+            conn = self._connect()
+        except Exception as exc:
+            raise SharedReferenceEncoderStartupError("shared reference encoder is not ready") from exc
+        with self._lock:
+            self._idle.append(conn)
+
+    def encode(self, wavs: list[torch.Tensor]) -> list:
+        """Codes (or an exception) per clip; raises when the host is unreachable."""
+        with self._lock:
+            conn = self._idle.pop() if self._idle else None
+        try:
+            if conn is None:
+                conn = self._connect()
+            arrays = [w.detach().to(torch.float32).cpu().numpy() for w in wavs]
+            _send(conn, {"count": len(arrays), "shapes": [list(a.shape) for a in arrays]}, arrays)
+            if not conn.poll(_REQUEST_TIMEOUT_S):
+                raise TimeoutError(f"no reply from the shared reference encoder in {_REQUEST_TIMEOUT_S}s")
+            header, payloads = _recv(conn)
+        except BaseException:
+            if conn is not None:
+                conn.close()
+            raise
+        with self._lock:
+            self._idle.append(conn)
+        codes = iter(
+            torch.from_numpy(np.frombuffer(data, dtype=np.int32).reshape(shape).copy()).long()
+            for data, shape in zip(payloads, header["shapes"])
+        )
+        return [next(codes) if error is None else RuntimeError(error) for error in header["errors"]]

vllm_omni/model_executor/models/moss_tts/slot_attention.py

@@ -0,0 +1,475 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Direct-slot ring attention for MOSS streaming codec.
+
+Inputs q/k/v are (execution_batch, heads, padded_T, head_dim). cache is
+contiguous (2, state_capacity, heads, ring_capacity, head_dim). end_offset
+contains absolute *pre-call* positions per state slot. slot_ids and valid_lengths
+contain one entry per execution row. valid_lengths=0 is a scratch/padding row.
+
+This op writes only valid tokens, attends to the resulting ring with the same
+position semantics as RingKVCache.complete, then advances only active offsets.
+Slots must be unique and in bounds; 0 <= valid_lengths <= padded_T. Chunks
+longer than the ring retain only their final ring_capacity tokens.
+These device-value preconditions are caller-owned to avoid host synchronization.
+The op deliberately preserves chunk-complete ring overwrite semantics; it does
+not recover entries overwritten by later tokens of the same input chunk.
+"""
+
+import torch
+from vllm.triton_utils import tl, triton
+
+
+@triton.jit
+def _write(
+    k_ptr,
+    v_ptr,
+    cache_ptr,
+    end_ptr,
+    slots_ptr,
+    lengths_ptr,
+    ks0: tl.constexpr,
+    ks1: tl.constexpr,
+    ks2: tl.constexpr,
+    ks3: tl.constexpr,
+    vs0: tl.constexpr,
+    vs1: tl.constexpr,
+    vs2: tl.constexpr,
+    vs3: tl.constexpr,
+    end_stride: tl.constexpr,
+    slot_stride: tl.constexpr,
+    length_stride: tl.constexpr,
+    state_capacity: tl.constexpr,
+    num_heads: tl.constexpr,
+    num_frames: tl.constexpr,
+    ring_capacity: tl.constexpr,
+    head_dim: tl.constexpr,
+    block_size: tl.constexpr,
+):
+    row = tl.program_id(0)
+    slot = tl.load(slots_ptr + row * slot_stride)
+    length = tl.load(lengths_ptr + row * length_stride)
+    start = tl.load(end_ptr + slot * end_stride)
+    x = tl.program_id(1) * block_size + tl.arange(0, block_size)
+    dim = x % head_dim
+    token = (x // head_dim) % num_frames
+    head = x // (num_frames * head_dim)
+    # For chunks longer than the ring, write only the final C tokens. This
+    # avoids duplicate physical addresses (and write races) in one launch.
+    skip = tl.maximum(length - ring_capacity, 0)
+    source_token = token + skip
+    mask = (head < num_heads) & (source_token < length)
+    pos = (start + source_token) % ring_capacity
+    dst = ((slot * num_heads + head) * ring_capacity + pos) * head_dim + dim
+    kval = tl.load(k_ptr + row * ks0 + head * ks1 + source_token * ks2 + dim * ks3, mask, 0)
+    vval = tl.load(v_ptr + row * vs0 + head * vs1 + source_token * vs2 + dim * vs3, mask, 0)
+    tl.store(cache_ptr + dst, kval, mask)
+    tl.store(cache_ptr + state_capacity * num_heads * ring_capacity * head_dim + dst, vval, mask)
+
+
+@triton.jit
+def _attend(
+    q_ptr,
+    cache_ptr,
+    end_ptr,
+    slots_ptr,
+    lengths_ptr,
+    out_ptr,
+    qs0: tl.constexpr,
+    qs1: tl.constexpr,
+    qs2: tl.constexpr,
+    qs3: tl.constexpr,
+    end_stride: tl.constexpr,
+    slot_stride: tl.constexpr,
+    length_stride: tl.constexpr,
+    state_capacity: tl.constexpr,
+    num_heads: tl.constexpr,
+    num_frames: tl.constexpr,
+    ring_capacity: tl.constexpr,
+    head_dim: tl.constexpr,
+    context_size: tl.constexpr,
+    block_m: tl.constexpr,
+    block_n: tl.constexpr,
+    skip_empty_tiles: tl.constexpr = False,
+):
+    rows = tl.program_id(0) * block_m + tl.arange(0, block_m)
+    bh = tl.program_id(1)
+    batch, head = bh // num_heads, bh % num_heads
+    slot = tl.load(slots_ptr + batch * slot_stride)
+    length = tl.load(lengths_ptr + batch * length_stride)
+    start = tl.load(end_ptr + slot * end_stride)
+    end = start + length
+    dims = tl.arange(0, head_dim)
+    q = tl.load(
+        q_ptr + batch * qs0 + head * qs1 + rows[:, None] * qs2 + dims[None, :] * qs3, rows[:, None] < num_frames, 0
+    )
+    acc = tl.full((block_m, head_dim), 0.0, tl.float32)
+    maximum = tl.full((block_m,), -float("inf"), tl.float32)
+    denominator = tl.full((block_m,), 0.0, tl.float32)
+    # Padding rows own scratch slots and must emit zeros without reading KV.
+    # Keep valid-row accumulation and the ring overwrite semantics unchanged.
+    if skip_empty_tiles:
+        block_start = tl.program_id(0) * block_m
+        live = (length > 0) & (block_start < length)
+        live = live & (block_start + block_m > tl.maximum(length - ring_capacity, 0))
+    else:
+        live = length > 0
+    if live:
+        # Physical-ring order matches the existing masked_attention accumulation.
+        last = end - 1
+        last_index = last % ring_capacity
+        if skip_empty_tiles:
+            # Before the first wrap, only physical keys [0, end) exist.
+            # Omitted blocks have zero probability for every query; keep the
+            # physical order and arithmetic of all remaining blocks unchanged.
+            key_blocks = tl.cdiv(tl.minimum(end, ring_capacity), block_n)
+        else:
+            key_blocks = tl.cdiv(ring_capacity, block_n)
+        for first in range(key_blocks):
+            cols = first * block_n + tl.arange(0, block_n)
+            delta = cols - last_index
+            positions = tl.where(delta <= 0, last + delta, last + delta - ring_capacity)
+            distance = start + rows[:, None] - positions[None, :]
+            allowed = (
+                (rows[:, None] < length)
+                & (cols[None, :] < ring_capacity)
+                & (positions[None, :] >= 0)
+                & (positions[None, :] < end)
+                & (distance >= 0)
+            )
+            if context_size > 0:
+                allowed = allowed & (distance < context_size)
+            base = (slot * num_heads + head) * ring_capacity * head_dim
+            k = tl.load(cache_ptr + base + cols[None, :] * head_dim + dims[:, None], cols[None, :] < ring_capacity, 0)
+            score = tl.dot(q, k).to(tl.float32) * (head_dim**-0.5)
+            score = tl.where(allowed, score, -float("inf"))
+            next_max = tl.maximum(maximum, tl.max(score, 1))
+            safe_max = tl.where(next_max == -float("inf"), 0.0, next_max)
+            rescale = tl.exp(maximum - safe_max)
+            p = tl.exp(score - safe_max[:, None])
+            denominator = denominator * rescale + tl.sum(p, 1)
+            v = tl.load(
+                cache_ptr
+                + state_capacity * num_heads * ring_capacity * head_dim
+                + base
+                + cols[:, None] * head_dim
+                + dims[None, :],
+                cols[:, None] < ring_capacity,
+                0,
+            )
+            acc = acc * rescale[:, None] + tl.dot(p.to(v.dtype), v)
+            maximum = next_max
+    out = acc / tl.where(denominator > 0, denominator, 1.0)[:, None]
+    tl.store(
+        out_ptr + ((batch * num_heads + head) * num_frames + rows[:, None]) * head_dim + dims[None, :],
+        out,
+        rows[:, None] < num_frames,
+    )
+
+
+@triton.jit
+def _advance(
+    end_ptr,
+    slots_ptr,
+    lengths_ptr,
+    end_stride: tl.constexpr,
+    slot_stride: tl.constexpr,
+    length_stride: tl.constexpr,
+    batch_size: tl.constexpr,
+    block_size: tl.constexpr,
+):
+    row = tl.arange(0, block_size)
+    slot = tl.load(slots_ptr + row * slot_stride, row < batch_size, 0)
+    length = tl.load(lengths_ptr + row * length_stride, row < batch_size, 0)
+    active = (row < batch_size) & (length > 0)
+    start = tl.load(end_ptr + slot * end_stride, active, 0)
+    tl.store(end_ptr + slot * end_stride, start + length, active)
+
+
+@triton.jit
+def _write_attend(
+    q_ptr,
+    k_ptr,
+    v_ptr,
+    cache_ptr,
+    starts_ptr,
+    slots_ptr,
+    lengths_ptr,
+    out_ptr,
+    qs0: tl.constexpr,
+    qs1: tl.constexpr,
+    qs2: tl.constexpr,
+    qs3: tl.constexpr,
+    ks0: tl.constexpr,
+    ks1: tl.constexpr,
+    ks2: tl.constexpr,
+    ks3: tl.constexpr,
+    vs0: tl.constexpr,
+    vs1: tl.constexpr,
+    vs2: tl.constexpr,
+    vs3: tl.constexpr,
+    start_stride: tl.constexpr,
+    slot_stride: tl.constexpr,
+    length_stride: tl.constexpr,
+    state_capacity: tl.constexpr,
+    num_heads: tl.constexpr,
+    num_frames: tl.constexpr,
+    ring_capacity: tl.constexpr,
+    head_dim: tl.constexpr,
+    context_size: tl.constexpr,
+    write_frames: tl.constexpr,
+    block_m: tl.constexpr,
+    block_n: tl.constexpr,
+    block_w: tl.constexpr,
+):
+    rows = tl.program_id(0) * block_m + tl.arange(0, block_m)
+    bh = tl.program_id(1)
+    batch, head = bh // num_heads, bh % num_heads
+    slot = tl.load(slots_ptr + batch * slot_stride)
+    length = tl.load(lengths_ptr + batch * length_stride)
+    start = tl.load(starts_ptr + batch * start_stride)
+    end = start + length
+    dims = tl.arange(0, head_dim)
+    base = (slot * num_heads + head) * ring_capacity * head_dim
+    v_base = state_capacity * num_heads * ring_capacity * head_dim + base
+    # Every program of this (row, head) writes the whole chunk for its head,
+    # so it only reads ring entries it wrote itself (duplicate writes across
+    # query blocks store identical values). Chunks longer than the ring keep
+    # only their final ring_capacity tokens, as in _write.
+    skip = tl.maximum(length - ring_capacity, 0)
+    for first in range(tl.cdiv(write_frames, block_w)):
+        token = first * block_w + tl.arange(0, block_w)
+        source = token + skip
+        mask = (source < length)[:, None]
+        pos = (start + source) % ring_capacity
+        kval = tl.load(k_ptr + batch * ks0 + head * ks1 + source[:, None] * ks2 + dims[None, :] * ks3, mask, 0)
+        vval = tl.load(v_ptr + batch * vs0 + head * vs1 + source[:, None] * vs2 + dims[None, :] * vs3, mask, 0)
+        dst = pos[:, None] * head_dim + dims[None, :]
+        tl.store(cache_ptr + base + dst, kval, mask)
+        tl.store(cache_ptr + v_base + dst, vval, mask)
+    tl.debug_barrier()
+    q = tl.load(
+        q_ptr + batch * qs0 + head * qs1 + rows[:, None] * qs2 + dims[None, :] * qs3, rows[:, None] < num_frames, 0
+    )
+    acc = tl.full((block_m, head_dim), 0.0, tl.float32)
+    maximum = tl.full((block_m,), -float("inf"), tl.float32)
+    denominator = tl.full((block_m,), 0.0, tl.float32)
+    if length > 0:
+        last = end - 1
+        last_index = last % ring_capacity
+        for first in range(tl.cdiv(ring_capacity, block_n)):
+            cols = first * block_n + tl.arange(0, block_n)
+            delta = cols - last_index
+            positions = tl.where(delta <= 0, last + delta, last + delta - ring_capacity)
+            distance = start + rows[:, None] - positions[None, :]
+            allowed = (
+                (rows[:, None] < length)
+                & (cols[None, :] < ring_capacity)
+                & (positions[None, :] >= 0)
+                & (positions[None, :] < end)
+                & (distance >= 0)
+            )
+            if context_size > 0:
+                allowed = allowed & (distance < context_size)
+            k = tl.load(cache_ptr + base + cols[None, :] * head_dim + dims[:, None], cols[None, :] < ring_capacity, 0)
+            score = tl.dot(q, k).to(tl.float32) * (head_dim**-0.5)
+            score = tl.where(allowed, score, -float("inf"))
+            next_max = tl.maximum(maximum, tl.max(score, 1))
+            safe_max = tl.where(next_max == -float("inf"), 0.0, next_max)
+            rescale = tl.exp(maximum - safe_max)
+            p = tl.exp(score - safe_max[:, None])
+            denominator = denominator * rescale + tl.sum(p, 1)
+            v = tl.load(
+                cache_ptr + v_base + cols[:, None] * head_dim + dims[None, :],
+                cols[:, None] < ring_capacity,
+                0,
+            )
+            acc = acc * rescale[:, None] + tl.dot(p.to(v.dtype), v)
+            maximum = next_max
+    out = acc / tl.where(denominator > 0, denominator, 1.0)[:, None]
+    tl.store(
+        out_ptr + ((batch * num_heads + head) * num_frames + rows[:, None]) * head_dim + dims[None, :],
+        out,
+        rows[:, None] < num_frames,
+    )
+
+
+@torch.library.custom_op("vllm_omni::moss_slot_ring_attention", mutates_args=("cache", "end_offset"))
+def slot_ring_attention(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    cache: torch.Tensor,
+    end_offset: torch.Tensor,
+    slot_ids: torch.Tensor,
+    valid_lengths: torch.Tensor,
+    context: int = -1,
+    skip_empty_tiles: bool = False,
+) -> torch.Tensor:
+    """Fused address/mask path, with three ordered launches and no KV gather.
+
+    q/k are already rotary-position encoded. The caller still advances its
+    separate MHA/RoPE offset state; this op advances only end_offset in place.
+    context=-1 means unbounded causal attention over available ring entries.
+    Outputs beyond each row's valid_lengths are exactly zero.
+    skip_empty_tiles skips wholly masked tiles for small (T <= 32) chunks.
+    Larger chunks keep the static loop, which benchmarks faster after wrap.
+    """
+    b, h, t, d = q.shape
+    if k.shape != q.shape or v.shape != q.shape:
+        raise ValueError("q/k/v must have identical shapes")
+    if q.dtype not in (torch.bfloat16, torch.float16) or k.dtype != q.dtype or v.dtype != q.dtype:
+        raise ValueError("q/k/v must share bf16 or fp16 dtype")
+    if d not in (32, 64, 128) or t < 1:
+        raise ValueError("requires head_dim 32/64/128 and positive padded_T")
+    if cache.ndim != 5 or cache.shape[0] != 2 or cache.shape[2] != h or cache.shape[4] != d:
+        raise ValueError("cache must be (2, state_capacity, heads, capacity, head_dim)")
+    s, c = cache.shape[1], cache.shape[3]
+    if not cache.is_contiguous() or cache.dtype != q.dtype:
+        raise ValueError("cache must be contiguous and share q dtype")
+    if end_offset.shape != (s,) or slot_ids.shape != (b,) or valid_lengths.shape != (b,):
+        raise ValueError("invalid metadata shapes")
+    if any(x.dtype not in (torch.int32, torch.int64) for x in (end_offset, slot_ids, valid_lengths)):
+        raise ValueError("metadata must be int32/int64")
+    if q.device.type != "cuda" or any(x.device != q.device for x in (k, v, cache, end_offset, slot_ids, valid_lengths)):
+        raise ValueError("all tensors must share one CUDA device")
+    if context == 0 or context < -1:
+        raise ValueError("context must be -1 or positive")
+    out = torch.empty(q.shape, device=q.device, dtype=q.dtype)
+    if b == 0:
+        return out
+    strides = (end_offset.stride(0), slot_ids.stride(0), valid_lengths.stride(0))
+    _write[(b, triton.cdiv(h * t * d, 256))](
+        k, v, cache, end_offset, slot_ids, valid_lengths, *k.stride(), *v.stride(), *strides, s, h, t, c, d, 256
+    )
+    bm = 16 if t <= 32 else 64
+    _attend[(triton.cdiv(t, bm), b * h)](
+        q,
+        cache,
+        end_offset,
+        slot_ids,
+        valid_lengths,
+        out,
+        *q.stride(),
+        *strides,
+        s,
+        h,
+        t,
+        c,
+        d,
+        context,
+        bm,
+        64,
+        skip_empty_tiles=skip_empty_tiles and t <= 32,
+        num_warps=8 if c <= 256 else 4,
+        num_stages=2,
+    )
+    _advance[(1,)](end_offset, slot_ids, valid_lengths, *strides, b, triton.next_power_of_2(b))
+    return out
+
+
+@slot_ring_attention.register_fake
+def _(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    cache: torch.Tensor,
+    end_offset: torch.Tensor,
+    slot_ids: torch.Tensor,
+    valid_lengths: torch.Tensor,
+    context: int = -1,
+    skip_empty_tiles: bool = False,
+) -> torch.Tensor:
+    return torch.empty(q.shape, device=q.device, dtype=q.dtype)
+
+
+def _check_inputs(q, k, v, cache, metadata, context):
+    b, h, t, d = q.shape
+    if k.shape != q.shape or v.shape != q.shape:
+        raise ValueError("q/k/v must have identical shapes")
+    if q.dtype not in (torch.bfloat16, torch.float16) or k.dtype != q.dtype or v.dtype != q.dtype:
+        raise ValueError("q/k/v must share bf16 or fp16 dtype")
+    if d not in (32, 64, 128) or t < 1:
+        raise ValueError("requires head_dim 32/64/128 and positive padded_T")
+    if cache.ndim != 5 or cache.shape[0] != 2 or cache.shape[2] != h or cache.shape[4] != d:
+        raise ValueError("cache must be (2, state_capacity, heads, capacity, head_dim)")
+    if not cache.is_contiguous() or cache.dtype != q.dtype:
+        raise ValueError("cache must be contiguous and share q dtype")
+    if any(x.dtype not in (torch.int32, torch.int64) for x in metadata):
+        raise ValueError("metadata must be int32/int64")
+    if q.device.type != "cuda" or any(x.device != q.device for x in (k, v, cache, *metadata)):
+        raise ValueError("all tensors must share one CUDA device")
+    if context == 0 or context < -1:
+        raise ValueError("context must be -1 or positive")
+
+
+@torch.library.custom_op("vllm_omni::moss_slot_ring_attention_rows", mutates_args=("cache",))
+def slot_ring_attention_rows(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    cache: torch.Tensor,
+    starts: torch.Tensor,
+    slot_ids: torch.Tensor,
+    valid_lengths: torch.Tensor,
+    context: int = -1,
+) -> torch.Tensor:
+    """``slot_ring_attention`` as one launch, with per-row start positions.
+
+    ``starts`` holds each execution row's pre-call absolute position (the
+    caller's gathered offset). The ring is written and attended exactly as by
+    ``slot_ring_attention``; no offset is advanced here, so the caller owns
+    every offset update and can fuse them with its own.
+    """
+    b, h, t, d = q.shape
+    _check_inputs(q, k, v, cache, (starts, slot_ids, valid_lengths), context)
+    if starts.shape != (b,) or slot_ids.shape != (b,) or valid_lengths.shape != (b,):
+        raise ValueError("invalid metadata shapes")
+    out = torch.empty(q.shape, device=q.device, dtype=q.dtype)
+    if b == 0:
+        return out
+    s, c = cache.shape[1], cache.shape[3]
+    bm = 16 if t <= 32 else 64
+    _write_attend[(triton.cdiv(t, bm), b * h)](
+        q,
+        k,
+        v,
+        cache,
+        starts,
+        slot_ids,
+        valid_lengths,
+        out,
+        *q.stride(),
+        *k.stride(),
+        *v.stride(),
+        starts.stride(0),
+        slot_ids.stride(0),
+        valid_lengths.stride(0),
+        s,
+        h,
+        t,
+        c,
+        d,
+        context,
+        min(t, c),
+        bm,
+        64,
+        16,
+        num_warps=8 if c <= 256 else 4,
+        num_stages=2,
+    )
+    return out
+
+
+@slot_ring_attention_rows.register_fake
+def _(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    cache: torch.Tensor,
+    starts: torch.Tensor,
+    slot_ids: torch.Tensor,
+    valid_lengths: torch.Tensor,
+    context: int = -1,
+) -> torch.Tensor:
+    return torch.empty(q.shape, device=q.device, dtype=q.dtype)

vllm_omni/model_executor/models/qwen3_omni/first_frame_decoder.py

@@ -0,0 +1,99 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""First codec frame -> PCM inside the Qwen3-Omni Talker process.
+
+Code2Wav decodes each streaming chunk statelessly from its codes and left
+context. With a one-frame first chunk (``codec_chunk_ramp`` starting at 1),
+chunk 0 is a context-free decode of the stream's first frame, which the Talker
+can run as soon as eager MTP completes that frame, removing the connector hop
+and the Code2Wav step from time to first audio. The orchestrator then trims
+Code2Wav's copy of those samples.
+
+The decoder is a second ``Qwen3OmniMoeCode2Wav`` built and loaded exactly like
+the Code2Wav stage's, and captured with the same one-frame input shape as that
+stage's ``(batch=1, frames=1)`` graph. Optional cuDNN autotuning can select
+different algorithms in the two processes and change floating-point rounding.
+Each bucket gets its own graph pool because it replays on a side stream while
+the Talker's own graphs may still run.
+"""
+
+from __future__ import annotations
+
+import os
+from collections.abc import Iterable, Sequence
+from typing import Any
+
+import torch
+import torch.nn as nn
+from torch.cuda import CUDAGraph
+from vllm.logger import init_logger
+
+from vllm_omni.model_executor.models.common.qwen3_code_predictor import CodePredictorWrapper
+from vllm_omni.model_executor.models.common.talker_first_audio import supports_talker_first_audio
+from vllm_omni.model_executor.stage_input_processors.chunk_size_utils import parse_chunk_ramp
+
+logger = init_logger(__name__)
+
+
+def talker_first_audio_enabled(vllm_config: Any) -> bool:
+    """Opt in only when chunk 0 carries the first frame in the same step."""
+    extra = CodePredictorWrapper._stage_connector_extra_config(vllm_config)
+    ramp = parse_chunk_ramp(extra)
+    initial_frames = ramp[0] if ramp is not None else int(extra.get("initial_codec_chunk_frames") or 0)
+    return (
+        CodePredictorWrapper._parse_bool_config(extra.get("talker_first_audio"))
+        and initial_frames == 1
+        and os.environ.get("VLLM_OMNI_TALKER_FIRST_AUDIO", "1") == "1"
+        and supports_talker_first_audio(vllm_config)
+    )
+
+
+class Qwen3OmniFirstFrameDecoder(nn.Module):
+    def __init__(self, code2wav: nn.Module, sample_rate: int) -> None:
+        super().__init__()
+        self.code2wav = code2wav
+        self.num_quantizers = int(code2wav.config.num_quantizers)
+        self.sample_rate = int(sample_rate)
+        self._graphs: dict[int, tuple[CUDAGraph, torch.Tensor, torch.Tensor]] = {}
+
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        """Load the checkpoint's ``code2wav.*`` weights as the Code2Wav stage does."""
+        loaded = self.code2wav.load_weights(weights)
+        if hasattr(self.code2wav, "precompute_snake_caches"):
+            self.code2wav.precompute_snake_caches()
+        return loaded
+
+    @torch.inference_mode()
+    def capture(self, batch_sizes: Sequence[int] = (1, 2, 4, 8)) -> None:
+        device = next(self.code2wav.parameters()).device
+        if device.type != "cuda":
+            return
+        for batch_size in sorted(set(batch_sizes)):
+            static_input = torch.zeros(batch_size, self.num_quantizers, 1, dtype=torch.long, device=device)
+            self.code2wav(static_input)
+            torch.accelerator.synchronize(device)
+            graph = CUDAGraph()
+            with torch.cuda.graph(graph, pool=torch.cuda.graph_pool_handle()):
+                static_output = self.code2wav(static_input)
+            self._graphs[batch_size] = (graph, static_input, static_output)
+        logger.info("Captured Qwen3-Omni Talker first-frame decoder graphs for batch sizes %s", sorted(self._graphs))
+
+    @torch.inference_mode()
+    def decode(self, codes: torch.Tensor) -> torch.Tensor:
+        """``codes`` [n, num_quantizers] long on device -> float32 PCM [n, samples] (new tensor)."""
+        if not self._graphs:
+            return self.code2wav(codes.to(torch.long).unsqueeze(-1))[:, 0, :].float()
+        outputs = []
+        largest = max(self._graphs)
+        for start in range(0, int(codes.shape[0]), largest):
+            chunk = codes[start : start + largest]
+            rows = int(chunk.shape[0])
+            batch_size = min(size for size in self._graphs if size >= rows)
+            graph, static_input, static_output = self._graphs[batch_size]
+            static_input.zero_()
+            static_input[:rows, :, 0].copy_(chunk)
+            graph.replay()
+            # A replay overwrites graph-owned storage; .float() aliases FP32 outputs.
+            outputs.append(static_output[:rows, 0, :].to(dtype=torch.float32, copy=True))
+        return torch.cat(outputs, dim=0)

vllm_omni/model_executor/models/qwen3_omni/pipeline.py

@@ -45,6 +45,8 @@ QWEN3_OMNI_PIPELINE = PipelineConfig(
             engine_output_type="latent",
             custom_process_next_stage_input_func=(f"{_PROC}.thinker2talker_full_payload"),
             async_chunk_process_next_stage_input_func=(f"{_PROC}.thinker2talker_async_chunk"),
+            # Takes effect only when the deploy selects model_runner v2 for this stage.
+            supports_native_mrv2_data_plane=True,
             sampling_constraints={"detokenize": True},
         ),
         StagePipelineConfig(
@@ -57,6 +59,8 @@ QWEN3_OMNI_PIPELINE = PipelineConfig(
             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"),
+            # Each stage can use MRv2 when its deploy selects model_runner v2.
+            supports_native_mrv2_data_plane=True,
             sampling_constraints={
                 "detokenize": False,
                 "stop_token_ids": [2150],
@@ -72,6 +76,7 @@ QWEN3_OMNI_PIPELINE = PipelineConfig(
             final_output_type="audio",
             hf_config_name="thinker_config",
             engine_output_type="audio",
+            supports_native_mrv2_data_plane=True,
             sampling_constraints={"detokenize": True},
             requires_full_payload_input=True,
         ),

vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py

@@ -4,7 +4,7 @@
 """Inference-only Qwen3-Omni-Moe unified model (thinker + talker + code2wav)."""
 
 import asyncio
-from collections.abc import AsyncGenerator, Iterable
+from collections.abc import AsyncGenerator, Iterable, Mapping
 from functools import cached_property
 from typing import Any
 
@@ -50,6 +50,7 @@ from vllm_omni.data_entry_keys import Embeddings, HiddenStates, Ids, OmniPayload
 from vllm_omni.metrics import definitions as defs
 from vllm_omni.model_executor.custom_process_mixin import CustomProcessMixin
 from vllm_omni.model_executor.models.output_templates import OmniOutput
+from vllm_omni.model_executor.models.qwen3_omni.first_frame_decoder import talker_first_audio_enabled
 from vllm_omni.model_executor.models.qwen3_omni.quantization import (
     apply_outer_quant_config_mapping,
 )
@@ -61,7 +62,9 @@ from vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_moe_thinker import (
     Qwen3OmniMoeThinkerProcessingInfo,
 )
 from vllm_omni.model_executor.models.utils import add_prefix_to_loaded_weights, safe_tensor_reshape
+from vllm_omni.model_executor.stage_input_processors.chunk_size_utils import parse_chunk_ramp, ramp_decode_windows
 from vllm_omni.platforms import current_omni_platform
+from vllm_omni.utils.device_copy import to_device_nonblocking
 
 # Special token IDs for Qwen3 Omni MoE
 # Reference: https://huggingface.co/Qwen/Qwen3-Omni-30B-A3B-Instruct/blob/main/tokenizer_config.json
@@ -88,6 +91,13 @@ TALKER_CODEC_THINK_EOS_ID = 4205  # Think mode end
 logger = init_logger(__name__)
 
 
+def _cpu_long(values: Any) -> torch.Tensor:
+    """Token ids as a host tensor (payload ids arrive as lists or CPU tensors)."""
+    if isinstance(values, torch.Tensor):
+        return values.detach().to("cpu", torch.long)
+    return torch.as_tensor(values, dtype=torch.long)
+
+
 @MULTIMODAL_REGISTRY.register_processor(
     Qwen3OmniMoeThinkerMultiModalProcessor,
     info=Qwen3OmniMoeThinkerProcessingInfo,
@@ -135,6 +145,11 @@ class Qwen3OmniMoeForConditionalGeneration(
 
     realtime_max_tokens = 64
 
+    # Set by the MRv2 model state once eager Talker-MTP frames are enabled.
+    eager_frames_active = False
+    # Postprocess only stores the last hidden for the deferred decode MTP.
+    eager_frames_need_postprocess = False
+
     def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
         super().__init__()
         self.have_multimodal_outputs = True
@@ -180,9 +195,18 @@ class Qwen3OmniMoeForConditionalGeneration(
         # thinker layers and the forward must return what stock vLLM expects.
         self.is_staged_run = getattr(vllm_config.model_config, "model_stage", None) is not None
         self._returns_tuple = self.model_stage == "thinker" and self.is_staged_run
+        # The captured layers are a fixed-structure tensor dict on the token
+        # axis, so MRv2 may replay decode in FULL graphs and return them too.
+        self.supports_mrv2_full_graph_aux_outputs = self._returns_tuple
 
         if self.model_stage == "thinker":
             self.use_async_omni_output = True
+            # MRv2: the Talker's prompt ends with the first generated token's
+            # embedding; publishing it with the prefill output lets the first
+            # Thinker->Talker chunk leave one decode step earlier.
+            # Speculative sampling can change the token width between steps.
+            # Keep the existing capture-stream handoff for the entire request.
+            self.publishes_sampled_embeddings = self.is_staged_run and vllm_config.speculative_config is None
             # Initialize thinker model (multimodal processing + text generation)
             # Create a new vllm_config with thinker_config as the hf_config
             thinker_vllm_config = vllm_config.with_hf_config(
@@ -233,6 +257,19 @@ class Qwen3OmniMoeForConditionalGeneration(
             )
             self.model = self.talker
             self.code2wav = None
+            # Residual codebooks (MTP) output key, shared by the V1 runner
+            # (talker_mtp_*) and the MRv2 model state (mtp_*).
+            self.talker_mtp_output_key = ("codes", "audio")
+            # MRv2: the step that samples CB0 may complete the frame itself,
+            # since the residual codebooks depend only on CB0 and that step's
+            # last hidden, and the next input is their embedding sum plus the
+            # next text step (see ``talker_mtp``).
+            self.mtp_eager_frames = True
+            self._codec_codebook_size = int(talker_config.code_predictor_config.vocab_size)
+            self._apply_subtalker_sampling_params(getattr(vllm_config.model_config, "subtalker_sampling_params", None))
+
+            # Decode each stream's first frame here (see talker_first_audio_enabled).
+            self.first_frame_decoder = self._build_first_frame_decoder(vllm_config, code2wav_config, prefix)
 
             # for CI: Initialize special tokens embeddings early to avoid AttributeError when loading dummy weights
             self._init_special_tokens_embeddings()
@@ -247,6 +284,9 @@ class Qwen3OmniMoeForConditionalGeneration(
                 # talker MTP codec codes must stay on GPU to avoid a per-step D2H
                 # sync stall; build_mm_cpu handles the eventual D2H at payload time.
                 ("codes", "audio"),
+                # The growing thinker-embed cache is read and extended every
+                # decode step; on CPU it cost a blocking D2H per request per step.
+                ("embed", "cached_decode"),
             }
 
         elif self.model_stage == "code2wav":
@@ -285,6 +325,74 @@ class Qwen3OmniMoeForConditionalGeneration(
                     ["hidden_states", "residual", *incoming_captures], thinker_config.text_config.hidden_size
                 )
 
+    @staticmethod
+    def _build_first_frame_decoder(
+        vllm_config: VllmConfig,
+        code2wav_config: Qwen3OmniMoeCode2WavConfig,
+        prefix: str,
+    ) -> nn.Module | None:
+        """A Code2Wav copy for the Talker's first-frame audio, or None when disabled.
+
+        Only a streaming (async-chunk) Talker hands frames to Code2Wav chunk by
+        chunk, so only there is a stream's first chunk a context-free decode.
+        The copy is built like the Code2Wav stage's module (same class, config,
+        dtype and device context). Optional cuDNN autotuning is also shared,
+        although per-process algorithm selection can change PCM rounding.
+        """
+        if not talker_first_audio_enabled(vllm_config):
+            return None
+        from vllm_omni.model_executor.models.qwen3_omni.first_frame_decoder import Qwen3OmniFirstFrameDecoder
+
+        code2wav = init_vllm_registered_model(
+            vllm_config=vllm_config.with_hf_config(code2wav_config, architectures=["Qwen3OmniMoeCode2Wav"]),
+            prefix=maybe_prefix(prefix, "first_frame_decoder.code2wav"),
+            hf_config=code2wav_config,
+            architectures=["Qwen3OmniMoeCode2Wav"],
+        )
+        logger.info("Qwen3-Omni Talker decodes first-frame audio in-stage")
+        return Qwen3OmniFirstFrameDecoder(code2wav, defs.resolve_audio_sample_rate(code2wav_config))
+
+    def capture_first_frame_graphs(self) -> None:
+        """Called by the MRv2 runner after its own graph capture."""
+        decoder = getattr(self, "first_frame_decoder", None)
+        if decoder is not None:
+            decoder.capture()
+
+    def _apply_subtalker_sampling_params(self, params: Any) -> None:
+        """Apply the deploy's ``subtalker_sampling_params`` to the residual-codebook predictor.
+
+        The predictor samples in its "stored" mode (top-k 50, top-p 0.8 by
+        default) and ignores per-call sampling arguments; ``do_sample: false``
+        makes it greedy.
+        """
+        if not isinstance(params, Mapping) or not params:
+            return
+        predictor = self.talker.code_predictor
+        predictor.set_sampling_params(
+            top_k=int(params.get("top_k", predictor._top_k)),
+            top_p=float(params.get("top_p", predictor._top_p)),
+            do_sample=bool(params.get("do_sample", True)),
+        )
+
+    # MRv2 capability names read the V1 canonical values, so a platform patch
+    # of ``talker_mtp_graph_safe`` applies to both runners.
+    @property
+    def mtp(self):
+        """MRv2 name for ``talker_mtp``; only the Talker stage has an MTP."""
+        return self.talker_mtp if self.model_stage == "talker" else None
+
+    @property
+    def mtp_output_key(self) -> tuple[str, str]:
+        return getattr(self, "talker_mtp_output_key", ("codes", "audio"))
+
+    @property
+    def mtp_graph_safe(self) -> bool:
+        return bool(getattr(self, "talker_mtp_graph_safe", False))
+
+    def mtp_frame_valid(self, layer0: torch.Tensor) -> torch.Tensor:
+        """Rows whose CB0 is a codec id rather than codec EOS (or another special id)."""
+        return (layer0 >= 0) & (layer0 < self._codec_codebook_size)
+
     def _thinker_capture_layer_indices(self) -> list[int]:
         accept_layer = getattr(self.talker_config, "accept_hidden_layer", None)
         return [0, int(accept_layer)] if self.is_staged_run and accept_layer is not None else []
@@ -496,38 +604,47 @@ class Qwen3OmniMoeForConditionalGeneration(
         # ========== Stage 3: Code2Wav ==========
         elif self.model_stage == "code2wav":
             seq_token_counts: list[int] | None = kwargs.get("seq_token_counts")
+            num_quantizers = int(self.code2wav.config.num_quantizers)
 
             # Extract codec codes from input
-            if input_ids.shape[0] % 16 == 0:
+            if input_ids.shape[0] % num_quantizers == 0:
                 if seq_token_counts is not None:
-                    max_seq_len = max(seq_token_counts) // 16
+                    max_seq_len = max(seq_token_counts) // num_quantizers
                     batch_size = len(seq_token_counts)
                     split_codes = torch.split(input_ids, seq_token_counts, dim=0)
-                    codes = torch.zeros((batch_size, 16, max_seq_len), device=input_ids.device, dtype=input_ids.dtype)
+                    codes = torch.zeros(
+                        (batch_size, num_quantizers, max_seq_len), device=input_ids.device, dtype=input_ids.dtype
+                    )
                     for idx, code in enumerate(split_codes):
-                        seq_len = code.shape[0] // 16
-                        codes[idx, :, :seq_len] = code.reshape(16, seq_len)
+                        seq_len = code.shape[0] // num_quantizers
+                        codes[idx, :, :seq_len] = code.reshape(num_quantizers, seq_len)
                 else:
-                    codes = input_ids.reshape(1, 16, -1)
+                    codes = input_ids.reshape(1, num_quantizers, -1)
             else:
                 if seq_token_counts is None:
                     logger.debug(
-                        "Code2Wav warmup input length %s is not divisible by 16; padding with zeros.",
+                        "Code2Wav warmup input length %s is not divisible by %s; padding with zeros.",
                         input_ids.shape[0],
+                        num_quantizers,
                     )
                 else:
                     logger.warning_once(
-                        "Code2Wav input length is not divisible by 16; padding with zeros. "
-                        "This is expected only during cudagraph warmup."
+                        "Code2Wav input length is not divisible by %s; padding with zeros. "
+                        "This is expected only during cudagraph warmup.",
+                        num_quantizers,
                     )
                 input_ids_flatten = input_ids.reshape(-1)
                 input_ids_flatten = torch.cat(
                     [
                         input_ids_flatten,
-                        torch.zeros(16 - input_ids.shape[0] % 16, dtype=torch.long, device=input_ids.device),
+                        torch.zeros(
+                            num_quantizers - input_ids.shape[0] % num_quantizers,
+                            dtype=torch.long,
+                            device=input_ids.device,
+                        ),
                     ]
                 )
-                codes = input_ids_flatten.reshape(1, 16, -1)
+                codes = input_ids_flatten.reshape(1, num_quantizers, -1)
 
             # Generate audio from codec codes
             # Get every request's left_context_size from runtime_additional_information (passed via kwargs)
@@ -541,6 +658,38 @@ class Qwen3OmniMoeForConditionalGeneration(
                 logger.debug("No additional_information provided to code2wav stage.")
             audio_tensors = self.generate_audio(codes, left_context_size, seq_token_counts)
 
+            # Keep the skip decision attached to this forward's output, not
+            # mutable model state: asynchronous materialization can overlap steps.
+            from vllm_omni.data_entry_keys import FIRST_AUDIO_REQUIRED_KEY
+
+            flags = [
+                bool(info.get("meta", {}).get("first_audio", False)) for info in (runtime_additional_information or [])
+            ]
+            if flags and any(flags):
+                if len(flags) != len(audio_tensors):
+                    raise ValueError("First-audio flags must align with Code2Wav requests")
+                sample_rate = defs.resolve_audio_sample_rate(self.code2wav_config)
+                frame_counts = (
+                    [count // num_quantizers for count in seq_token_counts] if seq_token_counts else [codes.shape[-1]]
+                )
+                # The causal decoder withholds a right-edge tail. The first
+                # frame has fewer samples than total_upsample; trim its actual
+                # prefix length, including when chunk 0 contains several frames.
+                prefix_lengths = [
+                    max(0, audio.numel() - (frames - 1) * int(self.code2wav.total_upsample)) if skip else 0
+                    for audio, frames, skip in zip(audio_tensors, frame_counts, flags, strict=True)
+                ]
+                return OmniOutput(
+                    text_hidden_states=None,
+                    multimodal_outputs={
+                        "model_outputs": [
+                            audio.reshape(1, -1)[..., prefix:]
+                            for audio, prefix in zip(audio_tensors, prefix_lengths, strict=True)
+                        ],
+                        "sr": [torch.tensor(sample_rate, dtype=torch.int32) for _ in audio_tensors],
+                        FIRST_AUDIO_REQUIRED_KEY: [torch.tensor(skip) for skip in flags],
+                    },
+                )
             return audio_tensors
 
         # Fallback (shouldn't reach here)
@@ -605,7 +754,13 @@ class Qwen3OmniMoeForConditionalGeneration(
 
             if "runtime_additional_information" in kwargs and "model_intermediate_buffer" not in kwargs:
                 logger.warning_once("runtime_additional_information is deprecated, use model_intermediate_buffer")
+            # MRv2 passes each request's token span.
+            spans = kwargs.get("request_token_spans")
+            if spans is not None and self.eager_frames_active and self.vllm_config.model_config.async_chunk:
+                return self._make_eager_talker_output(talker_hidden, spans)
             code_predictor_codes = [info.get("codes", {}).get("audio") for info in info_dicts]
+            if spans is not None:
+                code_predictor_codes = self._align_codes_to_spans(code_predictor_codes, spans, talker_hidden.device)
             audio_codes = torch.cat(code_predictor_codes, dim=0)
             multimodal_outputs: OmniPayload = {"codes": {"audio": audio_codes}}
             span_len = audio_codes.shape[0]
@@ -627,6 +782,44 @@ class Qwen3OmniMoeForConditionalGeneration(
 
         return model_outputs
 
+    def _align_codes_to_spans(
+        self,
+        codes: list[Any],
+        spans: list[tuple[int, int]],
+        device: torch.device,
+    ) -> list[torch.Tensor]:
+        """One ``codes.audio`` block per request with exactly its span's rows.
+
+        MRv2 also runs requests that never went through preprocess (vLLM's
+        ``_warmup_`` requests), whose buffers carry no codes; they get zero
+        rows, which the downstream processor skips like prefill rows.
+        """
+        q = int(self.talker.num_code_groups)
+        aligned: list[torch.Tensor] = []
+        for block, (start, end) in zip(codes, spans, strict=True):
+            if isinstance(block, torch.Tensor) and block.ndim == 2 and block.shape[0] == end - start:
+                aligned.append(block)
+            else:
+                aligned.append(torch.zeros((end - start, q), dtype=torch.long, device=device))
+        return aligned
+
+    def _make_eager_talker_output(self, hidden: torch.Tensor, spans: list[tuple[int, int]]) -> OmniOutput:
+        """Async-chunk output when the MRv2 runner completes frames after sampling.
+
+        ``run_eager_mtp`` writes each sampled row's frame and validity into the
+        last token row of its request span; every other row carries no frame.
+        """
+        num_tokens = spans[-1][1] if spans else 0
+        q = int(self.talker.num_code_groups)
+        mm: OmniPayload = {
+            "codes": {"audio": torch.zeros((num_tokens, q), dtype=torch.long, device=hidden.device)},
+            "meta": {
+                "codec_frame_valid": torch.zeros((num_tokens,), dtype=torch.int8, device=hidden.device),
+                "first_audio": torch.zeros((num_tokens,), dtype=torch.int8, device=hidden.device),
+            },
+        }
+        return OmniOutput(text_hidden_states=hidden, multimodal_outputs=mm)
+
     # ==================== Audio Generation ====================
 
     def generate_audio(
@@ -849,7 +1042,7 @@ class Qwen3OmniMoeForConditionalGeneration(
 
         def _proj_from_thinker(x_opt: torch.Tensor | None) -> torch.Tensor:
             if isinstance(x_opt, torch.Tensor) and x_opt.numel() > 0:
-                xin = _ensure_1x1(x_opt).to(module_device)
+                xin = to_device_nonblocking(_ensure_1x1(x_opt), module_device)
             else:
                 xin = torch.zeros(
                     (1, thinker_embed.shape[-1]),
@@ -882,27 +1075,26 @@ class Qwen3OmniMoeForConditionalGeneration(
             voice_type = str(voice_type).lower().strip()
         start_index = meta.get("num_processed_tokens", 0)
         end_index = start_index + input_embeds.shape[0]
-        # Read thinker outputs for prefill
-        thinker_sequence_embeds = embed["prefill"].to(
-            device=self._module_device(self.talker), dtype=torch.bfloat16
+        # Read thinker outputs for prefill. Every copy here is H2D from the
+        # connector payload; pageable copies would stall the whole Talker batch.
+        talker_device = self._module_device(self.talker)
+        thinker_sequence_embeds = to_device_nonblocking(embed["prefill"], talker_device).to(
+            torch.bfloat16
         )  # Tensor [P,H]
-        thinker_hidden_states = hs["output"].to(
-            device=self._module_device(self.talker), dtype=torch.bfloat16
-        )  # Tensor [K,H]
+        thinker_hidden_states = to_device_nonblocking(hs["output"], talker_device).to(torch.bfloat16)  # Tensor [K,H]
+        # Token ids stay on CPU too: prefill parsing reads them on the host.
+        thinker_sequences_cpu = None if ids.get("all") is None else _cpu_long(ids["all"])
+        thinker_chatml_ids_cpu = None if ids.get("prompt") is None else _cpu_long(ids["prompt"])
         thinker_sequences = (
-            ids.get("all")
-            if ids.get("all") is None
-            else torch.as_tensor(ids["all"], device=self._module_device(self.talker))
+            None if thinker_sequences_cpu is None else to_device_nonblocking(thinker_sequences_cpu, talker_device)
         )
         thinker_chatml_ids = (
-            ids.get("prompt")
-            if ids.get("prompt") is None
-            else torch.as_tensor(ids["prompt"], device=self._module_device(self.talker))
+            None if thinker_chatml_ids_cpu is None else to_device_nonblocking(thinker_chatml_ids_cpu, talker_device)
         )
 
-        tts_bos_thinker = embed["tts_bos"].to(device=self._module_device(self.talker), dtype=torch.bfloat16)
-        tts_eos_thinker = embed["tts_eos"].to(device=self._module_device(self.talker), dtype=torch.bfloat16)
-        tts_pad_thinker = embed["tts_pad"].to(device=self._module_device(self.talker), dtype=torch.bfloat16)
+        tts_bos_thinker = to_device_nonblocking(embed["tts_bos"], talker_device).to(torch.bfloat16)
+        tts_eos_thinker = to_device_nonblocking(embed["tts_eos"], talker_device).to(torch.bfloat16)
+        tts_pad_thinker = to_device_nonblocking(embed["tts_pad"], talker_device).to(torch.bfloat16)
 
         if thinker_sequence_embeds is None or thinker_hidden_states is None:
             raise ValueError(
@@ -924,6 +1116,8 @@ class Qwen3OmniMoeForConditionalGeneration(
             )
             if ids_chatml.ndim == 1:
                 ids_chatml = ids_chatml.unsqueeze(0)
+            if thinker_chatml_ids_cpu is not None and thinker_chatml_ids_cpu.ndim == 1:
+                thinker_chatml_ids_cpu = thinker_chatml_ids_cpu.unsqueeze(0)
         else:
             # Fallback: create dummy ids if not provided
             ids_chatml = torch.zeros(
@@ -932,6 +1126,7 @@ class Qwen3OmniMoeForConditionalGeneration(
                 device=self._module_device(self.talker),
             )
             thinker_sequences = ids_chatml
+            thinker_chatml_ids_cpu = thinker_sequences_cpu = None
 
         speaker_id = self._get_text_spk_token_id(voice_type)
         req_input_ids, req_embeds, trailing_text_hidden = self._thinker_to_talker_prefill(
@@ -944,6 +1139,8 @@ class Qwen3OmniMoeForConditionalGeneration(
             tts_bos_thinker=tts_bos_thinker,
             tts_eos_thinker=tts_eos_thinker,
             tts_pad_thinker=tts_pad_thinker,
+            input_ids_cpu=thinker_chatml_ids_cpu,
+            thinker_result_ids_cpu=thinker_sequences_cpu,
         )
 
         # Queue trailing_text_hidden for decode (drop first for next steps),
@@ -993,12 +1190,11 @@ class Qwen3OmniMoeForConditionalGeneration(
             if cached_thinker_decode_embeds is None:
                 update_dict.setdefault("embed", {})["cached_decode"] = thinker_decode_embeds
             else:
-                cached_thinker_decode_embeds = cached_thinker_decode_embeds.to(
-                    device=self._module_device(self.talker), dtype=torch.bfloat16
-                )
-                thinker_decode_embeds = thinker_decode_embeds.to(
-                    device=self._module_device(self.talker), dtype=torch.bfloat16
+                talker_device = self._module_device(self.talker)
+                cached_thinker_decode_embeds = to_device_nonblocking(cached_thinker_decode_embeds, talker_device).to(
+                    torch.bfloat16
                 )
+                thinker_decode_embeds = to_device_nonblocking(thinker_decode_embeds, talker_device).to(torch.bfloat16)
                 update_dict.setdefault("embed", {})["cached_decode"] = torch.cat(
                     [cached_thinker_decode_embeds, thinker_decode_embeds], dim=0
                 )
@@ -1015,26 +1211,31 @@ class Qwen3OmniMoeForConditionalGeneration(
         tts_bos_thinker: torch.Tensor | None = None,
         tts_eos_thinker: torch.Tensor | None = None,
         tts_pad_thinker: torch.Tensor | None = None,
+        input_ids_cpu: torch.Tensor | None = None,
+        thinker_result_ids_cpu: torch.Tensor | None = None,
     ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]:
         """
         Project thinker outputs to talker inputs during prefill stage.
 
+        ``input_ids_cpu`` / ``thinker_result_ids_cpu`` (host copies of the same
+        ids) move segment parsing off the GPU: each ``.item()`` / ``nonzero`` /
+        mask test on device ids is a host sync that stalls the Talker batch.
+
         Returns:
             (input_ids, input_embeds) for talker
         """
         target_len = thinker_result_ids.shape[-1]
-        im_start_indexes = torch.cat(
-            (
-                torch.nonzero(input_ids[0] == self.config.im_start_token_id).squeeze(-1),
-                torch.tensor([target_len], device=input_ids.device, dtype=input_ids.dtype),
-            ),
-            dim=-1,
-        )  # Shape [n_starts + 1]; Take batch 0 since batched inference is not supported here.
+        # Take batch 0 since batched inference is not supported here.
+        parse_ids = input_ids_cpu if input_ids_cpu is not None else input_ids
+        result_ids = thinker_result_ids_cpu if thinker_result_ids_cpu is not None else thinker_result_ids
+        im_start_indexes = torch.nonzero(parse_ids[0] == self.config.im_start_token_id).flatten().tolist()
+        im_start_indexes.append(target_len)
         multimodal_mask = (
-            (thinker_result_ids == self.thinker_config.audio_token_id) |
-            (thinker_result_ids == self.thinker_config.image_token_id) |
-            (thinker_result_ids == self.thinker_config.video_token_id)
-        ).to(input_ids.device)  # [t] # fmt: skip
+            (result_ids == self.thinker_config.audio_token_id) |
+            (result_ids == self.thinker_config.image_token_id) |
+            (result_ids == self.thinker_config.video_token_id)
+        )  # [t] # fmt: skip
+        role_tokens = parse_ids[0].tolist()
 
         tts_bos_embed, tts_eos_embed, tts_pad_embed = self._get_tts_embed(
             thinker_embed, tts_bos_thinker, tts_eos_thinker, tts_pad_thinker
@@ -1045,21 +1246,21 @@ class Qwen3OmniMoeForConditionalGeneration(
         trailing_text_hidden_all: torch.Tensor | None = None
         # For every chatml parts
         for i in range(len(im_start_indexes) - 1):
-            im_start_index = im_start_indexes[i].item()
-            segment_end_index = im_start_indexes[i + 1].item()
-            role_token = input_ids[0][im_start_index + 1]
+            im_start_index = im_start_indexes[i]
+            segment_end_index = im_start_indexes[i + 1]
+            role_token = role_tokens[im_start_index + 1]
             # Talker should ignore thinker system prompt
-            if (role_token == self.config.system_token_id).item():
+            if role_token == self.config.system_token_id:
                 continue
             # Talker takes word embeddings for tokens and hidden state from `accept_hidden_layer` for multimodal inputs
-            elif (role_token == self.config.user_token_id).item():
+            elif role_token == self.config.user_token_id:
                 talker_user_part = self._get_talker_user_parts(
                     im_start_index, segment_end_index, multimodal_mask, thinker_hidden, thinker_embed
                 )
                 talker_input_embeds.append(talker_user_part)
                 talker_input_ids.append(thinker_result_ids[im_start_index:segment_end_index])
             # Take assistant output (for now)
-            elif (role_token == self.config.assistant_token_id).item() and i == len(im_start_indexes) - 2:
+            elif role_token == self.config.assistant_token_id and i == len(im_start_indexes) - 2:
                 talker_assistant_embeds, talker_assistant_ids, trailing_text_hidden = self._get_talker_assistant_parts(
                     im_start_index,
                     segment_end_index,
@@ -1078,7 +1279,7 @@ class Qwen3OmniMoeForConditionalGeneration(
                 except Exception:
                     pass
             # History assistant output (ignore for now)
-            elif (role_token == self.config.assistant_token_id).item() and i != len(im_start_indexes) - 2:
+            elif role_token == self.config.assistant_token_id and i != len(im_start_indexes) - 2:
                 continue
             else:
                 raise AssertionError("Expect role id after <|im_start|> (assistant, user, system)")
@@ -1103,20 +1304,22 @@ class Qwen3OmniMoeForConditionalGeneration(
 
         cached_thinker_decode_embeds = embed.get("cached_decode", None)
         thinker_decode_embed = embed.get("decode", None)
+        if getattr(self.vllm_config.model_config, "use_v2_model_runner", False) and not meta.get("resumable", False):
+            return self._next_thinker_decode_text_step(
+                cached_thinker_decode_embeds, thinker_decode_embed, meta, device, update_dict
+            )
         start_index = meta.get("num_processed_tokens", 0)
 
         if cached_thinker_decode_embeds is not None and start_index < cached_thinker_decode_embeds.shape[0]:
-            cached_thinker_decode_embeds = cached_thinker_decode_embeds.to(device)
+            cached_thinker_decode_embeds = to_device_nonblocking(cached_thinker_decode_embeds, device)
             thinker_embed = cached_thinker_decode_embeds[start_index]
             if thinker_decode_embed is not None:
-                thinker_decode_embed = thinker_decode_embed.to(device)
+                thinker_decode_embed = to_device_nonblocking(thinker_decode_embed, device)
                 cached_thinker_decode_embeds = torch.cat([cached_thinker_decode_embeds, thinker_decode_embed], dim=0)
                 update_dict.setdefault("embed", {})["cached_decode"] = cached_thinker_decode_embeds
 
         elif thinker_decode_embed is not None:
-            thinker_embed = thinker_decode_embed
-            if thinker_embed.device != device:
-                thinker_embed = thinker_embed.to(device)
+            thinker_embed = to_device_nonblocking(thinker_decode_embed, device)
 
         else:
             # When the tokens output by the thinker are exhausted, an EOS token needs to be appended.
@@ -1129,6 +1332,41 @@ class Qwen3OmniMoeForConditionalGeneration(
         update_dict.setdefault("embed", {})["decode"] = None
         return self.talker.text_projection(thinker_embed).to(device)
 
+    def _next_thinker_decode_text_step(
+        self,
+        pending: torch.Tensor | None,
+        incoming: torch.Tensor | None,
+        meta: OmniPayloadMeta,
+        device: torch.device,
+        update_dict: OmniPayload,
+    ) -> torch.Tensor:
+        """Take the oldest unconsumed Thinker decode row as this step's text.
+
+        Rows arrive in order but not one per Talker step: the MRv2 receiver
+        hands over every row that arrived since the previous step (the V1
+        adapter delivers one chunk per step). New rows are appended to the
+        pending queue ``embed.cached_decode`` and each step consumes one. Once
+        the queue is empty the Thinker's text is exhausted: EOS once, then pad.
+        """
+        embed_update = update_dict.setdefault("embed", {})
+        embed_update["decode"] = None
+        if isinstance(incoming, torch.Tensor) and incoming.numel() > 0:
+            incoming = to_device_nonblocking(incoming, device).reshape(-1, incoming.shape[-1])
+            if isinstance(pending, torch.Tensor) and pending.shape[0] > 0:
+                pending = to_device_nonblocking(pending, device)
+                pending = torch.cat([pending, incoming.to(pending.dtype)], dim=0)
+            else:
+                pending = incoming
+        if not isinstance(pending, torch.Tensor) or pending.shape[0] == 0:
+            embed_update["cached_decode"] = None
+            if meta.get("eos_emitted", False):
+                return self.tts_pad_embed.to(device)
+            update_dict.setdefault("meta", {})["eos_emitted"] = True
+            return self.tts_eos_embed.to(device)
+        pending = to_device_nonblocking(pending, device)
+        embed_update["cached_decode"] = pending[1:] if pending.shape[0] > 1 else None
+        return self.talker.text_projection(pending[:1]).to(device)
+
     def talker_preprocess_decode(
         self, input_ids: torch.Tensor, input_embeds: torch.Tensor, update_dict: OmniPayload, payload: OmniPayload
     ):
@@ -1203,6 +1441,21 @@ class Qwen3OmniMoeForConditionalGeneration(
         )
 
         user_mm_mask = multimodal_mask[im_start_index:segment_end_index]
+        if user_mm_mask.device.type == "cpu":
+            # Host mask: no device sync for the test, and a text-only segment
+            # (the common TTS case) needs no masked gather at all.
+            if not bool(user_mm_mask.any()):
+                user_thinker_embed = thinker_embed[im_start_index:segment_end_index]
+                return self.talker.text_projection(user_thinker_embed).to(thinker_hidden.device, torch.bfloat16)
+            mm_index = to_device_nonblocking(torch.nonzero(user_mm_mask).flatten(), thinker_hidden.device)
+            text_index = to_device_nonblocking(torch.nonzero(~user_mm_mask).flatten(), thinker_hidden.device)
+            segment_hidden = thinker_hidden[im_start_index:segment_end_index]
+            segment_embed = thinker_embed[im_start_index:segment_end_index]
+            mm_hidden = self.talker.hidden_projection(segment_hidden.index_select(0, mm_index))
+            user_talker_part.index_copy_(0, mm_index, mm_hidden.to(user_talker_part.dtype))
+            text_hidden = self.talker.text_projection(segment_embed.index_select(0, text_index))
+            user_talker_part.index_copy_(0, text_index, text_hidden.to(user_talker_part.dtype))
+            return user_talker_part
         # Multimodal data exists
         if user_mm_mask.any():
             user_thinker_hidden_mm = thinker_hidden[im_start_index:segment_end_index][user_mm_mask]
@@ -1236,17 +1489,19 @@ class Qwen3OmniMoeForConditionalGeneration(
             ),
             dim=0,
         )
-        codec_special_tokens = torch.tensor(
-            [
-                self.config.talker_config.codec_nothink_id,
-                self.config.talker_config.codec_think_bos_id,
-                self.config.talker_config.codec_think_eos_id,
-                speaker_id,
-                self.config.talker_config.codec_pad_id,
-                self.config.talker_config.codec_bos_id,
-            ],
-            device=tts_pad_embed.device,
-            dtype=torch.long,
+        codec_special_tokens = to_device_nonblocking(
+            torch.tensor(
+                [
+                    self.config.talker_config.codec_nothink_id,
+                    self.config.talker_config.codec_think_bos_id,
+                    self.config.talker_config.codec_think_eos_id,
+                    speaker_id,
+                    self.config.talker_config.codec_pad_id,
+                    self.config.talker_config.codec_bos_id,
+                ],
+                dtype=torch.long,
+            ),
+            tts_pad_embed.device,
         )
         embed_input_ids = self.talker.embed_input_ids(codec_special_tokens).to(
             device=tts_pad_embed.device, dtype=torch.bfloat16
@@ -1381,18 +1636,31 @@ class Qwen3OmniMoeForConditionalGeneration(
 
     # ==================== Weight Loading ====================
 
-    def _get_codec_frame_config(self) -> tuple[int, int]:
-        """Extract codec_chunk_frames and codec_left_context_frames from stage connector config."""
+    def _get_codec_connector_extra(self) -> dict[str, Any]:
         model_cfg = getattr(self.vllm_config, "model_config", None)
         connector_cfg = getattr(model_cfg, "stage_connector_config", None)
         if isinstance(connector_cfg, dict):
             extra = connector_cfg.get("extra", {})
         else:
             extra = getattr(connector_cfg, "extra", None) or {}
+        return extra if isinstance(extra, Mapping) else {}
+
+    def _get_codec_frame_config(self) -> tuple[int, int]:
+        """Extract codec_chunk_frames and codec_left_context_frames from stage connector config."""
+        extra = self._get_codec_connector_extra()
         chunk_frames = int(extra.get("codec_chunk_frames", 0) or 0)
         left_frames = int(extra.get("codec_left_context_frames", 0) or 0)
         return chunk_frames, left_frames
 
+    def _get_codec_ramp_windows(self) -> list[int]:
+        """Decode windows of the connector's ``codec_chunk_ramp`` (empty without a ramp)."""
+        extra = self._get_codec_connector_extra()
+        chunk_frames = int(extra.get("codec_chunk_frames", 25) or 25)
+        ramp = parse_chunk_ramp(dict(extra), steady=chunk_frames)
+        if not ramp:
+            return []
+        return ramp_decode_windows(ramp, int(extra.get("codec_left_context_frames", 25) or 0))
+
     def _maybe_enable_code2wav_cudagraph(self) -> None:
         """Enable the inner Code2Wav CUDA graph unless this stage runs in eager mode."""
         if not self.code2wav or not hasattr(self.code2wav, "enable_cudagraph"):
@@ -1404,9 +1672,15 @@ class Qwen3OmniMoeForConditionalGeneration(
             return
 
         chunk_frames, left_frames = self._get_codec_frame_config()
+        streaming = bool(self.vllm_config.model_config.async_chunk)
+        # Opt-in multi-row graphs for streaming windows (connector extra
+        # ``codec_graph_batch_sizes``, e.g. [2, 3, 4, 6, 8, 12, 16]).
+        batch_sizes = self._get_codec_connector_extra().get("codec_graph_batch_sizes") or ()
         self.code2wav.enable_cudagraph(
             codec_chunk_frames=chunk_frames,
             codec_left_context_frames=left_frames,
+            extra_capture_sizes=self._get_codec_ramp_windows() if streaming else (),
+            streaming_batch_sizes=[int(b) for b in batch_sizes] if streaming else (),
         )
 
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
@@ -1439,6 +1713,12 @@ class Qwen3OmniMoeForConditionalGeneration(
             loaded_weights.update(talker_loaded)
             loaded_weights.update(self._init_special_tokens_embeddings())
 
+        # The Talker's first-frame decoder is a Code2Wav copy (same code2wav.* weights).
+        first_frame_decoder = getattr(self, "first_frame_decoder", None)
+        if first_frame_decoder is not None and code2wav_weights:
+            decoder_loaded = first_frame_decoder.load_weights(code2wav_weights)
+            loaded_weights.update(add_prefix_to_loaded_weights(decoder_loaded, "first_frame_decoder.code2wav"))
+
         # Load code2wav weights
         if self.code2wav and code2wav_weights:
             code2wav_loaded = self.code2wav.load_weights(code2wav_weights)

vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_code2wav.py

@@ -5,7 +5,8 @@
 
 from __future__ import annotations
 
-from collections.abc import Iterable
+from collections.abc import Callable, Iterable, Sequence
+from contextlib import nullcontext
 
 import numpy as np
 import torch
@@ -19,6 +20,7 @@ from transformers.models.qwen3_omni_moe.modeling_qwen3_omni_moe import (
     Qwen3OmniMoeCode2WavDecoderBlock,
     Qwen3OmniMoeCode2WavTransformerModel,
     Qwen3OmniMoeConvNeXtBlock,
+    Qwen3OmniMoeSnakeBeta,
 )
 from vllm.config import VllmConfig  # type: ignore
 from vllm.logger import init_logger  # type: ignore
@@ -31,9 +33,93 @@ from vllm_omni.model_executor.models.common.snake_activation import SnakeBeta
 from vllm_omni.model_executor.models.qwen3_omni.quantization import (
     Qwen3OmniNestedSupportsQuant,
 )
+from vllm_omni.platforms import current_omni_platform
 
 logger = init_logger(__name__)
 
+# Fixed cost of one streaming decode call (graph replay or eager launch), in
+# decoded-frame equivalents: a call on a single frame costs about as much GPU
+# time as decoding this many extra frames inside a larger call.
+_DECODE_CALL_COST_FRAMES = 32
+
+
+def use_fused_snake(module: nn.Module) -> int:
+    """Replace HF ``Qwen3OmniMoeSnakeBeta`` activations under ``module`` with the fused one.
+
+    The HF decoder blocks recompute ``exp(alpha)`` / ``exp(beta)`` and launch
+    about seven broadcast kernels per call; the shared ``SnakeBeta`` computes the
+    same ``x + 1/b * sin^2(a*x)`` in one kernel from precomputed caches.
+    Parameter names and shapes match, so this runs before weight loading.
+    """
+    count = 0
+    for parent in list(module.modules()):
+        for name, child in list(parent.named_children()):
+            if isinstance(child, Qwen3OmniMoeSnakeBeta):
+                setattr(parent, name, SnakeBeta(child.in_features))
+                count += 1
+    return count
+
+
+def plan_decode_groups(
+    lengths: Sequence[int],
+    bucket_of: Callable[[int], int | None],
+    batch_sizes_for: Callable[[int], Sequence[int]],
+    call_cost_frames: int = _DECODE_CALL_COST_FRAMES,
+) -> list[tuple[list[int], int]]:
+    """Split a streaming batch into decode calls of rows with similar lengths.
+
+    Decoding every row at the batch's longest window wastes most of the work
+    when short ramp chunks (1-15 frames) share a step with steady 50-frame
+    windows. Rows are sorted by length and cut into contiguous groups of
+    graph buckets (``bucket_of``) so that the modelled cost, per call
+    ``call_cost_frames`` plus padded rows (to the next captured row count of
+    that bucket, ``batch_sizes_for``) times the bucket length, is smallest.
+    Returns ``(row indices, decode length)`` per call; the decode length is
+    the group's longest row.
+    """
+    count = len(lengths)
+    if count == 0:
+        return []
+    order = sorted(range(count), key=lambda row: lengths[row])
+    buckets = [bucket_of(int(lengths[row])) or int(lengths[row]) for row in order]
+    distinct = sorted(set(buckets))
+    members = [[row for row, bucket in zip(order, buckets) if bucket == size] for size in distinct]
+
+    def sizes_for(size: int) -> list[int]:
+        return sorted({int(b) for b in batch_sizes_for(size) if int(b) > 0} | {1})
+
+    def cost(rows: int, size: int) -> int:
+        batch_sizes = sizes_for(size)
+        calls, rest = divmod(rows, batch_sizes[-1])
+        padded = calls * batch_sizes[-1]
+        if rest:
+            calls += 1
+            padded += next(b for b in batch_sizes if b >= rest)
+        return calls * call_cost_frames + padded * size
+
+    best = [0] + [None] * len(distinct)
+    cut = [0] * (len(distinct) + 1)
+    for end in range(1, len(distinct) + 1):
+        rows = 0
+        for start in range(end - 1, -1, -1):
+            rows += len(members[start])
+            total = best[start] + cost(rows, distinct[end - 1])
+            if best[end] is None or total < best[end]:
+                best[end], cut[end] = total, start
+    groups: list[tuple[list[int], int]] = []
+    end = len(distinct)
+    while end > 0:
+        start = cut[end]
+        groups.append(([row for bucket_rows in members[start:end] for row in bucket_rows], distinct[end - 1]))
+        end = start
+    calls: list[tuple[list[int], int]] = []
+    for rows, size in reversed(groups):
+        max_batch = sizes_for(size)[-1]
+        for offset in range(0, len(rows), max_batch):
+            part = rows[offset : offset + max_batch]
+            calls.append((part, max(int(lengths[row]) for row in part)))
+    return calls
+
 
 class Qwen3OmniMoeCode2Wav(nn.Module, Qwen3OmniNestedSupportsQuant):
     """
@@ -74,6 +160,14 @@ class Qwen3OmniMoeCode2Wav(nn.Module, Qwen3OmniNestedSupportsQuant):
         super().__init__()
 
         self.config: Qwen3OmniMoeCode2WavConfig = vllm_config.model_config.hf_config
+        connector = getattr(vllm_config.model_config, "stage_connector_config", None)
+        extra = connector.get("extra", {}) if isinstance(connector, dict) else getattr(connector, "extra", {})
+        # Both the separate stage and the Talker's first-frame copy receive
+        # this connector config. Keep algorithm selection consistent between
+        # them, and leave other models' cuDNN settings unchanged.
+        self._cudnn_benchmark = (extra or {}).get("codec_cudnn_benchmark", False)
+        if not isinstance(self._cudnn_benchmark, bool):
+            raise ValueError("codec_cudnn_benchmark must be a boolean")
 
         # Calculate total upsampling factor
         self.total_upsample = np.prod(self.config.upsample_rates + self.config.upsampling_ratios)
@@ -122,10 +216,19 @@ class Qwen3OmniMoeCode2Wav(nn.Module, Qwen3OmniNestedSupportsQuant):
             Qwen3OmniMoeCausalConvNet(output_dim, 1, kernel_size=7),
         ]
         self.decoder = nn.ModuleList(decoder)
+        # Keep existing V1 and non-CUDA decoder blocks unchanged.
+        if (
+            (extra or {}).get("codec_fused_snake", False)
+            and current_omni_platform.is_cuda()
+            and torch.device(vllm_config.device_config.device).type == "cuda"
+        ):
+            use_fused_snake(self.decoder)
 
         # CUDA Graph support — reuses CUDAGraphDecoderWrapper from Qwen3-TTS
         self._cudagraph_enabled = False
         self._cudagraph_wrapper = None
+        # Captured row counts per graph size; empty keeps one padded decode per batch.
+        self._streaming_batch_sizes: dict[int, list[int]] = {}
 
     def precompute_snake_caches(self):
         """Precompute exp(alpha) and 1/(exp(beta)+eps) for all SnakeBeta modules."""
@@ -142,8 +245,19 @@ class Qwen3OmniMoeCode2Wav(nn.Module, Qwen3OmniNestedSupportsQuant):
         device: torch.device | None = None,
         codec_chunk_frames: int = 0,
         codec_left_context_frames: int = 0,
+        extra_capture_sizes: Iterable[int] = (),
+        streaming_batch_sizes: Iterable[int] = (),
     ):
-        """Enable CUDA graph acceleration (same pattern as Qwen3-TTS Code2Wav)."""
+        """Enable CUDA graph acceleration (same pattern as Qwen3-TTS Code2Wav).
+
+        ``extra_capture_sizes`` adds exact-size graphs (e.g. a chunk ramp's
+        decode windows) to the default buckets, so those chunks replay without
+        zero padding. ``streaming_batch_sizes`` (e.g. ``[2, 4, 8, 16]``) also
+        captures multi-row graphs for every bucket up to one streaming window
+        (chunk + left context); a streaming batch is then decoded as a few
+        length groups (``plan_decode_groups``) instead of one eager call
+        padded to its longest row.
+        """
         from vllm_omni.model_executor.models.qwen3_tts.cuda_graph_decoder_wrapper import (
             CUDAGraphDecoderWrapper,
         )
@@ -154,8 +268,30 @@ class Qwen3OmniMoeCode2Wav(nn.Module, Qwen3OmniNestedSupportsQuant):
             logger.warning("Cannot enable CUDA Graph: not on CUDA device (got %s)", device)
             return
 
+        extra = {int(size) for size in extra_capture_sizes if int(size) > 0}
+        batch_sizes = sorted({int(b) for b in streaming_batch_sizes if int(b) > 1})
+        capture_sizes = None
+        extra_shapes: list[tuple[int, int]] = []
+        if extra or batch_sizes:
+            capture_sizes = sorted(
+                set(
+                    CUDAGraphDecoderWrapper.compute_capture_sizes(
+                        codec_chunk_frames=codec_chunk_frames,
+                        codec_left_context_frames=codec_left_context_frames,
+                    )
+                )
+                | extra
+            )
+            window = codec_chunk_frames + codec_left_context_frames
+            streaming_sizes = [size for size in capture_sizes if window <= 0 or size <= window]
+            # A multi-row graph never holds more frames than the largest
+            # single-row one, so the shared graph pool does not grow.
+            largest = max(capture_sizes)
+            extra_shapes = [(b, size) for b in batch_sizes for size in streaming_sizes if b * size <= largest]
         wrapper = CUDAGraphDecoderWrapper(
             decoder=self,
+            capture_sizes=capture_sizes,
+            extra_capture_shapes=extra_shapes,
             num_quantizers=self.config.num_quantizers,
             enabled=True,
         )
@@ -172,10 +308,18 @@ class Qwen3OmniMoeCode2Wav(nn.Module, Qwen3OmniNestedSupportsQuant):
             raise
         self._cudagraph_wrapper = wrapper
         self._cudagraph_enabled = True
+        if extra_shapes and torch.cuda.is_available():
+            # Return the eager warm-up activations: this stage shares the GPU.
+            torch.accelerator.empty_cache()
+        self._streaming_batch_sizes = {
+            size: [1, *sorted(b for b, shape_size in extra_shapes if shape_size == size)]
+            for size in {shape_size for _b, shape_size in extra_shapes}
+        }
         logger.info(
-            "CUDA Graph enabled for Code2Wav: num_quantizers=%d, sizes=%s",
+            "CUDA Graph enabled for Code2Wav: num_quantizers=%d, sizes=%s, multi-row shapes=%s",
             self.config.num_quantizers,
             self._cudagraph_wrapper.capture_sizes,
+            extra_shapes,
         )
 
     def forward(self, codes: torch.Tensor) -> torch.Tensor:
@@ -200,6 +344,30 @@ class Qwen3OmniMoeCode2Wav(nn.Module, Qwen3OmniNestedSupportsQuant):
         hidden = self.pre_transformer(inputs_embeds=hidden).last_hidden_state
         # Shape: [batch, seq_len, hidden_size]
 
+        # Algorithm search happens during eager warmup before graph capture.
+        # Restore process settings afterward, including on a failed decode.
+        # Captured replays use the selected kernels without this Python path.
+        cudnn = torch.backends.cudnn
+        context = (
+            cudnn.flags(
+                enabled=cudnn.enabled,
+                benchmark=True,
+                benchmark_limit=10,
+                deterministic=cudnn.deterministic,
+                allow_tf32=cudnn.allow_tf32,
+            )
+            if (
+                self._cudnn_benchmark
+                and hidden.is_cuda
+                and torch.version.hip is None
+                and not torch.cuda.is_current_stream_capturing()
+            )
+            else nullcontext()
+        )
+        with context:
+            return self._decode_waveform(hidden)
+
+    def _decode_waveform(self, hidden: torch.Tensor) -> torch.Tensor:
         # Stage 3: Upsampling
         hidden = hidden.permute(0, 2, 1)  # [batch, hidden_size, seq_len]
         for blocks in self.upsample:
@@ -305,28 +473,38 @@ class Qwen3OmniMoeCode2Wav(nn.Module, Qwen3OmniNestedSupportsQuant):
                 "defaulting to left_context_size=zeros(len(codes)). This is expected during cudagraph warmup."
             )
             left_context_size = [0] * codes.shape[0]
-        # Decode chunk
-        wavs = []
-        if self._cudagraph_enabled and self._cudagraph_wrapper is not None:
-            batch_wav = self._cudagraph_wrapper.decode(codes)
-        else:
-            batch_wav = self(codes)
         if seq_token_counts is not None:
             code_seq_lens = [n // self.config.num_quantizers for n in seq_token_counts]
         else:
             # Fallback: assume all batch elements share the same sequence length.
             code_seq_lens = [codes.shape[-1]] * codes.shape[0]
-        # The decoder trims the same right-edge tail from every batch row.
-        # Infer it from the decoded window, not a shorter request's length:
-        # padding that request with codec zeros does not provide valid context.
-        tail = max(0, int(codes.shape[-1] * self.total_upsample) - batch_wav.shape[-1])
-        for idx, code_seq_len in enumerate(code_seq_lens):
-            # Refill the previous chunk's tail using the re-decoded left context,
-            # and exclude the current row's padded tail just as singleton decode does.
-            start = max(0, left_context_size[idx] * self.total_upsample - tail)
-            end = max(0, code_seq_len * self.total_upsample - tail)
-            wav_chunk = batch_wav[idx, :, start:end]
-            wavs.append(wav_chunk)
+        batch, width = int(codes.shape[0]), int(codes.shape[-1])
+        graphs = self._cudagraph_wrapper if self._cudagraph_enabled else None
+        batch_sizes = getattr(self, "_streaming_batch_sizes", None)
+        if graphs is not None and batch_sizes and batch > 1:
+            groups = plan_decode_groups(
+                code_seq_lens, graphs._get_padded_size, lambda size: batch_sizes.get(size, (1,))
+            )
+        else:
+            groups = [(list(range(batch)), width)]
+        wavs: list[torch.Tensor | None] = [None] * batch
+        for rows, length in groups:
+            if len(rows) == batch and length == width:
+                group_codes = codes
+            else:
+                # Basic indexing only: a device index tensor would need a host copy.
+                group_codes = torch.stack([codes[row, :, :length] for row in rows])
+            batch_wav = graphs.decode(group_codes) if graphs is not None else self(group_codes)
+            # The decoder trims the same right-edge tail from every batch row.
+            # Infer it from the decoded window, not a shorter request's length:
+            # padding that request with codec zeros does not provide valid context.
+            tail = max(0, int(length * self.total_upsample) - batch_wav.shape[-1])
+            for position, row in enumerate(rows):
+                # Refill the previous chunk's tail using the re-decoded left context,
+                # and exclude the current row's padded tail just as singleton decode does.
+                start = max(0, left_context_size[row] * self.total_upsample - tail)
+                end = max(0, code_seq_lens[row] * self.total_upsample - tail)
+                wavs[row] = batch_wav[position, :, start:end]
         return wavs
 
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:

vllm_omni/model_executor/models/qwen3_tts/cached_code_predictor.py

@@ -18,8 +18,7 @@ from vllm_omni.model_executor.models.common.qwen3_code_predictor import (
     CodePredictorDecoderLayer,
     _rotate_half,
 )
-
-from .short_kv_attention import short_kv_attention
+from vllm_omni.model_executor.models.common.short_kv_attention import short_kv_attention
 
 
 def _attention_cached(

vllm_omni/model_executor/models/qwen3_tts/first_audio.py

@@ -4,9 +4,7 @@
 
 from typing import Any
 
-import torch
-
-from vllm_omni.platforms import current_omni_platform
+from vllm_omni.model_executor.models.common.talker_first_audio import supports_talker_first_audio
 
 
 def talker_first_audio_enabled(vllm_config: Any) -> bool:
@@ -18,16 +16,4 @@ def talker_first_audio_enabled(vllm_config: Any) -> bool:
     from .qwen3_tts_code_predictor_vllm import Qwen3TTSTalkerCodePredictorForConditionalGenerationVLLM as Predictor
 
     extra = Predictor._stage_connector_extra_config(vllm_config)
-    model = vllm_config.model_config
-    parallel = vllm_config.parallel_config
-    return (
-        Predictor._parse_bool_config(extra.get("talker_first_audio"))
-        and current_omni_platform.is_cuda()
-        and torch.device(vllm_config.device_config.device).type == "cuda"
-        and bool(getattr(model, "use_v2_model_runner", False))
-        and bool(getattr(model, "async_chunk", False))
-        and parallel.tensor_parallel_size == 1
-        and parallel.pipeline_parallel_size == 1
-        and parallel.distributed_executor_backend in (None, "uni")
-        and not vllm_config.cache_config.enable_prefix_caching
-    )
+    return Predictor._parse_bool_config(extra.get("talker_first_audio")) and supports_talker_first_audio(vllm_config)

vllm_omni/model_executor/models/qwen3_tts/fused_code_predictor.py

@@ -0,0 +1,312 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+# ruff: noqa: N803 - Triton constexpr parameters use kernel-style capitals.
+"""Fused residual-codebook predictor for Qwen3-TTS (CUDA, BF16).
+
+Same computation as the frame-local-KV predictor with per-call top-k Gumbel
+sampling, restructured to issue few kernels per residual step, since at serving
+batch sizes each of the 15 dependent steps is launch- and latency-bound:
+
+- Per layer: RMSNorm, one QKV GEMM, one attention kernel (q/k RMSNorm, RoPE,
+  this call's K/V write and attention over the frame's keys), O projection
+  accumulated into the residual stream (``addmm_``), RMSNorm, gate/up GEMM,
+  SiLU-mul, down projection accumulated into the stream.
+- The embedding of each sampled code followed by ``small_to_mtp_projection``
+  is linear, so it folds at load into one ``[vocab, hidden]`` table per
+  codebook; the sampling kernel writes the next step's input row directly.
+- Sampling finds the top-k threshold by bisection over the BF16 order keys
+  (exact k-th largest, ties kept) instead of a sort.
+"""
+
+from __future__ import annotations
+
+import torch
+from vllm.triton_utils import tl, tldevice, triton
+
+from .tokenizer_12hz.streaming_decoder import _silu_mul_kernel
+
+
+@triton.jit
+def _rms_rows_kernel(x_ptr, x_stride, w_ptr, out_ptr, eps, N: tl.constexpr):
+    row = tl.program_id(0).to(tl.int64)
+    offs = tl.arange(0, N)
+    x = tl.load(x_ptr + row * x_stride + offs).to(tl.float32)
+    var = tl.sum(x * x, axis=0) / N
+    dt = out_ptr.dtype.element_ty
+    y = (x * tl.rsqrt(var + eps)).to(dt).to(tl.float32)
+    w = tl.load(w_ptr + offs).to(tl.float32)
+    tl.store(out_ptr + row * N + offs, (w * y).to(dt))
+
+
+@triton.jit
+def _bf(x):
+    return x.to(tl.bfloat16).to(tl.float32)
+
+
+@triton.jit
+def _head_norm_rope(x1, x2, w1, w2, c1, c2, s1, s2, eps, HD: tl.constexpr):
+    """HF RMSNorm (``weight * normed.to(bf16)``) then RoPE in BF16 elementwise steps, on the two halves."""
+    var = (tl.sum(x1 * x1, axis=0) + tl.sum(x2 * x2, axis=0)) / HD
+    r = tl.rsqrt(var + eps)
+    n1 = _bf(w1 * _bf(x1 * r))
+    n2 = _bf(w2 * _bf(x2 * r))
+    # rotate_half(n) = [-n2, n1]
+    o1 = _bf(_bf(n1 * c1) + _bf(-n2 * s1))
+    o2 = _bf(_bf(n2 * c2) + _bf(n1 * s2))
+    return o1, o2
+
+
+@triton.jit
+def _cp_attention_kernel(
+    qkv_ptr,  # [B * NQ, (H + 2 * KVH) * HD]
+    kc_ptr,  # [Bmax, KVH, MAXP, HD]
+    vc_ptr,
+    qw_ptr,
+    kw_ptr,
+    cos_ptr,  # [MAXP, HD] bf16
+    sin_ptr,
+    out_ptr,  # [B * NQ, H * HD]
+    first,
+    eps,
+    scale,
+    NQ: tl.constexpr,
+    H: tl.constexpr,
+    KVH: tl.constexpr,
+    HD: tl.constexpr,
+    MAXP: tl.constexpr,
+    BK: tl.constexpr,
+):
+    b = tl.program_id(0)
+    h = tl.program_id(1)
+    qi = tl.program_id(2)
+    kvh = h // (H // KVH)
+    W: tl.constexpr = (H + 2 * KVH) * HD
+    HALF: tl.constexpr = HD // 2
+    d = tl.arange(0, HALF)
+    p = first + qi
+    row = (b * NQ + qi).to(tl.int64)
+    qw1 = tl.load(qw_ptr + d).to(tl.float32)
+    qw2 = tl.load(qw_ptr + HALF + d).to(tl.float32)
+    kw1 = tl.load(kw_ptr + d).to(tl.float32)
+    kw2 = tl.load(kw_ptr + HALF + d).to(tl.float32)
+    qb = qkv_ptr + row * W + h * HD
+    q1, q2 = _head_norm_rope(
+        tl.load(qb + d).to(tl.float32), tl.load(qb + HALF + d).to(tl.float32), qw1, qw2,
+        tl.load(cos_ptr + p * HD + d).to(tl.float32), tl.load(cos_ptr + p * HD + HALF + d).to(tl.float32),
+        tl.load(sin_ptr + p * HD + d).to(tl.float32), tl.load(sin_ptr + p * HD + HALF + d).to(tl.float32),
+        eps, HD,
+    )  # fmt: skip
+    j = tl.arange(0, BK)
+    cache = ((b * KVH + kvh) * MAXP).to(tl.int64) * HD
+    old = (j < first)[:, None]
+    kp = kc_ptr + cache + j[:, None] * HD + d[None, :]
+    vp = vc_ptr + cache + j[:, None] * HD + d[None, :]
+    k1 = tl.load(kp, mask=old, other=0.0).to(tl.float32)
+    k2 = tl.load(kp + HALF, mask=old, other=0.0).to(tl.float32)
+    v1 = tl.load(vp, mask=old, other=0.0).to(tl.float32)
+    v2 = tl.load(vp + HALF, mask=old, other=0.0).to(tl.float32)
+    for n in tl.static_range(NQ):
+        # this call's position first + n, computed from its own QKV row
+        on = first + n
+        kb = qkv_ptr + (b * NQ + n).to(tl.int64) * W + (H + kvh) * HD
+        vb = qkv_ptr + (b * NQ + n).to(tl.int64) * W + (H + KVH + kvh) * HD
+        kn1, kn2 = _head_norm_rope(
+            tl.load(kb + d).to(tl.float32), tl.load(kb + HALF + d).to(tl.float32), kw1, kw2,
+            tl.load(cos_ptr + on * HD + d).to(tl.float32), tl.load(cos_ptr + on * HD + HALF + d).to(tl.float32),
+            tl.load(sin_ptr + on * HD + d).to(tl.float32), tl.load(sin_ptr + on * HD + HALF + d).to(tl.float32),
+            eps, HD,
+        )  # fmt: skip
+        vn1 = tl.load(vb + d).to(tl.float32)
+        vn2 = tl.load(vb + HALF + d).to(tl.float32)
+        sel = (j == on)[:, None]
+        k1 = tl.where(sel, kn1[None, :], k1)
+        k2 = tl.where(sel, kn2[None, :], k2)
+        v1 = tl.where(sel, vn1[None, :], v1)
+        v2 = tl.where(sel, vn2[None, :], v2)
+        wm = (d < HALF) & (qi == n) & (h % (H // KVH) == 0)
+        dtc = kc_ptr.dtype.element_ty
+        tl.store(kc_ptr + cache + on * HD + d, kn1.to(dtc), mask=wm)
+        tl.store(kc_ptr + cache + on * HD + HALF + d, kn2.to(dtc), mask=wm)
+        tl.store(vc_ptr + cache + on * HD + d, vn1.to(dtc), mask=wm)
+        tl.store(vc_ptr + cache + on * HD + HALF + d, vn2.to(dtc), mask=wm)
+    km = j <= p
+    s = (tl.sum(k1 * q1[None, :], axis=1) + tl.sum(k2 * q2[None, :], axis=1)) * scale
+    s = tl.where(km, s, -float("inf"))
+    e = tl.exp(s - tl.max(s, axis=0))
+    e = tl.where(km, e, 0.0)
+    w = e / tl.sum(e, axis=0)
+    ob = out_ptr + row * (H * HD) + h * HD
+    dto = out_ptr.dtype.element_ty
+    tl.store(ob + d, tl.sum(v1 * w[:, None], axis=0).to(dto))
+    tl.store(ob + HALF + d, tl.sum(v2 * w[:, None], axis=0).to(dto))
+
+
+@triton.jit
+def _cp_sample_kernel(
+    logits_ptr,  # [B, V] bf16
+    u_ptr,  # uniforms, row stride u_stride, contiguous over V
+    u_stride,
+    codes_ptr,  # [B, G] int64
+    step,
+    G: tl.constexpr,
+    table_ptr,  # [V, HID] next-step input table, or dummy
+    next_ptr,  # [B, HID]
+    inv_temperature,
+    V: tl.constexpr,
+    TOPK: tl.constexpr,
+    HID: tl.constexpr,
+    HAS_NEXT: tl.constexpr,
+):
+    row = tl.program_id(0).to(tl.int64)
+    idx = tl.arange(0, V)
+    value = tl.load(logits_ptr + row * V + idx)
+    scaled = (value.to(tl.float32) * inv_temperature).to(value.dtype)
+    if TOPK > 0:
+        # order-preserving 16-bit keys of the BF16 values
+        bits = scaled.to(tl.int16, bitcast=True).to(tl.int32) & 0xFFFF
+        # Top-k keeps all numeric ties at the cutoff, including both zeros.
+        bits = tl.where((bits & 0x7FFF) == 0, 0, bits)
+        key = tl.where(bits >= 0x8000, 0xFFFF - bits, bits | 0x8000)
+        the = 0
+        for bit in tl.static_range(15, -1, -1):
+            cand = the | (1 << bit)
+            cnt = tl.sum((key >= cand).to(tl.int32), axis=0)
+            the = tl.where(cnt >= TOPK, cand, the)
+        sf = tl.where(key >= the, scaled.to(tl.float32), -float("inf"))
+    else:
+        sf = scaled.to(tl.float32)
+    u = tl.load(u_ptr + row * u_stride + idx)
+    score = sf - tldevice.log(-tldevice.log(u))
+    max_score = tl.max(score, 0)
+    chosen = tl.min(tl.where(score == max_score, idx, V), 0)
+    first_nan = tl.min(tl.where(score != score, idx, V), 0)
+    code = tl.where(first_nan < V, first_nan, chosen)
+    tl.store(codes_ptr + row * G + step, code.to(tl.int64))
+    if HAS_NEXT:
+        c = tl.arange(0, HID)
+        tl.store(next_ptr + row * HID + c, tl.load(table_ptr + code.to(tl.int64) * HID + c))
+
+
+class FusedCodePredictor:
+    """Execution object built from a loaded ``Qwen3TTSTalkerCodePredictorForConditionalGenerationVLLM``."""
+
+    def __init__(self, predictor, max_batch: int) -> None:
+        model = predictor.model
+        weight = next(model.parameters())
+        dev, dt = weight.device, weight.dtype
+        if dt != torch.bfloat16:
+            raise ValueError("fused code predictor requires BF16 weights")
+        cfg = predictor.config
+        self.groups = int(cfg.num_code_groups)
+        self.vocab = int(cfg.vocab_size)
+        self.hidden = int(cfg.hidden_size)
+        attn0 = model.layers[0].self_attn
+        self.heads, self.kv_heads, self.head_dim = attn0.num_heads, attn0.num_kv_heads, attn0.head_dim
+        self.scale = attn0.scaling
+        self.eps = float(model.norm.variance_epsilon)
+        self.max_pos = attn0.max_seq
+        self.max_batch = max_batch
+        f32 = lambda t: t.detach().to(device=dev, dtype=torch.float32)  # noqa: E731
+        c = lambda t: t.to(dt).contiguous()  # noqa: E731
+        proj = predictor.small_to_mtp_projection
+        if isinstance(proj, torch.nn.Identity):
+            pw, pb = None, None
+        else:
+            pw, pb = f32(proj.weight), f32(proj.bias)
+        self.proj_wt = c(pw.t()) if pw is not None else None
+        self.proj_b = c(pb) if pb is not None else None
+        tables = []
+        for emb in model.codec_embedding:
+            e = f32(emb.weight)
+            tables.append(c(e @ pw.t() + pb) if pw is not None else c(e))
+        self.tables = tables
+        self.layers = []
+        for layer in model.layers:
+            a = layer.self_attn
+            self.layers.append(
+                dict(
+                    ln1=c(f32(layer.input_layernorm.weight)),
+                    ln2=c(f32(layer.post_attention_layernorm.weight)),
+                    eps1=float(layer.input_layernorm.variance_epsilon),
+                    eps2=float(layer.post_attention_layernorm.variance_epsilon),
+                    qkv=c(f32(a.qkv_proj.weight).t()),
+                    qn=c(f32(a.q_norm.weight)),
+                    kn=c(f32(a.k_norm.weight)),
+                    head_eps=float(a.q_norm.variance_epsilon),
+                    o=c(f32(a.o_proj.weight).t()),
+                    gu=c(f32(layer.mlp.gate_up_proj.weight).t()),
+                    down=c(f32(layer.mlp.down_proj.weight).t()),
+                )
+            )
+            if a.qkv_proj.bias is not None:
+                raise ValueError("fused code predictor assumes bias-free attention")
+        self.norm_w = c(f32(model.norm.weight))
+        self.lm_heads = [c(f32(h.weight).t()) for h in predictor.lm_head]
+        rot = model.rotary_emb
+        self.cos = rot.cos_cached.to(device=dev, dtype=dt).contiguous()
+        self.sin = rot.sin_cached.to(device=dev, dtype=dt).contiguous()
+        kv = (len(self.layers), 2, max_batch, self.kv_heads, self.max_pos, self.head_dim)
+        self.cache = torch.zeros(kv, device=dev, dtype=dt)
+        self._dummy = torch.zeros(1, device=dev, dtype=dt)
+
+    def _layer(self, x: torch.Tensor, li: int, B: int, nq: int, first: int) -> None:
+        L = self.layers[li]
+        rows = x.shape[0]
+        h = torch.empty_like(x)
+        _rms_rows_kernel[(rows,)](x, x.stride(0), L["ln1"], h, L["eps1"], N=self.hidden)
+        qkv = torch.mm(h, L["qkv"])
+        att = torch.empty(rows, self.heads * self.head_dim, device=x.device, dtype=x.dtype)
+        _cp_attention_kernel[(B, self.heads, nq)](
+            qkv, self.cache[li, 0], self.cache[li, 1], L["qn"], L["kn"], self.cos, self.sin, att,
+            first, L["head_eps"], self.scale,
+            NQ=nq, H=self.heads, KVH=self.kv_heads, HD=self.head_dim, MAXP=self.max_pos,
+            BK=triton.next_power_of_2(self.max_pos), num_warps=1,
+        )  # fmt: skip
+        x.addmm_(att, L["o"])
+        _rms_rows_kernel[(rows,)](x, x.stride(0), L["ln2"], h, L["eps2"], N=self.hidden)
+        gu = torch.mm(h, L["gu"])
+        inter = gu.shape[1] // 2
+        a = torch.empty(rows, inter, device=x.device, dtype=x.dtype)
+        _silu_mul_kernel[(rows, triton.cdiv(inter, 1024))](gu, a, N=inter, BN=1024)
+        x.addmm_(a, L["down"])
+
+    @torch.inference_mode()
+    def __call__(
+        self,
+        layer0_code: torch.Tensor,
+        layer0_embed: torch.Tensor,
+        last_talker_hidden: torch.Tensor,
+        inv_temperature: float,
+        top_k: int,
+        sample_uniforms: torch.Tensor,
+    ) -> torch.Tensor:
+        B = int(layer0_code.shape[0])
+        if B > self.max_batch:
+            raise ValueError("fused code predictor batch exceeds its capacity")
+        dev, dt = layer0_embed.device, self.norm_w.dtype
+        inp = torch.stack([last_talker_hidden.reshape(B, -1), layer0_embed.reshape(B, -1)], 1).to(dt)
+        inp = inp.reshape(2 * B, -1)
+        x = torch.addmm(self.proj_b, inp, self.proj_wt) if self.proj_wt is not None else inp.clone()
+        codes = torch.empty(B, self.groups, device=dev, dtype=torch.int64)
+        codes[:, 0] = layer0_code.reshape(B)
+        u = sample_uniforms
+        nq, first = 2, 0
+        for step in range(1, self.groups):
+            for li in range(len(self.layers)):
+                self._layer(x, li, B, nq, first)
+            last = x[nq - 1 :: nq] if nq > 1 else x
+            hf = torch.empty(B, self.hidden, device=dev, dtype=dt)
+            _rms_rows_kernel[(B,)](last, last.stride(0), self.norm_w, hf, self.eps, N=self.hidden)
+            logits = torch.mm(hf, self.lm_heads[step - 1])
+            has_next = step < self.groups - 1
+            nxt = torch.empty(B, self.hidden, device=dev, dtype=dt) if has_next else self._dummy
+            us = u[:, step - 1]
+            _cp_sample_kernel[(B,)](
+                logits, us, us.stride(0), codes, step, self.groups,
+                self.tables[step - 1] if has_next else self._dummy, nxt, inv_temperature,
+                V=self.vocab, TOPK=top_k, HID=self.hidden, HAS_NEXT=has_next, num_warps=4,
+            )  # fmt: skip
+            x = nxt
+            first = step + 1
+            nq = 1
+        return codes

vllm_omni/model_executor/models/qwen3_tts/pipeline.py

@@ -72,3 +72,33 @@ QWEN3_TTS_PIPELINE = PipelineConfig(
         ),
     ),
 )
+
+
+# Single-stage variant: the Talker decodes each frame with the stateful
+# streaming codec decoder and emits PCM as the final output (enable with the
+# ``talker_stream_decode`` model option). No Code2Wav stage, no connector.
+QWEN3_TTS_FUSED_PIPELINE = PipelineConfig(
+    model_type="qwen3_tts_fused",
+    default_deploy_config_name="qwen3_tts_fused_single_gpu.yaml",
+    single_stage_async_chunk=True,
+    model_arch="Qwen3TTSTalkerForConditionalGeneration",
+    stages=(
+        StagePipelineConfig(
+            stage_id=0,
+            model_stage="qwen3_tts",
+            execution_type=StageExecutionType.LLM_AR,
+            input_sources=(),
+            owns_tokenizer=True,
+            final_output=True,
+            final_output_type="audio",
+            # The API output processor tags ``model_outputs`` by this modality.
+            engine_output_type="audio",
+            supports_running_prefix_cache_reset=False,
+            supports_native_mrv2_data_plane=True,
+            sampling_constraints={
+                "detokenize": False,
+                "stop_token_ids": [2150],
+            },
+        ),
+    ),
+)

vllm_omni/model_executor/models/qwen3_tts/prompt_embeds_builder.py

@@ -42,6 +42,7 @@ from vllm.logger import init_logger
 from vllm.multimodal.audio import AudioResampler
 
 from vllm_omni.utils.audio import mel_filter_bank
+from vllm_omni.utils.device_copy import index_to_device, to_device_nonblocking
 
 if TYPE_CHECKING:
     from .configuration_qwen3_tts import Qwen3TTSConfig, Qwen3TTSTalkerConfig
@@ -160,7 +161,7 @@ def coerce_token_ids(value: object, *, device: torch.device) -> torch.Tensor | N
         ids = ids.unsqueeze(0)
     if ids.ndim != 2 or ids.numel() == 0:
         return None
-    return ids.to(device=device, dtype=torch.long).contiguous()
+    return to_device_nonblocking(ids.to(dtype=torch.long).contiguous(), device)
 
 
 # ---------------------------------------------------------------------------
@@ -366,6 +367,9 @@ class Qwen3TTSPromptEmbedsBuilder:
         # Reusing them avoids launching the embedding/projection pair once per
         # request during prompt construction.
         self._projected_token_cache: dict[tuple[str, tuple[int, ...]], torch.Tensor] = {}
+        # CustomVoice prompt pieces that do not depend on the text, per
+        # (device, speaker, language): bounded by the model's voices.
+        self._custom_voice_prompt_cache: dict[tuple, tuple[torch.Tensor, torch.Tensor, torch.Tensor]] = {}
 
         self._ref_audio_artifact_cache_max_entries = int(ref_audio_artifact_cache_max_entries)
         self._ref_audio_artifact_cache: OrderedDict[str, dict[str, torch.Tensor | bool]] = OrderedDict()
@@ -1024,7 +1028,9 @@ class Qwen3TTSPromptEmbedsBuilder:
         dev = self._device()
         input_ids = coerce_token_ids(info_dict.pop(PRECOMPUTED_TEXT_IDS_KEY, None), device=dev)
         if input_ids is None:
-            input_ids = tok(build_assistant_text(text), return_tensors="pt", padding=False)["input_ids"].to(device=dev)
+            input_ids = to_device_nonblocking(
+                tok(build_assistant_text(text), return_tensors="pt", padding=False)["input_ids"], dev
+            )
 
         # Optional instruct prefix.
         instruct = (info_dict.get("instruct") or [""])[0]
@@ -1411,35 +1417,44 @@ class Qwen3TTSPromptEmbedsBuilder:
                 raise ValueError(f"Unsupported speaker: {speaker}")
             spk_id = spk_id_map[speaker]
             # Keep it at least 1D; embedding on a 0-d tensor can return 1D.
-            spk_tensor = torch.tensor([spk_id], device=input_ids.device, dtype=torch.long)
+            spk_tensor = index_to_device([spk_id], input_ids.device)
             spk_embed = codec_embed(spk_tensor)
             if spk_embed.ndim in (1, 2):
                 spk_embed = spk_embed.view(1, 1, -1)
             speaker_embed = spk_embed
             codec_input = torch.cat([codec_input_0, speaker_embed, codec_input_1], dim=1)
 
-            role_embed = text_projection(text_embedding(input_ids[:, :3]))
-            codec_prefix = torch.cat((tts_pad_embed.expand(-1, codec_input.shape[1] - 2, -1), tts_bos_embed), dim=1)
-            codec_prefix = codec_prefix + codec_input[:, :-1]
-            talker_prompt = torch.cat((role_embed, codec_prefix), dim=1)
+            # The role tokens, speaker/language codec prefix, codec pad row and
+            # closing row do not depend on the text: cache them per speaker and
+            # language so a request only embeds and projects its own text.
+            # A text starting with whitespace could merge with the template's
+            # newline token, so it keeps the per-request role embedding.
+            cache_key = None
+            if non_streaming_mode and instruct_embed is None and not text[:1].isspace():
+                cache_key = (
+                    str(input_ids.device),
+                    int(spk_id),
+                    codec_prefill_list[0][2] if len(codec_prefill_list[0]) == 4 else None,
+                )
+            cv_cache = self._custom_voice_prompt_cache
+            cached = cv_cache.get(cache_key) if cache_key is not None else None
+            if cached is None:
+                role_embed = text_projection(text_embedding(input_ids[:, :3]))
+                codec_prefix = torch.cat((tts_pad_embed.expand(-1, codec_input.shape[1] - 2, -1), tts_bos_embed), dim=1)
+                codec_prefix = codec_prefix + codec_input[:, :-1]
+                prefix = torch.cat((role_embed, codec_prefix), dim=1)
+                pad_row = codec_embed(self._long_tensor([talker_config.codec_pad_id], input_ids.device)).view(1, 1, -1)
+                tail = tts_pad_embed + codec_embed(self._long_tensor([talker_config.codec_bos_id], input_ids.device))
+                if cache_key is not None:
+                    cv_cache[cache_key] = (prefix, pad_row, tail)
+            else:
+                prefix, pad_row, tail = cached
+            talker_prompt = prefix
 
             if non_streaming_mode:
                 text_all = text_projection(text_embedding(input_ids[:, 3:-5]))
                 text_all = torch.cat([text_all, tts_eos_embed], dim=1)
-                pad_ids = torch.full(
-                    (1, int(text_all.shape[1])),
-                    int(talker_config.codec_pad_id),
-                    device=input_ids.device,
-                    dtype=torch.long,
-                )
-                talker_prompt = torch.cat(
-                    [
-                        talker_prompt,
-                        text_all + codec_embed(pad_ids),
-                        tts_pad_embed + codec_embed(self._long_tensor([talker_config.codec_bos_id], input_ids.device)),
-                    ],
-                    dim=1,
-                )
+                talker_prompt = torch.cat([talker_prompt, text_all + pad_row, tail], dim=1)
                 trailing_text_hidden = tts_pad_embed
             else:
                 first_text = text_projection(text_embedding(input_ids[:, 3:4])) + codec_input[:, -1:]

vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_code_predictor_vllm.py

@@ -55,6 +55,9 @@ class Qwen3TTSTalkerCodePredictorForConditionalGenerationVLLM(CodePredictorWrapp
             talker_hidden_size=int(talker_config.hidden_size),
             prefix=prefix,
         )
+        # TTS owns its frame-local cache and overrides _predict_step_logits.
+        # Do not also activate the shared wrapper's Omni incremental loop.
+        self._kv_cache_enabled = False
         # Store talker_config for backward compat (accessed by some callers)
         self.talker_config = talker_config
         self._vllm_config = vllm_config
@@ -62,6 +65,8 @@ class Qwen3TTSTalkerCodePredictorForConditionalGenerationVLLM(CodePredictorWrapp
         self._kv_requested = self._parse_bool_config(extra.get("code_predictor_kv_cache"))
         self._fused_sampling = self._parse_bool_config(extra.get("code_predictor_fused_sampling"))
         self._frame_cache: FrameLocalKVCache | None = None
+        self._fused_requested = self._parse_bool_config(extra.get("code_predictor_fused"))
+        self._fused = None
 
     def _setup_compile(self) -> None:
         if self._compiled_model_fwd is not None:
@@ -97,11 +102,73 @@ class Qwen3TTSTalkerCodePredictorForConditionalGenerationVLLM(CodePredictorWrapp
                         cache(self._proj_buf, batch, step)
         self._synchronize_warmup(weight.device)
         self._frame_cache = cache
+        if self._fused_requested:
+            from .fused_code_predictor import FusedCodePredictor
+
+            self._fused = FusedCodePredictor(self, max_batch)
+            # Compile every kernel variant outside any graph capture.
+            hidden = int(self.talker_config.hidden_size)
+            for batch in range(1, min(2, max_batch) + 1):
+                zeros = torch.zeros(batch, 1, hidden, device=weight.device, dtype=weight.dtype)
+                uniforms = torch.full(
+                    (batch, self._num_groups - 1, int(self.config.vocab_size)), 0.5, device=weight.device
+                )
+                codes = torch.zeros(batch, 1, device=weight.device, dtype=torch.long)
+                self._fused(codes, zeros, zeros, 1.0, 50, uniforms)
+            self._synchronize_warmup(weight.device)
+            logger.info("Qwen3-TTS fused code predictor enabled (capacity %d)", max_batch)
         # Preserve the shared wrapper's initialization sentinel. Its full
         # re-prefill callable is unused while _frame_cache is active.
         self._compiled_model_fwd = self.model.forward
         logger.info("Qwen3-TTS frame-local KV warmed for batches %s", self._bucket_sizes)
 
+    @torch.inference_mode()
+    def forward(
+        self,
+        layer0_code: torch.Tensor,
+        layer0_embed: torch.Tensor,
+        last_talker_hidden: torch.Tensor,
+        do_sample: bool = True,
+        temperature: float = 0.9,
+        top_k: int = 50,
+        top_p: float = 1.0,
+        generator: torch.Generator | None = None,
+        generators: Sequence[torch.Generator | None] | None = None,
+        sample_uniforms: torch.Tensor | None = None,
+    ):
+        self._validate_sampling_inputs(int(layer0_code.shape[0]), generators, sample_uniforms)
+        if self._fused_requested:
+            self._setup_compile()
+        if (
+            self._fused is not None
+            and do_sample
+            and temperature > 0
+            and top_p == 1.0
+            and sample_uniforms is not None
+            and sample_uniforms.dtype == torch.float32
+            and sample_uniforms.is_contiguous()
+        ):
+            return self._fused(
+                layer0_code,
+                layer0_embed,
+                last_talker_hidden,
+                1.0 / temperature,
+                int(top_k),
+                sample_uniforms,
+            )
+        return super().forward(
+            layer0_code,
+            layer0_embed,
+            last_talker_hidden,
+            do_sample=do_sample,
+            temperature=temperature,
+            top_k=top_k,
+            top_p=top_p,
+            generator=generator,
+            generators=generators,
+            sample_uniforms=sample_uniforms,
+        )
+
     def _predict_step_logits(
         self,
         proj_buf: torch.Tensor,

vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_talker.py

@@ -27,6 +27,7 @@ from vllm.sequence import IntermediateTensors
 
 from vllm_omni.data_entry_keys import OmniPayload
 from vllm_omni.model_executor.models.output_templates import OmniOutput, OwnedBatchTensor
+from vllm_omni.utils.device_copy import index_to_device, to_device_nonblocking
 from vllm_omni.utils.speaker_cache import (
     get_speaker_cache,
     iter_custom_voice_profiles,
@@ -39,6 +40,7 @@ from .configuration_qwen3_tts import Qwen3TTSConfig, Qwen3TTSSpeakerEncoderConfi
 from .first_audio import talker_first_audio_enabled
 from .prompt_embeds_builder import PRECOMPUTED_TEXT_IDS_KEY, Qwen3TTSPromptEmbedsBuilder, resolve_x_vector_only
 from .qwen3_tts_code_predictor_vllm import Qwen3TTSTalkerCodePredictorForConditionalGenerationVLLM
+from .stream_decode import stream_ref_context_frames, talker_stream_decode_enabled
 from .tokenizer_12hz.configuration_qwen3_tts_tokenizer_v2 import Qwen3TTSTokenizerV2Config
 from .tokenizer_12hz.modeling_qwen3_tts_tokenizer_v2 import Qwen3TTSTokenizerV2Encoder
 
@@ -411,7 +413,21 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module):
         # A frame's residual codebooks depend only on the Talker hidden and
         # CB0 of the step that sampled it, so MRV2 may complete the frame at
         # the end of that step instead of in the next step's preprocess.
-        self.mtp_eager_frames = talker_first_audio_enabled(vllm_config)
+        self.stream_decode = talker_stream_decode_enabled(vllm_config)
+        predictor = Qwen3TTSTalkerCodePredictorForConditionalGenerationVLLM
+        extra = predictor._stage_connector_extra_config(vllm_config)
+        self.stream_first_audio = self.stream_decode and predictor._parse_bool_config(
+            extra.get("talker_stream_first_audio")
+        )
+        # Model-local voice-clone priming boundary; each deployment selects
+        # its reference context independently of the two-stage codec profile.
+        self.stream_ref_context_frames = stream_ref_context_frames(vllm_config) if self.stream_decode else 0
+        self.mtp_eager_frames = talker_first_audio_enabled(vllm_config) or self.stream_decode
+        self.stream_decoder = None
+        self.stream_graphs = None
+        self.stream_prime_graphs = None
+        self.stream_sample_rate = 0
+        self.stream_chunk_frames = 25
         # The runners bypass only the outer whole-MTP graph when explicit
         # generators are present, so seeded requests can still share one raw
         # batched MTP call with independent per-row streams.
@@ -718,6 +734,11 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module):
 
     # -------------------- Omni multimodal output plumbing --------------------
 
+    @property
+    def mm_outputs_fresh_per_step(self) -> bool:
+        """Stream-decode outputs come from :meth:`_make_eager_omni_output`, allocated per step."""
+        return self.stream_decoder is not None and self.eager_frames_active
+
     def make_omni_output(self, model_outputs: torch.Tensor | OmniOutput, **kwargs: Any) -> OmniOutput:
         if isinstance(model_outputs, OmniOutput):
             return model_outputs
@@ -895,6 +916,13 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module):
                 "first_audio": torch.zeros((num_tokens,), dtype=torch.int8, device=hidden.device),
             },
         }
+        if self.stream_decoder is not None:
+            # Filled per sampled row by the runner's eager MTP; token-major so
+            # each request's slice is its frames of this step.
+            spf = int(self.stream_decoder.spf)
+            mm["model_outputs"] = torch.zeros((num_tokens, spf), dtype=torch.float32, device=hidden.device)
+            sr = torch.tensor(self.stream_sample_rate, dtype=torch.int32)
+            mm["sr"] = [sr] * len(info_dicts)
         ref_rows = [
             index
             for index, info in enumerate(info_dicts)
@@ -1442,8 +1470,72 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module):
             self.first_frame_decoder = Qwen3TTSFirstFrameDecoder(self.model_path)
             decoder_loaded = self.first_frame_decoder.load(self.vllm_config)
             loaded = set(loaded) | {f"first_frame_decoder.{name}" for name in decoder_loaded}
+        elif self.stream_decode:
+            # The in-stage decoder owns every frame. Optional first-audio
+            # delivery uses its PCM; later chunks use regular step outputs.
+            from .first_frame_decoder import Qwen3TTSFirstFrameDecoder
+            from .tokenizer_12hz.streaming_decoder import StreamingCodecDecoder
+
+            holder = Qwen3TTSFirstFrameDecoder(self.model_path)
+            decoder_loaded = holder.load(self.vllm_config)
+            loaded = set(loaded) | {f"stream_decoder.{name}" for name in decoder_loaded}
+            slots = int(self.vllm_config.scheduler_config.max_num_seqs)
+            self.stream_decoder = StreamingCodecDecoder(holder.decoder, num_slots=slots, dtype=torch.bfloat16)
+            self.stream_sample_rate = int(holder.sample_rate)
+            logger.info("Qwen3-TTS Talker stream decode enabled (%d slots)", slots)
         return loaded
 
+    def capture_stream_decode_graphs(self, batch_sizes: list[int]) -> None:
+        if self.stream_decoder is None or self.stream_graphs is not None:
+            return
+        from .tokenizer_12hz.streaming_decoder import StreamingDecodeGraphs
+
+        with torch.inference_mode():
+            self.stream_graphs = StreamingDecodeGraphs(self.stream_decoder, batch_sizes)
+            frames = min(self.stream_ref_context_frames, self.stream_chunk_frames, self.stream_decoder.max_frames)
+            if self.config.tts_model_type == "base" and frames > 1:
+                # Bound the multi-frame graph workspace by the configured
+                # token budget; larger priming groups keep the eager path.
+                budget = max(1, self.vllm_config.scheduler_config.max_num_batched_tokens // frames)
+                prime_sizes = sorted({1, *(size for size in batch_sizes if size <= budget)})
+                self.stream_prime_graphs = StreamingDecodeGraphs(self.stream_decoder, prime_sizes, frames=frames)
+        logger.info("Captured Talker stream decode graphs for batch sizes %s", self.stream_graphs.sizes)
+        if self.stream_prime_graphs is not None:
+            logger.info(
+                "Captured Talker reference priming graphs for %d frames and batch sizes %s",
+                self.stream_prime_graphs.frames,
+                self.stream_prime_graphs.sizes,
+            )
+
+    def get_stream_ref_context(self, info: OmniPayload) -> torch.Tensor | None:
+        """Model-owned reference-code layout and context boundary for a PCM slot."""
+        codes = info.get("codes")
+        ref = codes.get("ref") if isinstance(codes, dict) else None
+        if not isinstance(ref, torch.Tensor) or not ref.numel():
+            return None
+        return ref.reshape(-1, int(self.talker_config.num_code_groups))[-self.stream_ref_context_frames :]
+
+    def prime_stream_decoder(self, primes: list[tuple[int, torch.Tensor]]) -> None:
+        """Prime equal-length reference groups without padding or aliasing live slots."""
+        if not primes:
+            return
+        stream = self.stream_decoder
+        if stream is None:
+            raise RuntimeError("Reference priming requires the single-stage PCM decoder")
+        groups: dict[int, list[tuple[int, torch.Tensor]]] = {}
+        for req_idx, ref in primes:
+            groups.setdefault(int(ref.shape[0]), []).append((req_idx, ref))
+        for items in groups.values():
+            codes = to_device_nonblocking(torch.stack([ref for _idx, ref in items]), stream.device).to(torch.int32)
+            n = len(items)
+            slots = index_to_device([idx for idx, _ref in items], stream.device, dtype=torch.int32)
+            for t0 in range(0, int(codes.shape[1]), self.stream_chunk_frames):
+                pos = index_to_device([t0] * n, stream.device, dtype=torch.int32)
+                chunk = codes[:, t0 : t0 + self.stream_chunk_frames].contiguous()
+                graphs = self.stream_prime_graphs
+                decode = graphs if graphs is not None and chunk.shape[1] == graphs.frames else stream
+                decode(chunk, slots, pos)
+
     def _build_stacked_codec_embed(self) -> None:
         embeds = self.code_predictor.get_input_embeddings()
         if not embeds:

vllm_omni/model_executor/models/qwen3_tts/stream_decode.py

@@ -0,0 +1,60 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Eligibility and model options for the single-stage CUDA PCM pipeline."""
+
+import torch
+from vllm.config import VllmConfig
+
+from vllm_omni.platforms import current_omni_platform
+
+
+def talker_stream_decode_enabled(vllm_config: VllmConfig) -> bool:
+    """The Talker decodes every frame itself and emits PCM as its final output."""
+    from .qwen3_tts_code_predictor_vllm import Qwen3TTSTalkerCodePredictorForConditionalGenerationVLLM as Predictor
+
+    extra = Predictor._stage_connector_extra_config(vllm_config)
+    model = vllm_config.model_config
+    parallel = vllm_config.parallel_config
+    requested = Predictor._parse_bool_config(extra.get("talker_stream_decode"))
+    if not requested:
+        if getattr(model, "engine_output_type", None) == "audio":
+            raise ValueError("The single-stage Qwen3-TTS pipeline requires talker_stream_decode=True")
+        if Predictor._parse_bool_config(extra.get("talker_stream_first_audio")):
+            raise ValueError("talker_stream_first_audio requires the single-stage talker_stream_decode path")
+        return False
+    if getattr(model, "engine_output_type", None) != "audio" or not getattr(model, "final_output", False):
+        raise ValueError(
+            "talker_stream_decode requires a final audio-output Talker; "
+            "select qwen3_tts_fused_single_gpu.yaml instead of enabling it on the two-stage pipeline"
+        )
+    if getattr(model, "supports_running_prefix_cache_reset", True):
+        raise ValueError(
+            "The single-stage PCM pipeline must disable running prefix-cache reset; "
+            "select qwen3_tts_fused_single_gpu.yaml"
+        )
+    if not (
+        current_omni_platform.is_cuda()
+        and torch.device(vllm_config.device_config.device).type == "cuda"
+        and bool(getattr(model, "use_v2_model_runner", False))
+        and bool(getattr(model, "async_chunk", False))
+        and parallel.tensor_parallel_size == 1
+        and parallel.pipeline_parallel_size == 1
+        and parallel.distributed_executor_backend in (None, "uni")
+        and not vllm_config.cache_config.enable_prefix_caching
+    ):
+        raise ValueError(
+            "talker_stream_decode requires CUDA, MRv2, async_chunk=True, TP=PP=1, "
+            "an in-process worker and enable_prefix_caching=False; "
+            "use qwen3_tts_fused_single_gpu.yaml"
+        )
+    if Predictor._parse_bool_config(extra.get("talker_first_audio")):
+        raise ValueError("talker_stream_decode and talker_first_audio are mutually exclusive")
+    return True
+
+
+def stream_ref_context_frames(vllm_config: VllmConfig) -> int:
+    """Reference-code context frames for voice-clone priming in the Talker."""
+    from .qwen3_tts_code_predictor_vllm import Qwen3TTSTalkerCodePredictorForConditionalGenerationVLLM as Predictor
+
+    extra = Predictor._stage_connector_extra_config(vllm_config)
+    return int(extra.get("ref_code_context_frames") or extra.get("codec_left_context_frames", 25))

vllm_omni/model_executor/models/qwen3_tts/tokenizer_12hz/streaming_decoder.py

@@ -0,0 +1,799 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+# ruff: noqa: N803 - Triton constexpr parameters use kernel-style capitals.
+"""Stateful, fused streaming decoder for the Qwen3-TTS 12 Hz codec.
+
+Every layer of the codec decoder is causal, so what a new frame needs from the
+past is small and fixed: each causal conv's last ``(k - 1) * dilation`` input
+rows, each overlapping transposed conv's last GEMM row, and each attention
+layer's last ``sliding_window`` keys and values. A request owns a slot of
+per-layer state holding exactly that, and a call decodes ``T`` new frames for
+each of ``B`` requests, reading and advancing the state. The output equals the
+whole-utterance decode (``Qwen3TTSTokenizerV2Decoder._forward_exact``) frame by
+frame, up to floating point rounding: nothing is recomputed and nothing is
+truncated.
+
+Layout and fusion:
+
+- Activations are time-major ``[rows, C]``; every dense conv is a tap-major
+  im2col (one Triton kernel that also applies the pending bias and SnakeBeta,
+  and reads/writes the history) followed by one cuBLAS GEMM.
+- Linear chains fold at load: the RVQ output projections, ``pre_conv`` and the
+  transformer's ``input_proj`` become one ``k = 3`` conv over the concatenated
+  codebook embeddings; LayerScale folds into ``o_proj``/``down_proj``; ConvNeXt
+  ``gamma`` folds into ``pwconv2``.
+- A bias that feeds a residual stream is carried to that stream's consumer
+  (the next im2col or transposed-conv GEMM) instead of being added in place.
+- Residual branches accumulate into their stream with in-place ``addmm_``.
+
+History buffers are double-buffered per slot: a call reads half ``parity[slot]``
+and writes the other half, then flips the parity, so no kernel overwrites rows
+another program of the same launch still reads. A row whose first frame index
+``pos`` is 0 reads zeros, which are the convs' causal padding, and attention
+only looks back ``min(pos + 1, window)`` positions, so a new request needs no
+state reset.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import torch
+from torch import nn
+from vllm.triton_utils import tl, triton
+
+from vllm_omni.platforms import current_omni_platform
+
+_KERNEL = 7
+
+
+# --------------------------------------------------------------------------- kernels
+
+
+@triton.jit
+def _rvq_gather_kernel(codes_ptr, books_ptr, out_ptr, NQ: tl.constexpr, CB: tl.constexpr, HALF: tl.constexpr):
+    row = tl.program_id(0)
+    offs = tl.arange(0, HALF)
+    c0 = tl.load(codes_ptr + row * NQ).to(tl.int64)
+    c0 = tl.where((c0 >= 0) & (c0 < CB), c0, 0)  # EOS/special rows decode as code 0; callers drop them
+    first = tl.load(books_ptr + c0 * HALF + offs).to(tl.float32)
+    acc = tl.zeros([HALF], dtype=tl.float32)
+    for q in range(1, NQ):
+        c = tl.load(codes_ptr + row * NQ + q).to(tl.int64)
+        c = tl.where((c >= 0) & (c < CB), c, 0)
+        acc += tl.load(books_ptr + (q * CB + c) * HALF + offs).to(tl.float32)
+    dt = out_ptr.dtype.element_ty
+    tl.store(out_ptr + row * 2 * HALF + offs, first.to(dt))
+    tl.store(out_ptr + row * 2 * HALF + HALF + offs, acc.to(dt))
+
+
+@triton.jit
+def _fast_sin(y):
+    # Two-step reduction to [-pi, pi], then the hardware approximation
+    # (absolute error ~4e-7 there): accurate sin's instruction count, not
+    # memory, bounds every kernel that applies SnakeBeta.
+    k = tl.extra.cuda.libdevice.rint(y * 0.15915494309189535)
+    r = tl.fma(-k, -1.7484555314695172e-07, tl.fma(-k, 6.2831854820251465, y))
+    return tl.inline_asm_elementwise("sin.approx.f32 $0, $1;", "=r,r", [r], dtype=tl.float32, is_pure=True, pack=1)
+
+
+@triton.jit
+def _snake(v, a, ib):
+    s = _fast_sin(v * a)
+    return v + ib * s * s
+
+
+@triton.jit
+def _im2col_kernel(
+    x_ptr,  # [B * T, C] layer input rows
+    bias_ptr,  # [C] pending bias (added to x rows only) or dummy
+    alpha_ptr,  # [C] exp(alpha) or dummy
+    ibeta_ptr,  # [C] 1 / (exp(beta) + eps) or dummy
+    hist_ptr,  # [2, S, H, C] transformed history rows
+    slots_ptr,
+    pos_ptr,
+    par_ptr,
+    col_ptr,  # [B * T, K * C]
+    T,
+    S,
+    C: tl.constexpr,
+    K: tl.constexpr,
+    D: tl.constexpr,
+    H: tl.constexpr,
+    HAS_BIAS: tl.constexpr,
+    HAS_SNAKE: tl.constexpr,
+    BU: tl.constexpr,
+    BC: tl.constexpr,
+    EXT: tl.constexpr = False,
+):
+    """Source row ``u`` of ``[history (H) | x (T)]``: scatter to every tap that reads it, save the last H."""
+    b = tl.program_id(0)
+    u = tl.program_id(1) * BU + tl.arange(0, BU)
+    c = tl.program_id(2) * BC + tl.arange(0, BC)
+    cm = c < C
+    n_src = H + T
+    um = u < n_src
+    slot = tl.load(slots_ptr + b).to(tl.int64)
+    live = tl.load(pos_ptr + b) != 0
+    par = tl.load(par_ptr + slot).to(tl.int64)
+    # history rows
+    from_hist = u < H
+    hrow = tl.where(from_hist, u, 0)
+    hv = tl.load(
+        hist_ptr + ((par * S + slot) * H + hrow)[:, None] * C + c[None, :],
+        mask=(from_hist & um & live)[:, None] & cm[None, :],
+        other=0.0,
+    ).to(tl.float32)
+    xrow = tl.where(from_hist, 0, u - H)
+    xv = tl.load(
+        x_ptr + (b * T + xrow).to(tl.int64)[:, None] * C + c[None, :],
+        mask=((~from_hist) & um)[:, None] & cm[None, :],
+        other=0.0,
+    ).to(tl.float32)
+    if HAS_BIAS:
+        xv += tl.load(bias_ptr + c, mask=cm, other=0.0).to(tl.float32)[None, :]
+    if HAS_SNAKE:
+        a = tl.load(alpha_ptr + c, mask=cm, other=0.0).to(tl.float32)[None, :]
+        ib = tl.load(ibeta_ptr + c, mask=cm, other=0.0).to(tl.float32)[None, :]
+        xv = _snake(xv, a, ib)
+    dt = col_ptr.dtype.element_ty
+    v = tl.where(from_hist[:, None], hv, xv.to(dt).to(tl.float32)).to(dt)
+    if EXT:
+        tl.store(col_ptr + (b * n_src + u).to(tl.int64)[:, None] * C + c[None, :], v, mask=um[:, None] & cm[None, :])
+    else:
+        for j in tl.static_range(K):
+            t = u - j * D
+            tm = (t >= 0) & (t < T) & um
+            tl.store(
+                col_ptr + (b * T + tl.where(tm, t, 0)).to(tl.int64)[:, None] * (K * C) + j * C + c[None, :],
+                v,
+                mask=tm[:, None] & cm[None, :],
+            )
+    if H > 0:
+        hm = (u >= T) & um
+        tl.store(
+            hist_ptr + (((1 - par) * S + slot) * H + tl.where(hm, u - T, 0))[:, None] * C + c[None, :],
+            v,
+            mask=hm[:, None] & cm[None, :],
+        )
+
+
+@triton.jit
+def _rms_norm_kernel(x_ptr, w_ptr, out_ptr, eps, N: tl.constexpr):
+    row = tl.program_id(0).to(tl.int64)
+    offs = tl.arange(0, N)
+    x = tl.load(x_ptr + row * N + offs).to(tl.float32)
+    var = tl.sum(x * x, axis=0) / N
+    dt = out_ptr.dtype.element_ty
+    y = (x * tl.rsqrt(var + eps)).to(dt).to(tl.float32)
+    w = tl.load(w_ptr + offs).to(tl.float32)
+    tl.store(out_ptr + row * N + offs, (w * y).to(dt))
+
+
+@triton.jit
+def _rope_kv_write_kernel(
+    qkv_ptr,  # [B * T, 3 * NH * HD]
+    q_ptr,  # [B * T, NH * HD]
+    k_ring_ptr,  # [S, RING, NH * HD]
+    v_ring_ptr,
+    slots_ptr,
+    pos_ptr,
+    T,
+    theta,
+    NH: tl.constexpr,
+    HD: tl.constexpr,
+    RING: tl.constexpr,
+):
+    r = tl.program_id(0)
+    h = tl.program_id(1)
+    b = r // T
+    p = tl.load(pos_ptr + b) + r % T
+    slot = tl.load(slots_ptr + b).to(tl.int64)
+    half: tl.constexpr = HD // 2
+    i = tl.arange(0, half)
+    inv_freq = tl.exp(-(2.0 * i.to(tl.float32) / HD) * tl.log(theta))
+    ang = p.to(tl.float32) * inv_freq
+    dt = q_ptr.dtype.element_ty
+    cos = tl.cos(ang).to(dt).to(tl.float32)
+    sin = tl.sin(ang).to(dt).to(tl.float32)
+    W: tl.constexpr = NH * HD
+    base = qkv_ptr + r.to(tl.int64) * 3 * W + h * HD
+    q1 = tl.load(base + i).to(tl.float32)
+    q2 = tl.load(base + half + i).to(tl.float32)
+    k1 = tl.load(base + W + i).to(tl.float32)
+    k2 = tl.load(base + W + half + i).to(tl.float32)
+    vv = tl.arange(0, HD)
+    v = tl.load(base + 2 * W + vv)
+    qo = q_ptr + r.to(tl.int64) * W + h * HD
+    tl.store(qo + i, (q1 * cos - q2 * sin).to(dt))
+    tl.store(qo + half + i, (q2 * cos + q1 * sin).to(dt))
+    ring = (slot * RING + p % RING) * W + h * HD
+    tl.store(k_ring_ptr + ring + i, (k1 * cos - k2 * sin).to(dt))
+    tl.store(k_ring_ptr + ring + half + i, (k2 * cos + k1 * sin).to(dt))
+    tl.store(v_ring_ptr + ring + vv, v)
+
+
+@triton.jit
+def _attend_kernel(
+    q_ptr,  # [B * T, NH * HD]
+    k_ring_ptr,
+    v_ring_ptr,
+    out_ptr,  # [B * T, NH * HD]
+    slots_ptr,
+    pos_ptr,
+    T,
+    scale,
+    NH: tl.constexpr,
+    HD: tl.constexpr,
+    RING: tl.constexpr,
+    WINDOW: tl.constexpr,
+    BK: tl.constexpr,
+):
+    r = tl.program_id(0)
+    h = tl.program_id(1)
+    b = r // T
+    p = tl.load(pos_ptr + b) + r % T
+    slot = tl.load(slots_ptr + b).to(tl.int64)
+    W: tl.constexpr = NH * HD
+    d = tl.arange(0, HD)
+    q = tl.load(q_ptr + r.to(tl.int64) * W + h * HD + d).to(tl.float32)
+    j = tl.arange(0, BK)
+    n = tl.minimum(p + 1, WINDOW)
+    km = j < n
+    kp = p - j  # key positions p, p-1, ...
+    idx = (slot * RING + tl.where(km, kp, 0) % RING) * W + h * HD
+    k = tl.load(k_ring_ptr + idx[:, None] + d[None, :], mask=km[:, None], other=0.0).to(tl.float32)
+    s = tl.sum(k * q[None, :], axis=1) * scale
+    s = tl.where(km, s, -float("inf"))
+    m = tl.max(s, axis=0)
+    e = tl.exp(s - m)
+    e = tl.where(km, e, 0.0)
+    w = e / tl.sum(e, axis=0)
+    dt = out_ptr.dtype.element_ty
+    w = w.to(dt).to(tl.float32)
+    v = tl.load(v_ring_ptr + idx[:, None] + d[None, :], mask=km[:, None], other=0.0).to(tl.float32)
+    o = tl.sum(v * w[:, None], axis=0)
+    tl.store(out_ptr + r.to(tl.int64) * W + h * HD + d, o.to(dt))
+
+
+@triton.jit
+def _silu_mul_kernel(gu_ptr, out_ptr, N: tl.constexpr, BN: tl.constexpr):
+    row = tl.program_id(0).to(tl.int64)
+    offs = tl.program_id(1) * BN + tl.arange(0, BN)
+    m = offs < N
+    g = tl.load(gu_ptr + row * 2 * N + offs, mask=m).to(tl.float32)
+    u = tl.load(gu_ptr + row * 2 * N + N + offs, mask=m).to(tl.float32)
+    dt = out_ptr.dtype.element_ty
+    sg = (g / (1.0 + tl.exp(-g))).to(dt).to(tl.float32)
+    tl.store(out_ptr + row * N + offs, (sg * u).to(dt), mask=m)
+
+
+@triton.jit
+def _dwconv_ln_kernel(
+    z_ptr,  # [B * T, C] ConvNeXt input rows (true values)
+    w_ptr,  # [K, C]
+    b_ptr,  # [C]
+    lnw_ptr,
+    lnb_ptr,
+    hist_ptr,  # [2, S, K - 1, C]
+    slots_ptr,
+    pos_ptr,
+    par_ptr,
+    out_ptr,  # [B * T, C] LayerNorm output
+    T,
+    S,
+    eps,
+    C: tl.constexpr,
+    K: tl.constexpr,
+):
+    b = tl.program_id(0)
+    t = tl.program_id(1)  # 0 .. max(T, K - 1)
+    H: tl.constexpr = K - 1
+    c = tl.arange(0, C)
+    slot = tl.load(slots_ptr + b).to(tl.int64)
+    live = tl.load(pos_ptr + b) != 0
+    par = tl.load(par_ptr + slot).to(tl.int64)
+    dt = out_ptr.dtype.element_ty
+    hbase = hist_ptr + (par * S + slot) * H * C
+    acc = tl.load(b_ptr + c).to(tl.float32)
+    for j in tl.static_range(K):
+        u = t + j  # index into [hist (H) | z (T)]
+        fh = u < H
+        hv = tl.load(hbase + tl.where(fh, u, 0) * C + c, mask=(c < C) & fh & live, other=0.0).to(tl.float32)
+        zrow = (b * T + tl.where(fh, 0, u - H)).to(tl.int64)
+        zv = tl.load(z_ptr + zrow * C + c, mask=(c < C) & (~fh) & (t < T), other=0.0).to(tl.float32)
+        acc += tl.where(fh, hv, zv) * tl.load(w_ptr + j * C + c).to(tl.float32)
+    h = acc.to(dt).to(tl.float32)
+    mean = tl.sum(h, axis=0) / C
+    hc = h - mean
+    var = tl.sum(hc * hc, axis=0) / C
+    y = hc * tl.rsqrt(var + eps) * tl.load(lnw_ptr + c).to(tl.float32) + tl.load(lnb_ptr + c).to(tl.float32)
+    tl.store(out_ptr + (b * T + t).to(tl.int64) * C + c, y.to(dt), mask=(c < C) & (t < T))
+    # new history row t: source row T + t of [hist | z]
+    u = T + t
+    fh = u < H
+    hv = tl.load(hbase + tl.where(fh, u, 0) * C + c, mask=(c < C) & fh & live & (t < H), other=0.0)
+    zrow = (b * T + tl.where(fh, 0, u - H)).to(tl.int64)
+    zv = tl.load(z_ptr + zrow * C + c, mask=(c < C) & (~fh) & (t < H), other=0.0)
+    tl.store(
+        hist_ptr + (((1 - par) * S + slot) * H + tl.where(t < H, t, 0)) * C + c,
+        tl.where(fh, hv, zv).to(dt),
+        mask=(c < C) & (t < H),
+    )
+
+
+@triton.jit
+def _bias_gelu_kernel(x_ptr, b_ptr, N: tl.constexpr, BN: tl.constexpr):
+    row = tl.program_id(0).to(tl.int64)
+    offs = tl.program_id(1) * BN + tl.arange(0, BN)
+    m = offs < N
+    x = tl.load(x_ptr + row * N + offs, mask=m).to(tl.float32) + tl.load(b_ptr + offs, mask=m).to(tl.float32)
+    dt = x_ptr.dtype.element_ty
+    x = x.to(dt).to(tl.float32)
+    y = 0.5 * x * (1.0 + tl.erf(x * 0.7071067811865476))
+    tl.store(x_ptr + row * N + offs, y.to(dt), mask=m)
+
+
+@triton.jit
+def _col2im_kernel(
+    z_ptr,  # [B * T, 2 * R, C] transposed-conv GEMM rows
+    bias_ptr,  # [C]
+    prev_ptr,  # [2, S, R, C] previous call's last GEMM row, second half
+    slots_ptr,
+    pos_ptr,
+    par_ptr,
+    out_ptr,  # [B * T * R, C]
+    T,
+    S,
+    R: tl.constexpr,
+    C: tl.constexpr,
+    BC: tl.constexpr,
+):
+    row = tl.program_id(0)  # b * T + t
+    j = tl.program_id(1)  # 0 .. R - 1
+    c = tl.program_id(2) * BC + tl.arange(0, BC)
+    cm = c < C
+    b = row // T
+    t = row % T
+    slot = tl.load(slots_ptr + b).to(tl.int64)
+    live = tl.load(pos_ptr + b) != 0
+    par = tl.load(par_ptr + slot).to(tl.int64)
+    zr = z_ptr + row.to(tl.int64) * 2 * R * C
+    cur = tl.load(zr + j * C + c, mask=cm).to(tl.float32)
+    prev_in = tl.load(zr - 2 * R * C + (R + j) * C + c, mask=cm & (t > 0), other=0.0).to(tl.float32)
+    prev_st = tl.load(prev_ptr + ((par * S + slot) * R + j) * C + c, mask=cm & (t == 0) & live, other=0.0).to(
+        tl.float32
+    )
+    y = cur + tl.where(t > 0, prev_in, prev_st) + tl.load(bias_ptr + c, mask=cm).to(tl.float32)
+    dt = out_ptr.dtype.element_ty
+    tl.store(out_ptr + (row.to(tl.int64) * R + j) * C + c, y.to(dt), mask=cm)
+    last = tl.load(zr + (R + j) * C + c, mask=cm & (t == T - 1))
+    tl.store(prev_ptr + (((1 - par) * S + slot) * R + j) * C + c, last, mask=cm & (t == T - 1))
+
+
+@triton.jit
+def _conv_out_kernel(
+    ext_ptr,  # [B, H + T, C] transformed rows: history then this call's
+    w_ptr,  # [K, C]
+    wb,  # conv bias (float)
+    out_ptr,  # [B, T] float32 waveform
+    T,
+    C: tl.constexpr,
+    K: tl.constexpr,
+    BT: tl.constexpr,
+    CP: tl.constexpr,
+):
+    b = tl.program_id(0)
+    t = tl.program_id(1) * BT + tl.arange(0, BT)
+    c = tl.arange(0, CP)
+    cm = c < C
+    tm = t < T
+    base = ext_ptr + b.to(tl.int64) * (T + K - 1) * C
+    acc = tl.zeros([BT], dtype=tl.float32)
+    for j in tl.static_range(K):
+        v = tl.load(base + (t + j)[:, None] * C + c[None, :], mask=tm[:, None] & cm[None, :], other=0.0).to(tl.float32)
+        acc += tl.sum(v * tl.load(w_ptr + j * C + c, mask=cm, other=0.0).to(tl.float32)[None, :], axis=1)
+    dt = ext_ptr.dtype.element_ty
+    y = acc.to(dt).to(tl.float32) + wb
+    y = tl.minimum(tl.maximum(y.to(dt).to(tl.float32), -1.0), 1.0)
+    tl.store(out_ptr + b.to(tl.int64) * T + t, y, mask=tm)
+
+
+@triton.jit
+def _flip_kernel(par_ptr, slots_ptr):
+    slot = tl.load(slots_ptr + tl.program_id(0)).to(tl.int64)
+    tl.store(par_ptr + slot, 1 - tl.load(par_ptr + slot))
+
+
+# --------------------------------------------------------------------------- module
+
+
+@dataclass
+class _ConvSpec:
+    weight: torch.Tensor  # [K * C_in, C_out] tap-major
+    c_in: int
+    k: int
+    d: int
+    hist: torch.Tensor | None = None  # [2, S, H, C_in]
+
+
+def _conv_taps(w: torch.Tensor) -> torch.Tensor:
+    """``[C_out, C_in, K]`` -> ``[K * C_in, C_out]``."""
+    c_out, c_in, k = w.shape
+    return w.permute(2, 1, 0).reshape(k * c_in, c_out)
+
+
+def _snake_params(module) -> tuple[torch.Tensor, torch.Tensor]:
+    # SnakeBeta: x + 1 / (exp(beta) + 1e-9) * sin(exp(alpha) * x) ** 2
+    alpha = module.alpha.detach().float()
+    beta = module.beta.detach().float()
+    return torch.exp(alpha).contiguous(), (1.0 / (torch.exp(beta) + 1e-9)).contiguous()
+
+
+class StreamingCodecDecoder(nn.Module):
+    """Frame-exact stateful decoder built from a loaded ``Qwen3TTSTokenizerV2Decoder``."""
+
+    RING = 128  # KV ring positions; a call may carry up to RING - window + 1 frames
+
+    def __init__(self, decoder: nn.Module, num_slots: int, dtype: torch.dtype | None = None) -> None:
+        super().__init__()
+        cfg = decoder.config
+        dev = next(decoder.parameters()).device
+        dtype = dtype or next(decoder.parameters()).dtype
+        self.dtype, self.device = dtype, dev
+        self.num_slots = num_slots + 1  # last slot: scratch for padding rows
+        self.scratch_slot = num_slots
+        S = self.num_slots
+        self.nq = int(cfg.num_quantizers)
+        self.cb = int(cfg.codebook_size)
+        self.half = int(cfg.codebook_dim) // 2
+        self.hidden = int(cfg.hidden_size)
+        self.heads = int(cfg.num_attention_heads)
+        self.head_dim = int(cfg.head_dim)
+        self.window = int(cfg.sliding_window)
+        self.theta = float(cfg.rope_theta)
+        self.eps = float(cfg.rms_norm_eps)
+        self.max_frames = self.RING - self.window + 1
+        self.spf = int(decoder.total_upsample)
+        if self.heads != int(cfg.num_key_value_heads):
+            raise ValueError("streaming codec decoder assumes no grouped KV")
+
+        f32 = lambda t: t.detach().to(device=dev, dtype=torch.float32)  # noqa: E731
+        cast = lambda t: t.to(dtype).contiguous()  # noqa: E731
+
+        def hist(h: int, c: int) -> torch.Tensor:
+            return torch.zeros(2, S, h, c, device=dev, dtype=dtype)
+
+        # RVQ books, pre-divided by usage: [NQ * CB, HALF]
+        q = decoder.quantizer
+        books = []
+        for vq in [q.rvq_first.vq.layers[0]] + list(q.rvq_rest.vq.layers):
+            cbk = vq._codebook
+            books.append(f32(cbk.embedding_sum) / f32(cbk.cluster_usage).clamp(min=cbk.epsilon)[:, None])
+        self.books = cast(torch.cat(books))
+        w_rvq = torch.cat(
+            [f32(q.rvq_first.output_proj.weight)[:, :, 0], f32(q.rvq_rest.output_proj.weight)[:, :, 0]], 1
+        )
+        # pre_conv (k=3) and input_proj folded onto the embedding concat
+        pre_w, pre_b = f32(decoder.pre_conv.conv.weight), f32(decoder.pre_conv.conv.bias)
+        pt = decoder.pre_transformer
+        in_w, in_b = f32(pt.input_proj.weight), f32(pt.input_proj.bias)
+        taps = torch.stack([in_w @ pre_w[:, :, j] @ w_rvq for j in range(pre_w.shape[2])], dim=2)  # [hid, 2H, 3]
+        self.pre = _ConvSpec(cast(_conv_taps(taps)), 2 * self.half, 3, 1, hist(2, 2 * self.half))
+        self.pre_b = cast(in_w @ pre_b + in_b)
+
+        # transformer
+        self.layers = []
+        qd = self.heads * self.head_dim
+        for layer in pt.layers:
+            a = layer.self_attn
+            qkv = torch.cat([f32(a.q_proj.weight), f32(a.k_proj.weight), f32(a.v_proj.weight)], 0)
+            o = f32(a.o_proj.weight) * f32(layer.self_attn_layer_scale.scale)[:, None]
+            gu = torch.cat([f32(layer.mlp.gate_proj.weight), f32(layer.mlp.up_proj.weight)], 0)
+            down = f32(layer.mlp.down_proj.weight) * f32(layer.mlp_layer_scale.scale)[:, None]
+            self.layers.append(
+                dict(
+                    ln1=cast(f32(layer.input_layernorm.weight)),
+                    ln2=cast(f32(layer.post_attention_layernorm.weight)),
+                    qkv=cast(qkv.t()),
+                    o=cast(o.t()),
+                    gu=cast(gu.t()),
+                    down=cast(down.t()),
+                    k_ring=torch.zeros(S, self.RING, qd, device=dev, dtype=dtype),
+                    v_ring=torch.zeros(S, self.RING, qd, device=dev, dtype=dtype),
+                )
+            )
+        self.norm_w = cast(f32(pt.norm.weight))
+        self.out_w = cast(f32(pt.output_proj.weight).t())
+        self.out_b = cast(f32(pt.output_proj.bias))
+
+        # upsample: transposed conv (kernel = stride, no overlap) + ConvNeXt
+        self.ups = []
+        latent = int(cfg.latent_dim)
+        pending = torch.zeros(latent, device=dev, dtype=torch.float32)
+        for trans, block in decoder.upsample:
+            tw = f32(trans.conv.weight)  # [C_in, C_out, r]
+            r = tw.shape[2]
+            w = tw.permute(0, 2, 1).reshape(latent, r * latent)  # [C_in, r * C_out]
+            bias = (pending @ w).view(r, latent) + f32(trans.conv.bias)[None, :]
+            dw = block.dwconv.conv
+            gamma = f32(block.gamma)
+            self.ups.append(
+                dict(
+                    r=r,
+                    w=cast(w),
+                    b=cast(bias.reshape(-1)),
+                    dw_w=cast(f32(dw.weight)[:, 0, :].t()),
+                    dw_b=cast(f32(dw.bias)),
+                    ln_w=cast(f32(block.norm.weight)),
+                    ln_b=cast(f32(block.norm.bias)),
+                    ln_eps=float(block.norm.eps),
+                    pw1=cast(f32(block.pwconv1.weight).t()),
+                    pw1_b=cast(f32(block.pwconv1.bias)),
+                    pw2=cast((f32(block.pwconv2.weight) * gamma[:, None]).t()),
+                    hist=hist(_KERNEL - 1, latent),
+                )
+            )
+            pending = f32(block.pwconv2.bias) * gamma
+
+        # conv_in (pending bias of the last ConvNeXt added on load)
+        dec = decoder.decoder
+        cin = dec[0].conv
+        self.conv_in = _ConvSpec(cast(_conv_taps(f32(cin.weight))), latent, _KERNEL, 1, hist(_KERNEL - 1, latent))
+        self.conv_in_pb = cast(pending)
+        pending = f32(cin.bias)
+
+        # decoder blocks
+        self.blocks = []
+        for i, blk in enumerate(dec[1:-2]):
+            snake, trans, *units = blk.block
+            c_in = trans.conv.in_channels
+            c_out = trans.conv.out_channels
+            rate = trans.conv.stride[0]
+            tw = f32(trans.conv.weight)  # [C_in, C_out, 2 * rate]
+            entry = dict(
+                pb=cast(pending),
+                snake=_snake_params(snake),
+                rate=rate,
+                c_in=c_in,
+                c_out=c_out,
+                w=cast(tw.permute(0, 2, 1).reshape(c_in, 2 * rate * c_out)),
+                b=cast(f32(trans.conv.bias)),
+                prev=torch.zeros(2, S, rate, c_out, device=dev, dtype=dtype),
+                units=[],
+            )
+            pending = torch.zeros(c_out, device=dev, dtype=torch.float32)
+            for u in units:
+                c1, c2 = u.conv1.conv, u.conv2.conv
+                entry["units"].append(
+                    dict(
+                        pb=cast(pending),
+                        act1=_snake_params(u.act1),
+                        conv1=_ConvSpec(
+                            cast(_conv_taps(f32(c1.weight))),
+                            c_out,
+                            _KERNEL,
+                            c1.dilation[0],
+                            hist((_KERNEL - 1) * c1.dilation[0], c_out),
+                        ),
+                        b1=cast(f32(c1.bias)),
+                        conv1_w4=(
+                            cast(f32(c1.weight)).unsqueeze(2).contiguous(memory_format=torch.channels_last)
+                            if c1.dilation[0] == 1
+                            else None
+                        ),
+                        act2=_snake_params(u.act2),
+                        w2=cast(f32(c2.weight)[:, :, 0].t()),
+                    )
+                )
+                pending = pending + f32(c2.bias)
+            self.blocks.append(entry)
+        self.out_pb = cast(pending)
+        self.out_snake = _snake_params(dec[-2])
+        co = dec[-1].conv
+        self.out_c = co.in_channels
+        self.conv_out_w = cast(f32(co.weight)[0].t())  # [K, C]
+        self.conv_out_b = float(f32(co.bias)[0])
+        self.conv_out = _ConvSpec(self.conv_out_w, self.out_c, _KERNEL, 1, hist(_KERNEL - 1, self.out_c))
+        self.parity = torch.zeros(S, device=dev, dtype=torch.int32)
+        self._dummy = torch.zeros(1, device=dev, dtype=dtype)
+
+    def _slot_state(self, slot: int) -> list[torch.Tensor]:
+        state = [self.parity[slot : slot + 1]]
+        for layer in self.layers:
+            state.extend((layer["k_ring"][slot], layer["v_ring"][slot]))
+        for spec in (self.pre, self.conv_in, self.conv_out):
+            assert spec.hist is not None
+            state.append(spec.hist[:, slot])
+        for up in self.ups:
+            state.append(up["hist"][:, slot])
+        for block in self.blocks:
+            state.append(block["prev"][:, slot])
+            for unit in block["units"]:
+                state.append(unit["conv1"].hist[:, slot])
+        return state
+
+    def save_slot(self, slot: int) -> list[torch.Tensor]:
+        """Offload a preempted request before its runner slot is recycled."""
+        return [tensor.to(device="cpu", copy=True) for tensor in self._slot_state(slot)]
+
+    def restore_slot(self, slot: int, state: list[torch.Tensor]) -> None:
+        for target, saved in zip(self._slot_state(slot), state, strict=True):
+            target.copy_(saved)
+
+    # ------------------------------------------------------------------ helpers
+    def _im2col(self, x, spec: _ConvSpec, T, slots, pos, bias=None, snake=None, ext=False):
+        k = spec.k
+        d = spec.d
+        c = x.shape[1]
+        h = (k - 1) * d
+        B = slots.shape[0]
+        if ext:
+            col = torch.empty(B * (h + T), c, device=x.device, dtype=self.dtype)
+        else:
+            col = torch.empty(B * T, k * c, device=x.device, dtype=self.dtype)
+        bu = 32
+        bc = min(128, triton.next_power_of_2(c))
+        grid = (B, triton.cdiv(h + T, bu), triton.cdiv(c, bc))
+        a, ib = snake if snake is not None else (self._dummy, self._dummy)
+        _im2col_kernel[grid](
+            x, bias if bias is not None else self._dummy, a, ib,
+            spec.hist if h > 0 else self._dummy, slots, pos, self.parity, col,
+            T, self.num_slots, C=c, K=k, D=d, H=h,
+            HAS_BIAS=bias is not None, HAS_SNAKE=snake is not None, BU=bu, BC=bc, EXT=ext,
+        )  # fmt: skip
+        return col
+
+    def _conv_cudnn(self, x, u, T, slots, pos):
+        """Undilated unit conv: one source-row pass, then cuDNN on its channels-last view.
+
+        The im2col path writes a ``K``-times patch matrix; cuDNN reads the
+        ``[history | new]`` rows directly (same result up to summation order).
+        """
+        spec = u["conv1"]
+        ext = self._im2col(x, spec, T, slots, pos, bias=u["pb"], snake=u["act1"], ext=True)
+        b = slots.shape[0]
+        c = x.shape[1]
+        x4 = ext.view(b, -1, c).permute(0, 2, 1).unsqueeze(2)  # [B, C, 1, H + T], channels-last memory
+        y = torch.nn.functional.conv2d(x4, u["conv1_w4"])  # [B, N, 1, T], channels-last memory
+        return y.permute(0, 2, 3, 1).reshape(b * T, -1)
+
+    def _snake(self, x, snake, bias):
+        n, c = x.shape
+        out = torch.empty_like(x)
+        bu, bc = 32, min(128, triton.next_power_of_2(c))
+        one = torch.zeros(1, device=x.device, dtype=torch.int32)
+        _im2col_kernel[(1, triton.cdiv(n, bu), triton.cdiv(c, bc))](
+            x, bias, snake[0], snake[1], self._dummy, one, one, self.parity, out,
+            n, self.num_slots, C=c, K=1, D=1, H=0, HAS_BIAS=True, HAS_SNAKE=True, BU=bu, BC=bc,
+        )  # fmt: skip
+        return out
+
+    # ------------------------------------------------------------------ forward
+    @torch.no_grad()
+    def forward(self, codes: torch.Tensor, slots: torch.Tensor, pos: torch.Tensor) -> torch.Tensor:
+        """``codes`` [B, T, NQ] int, ``slots``/``pos`` [B] int32 -> waveform [B, T * spf] float32.
+
+        ``pos[b]`` is the index of the row's first frame in its stream; 0 starts
+        a new stream in ``slots[b]``. Frames beyond ``max_frames`` per call must
+        be split by the caller.
+        """
+        B, T, _ = codes.shape
+        if T > self.max_frames:
+            raise ValueError(f"at most {self.max_frames} frames per call, got {T}")
+        dev, dt = codes.device, self.dtype
+        codes = codes.reshape(B * T, self.nq).to(torch.int32).contiguous()
+        emb = torch.empty(B * T, 2 * self.half, device=dev, dtype=dt)
+        _rvq_gather_kernel[(B * T,)](codes, self.books, emb, NQ=self.nq, CB=self.cb, HALF=self.half)
+        x = torch.addmm(self.pre_b, self._im2col(emb, self.pre, T, slots, pos), self.pre.weight)
+
+        hid, qd = self.hidden, self.heads * self.head_dim
+        rows = B * T
+        h = torch.empty(rows, hid, device=dev, dtype=dt)
+        q = torch.empty(rows, qd, device=dev, dtype=dt)
+        att = torch.empty(rows, qd, device=dev, dtype=dt)
+        a = torch.empty(rows, self.layers[0]["down"].shape[0], device=dev, dtype=dt)
+        inter = a.shape[1]
+        scale = self.head_dim**-0.5
+        for L in self.layers:
+            _rms_norm_kernel[(rows,)](x, L["ln1"], h, self.eps, N=hid)
+            qkv = torch.mm(h, L["qkv"])
+            _rope_kv_write_kernel[(rows, self.heads)](
+                qkv, q, L["k_ring"], L["v_ring"], slots, pos, T, self.theta,
+                NH=self.heads, HD=self.head_dim, RING=self.RING,
+            )  # fmt: skip
+            _attend_kernel[(rows, self.heads)](
+                q, L["k_ring"], L["v_ring"], att, slots, pos, T, scale,
+                NH=self.heads, HD=self.head_dim, RING=self.RING, WINDOW=self.window,
+                BK=triton.next_power_of_2(self.window),
+            )  # fmt: skip
+            x.addmm_(att, L["o"])
+            _rms_norm_kernel[(rows,)](x, L["ln2"], h, self.eps, N=hid)
+            gu = torch.mm(h, L["gu"])
+            _silu_mul_kernel[(rows, triton.cdiv(inter, 1024))](gu, a, N=inter, BN=1024)
+            x.addmm_(a, L["down"])
+        _rms_norm_kernel[(rows,)](x, self.norm_w, h, self.eps, N=hid)
+        y = torch.addmm(self.out_b, h, self.out_w)  # [rows, latent]
+
+        t_rows = T
+        for U in self.ups:
+            r = U["r"]
+            z = torch.addmm(U["b"], y, U["w"]).view(rows * r, -1)
+            rows, t_rows = rows * r, t_rows * r
+            c = z.shape[1]
+            n = torch.empty_like(z)
+            _dwconv_ln_kernel[(B, max(t_rows, _KERNEL - 1))](
+                z, U["dw_w"], U["dw_b"], U["ln_w"], U["ln_b"], U["hist"], slots, pos, self.parity, n,
+                t_rows, self.num_slots, U["ln_eps"], C=c, K=_KERNEL,
+            )  # fmt: skip
+            p1 = torch.mm(n, U["pw1"])
+            _bias_gelu_kernel[(rows, triton.cdiv(p1.shape[1], 1024))](p1, U["pw1_b"], N=p1.shape[1], BN=1024)
+            z.addmm_(p1, U["pw2"])
+            y = z
+
+        x = torch.mm(self._im2col(y, self.conv_in, t_rows, slots, pos, bias=self.conv_in_pb), self.conv_in.weight)
+        for blk in self.blocks:
+            s = self._snake(x, blk["snake"], blk["pb"])
+            zt = torch.mm(s, blk["w"])  # [rows, 2 * rate * c_out]
+            r, c_out = blk["rate"], blk["c_out"]
+            x = torch.empty(rows * r, c_out, device=dev, dtype=dt)
+            bc = min(128, triton.next_power_of_2(c_out))
+            _col2im_kernel[(rows, r, triton.cdiv(c_out, bc))](
+                zt, blk["b"], blk["prev"], slots, pos, self.parity, x, t_rows, self.num_slots,
+                R=r, C=c_out, BC=bc,
+            )  # fmt: skip
+            rows, t_rows = rows * r, t_rows * r
+            for u in blk["units"]:
+                if u.get("conv1_w4") is not None:
+                    c1 = self._conv_cudnn(x, u, t_rows, slots, pos)
+                else:
+                    col = self._im2col(x, u["conv1"], t_rows, slots, pos, bias=u["pb"], snake=u["act1"])
+                    c1 = torch.mm(col, u["conv1"].weight)
+                s2 = self._snake(c1, u["act2"], u["b1"])
+                x.addmm_(s2, u["w2"])
+
+        ext = self._im2col(x, self.conv_out, t_rows, slots, pos, bias=self.out_pb, snake=self.out_snake, ext=True)
+        wav = torch.empty(B, t_rows, device=dev, dtype=torch.float32)
+        _conv_out_kernel[(B, triton.cdiv(t_rows, 256))](
+            ext, self.conv_out_w, self.conv_out_b, wav, t_rows,
+            C=self.out_c, K=_KERNEL, BT=256, CP=triton.next_power_of_2(self.out_c),
+        )  # fmt: skip
+        _flip_kernel[(B,)](self.parity, slots)
+        return wav
+
+
+class StreamingDecodeGraphs:
+    """CUDA graphs of ``StreamingCodecDecoder`` for a fixed frame count per batch bucket."""
+
+    def __init__(self, sd: StreamingCodecDecoder, batch_sizes: list[int], frames: int = 1) -> None:
+        self.sd, self.frames = sd, frames
+        self.graphs: dict[int, tuple] = {}
+        dev = sd.device
+        pool = torch.cuda.graph_pool_handle()
+        for bsz in sorted(set(batch_sizes)):
+            codes = torch.zeros(bsz, frames, sd.nq, dtype=torch.int32, device=dev)
+            slots = torch.full((bsz,), sd.scratch_slot, dtype=torch.int32, device=dev)
+            pos = torch.zeros(bsz, dtype=torch.int32, device=dev)
+            sd(codes, slots, pos)
+            current_omni_platform.synchronize()
+            graph = torch.cuda.CUDAGraph()
+            with torch.cuda.graph(graph, pool=pool):
+                out = sd(codes, slots, pos)
+            self.graphs[bsz] = (graph, codes, slots, pos, out)
+        self.sizes = sorted(self.graphs)
+
+    def __call__(self, codes: torch.Tensor, slots: torch.Tensor, pos: torch.Tensor) -> torch.Tensor:
+        """``codes`` [n, frames, NQ], ``slots``/``pos`` [n] on device -> waveform view [n, frames * spf]."""
+        n = int(codes.shape[0])
+        size = next((b for b in self.sizes if b >= n), None)
+        if size is None:
+            return self.sd(codes, slots.to(torch.int32), pos.to(torch.int32))
+        graph, s_codes, s_slots, s_pos, out = self.graphs[size]
+        s_codes[:n].copy_(codes)
+        s_slots[:n].copy_(slots)
+        s_pos[:n].copy_(pos)
+        if n < size:
+            s_slots[n:].fill_(self.sd.scratch_slot)
+            s_pos[n:].zero_()
+        graph.replay()
+        return out[:n]

vllm_omni/model_executor/models/registry.py

@@ -483,6 +483,12 @@ _OMNI_MODELS = {
         "acoustic",
         "MiniMaxMusic3AcousticForConditionalGeneration",
     ),
+    ## YuE2-3B (text-to-music; single-stage native-AR with terminal NAR/VAE)
+    "Yue2ForCausalLM": (
+        "yue2",
+        "yue2",
+        "Yue2ForCausalLM",
+    ),
 }
 
 

vllm_omni/model_executor/models/yue2/__init__.py

@@ -0,0 +1,7 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""YuE2-3B text-to-music model package."""
+
+from .pipeline import YUE2_PIPELINE
+
+__all__ = ["YUE2_PIPELINE"]

vllm_omni/model_executor/models/yue2/pipeline.py

@@ -0,0 +1,53 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""YuE2-3B pipeline topology.
+
+Single-stage native-AR text-to-music (gepard pattern). The whole
+Mixture-of-Transformers checkpoint — AR path, NAR path, projection heads —
+lives in one ``model.safetensors``, so both the autoregressive sampling and
+the terminal flow-matching + VAE pass run inside one stage and no weights are
+duplicated across stages. The VAE decoder is a separate repository
+(``m-a-p/YuE2-Vae``), loaded in ``load_weights`` from ``$YUE2_VAE`` or the
+default hub id.
+
+Sampling is model-owned. ``sampling_constraints`` pins ``detokenize`` and the
+union of both phase end tokens; ``extra_args`` carry the phase and the
+request-local sampling preset, and must be present in the stage's
+``default_sampling_params`` so ``has_sampling_extra_args`` turns on and the
+per-request args reach the model's forward.
+"""
+
+from vllm_omni.config.stage_config import (
+    PipelineConfig,
+    StageExecutionType,
+    StagePipelineConfig,
+)
+
+from .yue2 import STOP_TOKEN_IDS
+
+YUE2_PIPELINE = PipelineConfig(
+    model_type="yue2",
+    default_deploy_config_name="yue2.yaml",
+    model_arch="Yue2ForCausalLM",
+    stages=(
+        StagePipelineConfig(
+            stage_id=0,
+            model_stage="yue2",
+            execution_type=StageExecutionType.LLM_AR,
+            input_sources=(),
+            final_output=True,
+            final_output_type="audio",
+            owns_tokenizer=True,
+            engine_output_type="audio",
+            # Whole-song audio ships once, on the finishing step; the audio
+            # accumulator concatenates over time and a single tensor just
+            # passes through.
+            sampling_constraints={
+                "detokenize": False,
+                "stop_token_ids": list(STOP_TOKEN_IDS),
+            },
+        ),
+    ),
+)
+
+__all__ = ["YUE2_PIPELINE"]