From 26e003715c83cabebdc05ee3466b8359133f1074 Mon Sep 17 00:00:00 2001 From: Sy03 <1370724210@qq.com> Date: Tue, 18 Aug 2026 23:36:19 -0700 Subject: [PATCH] [Perf] Enable async schedule for MOSS-TTS (#6241) Signed-off-by: Sy03 <1370724210@qq.com> --- .../openai_api/test_tts_adapter.py | 18 +++++ .../models/moss_tts/test_compact_topk.py | 41 ++++++++++ .../moss_tts/test_local_depth_rope_cache.py | 44 +++++++++++ .../moss_tts/test_talker_logits_processor.py | 76 +++++++++++++++++++ .../openai/tts_adapters/moss_tts.py | 16 +++- .../moss_tts/modeling_moss_tts_local.py | 31 +++++--- .../moss_tts/modeling_moss_tts_local_depth.py | 53 +++++++++++-- .../moss_tts/modeling_moss_tts_talker.py | 22 ++++-- 8 files changed, 278 insertions(+), 23 deletions(-) create mode 100644 tests/model_executor/models/moss_tts/test_compact_topk.py create mode 100644 tests/model_executor/models/moss_tts/test_local_depth_rope_cache.py diff --git a/tests/entrypoints/openai_api/test_tts_adapter.py b/tests/entrypoints/openai_api/test_tts_adapter.py index 9ab744c83..e87cab05f 100644 --- a/tests/entrypoints/openai_api/test_tts_adapter.py +++ b/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) diff --git a/tests/model_executor/models/moss_tts/test_compact_topk.py b/tests/model_executor/models/moss_tts/test_compact_topk.py new file mode 100644 index 000000000..e2b06c824 --- /dev/null +++ b/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)) diff --git a/tests/model_executor/models/moss_tts/test_local_depth_rope_cache.py b/tests/model_executor/models/moss_tts/test_local_depth_rope_cache.py new file mode 100644 index 000000000..90c386ee0 --- /dev/null +++ b/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 diff --git a/tests/model_executor/models/moss_tts/test_talker_logits_processor.py b/tests/model_executor/models/moss_tts/test_talker_logits_processor.py index 35c309979..b5866781c 100644 --- a/tests/model_executor/models/moss_tts/test_talker_logits_processor.py +++ b/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 diff --git a/vllm_omni/entrypoints/openai/tts_adapters/moss_tts.py b/vllm_omni/entrypoints/openai/tts_adapters/moss_tts.py index ad43a1c01..81bda2078 100644 --- a/vllm_omni/entrypoints/openai/tts_adapters/moss_tts.py +++ b/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): diff --git a/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_local.py b/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_local.py index 2fd0beefd..31a9fced1 100644 --- a/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_local.py +++ b/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 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 7ee551f29..e2d46ce2c 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 @@ -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, :] 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 171e57328..3f9ebf42f 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 @@ -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))