From 5c3900960bc0d500c05619f02ee99211ecb81d73 Mon Sep 17 00:00:00 2001 From: Canlin Guo Date: Thu, 16 Jul 2026 15:02:15 +0800 Subject: [PATCH] [Feature] Support MOSS-TTS-Local-v1.5 vocoder graph (#4929) Signed-off-by: gcanlin Signed-off-by: sunnyxyli <1550693033@qq.com> Co-authored-by: sunnyxyli <1550693033@qq.com> --- .../benchmarks/data_modules/seed_tts_eval.py | 7 +- .../cuda_graph_streaming_decoder_wrapper.py | 164 ++++++++++++++++++ .../moss_tts/modeling_moss_tts_codec.py | 113 ++++++++++-- .../models/moss_tts/moss_codec_cudagraph.py | 19 +- .../worker/gpu_generation_model_runner.py | 11 +- 5 files changed, 283 insertions(+), 31 deletions(-) create mode 100644 vllm_omni/model_executor/models/moss_tts/cuda_graph_streaming_decoder_wrapper.py diff --git a/vllm_omni/benchmarks/data_modules/seed_tts_eval.py b/vllm_omni/benchmarks/data_modules/seed_tts_eval.py index 70ed7cfef..13c1f2c21 100644 --- a/vllm_omni/benchmarks/data_modules/seed_tts_eval.py +++ b/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: diff --git a/vllm_omni/model_executor/models/moss_tts/cuda_graph_streaming_decoder_wrapper.py b/vllm_omni/model_executor/models/moss_tts/cuda_graph_streaming_decoder_wrapper.py new file mode 100644 index 000000000..1daa5081e --- /dev/null +++ b/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"] diff --git a/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_codec.py b/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_codec.py index 60f455240..8f1a72e9b 100644 --- a/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_codec.py +++ b/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 diff --git a/vllm_omni/model_executor/models/moss_tts/moss_codec_cudagraph.py b/vllm_omni/model_executor/models/moss_tts/moss_codec_cudagraph.py index 71585f1c0..87f255072 100644 --- a/vllm_omni/model_executor/models/moss_tts/moss_codec_cudagraph.py +++ b/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 diff --git a/vllm_omni/worker/gpu_generation_model_runner.py b/vllm_omni/worker/gpu_generation_model_runner.py index c8fe92eba..ec41900b1 100644 --- a/vllm_omni/worker/gpu_generation_model_runner.py +++ b/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)