Merge branch 'main' into dev/vllm-align

tests/model_executor/models/moss_tts/test_local_depth_rope_cache.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 """CPU parity tests for MOSS-TTS Local Depth RoPE caching."""
 
 import pytest
@@ -42,3 +42,16 @@ def test_rope_cache_matches_reference_and_reuses_storage() -> None:
     # 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
+
+
+def test_frame_local_kv_matches_full_prefix_and_reuses_slots() -> None:
+    model = MossTTSLocalDepthTransformer(_test_config()).eval()
+    model.h[0].attn.prepare_rope_cache(12, torch.device("cpu"), torch.float32)
+    cache = (torch.empty(2, 4, 12, 8), torch.empty(2, 4, 12, 8))
+    with torch.inference_mode():
+        for _ in range(2):
+            inputs = torch.randn(2, 12, 32)
+            for position in range(12):
+                expected = model._forward_prefix(inputs[:, : position + 1])[:, -1:]
+                actual = model._forward_prefix(inputs[:, position : position + 1], cache, position)
+                torch.testing.assert_close(actual, expected, atol=1e-6, rtol=1e-5)

tests/model_executor/models/moss_tts/test_streaming_attention.py

@@ -0,0 +1,58 @@
+# 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", [(2, 20, 1, 125), (4, 20, 15, 125), (1, 12, 480, 400), (4, 12, 32, 250)]
+)
+def test_masked_attention_matches_sdpa_with_strides_and_invalid_rows(batch, heads, queries, keys):
+    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()
+    mask[0, :, 0] = False  # Cover an all-masked row alongside valid rows.
+    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)
+    assert actual[0, :, 0].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,89 @@
+# 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 (
+    MossTTSCodecDecoder,
+    _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=8, n_vq=2, vllm_config=None)
+
+
+def test_terminal_padding_preserves_audio_live_state_and_reused_slots():
+    s = session()
+    assert [s.acquire(), s.acquire()] == [0, 1]
+    plan = {0: torch.ones(2, 3), 1: torch.ones(2, 15)}
+    with pytest.raises(ValueError, match="Only terminal"):
+        s.step(plan)
+    assert s._codec.state.count_nonzero() == 0
+    s.step({0: torch.ones(2, 1), 1: torch.full((2, 1), 2)})
+    actual = s.step(plan, terminal_slots={0})
+    torch.testing.assert_close(actual[0], torch.tensor([[4.0, 6.0, 8.0]]))
+    torch.testing.assert_close(actual[1], torch.arange(6.0, 36.0, 2.0)[None])
+    s.release(0, state_already_reset=True)
+    assert s.acquire() == 0
+    resumed = s.step({0: torch.ones(2, 1), 1: torch.ones(2, 1)})
+    assert [resumed[i].item() for i in (0, 1)] == [2, 36]
+
+
+@pytest.mark.parametrize("graph_capacity, expected_calls", [(None, 5), (4, 5), (8, 3)])
+def test_only_existing_compatible_graph_buckets_are_coalesced(mocker, graph_capacity, expected_calls):
+    s = session()
+    s._cudagraph_wrapper = (
+        None
+        if graph_capacity is None
+        else SimpleNamespace(
+            batch_sizes=[1, graph_capacity],
+            _select_frame_size=lambda size, allow_padding: next((b for b in [1, 8, 15] if b >= size), None),
+            decode=lambda *args, **kwargs: None,
+        )
+    )
+    step = mocker.spy(s, "step")
+    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: s
+    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 step.call_count == expected_calls
+    assert [result[i].shape[-1] for i in range(5)] == lengths
+    if graph_capacity == 8:
+        assert list(step.call_args_list[0].args[0]) == [0, 1, 2]
+        assert step.call_args_list[0].kwargs["terminal_slots"] == {0, 1}
+        assert list(step.call_args_list[1].args[0]) == [3]  # Preserve the exact one-frame graph.