[Model] Optimize MOSS-TTS Local batched execution and streaming codec

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

tests/model_executor/models/moss_tts/test_streaming_attention.py

@@ -0,0 +1,61 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+import copy
+
+import pytest
+import torch
+import torch.nn.functional as F
+
+from vllm_omni.model_executor.models.moss_tts.audio_tokenizer_v2 import (
+    MossAudioTokenizerMultiheadAttention,
+    StreamingExecutionContext,
+)
+from vllm_omni.model_executor.models.moss_tts.streaming_attention import masked_attention
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cuda]
+
+
+@pytest.mark.parametrize(
+    "batch,heads,queries,keys", [(1, 20, 1, 125), (4, 20, 15, 125), (1, 12, 480, 400), (4, 12, 32, 250)]
+)
+@pytest.mark.parametrize("empty_mask", [False, True])
+def test_masked_attention_matches_sdpa_with_strides_and_invalid_rows(batch, heads, queries, keys, empty_mask):
+    torch.manual_seed(7)
+    q = torch.randn(batch, queries, 3, heads, 64, device="cuda", dtype=torch.bfloat16)[:, :, 0].permute(0, 2, 1, 3)
+    k = torch.randn(batch, heads, keys, 64, device="cuda", dtype=torch.bfloat16)
+    v = torch.randn_like(k)
+    positions = torch.arange(keys, device="cuda").roll(17)
+    query_positions = torch.arange(queries, device="cuda") + keys - queries
+    mask = (positions[None, :] <= query_positions[:, None])[None, None].expand(batch, 1, -1, -1).clone()
+    if empty_mask:
+        mask.zero_()
+    actual = masked_attention(q, k, v, mask)
+    expected = F.scaled_dot_product_attention(q, k, v, mask)
+    assert bool(actual.isfinite().all())
+    torch.testing.assert_close(actual, expected, atol=0.008, rtol=0.01)
+    if empty_mask:
+        assert actual.count_nonzero() == 0
+
+
+def test_streaming_attention_preserves_ring_wrap_reordering_and_slot_reset():
+    torch.manual_seed(9)
+    reference = MossAudioTokenizerMultiheadAttention(
+        128, 2, causal=True, context=9, device="cuda", dtype=torch.bfloat16
+    )
+    candidate = copy.deepcopy(reference)
+    candidate._streaming_attention = masked_attention
+    for model in (reference, candidate):
+        model._streaming_state = model._init_streaming_state(4)
+    for step in range(8):
+        slots = torch.tensor([1, 0, 3] if step % 2 else [0, 1, 3], device="cuda")
+        valid = torch.tensor([True, True, False], device="cuda")
+        context = StreamingExecutionContext(state_slot_ids=slots, valid_rows=valid)
+        x = torch.randn(3, 3, 128, device="cuda", dtype=torch.bfloat16)
+        with torch.inference_mode():
+            expected = reference(x, x, x, execution_context=context)
+            actual = candidate(x, x, x, execution_context=context)
+        torch.testing.assert_close(actual, expected, atol=0.01, rtol=0.02)
+        torch.testing.assert_close(candidate._streaming_state.offset, reference._streaming_state.offset, rtol=0, atol=0)
+        if step == 4:
+            for model in (reference, candidate):
+                model._streaming_state.reset_slots(torch.tensor([0], device="cuda"))

tests/model_executor/models/moss_tts/test_streaming_terminal_batch.py

@@ -0,0 +1,118 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+from types import SimpleNamespace
+
+import pytest
+import torch
+from torch import nn
+
+from vllm_omni.model_executor.models.moss_tts.modeling_moss_tts_codec import _MossCodecStreamSession
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+class CausalCodec(nn.Module):
+    downsample_rate = 1
+
+    def __init__(self):
+        super().__init__()
+        self.weight = nn.Parameter(torch.zeros(()))
+
+    def initialize_decoder_state_pool(self, capacity, scratch):
+        self.state = torch.zeros(capacity + scratch)
+
+    def reset_decoder_state_slots(self, slots):
+        self.state[slots] = 0
+
+    def decode_streaming_batch(self, codes, lengths, slots, valid_rows):
+        audio = codes.sum(0).float().cumsum(-1) + self.state[slots, None]
+        self.state[slots] = audio[:, -1]
+        return SimpleNamespace(audio=audio[:, None], audio_lengths=lengths)
+
+
+def session():
+    return _MossCodecStreamSession(CausalCodec(), state_capacity=3, n_vq=2, vllm_config=None)
+
+
+@pytest.mark.parametrize("tail_length", [1, 3, 7, 14])
+def test_terminal_padding_preserves_audio_live_state_and_reused_slots(tail_length):
+    batched, separate = session(), session()
+    for s in (batched, separate):
+        assert [s.acquire(), s.acquire()] == [0, 1]
+        s.step({0: torch.ones(2, 15), 1: torch.full((2, 15), 3)})
+    short = torch.arange(2 * tail_length).reshape(2, tail_length)
+    long = torch.full((2, 15), 5)
+    expected_short = separate.step({0: short}, terminal_slots={0})[0]
+    expected_long = separate.step({1: long})[1]
+    actual = batched.step({0: short, 1: long}, terminal_slots={0})
+    torch.testing.assert_close(actual[0], expected_short, rtol=0, atol=0)
+    torch.testing.assert_close(actual[1], expected_long, rtol=0, atol=0)
+    for s in (batched, separate):
+        s.release(0, state_already_reset=True)
+        assert s.acquire() == 0
+    next_step = {0: torch.ones(2, 1), 1: torch.ones(2, 1)}
+    a, b = batched.step(next_step), separate.step(next_step)
+    for slot in (0, 1):
+        torch.testing.assert_close(a[slot], b[slot], rtol=0, atol=0)
+
+
+def test_nonterminal_row_cannot_silently_advance_through_padding():
+    s = session()
+    s.acquire()
+    s.acquire()
+    with pytest.raises(ValueError, match="Only terminal"):
+        s.step({0: torch.ones(2, 3), 1: torch.ones(2, 15)})
+    assert s._codec.state.count_nonzero() == 0
+
+
+@pytest.mark.parametrize("graph_capacity, expected_calls", [(None, 5), (4, 5), (8, 3)])
+def test_only_existing_compatible_graph_buckets_are_coalesced(graph_capacity, expected_calls):
+    from vllm_omni.model_executor.models.moss_tts.modeling_moss_tts_codec import MossTTSCodecDecoder
+
+    calls = []
+    leased: list[int] = []
+
+    def acquire():
+        slot = len(leased)
+        leased.append(slot)
+        return slot
+
+    def step(plan, *, terminal_slots):
+        calls.append((dict(plan), set(terminal_slots)))
+        return {slot: torch.zeros(1, codes.shape[1]) for slot, codes in plan.items()}
+
+    graph = (
+        None
+        if graph_capacity is None
+        else SimpleNamespace(
+            batch_sizes=[1, graph_capacity],
+            frame_sizes=[1, 8, 15],
+            _select_frame_size=lambda size, allow_padding: next((b for b in [1, 8, 15] if b >= size), None),
+        )
+    )
+    fake_session = SimpleNamespace(
+        _cudagraph_wrapper=graph,
+        acquire=acquire,
+        step=step,
+        release=lambda slot, **kwargs: None,
+    )
+    decoder = object.__new__(MossTTSCodecDecoder)
+    nn.Module.__init__(decoder)
+    decoder._stream_max_step_frames = 15
+    decoder._stream_state_capacity = 8
+    decoder._stream_req_slots = {}
+    decoder._ensure_stream_session = lambda: fake_session
+    lengths = [3, 7, 8, 1, 15]
+    finished = [True, True, False, True, False]
+    result = decoder._decode_streaming_batch(
+        [
+            (i, str(i), torch.ones(2, length), done)
+            for i, (length, done) in enumerate(zip(lengths, finished, strict=True))
+        ]
+    )
+    assert len(calls) == expected_calls
+    assert [result[i].shape[-1] for i in range(5)] == lengths
+    if graph_capacity == 8:
+        assert set(calls[0][0]) == {0, 1, 2}
+        assert calls[0][1] == {0, 1}
+        assert set(calls[1][0]) == {3}  # Preserve the exact one-frame graph.