[Feature] Support MOSS-TTS-Local-v1.5 vocoder graph (#4929)

Signed-off-by: gcanlin <canlinguosdu@gmail.com>
Signed-off-by: sunnyxyli <1550693033@qq.com>
Co-authored-by: sunnyxyli <1550693033@qq.com>

vllm_omni/benchmarks/data_modules/seed_tts_eval.py

@@ -613,7 +613,12 @@ def compute_seed_tts_wer_metrics(
                 e,
             )
 
-        wav_16k = _pcm_s16le_to_f32_16k(pcm)
+        # Request functions normalize ``tts_output_pcm_bytes`` to Seed-TTS WER
+        # format before it reaches this evaluator: 24 kHz mono int16 PCM.
+        # Do not apply the streamed-response env format here again; for MOSS
+        # local that would reinterpret already-normalized 24 kHz mono audio as
+        # 48 kHz stereo and corrupt ASR/WER.
+        wav_16k = _pcm_s16le_to_f32_16k(pcm, pcm_sample_rate=24000, channels=1)
         if len(wav_16k) == 0:
             asr_failed += 1
             if include_per_item:

vllm_omni/model_executor/models/moss_tts/cuda_graph_streaming_decoder_wrapper.py

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

vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_codec.py

@@ -26,6 +26,9 @@ from vllm_omni.model_executor.models.moss_tts.audio_tokenizer_v2 import (
 from vllm_omni.model_executor.models.moss_tts.configuration_moss_audio_tokenizer_v2 import (
     MossAudioTokenizerConfig as MossAudioTokenizerV2Config,
 )
+from vllm_omni.model_executor.models.moss_tts.cuda_graph_streaming_decoder_wrapper import (
+    CUDAGraphStreamingDecoderWrapper,
+)
 from vllm_omni.model_executor.models.moss_tts.moss_codec_cudagraph import (
     MossTTSCUDAGraphCodecWrapper,
 )
@@ -43,6 +46,7 @@ class _MossCodecStreamSession:
         *,
         stream_slots: int,
         n_vq: int,
+        cudagraph_capture_sizes: list[int] | None = None,
     ) -> None:
         self._codec = codec
         self._stream_slots = int(stream_slots)
@@ -54,6 +58,19 @@ class _MossCodecStreamSession:
         self._closed = False
         with torch.no_grad():
             self._exit_stack.enter_context(codec.streaming(self._batch_size))
+        self._cudagraph_wrapper: CUDAGraphStreamingDecoderWrapper | None = None
+        capture_sizes = sorted({int(size) for size in (cudagraph_capture_sizes or []) if int(size) > 0})
+        if capture_sizes and self._device.type == "cuda":
+            self._cudagraph_wrapper = CUDAGraphStreamingDecoderWrapper(
+                codec,
+                batch_size=self._batch_size,
+                num_quantizers=self._n_vq,
+                reset_streaming_state=lambda: self.reset_slots(list(range(self._batch_size))),
+            )
+            self._cudagraph_wrapper.warmup(self._device, capture_sizes)
+            self.reset_slots(list(range(self._batch_size)))
+            if not self._cudagraph_wrapper.is_ready:
+                self._cudagraph_wrapper = None
 
     def acquire(self) -> int | None:
         if not self._free_stream_slots:
@@ -109,12 +126,23 @@ class _MossCodecStreamSession:
             codes_lengths[slot] = int(codes.shape[1])
             exec_mask[slot] = True
 
-        self._codec._set_streaming_exec_mask(exec_mask)
-        result = self._codec._decode_frame(codes_step, codes_lengths)
-        if result.audio is None:
-            return {}
-        audio = result.audio.detach().to("cpu", torch.float32)
-        lengths = result.audio_lengths.detach().to("cpu") if result.audio_lengths is not None else None
+        graph_output: tuple[torch.Tensor, torch.Tensor] | None = None
+        if self._cudagraph_wrapper is not None:
+            graph_output = self._cudagraph_wrapper.decode(codes_step, exec_mask)
+
+        used_cudagraph = graph_output is not None
+        if used_cudagraph:
+            audio_tensor, lengths_tensor = graph_output
+        else:
+            self._codec._set_streaming_exec_mask(exec_mask)
+            result = self._codec._decode_frame(codes_step, codes_lengths)
+            if result.audio is None:
+                return {}
+            audio_tensor = result.audio
+            lengths_tensor = result.audio_lengths
+
+        audio = audio_tensor.detach().to("cpu", torch.float32)
+        lengths = lengths_tensor.detach().to("cpu") if lengths_tensor is not None else None
         out: dict[int, torch.Tensor] = {}
         for slot in slot_codes:
             wav = audio[slot]
@@ -174,10 +202,17 @@ class MossTTSCodecDecoder(nn.Module):
         self._sr_tensor = torch.tensor(self._OUTPUT_SAMPLE_RATE, dtype=torch.int32)
         self._stream_session: _MossCodecStreamSession | None = None
         self._stream_slots: int = self._connector_int("codec_stream_slots", default=0)
-        self._stream_max_step_frames: int = self._connector_int("codec_max_step_frames", default=100)
+        self._stream_chunk_frames: int = self._connector_int("codec_chunk_frames", default=0)
+        self._stream_max_step_frames: int = self._stream_chunk_frames or 100
         self._stream_req_slots: dict[str, int] = {}
         self._stream_pending_codes: dict[str, list[torch.Tensor]] = {}
         self._stream_starved_reqs: set[str] = set()
+        self._codec_streaming: bool = self._connector_bool("codec_streaming", default=False)
+        default_graph_capture_max = self._stream_chunk_frames or min(self._stream_max_step_frames, 32)
+        default_graph_capture_sizes = list(range(1, max(1, default_graph_capture_max) + 1))
+        self._streaming_cudagraph_capture_sizes = self._streaming_cudagraph_capture_sizes_from_compilation_config(
+            default_graph_capture_sizes
+        )
 
     # ------------------------------------------------------------------
     # vLLM-Omni stubs (codec has no AR loop)
@@ -240,6 +275,15 @@ class MossTTSCodecDecoder(nn.Module):
                 },
             )
 
+        if self._codec_streaming and runtime_additional_information is None:
+            return OmniOutput(
+                text_hidden_states=None,
+                multimodal_outputs={
+                    "model_outputs": [empty] * num_req,
+                    "sr": [sr_tensor] * num_req,
+                },
+            )
+
         audios: list[torch.Tensor] = [empty] * num_req
         srs: list[torch.Tensor] = [sr_tensor] * num_req
         device = next(self._codec.parameters()).device
@@ -293,7 +337,7 @@ class MossTTSCodecDecoder(nn.Module):
                 continue
             meta = (info.get("meta", {}) if isinstance(info, dict) else {}) or {}
             finished = bool(meta.get("stream_finished", meta.get("finished", False)))
-            streaming_enabled = bool(meta.get("codec_streaming", False))
+            streaming_enabled = bool(meta.get("codec_streaming", self._codec_streaming))
             code_flat_numel = meta.get("code_flat_numel")
             if streaming_enabled and finished and code_flat_numel is not None and int(code_flat_numel) == 0:
                 for _, wav in self._finish_empty_streaming_requests([info]).items():
@@ -382,7 +426,7 @@ class MossTTSCodecDecoder(nn.Module):
             if not isinstance(info, dict):
                 continue
             meta = (info.get("meta", {}) or {}) if isinstance(info.get("meta", {}), dict) else {}
-            if not bool(meta.get("codec_streaming", False)):
+            if not bool(meta.get("codec_streaming", self._codec_streaming)):
                 continue
             finished = bool(meta.get("stream_finished", meta.get("finished", False)))
             if not finished:
@@ -446,6 +490,7 @@ class MossTTSCodecDecoder(nn.Module):
             self._codec,
             stream_slots=max(1, slots),
             n_vq=self._n_vq,
+            cudagraph_capture_sizes=self._streaming_cudagraph_capture_sizes,
         )
         return self._stream_session
 
@@ -597,6 +642,30 @@ class MossTTSCodecDecoder(nn.Module):
             return int(extra_cfg[name])
         return default
 
+    def _connector_bool(self, name: str, default: bool = False) -> bool:
+        model_cfg = getattr(self.vllm_config, "model_config", None)
+        connector_cfg = getattr(model_cfg, "stage_connector_config", None)
+        if isinstance(connector_cfg, dict):
+            extra_cfg: dict | None = connector_cfg.get("extra", connector_cfg)
+        else:
+            extra_cfg = getattr(connector_cfg, "extra", None)
+        if isinstance(extra_cfg, dict) and name in extra_cfg:
+            return bool(extra_cfg[name])
+        return default
+
+    def _streaming_cudagraph_capture_sizes_from_compilation_config(
+        self,
+        default: list[int],
+    ) -> list[int]:
+        if getattr(self.vllm_config.model_config, "enforce_eager", True):
+            return []
+        compilation_config = getattr(self.vllm_config, "compilation_config", None)
+        capture_sizes = getattr(compilation_config, "cudagraph_capture_sizes", None)
+        if capture_sizes:
+            max_step_frames = max(1, int(self._stream_max_step_frames))
+            return sorted({int(size) for size in capture_sizes if 0 < int(size) <= max_step_frames})
+        return list(default)
+
     # ------------------------------------------------------------------
     # Weight loading
     # ------------------------------------------------------------------
@@ -716,7 +785,9 @@ class MossTTSCodecDecoder(nn.Module):
             self._n_channels,
         )
 
-        self._maybe_enable_decoder_cudagraph(device)
+        self._configure_decoder_cudagraph(device)
+        if self._codec_streaming and self._streaming_cudagraph_capture_sizes:
+            self._ensure_stream_session()
 
         # vLLM's track_weights_loading() compares the returned set against
         # ``self.named_parameters()``. After ``self._codec = codec`` above,
@@ -739,12 +810,30 @@ class MossTTSCodecDecoder(nn.Module):
         codec = MossAudioTokenizerModel(codec_cfg)
         return codec_cfg, codec
 
-    def _maybe_enable_decoder_cudagraph(self, device: torch.device) -> None:
-        """Capture CUDA Graphs for the codec decoder if enforce_eager is False."""
+    def _configure_decoder_cudagraph(self, device: torch.device) -> None:
+        """Select the codec CUDA Graph path.
+
+        ``enforce_eager`` is the single graph on/off switch. If graphing is
+        enabled, ``codec_streaming`` decides whether decode uses the persistent
+        streaming-state wrapper or the offline full-chunk wrapper.
+        """
         if getattr(self.vllm_config.model_config, "enforce_eager", True):
+            self._streaming_cudagraph_capture_sizes = []
             return
         if self._codec is None:
             return
+        if self._codec_streaming:
+            logger.info(
+                "MOSS-TTS codec CUDA Graph selected streaming wrapper: capture_sizes=%s",
+                self._streaming_cudagraph_capture_sizes,
+            )
+            return
+
+        self._enable_non_streaming_decoder_cudagraph(device)
+
+    def _enable_non_streaming_decoder_cudagraph(self, device: torch.device) -> None:
+        if self._codec is None:
+            return
 
         # Read capture sizes from the connector's extra config (same convention
         # as Qwen3-TTS), falling back to a sensible default covering common

vllm_omni/model_executor/models/moss_tts/moss_codec_cudagraph.py

@@ -1,6 +1,6 @@
 """CUDA Graph acceleration for the MOSS Audio Tokenizer codec decoder.
 
-Captures the codec decode method (_decode_frame, or _decode on the vendored codec) for a set of fixed frame-count
+Captures MossAudioTokenizerModel._decode_frame for a set of fixed frame-count
 bucket sizes, then replays the captured graph at inference time to eliminate
 kernel-launch overhead.  Inputs that exceed all captured sizes fall back to
 eager execution transparently.
@@ -24,7 +24,7 @@ logger = init_logger(__name__)
 
 
 class MossTTSCUDAGraphCodecWrapper:
-    """CUDA Graph wrapper for the codec decode method (_decode_frame, or _decode on the vendored codec).
+    """CUDA Graph wrapper for MossAudioTokenizerModel._decode_frame.
 
     Graphs are keyed by padded_T (int).  On each call the actual T is
     bucket-matched to the smallest pre-captured size >= T.  The static code
@@ -51,15 +51,6 @@ class MossTTSCUDAGraphCodecWrapper:
         enabled: bool = True,
     ) -> None:
         self.model = model
-        # Upstream remote-code tokenizer exposes _decode_frame; the vendored
-        # MossAudioTokenizerModel still exposes _decode. Resolve once here so
-        # both codec implementations work.
-        decode_fn = getattr(model, "_decode_frame", None) or getattr(model, "_decode", None)
-        if decode_fn is None:
-            raise AttributeError(
-                f"{type(model).__module__}.{type(model).__name__} exposes neither _decode_frame nor _decode"
-            )
-        self._decode_fn = decode_fn
         self.capture_sizes: list[int] = sorted(capture_sizes)
         self.num_quantizers = num_quantizers
         self.enabled = enabled
@@ -108,7 +99,7 @@ class MossTTSCUDAGraphCodecWrapper:
             dummy_codes = torch.zeros(nq, 1, size, dtype=torch.long, device=device)
             dummy_lengths = torch.tensor([size], dtype=torch.long, device=device)
             with torch.no_grad():
-                _ = self._decode_fn(dummy_codes, dummy_lengths)
+                _ = self.model._decode_frame(dummy_codes, dummy_lengths)
 
         torch.accelerator.synchronize(device)
 
@@ -141,13 +132,13 @@ class MossTTSCUDAGraphCodecWrapper:
 
         # Extra eager warmup inside capture to ensure all kernels are compiled.
         with torch.no_grad():
-            _ = self._decode_fn(static_codes, static_lengths)
+            _ = self.model._decode_frame(static_codes, static_lengths)
         torch.accelerator.synchronize(device)
 
         graph = CUDAGraph()
         with torch.no_grad():
             with torch.cuda.graph(graph, pool=current_platform.get_global_graph_pool()):
-                static_out = self._decode_fn(static_codes, static_lengths)
+                static_out = self.model._decode_frame(static_codes, static_lengths)
 
         self.graphs[size] = graph
         self.static_codes[size] = static_codes

vllm_omni/worker/gpu_generation_model_runner.py

@@ -76,10 +76,6 @@ class GPUGenerationModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin
 
     def _update_request_states(self, scheduler_output: SchedulerOutput):
         # remove requests
-        # Some stateful vocoder model may need to clean the state
-        # to avoid the leak of slots when the requests have been aborted.
-        if scheduler_output.finished_req_ids and hasattr(self.model, "on_requests_finished"):
-            self.model.on_requests_finished(scheduler_output.finished_req_ids)
         for req_id in scheduler_output.finished_req_ids:
             self.input_batch.remove_request(req_id)
         scheduled_req_ids = scheduler_output.num_scheduled_tokens.keys()
@@ -157,6 +153,13 @@ class GPUGenerationModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin
             if self.model_config.async_chunk and num_scheduled_tokens:
                 self._update_request_states(scheduler_output)
             deferred_state_corrections_fn = self._update_states(scheduler_output)
+
+            # Notify stateful models of finished requests before the
+            # zero-token early return. Request cleanup is a model lifecycle
+            # concern and does not depend on the inter-stage transfer mode.
+            if scheduler_output.finished_req_ids and hasattr(self.model, "on_requests_finished"):
+                self.model.on_requests_finished(scheduler_output.finished_req_ids)
+
             if not scheduler_output.total_num_scheduled_tokens:
                 return self.attach_omni_connector_output(EMPTY_MODEL_RUNNER_OUTPUT)