Merge main into dev/vllm-align for vLLM 0.31 release

Signed-off-by: tzhouam <tzhouam@connect.ust.hk>

tests/model_executor/models/common/test_fused_sampling.py

@@ -0,0 +1,90 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+import pytest
+import torch
+
+pytestmark = [
+    pytest.mark.core_model,
+    pytest.mark.cuda,
+    pytest.mark.skipif(not torch.cuda.is_available() or torch.version.hip is not None, reason="requires CUDA"),
+]
+
+
+def _reference(logits, uniforms, top_k, top_p):
+    if top_k > 0:
+        threshold = logits.topk(top_k, dim=-1).values[:, -1:]
+        logits = logits.masked_fill(logits < threshold, -torch.inf)
+    if top_p < 1:
+        values, indices = logits.sort(dim=-1, descending=True)
+        probs = values.softmax(dim=-1, dtype=torch.float32)
+        values[(probs.cumsum(-1) - probs) >= top_p] = -torch.inf
+        logits = values.scatter(1, indices, values)
+    return (logits.float() - torch.log(-torch.log(uniforms))).argmax(-1, keepdim=True)
+
+
+@pytest.mark.parametrize("batch,vocab", [(1, 128), (7, 2048), (64, 2048)])
+@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16])
+@pytest.mark.parametrize("top_k,top_p", [(-1, 0.8), (0, 1.0), (50, 0.8), (50, 0.01)])
+def test_fixed_uniform_sampling_and_graph(batch, vocab, dtype, top_k, top_p):
+    from vllm_omni.model_executor.models.common.fused_sampling import sample_top_k_top_p_gumbel
+
+    torch.manual_seed(173)
+    logits = torch.randn(batch, vocab, device="cuda", dtype=dtype)
+    # The MTP caller supplies a noncontiguous row slice of [B, 15, vocab].
+    uniforms = torch.rand(batch, 15, vocab, device="cuda")[:, 3].clamp_min_(1e-7)
+    expected = _reference(logits, uniforms, top_k, top_p)
+    actual = sample_top_k_top_p_gumbel(logits, uniforms, top_k=top_k, top_p=top_p)
+    torch.testing.assert_close(actual, expected, rtol=0, atol=0)
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph):
+        captured = sample_top_k_top_p_gumbel(logits, uniforms, top_k=top_k, top_p=top_p)
+    logits.add_(torch.randn_like(logits) * 0.1)
+    uniforms.uniform_(1e-7, 1 - 1e-7)
+    graph.replay()
+    torch.testing.assert_close(captured, _reference(logits, uniforms, top_k, top_p), rtol=0, atol=0)
+
+
+@pytest.mark.parametrize("top_k,top_p", [(3, 0.5), (3, 1.0), (0, 0.8)])
+def test_threshold_ties_and_degenerate_rows(top_k, top_p):
+    from vllm_omni.model_executor.models.common.fused_sampling import sample_top_k_top_p_gumbel
+
+    logits = torch.tensor(
+        [[2, 2, 2, 2, 2, 1, 0, -1], [1] * 8, [-torch.inf] * 8, [torch.nan, 1, 2, 3, 4, 5, 6, 7]],
+        device="cuda",
+        dtype=torch.float32,
+    )
+    uniforms = torch.tensor([[0.4, 0.7, 0.8, 0.9, 0.2, 0.3, 0.6, 0.5]], device="cuda").expand(4, -1)
+    actual = sample_top_k_top_p_gumbel(logits, uniforms, top_k=top_k, top_p=top_p)
+    torch.testing.assert_close(actual, _reference(logits, uniforms, top_k, top_p), rtol=0, atol=0)
+
+
+@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16])
+def test_native_nucleus_boundaries(dtype):
+    from vllm_omni.model_executor.models.common.fused_sampling import sample_top_k_top_p_gumbel
+
+    torch.manual_seed(1923)
+    logits = torch.randn(1, 2048, device="cuda", dtype=dtype)
+    values, indices = logits.sort(descending=True)
+    probabilities = values.softmax(-1, dtype=torch.float32)
+    before = probabilities.cumsum(-1) - probabilities
+    for rank in [1, 19, 49, 1000, 2000]:
+        boundary = before[0, rank]
+        uniforms = torch.full_like(logits, 0.001, dtype=torch.float32)
+        uniforms[0, indices[0, rank]] = 1 - 1e-7
+        for direction in [-torch.inf, torch.inf]:
+            top_p = torch.nextafter(boundary, boundary.new_tensor(direction)).item()
+            actual = sample_top_k_top_p_gumbel(logits, uniforms, top_k=0, top_p=top_p)
+            torch.testing.assert_close(actual, _reference(logits, uniforms, 0, top_p), rtol=0, atol=0)
+
+
+@pytest.mark.parametrize("top_p", [0.8, 1.0])
+def test_strided_logits_and_non_power_of_two_vocabulary(top_p):
+    from vllm_omni.model_executor.models.common.fused_sampling import sample_top_k_top_p_gumbel
+
+    for seed in range(16):
+        torch.manual_seed(seed)
+        logits = torch.randn(128, 2050, device="cuda")[:, ::2]
+        uniforms = torch.rand(128, 3, 1025, device="cuda")[:, 1].clamp_min_(1e-7)
+        actual = sample_top_k_top_p_gumbel(logits, uniforms, top_k=-1, top_p=top_p)
+        torch.testing.assert_close(actual, _reference(logits, uniforms, -1, top_p), rtol=0, atol=0)

tests/model_executor/models/common/test_short_kv_attention.py

@@ -16,7 +16,7 @@ pytestmark = [
 @pytest.mark.parametrize("heads,kv_heads,dim", [(16, 8, 128), (8, 8, 64)])
 @pytest.mark.parametrize("positions", [[0, 1], [2], [8], [15]])
 def test_short_attention_causal_gqa_and_graph(batch, heads, kv_heads, dim, positions):
-    from vllm_omni.model_executor.models.qwen3_tts.short_kv_attention import short_kv_attention
+    from vllm_omni.model_executor.models.common.short_kv_attention import short_kv_attention
 
     with torch.inference_mode():
         torch.manual_seed(42)

tests/model_executor/models/cosyvoice3/test_cosyvoice3_hift_graph.py

@@ -0,0 +1,41 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""HiFT graphs consume live inputs and return independently owned waveforms."""
+
+import pytest
+import torch
+
+from tests.helpers.mark import hardware_test
+from tests.model_executor.models.cosyvoice3.test_cosyvoice3_incremental_hift import _make_hift
+
+pytestmark = [pytest.mark.core_model]
+
+
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+@pytest.mark.parametrize("finalize", [False, True])
+@torch.inference_mode()
+def test_hift_graph_preserves_current_and_previous_request(finalize):
+    hift = _make_hift().cuda().eval()
+    hift.remove_weight_norm()
+    hift.f0_predictor.remove_weight_norm()
+    hift.enable_decode_graphs()
+    generator = torch.Generator(device="cuda").manual_seed(17)
+    saved: list[tuple[torch.Tensor, torch.Tensor]] = []
+    for _ in range(5):
+        mel = torch.randn(1, 80, 30, device="cuda", generator=generator)
+        source = torch.randn(1, 1, 30 * 480, device="cuda", generator=generator)
+        expected = hift._decode_eager(mel, source, finalize=finalize)
+        actual = hift.decode(mel, source, finalize=finalize)
+        torch.testing.assert_close(actual, expected, rtol=1e-5, atol=1e-6)
+        for previous, snapshot in saved:
+            assert torch.equal(previous, snapshot)
+        if hift._decode_graphs.graphs:
+            saved.append((actual, actual.clone()))
+    assert len(hift._decode_graphs.graphs) == 1
+    hift._decode_graphs.max_graphs = 1
+    mel = torch.randn(1, 80, 31, device="cuda", generator=generator)
+    source = torch.randn(1, 1, 31 * 480, device="cuda", generator=generator)
+    actual = hift.decode(mel, source, finalize=finalize)
+    expected = hift._decode_eager(mel, source, finalize=finalize)
+    assert torch.equal(actual, expected)
+    assert len(hift._decode_graphs.graphs) == 1

tests/model_executor/models/higgs_audio_v3/test_mrv2_state.py

@@ -13,6 +13,7 @@ from vllm_omni.model_executor.models.higgs_audio_v3.higgs_audio_v3_talker import
     HiggsAudioV3TalkerForConditionalGeneration,
 )
 from vllm_omni.model_executor.models.higgs_audio_v3.model_state import HiggsModelState
+from vllm_omni.worker_v2.model_states.eager_mtp import EagerMTPState
 from vllm_omni.worker_v2.model_states.intermediate_buffer import OmniIntermediateBuffer
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
@@ -34,6 +35,8 @@ def make_state():
         _audio_continuation_id=99,
         update_decode_step_metadata=Mock(),
     )
+    state._stream_pos = {}
+    state._eager_state = EagerMTPState(state)
     return state
 
 

tests/model_executor/models/moss_tts/test_codec_async_output.py

@@ -0,0 +1,338 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Waveform ownership across graph replays and asynchronous output copies."""
+
+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
+from vllm_omni.model_executor.output_snapshot import PackedOutputSnapshot
+
+pytestmark = [pytest.mark.core_model]
+
+
+class CausalCodec(nn.Module):
+    downsample_rate = 1
+
+    def __init__(self, channels):
+        super().__init__()
+        self.weight = nn.Parameter(torch.zeros(()))
+        self.channels = channels
+        self.config = SimpleNamespace(codebook_size=1024)
+
+    def initialize_decoder_state_pool(self, capacity, scratch):
+        self.state = torch.zeros(capacity + scratch, device=self.weight.device)
+
+    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]
+        channels = torch.arange(1, self.channels + 1, device=audio.device)
+        return SimpleNamespace(audio=audio[:, None] * channels[None, :, None], audio_lengths=lengths)
+
+
+class ReplayGraph:
+    """A real CUDA graph with deliberately reused output storage."""
+
+    batch_sizes = [1]
+
+    def __init__(self, codec):
+        self.codes = torch.ones((2, 1, 3), device="cuda", dtype=torch.long)
+        self.slots = torch.zeros(1, device="cuda", dtype=torch.long)
+        self.lengths = torch.full((1,), 3, device="cuda", dtype=torch.long)
+        self.valid = torch.ones(1, device="cuda", dtype=torch.bool)
+        stream = torch.cuda.Stream()
+        stream.wait_stream(torch.cuda.current_stream())
+        with torch.cuda.stream(stream):
+            for _ in range(3):
+                codec.decode_streaming_batch(self.codes, self.lengths, self.slots, self.valid)
+        torch.cuda.current_stream().wait_stream(stream)
+        self.graph = torch.cuda.CUDAGraph()
+        with torch.cuda.graph(self.graph):
+            self.output = codec.decode_streaming_batch(self.codes, self.lengths, self.slots, self.valid)
+        codec.state.zero_()
+
+    def decode(self, codes, slots, *, allow_frame_padding):
+        if codes.shape[1:] != (1, 3):
+            return None
+        self.codes.copy_(codes)
+        self.slots.copy_(slots)
+        self.graph.replay()
+        return self.output.audio, self.output.audio_lengths, 1
+
+
+def make_decoder(*, v2=True, channels=2, device="cpu", model_type="moss_tts_local", streaming=True):
+    config = SimpleNamespace(
+        model_config=SimpleNamespace(
+            hf_config=SimpleNamespace(n_vq=2, model_type=model_type),
+            async_chunk=streaming,
+            use_v2_model_runner=v2,
+        ),
+        scheduler_config=SimpleNamespace(max_num_seqs=8),
+    )
+    decoder = MossTTSCodecDecoder(vllm_config=config)
+    decoder._codec = CausalCodec(channels).to(device)
+    decoder._n_channels = channels
+    decoder._stream_max_step_frames = 3
+    return decoder
+
+
+@pytest.mark.cpu
+@pytest.mark.parametrize("ramp,frames", [(None, [1, 15]), ([1, 4, 8, 15], [1, 4, 8, 15]), ([1, 4, 4, 15], [1, 4, 15])])
+def test_ramp_chunks_are_captured_at_exact_frame_counts(ramp, frames):
+    config = SimpleNamespace(
+        model_config=SimpleNamespace(
+            hf_config=SimpleNamespace(n_vq=2, model_type="moss_tts_local"),
+            async_chunk=True,
+            use_v2_model_runner=True,
+            stage_connector_config={
+                "extra": {
+                    "initial_codec_chunk_frames": 1,
+                    "codec_chunk_frames": 15,
+                    "codec_chunk_ramp": ramp,
+                }
+            },
+        ),
+        scheduler_config=SimpleNamespace(max_num_seqs=8),
+    )
+    decoder = MossTTSCodecDecoder(vllm_config=config)
+    assert decoder._streaming_graph_frame_sizes == frames
+    assert decoder._stream_max_step_frames == 15
+
+
+@pytest.mark.cpu
+@pytest.mark.parametrize(
+    "v2,streaming,model_type",
+    [
+        (False, True, "moss_tts_local"),
+        (True, False, "moss_tts_local"),
+        (True, True, "moss_tts_realtime"),
+        (True, True, "moss_tts_local"),
+    ],
+)
+def test_gpu_output_requires_local_streaming_mrv2_and_cuda(v2, streaming, model_type):
+    decoder = make_decoder(v2=v2, streaming=streaming, model_type=model_type)
+    assert decoder._gpu_stream_output == (v2 and streaming and model_type == "moss_tts_local")
+    session = decoder._ensure_stream_session()
+    slot = session.acquire()
+    assert session.step({slot: torch.ones(2, 3)})[slot].device.type == "cpu"
+
+
+@pytest.mark.cuda
+@pytest.mark.parametrize("v2", [False, True])
+@pytest.mark.parametrize("channels", [1, 2])
+@pytest.mark.parametrize("graph", [False, True])
+@torch.no_grad()
+def test_replays_terminal_tails_and_slot_reuse_preserve_pending_waveforms(v2, channels, graph):
+    from vllm_omni.outputs import OmniModelRunnerOutput
+    from vllm_omni.worker_v2.omni_generation_model_runner import (
+        OmniGenerationAsyncOutput,
+        OmniGenerationModelRunner,
+        _contains_cuda_tensor,
+    )
+
+    decoder = make_decoder(v2=v2, channels=channels, device="cuda")
+    # Real Local-v1.5 weights replace the class's 24 kHz default at load time.
+    decoder._sr_tensor = torch.tensor(48_000, dtype=torch.int32)
+    session = decoder._ensure_stream_session()
+    if graph:
+        session._cudagraph_wrapper = ReplayGraph(decoder._codec)
+    copy_stream = torch.cuda.Stream()
+    pending = []
+    # Do not materialize any output until later forwards have replayed the
+    # graph, reset terminal/cancelled state, and reused those state slots.
+    for round_index in range(4):
+        value = round_index + 1
+        codes = torch.cat(
+            [
+                torch.full((14,), value, device="cuda", dtype=torch.long),
+                torch.full((6,), value + 10, device="cuda", dtype=torch.long),
+            ]
+        )
+        info = [{"request_id": rid, "meta": {"finished": rid == "tail"}} for rid in ["tail", "empty", "live"]]
+        output = decoder.forward(codes, runtime_additional_information=info, seq_token_counts=[14, 0, 6])
+        # The runner may overwrite its input buffer before codec completion.
+        codes.fill_(-100)
+        payload = output.multimodal_outputs
+        assert _contains_cuda_tensor(payload) == v2
+        assert isinstance(payload, PackedOutputSnapshot) == v2
+        if v2:
+            # Exactly one waveform slab, regardless of request/chunk count.
+            assert sum(slab.is_cuda for _, slab in payload._groups) == 1
+            result = OmniModelRunnerOutput(
+                req_ids=["tail", "empty", "live"],
+                req_id_to_index={"tail": 0, "empty": 1, "live": 2},
+                sampled_token_ids=[[], [], []],
+            )
+            pending.append(
+                OmniGenerationAsyncOutput(
+                    model_runner_output=result,
+                    multimodal_outputs=payload,
+                    num_reqs=3,
+                    main_stream=torch.cuda.current_stream(),
+                    copy_stream=copy_stream,
+                )
+            )
+        else:
+            pending.append(OmniGenerationModelRunner._build_pooler_output(output, 3))
+        decoder.on_requests_finished(["live"])
+        assert decoder._stream_req_slots == {}
+        assert not session._leased_slots
+
+    for round_index, item in enumerate(pending):
+        rows = item.get_output().multimodal_outputs if v2 else item
+        for index, (frames, value) in enumerate([(7, round_index + 1), (0, 0), (3, round_index + 11)]):
+            expected = 2 * value * torch.arange(1, frames + 1, dtype=torch.float32)
+            if channels > 1:
+                expected = expected[None] * torch.arange(1, channels + 1)[:, None]
+            torch.testing.assert_close(rows[index]["model_outputs"], expected, rtol=0, atol=0)
+            assert rows[index]["sr"].item() == 48_000
+
+
+@pytest.mark.cuda
+@torch.no_grad()
+def test_owned_input_copy_clamps_entire_ragged_batch():
+    from vllm_omni.worker_v2.omni_ar_model_runner import _async_copy_mm
+
+    decoder = make_decoder(device="cuda", channels=2)
+    codes = torch.tensor([-1, 1024, 1, -10, 2, 1025, 5000, -2], device="cuda")
+    output = decoder.forward(
+        codes,
+        runtime_additional_information=[{"request_id": name} for name in ("a", "empty", "b")],
+        seq_token_counts=[6, 0, 2],
+    )
+    codes.fill_(-100)
+    cpu = _async_copy_mm(output.multimodal_outputs, 0, pin_memory=True)
+    torch.accelerator.synchronize()
+    expected = torch.tensor([[0.0, 1025.0, 2049.0], [0.0, 2050.0, 4098.0]])
+    torch.testing.assert_close(cpu["model_outputs"][0], expected, rtol=0, atol=0)
+    assert cpu["model_outputs"][1].numel() == 0
+    torch.testing.assert_close(cpu["model_outputs"][2], torch.tensor([[1023.0], [2046.0]]), rtol=0, atol=0)
+    decoder.on_requests_finished(["a", "b"])
+
+
+@pytest.mark.cuda
+@torch.no_grad()
+def test_pinned_metadata_reuse_waits_for_dma_and_preserves_mixed_terminals():
+    decoder = make_decoder(device="cuda")
+    session = decoder._ensure_stream_session()
+    # Force reuse every call, with an outstanding transfer on the first call.
+    # This delay tests host-buffer ownership, never performance or overlap.
+    session._metadata_ring = session._metadata_ring[:1]
+    slots = [session.acquire() for _ in range(3)]
+    pending = []
+    for step in range(8):
+        order = slots[step % 3 :] + slots[: step % 3]
+        terminal = {order[1]}
+        torch.cuda._sleep(2_000_000)
+        ids, tails = session._stage_slot_ids(order, terminal)
+        pending.append((ids.clone(), tails.clone(), order, list(terminal)))
+    session.close()
+    for ids, tails, expected_ids, expected_tails in pending:
+        assert ids.cpu().tolist() == expected_ids
+        assert tails.cpu().tolist() == expected_tails
+
+
+@pytest.mark.cuda
+@torch.no_grad()
+def test_mixed_terminal_reset_does_not_reset_live_slot():
+    decoder = make_decoder(device="cuda", channels=1)
+    session = decoder._ensure_stream_session()
+    first, second = session.acquire(), session.acquire()
+    codes = torch.ones(2, 3, device="cuda")
+    initial = session.step({first: codes, second: 2 * codes}, terminal_slots={first})
+    following = session.step({second: codes, first: codes})
+    torch.testing.assert_close(initial[first].cpu(), torch.tensor([[2.0, 4.0, 6.0]]), rtol=0, atol=0)
+    torch.testing.assert_close(following[first].cpu(), torch.tensor([[2.0, 4.0, 6.0]]), rtol=0, atol=0)
+    torch.testing.assert_close(following[second].cpu(), torch.tensor([[14.0, 16.0, 18.0]]), rtol=0, atol=0)
+    session.close()
+
+
+@pytest.mark.cuda
+@pytest.mark.parametrize("explicit_copy_stream", [False, True])
+@torch.no_grad()
+def test_output_copy_waits_for_private_codec_stream(explicit_copy_stream):
+    from vllm_omni.worker_v2.omni_ar_model_runner import _async_copy_mm
+
+    decoder = make_decoder(device="cuda", channels=1)
+    decoder._codec_stream = torch.cuda.Stream()
+    with torch.cuda.stream(decoder._codec_stream):
+        torch.cuda._sleep(20_000_000)
+    codes = torch.ones(6, device="cuda", dtype=torch.long)
+    output = decoder.forward(
+        codes,
+        runtime_additional_information=[{"request_id": "live"}],
+        seq_token_counts=[6],
+    )
+    codes.fill_(-100)
+    payload = output.multimodal_outputs
+    assert isinstance(payload, PackedOutputSnapshot)
+    assert payload.producer_event is not None
+    # No explicit consumer wait here: the shared copier must honor readiness.
+    copy_stream = torch.cuda.Stream()
+    with torch.cuda.stream(copy_stream):
+        cpu = _async_copy_mm(payload, 0, copy_stream=copy_stream if explicit_copy_stream else None, pin_memory=True)
+    decoder.on_requests_finished(["live"])
+    copy_stream.synchronize()
+    torch.testing.assert_close(cpu["model_outputs"][0], torch.tensor([2.0, 4.0, 6.0]), rtol=0, atol=0)
+
+
+@pytest.mark.cuda
+@pytest.mark.parametrize("channels", [1, 2])
+@torch.no_grad()
+def test_short_terminal_row_crops_each_channel_into_output_slab(channels):
+    decoder = make_decoder(device="cuda", channels=channels)
+    session = decoder._ensure_stream_session()
+    short, live = session.acquire(), session.acquire()
+    codes = {short: torch.ones(2, 2, device="cuda"), live: 3 * torch.ones(2, 5, device="cuda")}
+    slab = torch.empty(channels * 7, device="cuda")
+    buffers = {
+        short: slab[: channels * 2].view(channels, 2),
+        live: slab[channels * 2 :].view(channels, 5),
+    }
+    result = session.step(codes, terminal_slots={short}, output_buffers=buffers)
+    for slot, frames, value in [(short, 2, 1), (live, 5, 3)]:
+        expected = 2 * value * torch.arange(1, frames + 1, dtype=torch.float32)[None]
+        expected = expected * torch.arange(1, channels + 1)[:, None]
+        assert result[slot].data_ptr() == buffers[slot].data_ptr()
+        torch.testing.assert_close(result[slot].cpu(), expected, rtol=0, atol=0)
+    session.close()
+
+
+@pytest.mark.cuda
+@pytest.mark.parametrize("channels", [1, 2])
+@pytest.mark.parametrize("layout", ["adjacent", "strided", "gapped"])
+@torch.no_grad()
+def test_batched_output_copy_preserves_rows_and_slab_ownership(channels, layout):
+    decoder = make_decoder(device="cuda", channels=channels)
+    session = decoder._ensure_stream_session()
+    slots = [session.acquire() for _ in range(4)]
+    slab = torch.full((6 if layout == "gapped" else 4, channels, 6), -1.0, device="cuda")
+    if layout == "gapped":
+        buffers = {slot: slab[row] for row, slot in zip([0, 1, 3, 4], slots, strict=True)}
+        frames = 6
+    elif layout == "adjacent":
+        buffers = {slot: slab.reshape(4, channels * 6)[row].view(channels, 6) for row, slot in enumerate(slots)}
+        frames = 6
+    else:
+        # Interleaved with other requests' output: the short views must not
+        # overwrite the untouched part of another row/channel.
+        buffers = {slot: slab[row, :, :3] for row, slot in enumerate(slots)}
+        frames = 3
+    codes = {slot: torch.full((2, frames), row + 1, device="cuda") for row, slot in enumerate(slots)}
+    result = session.step(codes, output_buffers=buffers)
+    for row, slot in enumerate(slots):
+        expected = 2 * (row + 1) * torch.arange(1, frames + 1, dtype=torch.float32)[None]
+        expected = expected * torch.arange(1, channels + 1)[:, None]
+        assert result[slot].data_ptr() == buffers[slot].data_ptr()
+        torch.testing.assert_close(result[slot].cpu(), expected, rtol=0, atol=0)
+    if layout == "strided":
+        assert torch.all(slab[:, :, 3:] == -1)
+    elif layout == "gapped":
+        assert torch.all(slab[[2, 5]] == -1)
+    session.close()

tests/model_executor/models/moss_tts/test_codec_ring_headroom.py

@@ -0,0 +1,146 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Chunked streaming decode must match whole-sequence decode.
+
+``RingKVCache.complete`` writes a whole chunk before attending. With a ring of
+exactly ``context`` entries, a chunk of T tokens evicts up to T keys that its
+own queries should still see; for T > context it evicts part of the chunk
+itself. ``initialize_decoder_state_pool(chunk_frames=...)`` adds one chunk of
+headroom per layer so streaming reproduces the non-streaming decoder.
+"""
+
+import pytest
+import torch
+
+from vllm_omni.model_executor.models.moss_tts.audio_tokenizer_v2 import MossAudioTokenizerModel
+from vllm_omni.model_executor.models.moss_tts.configuration_moss_audio_tokenizer_v2 import (
+    MossAudioTokenizerConfig,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+TRANSFORMER = {
+    "module_type": "Transformer",
+    "d_model": 16,
+    "num_heads": 2,
+    "num_layers": 1,
+    "dim_feedforward": 32,
+    "causal": True,
+    "norm": "layer_norm",
+    "positional_embedding": "rope",
+    "max_period": 10000,
+    "gating": "none",
+    "layer_scale": 1.0,
+    "conv_layout": True,
+}
+
+
+def tiny_codec() -> MossAudioTokenizerModel:
+    # 64 Hz mono input, 4-sample code frames (16 code frames per second).
+    # Decoder: transformer @16 Hz (context 0.5 s = 8 tokens) -> x2 -> transformer
+    # @32 Hz (context 0.25 s = 8 tokens) -> x2 -> 64 Hz. A 6-frame chunk is 12
+    # tokens at the second transformer, more than its 8-token ring.
+    config = MossAudioTokenizerConfig(
+        sampling_rate=64,
+        downsample_rate=4,
+        number_channels=1,
+        enable_channel_interleave=False,
+        encoder_kwargs=[
+            {"module_type": "PatchedPretransform", "patch_size": 4},
+            {**TRANSFORMER, "input_dimension": 4, "output_dimension": 8, "context_duration": 0.5},
+        ],
+        decoder_kwargs=[
+            {**TRANSFORMER, "input_dimension": 8, "output_dimension": 16, "context_duration": 0.5},
+            {"module_type": "PatchedPretransform", "patch_size": 2},
+            {**TRANSFORMER, "input_dimension": 8, "output_dimension": 2, "context_duration": 0.25},
+            {"module_type": "PatchedPretransform", "patch_size": 2},
+        ],
+        quantizer_kwargs={
+            "input_dim": 8,
+            "rvq_dim": 8,
+            "output_dim": 8,
+            "num_quantizers": 2,
+            "codebook_size": 16,
+            "codebook_dim": 4,
+            "quantizer_type": "rlfq",
+        },
+    )
+    torch.manual_seed(0)
+    model = MossAudioTokenizerModel(config).eval()
+    with torch.no_grad():
+        for parameter in model.parameters():
+            parameter.copy_(torch.randn_like(parameter) * 0.5)
+    model.quantizer.build_decode_lut(2, dtype=torch.float32)
+    return model
+
+
+def stream_decode(model, codes, chunk_frames, *, headroom):
+    n_q, _, total = codes.shape
+    model.initialize_decoder_state_pool(1, 0, chunk_frames=chunk_frames if headroom else 0)
+    try:
+        slot = torch.zeros(1, dtype=torch.long)
+        valid = torch.ones(1, dtype=torch.bool)
+        parts = []
+        for start in range(0, total, chunk_frames):
+            chunk = codes[:, :, start : start + chunk_frames]
+            lengths = torch.full((1,), chunk.shape[-1], dtype=torch.long)
+            parts.append(model.decode_streaming_batch(chunk, lengths, slot, valid).audio)
+        return torch.cat(parts, dim=-1)
+    finally:
+        model.close_decoder_state_pool()
+
+
+@torch.no_grad()
+def test_ring_headroom_sizes_each_layer_for_one_chunk():
+    model = tiny_codec()
+    transformers = [m for m in model.decoder if hasattr(m, "tokens_per_input_frame")]
+    assert [m.tokens_per_input_frame for m in transformers] == [1, 2]
+    assert [m.transformer.layers[0].self_attn.context for m in transformers] == [8, 8]
+    model.initialize_decoder_state_pool(1, 0, chunk_frames=6)
+    try:
+        rings = [m.transformer.layers[0].self_attn._streaming_state.kv_cache.capacity for m in transformers]
+        assert rings == [8 + 6, 8 + 12]
+    finally:
+        model.close_decoder_state_pool()
+    model.initialize_decoder_state_pool(1, 0)
+    try:
+        rings = [m.transformer.layers[0].self_attn._streaming_state.kv_cache.capacity for m in transformers]
+        assert rings == [8, 8]
+    finally:
+        model.close_decoder_state_pool()
+
+
+@torch.no_grad()
+def test_chunked_streaming_matches_whole_sequence_only_with_headroom():
+    model = tiny_codec()
+    torch.manual_seed(1)
+    codes = torch.randint(0, 16, (2, 1, 18))
+    reference = model._decode_frame(codes).audio
+    with_headroom = stream_decode(model, codes, 6, headroom=True)
+    legacy = stream_decode(model, codes, 6, headroom=False)
+    assert with_headroom.shape == reference.shape == legacy.shape
+    torch.testing.assert_close(with_headroom, reference, rtol=1e-5, atol=1e-5)
+    # The legacy context-sized ring lets the 12-token chunk evict keys inside
+    # its own window, so streaming diverges from whole-sequence decoding.
+    assert not torch.allclose(legacy, reference, rtol=1e-3, atol=1e-3)
+
+
+@torch.no_grad()
+def test_padded_terminal_tail_matches_exact_tail_with_headroom():
+    model = tiny_codec()
+    torch.manual_seed(2)
+    codes = torch.randint(0, 16, (2, 1, 14))
+    history, tail = codes[:, :, :12], codes[:, :, 12:]
+    padded_tail = torch.nn.functional.pad(tail, (0, 4))  # 2 real frames padded to a 6-frame step
+    outputs = {}
+    for name, last in [("exact", tail), ("padded", padded_tail)]:
+        model.initialize_decoder_state_pool(1, 0, chunk_frames=6)
+        try:
+            slot, valid = torch.zeros(1, dtype=torch.long), torch.ones(1, dtype=torch.bool)
+            for chunk in (history[:, :, :6], history[:, :, 6:]):
+                model.decode_streaming_batch(chunk, torch.tensor([6]), slot, valid)
+            audio = model.decode_streaming_batch(last, torch.tensor([last.shape[-1]]), slot, valid).audio
+            outputs[name] = audio[..., : 2 * model.downsample_rate]
+        finally:
+            model.close_decoder_state_pool()
+    torch.testing.assert_close(outputs["padded"], outputs["exact"], rtol=1e-5, atol=1e-5)

tests/model_executor/models/moss_tts/test_first_chunk_fast_path.py

@@ -0,0 +1,547 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""The first-chunk fast path emits audio out of band; the main path resumes the stream."""
+
+import queue
+import threading
+from types import SimpleNamespace
+
+import pytest
+import torch
+from torch import nn
+
+from tests.model_executor.models.moss_tts.test_streaming_terminal_batch import session as cpu_session
+from vllm_omni.model_executor.models.moss_tts.first_chunk_fast_path import MossFirstChunkFastPath, _SlotHandoff
+from vllm_omni.model_executor.models.moss_tts.modeling_moss_tts_codec import MossTTSCodecDecoder
+from vllm_omni.worker_v2.first_audio_sender import engine_output_queue_sink
+
+pytestmark = pytest.mark.core_model
+
+N_VQ, SAMPLES, CHANNELS = 2, 4, 2
+
+
+class _Session:
+    def __init__(self, capacity):
+        self.free = list(range(capacity))
+        self.reset_waits = []
+
+    def acquire(self, *, oldest=False):
+        return self.free.pop(0 if oldest else -1) if self.free else None
+
+    def order_after_reset(self, slot, stream):
+        self.reset_waits.append(slot)
+
+    def release(self, slot, **kwargs):
+        self.free.append(slot)
+
+
+class _Wrapper:
+    batch_sizes = [1, 2, 4, 8]
+
+    def __init__(self, device):
+        self.state = torch.zeros(16, device=device)
+        self.calls = []
+        self.gate = threading.Event()
+        self.gate.set()
+
+    def decode(self, codes, slots):
+        self.gate.wait()
+        self.calls.append(int(codes.shape[1]))
+        base = codes.sum(0)[:, 0].float() + 100 * slots.float()
+        channel = torch.arange(CHANNELS, device=codes.device).float()
+        audio = base[:, None, None] + channel[None, :, None] + torch.zeros(1, 1, SAMPLES, device=codes.device)
+        self.state[slots] = 1
+        return audio.to(torch.bfloat16), None, int(codes.shape[1])
+
+
+@pytest.fixture
+def cuda():
+    if not torch.cuda.is_available():
+        pytest.skip("CUDA required")
+    return torch.device("cuda")
+
+
+def _fast(device, capacity=4):
+    session, wrapper = _Session(capacity), _Wrapper(device)
+    fast = MossFirstChunkFastPath(
+        session,
+        wrapper,
+        n_vq=N_VQ,
+        frames=1,
+        codebook_size=1024,
+        samples_per_frame=SAMPLES,
+        n_channels=CHANNELS,
+        sample_rate=torch.tensor(24000, dtype=torch.int32),
+        device=device,
+    )
+    return fast, session, wrapper
+
+
+class _CallbackSink:
+    def __init__(self, callback):
+        self.callback = callback
+
+    def prepare(self, request_ids):
+        callback = self.callback
+
+        class Delivery:
+            routes = {rid: (3 if rid == "r0" else 0) for rid in request_ids}
+
+            def __call__(self, ids, rows, sr):
+                for rid, row in zip(ids, rows, strict=True):
+                    callback(self.routes[rid], rid, {"model_outputs": row, "sr": sr})
+
+            def fail(self, ids):
+                pass
+
+        return Delivery()
+
+
+@pytest.mark.cuda
+def test_emits_each_first_chunk_and_hands_slot_to_main_path(cuda):
+    fast, session, wrapper = _fast(cuda)
+    emitted = {}
+    done = threading.Event()
+
+    def sink(client, req_id, payload):
+        emitted[req_id] = (client, payload["model_outputs"], int(payload["sr"]))
+        if len(emitted) == 2:
+            done.set()
+
+    slots: dict[str, int] = {}
+    assert not fast.submit("r0", "k0", [1, 2], slots)  # unbound: main path owns it
+    fast.bind(_CallbackSink(sink))
+    wrapper.gate.clear()
+    assert fast.submit("r0", "k0", torch.tensor([1, 2]), slots)
+    assert fast.submit("r1", "k1", [5, 7], slots)
+    assert not fast.submit("r0", "k0", [1, 2], slots)  # already has a slot
+    assert not fast.submit("r2", "k2", [1, 2, 3, 4], slots)  # not a 1-frame chunk
+    assert slots == {"k0": 0, "k1": 1} and session.free == [2, 3]
+    wrapper.gate.set()
+    assert done.wait(10)
+    fast.close()
+
+    assert wrapper.calls in ([1, 1], [2])
+    for req_id, client, code_sum, slot in (("r0", 3, 3, 0), ("r1", 0, 12, 1)):
+        got_client, wav, sr = emitted[req_id]
+        assert got_client == client and sr == 24000 and wav.device.type == "cpu"
+        expected = (code_sum + 100 * slot + torch.arange(CHANNELS).float())[:, None].expand(CHANNELS, SAMPLES)
+        torch.testing.assert_close(wav, expected)
+    assert sorted(session.reset_waits) == [0, 1]
+
+    # The main path consumes the marker once and is ordered after the decode.
+    assert fast.take_decoded("k0") and not fast.take_decoded("k0")
+    fast.order_after(0)
+    fast.order_after(0)  # no pending handoff: no-op
+    assert wrapper.state[:2].tolist() == [1, 1]
+
+
+@pytest.mark.cuda
+def test_order_after_blocks_until_decode_finishes(cuda):
+    fast, _, wrapper = _fast(cuda)
+    fast.bind(_CallbackSink(lambda *args: None))
+    wrapper.gate.clear()
+    assert fast.submit("r0", "k0", [1, 2], {})
+    finished = threading.Event()
+
+    def wait_for_slot():
+        fast.order_after(0)
+        finished.set()
+
+    waiter = threading.Thread(target=wait_for_slot)
+    waiter.start()
+    assert not finished.wait(0.2)
+    wrapper.gate.set()
+    assert finished.wait(10)
+    waiter.join()
+    fast.close()
+
+
+@pytest.mark.cuda
+def test_no_free_slot_leaves_chunk_to_main_path(cuda):
+    fast, _, _ = _fast(cuda, capacity=0)
+    fast.bind(_CallbackSink(lambda *args: None))
+    slots: dict[str, int] = {}
+    assert not fast.submit("r0", "k0", [1, 2], slots)
+    assert slots == {} and not fast.take_decoded("k0")
+    fast.close()
+
+
+class _FakeFast:
+    frames = 1
+
+    def __init__(self, decoded, unsent=None):
+        self.decoded = set(decoded)
+        self.unsent = dict(unsent or {})
+        self.ordered = []
+        self.forgotten = []
+
+    def take_decoded(self, key):
+        if key in self.decoded:
+            self.decoded.remove(key)
+            return True
+        return False
+
+    def get_request_slot(self, key, slots):
+        return slots.get(key)
+
+    def order_after(self, slot):
+        self.ordered.append(slot)
+
+    def take_unsent_audio(self, key):
+        return self.unsent.pop(key, None)
+
+    def forget(self, key):
+        self.forgotten.append(key)
+
+    def gate(self):
+        self.gated = getattr(self, "gated", 0) + 1
+
+
+def _decoder(s, fast):
+    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
+    decoder._first_chunk_fast_path = fast
+    decoder._stream_session = s
+    decoder._codec_stream = None
+    return decoder
+
+
+@pytest.mark.cpu
+def test_main_path_skips_fast_decoded_chunk_and_resumes_state():
+    s = cpu_session()
+    reference = cpu_session()
+    fast = _FakeFast({"a"})
+    decoder = _decoder(s, fast)
+    # The fast path leased slot 0 and decoded frame 0 of "a" into it.
+    slot = s.acquire(oldest=True)
+    decoder._stream_req_slots["a"] = slot
+    s.step({slot: torch.full((2, 1), 3.0)})
+    ref_slot = reference.acquire()
+    reference.step({ref_slot: torch.full((2, 1), 3.0)})
+
+    out = decoder._decode_streaming_batch([(0, "a", torch.full((2, 1), 3.0), False), (1, "b", torch.ones(2, 1), False)])
+    assert 0 not in out and out[1].tolist() == [[2.0]]
+    assert fast.ordered == [slot]
+
+    out = decoder._decode_streaming_batch([(0, "a", torch.ones(2, 3), True)])
+    expected = reference.step({ref_slot: torch.ones(2, 3)})[ref_slot]
+    torch.testing.assert_close(out[0], expected)
+    assert "a" not in decoder._stream_req_slots and fast.ordered == [slot, slot]
+
+
+@pytest.mark.cpu
+def test_abort_orders_release_after_fast_decode():
+    s = cpu_session()
+    fast = _FakeFast({"a"})
+    decoder = _decoder(s, fast)
+    slot = decoder._stream_req_slots["a"] = s.acquire(oldest=True)
+    decoder.on_requests_finished({"a"})
+    assert fast.ordered == [slot] and fast.forgotten == ["a"] and decoder._stream_req_slots == {}
+    assert slot in s._free_stream_slots
+
+
+@pytest.mark.cpu
+def test_main_path_returns_unsent_first_chunk_without_decoding_twice():
+    s = cpu_session()
+    first = torch.full((1, 1), 7.0)
+    fast = _FakeFast({"a"}, {"a": first})
+    decoder = _decoder(s, fast)
+    slot = s.acquire(oldest=True)
+    decoder._stream_req_slots["a"] = slot
+    s.step({slot: torch.full((2, 1), 3.0)})
+
+    out = decoder._decode_streaming_batch([(0, "a", torch.full((2, 1), 3.0), False)])
+    torch.testing.assert_close(out[0], first)
+    assert fast.unsent == {} and fast.ordered == [slot]
+
+    # A payload containing a later frame must concatenate it after the
+    # preserved first samples, while advancing state only for the later frame.
+    first = torch.full((1, 1), 7.0)
+    fast.decoded.add("b")
+    fast.unsent["b"] = first
+    slot_b = s.acquire(oldest=True)
+    decoder._stream_req_slots["b"] = slot_b
+    s.step({slot_b: torch.full((2, 1), 3.0)})
+    out = decoder._decode_streaming_batch([(0, "b", torch.full((2, 2), 3.0), False)])
+    torch.testing.assert_close(out[0][..., :1], first)
+    assert out[0].shape[-1] == 2
+
+
+@pytest.mark.cpu
+def test_oldest_acquire_reuses_least_recently_released_slot():
+    s = cpu_session()
+    a, b = s.acquire(), s.acquire()
+    s.release(a)
+    s.release(b)
+    assert s.acquire() == b
+    assert s.acquire(oldest=True) == 7  # never-used slots were freed before any release
+    assert s._free_stream_slots[-1] == a
+
+
+@pytest.mark.cuda
+def test_failed_decode_reports_error_without_replaying_state(cuda):
+    fast, _, wrapper = _fast(cuda)
+
+    def fail(codes, slots):
+        raise RuntimeError("boom")
+
+    wrapper.decode = fail
+    fast.bind(_CallbackSink(lambda *args: pytest.fail("no output expected")))
+    slots: dict[str, int] = {}
+    assert fast.submit("r0", "k0", [1, 2], slots)
+    with pytest.raises(RuntimeError, match="first-chunk decode failed"):
+        fast.order_after(slots["k0"])
+    assert fast.take_decoded("k0")
+    fast.close()
+
+
+@pytest.mark.cuda
+def test_failed_output_handoff_retains_decoded_samples(cuda):
+    fast, _, _ = _fast(cuda)
+
+    def reject(*args):
+        raise RuntimeError("output queue closed")
+
+    fast.bind(_CallbackSink(reject))
+    slots: dict[str, int] = {}
+    assert fast.submit("r0", "k0", [1, 2], slots)
+    fast.order_after(slots["k0"])
+    assert fast.take_decoded("k0")
+    wav = fast.take_unsent_audio("k0")
+    assert wav is not None and wav.shape == (CHANNELS, SAMPLES)
+    assert fast.take_unsent_audio("k0") is None
+    fast.close()
+
+
+@pytest.mark.cuda
+def test_gate_orders_main_stream_after_inflight_fast_decode(cuda):
+    fast, _, _ = _fast(cuda)
+    side = torch.cuda.Stream()
+    with torch.cuda.stream(side):
+        torch.cuda._sleep(200_000_000)
+        pending = torch.cuda.Event()
+        pending.record(side)
+    fast._inflight = pending
+    fast.gate()  # disabled: no wait
+    unordered = torch.cuda.Event()
+    unordered.record()
+    fast._gate_main = True
+    fast.gate()
+    ordered = torch.cuda.Event()
+    ordered.record()
+    unordered.synchronize()
+    assert not pending.query() and not ordered.query()
+    ordered.synchronize()
+    assert pending.query()
+
+
+def _cpu_fast(monkeypatch, *, scheduler=None):
+    # Replace only CUDA allocations, retaining the real lifecycle and worker.
+    import vllm_omni.model_executor.models.moss_tts.first_chunk_fast_path as module
+
+    monkeypatch.setattr(torch.Tensor, "pin_memory", lambda tensor: tensor)
+    monkeypatch.setattr(torch.cuda, "Stream", lambda **kwargs: None)
+    monkeypatch.setattr(torch.cuda, "Event", lambda: None)
+    monkeypatch.setattr(module, "current_omni_platform", SimpleNamespace(set_device=lambda device: None))
+    session = _Session(2)
+    fast = MossFirstChunkFastPath(
+        session,
+        SimpleNamespace(batch_sizes=[1, 2]),
+        n_vq=2,
+        frames=1,
+        codebook_size=1024,
+        samples_per_frame=4,
+        n_channels=1,
+        sample_rate=torch.tensor(24000),
+        device=torch.device("cpu:0"),
+        handoff_timeout_s=0.01,
+    )
+    outputs: queue.Queue = queue.Queue()
+    scheduler = scheduler or SimpleNamespace(requests={"r": SimpleNamespace(client_index=3)})
+    sink = engine_output_queue_sink(outputs, scheduler, upstream_first_audio=False)
+    return fast, session, sink, outputs
+
+
+def _cancel_during_admission(fast, session, scheduler, sink):
+    entered, release = threading.Event(), threading.Event()
+
+    class PausedSink:
+        def prepare(self, ids):
+            delivery = sink.prepare(ids)
+            if delivery.routes:
+                entered.set()
+                assert release.wait(5)
+            return delivery
+
+    fast.bind(PausedSink())
+    decoder = _decoder(session, fast)
+    errors = []
+    finished = threading.Event()
+    finishing = threading.Event()
+
+    def submit():
+        try:
+            assert fast.submit("r", "r", [1, 2], decoder._stream_req_slots)
+        except BaseException as error:
+            errors.append(error)
+
+    def cancel():
+        try:
+            finishing.set()
+            decoder.on_requests_finished({"r"})
+            finished.set()
+        except BaseException as error:
+            errors.append(error)
+
+    producer = threading.Thread(target=submit)
+    cleaner = threading.Thread(target=cancel)
+    producer.start()
+    try:
+        assert entered.wait(5)
+        assert decoder._stream_req_slots == {}
+        # Scheduler cancellation retires the route before the runner hook.
+        scheduler.requests.pop("r")
+        cleaner.start()
+        assert finishing.wait(5)
+        assert not finished.wait(0.1)
+        release.set()
+        producer.join(5)
+        cleaner.join(5)
+        assert not producer.is_alive() and not cleaner.is_alive()
+        assert not errors, errors
+        assert finished.is_set()
+        assert decoder._stream_req_slots == {}
+        assert fast._decoded == {} and fast._handoffs == {}
+        assert sorted(session.free) == [0, 1]
+        assert not fast.submit("r", "r", [1, 2], decoder._stream_req_slots)
+
+        # A fresh request can use and return the same capacity immediately.
+        scheduler.requests["reuse"] = SimpleNamespace(client_index=3)
+        release.set()
+        assert fast.submit("reuse", "reuse", [1, 2], decoder._stream_req_slots)
+        scheduler.requests.pop("reuse")
+        decoder.on_requests_finished({"reuse"})
+        assert decoder._stream_req_slots == {} and sorted(session.free) == [0, 1]
+    finally:
+        release.set()
+        producer.join(5)
+        if cleaner.ident is not None:
+            cleaner.join(5)
+        fast.close()
+
+
+@pytest.mark.cpu
+def test_cancel_waits_for_first_chunk_admission_and_reclaims_slot(monkeypatch):
+    scheduler = SimpleNamespace(requests={"r": SimpleNamespace(client_index=3)})
+    fast, session, sink, _ = _cpu_fast(monkeypatch, scheduler=scheduler)
+
+    # Keep the actual queue/worker/handoff lifecycle, omitting CUDA decode.
+    def decode(jobs):
+        for job in jobs:
+            job.handoff.done.set()
+
+    monkeypatch.setattr(fast, "_decode", decode)
+    fast._handoff_timeout_s = 5.0
+    _cancel_during_admission(fast, session, scheduler, sink)
+
+
+@pytest.mark.cuda
+def test_cancel_during_first_chunk_admission_orders_cuda_handoff_and_reuse(cuda):
+    fast, session, _ = _fast(cuda, capacity=2)
+    scheduler = SimpleNamespace(requests={"r": SimpleNamespace(client_index=3)})
+    sink = engine_output_queue_sink(queue.Queue(), scheduler, upstream_first_audio=False)
+    _cancel_during_admission(fast, session, scheduler, sink)
+
+
+@pytest.mark.cpu
+def test_slot_timeout_keeps_handoff_and_prevents_slot_reuse(monkeypatch):
+    fast, _, _, _ = _cpu_fast(monkeypatch)
+    handoff = _SlotHandoff()
+    fast._handoffs[0] = handoff
+    for _ in range(2):
+        with pytest.raises(TimeoutError, match="slot 0"):
+            fast.order_after(0)
+        assert fast._handoffs[0] is handoff
+    handoff.done.set()
+    fast.order_after(0)
+    assert fast._handoffs == {}
+
+
+@pytest.mark.cpu
+def test_closed_decoder_rejects_jobs_and_cannot_restart(monkeypatch):
+    fast, session, sink, _ = _cpu_fast(monkeypatch)
+    fast.bind(sink)
+    fast.close()
+    slots: dict[str, int] = {}
+    assert not fast.submit("r", "external", [1, 2], slots)
+    assert slots == {} and session.free == [0, 1]
+    with pytest.raises(RuntimeError, match="closed"):
+        fast.bind(sink)
+
+
+@pytest.mark.cpu
+def test_missing_route_does_not_claim_slot(monkeypatch):
+    fast, session, sink, outputs = _cpu_fast(monkeypatch, scheduler=SimpleNamespace(requests={}))
+    fast.bind(sink)
+    try:
+        assert not fast.submit("gone", "external", [1, 2], {})
+        assert session.free == [0, 1] and outputs.empty()
+    finally:
+        fast.close()
+
+
+@pytest.mark.cpu
+def test_decode_failure_wakes_current_and_queued_handoffs(monkeypatch):
+    from vllm.v1.engine import FinishReason
+
+    fast, _, sink, outputs = _cpu_fast(monkeypatch)
+    entered, release = threading.Event(), threading.Event()
+
+    def fail(jobs):
+        entered.set()
+        assert release.wait(2)
+        raise RuntimeError("partial replay failed")
+
+    fast._decode = fail
+    fast.bind(sink)
+    slots: dict[str, int] = {}
+    try:
+        assert fast.submit("r", "a", [1, 2], slots)
+        assert entered.wait(2)
+        assert fast.submit("r", "b", [3, 4], slots)
+        release.set()
+        fast._thread.join(2)
+        assert not fast._thread.is_alive()
+        for slot in slots.values():
+            with pytest.raises(RuntimeError, match="decode failed"):
+                fast.order_after(slot)
+        assert all(fast.take_decoded(key) for key in slots)
+        assert not fast.submit("r", "c", [1, 2], slots)
+        assert outputs.qsize() == 2
+        while not outputs.empty():
+            assert outputs.get_nowait()[1].outputs[0].finish_reason == FinishReason.ERROR
+    finally:
+        release.set()
+        fast.close()
+
+
+@pytest.mark.cpu
+def test_platform_failure_wakes_pending_handoff(monkeypatch):
+    import vllm_omni.model_executor.models.moss_tts.first_chunk_fast_path as module
+
+    fast, _, _, _ = _cpu_fast(monkeypatch)
+    fast._handoffs[0] = _SlotHandoff()
+
+    def fail(device):
+        raise RuntimeError("device unavailable")
+
+    monkeypatch.setattr(module, "current_omni_platform", SimpleNamespace(set_device=fail))
+    fast._run()
+    assert fast._closed
+    with pytest.raises(RuntimeError, match="decode failed"):
+        fast.order_after(0)

tests/model_executor/models/moss_tts/test_local_code_rows_snapshot.py

@@ -0,0 +1,32 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Per-request code rows keep the runner payload layout with one host copy."""
+
+import pytest
+import torch
+
+from vllm_omni.model_executor.models.moss_tts.local_model_state import _CodeRowsSnapshot
+from vllm_omni.model_executor.output_snapshot import PackedOutputSnapshot
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def test_rows_map_to_batch_positions_and_copy_once():
+    codes = torch.arange(6, dtype=torch.long).reshape(2, 3)
+    snapshot = _CodeRowsSnapshot(codes, [2, 0], 4)
+    assert isinstance(snapshot, PackedOutputSnapshot) and snapshot.producer_event is None
+    rows = snapshot["codes"]["audio"]
+    assert [tuple(row.shape) for row in rows] == [(1, 3), (0, 3), (1, 3), (0, 3)]
+    assert torch.equal(rows[2], codes[:1]) and torch.equal(rows[0], codes[1:])
+
+    copies = []
+
+    def copy(tensor):
+        copies.append(tensor)
+        return tensor.clone()
+
+    host = snapshot.copy_to_cpu(copy)["codes"]["audio"]
+    assert len(copies) == 1 and copies[0] is codes
+    codes.zero_()
+    assert host[2].tolist() == [[0, 1, 2]] and host[0].tolist() == [[3, 4, 5]]
+    assert host[1].numel() == 0 and host[1].dtype == torch.long

tests/model_executor/models/moss_tts/test_local_compiled_sampler.py

@@ -0,0 +1,37 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+import pytest
+import torch
+
+pytestmark = [
+    pytest.mark.core_model,
+    pytest.mark.cuda,
+    pytest.mark.tts,
+    pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required"),
+]
+
+
+def test_explicit_generator_bypasses_compiled_audio_sampler():
+    from transformers import GPT2Config
+
+    from vllm_omni.model_executor.models.moss_tts.modeling_moss_tts_local_depth import (
+        MossTTSLocalDepthTransformer,
+    )
+
+    model = MossTTSLocalDepthTransformer(GPT2Config(n_embd=80, n_head=1, n_inner=160)).cuda().bfloat16().eval()
+    heads = torch.nn.ModuleList([torch.nn.Linear(80, 32) for _ in range(2)]).cuda().bfloat16()
+    embeddings = torch.nn.ModuleList([torch.nn.Embedding(32, 80) for _ in range(2)]).cuda().bfloat16()
+    binary = torch.nn.Linear(80, 2).cuda().bfloat16()
+    hidden = torch.randn(1, 80, device="cuda", dtype=torch.bfloat16)
+    generator = torch.Generator(device="cuda").manual_seed(17)
+    expected = model.generate_frame(hidden, heads, embeddings, binary, n_vq=2, generator=generator)
+
+    def forbidden(*args, **kwargs):
+        raise AssertionError("compiled path selected")
+
+    model._compiled_audio_sampler = forbidden
+    generator.manual_seed(17)
+    actual = model.generate_frame(hidden, heads, embeddings, binary, n_vq=2, generator=generator)
+    assert all(torch.equal(a, b) for a, b in zip(expected, actual))
+    with pytest.raises(AssertionError, match="compiled path selected"):
+        model.generate_frame(hidden, heads, embeddings, binary, n_vq=2)

tests/model_executor/models/moss_tts/test_local_determined_tokens.py

@@ -0,0 +1,124 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Determined text tokens retain the upstream sampling/counting contract."""
+
+from types import SimpleNamespace
+
+import numpy as np
+import pytest
+import torch
+from vllm.sampling_params import SamplingParams
+from vllm.v1.worker.gpu.sample.sampler import Sampler
+from vllm.v1.worker.gpu.states import RequestState
+
+from tests.model_executor.models.moss_tts.test_local_model_state import _batch, _state
+from vllm_omni.model_executor.models.moss_tts.local_model_state import MossLocalModelState
+from vllm_omni.worker_v2.omni_ar_model_runner import OmniARModelRunner
+from vllm_omni.worker_v2.omni_model_runner import OmniGPUModelRunner
+
+pytestmark = pytest.mark.core_model
+
+
+@pytest.mark.cuda
+@pytest.mark.parametrize("seed", [None, 17])
+@pytest.mark.parametrize("temperature", [0.0, 1.7])
+def test_matches_real_sampler_for_mixed_prefill_decode_stop_and_reordering(seed, temperature):
+    if not torch.cuda.is_available():
+        pytest.skip("CUDA required")
+    device = torch.device("cuda")
+    reqs = RequestState(5, 32, 16, 0, 8, device)
+    sampler = Sampler(5, 8, device, reqs)
+    for i, length in enumerate([4, 1, 3]):
+        reqs.add_request(str(i), length, [2] * length, 0, 8)
+        slot = reqs.req_id_to_index[str(i)]
+        sampler.add_request(slot, length, SamplingParams(temperature=temperature, top_k=3, top_p=0.8, seed=seed))
+    reqs.apply_staged_writes()
+    sampler.apply_staged_writes()
+    state = _state(MossLocalModelState, device)
+    state._direct_tokens = True
+    state.model._batch_state = None
+    batch = _batch(device, [2, 4, 3], [1, 2, 1])
+    batch.seq_lens = torch.tensor([5, 2, 2], device=device, dtype=torch.int32)
+    batch.cu_num_logits_np = np.arange(4, dtype=np.int32)
+    batch.cu_num_logits = torch.tensor(batch.cu_num_logits_np, device=device)
+    batch.expanded_idx_mapping = batch.idx_mapping
+    batch.expanded_local_pos = torch.zeros(3, dtype=torch.int32, device=device)
+    batch.logits_indices = batch.query_start_loc[1:].long() - 1
+    batch.positions = torch.arange(batch.num_tokens, device=device)
+    state.model._batch_should_continue = torch.tensor([False, True, True], device=device)
+    with torch.inference_mode():
+        logits = state.model.compute_logits(torch.ones(3, 4, device=device))
+        expected = sampler(logits, batch)
+        actual = state.sample_determined_tokens(batch, sampler)
+    assert actual is not None
+    for field in ["sampled_token_ids", "num_sampled", "num_rejected"]:
+        torch.testing.assert_close(getattr(actual, field), getattr(expected, field), rtol=0, atol=0)
+    assert actual.num_sampled.tolist() == [1, 0, 1]
+    snapshot = actual.sampled_token_ids.clone()
+    state.model._batch_should_continue.logical_not_()
+    torch.testing.assert_close(actual.sampled_token_ids, snapshot, rtol=0, atol=0)
+
+
+@pytest.mark.cuda
+@pytest.mark.parametrize(
+    "constraint",
+    [
+        {"min_tokens": 2, "stop_token_ids": [1]},
+        {"allowed_token_ids": [1, 2]},
+        {"logit_bias": {1: -10.0}},
+    ],
+)
+def test_real_sampler_constraints_disable_determined_tokens(constraint):
+    if not torch.cuda.is_available():
+        pytest.skip("CUDA required")
+    device = torch.device("cuda")
+    reqs = RequestState(5, 32, 16, 0, 8, device)
+    sampler = Sampler(5, 8, device, reqs)
+    reqs.add_request("constrained", 1, [2], 0, 8)
+    slot = reqs.req_id_to_index["constrained"]
+    sampler.add_request(slot, 1, SamplingParams(**constraint))
+    state = _state(MossLocalModelState, device)
+    state._direct_tokens = True
+    # Use the actual upstream request registration, including its combined
+    # min-token/allowed-ID/bias flag, rather than mocking that flag ourselves.
+    assert state.sample_determined_tokens(SimpleNamespace(idx_mapping_np=np.array([slot])), sampler) is None
+
+
+@pytest.mark.cpu
+@pytest.mark.parametrize("feature", ["logprobs", "bias", "penalty", "bad_words", "thinking", "trace", "mask", "nans"])
+def test_distribution_features_fall_back_before_using_gpu(feature):
+    sampler = object.__new__(Sampler)
+    sampler.compute_nans = feature == "nans"
+    sampler.return_sampling_mask = feature == "mask"
+    sampler.trace_replay_state = object() if feature == "trace" else None
+    sampler.get_logprobs_dims = lambda rows: (1, 0) if feature == "logprobs" else None
+    sampler.logit_bias_state = SimpleNamespace(use_logit_bias=np.array([feature == "bias"]))
+    sampler.penalties_state = SimpleNamespace(use_penalty=np.array([feature == "penalty"]))
+    sampler.bad_words_state = SimpleNamespace(num_bad_words=SimpleNamespace(np=np.array([feature == "bad_words"])))
+    sampler.thinking_budget_state = SimpleNamespace(enabled=True, use_thinking_budget=np.array([feature == "thinking"]))
+    state = object.__new__(MossLocalModelState)
+    state._direct_tokens = True
+    assert state.sample_determined_tokens(SimpleNamespace(idx_mapping_np=np.array([0])), sampler) is None
+
+
+@pytest.mark.cpu
+@pytest.mark.parametrize("feature", ["grammar", "sharding", "pipeline", "speculative", "empty", "declined"])
+def test_runner_retains_upstream_path_for_unsupported_execution(mocker, feature):
+    hook = mocker.Mock(return_value=None)
+
+    class State:
+        def sample_determined_tokens(self, batch, sampler):
+            return hook(batch, sampler)
+
+    runner = object.__new__(OmniARModelRunner)
+    runner.model_state = State()
+    runner.sampler = object()
+    runner.batch_sharder = object() if feature == "sharding" else None
+    runner.pp_handler = object() if feature == "pipeline" else None
+    batch = SimpleNamespace(num_draft_tokens=int(feature == "speculative"), num_reqs=0 if feature == "empty" else 1)
+    grammar = object() if feature == "grammar" else None
+    fallback = mocker.patch.object(OmniGPUModelRunner, "sample", return_value="upstream")
+    hidden = object()
+    assert runner.sample(hidden, batch, grammar) == "upstream"
+    fallback.assert_called_once_with(hidden, batch, grammar)
+    assert hook.call_count == int(feature == "declined")

tests/model_executor/models/moss_tts/test_local_eager_mtp.py

@@ -0,0 +1,166 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Eager MTP emits the same frames one step earlier with identical backbone inputs."""
+
+from types import SimpleNamespace
+
+import numpy as np
+import pytest
+import torch
+
+from tests.model_executor.models.moss_tts.test_local_model_state import _admit, _batch, _state, _step
+from vllm_omni.model_executor.models.moss_tts.local_model_state import MossLocalModelState
+
+pytestmark = pytest.mark.core_model
+
+
+@pytest.fixture(params=[pytest.param("cpu", marks=pytest.mark.cpu), pytest.param("cuda", marks=pytest.mark.cuda)])
+def device(request):
+    if request.param == "cuda" and not torch.cuda.is_available():
+        pytest.skip("CUDA required")
+    return torch.device(request.param)
+
+
+def _run(state, device, schedule, stop_step):
+    computed, prompt = np.zeros(5, dtype=np.int32), np.zeros(5, dtype=np.int32)
+    prompt[3], prompt[0] = 4, 2
+    req_states = SimpleNamespace(prompt_len=prompt, num_computed_tokens=computed)
+    embeds = []
+    frames: dict[int, list[torch.Tensor]] = {3: [], 0: []}
+    for index, (slots, counts) in enumerate(schedule):
+        state.model.force_stop = index == stop_step
+        embed, payload, _ = _step(state, _batch(device, slots, counts), req_states)
+        embeds.append(embed)
+        rows = payload.get("codes", {}).get("audio", [])
+        for slot, row in zip(slots, rows):
+            if row.numel() and bool(row.ne(state.model.audio_pad_token_id).any()):
+                frames[slot].append((index, row.clone()))
+        for slot, count in zip(slots, counts):
+            computed[slot] += count
+    return embeds, frames
+
+
+def test_eager_frames_match_canonical_one_step_earlier(device):
+    schedule = [([3, 0], [2, 2]), ([3, 0], [2, 1]), ([3, 0], [1, 1]), ([0, 3], [1, 1]), ([3, 0], [1, 1])]
+    canonical, eager = (_state(MossLocalModelState, device) for _ in range(2))
+    eager._local_eager_mtp = True
+    for state in (canonical, eager):
+        _admit(state, 3, "long", 17)
+        _admit(state, 0, "short", 17)
+    # Stop both streams at the canonical step that would draw the 4th step's frames.
+    c_embeds, c_frames = _run(canonical, device, schedule, stop_step=4)
+    e_embeds, e_frames = _run(eager, device, schedule, stop_step=3)
+
+    for c, e in zip(c_embeds, e_embeds):
+        torch.testing.assert_close(c, e, rtol=0, atol=0)
+    for slot in (3, 0):
+        assert [f for _, f in c_frames[slot]] and len(c_frames[slot]) == len(e_frames[slot])
+        for (ci, cf), (ei, ef) in zip(c_frames[slot], e_frames[slot]):
+            assert ei == ci - 1
+            torch.testing.assert_close(cf, ef, rtol=0, atol=0)
+
+
+def test_stopped_stream_does_not_emit_again(device):
+    state = _state(MossLocalModelState, device)
+    state._local_eager_mtp = True
+    _admit(state, 0, "short", 17)
+    computed, prompt = np.zeros(5, dtype=np.int32), np.zeros(5, dtype=np.int32)
+    prompt[0] = 2
+    req_states = SimpleNamespace(prompt_len=prompt, num_computed_tokens=computed)
+    emitted = []
+    for index, count in enumerate([2, 1, 1, 1]):
+        state.model.force_stop = index == 1
+        _, payload, _ = _step(state, _batch(device, [0], [count]), req_states)
+        row = payload["codes"]["audio"][0]
+        emitted.append(bool(row.ne(state.model.audio_pad_token_id).any()))
+        computed[0] += count
+        state.model.force_stop = False
+    # Frame at the prefill step, stop at step 1, then asynchronous overrun steps.
+    assert emitted == [True, False, False, False]
+
+
+def test_unmarked_decode_row_ignores_stale_slot_state(device):
+    canonical, eager = (_state(MossLocalModelState, device) for _ in range(2))
+    eager._local_eager_mtp = True
+    for state in (canonical, eager):
+        _admit(state, 0, "short", 17)
+    # A decode row whose completing prefill was not observed in this state.
+    eager._keep_pool[0] = False
+    eager._eager_emb_pool[0].fill_(3.0)
+    computed, prompt = np.zeros(5, dtype=np.int32), np.zeros(5, dtype=np.int32)
+    prompt[0], computed[0] = 2, 2
+    req_states = SimpleNamespace(prompt_len=prompt, num_computed_tokens=computed)
+    for state in (canonical, eager):
+        state.intermediate_buffer.buffers[0]["audio_state"] = {"is_stopping": False}
+    c_embed, c_payload, _ = _step(canonical, _batch(device, [0], [1]), req_states)
+    e_embed, e_payload, _ = _step(eager, _batch(device, [0], [1]), req_states)
+    torch.testing.assert_close(c_embed, e_embed, rtol=0, atol=0)
+    torch.testing.assert_close(c_payload["codes"]["audio"][0], e_payload["codes"]["audio"][0], rtol=0, atol=0)
+
+
+@pytest.mark.cpu
+def test_first_prefill_through_runner_preserves_parent_eager_contract(monkeypatch, mocker):
+    """The Local config must not enable the parent's incompatible eager hook."""
+    from vllm_omni.worker_v2 import omni_ar_model_runner as runner_module
+    from vllm_omni.worker_v2.model_states.eager_mtp import EagerMTPState
+    from vllm_omni.worker_v2.model_states.omni_model_state import OmniModelState
+    from vllm_omni.worker_v2.omni_ar_model_runner import OmniARModelRunner
+
+    device = torch.device("cpu")
+    template = _state(MossLocalModelState, device)
+    model = template.model
+    model.config = SimpleNamespace(mrv2_eager_mtp=True)
+
+    def initialize_parent(state, *args):
+        state.__dict__.update(template.__dict__)
+        state._eager_mtp, state._eager_rows = False, None
+        state._eager_state = EagerMTPState(state)
+
+    monkeypatch.setattr(OmniModelState, "__init__", initialize_parent)
+    state = MossLocalModelState(SimpleNamespace(), model, None, device)
+    _admit(state, 0, "req", 17)
+    batch = _batch(device, [0], [2])
+    batch.req_ids = ["req"]
+    computed, prompt = np.zeros(5, dtype=np.int32), np.zeros(5, dtype=np.int32)
+    prompt[0] = 2
+    _step(state, batch, SimpleNamespace(prompt_len=prompt, num_computed_tokens=computed))
+    hook = mocker.spy(state, "run_eager_mtp")
+
+    runner = OmniARModelRunner.__new__(OmniARModelRunner)
+    runner.execute_model_state = SimpleNamespace(
+        input_batch=batch,
+        hidden_states=torch.zeros(2, 4),
+        finished_req_ids=set(),
+        ec_connector_output=None,
+        routed_experts=None,
+    )
+    runner._kv_extracted_req_ids = runner._last_aux_output = runner._last_multimodal_outputs = None
+    runner.is_last_pp_rank, runner.pp_handler, runner.check_ep_fault = True, None, False
+    runner.model_config = SimpleNamespace(async_chunk=False)
+    runner.vllm_config = SimpleNamespace(model_config=SimpleNamespace(engine_output_type="text"))
+    runner.model_state, runner.model = state, model
+    runner.req_states = SimpleNamespace(
+        all_token_ids=SimpleNamespace(gpu=torch.zeros(5, 2)),
+        num_computed_tokens=SimpleNamespace(gpu=torch.from_numpy(computed)),
+        prompt_len=SimpleNamespace(np=prompt, gpu=torch.from_numpy(prompt)),
+    )
+    runner.main_stream = runner.output_copy_stream = mocker.Mock()
+    runner.eplb = runner._finalize_native_data_plane_output = runner._reserve_native_data_plane_outputs = mocker.Mock()
+    runner.sampler = None
+    runner.sample = mocker.Mock(
+        return_value=(
+            SimpleNamespace(sampled_token_ids=torch.tensor([[2]])),
+            torch.ones(1),
+            torch.zeros(1),
+        )
+    )
+    runner.prompt_logprobs_worker = SimpleNamespace(compute_prompt_logprobs=mocker.Mock(return_value={}))
+    runner.postprocess_sampled = mocker.Mock()
+    runner.kv_connector = SimpleNamespace(post_forward=mocker.Mock(return_value=None))
+    mock_out = SimpleNamespace(copy_event=None)
+    monkeypatch.setattr(runner_module, "OmniAsyncOutput", mocker.Mock(return_value=mock_out))
+
+    assert runner.sample_tokens(None) is mock_out
+    hook.assert_called_once()
+    assert state._eager_mtp is False and state._eager_rows is None
+    assert state._local_eager_mtp is True

tests/model_executor/models/moss_tts/test_local_model_state.py

@@ -0,0 +1,286 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Compare slot lifecycle with canonical Local hooks, including output ownership."""
+
+from contextlib import nullcontext
+from types import SimpleNamespace
+from unittest.mock import MagicMock, patch
+
+import numpy as np
+import pytest
+import torch
+from vllm.v1.worker.gpu.model_states.default import DefaultModelState
+
+from vllm_omni.model_executor.models.moss_tts.local_model_state import MossLocalModelState
+from vllm_omni.model_executor.models.moss_tts.modeling_moss_tts_talker import MossTTSLocalTalkerForGeneration
+from vllm_omni.worker_v2.model_states import init_omni_model_state
+from vllm_omni.worker_v2.model_states.eager_mtp import EagerMTPState
+from vllm_omni.worker_v2.model_states.intermediate_buffer import OmniIntermediateBuffer
+from vllm_omni.worker_v2.model_states.omni_model_state import OmniModelState
+
+pytestmark = pytest.mark.core_model
+
+
+@pytest.fixture(params=[pytest.param("cpu", marks=pytest.mark.cpu), pytest.param("cuda", marks=pytest.mark.cuda)])
+def device(request):
+    if request.param == "cuda" and not torch.cuda.is_available():
+        pytest.skip("CUDA required")
+    return torch.device(request.param)
+
+
+def _model(device):
+    model = MossTTSLocalTalkerForGeneration.__new__(MossTTSLocalTalkerForGeneration)
+    torch.nn.Module.__init__(model)
+    dtype = torch.bfloat16 if device.type == "cuda" else torch.float32
+    model.model = torch.nn.Module()
+    weight = torch.arange(32, dtype=dtype, device=device).reshape(8, 4) / 16
+    model.model.embed_tokens = torch.nn.Embedding(8, 4, _weight=weight)
+    model.hidden_size, model.n_vq = 4, 2
+    model.audio_pad_token_id, model.text_vocab_size = 8, 8
+    model.audio_assistant_slot_token_id, model.im_end_token_id = 2, 3
+    model.talker_mtp_output_key = ("audio_codes", "current")
+    model.talker_mtp_graph_safe = False
+    model.gpu_resident_buffer_keys = {("hidden_states", "last"), ("audio_codes", "current")}
+    model.audio_lm_heads = model.audio_embeddings = model.local_text_lm_head = None
+    model._audio_embed = lambda codes: codes[:, :1].expand(-1, 4).to(dtype) / 8
+    model.force_stop = False
+
+    def frame(hidden, *_args, generator=None, temperature=None, top_k=None, top_p=None, **_kwargs):
+        assert (temperature, top_k, top_p) == (1.7, 25, 0.8)
+        noise = torch.randint(0, 8, (hidden.shape[0], 2), device=device, generator=generator)
+        codes = (hidden[:, :2].mul(16).long() + noise) % 8
+        return torch.full((hidden.shape[0],), not model.force_stop, device=device), codes
+
+    model.local_transformer = SimpleNamespace(generate_frame=frame)
+    return model
+
+
+def _state(cls, device):
+    state = object.__new__(cls)
+    state.model = _model(device)
+    state.device, state.dtype = device, state.model.model.embed_tokens.weight.dtype
+    state.has_preprocess = state.has_postprocess = state.have_multimodal_outputs = True
+    state.scheduler_config = SimpleNamespace(max_num_seqs=5)
+    state.vllm_config = SimpleNamespace()
+    state.intermediate_buffer = OmniIntermediateBuffer(5)
+    state._static_inputs_embeds = torch.zeros(16, 4, device=device, dtype=state.dtype)
+    state._mtp_input_ids = torch.zeros(5, device=device, dtype=torch.long)
+    state._mtp_offsets = torch.zeros(5, device=device, dtype=torch.long)
+    for name in ("_mtp_input_embeds", "_mtp_hidden", "_mtp_text_step"):
+        setattr(state, name, torch.zeros(5, 4, device=device, dtype=state.dtype))
+    state._mtp_runner = state._mtp_sample_uniforms = None
+    state._mtp_generators = {}
+    state._stream_pos = {}
+    state._stream_decode_event = None
+    state._eager_state = EagerMTPState(state)
+    if cls is MossLocalModelState:
+        state._init_slot_buffers(5, 4, device, state.dtype)
+    return state
+
+
+def _admit(state, slot, name, seed):
+    params = SimpleNamespace(extra_args={"tts_local_seed": seed} if seed is not None else {}, seed=None)
+    with patch.object(DefaultModelState, "add_request", return_value=None):
+        state.add_request(slot, SimpleNamespace(req_id=name, mm_features=[], sampling_params=params))
+    state.intermediate_buffer.buffers[slot]["codes"] = {"ref": torch.tensor([[1, 2], [3, 4], [5, 6], [2, 3]])}
+
+
+def _batch(device, slots, counts):
+    starts = np.array([0, *np.cumsum(counts)], dtype=np.int32)
+    return SimpleNamespace(
+        idx_mapping_np=np.array(slots),
+        idx_mapping=torch.tensor(slots, device=device),
+        num_reqs=len(slots),
+        num_tokens=int(starts[-1]),
+        num_scheduled_tokens=np.array(counts),
+        query_start_loc_np=starts,
+        query_start_loc=torch.tensor(starts, device=device),
+        # Upstream MRV2 uses int32 token buffers, unlike the MTP graph inputs.
+        input_ids=torch.full((int(starts[-1]),), 2, dtype=torch.int32, device=device),
+    )
+
+
+def _step(state, batch, req_states, dispatcher=None):
+    inputs = {"input_ids": batch.input_ids, "inputs_embeds": state._static_inputs_embeds[: batch.num_tokens]}
+    with (
+        patch("vllm.forward_context.set_forward_context", return_value=nullcontext()),
+        patch(
+            "vllm_omni.model_executor.models.moss_tts.local_model_state.set_forward_context", return_value=nullcontext()
+        ),
+        torch.inference_mode(),
+    ):
+        state.run_preprocess(batch, inputs, req_states, dispatcher)
+        hidden = inputs["inputs_embeds"] + 0.25
+        state.run_postprocess(hidden, batch)
+        _, payload = state.postprocess_model_output(hidden, batch, req_states)
+        last = batch.query_start_loc[1:] - 1
+        tokens = state.model.compute_logits(hidden.index_select(0, last)).argmax(-1)
+        result = inputs["inputs_embeds"].clone(), payload, tokens
+        hidden.fill_(-100)  # next graph replay can overwrite its output immediately
+    return result
+
+
+@pytest.mark.parametrize("seed", [None, 17])
+@pytest.mark.parametrize("batch_prefill", [False, True])
+def test_slot_matches_canonical_mixed_prefill_reorder_stop_and_reuse(device, seed, batch_prefill):
+    reference, candidate = (_state(cls, device) for cls in (OmniModelState, MossLocalModelState))
+    candidate._batch_prefill = batch_prefill
+    for state in (reference, candidate):
+        _admit(state, 3, "long", seed)
+        _admit(state, 0, "short", seed)
+    computed, prompt = np.zeros(5, dtype=np.int32), np.zeros(5, dtype=np.int32)
+    prompt[3], prompt[0] = 4, 2
+    req_states = SimpleNamespace(prompt_len=prompt, num_computed_tokens=computed)
+    saved = []
+    for index, (slots, counts) in enumerate(
+        [
+            ([3, 0], [2, 2]),
+            ([0, 3], [1, 1]),
+            ([3, 0], [1, 1]),
+            ([3, 0], [1, 1]),
+            ([0, 3], [1, 1]),
+            ([3, 0], [1, 1]),
+        ]
+    ):
+        if index == 4:
+            # Cancellation/removal followed by reuse must not retain hidden,
+            # reference offset, control, or the old request's RNG stream.
+            for state in (reference, candidate):
+                state.remove_request("short")
+                assert "short" not in state._mtp_generators
+                _admit(state, 0, "replacement", seed)
+            computed[0], prompt[0] = 0, 1
+        outputs = []
+        for state in (reference, candidate):
+            state.model.force_stop = index == 3
+            torch.manual_seed(91 + index)
+            outputs.append(_step(state, _batch(device, slots, counts), req_states))
+        a, b = outputs
+        torch.testing.assert_close(a[0], b[0], rtol=0, atol=0)
+        torch.testing.assert_close(a[2], b[2], rtol=0, atol=0)
+        if index == 3:
+            assert b[2].tolist() == [3, 3]
+        assert a[1].keys() == b[1].keys()
+        for ar, br in zip(a[1].get("codes", {}).get("audio", []), b[1].get("codes", {}).get("audio", [])):
+            torch.testing.assert_close(ar, br, rtol=0, atol=0)
+            saved.append((br, br.clone()))
+        for slot, count in zip(slots, counts):
+            computed[slot] += count
+            old = reference.intermediate_buffer.buffers[slot]
+            new = candidate.intermediate_buffer.buffers[slot]
+            assert old["ref_offset"] == new["ref_offset"]
+            torch.testing.assert_close(candidate._hidden_pool[slot], old["hidden_states"]["last"], rtol=0, atol=0)
+            assert "hidden_states" not in new and "audio_codes" not in new
+    candidate._codes_pool.fill_(8)
+    candidate._hidden_pool.zero_()
+    for current, snapshot in saved:
+        torch.testing.assert_close(current, snapshot, rtol=0, atol=0)
+
+
+def test_slot_honors_external_stopping_control(device):
+    states = [_state(cls, device) for cls in (OmniModelState, MossLocalModelState)]
+    for state in states:
+        _admit(state, 0, "stop", None)
+        req = SimpleNamespace(prompt_len=np.ones(5), num_computed_tokens=np.zeros(5))
+        _step(state, _batch(device, [0], [1]), req)
+        state.intermediate_buffer.buffers[0]["audio_state"]["is_stopping"] = True
+        req.num_computed_tokens[:] = 1
+        output = _step(state, _batch(device, [0], [1]), req)
+        assert output[2].tolist() == [3]
+        assert output[1]["codes"]["audio"][0].tolist() == [[8, 8]]
+
+
+@pytest.mark.cuda
+def test_slot_mtp_graph_padding_reorder_and_owned_output():
+    if not torch.cuda.is_available():
+        pytest.skip("CUDA required")
+    states = [_state(cls, torch.device("cuda")) for cls in (OmniModelState, MossLocalModelState)]
+    for state in states:
+
+        def frame(hidden, *_args, **_kwargs):
+            return hidden[:, 0] < 3, hidden[:, :2].mul(16).long() % 8
+
+        state.model.local_transformer.generate_frame = frame
+        for slot in [0, 3, 4]:
+            _admit(state, slot, f"req{slot}", None)
+    candidate = states[1]
+    buffers = [
+        candidate._mtp_input_ids[:4],
+        candidate._mtp_input_embeds[:4],
+        candidate._mtp_hidden[:4],
+        candidate._mtp_text_step[:4],
+    ]
+    stream = torch.cuda.Stream()
+    stream.wait_stream(torch.cuda.current_stream())
+    with torch.cuda.stream(stream), torch.inference_mode():
+        for _ in range(3):
+            candidate.model.mtp(*buffers)
+    torch.cuda.current_stream().wait_stream(stream)
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph), torch.inference_mode():
+        graph_output = candidate.model.mtp(*buffers)
+
+    def replay(*args, **kwargs):
+        assert all(a.data_ptr() == b.data_ptr() for a, b in zip(args, buffers))
+        assert args[0].shape[0] == 4
+        graph.replay()
+        return graph_output
+
+    candidate._mtp_runner = replay
+    candidate._is_mtp_graph_runner = lambda: True
+    from vllm.config import CUDAGraphMode
+
+    def dispatcher(bsz):
+        return SimpleNamespace(num_tokens=4, cg_mode=CUDAGraphMode.FULL)
+
+    req = SimpleNamespace(prompt_len=np.ones(5, dtype=np.int32), num_computed_tokens=np.zeros(5, dtype=np.int32))
+    retained = []
+    for slots in [[0, 3, 4], [3, 0], [4, 0, 3], [3, 4]]:
+        out = [
+            _step(
+                state,
+                _batch(torch.device("cuda"), slots, [1] * len(slots)),
+                req,
+                dispatcher if state is candidate else None,
+            )
+            for state in states
+        ]
+        torch.testing.assert_close(out[0][0], out[1][0], rtol=0, atol=0)
+        torch.testing.assert_close(out[0][2], out[1][2], rtol=0, atol=0)
+        for a, b in zip(out[0][1].get("codes", {}).get("audio", []), out[1][1].get("codes", {}).get("audio", [])):
+            torch.testing.assert_close(a, b, rtol=0, atol=0)
+            retained.append((b, b.clone()))
+        req.num_computed_tokens[slots] += 1
+    with torch.inference_mode():
+        for tensor in graph_output:
+            tensor.fill_(7)
+    for tensor, copy in retained:
+        torch.testing.assert_close(tensor, copy, rtol=0, atol=0)
+
+
+@pytest.mark.cpu
+def test_factory_requires_explicit_class_capability_and_valid_state(monkeypatch):
+    monkeypatch.setattr(OmniModelState, "__init__", lambda *args: None)
+    dynamic = MagicMock(has_preprocess=True, create_omni_model_state=None)
+    assert isinstance(init_omni_model_state(None, dynamic, None, torch.device("cpu")), OmniModelState)
+    dynamic.create_mrv2_model_state.assert_not_called()
+
+    class Invalid:
+        has_preprocess = True
+
+        def create_mrv2_model_state(self, *args):
+            return object()
+
+    with pytest.raises(TypeError, match="must return a ModelState"):
+        init_omni_model_state(None, Invalid(), None, torch.device("cpu"))
+
+
+@pytest.mark.cpu
+@pytest.mark.parametrize("enabled", [False, True])
+def test_local_factory_default_and_option(monkeypatch, enabled):
+    monkeypatch.setattr(OmniModelState, "__init__", lambda *args: None)
+    monkeypatch.setattr(MossLocalModelState, "__init__", lambda *args: None)
+    model = _model(torch.device("cpu"))
+    model.config = SimpleNamespace(mrv2_gpu_slot_state=enabled)
+    state = init_omni_model_state(None, model, None, torch.device("cpu"))
+    assert type(state) is (MossLocalModelState if enabled else OmniModelState)

tests/model_executor/models/moss_tts/test_local_prefill_batch.py

@@ -0,0 +1,86 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Reference batching preserves coordinates and asynchronous source ownership."""
+
+from types import SimpleNamespace
+
+import numpy as np
+import pytest
+import torch
+
+from tests.model_executor.models.moss_tts.test_local_model_state import _batch, _state
+from vllm_omni.model_executor.models.moss_tts.local_model_state import MossLocalModelState
+from vllm_omni.model_executor.models.moss_tts.modeling_moss_tts_talker import MossTTSLocalTalkerForGeneration
+
+pytestmark = pytest.mark.core_model
+
+
+@pytest.mark.cpu
+@pytest.mark.parametrize("flat", [False, True])
+def test_prefill_reuses_text_embeddings_and_noncontiguous_reference_offsets(mocker, flat):
+    state = _state(MossLocalModelState, torch.device("cpu"))
+    state._batch_prefill = True
+    codes = torch.tensor([[1, 2], [3, 4], [5, 6], [2, 3]])
+    for slot, offset in [(3, 1), (0, 2)]:
+        state.intermediate_buffer.buffers[slot] = {
+            "req_id": str(slot),
+            "codes": {"ref": codes.flatten() if flat else codes},
+            "ref_offset": offset,
+        }
+    batch = _batch(torch.device("cpu"), [3, 0], [2, 1])
+    req = SimpleNamespace(prompt_len=np.full(5, 8), num_computed_tokens=np.zeros(5))
+    embeds = state._static_inputs_embeds[:3]
+    expected = state.model.embed_input_ids(batch.input_ids) + state.model._audio_embed(codes[[1, 2, 2]])
+    embed = mocker.spy(state.model.model.embed_tokens, "forward")
+    with torch.inference_mode():
+        state.run_preprocess(batch, {"input_ids": batch.input_ids, "inputs_embeds": embeds}, req)
+    torch.testing.assert_close(embeds, expected, rtol=0, atol=0)
+    assert embed.call_count == 1
+    assert state.intermediate_buffer.buffers[3]["ref_offset"] == 3
+    assert state.intermediate_buffer.buffers[0]["ref_offset"] == 3
+
+
+@pytest.mark.cuda
+def test_pending_reference_uploads_survive_staging_reuse_and_overflow():
+    if not torch.cuda.is_available():
+        pytest.skip("CUDA required")
+    state = _state(MossLocalModelState, torch.device("cuda"))
+    stream = torch.cuda.Stream()
+    stream.wait_stream(torch.cuda.current_stream())
+    outputs = []
+    with torch.cuda.stream(stream), torch.inference_mode():
+        # Keep the first H2D pending while CPU prepares all subsequent batches.
+        torch.cuda._sleep(100_000_000)
+        for value in range(1, 8):
+            embeds = torch.zeros(4, 4, device="cuda", dtype=state.dtype)
+            codes = torch.full((2, 2), value)
+            state._apply_reference_batch(embeds, [(1, codes)])
+            outputs.append((embeds, value))
+            codes.fill_(-1)
+    stream.synchronize()
+    assert len(state._prefill_staging) <= 3
+    for output, value in outputs:
+        expected = torch.zeros_like(output)
+        expected[1:3] = value / 8
+        torch.testing.assert_close(output, expected, rtol=0, atol=0)
+
+
+@pytest.mark.cuda
+def test_batched_real_audio_embedding_matches_scalar_bfloat16_reduction():
+    if not torch.cuda.is_available():
+        pytest.skip("CUDA required")
+    state = _state(MossLocalModelState, torch.device("cuda"))
+    model = state.model
+    model.n_vq = 12
+    model.audio_vocab_size = model.audio_pad_token_id = 1024
+    model._stacked_audio_emb_w = torch.randn(12, 1024, 3072, device="cuda", dtype=torch.bfloat16)
+    model._audio_embed = MossTTSLocalTalkerForGeneration._audio_embed.__get__(model)
+    state._static_inputs_embeds = torch.zeros(512, 3072, device="cuda", dtype=torch.bfloat16)
+    references = [(start, torch.randint(0, 1025, (count, 12))) for start, count in [(0, 169), (201, 11), (220, 257)]]
+    with torch.inference_mode():
+        actual = torch.randn_like(state._static_inputs_embeds)
+        expected = actual.clone()
+        for start, codes in references:
+            expected[start : start + len(codes)] += model._audio_embed(codes.to("cuda"))
+        state._apply_reference_batch(actual, references)
+    torch.testing.assert_close(actual, expected, rtol=0, atol=0)

tests/model_executor/models/moss_tts/test_local_short_attention.py

@@ -0,0 +1,82 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+import pytest
+import torch
+import torch.nn.functional as F
+
+pytestmark = [
+    pytest.mark.core_model,
+    pytest.mark.cuda,
+    pytest.mark.tts,
+    pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required"),
+]
+
+
+@pytest.mark.parametrize("batch,heads,dim", [(1, 4, 64), (16, 32, 80), (128, 32, 80), (3, 8, 128)])
+def test_all_depth_positions_with_strided_cache_and_dirty_future(batch, heads, dim):
+    from vllm_omni.model_executor.models.moss_tts.local_short_attention import local_short_attention
+
+    torch.manual_seed(43)
+    query = torch.randn(batch, heads, 1, dim, device="cuda", dtype=torch.bfloat16)
+    key = torch.randn(batch, heads, 12, dim, device="cuda", dtype=torch.bfloat16)
+    value = torch.randn_like(key)
+    for position in range(12):
+        # Non-contiguous prefixes retain the stride of the full 12-slot cache.
+        k, v = key[:, :, : position + 1], value[:, :, : position + 1]
+        before_k, before_v = key.clone(), value.clone()
+        expected = F.scaled_dot_product_attention(query, k, v)
+        actual = local_short_attention(query, k, v)
+        assert torch.equal(key, before_k) and torch.equal(value, before_v)
+        if position == 0:
+            assert torch.equal(actual, v)
+        torch.testing.assert_close(actual, expected, atol=0.015625, rtol=0.01)
+        # Future cache contents must never affect an earlier position.
+        key[:, :, position + 1 :] += 100
+        value[:, :, position + 1 :] -= 100
+        assert torch.equal(local_short_attention(query, k, v), actual)
+        key.copy_(before_k)
+        value.copy_(before_v)
+
+
+def test_compile_and_graph_replay_reads_fresh_inputs():
+    from vllm_omni.model_executor.models.moss_tts.local_short_attention import local_short_attention
+
+    query = torch.randn(16, 32, 1, 96, device="cuda", dtype=torch.bfloat16)
+    key = torch.randn(16, 32, 12, 96, device="cuda", dtype=torch.bfloat16)
+    value = torch.randn_like(key)
+    compiled = torch.compile(local_short_attention, fullgraph=True)
+    for _ in range(3):
+        compiled(query, key, value)
+    torch.accelerator.synchronize()
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph):
+        result = compiled(query, key, value)
+    for _ in range(3):
+        query.normal_()
+        key.normal_()
+        value.normal_()
+        graph.replay()
+        torch.testing.assert_close(result, F.scaled_dot_product_attention(query, key, value), atol=0.015625, rtol=0.01)
+
+
+def test_depth_block_cache_overwrite_and_compiled_positions(monkeypatch):
+    from transformers import GPT2Config
+
+    from vllm_omni.model_executor.models.moss_tts.modeling_moss_tts_local_depth import (
+        MossTTSLocalDepthTransformer,
+    )
+
+    monkeypatch.setenv("VLLM_OMNI_MOSS_LOCAL_SHORT_ATTN", "1")
+    model = MossTTSLocalDepthTransformer(GPT2Config(n_embd=320, n_head=4, n_inner=640)).cuda().bfloat16().eval()
+    attn = model.h[0].attn
+    attn.prepare_rope_cache(12, torch.device("cuda"), torch.bfloat16)
+    cache = tuple(torch.empty(4, 4, 12, 80, device="cuda", dtype=torch.bfloat16) for _ in range(2))
+    forward = torch.compile(model._forward_prefix, fullgraph=True, dynamic=True)
+    with torch.inference_mode():
+        for _ in range(2):
+            inputs = torch.randn(4, 12, 320, device="cuda", dtype=torch.bfloat16)
+            # The second frame overwrites the previous frame's K/V before use.
+            for position in range(12):
+                expected = model._forward_prefix(inputs[:, : position + 1])[:, -1:]
+                actual = forward(inputs[:, position : position + 1], cache, position)
+                torch.testing.assert_close(actual, expected, atol=0.03125, rtol=0.02)

tests/model_executor/models/moss_tts/test_reference_encoder.py

@@ -6,6 +6,7 @@ single-flight, and micro-batched encoding."""
 import asyncio
 import threading
 import time
+from tempfile import TemporaryDirectory
 
 import pytest
 import torch
@@ -598,3 +599,239 @@ def test_reference_resampler_is_cached_and_matches_functional_resample():
         assert info.misses == 1 and info.hits == 1
     finally:
         _reference_resampler.cache_clear()
+
+
+# --------------------------------------------------------------------------- #
+# Inline data: URI fast path                                                  #
+# --------------------------------------------------------------------------- #
+
+
+class _DigestAudio(_FakeAudio):
+    """Resolver whose key for data: URIs is the SHA-1 digest, like serving."""
+
+    def __init__(self, *, digest_keys=True):
+        super().__init__()
+        self._digest_keys = digest_keys
+
+    async def resolve(self, ref_str):
+        wav_list, sr, _ = await super().resolve(ref_str)
+        if not self._digest_keys:
+            return wav_list, sr, "rk:" + ref_str
+        import hashlib
+
+        return wav_list, sr, hashlib.sha1(ref_str.encode("utf-8")).hexdigest()
+
+    def artifact_key(self, cache_key):
+        return None
+
+
+async def test_inline_reference_skips_repeat_resolution(make_encoder):
+    proc, audio = _FakeProcessor(), _DigestAudio()
+    ref = "data:audio/wav;base64,QUJD"
+    audio.register(ref, 5)
+    enc = make_encoder(proc)
+
+    first, key1 = await enc.encode(ref, resolve_ref_audio=audio.resolve, get_artifact_key=audio.artifact_key)
+    first.fill_(0)  # callers own their copy
+    second, key2 = await enc.encode(ref, resolve_ref_audio=audio.resolve, get_artifact_key=audio.artifact_key)
+
+    assert audio.resolve_calls == [ref]
+    assert key1 == key2
+    assert torch.equal(second, torch.full((3, _N_VQ), 5, dtype=torch.long))
+
+
+async def test_inline_index_requires_matching_resolve_key(make_encoder):
+    proc, audio = _FakeProcessor(), _DigestAudio(digest_keys=False)
+    ref = "data:audio/wav;base64,REVG"
+    audio.register(ref, 6)
+    enc = make_encoder(proc)
+
+    for _ in range(2):
+        await enc.encode(ref, resolve_ref_audio=audio.resolve, get_artifact_key=audio.artifact_key)
+    assert audio.resolve_calls == [ref, ref]
+
+
+async def test_inline_reference_shared_across_encoders(make_encoder, monkeypatch):
+    # Keep the AF_UNIX socket path independent of pytest's nested base path.
+    with TemporaryDirectory(prefix="moss-ref-") as shared_dir:
+        monkeypatch.setenv("VLLM_OMNI_MOSS_REF_CODES_SHARED_DIR", shared_dir)
+        ref = "data:audio/wav;base64,R0hJ"
+        first_audio, second_audio = _DigestAudio(), _DigestAudio()
+        first_audio.register(ref, 9)
+        second_audio.register(ref, 9)
+        first, second = make_encoder(_FakeProcessor()), make_encoder(_FakeProcessor())
+
+        a, key_a = await first.encode(
+            ref, resolve_ref_audio=first_audio.resolve, get_artifact_key=first_audio.artifact_key
+        )
+        b, key_b = await second.encode(
+            ref, resolve_ref_audio=second_audio.resolve, get_artifact_key=second_audio.artifact_key
+        )
+
+        assert first_audio.resolve_calls == [ref] and second_audio.resolve_calls == []
+        assert key_a == key_b and torch.equal(a, b)
+
+
+# --------------------------------------------------------------------------- #
+# Reference-encoder CUDA graphs: selection and processor-equivalent prep      #
+# --------------------------------------------------------------------------- #
+
+
+class _GraphsStub:
+    def __init__(self, result):
+        self.result = result
+        self.calls: list[list[torch.Tensor]] = []
+
+    def encode(self, wavs):
+        self.calls.append(wavs)
+        return self.result
+
+
+class _LoudnessProcessor(_FakeProcessor):
+    @staticmethod
+    def loudness_normalize(wav):
+        return wav * 2.0
+
+
+async def test_encode_prepared_prefers_graphs_and_falls_back_to_processor(make_encoder):
+    proc = _LoudnessProcessor()
+    enc = make_encoder(proc)
+    graphs_codes = [torch.full((3, _N_VQ), 7, dtype=torch.long)]
+    enc._graphs_enabled, enc._graphs = True, _GraphsStub(graphs_codes)
+    prepared = [torch.ones(1, 100)]
+    assert enc._encode_prepared(prepared) is graphs_codes and proc.attempt_sizes == []
+    enc._graphs = _GraphsStub([None])  # no graph fits this clip
+    out = enc._encode_prepared(prepared)
+    assert proc.attempt_sizes == [1] and out[0].tolist() == [[1] * _N_VQ] * 3
+
+
+async def test_processor_prepare_matches_processor_channel_and_loudness_handling(make_encoder):
+    enc = make_encoder(_LoudnessProcessor())
+    mono = enc._processor_prepare(torch.ones(100))
+    assert mono.shape == (2, 100) and torch.equal(mono, torch.full((2, 100), 2.0))
+    many = enc._processor_prepare(torch.arange(3 * 4, dtype=torch.float32).reshape(3, 4))
+    assert torch.equal(many, torch.arange(8, dtype=torch.float32).reshape(2, 4) * 2.0)
+
+
+async def test_processor_prepare_resamples_to_the_tokenizer_rate(make_encoder):
+    torchaudio = pytest.importorskip("torchaudio")
+    proc = _LoudnessProcessor(sampling_rate=48000)
+    enc = make_encoder(proc)  # works at 24 kHz, the Local-v1.5 reference rate
+    wav = torch.randn(1, 2400)
+    expected = torchaudio.functional.resample(wav.repeat(2, 1), 24000, 48000) * 2.0
+    # The encoder resamples the mono clip before duplicating it. Some CPU conv
+    # kernels round a one-row batch differently from a two-row one, so the two
+    # orders agree to float32 precision, not bit for bit.
+    torch.testing.assert_close(enc._processor_prepare(wav), expected, rtol=0, atol=1e-5)
+
+
+@pytest.mark.parametrize("role", ["", "host", "client"])
+async def test_prepare_captures_graphs_or_starts_the_shared_host(make_encoder, monkeypatch, tmp_path, role):
+    from types import SimpleNamespace
+
+    from vllm_omni.model_executor.models.moss_tts import reference_encoder as re_mod
+
+    enc = make_encoder(_FakeProcessor())
+    enc._shares_encoder, enc._shared_codes_dir = True, str(tmp_path)
+    monkeypatch.setattr(re_mod, "shared_encoder_role", lambda: role)
+    captured, hosts, ready = [], [], []
+    monkeypatch.setattr(enc, "_reference_graphs", lambda: captured.append(True))
+    monkeypatch.setattr(enc, "_host_workers", lambda: [])
+    helper = SimpleNamespace(wait_until_ready=lambda: ready.append(True))
+
+    def make_host(*args, **kwargs):
+        hosts.append(args)
+        return helper
+
+    monkeypatch.setattr(re_mod, "SharedReferenceEncoderHost", make_host)
+    monkeypatch.setattr(re_mod, "SharedReferenceEncoderClient", lambda *args: helper)
+    enc.prepare()
+    enc.prepare()
+    # Every API process waits for the shared host before startup completes.
+    assert (len(captured), len(hosts)) == {"": (2, 0), "host": (0, 1), "client": (0, 0)}[role]
+    assert len(ready) == (2 if role else 0)
+
+
+async def test_clips_reach_the_encoder_at_the_tokenizer_rate(make_encoder, monkeypatch):
+    torchaudio = pytest.importorskip("torchaudio")
+    proc = _LoudnessProcessor(sampling_rate=48000)
+    enc = make_encoder(proc)  # works at 24 kHz
+    sent = []
+
+    def record_prepared(prepared):
+        sent.extend(prepared)
+        return [None] * len(prepared)
+
+    monkeypatch.setattr(enc, "_encode_shared_or_local", record_prepared)
+    wav = [float(v) for v in torch.randn(2400)]
+    enc._encode_batch_sync([(wav, _SR)])
+    # Resampled where the request arrived (a shared encoder receives it this way); still one channel.
+    expected = torchaudio.functional.resample(torch.tensor(wav)[None], 24000, 48000)
+    torch.testing.assert_close(sent[0], expected, rtol=0, atol=0)
+    # The eager path is told the clip is already at the tokenizer rate.
+    calls = []
+
+    def record_encoding(wavs, sampling_rate, n_vq=None):
+        calls.append(sampling_rate)
+        return [None]
+
+    monkeypatch.setattr(proc, "encode_audios_from_wav", record_encoding)
+    enc._graphs_enabled = False
+    enc._encode_prepared(sent)
+    assert calls == [48000]
+
+
+async def test_batcher_runs_up_to_max_inflight_batches_at_once():
+    inside = threading.Barrier(2, timeout=5)
+    seen: list[int] = []
+
+    def encode_batch(payload):
+        inside.wait()  # two batches are encoding at the same time
+        seen.append(len(payload))
+        return [torch.zeros(1) for _ in payload]
+
+    batcher = _RefEncodeBatcher(encode_batch, window_ms=0.0, max_batch=8, max_inflight=2)
+    try:
+        first = asyncio.ensure_future(batcher.submit([1.0], _SR))
+        await asyncio.sleep(0.05)  # the first batch is in flight before the second clip arrives
+        await asyncio.gather(first, batcher.submit([2.0], _SR))
+    finally:
+        await batcher.aclose()
+    assert seen == [1, 1]
+
+
+def test_reference_stream_priority(monkeypatch):
+    from vllm_omni.model_executor.models.moss_tts import reference_encoder as re_mod
+
+    monkeypatch.delenv(re_mod._REF_ENCODE_PRIORITY_ENV, raising=False)
+    assert re_mod.reference_stream_priority() == 0
+    monkeypatch.setenv(re_mod._REF_ENCODE_PRIORITY_ENV, "high")
+    expected = torch.cuda.Stream.priority_range()[1] if torch.cuda.is_available() else 0
+    assert re_mod.reference_stream_priority() == expected
+
+
+async def test_device_prep_resamples_on_the_tokenizer_device(make_encoder, monkeypatch):
+    torchaudio = pytest.importorskip("torchaudio")
+    from vllm_omni.model_executor.models.moss_tts import reference_encoder as re_mod
+
+    proc = _LoudnessProcessor(sampling_rate=48000)
+    monkeypatch.setattr(proc, "audio_tokenizer", torch.nn.Linear(1, 1), raising=False)  # CPU placement
+    enc = make_encoder(proc)
+    enc._gpu_prep = True
+    wav = torch.randn(1, 2400)
+    clip = enc._encoder_input(wav)
+    assert clip.shape == (1, 2400)  # left at the working rate for the encoder's device
+    expected = torchaudio.functional.resample(wav, 24000, 48000).repeat(2, 1) * 2.0
+    torch.testing.assert_close(enc._finish_prepare(clip), expected, rtol=1e-5, atol=1e-6)
+    # The eager fallback is told the clip's own rate.
+    calls = []
+
+    def record_encoding(wavs, sampling_rate, n_vq=None):
+        calls.append(sampling_rate)
+        return [None]
+
+    monkeypatch.setattr(proc, "encode_audios_from_wav", record_encoding)
+    enc._graphs_enabled = False
+    enc._encode_prepared([clip])
+    assert calls == [24000]
+    assert re_mod._REF_ENCODE_GPU_PREP_ENV == "VLLM_OMNI_MOSS_REF_GPU_PREP"

tests/model_executor/models/moss_tts/test_reference_encoder_graphs.py

@@ -0,0 +1,320 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Reference-encoder CUDA graphs reproduce the tokenizer's eager encode."""
+
+from contextlib import nullcontext
+
+import pytest
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+from vllm_omni.model_executor.models.moss_tts.reference_encoder_graphs import (
+    MossReferenceEncoderGraphs,
+    _cached_reference_rope,
+    _reference_rope_factors,
+    _validate_unmasked_reference_encoder,
+    _windowed_attention,
+    install_windowed_attention,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cuda]
+
+FRAME, N_VQ = 8, 3
+
+
+class _Stage(nn.Module):
+    """Causal per-frame stage: frames never mix, like a causal encoder layer."""
+
+    def __init__(self, dim_in, dim_out):
+        super().__init__()
+        self.proj = nn.Linear(dim_in, dim_out)
+
+    def forward(self, x, lengths):
+        return torch.tanh(self.proj(x)), lengths
+
+
+class _Quantizer(nn.Module):
+    def forward(self, hidden, lengths, n_q):
+        # hidden (B, frames, D): codes (n_q, B, frames) from per-frame argmax slices.
+        codes = torch.stack([hidden[..., q::n_q].argmax(-1) for q in range(n_q)])
+        return hidden, codes, lengths
+
+
+class _Tokenizer(nn.Module):
+    number_channels = 2
+    downsample_rate = FRAME
+    sampling_rate = 64
+
+    def __init__(self):
+        super().__init__()
+        self.encoder = nn.ModuleList([_Stage(FRAME * 2, 16), _Stage(16, 12)])
+        self.quantizer = _Quantizer()
+
+    def _flatten_channels_for_codec(self, audio, lengths):
+        batch, channels, samples = audio.shape
+        frames = samples // FRAME
+        values = audio.reshape(batch, channels, frames, FRAME).permute(0, 2, 1, 3).reshape(batch, frames, -1)
+        return values, lengths // FRAME
+
+    def _codec_inference_autocast(self):
+        return nullcontext()
+
+    def eager(self, wavs):
+        out = []
+        for wav in wavs:
+            values, lengths = self._flatten_channels_for_codec(wav[None], torch.tensor([wav.shape[-1]]))
+            for module in self.encoder:
+                values, lengths = module(values, lengths)
+            _, codes, _ = self.quantizer(values, lengths, N_VQ)
+            out.append(codes[:, 0, : wav.shape[-1] // FRAME].transpose(0, 1).long())
+        return out
+
+
+@pytest.fixture
+def cuda():
+    if not torch.cuda.is_available():
+        pytest.skip("CUDA required")
+    return torch.device("cuda", torch.accelerator.current_device_index())
+
+
+def test_cached_rope_reuses_longest_table_without_modifying_inputs(cuda):
+    from types import SimpleNamespace
+
+    module = SimpleNamespace(rope=SimpleNamespace(max_period=10000.0))
+    cache: dict[tuple[int, float, torch.device], tuple[torch.Tensor, torch.Tensor]] = {}
+    torch.manual_seed(17)
+    # Prime once; all subsequent sequences and batches use an immutable prefix.
+    q = torch.randn(1, 3, 128, 32, device=cuda, dtype=torch.bfloat16)
+    k = torch.randn_like(q)
+    _cached_reference_rope(module, q, k, cache)
+    stored = next(iter(cache.values()))
+    for batch, length in [(1, 17), (4, 64), (2, 128)]:
+        q = torch.randn(batch, 3, length, 32, device=cuda, dtype=torch.bfloat16)
+        k = torch.randn_like(q)
+        before = (q.clone(), k.clone())
+        got = _cached_reference_rope(module, q, k, cache)
+        cos, sin = _reference_rope_factors(length, 32, 10000.0, cuda)
+        torch.testing.assert_close(stored[0][:length], cos, rtol=0, atol=0)
+        torch.testing.assert_close(stored[1][:length], sin, rtol=0, atol=0)
+        for value, source, saved in zip(got, (q, k), before):
+            real, imag = source.float().reshape(batch, 3, length, 16, 2).unbind(-1)
+            expected = torch.stack((real * cos - imag * sin, real * sin + imag * cos), -1).to(source.dtype)
+            torch.testing.assert_close(value, expected.reshape_as(source), rtol=0, atol=0)
+            torch.testing.assert_close(source, saved, rtol=0, atol=0)
+        assert len(cache) == 1
+        assert next(iter(cache.values())) is stored
+
+
+def test_cached_rope_distinguishes_period_and_head_width(cuda):
+    from types import SimpleNamespace
+
+    cache: dict[tuple[int, float, torch.device], tuple[torch.Tensor, torch.Tensor]] = {}
+    for width, period in [(32, 10000.0), (64, 10000.0), (32, 1000.0)]:
+        module = SimpleNamespace(rope=SimpleNamespace(max_period=period))
+        q = torch.randn(1, 2, 32, width, device=cuda, dtype=torch.bfloat16)
+        _cached_reference_rope(module, q, q, cache)
+    assert len(cache) == 3
+
+
+@pytest.mark.parametrize("batched_transfer", [False, True])
+def test_graph_codes_match_eager_per_clip_and_fall_back_outside_buckets(cuda, batched_transfer):
+    torch.manual_seed(0)
+    tokenizer = _Tokenizer().to(cuda)
+    # Buckets of 2 s and 4 s at 64 Hz: 128 and 256 samples.
+    graphs = MossReferenceEncoderGraphs(
+        tokenizer, n_vq=N_VQ, batch_sizes=(1, 2, 4), bucket_seconds=(2.0, 4.0), batched_transfer=batched_transfer
+    )
+    graphs.capture()
+    assert graphs.captured == [(1, 128), (1, 256), (2, 128), (2, 256), (4, 128), (4, 256)]
+    assert graphs.max_batch == 4
+
+    for count, frames in [(1, [5]), (3, [16, 3, 9]), (4, [32, 1, 20, 17]), (2, [16, 16])]:
+        wavs = [torch.randn(2, n * FRAME) for n in frames]
+        codes = graphs.encode(wavs)
+        expected = tokenizer.eager([w.to(cuda) for w in wavs])
+        assert len(codes) == count
+        for got, want in zip(codes, expected):
+            assert got.device.type == "cpu" and got.dtype == torch.long
+            torch.testing.assert_close(got, want.cpu(), rtol=0, atol=0)
+
+    # Longer than every bucket: left to the caller; the rest still use graphs.
+    long_clip, short_clip = torch.randn(2, 33 * FRAME), torch.randn(2, 4 * FRAME)
+    codes = graphs.encode([long_clip, short_clip])
+    assert codes[0] is None
+    torch.testing.assert_close(codes[1], tokenizer.eager([short_clip.to(cuda)])[0].cpu(), rtol=0, atol=0)
+
+    # Wider than every batch: split, longest first, and returned in input order.
+    wavs = [torch.randn(2, n * FRAME) for n in (3, 30, 7, 12, 1)]
+    for got, want in zip(graphs.encode(wavs), tokenizer.eager([w.to(cuda) for w in wavs])):
+        torch.testing.assert_close(got, want.cpu(), rtol=0, atol=0)
+
+    # Result ownership includes the degenerate single-frame shape. Reusing
+    # the pinned output slab must not overwrite a previous caller's codes.
+    previous = graphs.encode([torch.randn(2, FRAME)])[0]
+    snapshot = previous.clone()
+    graphs.encode([torch.randn(2, 16 * FRAME) for _ in range(4)])
+    torch.testing.assert_close(previous, snapshot, rtol=0, atol=0)
+
+
+def test_batch_audio_budget_limits_captured_graphs(cuda):
+    tokenizer = _Tokenizer().to(cuda)
+    graphs = MossReferenceEncoderGraphs(
+        tokenizer, n_vq=N_VQ, batch_sizes=(1, 2, 4), bucket_seconds=(2.0, 4.0), max_batch_seconds=8.0
+    )
+    graphs.capture()
+    # 4 x 4 s exceeds the 8 s budget; a single clip is always captured.
+    assert graphs.captured == [(1, 128), (1, 256), (2, 128), (2, 256), (4, 128)]
+    wavs = [torch.randn(2, 30 * FRAME) for _ in range(3)]
+    for got, want in zip(graphs.encode(wavs), tokenizer.eager([w.to(cuda) for w in wavs])):
+        torch.testing.assert_close(got, want.cpu(), rtol=0, atol=0)
+
+
+def test_finer_singleton_buckets_preserve_multiclip_batching(cuda):
+    tokenizer = _Tokenizer().to(cuda)
+    graphs = MossReferenceEncoderGraphs(
+        tokenizer,
+        n_vq=N_VQ,
+        batch_sizes=(1, 2, 4),
+        bucket_seconds=(2.0, 4.0),
+        singleton_bucket_seconds=(1.0, 3.0),
+    )
+    graphs.capture()
+    assert (1, 64) in graphs.captured and (1, 192) in graphs.captured
+    assert (2, 64) not in graphs.captured and (4, 192) not in graphs.captured
+    replay = graphs._replay
+    shapes = []
+
+    def record(entry, wavs):
+        shapes.append(tuple(entry.audio.shape))
+        return replay(entry, wavs)
+
+    graphs._replay = record
+    one = torch.randn(2, 20 * FRAME)
+    out = graphs.encode([one])
+    assert shapes == [(1, 2, 192)]
+    torch.testing.assert_close(out[0], tokenizer.eager([one.to(cuda)])[0].cpu(), rtol=0, atol=0)
+    shapes.clear()
+    wavs = [one, one[:, : 8 * FRAME], one[:, : 12 * FRAME]]
+    out = graphs.encode(wavs)
+    assert shapes == [(4, 2, 256)]
+    for got, want in zip(out, tokenizer.eager([w.to(cuda) for w in wavs])):
+        torch.testing.assert_close(got, want.cpu(), rtol=0, atol=0)
+
+
+def test_compiled_capture_matches_eager(cuda):
+    torch.manual_seed(2)
+    tokenizer = _Tokenizer().to(cuda)
+    graphs = MossReferenceEncoderGraphs(
+        tokenizer, n_vq=N_VQ, batch_sizes=(1, 4), bucket_seconds=(2.0, 4.0), compile_core=True
+    )
+    graphs.capture()
+    assert graphs.captured == [(1, 128), (1, 256), (4, 128), (4, 256)]
+    wavs = [torch.randn(2, n * FRAME) for n in (5, 30, 12)]
+    for got, want in zip(graphs.encode(wavs), tokenizer.eager([w.to(cuda) for w in wavs])):
+        torch.testing.assert_close(got, want.cpu(), rtol=0, atol=0)
+
+
+class _Attention(nn.Module):
+    """The attention surface the tokenizer's MHA exposes to the windowed kernel."""
+
+    def __init__(self, embed_dim=128, heads=2, context=5):
+        super().__init__()
+        self.embed_dim, self.num_heads, self.context, self.causal = embed_dim, heads, context, True
+        self.in_proj = nn.Linear(embed_dim, 3 * embed_dim, bias=False)
+
+    def _project_qkv(self, x):
+        b, t, _ = x.shape
+        return self.in_proj(x).reshape(b, t, 3, self.num_heads, -1).permute(2, 0, 3, 1, 4).unbind(0)
+
+    def _apply_dense_rope(self, q, k):
+        return q, k
+
+    def _forward_non_streaming_sdpa(self, x, input_lengths):
+        # Reference semantics: causal, distance < context, padded rows zeroed.
+        b, t, _ = x.shape
+        q, k, v = self._project_qkv(x)
+        pos = torch.arange(t, device=x.device)
+        delta = pos.view(-1, 1) - pos.view(1, -1)
+        mask = (delta >= 0) & (delta < self.context)
+        out = F.scaled_dot_product_attention(q, k, v, mask[None, None])
+        valid = (pos.view(1, t) < input_lengths.view(-1, 1)).view(b, 1, t, 1)
+        out = torch.where(valid, out, torch.zeros((), device=out.device, dtype=out.dtype))
+        return out.transpose(1, 2).reshape(b, t, self.embed_dim)
+
+
+@pytest.mark.parametrize("fa_version", [2, 3])
+def test_windowed_attention_matches_masked_sdpa(cuda, fa_version):
+    if fa_version == 3 and torch.cuda.get_device_capability(cuda)[0] != 9:
+        pytest.skip("FA3 requires Hopper")
+    torch.manual_seed(1)
+    model = nn.Module()
+    model.encoder = nn.ModuleList([_Attention().to(cuda, torch.bfloat16)])
+    attention = model.encoder[0]
+    x = torch.randn(3, 40, 128, device=cuda, dtype=torch.bfloat16)
+    lengths = torch.tensor([40, 17, 3], device=cuda)
+    expected = attention._forward_non_streaming_sdpa(x, lengths)
+    assert install_windowed_attention(model, fa_version=fa_version) == 1
+    actual = attention._forward_non_streaming_sdpa(x, lengths)
+    torch.testing.assert_close(actual, expected, rtol=2e-2, atol=2e-2)
+    assert actual[1, 17:].count_nonzero() == 0 and actual[2, 3:].count_nonzero() == 0
+    # Non-bf16 inputs keep the tokenizer's own SDPA path.
+    x32 = x.float()
+    attention.float()
+    torch.testing.assert_close(
+        attention._forward_non_streaming_sdpa(x32, lengths), expected.float(), rtol=3e-2, atol=3e-2
+    )
+
+
+@pytest.mark.parametrize("fa_version", [2, 3])
+def test_unmasked_causal_queries_preserve_valid_prefix_and_ignore_future_padding(cuda, fa_version):
+    if fa_version == 3 and torch.cuda.get_device_capability(cuda)[0] != 9:
+        pytest.skip("FA3 requires Hopper")
+    torch.manual_seed(44)
+    attention = _Attention().to(cuda, torch.bfloat16)
+    x = torch.randn(4, 40, 128, device=cuda, dtype=torch.bfloat16)
+    lengths = torch.tensor([40, 17, 1, 0], device=cuda)
+    kwargs = dict(sdpa=attention._forward_non_streaming_sdpa, fa_version=fa_version)
+    with torch.no_grad():
+        masked = _windowed_attention(attention, x, lengths, **kwargs)
+        unmasked = _windowed_attention(attention, x, lengths, skip_padded_query_mask=True, **kwargs)
+        poisoned = x.clone()
+        for row, length in enumerate([40, 17, 1, 0]):
+            poisoned[row, length:] = torch.randn_like(poisoned[row, length:]) * 30
+        changed = _windowed_attention(attention, poisoned, lengths, skip_padded_query_mask=True, **kwargs)
+        for row, length in enumerate([40, 17, 1, 0]):
+            torch.testing.assert_close(unmasked[row, :length], masked[row, :length], rtol=0, atol=0)
+            torch.testing.assert_close(changed[row, :length], masked[row, :length], rtol=0, atol=0)
+            assert masked[row, length:].count_nonzero() == 0
+        assert unmasked[1, 17:].count_nonzero() > 0
+        # The FP32 fallback retains its original masking contract.
+        attention.float()
+        got = _windowed_attention(attention, x.float(), lengths, skip_padded_query_mask=True, **kwargs)
+        expected = attention._forward_non_streaming_sdpa(x.float(), lengths)
+        torch.testing.assert_close(got, expected, rtol=0, atol=0)
+
+
+def test_unmasked_queries_reject_unknown_encoder_before_patching():
+    tokenizer = _Tokenizer()
+    with pytest.raises(ValueError, match="patch-downsample"):
+        install_windowed_attention(tokenizer, skip_padded_query_mask=True)
+
+
+@pytest.mark.parametrize("bad_part", ["upsample", "noncausal", "quantizer", "convolution"])
+def test_unmasked_queries_reject_structures_that_can_read_invalid_rows(bad_part):
+    from types import SimpleNamespace
+
+    patch = type("MossAudioTokenizerPatchedPretransform", (), {})()
+    patch.is_downsample = bad_part != "upsample"
+    projected = type("MossAudioTokenizerProjectedTransformer", (), {})()
+    projected.transformer = SimpleNamespace(
+        layers=[SimpleNamespace(self_attn=SimpleNamespace(causal=bad_part != "noncausal"))]
+    )
+    quantizer_cls = type(
+        "MossAudioTokenizerResidualVQ" if bad_part != "quantizer" else "UnknownQuantizer", (nn.Module,), {}
+    )
+    quantizer = quantizer_cls()
+    quantizer.proj = nn.Conv1d(2, 2, 3 if bad_part == "convolution" else 1)
+    tokenizer = SimpleNamespace(encoder=[patch, projected], quantizer=quantizer)
+    with pytest.raises(ValueError):
+        _validate_unmasked_reference_encoder(tokenizer)

tests/model_executor/models/moss_tts/test_shared_reference_encoder.py

@@ -0,0 +1,341 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""API processes share one reference encoder through a Unix socket."""
+
+import threading
+from tempfile import TemporaryDirectory
+
+import pytest
+import torch
+
+from vllm_omni.model_executor.models.moss_tts import shared_reference_encoder as shared
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def _codes(wav):
+    # Deterministic per clip: frames from the clip length, values from its first sample.
+    return torch.full((wav.shape[-1] // 10, 4), int(wav.reshape(-1)[0]), dtype=torch.long)
+
+
+class _Recorder:
+    def __init__(self, gate=None):
+        self.batches = []
+        self.gate = gate
+
+    def __call__(self, wavs):
+        if self.gate is not None:
+            self.gate.wait(5)
+        self.batches.append(len(wavs))
+        return [RuntimeError("bad clip") if int(w.reshape(-1)[0]) < 0 else _codes(w) for w in wavs]
+
+
+@pytest.fixture
+def host_dir(monkeypatch):
+    monkeypatch.setattr(shared, "_host_lock_handle", None)
+    # AF_UNIX limits the full socket path, including pytest directory names.
+    with TemporaryDirectory(prefix="moss-ref-") as path:
+        yield path
+
+
+def test_only_one_process_per_directory_hosts(host_dir):
+    assert shared.elect_host(host_dir) is True
+    assert shared.elect_host(host_dir) is False  # a second open file description cannot take the lock
+
+
+def test_close_releases_queued_and_active_callers(host_dir):
+    entered, release = threading.Event(), threading.Event()
+
+    def encode(wavs):
+        entered.set()
+        assert release.wait(5)
+        return [_codes(wav) for wav in wavs]
+
+    host = shared.SharedReferenceEncoderHost(host_dir, [(encode, None)])
+    results: dict[str, list[torch.Tensor | BaseException]] = {}
+    threads = [
+        threading.Thread(target=lambda key=key: results.setdefault(key, host.encode([torch.ones(1, 20)])), daemon=True)
+        for key in ("active", "queued")
+    ]
+    try:
+        threads[0].start()
+        assert entered.wait(2)
+        threads[1].start()
+        with host._ready:
+            assert host._ready.wait_for(lambda: len(host._pending) == 1, timeout=2)
+        host.close()
+        for thread in threads:
+            thread.join(1)
+            assert not thread.is_alive(), "closed encoder left a request waiting"
+        assert all(isinstance(results[key][0], RuntimeError) for key in results)
+    finally:
+        release.set()
+        host.close()
+        for thread in threads:
+            if thread.ident is not None:
+                thread.join(1)
+
+
+def test_worker_result_count_mismatch_fails_entire_batch(host_dir):
+    host = shared.SharedReferenceEncoderHost(host_dir, [(lambda wavs: [_codes(wavs[0])], None)])
+    results = []
+    caller = threading.Thread(target=lambda: results.extend(host.encode([torch.ones(1, 20)] * 2)), daemon=True)
+    try:
+        caller.start()
+        caller.join(2)
+        assert not caller.is_alive(), "missing encoder result left a request waiting"
+        assert len(results) == 2
+        assert all(isinstance(result, RuntimeError) for result in results)
+    finally:
+        host.close()
+        caller.join(1)
+
+
+def test_client_round_trip_and_per_clip_errors(host_dir):
+    host = shared.SharedReferenceEncoderHost(host_dir, [(_Recorder(), None)])
+    try:
+        client = shared.SharedReferenceEncoderClient(host_dir)
+        wavs = [torch.full((1, 40), 3.0), torch.full((1, 70), -1.0), torch.full((1, 20), 5.0)]
+        out = client.encode(wavs)
+        assert out[0].dtype == torch.long and torch.equal(out[0], _codes(wavs[0]))
+        assert isinstance(out[1], RuntimeError) and "bad clip" in str(out[1])
+        assert torch.equal(out[2], _codes(wavs[2]))
+        # The host process's own clips skip the socket.
+        assert torch.equal(host.encode([wavs[2]])[0], _codes(wavs[2]))
+    finally:
+        host.close()
+
+
+def test_host_merges_clips_from_concurrent_clients(host_dir):
+    gate = threading.Event()
+    recorder = _Recorder(gate)
+    host = shared.SharedReferenceEncoderHost(host_dir, [(recorder, None)])
+    try:
+        results = {}
+
+        def send(name, value):
+            results[name] = shared.SharedReferenceEncoderClient(host_dir).encode([torch.full((1, 30), value)] * 2)
+
+        # Encoding is held until both clients' clips are queued.
+        warm = threading.Thread(target=lambda: host.encode([torch.full((1, 10), 1.0)]))
+        warm.start()
+        threads = [threading.Thread(target=send, args=(n, v)) for n, v in (("a", 2.0), ("b", 4.0))]
+        for thread in threads:
+            thread.start()
+        threading.Timer(0.5, gate.set).start()
+        for thread in (warm, *threads):
+            thread.join(10)
+        # Both clients' clips were encoded together, not one request at a time.
+        assert sum(recorder.batches) == 5 and max(recorder.batches) >= 4
+        assert all(int(c[0, 0]) == 2 for c in results["a"]) and all(int(c[0, 0]) == 4 for c in results["b"])
+    finally:
+        host.close()
+
+
+def test_client_without_host_raises(host_dir, monkeypatch):
+    monkeypatch.setattr(shared, "_CONNECT_TIMEOUT_S", 0.3)
+    with pytest.raises((FileNotFoundError, ConnectionRefusedError)):
+        shared.SharedReferenceEncoderClient(host_dir).encode([torch.zeros(1, 10)])
+
+
+def test_host_window_collects_across_workers_without_reserving_singletons(host_dir):
+    recorder = _Recorder()
+    host = shared.SharedReferenceEncoderHost(host_dir, [(recorder, None)] * 3, batch_window_ms=500)
+    try:
+        barrier = threading.Barrier(3)
+        results = []
+
+        def send(value):
+            barrier.wait()
+            results.extend(host.encode([torch.full((1, 20), value)]))
+
+        threads = [threading.Thread(target=send, args=(v,)) for v in (1.0, 2.0)]
+        for thread in threads:
+            thread.start()
+        barrier.wait()
+        for thread in threads:
+            thread.join(5)
+            assert not thread.is_alive()
+        assert recorder.batches == [2]
+        assert sorted(int(x[0, 0]) for x in results) == [1, 2]
+    finally:
+        host.close()
+
+
+def test_reference_encoder_client_falls_back_to_local_encode(host_dir, monkeypatch):
+    from tests.model_executor.models.moss_tts.test_reference_encoder import _N_VQ, _SR, _FakeProcessor
+    from vllm_omni.model_executor.models.moss_tts import reference_encoder as re_mod
+    from vllm_omni.utils.speaker_cache import SpeakerEmbeddingCache
+
+    monkeypatch.setenv(re_mod._SHARED_CODES_DIR_ENV, host_dir)
+    monkeypatch.setattr(re_mod, "shared_encoder_role", lambda: "client")
+    monkeypatch.setattr(shared, "_CONNECT_TIMEOUT_S", 0.3)
+    proc = _FakeProcessor()
+    enc = re_mod.MossReferenceEncoder(
+        proc, variant="local", n_vq=_N_VQ, sr_target=_SR, speaker_cache=SpeakerEmbeddingCache(max_bytes=1 << 20)
+    )
+    placed = []
+    monkeypatch.setattr(enc, "_place_tokenizer_for_local_encode", lambda: placed.append(True))
+    out = enc._encode_batch_sync([([7.0] * 50, _SR)])
+    assert placed == [True] and proc.attempt_sizes == [1] and int(out[0][0, 0]) == 7
+
+
+def test_workers_encode_concurrently_after_their_warmups(host_dir):
+    inside, release = threading.Barrier(3, timeout=5), threading.Event()
+    warmed = []
+
+    def worker(name):
+        def encode(wavs):
+            inside.wait()  # both workers are inside an encode at the same time
+            release.wait(5)
+            return [_codes(w) for w in wavs]
+
+        return encode, lambda: warmed.append(name)
+
+    host = shared.SharedReferenceEncoderHost(host_dir, [worker("a"), worker("b")])
+    try:
+        out = []
+
+        def submit(value):
+            out.append(host.encode([torch.full((1, 20), value)]))
+
+        first = threading.Thread(target=submit, args=(1.0,))
+        first.start()
+        # The first worker is busy by now, so the second clip goes to the other one.
+        threading.Event().wait(0.1)
+        second = threading.Thread(target=submit, args=(2.0,))
+        second.start()
+        inside.wait()
+        release.set()
+        for thread in (first, second):
+            thread.join(5)
+        assert sorted(warmed) == ["a", "b"] and sorted(int(o[0][0, 0]) for o in out) == [1, 2]
+    finally:
+        host.close()
+
+
+def test_client_requests_in_flight_use_their_own_connections(host_dir):
+    entered, release = threading.Semaphore(0), threading.Event()
+
+    def worker():
+        def encode(wavs):
+            entered.release()
+            release.wait(5)
+            return [_codes(w) for w in wavs]
+
+        return encode, None
+
+    host = shared.SharedReferenceEncoderHost(host_dir, [worker(), worker()])
+    try:
+        client = shared.SharedReferenceEncoderClient(host_dir)
+        out = []
+        threads = [
+            threading.Thread(target=lambda v=v: out.append(client.encode([torch.full((1, 20), v)]))) for v in (1.0, 2.0)
+        ]
+        for thread in threads:
+            thread.start()
+            # A worker holds this request, still unanswered, before the next is
+            # sent. Sent together, one free worker may take both as one batch.
+            assert entered.acquire(timeout=5)
+        release.set()  # both requests reached the host at once
+        for thread in threads:
+            thread.join(5)
+        assert sorted(int(o[0][0, 0]) for o in out) == [1, 2]
+        assert len(client._idle) == 2  # both connections are kept for reuse
+    finally:
+        host.close()
+
+
+def test_readiness_waits_for_every_worker_warmup(host_dir, monkeypatch):
+    entered, release = threading.Event(), threading.Event()
+    recorder = _Recorder()
+    warmed = []
+
+    def last_warmup():
+        entered.set()
+        assert release.wait(5)
+        warmed.append("last")
+
+    host = shared.SharedReferenceEncoderHost(
+        host_dir, [(recorder, lambda: warmed.append("first")), (recorder, last_warmup)]
+    )
+    client = shared.SharedReferenceEncoderClient(host_dir)
+    try:
+        assert entered.wait(5)
+        monkeypatch.setattr(shared, "_STARTUP_TIMEOUT_S", 0.03)
+        with pytest.raises(shared.SharedReferenceEncoderStartupError):
+            host.wait_until_ready()
+        with pytest.raises(shared.SharedReferenceEncoderStartupError):
+            client.wait_until_ready()
+        assert warmed == ["first"] and recorder.batches == [] and not client._idle
+        release.set()
+        monkeypatch.setattr(shared, "_STARTUP_TIMEOUT_S", 5)
+        host.wait_until_ready()
+        client.wait_until_ready()
+        assert warmed == ["first", "last"] and recorder.batches == []
+        assert len(client._idle) == 1
+        assert torch.equal(client.encode([torch.ones(1, 20)])[0], _codes(torch.ones(1, 20)))
+    finally:
+        release.set()
+        host.close()
+
+
+def test_slow_startup_does_not_use_encode_timeout(host_dir, monkeypatch):
+    release = threading.Event()
+    monkeypatch.setattr(shared, "_REQUEST_TIMEOUT_S", 0.05)
+    monkeypatch.setattr(shared, "_STARTUP_TIMEOUT_S", 5, raising=False)
+    host = shared.SharedReferenceEncoderHost(host_dir, [(_Recorder(), lambda: release.wait(5))])
+    timer = threading.Timer(0.2, release.set)
+    timer.start()
+    try:
+        # The first socket waits for readiness before starting its much shorter
+        # encode deadline. No API warmup call is required to protect this path.
+        client = shared.SharedReferenceEncoderClient(host_dir)
+        assert torch.equal(client.encode([torch.ones(1, 20)])[0], _codes(torch.ones(1, 20)))
+    finally:
+        release.set()
+        timer.join(5)
+        host.close()
+
+
+def test_failed_host_warmup_is_not_ready(host_dir):
+    def fail():
+        raise RuntimeError("warmup failed")
+
+    recorder = _Recorder()
+    host = shared.SharedReferenceEncoderHost(host_dir, [(recorder, fail)])
+    try:
+        with pytest.raises(shared.SharedReferenceEncoderStartupError):
+            host.wait_until_ready()
+        with pytest.raises(shared.SharedReferenceEncoderStartupError):
+            shared.SharedReferenceEncoderClient(host_dir).wait_until_ready()
+        with pytest.raises(shared.SharedReferenceEncoderStartupError):
+            host.encode([torch.ones(1, 20)])
+        assert recorder.batches == []
+    finally:
+        host.close()
+
+
+def test_startup_failure_does_not_trigger_local_capture(host_dir, monkeypatch):
+    from unittest.mock import Mock
+
+    from tests.model_executor.models.moss_tts.test_reference_encoder import _N_VQ, _SR, _FakeProcessor
+    from vllm_omni.model_executor.models.moss_tts import reference_encoder as re_mod
+    from vllm_omni.utils.speaker_cache import SpeakerEmbeddingCache
+
+    monkeypatch.setenv(re_mod._SHARED_CODES_DIR_ENV, host_dir)
+    monkeypatch.setattr(re_mod, "shared_encoder_role", lambda: "client")
+    enc = re_mod.MossReferenceEncoder(
+        _FakeProcessor(),
+        variant="local",
+        n_vq=_N_VQ,
+        sr_target=_SR,
+        speaker_cache=SpeakerEmbeddingCache(max_bytes=1 << 20),
+    )
+    enc._shared_client = Mock()
+    enc._shared_client.encode.side_effect = shared.SharedReferenceEncoderStartupError("still compiling")
+    local = Mock()
+    monkeypatch.setattr(enc, "_place_tokenizer_for_local_encode", local)
+    with pytest.raises(shared.SharedReferenceEncoderStartupError):
+        enc._encode_shared_or_local([torch.ones(1, 20)])
+    local.assert_not_called()

tests/model_executor/models/moss_tts/test_slot_attention.py

@@ -0,0 +1,301 @@
+# 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,
+    MossAudioTokenizerRotaryEmbedding,
+    RingKVCache,
+    StreamingExecutionContext,
+)
+from vllm_omni.model_executor.models.moss_tts.slot_attention import slot_ring_attention, slot_ring_attention_rows
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cuda]
+
+
+@pytest.mark.parametrize(
+    "frames,capacity", [(1, 125), (2, 250), (15, 125), (30, 250), (32, 400), (60, 400), (480, 400)]
+)
+@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
+@pytest.mark.parametrize("execution", ["graph", "compiled"])
+def test_empty_tile_skip_matches_static_ring_across_reuse(frames, capacity, dtype, execution):
+    torch.manual_seed(45)
+    batch, heads, states, dim = 3, 2, 5, 64
+    projected = torch.randn(batch, frames, 3, heads, dim, device="cuda", dtype=dtype)
+    q, k, v = projected.permute(2, 0, 3, 1, 4).unbind(0)
+    cache = torch.randn(2, states, heads, capacity, dim, device="cuda", dtype=dtype)
+    expected_cache = cache.clone()
+    end = torch.zeros(states, 2, device="cuda", dtype=torch.long)[:, 0]
+    expected_end = end.clone()
+    slots = torch.tensor([0, 4, 2], device="cuda")
+    lengths = torch.full((batch, 2), frames, device="cuda", dtype=torch.int32)[:, 0]
+
+    def candidate():
+        return slot_ring_attention(q, k, v, cache, end, slots, lengths, 17, skip_empty_tiles=True)
+
+    graph = None
+    if execution == "compiled":
+        run = torch.compile(candidate, fullgraph=True)
+    else:
+        stream = torch.cuda.Stream()
+        stream.wait_stream(torch.cuda.current_stream())
+        with torch.cuda.stream(stream):
+            candidate()
+        torch.cuda.current_stream().wait_stream(stream)
+        graph = torch.cuda.CUDAGraph()
+        with torch.cuda.graph(graph, stream=stream):
+            actual = candidate()
+
+    for step, start in enumerate([0, 1, capacity - 1, capacity + 9, 0, 3 * capacity]):
+        end.fill_(start)
+        expected_end.fill_(start)
+        cache.copy_(expected_cache)
+        projected.normal_()
+        slots.copy_(torch.roll(torch.arange(states, device="cuda"), step)[:batch])
+        lengths.copy_(torch.tensor([frames, max(1, frames - 3), 0], device="cuda", dtype=torch.int32))
+        expected = slot_ring_attention(q, k, v, expected_cache, expected_end, slots, lengths, 17)
+        if graph is not None:
+            graph.replay()
+        else:
+            actual = run()
+        torch.testing.assert_close(actual, expected, rtol=0, atol=0)
+        torch.testing.assert_close(cache, expected_cache, rtol=0, atol=0)
+        torch.testing.assert_close(end, expected_end, rtol=0, atol=0)
+
+
+@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
+@pytest.mark.parametrize("frames,capacity", [(1, 125), (15, 140), (480, 400), (480, 880)])
+def test_all_padding_preserves_state_and_replay_can_become_active(dtype, frames, capacity):
+    q, k, v = [torch.randn(2, 2, frames, 64, device="cuda", dtype=dtype) for _ in range(3)]
+    cache = torch.randn(2, 4, 2, capacity, 64, device="cuda", dtype=dtype)
+    end = torch.tensor([0, capacity + 3, 9, 0], device="cuda")
+    slots = torch.tensor([1, 3], device="cuda")
+    lengths = torch.zeros(2, device="cuda", dtype=torch.int32)
+    before_cache, before_end = cache.clone(), end.clone()
+    stream = torch.cuda.Stream()
+    stream.wait_stream(torch.cuda.current_stream())
+    with torch.cuda.stream(stream):
+        for _ in range(3):
+            slot_ring_attention(q, k, v, cache, end, slots, lengths, capacity)
+    torch.cuda.current_stream().wait_stream(stream)
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph):
+        out = slot_ring_attention(q, k, v, cache, end, slots, lengths, capacity)
+    for active in [False, True, False, True]:
+        lengths.fill_(frames if active else 0)
+        expected = _ring_reference(q, k, v, before_cache, before_end, slots, lengths, capacity)
+        graph.replay()
+        torch.testing.assert_close(cache, before_cache, rtol=0, atol=0)
+        torch.testing.assert_close(end, before_end, rtol=0, atol=0)
+        torch.testing.assert_close(out, expected, rtol=0.025, atol=0.025)
+        if not active:
+            assert out.count_nonzero() == 0
+
+
+def test_dense_ring_retains_last_tokens_when_chunk_exceeds_capacity():
+    torch.manual_seed(19)
+    capacity, frames = 400, 480
+    state = RingKVCache(3, 2, 64, capacity, device=torch.device("cuda"), dtype=torch.float32)
+    slots = torch.tensor([1, 0], device="cuda")
+    context = StreamingExecutionContext(state_slot_ids=slots, valid_rows=torch.ones(2, device="cuda", dtype=torch.bool))
+    for start in [0, frames]:
+        k = torch.randn(2, 2, frames, 64, device="cuda")
+        v = torch.randn_like(k)
+        state.complete(k, v, execution_context=context)
+        retained = torch.arange(frames - capacity, frames, device="cuda")
+        destinations = (start + retained) % capacity
+        torch.testing.assert_close(state.cache[0, slots][:, :, destinations], k[:, :, retained], rtol=0, atol=0)
+        torch.testing.assert_close(state.cache[1, slots][:, :, destinations], v[:, :, retained], rtol=0, atol=0)
+        torch.testing.assert_close(state.end_offset[slots], torch.full_like(slots, start + frames), rtol=0, atol=0)
+
+
+def _ring_reference(q, k, v, cache, end, slots, lengths, context):
+    """Independent chronological ring oracle with FP32 attention arithmetic."""
+    capacity = cache.shape[-2]
+    output = torch.zeros_like(q)
+    for row, (slot, length) in enumerate(zip(slots.tolist(), lengths.tolist())):
+        if length == 0:
+            continue
+        start = int(end[slot])
+        for token in range(max(0, length - capacity), length):
+            cache[0, slot, :, (start + token) % capacity] = k[row, :, token]
+            cache[1, slot, :, (start + token) % capacity] = v[row, :, token]
+        stop = start + length
+        physical_positions = torch.full((capacity,), -1, device=q.device, dtype=torch.long)
+        available = torch.arange(max(0, stop - capacity), stop, device=q.device)
+        physical_positions[available % capacity] = available
+        queries = torch.arange(start, stop, device=q.device)
+        distance = queries[:, None] - physical_positions[None]
+        mask = (physical_positions[None] >= 0) & (distance >= 0)
+        if context > 0:
+            mask &= distance < context
+        output[row, :, :length] = F.scaled_dot_product_attention(
+            q[row, :, :length].float(), cache[0, slot].float(), cache[1, slot].float(), mask[None]
+        ).to(q.dtype)
+        end[slot] = stop
+    return output
+
+
+@pytest.mark.parametrize("frames,capacity,head_dim", [(1, 125, 64), (15, 125, 64), (480, 400, 64), (33, 65, 32)])
+@pytest.mark.parametrize("execution", ["eager", "graph", "compiled"])
+def test_slot_attention_preserves_ring_state_and_padding(frames, capacity, head_dim, execution):
+    torch.manual_seed(29)
+    batch, heads, state_capacity = 3, 2, 5
+    projected = torch.randn(batch, frames, 3, heads, head_dim, device="cuda", dtype=torch.bfloat16)
+    q, k, v = projected.permute(2, 0, 3, 1, 4).unbind(0)
+    cache = torch.randn(2, state_capacity, heads, capacity, head_dim, device="cuda", dtype=q.dtype)
+    # Strided metadata is allowed by the kernel contract.
+    end = torch.zeros(state_capacity, 2, device="cuda", dtype=torch.long)[:, 0]
+    slots = torch.arange(batch * 2, device="cuda")[::2]
+    lengths = torch.full((batch, 2), frames, device="cuda", dtype=torch.int32)[:, 0]
+    context = capacity
+    attention = slot_ring_attention
+    if execution == "compiled":
+        # Functionalization must preserve mutations of the persistent ring
+        # and its offsets between calls, including slot reuse.
+        attention = torch.compile(slot_ring_attention, backend="inductor", fullgraph=True)
+    if execution == "graph":
+        stream = torch.cuda.Stream()
+        stream.wait_stream(torch.cuda.current_stream())
+        with torch.cuda.stream(stream):
+            for _ in range(3):
+                slot_ring_attention(q, k, v, cache, end, slots, lengths, context)
+        torch.cuda.current_stream().wait_stream(stream)
+        graph = torch.cuda.CUDAGraph()
+        with torch.cuda.graph(graph):
+            actual = slot_ring_attention(q, k, v, cache, end, slots, lengths, context)
+    end.zero_()
+    expected_cache, expected_end = cache.clone(), end.clone()
+    for step in range(8):
+        slots.copy_(torch.roll(torch.arange(state_capacity, device="cuda"), step)[:batch])
+        lengths.copy_(torch.tensor([frames, max(1, frames - 3), 0], device="cuda", dtype=torch.int32))
+        projected.normal_()
+        if step == 4:
+            end[1:3].zero_()
+            expected_end[1:3].zero_()
+        expected = _ring_reference(q, k, v, expected_cache, expected_end, slots, lengths, context)
+        if execution == "graph":
+            graph.replay()
+        else:
+            actual = attention(q, k, v, cache, end, slots, lengths, context)
+        torch.testing.assert_close(cache, expected_cache, rtol=0, atol=0)
+        torch.testing.assert_close(end, expected_end, rtol=0, atol=0)
+        torch.testing.assert_close(actual, expected, rtol=0.025, atol=0.025)
+        padding = torch.arange(frames, device="cuda")[None] >= lengths[:, None]
+        assert actual.permute(0, 2, 1, 3)[padding].count_nonzero() == 0
+
+
+@pytest.mark.parametrize("frames,context", [(1, 9), (3, 9), (15, 125), (480, 400)])
+def test_slot_attention_integrates_with_current_mha(frames, context):
+    torch.manual_seed(9)
+    reference = MossAudioTokenizerMultiheadAttention(
+        128,
+        2,
+        causal=True,
+        context=context,
+        rope=MossAudioTokenizerRotaryEmbedding(),
+        device="cuda",
+        dtype=torch.bfloat16,
+    )
+    candidate = copy.deepcopy(reference)
+    candidate._slot_attention = slot_ring_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")
+        execution_context = StreamingExecutionContext(state_slot_ids=slots, valid_rows=valid)
+        x = torch.randn(3, frames, 128, device="cuda", dtype=torch.bfloat16)
+        with torch.inference_mode():
+            expected = reference(x, x, x, execution_context=execution_context)
+            actual = candidate(x, x, x, execution_context=execution_context)
+        torch.testing.assert_close(actual, expected, atol=0.01, rtol=0.02)
+        cs, rs = candidate._streaming_state, reference._streaming_state
+        torch.testing.assert_close(cs.offset, rs.offset, rtol=0, atol=0)
+        torch.testing.assert_close(cs.kv_cache.end_offset, rs.kv_cache.end_offset, rtol=0, atol=0)
+        torch.testing.assert_close(cs.kv_cache.cache[:, :3], rs.kv_cache.cache[:, :3], rtol=0, atol=0)
+        if step == 4:
+            for model in (reference, candidate):
+                model._streaming_state.reset_slots(torch.tensor([0], device="cuda"))
+
+
+@pytest.mark.parametrize(
+    "frames,capacity,head_dim", [(1, 125, 64), (15, 140, 64), (40, 125, 64), (480, 400, 64), (33, 65, 32)]
+)
+@pytest.mark.parametrize("execution", ["eager", "graph", "compiled"])
+def test_fused_rows_attention_matches_three_launch_op_exactly(frames, capacity, head_dim, execution):
+    torch.manual_seed(31)
+    batch, heads, state_capacity = 3, 2, 6
+    projected = torch.randn(batch, frames, 3, heads, head_dim, device="cuda", dtype=torch.bfloat16)
+    q, k, v = projected.permute(2, 0, 3, 1, 4).unbind(0)
+    cache = torch.randn(2, state_capacity, heads, capacity, head_dim, device="cuda", dtype=q.dtype)
+    reference_cache = cache.clone()
+    end = torch.zeros(state_capacity, device="cuda", dtype=torch.long)
+    slots = torch.zeros(batch, device="cuda", dtype=torch.long)
+    lengths = torch.zeros(batch, device="cuda", dtype=torch.int32)
+    starts = torch.zeros(batch, device="cuda", dtype=torch.long)
+    attention = slot_ring_attention_rows
+    if execution == "compiled":
+        attention = torch.compile(slot_ring_attention_rows, backend="inductor", fullgraph=True)
+    if execution == "graph":
+        stream = torch.cuda.Stream()
+        stream.wait_stream(torch.cuda.current_stream())
+        with torch.cuda.stream(stream):
+            for _ in range(3):
+                slot_ring_attention_rows(q, k, v, cache, starts, slots, lengths, capacity)
+        torch.cuda.current_stream().wait_stream(stream)
+        graph = torch.cuda.CUDAGraph()
+        with torch.cuda.graph(graph):
+            actual = slot_ring_attention_rows(q, k, v, cache, starts, slots, lengths, capacity)
+        cache.copy_(reference_cache)
+    for step in range(8):
+        slots.copy_(torch.roll(torch.arange(state_capacity, device="cuda"), step)[:batch])
+        lengths.copy_(torch.tensor([frames, max(1, frames - 3), 0], device="cuda", dtype=torch.int32))
+        starts.copy_(end[slots])
+        projected.normal_()
+        expected = slot_ring_attention(q, k, v, reference_cache, end, slots, lengths, capacity)
+        if execution == "graph":
+            graph.replay()
+        else:
+            actual = attention(q, k, v, cache, starts, slots, lengths, capacity)
+        torch.testing.assert_close(actual, expected, rtol=0, atol=0)
+        torch.testing.assert_close(cache, reference_cache, rtol=0, atol=0)
+
+
+@pytest.mark.parametrize("frames,context", [(1, 9), (15, 125), (480, 400)])
+def test_fused_rows_attention_integrates_with_mha_offsets(frames, context):
+    torch.manual_seed(11)
+    reference = MossAudioTokenizerMultiheadAttention(
+        128,
+        2,
+        causal=True,
+        context=context,
+        rope=MossAudioTokenizerRotaryEmbedding(),
+        device="cuda",
+        dtype=torch.bfloat16,
+    )
+    candidate = copy.deepcopy(reference)
+    reference._slot_attention = slot_ring_attention
+    candidate._slot_attention_rows = slot_ring_attention_rows
+    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")
+        execution_context = StreamingExecutionContext(state_slot_ids=slots, valid_rows=valid)
+        x = torch.randn(3, frames, 128, device="cuda", dtype=torch.bfloat16)
+        with torch.inference_mode():
+            expected = reference(x, x, x, execution_context=execution_context)
+            actual = candidate(x, x, x, execution_context=execution_context)
+        torch.testing.assert_close(actual, expected, rtol=0, atol=0)
+        cs, rs = candidate._streaming_state, reference._streaming_state
+        torch.testing.assert_close(cs.offset, rs.offset, rtol=0, atol=0)
+        torch.testing.assert_close(cs.kv_cache.end_offset, rs.kv_cache.end_offset, rtol=0, atol=0)
+        torch.testing.assert_close(cs.kv_cache.cache, rs.kv_cache.cache, 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_decoder_wrapper_concurrency.py

@@ -0,0 +1,86 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""A second streaming-decoder graph set can replay concurrently with the first."""
+
+import threading
+
+import pytest
+import torch
+from torch import nn
+
+from vllm_omni.model_executor.models.moss_tts.cuda_graph_streaming_decoder_wrapper import (
+    CUDAGraphStreamingDecoderWrapper,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cuda]
+
+CAP, DIM = 8, 64
+
+
+class _Codec(nn.Module):
+    def __init__(self, device):
+        super().__init__()
+        self.weight = nn.Parameter(torch.randn(DIM, DIM, device=device) / DIM)
+        self.state = torch.zeros(CAP + 16, DIM, device=device)
+
+    def decode_streaming_tensors(self, codes, lengths, slots, valid_rows):
+        x = codes.float().sum(0)[..., None].expand(-1, -1, DIM) + self.state[slots][:, None]
+        y = x @ self.weight
+        self.state.index_copy_(0, slots, y[:, -1])
+        return y.transpose(1, 2).contiguous(), lengths
+
+    def reset_decoder_state_slots(self, slots):
+        self.state.index_fill_(0, slots.to(self.state.device), 0)
+
+
+def _wrapper(codec, **kwargs):
+    return CUDAGraphStreamingDecoderWrapper(
+        codec,
+        state_capacity=CAP,
+        num_quantizers=2,
+        vllm_config=None,
+        compiled_decode=codec.decode_streaming_tensors,
+        **kwargs,
+    )
+
+
+def test_private_wrapper_captures_on_its_own_stream_and_replays_concurrently():
+    if not torch.cuda.is_available():
+        pytest.skip("CUDA required")
+    device = torch.device("cuda", torch.accelerator.current_device_index())
+    codec = _Codec(device)
+    main = _wrapper(codec, batch_sizes=[4], frame_sizes=[3])
+    fast = _wrapper(codec, batch_sizes=[1, 2], frame_sizes=[1], scratch_base=CAP + 4, private_pool=True)
+    main.warmup(device)
+    fast.warmup(device)
+    assert main.is_ready and fast.is_ready
+    assert main._capture_stream is None and fast._capture_stream is not None
+
+    codes_main = torch.randint(0, 5, (2, 4, 3), device=device)
+    codes_fast = torch.randint(0, 5, (2, 2, 1), device=device)
+    errors = []
+
+    def run(wrapper, codes, slots, iters):
+        stream = torch.cuda.Stream(device=device)
+        slot_ids = torch.tensor(slots, device=device)
+        try:
+            with torch.cuda.stream(stream):
+                for _ in range(iters):
+                    codec.reset_decoder_state_slots(slot_ids)
+                    audio = wrapper.decode(codes, slot_ids)[0][: len(slots)].clone()
+                    expected = (codes.float().sum(0)[..., None].expand(-1, -1, DIM) @ codec.weight).transpose(1, 2)
+                    stream.synchronize()
+                    torch.testing.assert_close(audio, expected, rtol=1e-3, atol=1e-3)
+        except Exception as error:  # surfaced below
+            errors.append(error)
+
+    threads = [
+        threading.Thread(target=run, args=(main, codes_main, [0, 1, 2, 3], 200)),
+        threading.Thread(target=run, args=(fast, codes_fast, [4, 5], 200)),
+    ]
+    for thread in threads:
+        thread.start()
+    for thread in threads:
+        thread.join(timeout=60)
+    assert not any(thread.is_alive() for thread in threads), "concurrent graph replay did not finish"
+    assert not errors, errors

tests/model_executor/models/moss_tts/test_streaming_terminal_batch.py

@@ -54,8 +54,19 @@ def test_terminal_padding_preserves_audio_live_state_and_reused_slots():
     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):
+@pytest.mark.parametrize(
+    "graph_capacity,expected_groups",
+    [
+        # No graph, or graphs that do not cover the state capacity: exact tails.
+        (None, {3: [0], 7: [1], 8: [2], 1: [3], 15: [4]}),
+        (4, {3: [0], 7: [1], 8: [2], 1: [3], 15: [4]}),
+        # Full graph coverage: finished tails (3, 7) join the padded 8-frame
+        # bucket with the live 8-frame row; the finished 1-frame row keeps the
+        # exact 1-frame graph and the live 15-frame row stays alone.
+        (8, {8: [0, 1, 2], 1: [3], 15: [4]}),
+    ],
+)
+def test_terminal_tails_coalesce_into_padded_frame_buckets(mocker, graph_capacity, expected_groups):
     s = session()
     s._cudagraph_wrapper = (
         None
@@ -81,9 +92,16 @@ def test_only_existing_compatible_graph_buckets_are_coalesced(mocker, graph_capa
             for i, (length, done) in enumerate(zip(lengths, finished, strict=True))
         ]
     )
-    assert step.call_count == expected_calls
+    # Slots are leased in request order, so slot == row here.
+    groups = {}
+    for call in step.call_args_list:
+        plan = call.args[0]
+        groups[max(int(codes.shape[1]) for codes in plan.values())] = sorted(plan)
+        terminal = call.kwargs["terminal_slots"] or set()
+        assert terminal == {slot for slot in plan if finished[slot]}
+    assert groups == expected_groups
+    assert step.call_count == len(expected_groups)
+    # Output is always cropped to the exact frame count, padded or not.
     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.
+    # Finished rows released their slots; live rows keep theirs.
+    assert set(decoder._stream_req_slots) == {"2", "4"}

tests/model_executor/models/qwen3_omni/test_first_frame_cuda.py

@@ -0,0 +1,32 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from vllm_omni.model_executor.models.qwen3_omni.first_frame_decoder import Qwen3OmniFirstFrameDecoder
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cuda]
+
+
+@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16])
+@pytest.mark.parametrize("rows", [16, 19])
+@torch.inference_mode()
+def test_first_frame_graph_outputs_own_storage(dtype, rows):
+    class Decoder(torch.nn.Linear):
+        config = SimpleNamespace(num_quantizers=1)
+
+        def forward(self, codes):
+            return codes.to(dtype) * self.weight
+
+    model = Decoder(1, 1, bias=False, device="cuda", dtype=dtype)
+    model.weight.fill_(1)
+    decoder = Qwen3OmniFirstFrameDecoder(model, sample_rate=24000)
+    decoder.capture()
+    codes = torch.arange(rows, device="cuda").reshape(-1, 1)
+    first = decoder.decode(codes)
+    second = decoder.decode(codes + 32)
+    assert first.dtype == second.dtype == torch.float32
+    torch.testing.assert_close(first, codes.float(), rtol=0, atol=0)
+    torch.testing.assert_close(second, (codes + 32).float(), rtol=0, atol=0)

tests/model_executor/models/qwen3_omni/test_qwen3_omni_code2wav_cudnn.py

@@ -0,0 +1,87 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+from transformers.models.qwen3_omni_moe.configuration_qwen3_omni_moe import Qwen3OmniMoeCode2WavConfig
+
+from vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_code2wav import Qwen3OmniMoeCode2Wav
+
+pytestmark = [
+    pytest.mark.core_model,
+    pytest.mark.cuda,
+    pytest.mark.skipif(not torch.cuda.is_available() or torch.version.hip is not None, reason="requires CUDA"),
+]
+
+
+def _codec(enabled):
+    config = Qwen3OmniMoeCode2WavConfig(
+        hidden_size=32,
+        intermediate_size=64,
+        num_hidden_layers=1,
+        num_attention_heads=4,
+        num_key_value_heads=4,
+        head_dim=8,
+        decoder_dim=32,
+        upsample_rates=[2, 2],
+        upsampling_ratios=[2],
+        codebook_size=16,
+        num_quantizers=2,
+    )
+    config._attn_implementation = "sdpa"
+    vllm_config = SimpleNamespace(
+        model_config=SimpleNamespace(
+            hf_config=config, stage_connector_config={"extra": {"codec_cudnn_benchmark": enabled}}
+        )
+    )
+    model = Qwen3OmniMoeCode2Wav(vllm_config=vllm_config).cuda().eval()
+    model.precompute_snake_caches()
+    return model
+
+
+def _flags():
+    c = torch.backends.cudnn
+    return c.enabled, c.benchmark, c.benchmark_limit, c.deterministic, c.allow_tf32
+
+
+@torch.inference_mode()
+def test_autotuned_codec_graph_replay_restores_flags():
+    torch.manual_seed(21)
+    model = _codec(True)
+    codes = torch.randint(0, 16, (2, 2, 8), device="cuda")
+    original = _flags()
+    for _ in range(3):
+        model(codes)
+    assert _flags() == original
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph):
+        captured = model(codes)
+    assert _flags() == original
+    codes.random_(0, 16)
+    graph.replay()
+    expected = model(codes)
+    torch.testing.assert_close(captured, expected)
+    assert _flags() == original
+
+
+@pytest.mark.parametrize("enabled", [False, True])
+@torch.inference_mode()
+def test_codec_failure_restores_cudnn_settings(mocker, enabled):
+    model = _codec(enabled)
+    codes = torch.randint(0, 16, (1, 2, 8), device="cuda")
+    original = _flags()
+
+    def fail(hidden):
+        if enabled:
+            assert torch.backends.cudnn.benchmark
+            assert torch.backends.cudnn.benchmark_limit == 10
+        else:
+            assert _flags() == original
+        raise RuntimeError("decode failed")
+
+    mocker.patch.object(model, "_decode_waveform", side_effect=fail)
+    with pytest.raises(RuntimeError, match="decode failed"):
+        model(codes)
+    assert _flags() == original

tests/model_executor/models/qwen3_omni/test_qwen3_omni_code2wav_snake.py

@@ -0,0 +1,88 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+import copy
+from types import SimpleNamespace
+
+import pytest
+import torch
+from transformers.models.qwen3_omni_moe.configuration_qwen3_omni_moe import Qwen3OmniMoeCode2WavConfig
+from transformers.models.qwen3_omni_moe.modeling_qwen3_omni_moe import (
+    Qwen3OmniMoeCode2WavDecoderBlock,
+    Qwen3OmniMoeSnakeBeta,
+)
+
+from vllm_omni.model_executor.models.common.snake_activation import SnakeBeta
+from vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_code2wav import Qwen3OmniMoeCode2Wav, use_fused_snake
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def _blocks():
+    config = Qwen3OmniMoeCode2WavConfig(decoder_dim=32, upsample_rates=[2, 2])
+    torch.manual_seed(0)
+    blocks = torch.nn.ModuleList(Qwen3OmniMoeCode2WavDecoderBlock(config, i) for i in range(2)).eval()
+    with torch.no_grad():
+        for module in blocks.modules():
+            if isinstance(module, Qwen3OmniMoeSnakeBeta):
+                module.alpha.uniform_(-0.5, 0.5)
+                module.beta.uniform_(-0.5, 0.5)
+    return blocks
+
+
+def test_decoder_blocks_swap_every_hf_snake_and_keep_weight_names():
+    blocks = _blocks()
+    hf_count = sum(isinstance(m, Qwen3OmniMoeSnakeBeta) for m in blocks.modules())
+    keys = set(blocks.state_dict())
+
+    assert use_fused_snake(blocks) == hf_count > 0
+    assert not any(isinstance(m, Qwen3OmniMoeSnakeBeta) for m in blocks.modules())
+    assert sum(isinstance(m, SnakeBeta) for m in blocks.modules()) == hf_count
+    assert set(blocks.state_dict()) == keys  # checkpoint weights load unchanged
+
+
+def test_fused_snake_blocks_match_hf_blocks():
+    reference = _blocks()
+    fused = copy.deepcopy(reference)
+    use_fused_snake(fused)
+    fused.load_state_dict(reference.state_dict())
+    x = torch.randn(2, 32, 20)
+    with torch.no_grad():
+        expected = x
+        actual = x
+        for ref_block, fused_block in zip(reference, fused):
+            expected = ref_block(expected)
+            actual = fused_block(actual)
+    torch.testing.assert_close(actual, expected, rtol=1e-5, atol=1e-5)
+
+
+@pytest.mark.parametrize(
+    ("option", "device", "cuda", "fused"),
+    [
+        (False, "cuda", True, False),
+        (True, "cuda", True, True),
+        (True, "cpu", True, False),
+        (True, "cuda", False, False),
+    ],
+)
+def test_decoder_block_fusion_requires_cuda_opt_in(monkeypatch, option, device, cuda, fused):
+    from vllm_omni.platforms import current_omni_platform
+
+    monkeypatch.setattr(current_omni_platform, "is_cuda", lambda: cuda)
+    config = Qwen3OmniMoeCode2WavConfig(
+        hidden_size=32,
+        intermediate_size=64,
+        num_hidden_layers=1,
+        num_attention_heads=2,
+        num_key_value_heads=2,
+        decoder_dim=32,
+    )
+    model = Qwen3OmniMoeCode2Wav(
+        vllm_config=SimpleNamespace(
+            model_config=SimpleNamespace(
+                hf_config=config, stage_connector_config={"extra": {"codec_fused_snake": option}}
+            ),
+            device_config=SimpleNamespace(device=device),
+        )
+    )
+    assert any(isinstance(m, Qwen3OmniMoeSnakeBeta) for m in model.decoder.modules()) is (not fused)

tests/model_executor/models/qwen3_omni/test_qwen3_omni_first_audio.py

@@ -0,0 +1,208 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from types import SimpleNamespace
+from unittest.mock import MagicMock
+
+import pytest
+import torch
+from transformers.models.qwen3_omni_moe.configuration_qwen3_omni_moe import Qwen3OmniMoeCode2WavConfig
+
+from vllm_omni.model_executor.models.qwen3_omni.first_frame_decoder import Qwen3OmniFirstFrameDecoder
+from vllm_omni.model_executor.models.qwen3_omni.qwen3_omni import Qwen3OmniMoeForConditionalGeneration
+from vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_code2wav import (
+    Qwen3OmniMoeCode2Wav,
+    plan_decode_groups,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+_Q = 4
+_CODEBOOK = 16
+
+
+def _code2wav(quantizers=_Q) -> Qwen3OmniMoeCode2Wav:
+    torch.manual_seed(0)
+    config = Qwen3OmniMoeCode2WavConfig(
+        codebook_size=_CODEBOOK,
+        num_quantizers=quantizers,
+        hidden_size=32,
+        intermediate_size=64,
+        num_hidden_layers=2,
+        num_attention_heads=2,
+        num_key_value_heads=2,
+        decoder_dim=32,
+        sliding_window=8,
+    )
+    model = Qwen3OmniMoeCode2Wav(vllm_config=SimpleNamespace(model_config=SimpleNamespace(hf_config=config)))
+    return model.eval()
+
+
+def test_first_frame_pcm_equals_code2wav_streaming_chunk0():
+    code2wav = _code2wav()
+    decoder = Qwen3OmniFirstFrameDecoder(code2wav, sample_rate=24000)
+    codes = torch.randint(0, _CODEBOOK, (3, _Q))
+    with torch.inference_mode():
+        for row in range(codes.shape[0]):
+            pcm = decoder.decode(codes[row : row + 1])
+            chunk0 = code2wav.chunked_decode_streaming(
+                codes[row].reshape(1, _Q, 1), left_context_size=[0], seq_token_counts=[_Q]
+            )[0]
+            assert pcm.dtype == torch.float32
+            assert torch.equal(pcm[0], chunk0.reshape(-1).float())
+            assert 0 < pcm.shape[-1] < int(code2wav.total_upsample)
+        batched = decoder.decode(codes)
+    assert batched.shape == (3, pcm.shape[-1])
+
+
+def test_first_frame_is_a_prefix_of_a_longer_first_chunk():
+    code2wav = _code2wav()
+    decoder = Qwen3OmniFirstFrameDecoder(code2wav, sample_rate=24000)
+    codes = torch.randint(0, _CODEBOOK, (_Q, 4))
+    with torch.inference_mode():
+        first = decoder.decode(codes[:, :1].T)[0]
+        chunk0 = code2wav.chunked_decode_streaming(
+            codes.reshape(1, _Q, 4), left_context_size=[0], seq_token_counts=[4 * _Q]
+        )[0]
+    torch.testing.assert_close(chunk0.reshape(-1)[: first.numel()].float(), first)
+
+
+def test_chunk_ramp_adds_exact_code2wav_graph_sizes(monkeypatch):
+    import vllm_omni.model_executor.models.qwen3_tts.cuda_graph_decoder_wrapper as wrapper_module
+
+    real = wrapper_module.CUDAGraphDecoderWrapper
+    wrapper = MagicMock(compute_capture_sizes=real.compute_capture_sizes)
+    monkeypatch.setattr(wrapper_module, "CUDAGraphDecoderWrapper", wrapper)
+    extra = {"codec_chunk_frames": 25, "codec_left_context_frames": 25, "codec_chunk_ramp": [1, 2, 4, 8, 16, 25]}
+    code2wav = SimpleNamespace(config=SimpleNamespace(num_quantizers=_Q))
+    code2wav.enable_cudagraph = lambda **kw: Qwen3OmniMoeCode2Wav.enable_cudagraph(
+        code2wav, device=torch.device("cuda"), **kw
+    )
+    model = object.__new__(Qwen3OmniMoeForConditionalGeneration)
+    torch.nn.Module.__init__(model)
+    model.code2wav = code2wav
+    model.vllm_config = SimpleNamespace(
+        model_config=SimpleNamespace(async_chunk=True, enforce_eager=False, stage_connector_config={"extra": extra})
+    )
+    for batches in ([], [2, 4, 8]):
+        extra["codec_graph_batch_sizes"] = batches
+        Qwen3OmniMoeForConditionalGeneration._maybe_enable_code2wav_cudagraph(model)
+        args = wrapper.call_args.kwargs
+        sizes = args["capture_sizes"]
+        assert {1, 3, 7, 15, 31, 50} <= set(sizes)
+        assert set(real.compute_capture_sizes(codec_chunk_frames=25, codec_left_context_frames=25)) <= set(sizes)
+        assert set(args["extra_capture_shapes"]) == {
+            (b, size) for b in batches for size in sizes if size <= 50 and b * size <= 325
+        }
+    assert code2wav._streaming_batch_sizes[50] == [1, 2, 4]
+    extra.pop("codec_chunk_ramp")
+    extra.pop("codec_graph_batch_sizes")
+    Qwen3OmniMoeForConditionalGeneration._maybe_enable_code2wav_cudagraph(model)
+    assert wrapper.call_args.kwargs["capture_sizes"] is None
+
+
+_SIZES = [1, 2, 3, 4, 7, 8, 15, 16, 25, 31, 32, 50, 64]
+
+
+def _bucket(length):
+    return next((size for size in _SIZES if size >= length), None)
+
+
+def test_decode_groups_keep_a_uniform_batch_whole_and_split_by_the_largest_graph():
+    assert plan_decode_groups([50] * 6, _bucket, lambda size: [1, 2, 4, 8]) == [(list(range(6)), 50)]
+    calls = plan_decode_groups([25] * 20, _bucket, lambda size: [1, 2, 4, 8, 16])
+    assert [len(rows) for rows, _length in calls] == [16, 4]
+    calls = plan_decode_groups([1] * 8 + [50] * 8, _bucket, lambda size: [1, 2, 4, 8] if size < 50 else [1, 2, 4])
+    assert [(len(rows), length) for rows, length in calls] == [(8, 1), (4, 50), (4, 50)]
+    assert plan_decode_groups([], _bucket, lambda size: [1, 2]) == []
+
+
+@pytest.mark.parametrize("frames", [65, 129])
+def test_decode_groups_beyond_largest_graph_use_eager(monkeypatch, frames):
+    from vllm_omni.model_executor.models.qwen3_tts.cuda_graph_decoder_wrapper import CUDAGraphDecoderWrapper
+
+    wrapper = CUDAGraphDecoderWrapper.__new__(CUDAGraphDecoderWrapper)
+    wrapper.enabled = wrapper._warmed_up = True
+    wrapper._bucket_sizes = [8, 32, 64]
+    wrapper._compiled_shapes = {(1, 64), (8, 64)}
+    wrapper._compiled_graphs = wrapper.graphs = dict.fromkeys(wrapper._compiled_shapes, MagicMock())
+    wrapper.decoder = MagicMock(side_effect=lambda codes: codes.float() * 2)
+    monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: False)
+    groups = plan_decode_groups([frames, frames], wrapper._get_padded_size, lambda _: (1,))
+    assert groups == [([0], frames), ([1], frames)]
+    codes = torch.arange(2 * frames).reshape(2, 1, frames)
+    for rows, width in groups:
+        actual = wrapper.decode(codes[rows, :, :width])
+        torch.testing.assert_close(actual, codes[rows].float() * 2, rtol=0, atol=0)
+    assert wrapper.decoder.call_count == 2
+    for graph in (*wrapper.graphs.values(), *wrapper._compiled_graphs.values()):
+        graph.replay.assert_not_called()
+
+
+class _PaddingGraphs:
+    def __init__(self, model, batch_sizes):
+        self.model = model
+        self.batch_sizes = sorted(batch_sizes)
+        self.calls = []
+
+    def _get_padded_size(self, length):
+        return _bucket(length)
+
+    def decode(self, codes):
+        rows, _q, frames = codes.shape
+        size = _bucket(frames)
+        batch = next(b for b in self.batch_sizes if b >= rows)
+        self.calls.append((rows, frames))
+        padded = torch.zeros(batch, codes.shape[1], size, dtype=codes.dtype)
+        padded[:rows, :, :frames] = codes
+        out = self.model(padded)
+        return out[:rows, :, : out.shape[-1] - (size - frames) * int(self.model.total_upsample)].clone()
+
+
+def test_grouped_streaming_decode_matches_one_padded_batch():
+    code2wav = _code2wav()
+    lengths = [1, 3, 50, 7, 31]
+    lefts = [0, 1, 25, 3, 15]
+    codes = torch.zeros(len(lengths), _Q, max(lengths), dtype=torch.long)
+    for row, length in enumerate(lengths):
+        codes[row, :, :length] = torch.randint(0, _CODEBOOK, (_Q, length))
+    counts = [length * _Q for length in lengths]
+    with torch.inference_mode():
+        reference = code2wav.chunked_decode_streaming(codes, left_context_size=lefts, seq_token_counts=counts)
+        graphs = _PaddingGraphs(code2wav, [1, 2, 4])
+        code2wav._cudagraph_wrapper, code2wav._cudagraph_enabled = graphs, True
+        code2wav._streaming_batch_sizes = {size: [1, 2, 4] for size in _SIZES}
+        grouped = code2wav.chunked_decode_streaming(codes, left_context_size=lefts, seq_token_counts=counts)
+
+    for want, got in zip(reference, grouped, strict=True):
+        assert want.shape == got.shape
+        torch.testing.assert_close(got, want, rtol=1e-4, atol=1e-5)
+
+
+@pytest.mark.parametrize(("frames", "quantizers"), [(1, 16), (4, 16), (4, 4)])
+def test_delivered_first_frame_is_removed_from_code2wav_output(frames, quantizers):
+    from vllm_omni.data_entry_keys import FIRST_AUDIO_REQUIRED_KEY
+
+    codec = _code2wav(quantizers)
+    model = object.__new__(Qwen3OmniMoeForConditionalGeneration)
+    torch.nn.Module.__init__(model)
+    model.model_stage, model.code2wav, model.code2wav_config = "code2wav", codec, codec.config
+    model.generate_audio = lambda codes, left, counts: codec.chunked_decode_streaming(
+        codes, left_context_size=left, seq_token_counts=counts
+    )
+    codes = torch.randint(0, _CODEBOOK, (2, quantizers, frames))
+    with torch.inference_mode():
+        reference = model.generate_audio(codes, [0, 0], [quantizers * frames] * 2)
+        first = Qwen3OmniFirstFrameDecoder(codec, sample_rate=24000).decode(codes[0, :, :1].T)[0]
+        out = model.forward(
+            input_ids=codes.flatten(),
+            positions=None,
+            seq_token_counts=[quantizers * frames] * 2,
+            runtime_additional_information=[
+                {"meta": {"left_context_size": 0, "first_audio": True}},
+                {"meta": {"left_context_size": 0, "first_audio": False}},
+            ],
+        ).multimodal_outputs
+    torch.testing.assert_close(torch.cat([first, out["model_outputs"][0].flatten()]), reference[0].flatten())
+    torch.testing.assert_close(out["model_outputs"][1], reference[1].reshape(1, -1))
+    assert [bool(flag) for flag in out[FIRST_AUDIO_REQUIRED_KEY]] == [True, False]

tests/model_executor/models/qwen3_omni/test_qwen3_omni_mrv2.py

@@ -0,0 +1,116 @@
+# 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.qwen3_omni.qwen3_omni import Qwen3OmniMoeForConditionalGeneration
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+_H = 4
+_Q = 3
+
+
+class _Talker(nn.Module):
+    num_code_groups = _Q
+
+    def __init__(self):
+        super().__init__()
+        self.text_projection = nn.Identity()
+        self.code_predictor = SimpleNamespace(
+            _top_k=50,
+            _top_p=0.8,
+            set_sampling_params=lambda **kwargs: setattr(self, "sampling", kwargs),
+        )
+
+
+def _talker(*, async_chunk: bool = True) -> Qwen3OmniMoeForConditionalGeneration:
+    model = object.__new__(Qwen3OmniMoeForConditionalGeneration)
+    nn.Module.__init__(model)
+    model.model_stage = "talker"
+    model.talker = _Talker()
+    model.vllm_config = SimpleNamespace(model_config=SimpleNamespace(async_chunk=async_chunk, use_v2_model_runner=True))
+    model.tts_eos_embed = torch.full((_H,), -1.0)
+    model.tts_pad_embed = torch.full((_H,), -2.0)
+    model._codec_codebook_size = 2048
+    return model
+
+
+def _rows(*values: float) -> torch.Tensor:
+    return torch.tensor([[value] * _H for value in values])
+
+
+def _step(model, payload: dict) -> torch.Tensor:
+    update: dict = {}
+    text_step = model._thinker_decode_to_talker_decode(payload, torch.device("cpu"), update)
+    for key, value in update.items():
+        payload.setdefault(key, {}).update(value)
+    return text_step.reshape(-1)
+
+
+def test_mtp_frame_valid_rejects_codec_eos_and_special_ids():
+    layer0 = torch.tensor([0, 2047, 2048, 2150, -1])
+    assert _talker().mtp_frame_valid(layer0).tolist() == [True, True, False, False, False]
+
+
+def test_decode_rows_are_consumed_one_per_step_in_arrival_order():
+    model = _talker()
+    # Two rows were already pending (cached at prefill); two more arrive at once.
+    payload = {"embed": {"cached_decode": _rows(1, 2), "decode": _rows(3, 4)}, "meta": {}}
+    steps = [_step(model, payload)[0].item()]
+    for arrival in (None, _rows(5), None, None, None, None):
+        payload["embed"]["decode"] = arrival
+        steps.append(_step(model, payload)[0].item())
+    assert steps == [1, 2, 3, 4, 5, -1, -2]
+    assert payload["embed"]["cached_decode"] is None
+
+
+@pytest.mark.parametrize(("v2", "resumable"), [(True, True), (False, False), (False, True)])
+def test_v1_and_resumable_requests_keep_the_indexed_cache_path(v2, resumable):
+    model = _talker()
+    model.vllm_config.model_config.use_v2_model_runner = v2
+    payload = {"embed": {"cached_decode": _rows(1, 2, 3), "decode": None}, "meta": {"resumable": resumable}}
+    payload["meta"]["num_processed_tokens"] = 2
+    assert _step(model, payload)[0].item() == 3
+
+
+def test_make_omni_output_aligns_codes_to_request_spans():
+    model = _talker()
+    hidden = torch.zeros((8, _H))  # 5 real tokens + graph padding
+    decode_codes = torch.tensor([[5, 6, 7]])
+    buffers = [
+        {"req_id": "_warmup_0_"},  # no preprocess, no codes
+        {"req_id": "a", "codes": {"audio": torch.zeros((3, _Q), dtype=torch.long)}},
+        {"req_id": "b", "codes": {"audio": decode_codes}},
+    ]
+    out = model.make_omni_output(
+        hidden, model_intermediate_buffer=buffers, request_token_spans=[(0, 1), (1, 4), (4, 5)]
+    )
+    codes = out.multimodal_outputs["codes"]["audio"]
+    assert codes.shape == (5, _Q)
+    assert codes[4].tolist() == [5, 6, 7]
+    assert out.text_hidden_states.shape[0] == 5
+
+
+def test_make_omni_output_without_spans_keeps_v1_concat():
+    model = _talker()
+    buffers = [{"codes": {"audio": torch.tensor([[1, 2, 3]])}}, {"codes": {"audio": torch.tensor([[4, 5, 6]])}}]
+    out = model.make_omni_output(torch.zeros((4, _H)), model_intermediate_buffer=buffers)
+    assert out.multimodal_outputs["codes"]["audio"].tolist() == [[1, 2, 3], [4, 5, 6]]
+    assert out.text_hidden_states.shape[0] == 2
+
+
+def test_eager_output_is_token_major_with_invalid_rows():
+    model = _talker()
+    model.eager_frames_active = True
+    hidden = torch.zeros((6, _H))
+    buffers = [{"req_id": "a"}, {"req_id": "b"}]
+    out = model.make_omni_output(hidden, model_intermediate_buffer=buffers, request_token_spans=[(0, 3), (3, 4)])
+    mm = out.multimodal_outputs
+    assert mm["codes"]["audio"].shape == (4, _Q)
+    assert mm["meta"]["codec_frame_valid"].tolist() == [0, 0, 0, 0]
+    assert out.text_hidden_states is hidden

tests/model_executor/models/qwen3_omni/test_qwen3_omni_pp_capture.py

@@ -104,6 +104,7 @@ def test_factory_declares_only_captures_from_earlier_ranks(monkeypatch, start, e
     if staged:
         model_config.model_stage = "thinker"
     vllm_config = SimpleNamespace(model_config=model_config, with_hf_config=lambda *args, **kwargs: None)
+    vllm_config.speculative_config = None
     thinker = nn.Module()
     thinker.weight = nn.Parameter(torch.ones(1))
     thinker.language_model = SimpleNamespace(model=SimpleNamespace(start_layer=start))

tests/model_executor/models/qwen3_tts/test_code_predictor_dtype.py

@@ -569,6 +569,27 @@ class TestCodePredictorPerRowGenerators:
             )
 
 
+@pytest.mark.parametrize("fused_requested", [False, True])
+@pytest.mark.parametrize("invalid_input", ["uniforms", "generators"])
+def test_sampling_contract_rejected_before_dispatch(mocker, loaded_target_classes, fused_requested, invalid_input):
+    predictor, talker_config = TestCodePredictorPerRowGenerators()._make_predictor(mocker, loaded_target_classes)
+    predictor._fused_requested = fused_requested
+    predictor._fused = mocker.Mock()
+    setup = mocker.patch.object(predictor, "_setup_compile")
+    uniforms = torch.full((2, 3, 1 if invalid_input == "uniforms" else 64), 0.5)
+    generators = [torch.Generator()] if invalid_input == "generators" else None
+    with pytest.raises(ValueError, match="sample_uniforms must have shape|one entry per row"):
+        predictor(
+            torch.zeros(2, dtype=torch.long),
+            torch.randn(2, talker_config.hidden_size),
+            torch.randn(2, talker_config.hidden_size),
+            sample_uniforms=uniforms,
+            generators=generators,
+        )
+    setup.assert_not_called()
+    predictor._fused.assert_not_called()
+
+
 class TestCodePredictorModelDtype:
     """Test the inner model forward with different dtypes."""
 
@@ -663,6 +684,8 @@ class TestCodePredictorGraphReplay:
 
         predictor = object.__new__(code_predictor_wrapper)
         torch.nn.Module.__init__(predictor)
+        predictor._fused_requested = False
+        predictor._fused = None
         predictor._num_groups = 3
         predictor._model_dtype = torch.float32
         predictor._setup_compile = mocker.Mock()
@@ -1432,3 +1455,24 @@ def test_shared_predictor_ignores_tts_fast_path_options(mocker, loaded_target_cl
     expected = (logits.float() - torch.log(-torch.log(uniforms))).argmax(-1, keepdim=True)
     actual = predictor._sample_per_call(logits, 1.0, 0, None, uniforms)
     torch.testing.assert_close(actual, expected)
+
+
+def test_tts_cached_predictor_keeps_its_own_step_dispatch(mocker: MockerFixture, loaded_target_classes) -> None:
+    """The TTS cache must not enter Omni's uninitialized incremental loop."""
+    _, _, wrapper, _, _ = loaded_target_classes
+    common_mod = sys.modules["vllm_omni.model_executor.models.common.qwen3_code_predictor"]
+    mocker.patch.object(common_mod.current_omni_platform, "is_cuda", return_value=True)
+    mocker.patch.object(common_mod.current_omni_platform, "is_npu", return_value=False)
+    mocker.patch.object(wrapper, "_stage_connector_extra_config", return_value={"code_predictor_kv_cache": True})
+    cp_config, talker_config = _make_tiny_config(loaded_target_classes)
+    predictor = wrapper(vllm_config=_make_vllm_config(mocker), config=cp_config, talker_config=talker_config)
+    assert predictor._kv_requested
+    predictor._model_dtype = torch.float32
+    predictor._codec_embeds_list = list(predictor.model.codec_embedding)
+    predictor._bucket_sizes = [1]
+    mocker.patch.object(predictor, "_setup_compile")
+    step = mocker.patch.object(predictor, "_predict_step_logits", return_value=torch.ones(1, cp_config.vocab_size))
+    codes = predictor(torch.zeros(1, dtype=torch.long), torch.zeros(1, 32), torch.zeros(1, 32), do_sample=False)
+    assert step.call_count == cp_config.num_code_groups - 1
+    assert codes.shape == (1, cp_config.num_code_groups)
+    assert predictor._kv_buf is None

tests/model_executor/models/qwen3_tts/test_code_predictor_sampling.py

@@ -2,6 +2,8 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 """Supplied-uniform samples must match the existing ATen path exactly."""
 
+from types import SimpleNamespace
+
 import pytest
 import torch
 
@@ -12,6 +14,79 @@ pytestmark = [
 ]
 
 
+@pytest.mark.parametrize("top_k", [0, 1, 50, 2048])
+@torch.inference_mode()
+def test_fused_sample_kernel_preserves_signed_zero_cutoff_ties(top_k):
+    from vllm_omni.model_executor.models.qwen3_tts.fused_code_predictor import _cp_sample_kernel
+
+    vocab, hidden, groups = 2048, 64, 4
+    logits = torch.zeros(2, vocab, device="cuda", dtype=torch.bfloat16)
+    logits[0, 1] = -0.0
+    logits[1, ::2] = -0.0
+    uniforms = torch.full((2, groups - 1, vocab), 0.01, device="cuda")[:, 1, :]
+    uniforms[0, 1] = uniforms[1, 2] = 0.99
+    codes = torch.full((2, groups), -1, device="cuda", dtype=torch.int64)
+    table = torch.randn(vocab, hidden, device="cuda", dtype=torch.bfloat16)
+    next_input = torch.empty(2, hidden, device="cuda", dtype=table.dtype)
+
+    def sample():
+        _cp_sample_kernel[(2,)](
+            logits, uniforms, uniforms.stride(0), codes, 1, groups, table, next_input, 1.0,
+            V=vocab, TOPK=top_k, HID=hidden, HAS_NEXT=True, num_warps=4,
+        )  # fmt: skip
+
+    scaled = logits.float()
+    if top_k > 0:
+        scaled = scaled.masked_fill(scaled < scaled.topk(top_k).values[:, -1:], -torch.inf)
+    expected = (scaled - torch.log(-torch.log(uniforms))).argmax(-1)
+    assert expected.tolist() == [1, 2]
+    sample()
+    torch.testing.assert_close(codes[:, 1], expected)
+    torch.testing.assert_close(next_input, table[expected])
+    assert (codes[:, [0, 2, 3]] == -1).all()
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph):
+        sample()
+    codes.fill_(-1)
+    graph.replay()
+    torch.testing.assert_close(codes[:, 1], expected)
+    torch.testing.assert_close(next_input, table[expected])
+
+
+@torch.inference_mode()
+def test_fused_predictor_warmup_with_capacity_one():
+    from vllm_omni.model_executor.models.qwen3_tts.configuration_qwen3_tts import (
+        Qwen3TTSTalkerCodePredictorConfig,
+        Qwen3TTSTalkerConfig,
+    )
+    from vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_code_predictor_vllm import (
+        Qwen3TTSTalkerCodePredictorForConditionalGenerationVLLM,
+    )
+
+    cp = Qwen3TTSTalkerCodePredictorConfig(
+        vocab_size=64, hidden_size=64, intermediate_size=128, num_hidden_layers=1,
+        num_attention_heads=4, num_key_value_heads=2, head_dim=16, num_code_groups=4,
+    )  # fmt: skip
+    talker = Qwen3TTSTalkerConfig(hidden_size=64, num_code_groups=4)
+    config = SimpleNamespace(
+        model_config=SimpleNamespace(stage_connector_config={}),
+        additional_config={"code_predictor_kv_cache": True, "code_predictor_fused": True},
+        scheduler_config=SimpleNamespace(max_num_seqs=1),
+    )
+    predictor = Qwen3TTSTalkerCodePredictorForConditionalGenerationVLLM(
+        vllm_config=config, config=cp, talker_config=talker
+    ).to(device="cuda", dtype=torch.bfloat16)
+    predictor._setup_compile()
+    assert predictor._fused is not None
+    hidden = torch.randn(1, 1, 64, device="cuda", dtype=torch.bfloat16)
+    result = predictor(
+        torch.zeros(1, 1, device="cuda", dtype=torch.long), hidden, hidden,
+        sample_uniforms=torch.full((1, 3, 64), 0.5, device="cuda"),
+    )  # fmt: skip
+    assert result.shape == (1, 4)
+    assert ((result >= 0) & (result < 64)).all()
+
+
 @pytest.mark.parametrize("batch", [1, 64])
 @pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16, torch.float32])
 @pytest.mark.parametrize("top_k", [-1, 1, 50, 2048])
@@ -66,3 +141,47 @@ def test_topk_gumbel_accepts_strided_supplied_uniforms(top_k):
     expected = (scaled.float() - torch.log(-torch.log(uniforms))).argmax(-1, keepdim=True)
     actual = sample_code_topk_gumbel(logits, uniforms, top_k, 1 / 0.9)
     torch.testing.assert_close(actual, expected)
+
+
+@pytest.mark.parametrize("batch", [1, 4])
+@pytest.mark.parametrize("top_k", [1, 50])
+@torch.inference_mode()
+def test_fused_predictor_matches_full_reference_with_fixed_uniforms(mocker, batch, top_k):
+    from vllm.config import VllmConfig
+
+    from vllm_omni.model_executor.models.qwen3_tts.configuration_qwen3_tts import (
+        Qwen3TTSTalkerCodePredictorConfig,
+        Qwen3TTSTalkerConfig,
+    )
+    from vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_code_predictor_vllm import (
+        Qwen3TTSTalkerCodePredictorForConditionalGenerationVLLM,
+    )
+
+    torch.manual_seed(42)
+    cp = Qwen3TTSTalkerCodePredictorConfig(
+        vocab_size=64, hidden_size=64, intermediate_size=128, num_hidden_layers=1,
+        num_attention_heads=4, num_key_value_heads=2, head_dim=16, num_code_groups=4,
+    )  # fmt: skip
+    config = mocker.Mock(
+        spec=VllmConfig,
+        model_config=mocker.Mock(stage_connector_config={}),
+        additional_config={"code_predictor_kv_cache": True, "code_predictor_fused": True},
+        scheduler_config=mocker.Mock(max_num_seqs=4),
+    )
+    predictor = Qwen3TTSTalkerCodePredictorForConditionalGenerationVLLM(
+        vllm_config=config, config=cp, talker_config=Qwen3TTSTalkerConfig(hidden_size=64, num_code_groups=4)
+    ).to(device="cuda", dtype=torch.bfloat16)
+    predictor._setup_compile()
+    assert predictor._fused is not None
+    ids = torch.randint(0, 64, (batch, 1), device="cuda")
+    embeds = torch.randn(batch, 1, 64, device="cuda", dtype=torch.bfloat16)
+    hidden = torch.randn_like(embeds)
+    uniforms = torch.rand(batch, 3, 64, device="cuda").clamp_(1e-6, 1 - 1e-6)
+    actual = predictor(ids, embeds, hidden, top_k=top_k, sample_uniforms=uniforms).clone()
+    # Use identical parameters through the full re-prefill and ATen sampling
+    # path, bypassing both the fused predictor and its frame-local KV cache.
+    predictor._fused_requested = False
+    predictor._fused = None
+    predictor._frame_cache = None
+    expected = predictor(ids, embeds, hidden, top_k=top_k, sample_uniforms=uniforms)
+    torch.testing.assert_close(actual, expected, rtol=0, atol=0)

tests/model_executor/models/qwen3_tts/test_first_audio_graphs.py

@@ -2,6 +2,7 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 """Prefix graphs must follow per-request first-audio delivery, not a capture hint."""
 
+import math
 from types import SimpleNamespace
 
 import pytest
@@ -18,6 +19,38 @@ pytestmark = [pytest.mark.core_model]
 DEVICE = torch.device("cuda:0")
 
 
+@pytest.mark.cpu
+@pytest.mark.parametrize("capacity,model_type", [(1, "base"), (128, "base"), (128, "custom_voice")])
+def test_stream_priming_graph_workspace_respects_token_budget(monkeypatch, capacity, model_type):
+    from vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_talker import Qwen3TTSTalkerForConditionalGeneration
+    from vllm_omni.model_executor.models.qwen3_tts.tokenizer_12hz import streaming_decoder
+
+    captures = []
+
+    def capture(stream, sizes, frames=1):
+        captures.append((sizes, frames))
+        return SimpleNamespace(sizes=sizes, frames=frames)
+
+    monkeypatch.setattr(streaming_decoder, "StreamingDecodeGraphs", capture)
+    model = SimpleNamespace(
+        stream_decoder=SimpleNamespace(max_frames=33),
+        stream_graphs=None,
+        stream_prime_graphs=None,
+        stream_ref_context_frames=25,
+        stream_chunk_frames=25,
+        config=SimpleNamespace(tts_model_type=model_type),
+        vllm_config=SimpleNamespace(scheduler_config=SimpleNamespace(max_num_batched_tokens=512)),
+    )
+    sizes = [size for size in [1, 2, 4, 8, 16, 32, 64, 128] if size <= capacity]
+    Qwen3TTSTalkerForConditionalGeneration.capture_stream_decode_graphs(model, sizes)
+    expected = [(sizes, 1)]
+    if model_type == "base":
+        expected.append(([size for size in sizes if size * 25 <= 512], 25))
+    assert captures == expected
+    Qwen3TTSTalkerForConditionalGeneration.capture_stream_decode_graphs(model, sizes)
+    assert captures == expected
+
+
 @pytest.mark.cpu
 @pytest.mark.parametrize("skip_flags", [(False, False), (True, True), (False, True), (True, False)])
 @pytest.mark.parametrize("initial_frames", [1, 3])
@@ -124,3 +157,134 @@ def test_xvec_graphs_follow_request_first_audio_flags(initial_frames, time_major
             torch.testing.assert_close(got, want, atol=1e-4, rtol=1e-4)
             for key in ("ref_hidden", "ref_conv", "prefix_hidden", "suffix_quantized", "suffix_conv"):
                 torch.testing.assert_close(cache[key], expected_cache[key], atol=1e-5, rtol=1e-5)
+
+
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+@torch.inference_mode()
+def test_stream_decoder_preemption_restores_all_state_into_a_different_slot():
+    from tests.model_executor.models.qwen3_tts.test_time_major_decoder import _make_decoder
+    from vllm_omni.model_executor.models.qwen3_tts.tokenizer_12hz.streaming_decoder import StreamingCodecDecoder
+
+    decoder = _make_decoder().to(device=DEVICE, dtype=torch.bfloat16)
+    for name, param in decoder.named_parameters():
+        if name.endswith("embedding_sum"):
+            param.normal_()
+    decoder.config.head_dim = decoder.config.hidden_size // decoder.config.num_attention_heads
+    stream = StreamingCodecDecoder(decoder, num_slots=2, dtype=torch.bfloat16)
+    codes = torch.randint(0, 32, (1, 8, 2), device=DEVICE)
+    slot = torch.tensor([0], device=DEVICE, dtype=torch.int32)
+    for pos in range(7):
+        stream(codes[:, pos : pos + 1], slot, torch.tensor([pos], device=DEVICE, dtype=torch.int32))
+    saved = stream.save_slot(0)
+    expected = stream(codes[:, 7:], slot, torch.tensor([7], device=DEVICE, dtype=torch.int32)).clone()
+    # Both the freed slot and the destination can hold an unrelated request.
+    other_codes = (codes[:, :3] + 1) % 32
+    for pos in range(3):
+        stream(
+            other_codes[:, pos : pos + 1].expand(2, -1, -1).contiguous(),
+            torch.tensor([0, 1], device=DEVICE, dtype=torch.int32),
+            torch.full((2,), pos, device=DEVICE, dtype=torch.int32),
+        )
+    without_restore = stream(
+        codes[:, 7:],
+        torch.tensor([1], device=DEVICE, dtype=torch.int32),
+        torch.tensor([7], device=DEVICE, dtype=torch.int32),
+    ).clone()
+    assert not torch.equal(without_restore, expected)
+    stream.restore_slot(1, saved)
+    actual = stream(
+        codes[:, 7:],
+        torch.tensor([1], device=DEVICE, dtype=torch.int32),
+        torch.tensor([7], device=DEVICE, dtype=torch.int32),
+    )
+    torch.testing.assert_close(actual, expected, rtol=0, atol=0)
+
+
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+@pytest.mark.parametrize("use_graph", [False, True])
+@torch.inference_mode()
+def test_batched_reference_priming_matches_serial_decoder_continuation(use_graph):
+    from tests.model_executor.models.qwen3_tts.test_time_major_decoder import _make_decoder
+    from vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_talker import Qwen3TTSTalkerForConditionalGeneration
+    from vllm_omni.model_executor.models.qwen3_tts.tokenizer_12hz.streaming_decoder import (
+        StreamingCodecDecoder,
+        StreamingDecodeGraphs,
+    )
+
+    decoder = _make_decoder().to(device=DEVICE, dtype=torch.bfloat16)
+    for name, param in decoder.named_parameters():
+        if name.endswith("embedding_sum"):
+            param.normal_()
+    decoder.config.head_dim = decoder.config.hidden_size // decoder.config.num_attention_heads
+    stream = StreamingCodecDecoder(decoder, num_slots=5, dtype=torch.bfloat16)
+    graphs = StreamingDecodeGraphs(stream, [2], frames=25) if use_graph else None
+    model = Qwen3TTSTalkerForConditionalGeneration.__new__(Qwen3TTSTalkerForConditionalGeneration)
+    torch.nn.Module.__init__(model)
+    model.stream_decoder = stream
+    model.stream_prime_graphs = graphs
+    model.stream_chunk_frames = 25
+    refs = [torch.randint(0, 32, (length, 2), device=DEVICE) for length in [25, 3, 51, 25]]
+    slots = [2, 0, 3, 4]
+    frame = torch.randint(0, 32, (4, 1, 2), device=DEVICE)
+    slot_tensor = torch.tensor(slots, device=DEVICE, dtype=torch.int32)
+    positions = torch.tensor([len(ref) for ref in refs], device=DEVICE, dtype=torch.int32)
+    for slot, ref in zip(slots, refs, strict=True):
+        for t0 in range(0, len(ref), 25):
+            stream(
+                ref[None, t0 : t0 + 25].to(torch.int32).contiguous(),
+                torch.tensor([slot], device=DEVICE, dtype=torch.int32),
+                torch.tensor([t0], device=DEVICE, dtype=torch.int32),
+            )
+    expected = stream(frame, slot_tensor, positions).clone()
+    untouched = stream.save_slot(1)
+    # Reuse dirty slots; pos=0 must reset every causal layer for each group.
+    model.prime_stream_decoder(list(zip(slots, refs, strict=True)))
+    actual = stream(frame, slot_tensor, positions)
+    torch.testing.assert_close(actual, expected, rtol=0, atol=3e-5)
+    for got, want in zip(stream.save_slot(1), untouched, strict=True):
+        torch.testing.assert_close(got, want, rtol=0, atol=0)
+
+
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+@pytest.mark.parametrize("use_graph", [False, True])
+@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16])
+@torch.inference_mode()
+def test_stream_frames_match_fp32_exact_decoder(use_graph, dtype):
+    from tests.model_executor.models.qwen3_tts.test_time_major_decoder import _make_decoder
+    from vllm_omni.model_executor.models.qwen3_tts.tokenizer_12hz.streaming_decoder import (
+        StreamingCodecDecoder,
+        StreamingDecodeGraphs,
+    )
+
+    decoder = _make_decoder().to(device=DEVICE, dtype=dtype)
+    # The tiny HF default initialization shrinks the deep conv stack to
+    # near-zero PCM. Keep the oracle signal large enough to catch lost state.
+    for module in decoder.modules():
+        if isinstance(module, (torch.nn.Conv1d, torch.nn.ConvTranspose1d, torch.nn.Linear)):
+            torch.nn.init.kaiming_uniform_(module.weight, a=math.sqrt(5))
+            if module.bias is not None:
+                module.bias.uniform_(-0.01, 0.01)
+    for name, param in decoder.named_parameters():
+        if name.endswith("embedding_sum"):
+            param.normal_()
+    decoder.config.head_dim = decoder.config.hidden_size // decoder.config.num_attention_heads
+    reference = _make_decoder().to(device=DEVICE, dtype=torch.float32)
+    reference.load_state_dict(decoder.state_dict())
+    reference.precompute_snake_caches()
+    # Cross both the attention window and KV ring wrap, checking every frame.
+    codes = torch.randint(0, 32, (2, 136, 2), device=DEVICE)
+    expected = reference._forward_exact(codes.transpose(1, 2))[:, 0]
+    assert expected.square().sum() > 1e-3
+    stream = StreamingCodecDecoder(decoder, num_slots=3, dtype=dtype)
+    decode = StreamingDecodeGraphs(stream, [2]) if use_graph else stream
+    slots = torch.tensor([2, 0], device=DEVICE, dtype=torch.int32)
+    spf = int(stream.spf)
+    for frame in range(codes.shape[1]):
+        positions = torch.full((2,), frame, device=DEVICE, dtype=torch.int32)
+        actual = decode(codes[:, frame : frame + 1].contiguous().to(torch.int32), slots, positions).reshape(2, -1)
+        exact_frame = expected[:, frame * spf : (frame + 1) * spf]
+        relative_rms = (actual - exact_frame).square().mean(1).sqrt() / exact_frame.square().mean(1).sqrt()
+        # BF16 kernels are approximate; compare every frame to the independent
+        # FP32 full decoder, not only one streaming implementation to itself.
+        tolerance = 0.05 if dtype == torch.bfloat16 else 5e-4
+        assert torch.all(relative_rms < tolerance), (frame, relative_rms)

tests/model_executor/models/qwen3_tts/test_streaming_col2im_cuda.py

@@ -0,0 +1,57 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Streaming col2im preserves PCM inputs and ping-pong slot state."""
+
+import pytest
+import torch
+from vllm.triton_utils import triton
+
+from tests.helpers.mark import hardware_test
+from vllm_omni.model_executor.models.qwen3_tts.tokenizer_12hz.streaming_decoder import _col2im_kernel
+
+pytestmark = [pytest.mark.core_model, pytest.mark.tts]
+
+
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+@pytest.mark.parametrize("batch,frames,channels,rate", [(1, 1, 5, 2), (3, 3, 129, 3), (9, 1, 64, 4)])
+def test_col2im_partial_tiles_and_reordered_slot_state(batch: int, frames: int, channels: int, rate: int):
+    torch.manual_seed(23)
+    capacity = batch + 2
+    z = torch.randn(batch, frames, 2 * rate, channels).to(torch.bfloat16)
+    bias = torch.randn(channels).to(torch.bfloat16)
+    state = torch.randn(2, capacity, rate, channels).to(torch.bfloat16)
+    expected_state = state.clone()
+    slots = torch.arange(batch - 1, -1, -1, dtype=torch.int32)
+    positions = torch.arange(batch, dtype=torch.int32)
+    parity = torch.arange(capacity, dtype=torch.int32) % 2
+    expected = torch.empty(batch, frames, rate, channels, dtype=torch.bfloat16)
+    for b in range(batch):
+        slot = int(slots[b])
+        side = int(parity[slot])
+        for t in range(frames):
+            previous = z[b, t - 1, rate:] if t else state[side, slot] if positions[b] else torch.zeros(rate, channels)
+            expected[b, t] = (z[b, t, :rate].float() + previous.float() + bias.float()).to(torch.bfloat16)
+        expected_state[1 - side, slot] = z[b, -1, rate:]
+
+    z_gpu, bias_gpu, state_gpu = z.cuda(), bias.cuda(), state.cuda()
+    slots_gpu, positions_gpu, parity_gpu = slots.cuda(), positions.cuda(), parity.cuda()
+    actual = torch.empty_like(expected, device="cuda")
+    width = min(128, triton.next_power_of_2(channels))
+    rows = batch * frames
+    _col2im_kernel[(rows, rate, triton.cdiv(channels, width))](
+        z_gpu,
+        bias_gpu,
+        state_gpu,
+        slots_gpu,
+        positions_gpu,
+        parity_gpu,
+        actual,
+        frames,
+        capacity,
+        R=rate,
+        C=channels,
+        BC=width,
+        num_warps=4,
+    )
+    torch.testing.assert_close(actual.cpu(), expected, rtol=0, atol=0)
+    torch.testing.assert_close(state_gpu.cpu(), expected_state, rtol=0, atol=0)

tests/model_executor/models/test_breeze_tts_2_first_code_sampler.py

@@ -4,6 +4,7 @@
 
 import pytest
 import torch
+from vllm.platforms import current_platform
 
 from tests.helpers.mark import hardware_test
 from vllm_omni.model_executor.models.breeze_tts_2.depth_decoder import sample_logits
@@ -16,6 +17,10 @@ VOCAB = 2052
 EOS = 2051
 
 
+def _captured_generator_state_supported() -> bool:
+    return current_platform.is_cuda() and hasattr(torch.cuda.CUDAGraph, "register_generator_state")
+
+
 def _logits(count: int, *, device: str, seed: int = 100, vocab: int = VOCAB) -> list[torch.Tensor]:
     source = torch.Generator(device="cpu").manual_seed(seed)
     values = torch.randn((count, vocab), generator=source, dtype=torch.float32) * 3
@@ -164,6 +169,14 @@ def test_cuda_replay_preserves_tokens_rng_and_owned_outputs(monkeypatch) -> None
                     right[codebook].exponential_(generator=live_expected[index])
                 torch.testing.assert_close(left, right, atol=0, rtol=0)
             _assert_states(actual_generators, expected_generators)
+    if not _captured_generator_state_supported():
+        # ROCm deliberately uses the eager sampler because CUDA graph generator-state registration is an
+        # NVIDIA-only path. The loop above still verifies GPU tokens and every request-owned RNG stream.
+        assert not sampler._graphs
+        for value, expected in retained:
+            torch.testing.assert_close(value, expected, atol=0, rtol=0)
+        torch.testing.assert_close(torch.cuda.get_rng_state(), global_state, atol=0, rtol=0)
+        return
     # B1/3/16/17/31/32 must use exactly buckets 1/4/16/32.
     assert len(sampler._graphs) == 4
     rows = _logits(3, device="cuda", seed=2026)
@@ -189,6 +202,7 @@ def test_cuda_replay_preserves_tokens_rng_and_owned_outputs(monkeypatch) -> None
 @torch.inference_mode()
 def test_cuda_parameters_vocab_and_missing_graph_api(monkeypatch) -> None:
     sampler = BreezeFirstCodeSampler()
+    captures_generator_state = _captured_generator_state_supported()
     actual_generators = _generators([42, 43, 44], "cuda")
     expected_generators = _generators([42, 43, 44], "cuda")
     cases = ((0.9, 50, 1.0), (0.7, 100, 0.8), (1.2, 0, 0.95), (1.0, 1, 1.0), (1.0, 0, 0.5))
@@ -201,12 +215,12 @@ def test_cuda_parameters_vocab_and_missing_graph_api(monkeypatch) -> None:
         expected = _reference(rows, parameters, expected_generators)
         _assert_outputs(sampler.sample(rows, parameters, actual_generators), expected, rows)
         _assert_states(actual_generators, expected_generators)
-    assert len(sampler._graphs) == len(cases)
+    assert len(sampler._graphs) == (len(cases) if captures_generator_state else 0)
     rows = _logits(3, device="cuda", vocab=17)
     expected = _reference(rows, PARAMETERS, expected_generators)
     _assert_outputs(sampler.sample(rows, PARAMETERS, actual_generators), expected, rows)
     _assert_states(actual_generators, expected_generators)
-    assert len(sampler._graphs) == len(cases) + 1
+    assert len(sampler._graphs) == (len(cases) + 1 if captures_generator_state else 0)
 
     class UnavailableGraph:
         def __init__(self, *_args, **_kwargs):
@@ -243,7 +257,7 @@ def test_cuda_shared_generator_uses_sequential_eager_fallback(monkeypatch) -> No
     # The duplicate-generator guard must precede lookup of an existing graph.
     sampler.sample(rows, PARAMETERS, _generators([100, 101, 102], "cuda"))
     cached = dict(sampler._graphs)
-    assert len(cached) == 1
+    assert len(cached) == (1 if _captured_generator_state_supported() else 0)
     actual = torch.Generator(device="cuda").manual_seed(2**64 - 23)
     reference = torch.Generator(device="cuda").manual_seed(2**64 - 23)
     multinomial = torch.multinomial
@@ -270,6 +284,7 @@ def test_cuda_shared_generator_uses_sequential_eager_fallback(monkeypatch) -> No
 @torch.inference_mode()
 def test_cuda_cache_limit_falls_back_without_evicting_warmed_graphs(monkeypatch) -> None:
     sampler = BreezeFirstCodeSampler()
+    captures_generator_state = _captured_generator_state_supported()
     rows = _logits(1, device="cuda")
     actual_generators = _generators([42], "cuda")
     expected_generators = _generators([42], "cuda")
@@ -279,7 +294,7 @@ def test_cuda_cache_limit_falls_back_without_evicting_warmed_graphs(monkeypatch)
         _assert_outputs(sampler.sample(rows, parameters, actual_generators), expected, rows)
         _assert_states(actual_generators, expected_generators)
     cached = dict(sampler._graphs)
-    assert len(cached) == 32
+    assert len(cached) == (32 if captures_generator_state else 0)
     parameters = (0.9, 33, 1.0)
     expected = _reference(rows, parameters, expected_generators)
     multinomial = torch.multinomial
@@ -298,6 +313,11 @@ def test_cuda_cache_limit_falls_back_without_evicting_warmed_graphs(monkeypatch)
     assert all(sampler._graphs[key] is value for key, value in cached.items())
     expected = _reference(rows, (0.9, 1, 1.0), expected_generators)
 
+    if not captures_generator_state:
+        _assert_outputs(sampler.sample(rows, (0.9, 1, 1.0), actual_generators), expected, rows)
+        _assert_states(actual_generators, expected_generators)
+        return
+
     def unexpected_multinomial(*_args, **_kwargs):
         raise AssertionError("Cache exhaustion must retain existing captured entries")
 

tests/model_executor/models/test_breeze_tts_2_graphs.py

@@ -9,7 +9,11 @@ from transformers.models.t5gemma2.modeling_t5gemma2 import T5Gemma2TextEncoder
 from vllm.platforms import current_platform
 
 from tests.helpers.mark import hardware_test
-from vllm_omni.model_executor.models.breeze_tts_2.depth_decoder import BreezeDepthDecoder, sample_logits
+from vllm_omni.model_executor.models.breeze_tts_2.depth_decoder import (
+    BreezeDepthDecoder,
+    _depth_layer_compile_options,
+    sample_logits,
+)
 from vllm_omni.model_executor.models.breeze_tts_2.modeling_breeze import BreezeForConditionalGeneration
 from vllm_omni.model_executor.models.breeze_tts_2.text_encoder_graph import (
     BreezeTextEncoderCompiled,
@@ -19,6 +23,17 @@ from vllm_omni.model_executor.models.breeze_tts_2.text_encoder_graph import (
 pytestmark = [pytest.mark.core_model]
 
 
+@pytest.mark.cpu
+@pytest.mark.parametrize(("is_rocm", "expected"), [(True, False), (False, True)])
+def test_depth_layer_compile_options(monkeypatch, is_rocm: bool, expected: bool) -> None:
+    monkeypatch.setattr(current_platform, "is_rocm", lambda: is_rocm)
+
+    assert _depth_layer_compile_options() == {
+        "epilogue_fusion": False,
+        "max_autotune": expected,
+    }
+
+
 @pytest.fixture
 def full_precision_matmul():
     # Different padded shapes can select different TF32 kernels. Compare
@@ -241,6 +256,7 @@ def test_depth_graph_buckets_bound_shapes_and_preserve_requests(full_precision_m
     seeds[0], seeds[-1] = 2**63 + 17, 2**64 - 1
     generators = [torch.Generator(device="cuda").manual_seed(seed) for seed in seeds]
     reference_generators = [torch.Generator(device="cuda").manual_seed(seed) for seed in seeds]
+    captures_generator_state = current_platform.is_cuda() and hasattr(torch.cuda.CUDAGraph, "register_generator_state")
 
     expected_greedy = torch.cat(
         [
@@ -314,7 +330,7 @@ def test_depth_graph_buckets_bound_shapes_and_preserve_requests(full_precision_m
             branches = 2 if guidance_scale != 1 else 1
             graph_key = (request_bucket * branches, branches, temperature == 0)
             with monkeypatch.context() as rng_guard:
-                if temperature > 0 and graph_key in depth._graphs:
+                if temperature > 0 and graph_key in depth._graphs and captures_generator_state:
                     # A warmed replay must generate noise within the CUDA
                     # graph, without dispatching new RNG kernels from Python.
                     rng_guard.setattr(torch.Tensor, "exponential_", reject_eager_noise)
@@ -349,7 +365,7 @@ def test_depth_graph_buckets_bound_shapes_and_preserve_requests(full_precision_m
     expected_keys.update((batch, 2, greedy) for batch in (2, 8, 16, 32) for greedy in (False, True))
     assert set(depth._graphs) == expected_keys
     for (batch_bucket, branches, greedy), entry in depth._graphs.items():
-        if greedy:
+        if greedy or not captures_generator_state:
             assert entry.noise_generators is None
         else:
             assert entry.noise_generators is not None

tests/model_executor/models/test_talker_first_audio.py

@@ -0,0 +1,70 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from types import SimpleNamespace
+
+import pytest
+
+from vllm_omni.model_executor.models.qwen3_omni.first_frame_decoder import talker_first_audio_enabled
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+@pytest.mark.parametrize(
+    ("env", "extra", "async_chunk", "v2", "expected"),
+    [
+        (None, {}, True, True, False),
+        ("1", {}, True, True, False),
+        (None, {"talker_first_audio": True, "codec_chunk_ramp": [1, 2]}, True, True, True),
+        ("0", {"talker_first_audio": True, "codec_chunk_ramp": [1, 2]}, True, True, False),
+        (None, {"talker_first_audio": True, "codec_chunk_ramp": [4, 8]}, True, True, False),
+        (None, {"talker_first_audio": True, "initial_codec_chunk_frames": 1}, True, True, True),
+        (None, {"talker_first_audio": True, "initial_codec_chunk_frames": 4}, True, True, False),
+        (None, {"talker_first_audio": True, "codec_chunk_ramp": [1]}, True, True, False),
+        (None, {"talker_first_audio": True, "codec_chunk_ramp": [1, 2]}, False, True, False),
+        (None, {"talker_first_audio": True, "codec_chunk_ramp": [1, 2]}, True, False, False),
+    ],
+)
+def test_talker_first_audio_default_and_conditions(monkeypatch, env, extra, async_chunk, v2, expected):
+    if env is None:
+        monkeypatch.delenv("VLLM_OMNI_TALKER_FIRST_AUDIO", raising=False)
+    else:
+        monkeypatch.setenv("VLLM_OMNI_TALKER_FIRST_AUDIO", env)
+    from vllm_omni.platforms import current_omni_platform
+
+    monkeypatch.setattr(current_omni_platform, "is_cuda", lambda: True)
+    config = SimpleNamespace(
+        model_config=SimpleNamespace(
+            async_chunk=async_chunk, use_v2_model_runner=v2, stage_connector_config={"extra": extra}
+        ),
+        device_config=SimpleNamespace(device="cuda"),
+        parallel_config=SimpleNamespace(
+            tensor_parallel_size=1, pipeline_parallel_size=1, distributed_executor_backend=None
+        ),
+        cache_config=SimpleNamespace(enable_prefix_caching=False),
+    )
+    assert talker_first_audio_enabled(config) is expected
+    for section, field, value in [
+        ("device_config", "device", "cpu"),
+        ("parallel_config", "tensor_parallel_size", 2),
+        ("parallel_config", "pipeline_parallel_size", 2),
+        ("parallel_config", "distributed_executor_backend", "mp"),
+        ("cache_config", "enable_prefix_caching", True),
+    ]:
+        with monkeypatch.context() as patch:
+            patch.setattr(getattr(config, section), field, value)
+            assert not talker_first_audio_enabled(config)
+    monkeypatch.setattr(current_omni_platform, "is_cuda", lambda: False)
+    assert not talker_first_audio_enabled(config)
+
+
+@pytest.mark.parametrize("option", [False, True])
+def test_tts_requires_connector_opt_in(monkeypatch, option):
+    from vllm_omni.model_executor.models.qwen3_tts import first_audio
+    from vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_code_predictor_vllm import (
+        Qwen3TTSTalkerCodePredictorForConditionalGenerationVLLM as Predictor,
+    )
+
+    monkeypatch.setattr(Predictor, "_stage_connector_extra_config", lambda _: {"talker_first_audio": option})
+    monkeypatch.setattr(first_audio, "supports_talker_first_audio", lambda _: True)
+    assert first_audio.talker_first_audio_enabled(None) is option