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

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

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