[Perf] Enable compile for MOSS-TTS-Local v1.5 talker and codec (#5530)

Signed-off-by: Canlin Guo <canlinguosdu@gmail.com>

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:

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

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,

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:

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

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

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),