[Perf] Enable async schedule for MOSS-TTS (#6241)

Signed-off-by: Sy03 <1370724210@qq.com>

tests/entrypoints/openai_api/test_tts_adapter.py

@@ -22,6 +22,10 @@ from vllm_omni.entrypoints.openai.tts_adapters.indextts2 import (
     IndexTTS25Adapter,
     indextts2_conditioning_cache_salt,
 )
+from vllm_omni.entrypoints.openai.tts_adapters.moss_tts import (
+    MossTTSAdapter,
+    MossTTSNanoAdapter,
+)
 from vllm_omni.entrypoints.openai.tts_adapters.qwen3_tts import Qwen3TTSAdapter
 from vllm_omni.model_executor.models.indextts2 import prompt_utils
 from vllm_omni.model_executor.models.indextts2.tokenizer_v2_5 import (
@@ -94,6 +98,20 @@ def test_all_adapters_are_ar_or_diffusion():
         assert cls.backend in ("ar", "diffusion")
 
 
+@pytest.mark.parametrize("adapter_cls", [MossTTSAdapter, MossTTSNanoAdapter])
+def test_moss_tts_applies_request_max_new_tokens(adapter_cls):
+    adapter = adapter_cls(SimpleNamespace(server=object()))
+    stage_defaults = [SimpleNamespace(max_tokens=4096)]
+
+    overridden = adapter.apply_sampling_overrides(
+        stage_defaults,
+        SimpleNamespace(max_new_tokens=512),
+    )
+
+    assert overridden[0].max_tokens == 512
+    assert stage_defaults[0].max_tokens == 4096
+
+
 def test_qwen3_tts_metadata():
     assert Qwen3TTSAdapter.backend == "ar"
     assert issubclass(Qwen3TTSAdapter, ARTTSAdapter)

tests/model_executor/models/moss_tts/test_compact_topk.py

@@ -0,0 +1,41 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""CPU regression tests for MOSS-TTS compact top-k sampling."""
+
+import pytest
+import torch
+
+from vllm_omni.model_executor.models.moss_tts import modeling_moss_tts_local
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.tts]
+
+
+def test_compact_topk_only_samples_retained_candidates() -> None:
+    logits = torch.randn(32, 1024, generator=torch.Generator().manual_seed(7))
+
+    sampled = modeling_moss_tts_local._sample_token(
+        logits,
+        temperature=1.7,
+        top_k=25,
+        top_p=0.8,
+        do_sample=True,
+        generator=torch.Generator().manual_seed(42),
+    )
+
+    retained = logits.topk(25, dim=-1).indices
+    assert torch.all((retained == sampled.unsqueeze(-1)).any(dim=-1))
+
+
+def test_compact_topk_tiny_nucleus_selects_argmax() -> None:
+    logits = torch.randn(32, 1024, generator=torch.Generator().manual_seed(17))
+
+    sampled = modeling_moss_tts_local._sample_token(
+        logits,
+        temperature=1.7,
+        top_k=25,
+        top_p=1e-9,
+        do_sample=True,
+        generator=torch.Generator().manual_seed(42),
+    )
+
+    torch.testing.assert_close(sampled, logits.argmax(dim=-1))

tests/model_executor/models/moss_tts/test_local_depth_rope_cache.py

@@ -0,0 +1,44 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""CPU parity tests for MOSS-TTS Local Depth RoPE caching."""
+
+import pytest
+import torch
+from transformers import GPT2Config
+
+from vllm_omni.model_executor.models.moss_tts.modeling_moss_tts_local_depth import (
+    MossTTSLocalDepthTransformer,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.tts]
+
+
+def _test_config() -> GPT2Config:
+    return GPT2Config(
+        n_embd=32,
+        n_head=4,
+        n_inner=64,
+        layer_norm_epsilon=1e-6,
+        rope_base=10000.0,
+    )
+
+
+def test_rope_cache_matches_reference_and_reuses_storage() -> None:
+    model = MossTTSLocalDepthTransformer(_test_config()).eval()
+    attn = model.h[0].attn
+    device = torch.device("cpu")
+
+    attn.prepare_rope_cache(12, device, torch.float32)
+    cache_ptr = attn._rope_cos_cache.data_ptr()
+    cos, sin = attn._rope_cos_sin(4, device, torch.float32)
+
+    positions = torch.arange(4, dtype=torch.float32)
+    freqs = torch.einsum("s,d->sd", positions, attn.inv_freq.float())
+    expected_cos = freqs.cos().repeat_interleave(2, dim=-1).view(1, 4, 1, attn.head_dim)
+    expected_sin = freqs.sin().repeat_interleave(2, dim=-1).view(1, 4, 1, attn.head_dim)
+    torch.testing.assert_close(cos, expected_cos)
+    torch.testing.assert_close(sin, expected_sin)
+
+    # Smaller slices retain the fixed allocation used by CUDA Graph.
+    attn.prepare_rope_cache(6, device, torch.float32)
+    assert attn._rope_cos_cache.data_ptr() == cache_ptr

tests/model_executor/models/moss_tts/test_talker_logits_processor.py

@@ -23,6 +23,11 @@ class _RecordingLogitsProcessor:
         return torch.zeros((1, 4))
 
 
+class _EmbeddingBackbone(nn.Module):
+    def embed_tokens(self, input_ids: torch.Tensor) -> torch.Tensor:
+        return input_ids.to(dtype=torch.float32).unsqueeze(-1).repeat(1, 4)
+
+
 def test_moss_tts_delay_compute_logits_does_not_forward_sampling_metadata() -> None:
     from vllm_omni.model_executor.models.moss_tts.modeling_moss_tts_talker import (
         MossTTSDelayTalkerForGeneration,
@@ -41,3 +46,74 @@ def test_moss_tts_delay_compute_logits_does_not_forward_sampling_metadata() -> N
     assert logits.shape == (1, 4)
     assert model.logits_processor.args == (model.text_lm_head, hidden_states)
     assert model.logits_processor.kwargs == {}
+
+
+def test_moss_tts_local_keeps_fixed_shape_rows_and_routes_stop_logits() -> None:
+    from vllm_omni.model_executor.models.moss_tts.modeling_moss_tts_talker import (
+        MossTTSLocalTalkerForGeneration,
+    )
+
+    model = MossTTSLocalTalkerForGeneration.__new__(MossTTSLocalTalkerForGeneration)
+    nn.Module.__init__(model)
+    model.audio_pad_token_id = 1024
+    model.audio_assistant_slot_token_id = 2
+    model.im_end_token_id = 3
+    model.text_vocab_size = 8
+    model.n_vq = 4
+    model._batch_state = None
+    model._batch_state_spans = None
+    model._batch_should_continue = None
+
+    valid = torch.tensor([[1, 2, 3, 4]], dtype=torch.long)
+    stopped = torch.full((1, 4), model.audio_pad_token_id, dtype=torch.long)
+    infos = [
+        {"audio_state": {}, "audio_codes": {"current": valid}},
+        {"audio_state": {}, "audio_codes": {"current": stopped}},
+    ]
+    spans = [(0, 3), (3, 4)]
+
+    output = model.make_omni_output(
+        torch.randn(4, 8),
+        model_intermediate_buffer=infos,
+        request_token_spans=spans,
+    )
+
+    audio = output.multimodal_outputs["codes"]["audio"]
+    assert len(audio) == 2
+    torch.testing.assert_close(audio[0], valid)
+    # Keeping this row avoids CUDA boolean-index shape discovery; the
+    # talker2codec raw async-chunk processor drops all-pad rows on CPU.
+    torch.testing.assert_close(audio[1], stopped)
+    assert model._batch_should_continue.tolist() == [True, False]
+
+    logits = model.compute_logits(torch.randn(4, 8))
+
+    assert logits is not None
+    assert logits.argmax(dim=-1).tolist() == [2, 2, 2, 3]
+
+
+def test_moss_tts_local_single_token_prefill_does_not_advance_audio() -> None:
+    from vllm_omni.model_executor.models.moss_tts.modeling_moss_tts_talker import (
+        MossTTSLocalTalkerForGeneration,
+    )
+
+    model = MossTTSLocalTalkerForGeneration.__new__(MossTTSLocalTalkerForGeneration)
+    nn.Module.__init__(model)
+    model.model = _EmbeddingBackbone()
+    model.n_vq = 4
+
+    input_ids = torch.tensor([7], dtype=torch.long)
+    _, embeds, update = model.preprocess(
+        input_ids,
+        None,
+        _omni_is_prefill=True,
+        audio_state={"is_stopping": True},
+        ref_offset=5,
+    )
+
+    torch.testing.assert_close(embeds, torch.full((1, 4), 7.0))
+    assert update == {
+        "audio_state": {"is_stopping": False},
+        "ref_offset": 6,
+    }
+    assert "mtp_inputs" not in update

vllm_omni/entrypoints/openai/tts_adapters/moss_tts.py

@@ -10,7 +10,12 @@ from typing import TYPE_CHECKING
 from vllm.inputs import tokens_input
 
 from vllm_omni.entrypoints.openai.tts_adapters import register_tts_adapter
-from vllm_omni.entrypoints.openai.tts_adapters.base import ARTTSAdapter, PreparedRequest, conditioning_cache_salt
+from vllm_omni.entrypoints.openai.tts_adapters.base import (
+    ARTTSAdapter,
+    PreparedRequest,
+    apply_max_new_tokens,
+    conditioning_cache_salt,
+)
 
 if TYPE_CHECKING:
     from vllm_omni.entrypoints.openai.protocol.audio import OpenAICreateSpeechRequest
@@ -104,6 +109,15 @@ class _MossTTSAdapterBase(ARTTSAdapter):
         prompt["cache_salt"] = conditioning_cache_salt(request, tts_params)
         return PreparedRequest(prompt=prompt, tts_params=tts_params, model_type=self.name)
 
+    def apply_sampling_overrides(
+        self,
+        sampling_params_list: list,
+        request: "OpenAICreateSpeechRequest",
+        prompt: dict | None = None,
+        request_id: str | None = None,
+    ) -> list:
+        return apply_max_new_tokens(sampling_params_list, request)
+
 
 @register_tts_adapter
 class MossTTSNanoAdapter(_MossTTSAdapterBase):

vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_local.py

@@ -128,32 +128,45 @@ def _sample_token(
     do_sample: bool,
     generator: torch.Generator | None = None,
 ) -> torch.Tensor:
-    """Top-k + top-p sampling (matches upstream's ``sample_token`` for the
-    inference branch).
+    """Top-k + top-p sampling for the upstream inference branch.
+
+    Nucleus filtering and multinomial sampling operate only on the retained
+    top-k candidates before the result is mapped back to the original
+    vocabulary. This preserves the categorical distribution when the top-k
+    boundary has no ties, but it is not seed/bit equivalent to multinomial over
+    a full-vocabulary tensor because random-number mapping depends on width.
     """
     if not do_sample or temperature <= 0:
         return logits.argmax(dim=-1)
 
     logits = logits / max(temperature, 1e-6)
+    compact_indices = None
     if top_k and top_k > 0 and top_k < logits.shape[-1]:
-        top_vals, _ = torch.topk(logits, top_k, dim=-1)
-        thresh = top_vals[..., -1:].expand_as(logits)
-        logits = torch.where(logits < thresh, torch.full_like(logits, float("-inf")), logits)
+        top_vals, top_indices = torch.topk(logits, top_k, dim=-1, sorted=True)
+        # topk is already descending: keep the following nucleus filter,
+        # softmax, and multinomial at width k instead of the full vocab.
+        logits = top_vals
+        compact_indices = top_indices
 
     if 0.0 < top_p < 1.0:
-        sorted_logits, sorted_idx = torch.sort(logits, descending=True, dim=-1)
-        probs = F.softmax(sorted_logits, dim=-1)
+        sorted_indices = None
+        if compact_indices is None:
+            logits, sorted_indices = torch.sort(logits, descending=True, dim=-1)
+        probs = F.softmax(logits, dim=-1)
         cum = probs.cumsum(dim=-1)
         # Drop tail beyond top_p (keep at least one token).
         drop = cum > top_p
         drop[..., 1:] = drop[..., :-1].clone()
         drop[..., 0] = False
-        sorted_logits = sorted_logits.masked_fill(drop, float("-inf"))
-        logits = torch.full_like(logits, float("-inf")).scatter_(-1, sorted_idx, sorted_logits)
+        logits = logits.masked_fill(drop, float("-inf"))
+        if sorted_indices is not None:
+            logits = torch.full_like(logits, float("-inf")).scatter_(-1, sorted_indices, logits)
 
     probs = F.softmax(logits, dim=-1)
     flat = probs.reshape(-1, probs.shape[-1])
     sampled = torch.multinomial(flat, num_samples=1, generator=generator).reshape(probs.shape[:-1])
+    if compact_indices is not None:
+        return compact_indices.gather(-1, sampled.unsqueeze(-1)).squeeze(-1)
     return sampled
 
 

vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_local_depth.py

@@ -52,15 +52,48 @@ class _MossTTSLocalAttention(nn.Module):
         self.c_proj = nn.Linear(hidden_size, hidden_size, bias=True)
         inv_freq = 1.0 / (rope_base ** (torch.arange(0, self.head_dim, 2, dtype=torch.float32) / self.head_dim))
         self.register_buffer("inv_freq", inv_freq, persistent=False)
+        self.register_buffer("_rope_cos_cache", torch.empty(0), persistent=False)
+        self.register_buffer("_rope_sin_cache", torch.empty(0), persistent=False)
 
-    def _rope_cos_sin(
-        self, seq_len: int, device: torch.device, dtype: torch.dtype
-    ) -> tuple[torch.Tensor, torch.Tensor]:
-        position_ids = torch.arange(seq_len, device=device, dtype=torch.float32)
-        freqs = torch.einsum("s,d->sd", position_ids, self.inv_freq.to(device=device))
+    def prepare_rope_cache(
+        self,
+        max_seq_len: int,
+        device: torch.device,
+        dtype: torch.dtype,
+    ) -> None:
+        """Materialize frame-local RoPE values once per device/dtype.
+
+        Local Depth positions always restart at zero and never exceed the
+        number of RVQ codebooks (12 for MOSS-TTS Local v1.5).  Building
+        arange/einsum/cos/sin inside every one of the 12 depth steps creates
+        many tiny kernels and also bloats the enclosing talker CUDA graph.
+        Keep one non-persistent cache and slice it for prefix execution.
+        """
+        if max_seq_len <= 0:
+            return
+        cache = self._rope_cos_cache
+        if cache.numel() > 0 and cache.device == device and cache.dtype == dtype and int(cache.shape[1]) >= max_seq_len:
+            return
+
+        position_ids = torch.arange(max_seq_len, device=device, dtype=torch.float32)
+        inv_freq = self.inv_freq.to(device=device, dtype=torch.float32)
+        freqs = torch.einsum("s,d->sd", position_ids, inv_freq)
         cos = freqs.cos().repeat_interleave(2, dim=-1).to(dtype)
         sin = freqs.sin().repeat_interleave(2, dim=-1).to(dtype)
-        return cos.view(1, seq_len, 1, self.head_dim), sin.view(1, seq_len, 1, self.head_dim)
+        self._rope_cos_cache = cos.view(1, max_seq_len, 1, self.head_dim)
+        self._rope_sin_cache = sin.view(1, max_seq_len, 1, self.head_dim)
+
+    def _rope_cos_sin(
+        self,
+        seq_len: int,
+        device: torch.device,
+        dtype: torch.dtype,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        self.prepare_rope_cache(seq_len, device, dtype)
+        return (
+            self._rope_cos_cache[:, :seq_len],
+            self._rope_sin_cache[:, :seq_len],
+        )
 
     @staticmethod
     def _rotate_half(x: torch.Tensor) -> torch.Tensor:
@@ -200,9 +233,15 @@ class MossTTSLocalDepthTransformer(nn.Module):
         batch_size = backbone_last_hidden.shape[0]
         dtype = self.ln_f.weight.dtype
 
+        # Populate the complete [0, n_vq) table before torch.compile traces
+        # _forward_prefix or the runner captures talker_mtp.  The compiled /
+        # captured body then contains only fixed-address cache slices rather
+        # than cache construction or buffer replacement.
+        for block in self.h:
+            block.attn.prepare_rope_cache(n_vq, backbone_last_hidden.device, dtype)
+
         embeds = backbone_last_hidden.new_zeros((batch_size, n_vq, self.hidden_size), dtype=dtype)
         embeds[:, 0, :] = backbone_last_hidden.to(dtype)
-
         hidden = self._run_prefix(embeds[:, :1, :])
         local_hidden = hidden[:, 0, :]
 

vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_talker.py

@@ -1323,7 +1323,11 @@ class MossTTSLocalTalkerForGeneration(nn.Module):
         self.mtp_hidden_size = hidden_size
         self.talker_mtp_graph_safe = True
         self.talker_mtp_output_key = ("audio_codes", "current")
-        self.use_async_omni_output = False
+        # ``make_omni_output`` keeps code rows fixed-shape and performs all
+        # state updates eagerly, so the runner can safely pack/snapshot them
+        # before CUDA-graph buffers are reused.  This moves the pinned D2H and
+        # output materialization off the Stage-0 execution thread.
+        self.use_async_omni_output = True
         self.eager_omni_postprocess_before_async_output = True
         self.omni_pooler_payload_include_hidden = False
         self.postprocess_uses_multimodal_outputs = False
@@ -1466,7 +1470,7 @@ class MossTTSLocalTalkerForGeneration(nn.Module):
     ) -> tuple[torch.Tensor, torch.Tensor, dict[str, Any]]:
         """Build per-step input embeddings (text + audio additive fusion).
 
-        Prefill: initialise the per-request ``{is_stopping, step}`` state.
+        Prefill: initialise the per-request ``{is_stopping}`` state.
         Decode: combine the forced text token (from ``compute_logits``) with
         the audio codes the local transformer produced last step.
         """
@@ -1474,8 +1478,9 @@ class MossTTSLocalTalkerForGeneration(nn.Module):
         span_len = int(input_ids.shape[0])
         audio_state = info_dict.get("audio_state")
         is_first_call = not isinstance(audio_state, dict)
+        is_prefill = bool(info_dict.get("_omni_is_prefill", False))
 
-        if span_len > 1 or is_first_call:
+        if span_len > 1 or is_first_call or is_prefill:
             embeds = self.model.embed_tokens(input_ids)
 
             ref_codes = (info_dict.get("codes", {}) or {}).get("ref")
@@ -1668,9 +1673,14 @@ class MossTTSLocalTalkerForGeneration(nn.Module):
                 should_continue = codes.ne(self.audio_pad_token_id).any(dim=-1)
                 should_values.append(should_continue.reshape(-1)[:1])
                 have_mtp_output = True
-                emitted_codes = codes[should_continue]
-                per_req_codes.append(emitted_codes)
-                have_codes = have_codes or emitted_codes.numel() > 0
+                # Do not boolean-index a CUDA tensor here.  Even for the
+                # common single-row case, ``codes[should_continue]`` must
+                # discover a dynamic output size on the host and therefore
+                # synchronizes the CUDA stream once per request.  Keep the
+                # fixed-shape row (including an all-pad stop row); the async
+                # chunk processor already filters all-pad rows on CPU.
+                per_req_codes.append(codes)
+                have_codes = True
             else:
                 should_values.append(torch.ones((1,), device=hidden.device, dtype=torch.bool))
                 per_req_codes.append(hidden.new_empty((0, self.n_vq), dtype=torch.long))