diff --git a/vllm_omni/deploy/moss_tts_local.yaml b/vllm_omni/deploy/moss_tts_local.yaml index 7d3ae8601..2ed937e0a 100644 --- a/vllm_omni/deploy/moss_tts_local.yaml +++ b/vllm_omni/deploy/moss_tts_local.yaml @@ -54,7 +54,7 @@ stages: # torch_profiler_dir: /root/vllm-omni-workspace/vllm-omni/results/moss_local_stage1/batch_8 max_num_batched_tokens: 65536 max_model_len: 65536 - devices: "1" + devices: "0" input_connectors: from_stage_0: shm default_sampling_params: diff --git a/vllm_omni/entrypoints/openai/serving_speech.py b/vllm_omni/entrypoints/openai/serving_speech.py index 1ee2ce668..9cc6118e9 100644 --- a/vllm_omni/entrypoints/openai/serving_speech.py +++ b/vllm_omni/entrypoints/openai/serving_speech.py @@ -1764,7 +1764,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): # headline improvement is multilingual synthesis when the language is # given (build_user_message(..., language=...)); 1.0 ignores it # gracefully. Sound-effect output is non-verbal, so skip it there. - if v in ("tts", "ttsd", "voice_generator") and getattr(request, "language", None): + if v in ("tts", "local", "ttsd", "voice_generator") and getattr(request, "language", None): user_kwargs["language"] = request.language # Build the unified-codes prompt: (L, 1+n_vq) where col 0 is text/special diff --git a/vllm_omni/model_executor/models/moss_tts/audio_tokenizer_v2.py b/vllm_omni/model_executor/models/moss_tts/audio_tokenizer_v2.py index 26ea9c567..098f53792 100644 --- a/vllm_omni/model_executor/models/moss_tts/audio_tokenizer_v2.py +++ b/vllm_omni/model_executor/models/moss_tts/audio_tokenizer_v2.py @@ -484,7 +484,6 @@ class RingKVCache: return slots = state_slot_ids.to(device=self.end_offset.device, dtype=torch.long) self.end_offset.index_fill_(0, slots, 0) - self.cache.index_fill_(1, slots, 0) def complete( self, @@ -1607,6 +1606,7 @@ class MossAudioTokenizerModel(MossAudioTokenizerPreTrainedModel): self._streaming_modules: list[StreamingModule] = [] self._streaming_exec_mask: torch.Tensor | None = None self._decoder_state_capacity = 0 + self._decoder_slot_offsets: torch.Tensor | None = None self.post_init() def _start_streaming(self, batch_size: int, *, decoder_only: bool = False) -> None: @@ -1634,6 +1634,49 @@ class MossAudioTokenizerModel(MossAudioTokenizerPreTrainedModel): root.apply(_start) self._streaming_modules = streaming_modules self._streaming_exec_mask = shared_exec_mask + if decoder_only: + self._centralize_decoder_slot_offsets(batch_size) + + def _centralize_decoder_slot_offsets(self, state_capacity: int) -> None: + """Pack per-module slot offsets so reset needs one CUDA kernel.""" + state_tensors: list[tuple[StreamingState, list[tuple[object, str]]]] = [] + for module in self._streaming_modules: + state = module._streaming_state + fields: list[tuple[object, str]] = [] + if isinstance(state, MHAState): + fields.append((state, "offset")) + if state.kv_cache is not None: + fields.append((state.kv_cache, "end_offset")) + elif isinstance(state, TransformerState): + fields.append((state, "offsets")) + else: + continue + + if not all( + isinstance(getattr(owner, name), torch.Tensor) + and getattr(owner, name).shape == (state_capacity,) + and getattr(owner, name).dtype == torch.long + for owner, name in fields + ): + raise RuntimeError( + "Dynamic decoder state slots require every per-slot offset to use the full decoder state capacity." + ) + state_tensors.append((state, fields)) + + num_offsets = sum(len(fields) for _, fields in state_tensors) + if num_offsets == 0: + raise RuntimeError("Dynamic decoder state pool has no per-slot offsets.") + storage = torch.zeros( + (num_offsets, state_capacity), + dtype=torch.long, + device=next(self.parameters()).device, + ) + row = 0 + for state, fields in state_tensors: + for owner, name in fields: + setattr(owner, name, storage[row]) + row += 1 + self._decoder_slot_offsets = storage def _stop_streaming(self) -> None: """Stop streaming mode for all modules.""" @@ -1642,6 +1685,7 @@ class MossAudioTokenizerModel(MossAudioTokenizerPreTrainedModel): self._streaming_modules = [] self._streaming_exec_mask = None 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.""" @@ -1658,18 +1702,16 @@ class MossAudioTokenizerModel(MossAudioTokenizerPreTrainedModel): raise RuntimeError("MOSS Audio Tokenizer decoder state pool is not initialized.") if state_slot_ids.numel() == 0: return - slots = state_slot_ids.to(device=next(self.parameters()).device, dtype=torch.long) - if not torch.cuda.is_current_stream_capturing(): - slots_cpu = slots.detach().to("cpu") - if int(slots_cpu.min()) < 0 or int(slots_cpu.max()) >= self._decoder_state_capacity: + slot_offsets = self._decoder_slot_offsets + if slot_offsets is None: + raise RuntimeError("MOSS Audio Tokenizer decoder slot offsets are not initialized.") + if state_slot_ids.device.type == "cpu": + if int(state_slot_ids.min()) < 0 or int(state_slot_ids.max()) >= self._decoder_state_capacity: raise ValueError( - f"Decoder state slots must be in [0, {self._decoder_state_capacity}), got {slots_cpu.tolist()}" + f"Decoder state slots must be in [0, {self._decoder_state_capacity}), got {state_slot_ids.tolist()}" ) - for module in self._streaming_modules: - state = module._streaming_state - reset_slots = getattr(state, "reset_slots", None) - if callable(reset_slots): - reset_slots(slots) + slots = state_slot_ids.to(device=slot_offsets.device, dtype=torch.long) + slot_offsets.index_fill_(1, slots, 0) def decode_streaming_batch( self, @@ -1686,7 +1728,27 @@ class MossAudioTokenizerModel(MossAudioTokenizerPreTrainedModel): state_capacity=self._decoder_state_capacity, device=codes.device, ) - return self._decode_frame(codes, codes_lengths, execution_context=execution_context) + audio, audio_lengths = self.decode_streaming_tensors( + codes, + codes_lengths, + state_slot_ids, + valid_rows, + ) + return MossAudioTokenizerDecoderOutput(audio=audio, audio_lengths=audio_lengths) + + def decode_streaming_tensors( + self, + codes: torch.Tensor, + codes_lengths: torch.Tensor, + state_slot_ids: torch.Tensor, + valid_rows: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Tensor-only streaming decode boundary for vLLM compilation.""" + execution_context = StreamingExecutionContext( + state_slot_ids=state_slot_ids, + valid_rows=valid_rows, + ) + return self._decode_frame_tensors(codes, codes_lengths, execution_context=execution_context) def _set_streaming_exec_mask(self, exec_mask: torch.Tensor) -> None: """Update the shared active-slot mask with one device copy.""" @@ -1925,13 +1987,13 @@ class MossAudioTokenizerModel(MossAudioTokenizerPreTrainedModel): ) @torch.no_grad() - def _decode_frame( + def _decode_frame_tensors( self, codes: torch.Tensor, codes_lengths: torch.Tensor | None = None, execution_context: StreamingExecutionContext | None = None, - ) -> MossAudioTokenizerDecoderOutput: - """Detokenize discrete tokens into audio waveform.""" + ) -> tuple[torch.Tensor, torch.Tensor]: + """Detokenize codes and return tensors without Python output wrappers.""" nq, B, T = codes.shape device = codes.device @@ -1955,7 +2017,22 @@ class MossAudioTokenizerModel(MossAudioTokenizerPreTrainedModel): d, d_lengths = decoder_module(d, d_lengths) d, d_lengths = self._restore_channels_from_codec(d, d_lengths) - return MossAudioTokenizerDecoderOutput(audio=d, audio_lengths=d_lengths) + return d, d_lengths + + @torch.no_grad() + def _decode_frame( + self, + codes: torch.Tensor, + codes_lengths: torch.Tensor | None = None, + execution_context: StreamingExecutionContext | None = None, + ) -> MossAudioTokenizerDecoderOutput: + """Compatibility wrapper around the tensor-only decoder.""" + audio, audio_lengths = self._decode_frame_tensors( + codes, + codes_lengths, + execution_context=execution_context, + ) + return MossAudioTokenizerDecoderOutput(audio=audio, audio_lengths=audio_lengths) def encode( # type: ignore[override] self, 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 index 4aeb17785..52815dd3b 100644 --- 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 @@ -5,18 +5,52 @@ from __future__ import annotations +import copy import time from dataclasses import dataclass import torch import torch.nn as nn from torch.cuda import CUDAGraph +from vllm.compilation.decorators import support_torch_compile +from vllm.config import CUDAGraphMode, VllmConfig +from vllm.config.vllm import set_current_vllm_config from vllm.logger import init_logger from vllm.platforms import current_platform logger = init_logger(__name__) +@support_torch_compile( + dynamic_arg_dims={ + "codes": {1: "batch", 2: "frames"}, + "codes_lengths": {0: "batch"}, + "state_slot_ids": {0: "batch"}, + "valid_rows": {0: "batch"}, + } +) +class _MossStreamingDecodeCompileAdapter(nn.Module): + """Expose the BF16 codec streaming hot path to vLLM compile. + + The dtype is part of the class identity intentionally: vLLM's AOT cache + drops runtime guards, so an artifact traced with the old FP32 decoder + weights must never be reused after the decoder is materialized as BF16. + """ + + def __init__(self, codec: nn.Module, *, vllm_config: VllmConfig) -> None: + super().__init__() + self.codec = codec + + def forward( + self, + codes: torch.Tensor, + codes_lengths: torch.Tensor, + state_slot_ids: torch.Tensor, + valid_rows: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + return self.codec.decode_streaming_tensors(codes, codes_lengths, state_slot_ids, valid_rows) + + @dataclass class _CapturedStreamingDecodeGraph: graph: CUDAGraph @@ -44,6 +78,7 @@ class CUDAGraphStreamingDecoderWrapper: batch_sizes: list[int], frame_sizes: list[int], num_quantizers: int, + vllm_config: VllmConfig, ) -> None: self.codec = codec self.state_capacity = int(state_capacity) @@ -53,6 +88,18 @@ class CUDAGraphStreamingDecoderWrapper: self.graphs: dict[tuple[int, int], _CapturedStreamingDecodeGraph] = {} self._pool = None self._warmed_up = False + # 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. + compile_config = copy.copy(vllm_config) + compile_config.compilation_config = copy.copy(vllm_config.compilation_config) + 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( + codec, + vllm_config=compile_config, + ) @property def is_ready(self) -> bool: @@ -88,8 +135,12 @@ class CUDAGraphStreamingDecoderWrapper: for batch_size, frame_size in capture_keys: key = (batch_size, frame_size) try: - self._capture(batch_size, frame_size, device) - logger.info(" Captured MOSS-TTS streaming decoder CUDA graph for (B,T)=%s", key) + compiled = self._capture(batch_size, frame_size, device) + logger.info( + " Captured %s MOSS-TTS streaming decoder CUDA graph for (B,T)=%s", + "compiled" if compiled else "plain", + key, + ) except Exception: self.graphs.pop(key, None) logger.warning( @@ -105,7 +156,46 @@ class CUDAGraphStreamingDecoderWrapper: ) @torch.no_grad() - def _capture(self, batch_size: int, frame_size: int, device: torch.device) -> None: + def _capture(self, batch_size: int, frame_size: int, device: torch.device) -> bool: + if self._compiled_decode is not None: + try: + self._capture_with_decode( + batch_size, + frame_size, + device, + self._compiled_decode, + ) + return True + except Exception: + logger.warning( + "vLLM compile/capture failed for MOSS-TTS (B,T)=(%d,%d); falling back to a plain CUDA Graph", + batch_size, + frame_size, + exc_info=True, + ) + scratch_slots = self.state_capacity + torch.arange( + batch_size, + dtype=torch.long, + device=device, + ) + self.codec.reset_decoder_state_slots(scratch_slots) + self._compiled_decode = None + self._capture_with_decode( + batch_size, + frame_size, + device, + self.codec.decode_streaming_tensors, + ) + return False + + @torch.no_grad() + def _capture_with_decode( + self, + batch_size: int, + frame_size: int, + device: torch.device, + decode, + ) -> None: codes = torch.zeros( self.num_quantizers, batch_size, @@ -120,7 +210,8 @@ class CUDAGraphStreamingDecoderWrapper: stream = torch.cuda.Stream() stream.wait_stream(torch.cuda.current_stream()) with torch.cuda.stream(stream): - _ = self.codec.decode_streaming_batch(codes, lengths, scratch_slots, valid_rows) + for _ in range(3): + _ = decode(codes, lengths, scratch_slots, valid_rows) torch.cuda.current_stream().wait_stream(stream) torch.accelerator.synchronize(device) self.codec.reset_decoder_state_slots(scratch_slots) @@ -129,18 +220,16 @@ class CUDAGraphStreamingDecoderWrapper: 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_streaming_batch(codes, lengths, scratch_slots, valid_rows) + audio, audio_lengths = decode(codes, lengths, scratch_slots, valid_rows) - 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[(batch_size, frame_size)] = _CapturedStreamingDecodeGraph( graph=graph, static_codes=codes, static_lengths=lengths, static_state_slot_ids=scratch_slots, static_valid_rows=valid_rows, - static_audio=output.audio, - static_audio_lengths=output.audio_lengths, + static_audio=audio, + static_audio_lengths=audio_lengths, ) def _select_batch_size(self, actual_batch_size: int) -> int | None: 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 d6bc525ac..ea45b9c02 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 @@ -14,6 +14,7 @@ from vllm.config import VllmConfig from vllm.logger import init_logger from vllm.model_executor.model_loader import DefaultModelLoader from vllm.model_executor.model_loader.weight_utils import default_weight_loader +from vllm.utils.torch_utils import set_default_torch_dtype from vllm_omni.model_executor.models.moss_tts.audio_tokenizer import ( MossAudioTokenizerConfig, @@ -45,6 +46,7 @@ class _MossCodecStreamSession: *, state_capacity: int, n_vq: int, + vllm_config: VllmConfig, graph_batch_sizes: list[int] | None = None, graph_frame_sizes: list[int] | None = None, ) -> None: @@ -59,6 +61,13 @@ class _MossCodecStreamSession: 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 == "cuda" else 0 + self._total_state_capacity = self._state_capacity + scratch_capacity + self._state_slot_ids = torch.arange( + self._total_state_capacity, + device=self._device, + dtype=torch.long, + ) + self._samples_per_frame = int(codec.downsample_rate) initialize_state_pool = getattr(codec, "initialize_decoder_state_pool", None) if not callable(initialize_state_pool): raise RuntimeError("The streaming codec does not implement a decoder state pool.") @@ -71,6 +80,7 @@ class _MossCodecStreamSession: batch_sizes=batch_sizes, frame_sizes=frame_sizes, num_quantizers=self._n_vq, + vllm_config=vllm_config, ) self._cudagraph_wrapper.warmup(self._device) self.reset_slots(list(range(self._state_capacity + scratch_capacity))) @@ -84,24 +94,32 @@ class _MossCodecStreamSession: self._leased_slots.add(slot) return slot - def release(self, slot: int) -> None: + def release(self, slot: int, *, state_already_reset: bool = False) -> None: if self._closed: return if slot not in self._leased_slots: return - self.reset_slots([slot]) + if not state_already_reset: + self.reset_slots([slot]) self._leased_slots.remove(slot) self._free_stream_slots.append(slot) + def _reset_slot_ids(self, slot_ids: torch.Tensor) -> None: + reset_streaming_slots = getattr(self._codec, "reset_decoder_state_slots", None) + if not callable(reset_streaming_slots): + raise RuntimeError("The streaming codec does not implement state-slot reset.") + with torch.no_grad(): + reset_streaming_slots(slot_ids) + def reset_slots(self, slots: list[int]) -> None: if not slots: return - reset_streaming_slots = getattr(self._codec, "reset_decoder_state_slots", None) - if callable(reset_streaming_slots): - with torch.no_grad(): - reset_streaming_slots(torch.tensor(slots, device=self._device, dtype=torch.long)) - return - raise RuntimeError("The streaming codec does not implement state-slot reset.") + if min(slots) < 0 or max(slots) >= self._total_state_capacity: + raise ValueError(f"Decoder state slots must be in [0, {self._total_state_capacity}), got {slots}") + start = slots[0] + if slots != list(range(start, start + len(slots))): + raise ValueError(f"Decoder state slots must be contiguous, got {slots}") + self._reset_slot_ids(self._state_slot_ids[start : start + len(slots)]) def close(self) -> None: if self._closed: @@ -150,9 +168,8 @@ class _MossCodecStreamSession: used_cudagraph = graph_output is not None if used_cudagraph: - audio_tensor, lengths_tensor, actual_batch_size = graph_output + audio_tensor, _, actual_batch_size = graph_output audio_tensor = audio_tensor[:actual_batch_size] - lengths_tensor = lengths_tensor[:actual_batch_size] else: codes_lengths = torch.full( (len(slots),), @@ -170,16 +187,17 @@ class _MossCodecStreamSession: if result.audio is None: return {} audio_tensor = result.audio - lengths_tensor = result.audio_lengths + + 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] + self._reset_slot_ids(terminal_slot_ids) audio = audio_tensor.detach().to("cpu", torch.float32) - lengths = lengths_tensor.detach().to("cpu") if lengths_tensor is not None else None + audio_length = step_t * self._samples_per_frame out: dict[int, torch.Tensor] = {} for row, slot in enumerate(slots): - wav = audio[row] - if lengths is not None: - wav = wav[..., : int(lengths[row].item())] - out[slot] = wav.contiguous() + out[slot] = audio[row, ..., :audio_length].contiguous() return out @@ -479,6 +497,7 @@ class MossTTSCodecDecoder(nn.Module): self._codec, state_capacity=self._stream_state_capacity, n_vq=self._n_vq, + vllm_config=self.vllm_config, graph_batch_sizes=self._streaming_graph_batch_sizes, graph_frame_sizes=self._streaming_graph_frame_sizes, ) @@ -517,7 +536,12 @@ class MossTTSCodecDecoder(nn.Module): if wav is not None: outputs[output_index] = wav if finished: - self._finish_stream_request(request_id, session, slot) + self._finish_stream_request( + request_id, + session, + slot, + state_already_reset=True, + ) continue grouped.setdefault(int(codes_nq_t.shape[1]), []).append( @@ -533,7 +557,12 @@ class MossTTSCodecDecoder(nn.Module): if wav is not None: outputs[output_index] = wav if finished: - self._finish_stream_request(request_id, session, slot) + self._finish_stream_request( + request_id, + session, + slot, + state_already_reset=True, + ) return outputs @@ -568,9 +597,11 @@ class MossTTSCodecDecoder(nn.Module): request_id: str, session: _MossCodecStreamSession, slot: int | None, + *, + state_already_reset: bool = False, ) -> None: if slot is not None: - session.release(slot) + session.release(slot, state_already_reset=state_already_reset) self._stream_req_slots.pop(request_id, None) def on_requests_finished(self, finished_req_ids: set[str] | list[str]) -> None: @@ -630,9 +661,17 @@ class MossTTSCodecDecoder(nn.Module): pass codec_path = self._codec_path - logger.info("Loading MOSS Audio Tokenizer from %s", codec_path) + device = self.vllm_config.device_config.device + logger.info("Loading MOSS Audio Tokenizer from %s directly on %s", codec_path, device) - codec_cfg, codec = self._build_codec(codec_path) + # This codec comes from a secondary checkpoint, so it cannot be built + # by the outer stage's normal initialize_model() call. Re-enter the + # same target-device/default-dtype contexts used by vLLM's + # BaseModelLoader.load_model() instead of constructing an 8 GiB FP32 + # model on CPU and copying the whole module to the GPU afterwards. + with set_default_torch_dtype(torch.float32): + with device: + codec_cfg, codec = self._build_codec(codec_path) model_loader = DefaultModelLoader(self.vllm_config.load_config) source = DefaultModelLoader.Source( @@ -710,9 +749,9 @@ class MossTTSCodecDecoder(nn.Module): skipped[:3] if skipped else "none", ) - device = self.vllm_config.device_config.device - codec.to(device=device, dtype=torch.float32) codec.eval() + if device.type != "cpu": + codec.decoder.to(dtype=torch.bfloat16) 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 diff --git a/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_local_depth.py b/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_local_depth.py index e0d871b46..7ee551f29 100644 --- a/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_local_depth.py +++ b/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_local_depth.py @@ -22,10 +22,14 @@ from __future__ import annotations import torch import torch.nn as nn import torch.nn.functional as F +from vllm.logger import init_logger from vllm_omni.model_executor.models.moss_tts.modeling_moss_tts_local import ( _sample_token, ) +from vllm_omni.platforms import current_omni_platform + +logger = init_logger(__name__) class _MossTTSLocalAttention(nn.Module): @@ -135,6 +139,7 @@ class MossTTSLocalDepthTransformer(nn.Module): rope_base = float(getattr(gpt2_config, "rope_base", 1_000_000.0)) self.h = nn.ModuleList([_MossTTSLocalBlock(self.hidden_size, n_head, inner_size, rope_base, eps)]) self.ln_f = nn.LayerNorm(self.hidden_size, eps=eps) + self._compiled_forward_prefix = None def _forward_prefix(self, seq_embeds: torch.Tensor) -> torch.Tensor: hidden_states = seq_embeds @@ -142,6 +147,24 @@ class MossTTSLocalDepthTransformer(nn.Module): hidden_states = block(hidden_states) return self.ln_f(hidden_states) + def setup_compile(self) -> None: + if self._compiled_forward_prefix is not None: + return + if not current_omni_platform.supports_torch_inductor(): + self._compiled_forward_prefix = self._forward_prefix + logger.warning_once("MOSS-TTS local depth torch.compile disabled on this platform") + return + self._compiled_forward_prefix = torch.compile( + self._forward_prefix, + dynamic=False, + options={"epilogue_fusion": False}, + ) + logger.info("MOSS-TTS local depth prefix enabled with torch.compile") + + def _run_prefix(self, seq_embeds: torch.Tensor) -> torch.Tensor: + forward_prefix = self._compiled_forward_prefix or self._forward_prefix + return forward_prefix(seq_embeds) + @torch.no_grad() def generate_frame( self, @@ -180,7 +203,7 @@ class MossTTSLocalDepthTransformer(nn.Module): embeds = backbone_last_hidden.new_zeros((batch_size, n_vq, self.hidden_size), dtype=dtype) embeds[:, 0, :] = backbone_last_hidden.to(dtype) - hidden = self._forward_prefix(embeds[:, :1, :]) + hidden = self._run_prefix(embeds[:, :1, :]) local_hidden = hidden[:, 0, :] binary_logits = local_text_lm_head(local_hidden).float() @@ -231,7 +254,7 @@ class MossTTSLocalDepthTransformer(nn.Module): if channel_index + 1 < n_vq: embeds[:, channel_index + 1, :] = audio_embeddings[channel_index](channel_token).to(dtype) - hidden = self._forward_prefix(embeds[:, : channel_index + 2, :]) + hidden = self._run_prefix(embeds[:, : channel_index + 2, :]) local_hidden = hidden[:, channel_index + 1, :] return should_continue, codes diff --git a/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_talker.py b/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_talker.py index 59313ecdf..171e57328 100644 --- a/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_talker.py +++ b/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_talker.py @@ -1513,6 +1513,59 @@ class MossTTSLocalTalkerForGeneration(nn.Module): mtp_control[:, :1] = 1.0 if active else 0.0 return input_ids, text_embed, {"mtp_inputs": (mtp_hidden, mtp_control)} + def preprocess_decode_batch( + self, + *, + input_ids: torch.Tensor, + req_infos: list[dict[str, Any]], + ) -> tuple[ + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + list[dict[str, Any]], + ]: + """Batch the tensor work from the decode branch of ``preprocess``.""" + input_ids = input_ids.reshape(-1) + batch_size = len(req_infos) + if input_ids.numel() != batch_size: + raise ValueError( + f"MOSS-TTS Local decode preprocess received {input_ids.numel()} token ids for {batch_size} requests" + ) + + text_embeds = self.model.embed_tokens(input_ids) + hidden_rows: list[torch.Tensor] = [] + active_rows: list[bool] = [] + for info in req_infos: + last_hidden = (info.get("hidden_states", {}) or {}).get("last") + if isinstance(last_hidden, torch.Tensor) and last_hidden.numel() > 0: + hidden_rows.append( + last_hidden.to( + device=text_embeds.device, + dtype=text_embeds.dtype, + ).reshape(1, -1) + ) + else: + hidden_rows.append(text_embeds.new_zeros((1, self.hidden_size))) + + state = info.get("audio_state", {}) or {} + active_rows.append(isinstance(state, dict) and not bool(state.get("is_stopping"))) + + last_talker_hidden = torch.cat(hidden_rows, dim=0) + text_step = torch.zeros_like(last_talker_hidden) + text_step[:, 0] = torch.tensor( + active_rows, + device=text_step.device, + dtype=text_step.dtype, + ) + return ( + input_ids, + text_embeds, + last_talker_hidden, + text_step, + [{} for _ in range(batch_size)], + ) + def postprocess(self, hidden_states: torch.Tensor, **_: Any) -> dict[str, Any]: if hidden_states.numel() == 0: return {} @@ -1681,6 +1734,9 @@ class MossTTSLocalTalkerForGeneration(nn.Module): [e.weight.detach() for e in self.audio_embeddings], dim=0 ) # (n_vq, audio_vocab_size, hidden_size) + if not self.vllm_config.model_config.enforce_eager: + self.local_transformer.setup_compile() + logger.info( "[MossTTSLocal] loaded %d/%d params; skipped=%d (first 5: %s)", len(loaded),