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

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

tests/worker_v2/test_capture_query_bound_gpu.py

@@ -45,6 +45,8 @@ def _check_prefill():
     kvconfig = KVCacheConfig(num_blocks=1024, kv_cache_tensors=[], kv_cache_groups=[KVCacheGroupSpec(["attn"], spec)])
     state = OmniModelState.__new__(OmniModelState)
     state.model_config = SimpleNamespace(max_model_len=256)
+    state.vllm_config = c
+    state.max_model_len = 256
     state.supports_mm_inputs = False
     with set_current_vllm_config(c):
         group = AttentionGroup(FlashAttentionBackend, ["attn"], spec, 0)

tests/worker_v2/test_decode_query_bound_gpu.py

@@ -0,0 +1,99 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Real FA3 graph replay must preserve the decode GQA layout."""
+
+from dataclasses import dataclass, replace
+
+import numpy as np
+import pytest
+import torch
+from transformers import Qwen2Config
+from vllm.config import ModelConfig, VllmConfig, set_current_vllm_config
+from vllm.config.compilation import CUDAGraphMode
+from vllm.v1.attention.backends.flash_attn import FlashAttentionBackend, FlashAttentionImpl
+from vllm.v1.kv_cache_interface import FullAttentionSpec, KVCacheConfig, KVCacheGroupSpec
+from vllm.v1.worker.gpu.input_batch import InputBatch, InputBuffers
+from vllm.v1.worker.utils import AttentionGroup
+
+from tests.helpers.mark import hardware_test
+from vllm_omni.worker_v2.model_states.omni_model_state import OmniModelState
+
+pytestmark = [pytest.mark.core_model]
+
+
+@dataclass
+class _AttentionScales:
+    _q_scale: torch.Tensor
+    _k_scale: torch.Tensor
+    _v_scale: torch.Tensor
+
+
+@hardware_test(res={"cuda": "H100"}, num_cards=1)
+@pytest.mark.parametrize("graph_mode", [CUDAGraphMode.FULL_AND_PIECEWISE, CUDAGraphMode.FULL_DECODE_ONLY])
+def test_full_decode_graph_matches_fa3_aot_schedule(graph_mode, tmp_path):
+    if not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] < 9:
+        pytest.skip("requires FA3 on Hopper or newer")
+    torch.manual_seed(17)
+    device = torch.device("cuda:0")
+    c = VllmConfig()
+    Qwen2Config(
+        architectures=["Qwen2ForCausalLM"],
+        hidden_size=1024,
+        intermediate_size=2048,
+        num_hidden_layers=1,
+        num_attention_heads=16,
+        num_key_value_heads=2,
+        max_position_embeddings=2048,
+    ).save_pretrained(tmp_path)
+    c.model_config = ModelConfig(model=str(tmp_path), skip_tokenizer_init=True, dtype="bfloat16")
+    c.scheduler_config.max_num_seqs = 64
+    c.compilation_config.cudagraph_mode = graph_mode
+    c.compilation_config.max_cudagraph_capture_size = 128
+    c.attention_config.flash_attn_version = 3
+    spec = FullAttentionSpec(block_size=16, num_kv_heads=2, head_size=64, dtype=torch.bfloat16)
+    kvconfig = KVCacheConfig(num_blocks=1024, kv_cache_tensors=[], kv_cache_groups=[KVCacheGroupSpec(["attn"], spec)])
+    state = OmniModelState.__new__(OmniModelState)
+    state.vllm_config = c
+    state.max_model_len = 2048
+    state.supports_mm_inputs = False
+    with set_current_vllm_config(c):
+        group = AttentionGroup(FlashAttentionBackend, ["attn"], spec, 0)
+        group.create_metadata_builders(c, device)
+        impl = FlashAttentionImpl(16, 64, 64**-0.5, 2, None, None, "auto")
+        assert impl.vllm_flash_attn_version == 3
+        buffers = InputBuffers(64, 128, device)
+        batch = InputBatch.make_dummy(8, 8, buffers, max_query_len=None)
+        tables = torch.arange(1024, dtype=torch.int32, device=device).reshape(8, 128)
+        slots = torch.zeros((1, 8), dtype=torch.int64, device=device)
+        cache = torch.randn(1024, 2, 16, 128, device=device, dtype=torch.bfloat16)
+        q = torch.randn(8, 16, 64, device=device, dtype=torch.bfloat16)
+        k = v = torch.empty(8, 2, 64, device=device, dtype=torch.bfloat16)
+        layer = _AttentionScales(
+            _q_scale=torch.ones((), device=device),
+            _k_scale=torch.ones((), device=device),
+            _v_scale=torch.ones((), device=device),
+        )
+        md = state.prepare_attn(batch, CUDAGraphMode.FULL, (tables,), slots, [[group]], kvconfig, for_capture=True)[
+            "attn"
+        ]
+        out = torch.full_like(q, float("nan"))
+        stream = torch.cuda.Stream()
+        stream.wait_stream(torch.cuda.current_stream())
+        with torch.cuda.stream(stream):
+            for _ in range(3):
+                impl.forward(layer, q, k, v, cache, md, out)
+        stream.synchronize()
+        graph = torch.cuda.CUDAGraph()
+        with torch.cuda.graph(graph, stream=stream):
+            impl.forward(layer, q, k, v, cache, md, out)
+        lengths = np.arange(37, 45, dtype=np.int32)
+        batch.seq_lens.copy_(torch.tensor(lengths, device=device))
+        runtime = replace(batch, max_query_len=1, seq_lens_cpu_upper_bound=torch.tensor(lengths))
+        eager_md = state.prepare_attn(runtime, CUDAGraphMode.FULL, (tables,), slots, [[group]], kvconfig)["attn"]
+        expected = torch.empty_like(q)
+        impl.forward(layer, q, k, v, cache, eager_md, expected)
+        state.prepare_attn(runtime, CUDAGraphMode.FULL, (tables,), slots, [[group]], kvconfig)
+        graph.replay()
+        torch.accelerator.synchronize()
+        assert md.max_query_len == 1
+        torch.testing.assert_close(out, expected, atol=0.002, rtol=0.002)

tests/worker_v2/test_first_audio_sender.py

@@ -235,3 +235,19 @@ def test_sender_copy_failure_reports_all_prepared_routes():
     assert client == 3
     assert [(o.request_id, o.finish_reason) for o in outputs.outputs] == [("r", FinishReason.ERROR)]
     assert output_queue.empty()
+
+
+def test_final_codec_delivery_omits_upstream_marker_and_freezes_route():
+    from vllm_omni.data_entry_keys import FIRST_AUDIO_KEY
+
+    outputs: queue.Queue = queue.Queue()
+    scheduler = _scheduler(r=3)
+    delivery = engine_output_queue_sink(outputs, scheduler, upstream_first_audio=False).prepare(["r"])
+    scheduler.requests.clear()
+    pcm = torch.ones(4)
+    delivery(["r"], [pcm], torch.tensor(24000))
+    client, batch = outputs.get_nowait()
+    assert client == 3
+    assert FIRST_AUDIO_KEY not in batch.outputs[0].multimodal_output
+    assert batch.outputs[0].multimodal_output["model_outputs"] is pcm
+    assert batch.outputs[0].finish_reason is None

tests/worker_v2/test_index_staging.py

@@ -0,0 +1,21 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+import pytest
+import torch
+
+from vllm_omni.utils.device_copy import index_to_device
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cuda]
+
+
+def test_index_staging_survives_later_copies_and_graph_replay():
+    stream = torch.cuda.Stream()
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.stream(stream):
+        with torch.cuda.graph(graph, stream=stream):
+            captured = index_to_device([17] * 6, "cuda")
+    for _ in range(32):
+        index_to_device([99] * 6, "cuda")
+    graph.replay()
+    torch.accelerator.synchronize()
+    assert captured.tolist() == [17] * 6

tests/worker_v2/test_omni_ar_model_runner.py

@@ -18,6 +18,7 @@ from vllm.v1.worker.gpu.sample.prompt_logprob import PromptLogprobsWorker
 
 import vllm_omni.worker_v2.omni_ar_model_runner as omni_ar_model_runner
 from tests.helpers.mark import hardware_test
+from vllm_omni.model_executor.output_snapshot import PackedOutputSnapshot
 from vllm_omni.worker_v2.omni_ar_model_runner import OmniARModelRunner, OmniAsyncOutput
 from vllm_omni.worker_v2.output_snapshot import pack_output_snapshot
 
@@ -28,6 +29,9 @@ class _FakeStream:
     def wait_stream(self, _stream) -> None:
         pass
 
+    def wait_event(self, _event) -> None:
+        pass
+
 
 class _FakeEvent:
     def record(self, _stream) -> None:
@@ -224,8 +228,10 @@ def test_chunked_fixed_token_scores_are_copied_from_cuda() -> None:
 def test_async_mm_snapshot_owns_output_until_copy_finishes() -> None:
     runner = OmniARModelRunner.__new__(OmniARModelRunner)
     runner.model_config = SimpleNamespace(async_chunk=True)
+    runner.model = SimpleNamespace()
     runner._async_mm_snapshot_slots, runner._async_mm_snapshot_events = [{}], [None]
     runner._async_mm_snapshot_pending, runner._async_mm_snapshot_cursor = [False], 0
+    runner._last_multimodal_snapshot_slot = None
     waited: list[object] = []
     runner.main_stream = SimpleNamespace(wait_event=waited.append)
     source = torch.tensor([[7, 8]], dtype=torch.long)
@@ -243,6 +249,35 @@ def test_async_mm_snapshot_owns_output_until_copy_finishes() -> None:
     assert waited == [copy_event]
 
 
+def test_producer_snapshot_is_not_repacked() -> None:
+    runner = OmniARModelRunner.__new__(OmniARModelRunner)
+    runner.model_config = SimpleNamespace(async_chunk=True)
+    runner.model = SimpleNamespace()
+    runner._async_mm_snapshot_slots, runner._async_mm_snapshot_events = [{}], [None]
+    runner._async_mm_snapshot_pending, runner._async_mm_snapshot_cursor = [False], 0
+    runner._last_multimodal_snapshot_slot = None
+    source = torch.tensor([[7, 8]], dtype=torch.float32)
+    snapshot = pack_output_snapshot({"audio": source}, {}, max_buckets=1)
+    assert isinstance(snapshot, PackedOutputSnapshot)
+
+    retained = runner._retain_multimodal_outputs(snapshot)
+
+    assert retained is snapshot
+    assert runner._last_multimodal_snapshot_slot is None
+    assert runner._async_mm_snapshot_pending == [False]
+
+
+def test_fresh_per_step_outputs_are_not_repacked() -> None:
+    runner = OmniARModelRunner.__new__(OmniARModelRunner)
+    runner.model_config = SimpleNamespace(async_chunk=True)
+    runner.model = SimpleNamespace(mm_outputs_fresh_per_step=True)
+    runner._last_multimodal_snapshot_slot = None
+    outputs = {"model_outputs": torch.ones(2, 4)}
+
+    assert runner._retain_multimodal_outputs(outputs) is outputs
+    assert runner._last_multimodal_snapshot_slot is None
+
+
 def test_snapshot_slots_bounded_by_shape_and_packed_grouping_isolation(monkeypatch) -> None:
     slot: dict[tuple[Any, ...], torch.Tensor] = {}
     omni_ar_model_runner._copy_mm_to_snapshot_slot(torch.ones(1, 2), slot)
@@ -399,6 +434,231 @@ def test_request_reference_codes_preserve_local_axis(prefill_first, padded):
         assert torch.equal(outputs[i]["codes.audio"], codes[offsets[i] : offsets[i + 1]])
 
 
+def test_stream_audio_cancel_rejects_inflight_output_and_id_reuse() -> None:
+    from vllm_omni.worker_v2.streaming_audio import StreamingAudioBuffer
+
+    buffer = StreamingAudioBuffer(samples_per_frame=4, chunk_frames=3)
+    cancelled = buffer.add("request")
+    buffer.finish({"request"})
+    replacement = buffer.add("request")
+    assert cancelled.push(torch.ones(4), False) is None
+    assert not cancelled.pending
+    assert replacement.push(torch.full((4,), 2.0), False).tolist() == [2.0] * 4
+    buffer.finish({"request"})
+    assert not buffer.requests
+
+
+def test_stream_audio_length_end_flushes_first_and_partial_chunks() -> None:
+    from vllm_omni.worker_v2.streaming_audio import StreamingAudioBuffer, StreamingAudioOutput
+
+    buffer = StreamingAudioBuffer(samples_per_frame=4, chunk_frames=3)
+    state = buffer.add("request")
+    output = StreamingAudioOutput(
+        torch.ones(1, 4),
+        torch.tensor([True]),
+        torch.tensor(24000),
+        None,
+        np.array([0, 1]),
+        [state],
+        np.array([True]),
+    )
+    assert output.get_output()[0]["model_outputs"].numel() == 4
+    assert not state.active
+    state = buffer.add("another")
+    assert state.push(torch.ones(4), False).numel() == 4
+    assert state.push(torch.ones(4), False) is None
+    assert state.push(torch.ones(4), True).numel() == 8
+
+
+@pytest.mark.parametrize("dtype", [torch.float16, torch.float32, torch.float64, torch.bfloat16])
+def test_stream_audio_batch_views_retain_pcm_after_output_release(dtype):
+    from vllm_omni.worker_v2.streaming_audio import StreamingAudioBuffer, StreamingAudioOutput
+
+    buffer = StreamingAudioBuffer(samples_per_frame=4, chunk_frames=3)
+    states = [buffer.add("a"), buffer.add("b")]
+    expected: list[list[torch.Tensor]] = [[], []]
+    emitted: list[list[torch.Tensor]] = [[], []]
+    for step in range(5):
+        # Mixed prefill/decode spans, invalid rows, and a final partial chunk.
+        wav = (torch.arange(16).reshape(4, 4) + 100 * step).to(dtype)
+        valid = torch.tensor([False, True, True, step != 4])
+        output = StreamingAudioOutput(
+            wav,
+            valid,
+            torch.tensor(24000),
+            None,
+            np.array([0, 3, 4]),
+            states,
+            np.array([step == 4, step == 4]),
+        )
+        for i, (start, end) in enumerate([(0, 3), (3, 4)]):
+            expected[i].append(wav[start:end][valid[start:end]].reshape(-1).clone())
+        for i, payload in enumerate(output.get_output()):
+            if payload is not None:
+                emitted[i].append(payload["model_outputs"])
+        del output, wav  # Pending views must keep the old host allocation alive.
+    for i in range(2):
+        assert torch.equal(torch.cat(emitted[i]), torch.cat(expected[i]))
+        assert torch.cat(emitted[i]).dtype == dtype
+        assert not states[i].active and not states[i].pending
+        assert states[i].pending_samples == 0
+
+
+def test_stream_audio_numpy_direct_first_cancel_and_mixed_fallback():
+    from vllm_omni.worker_v2.streaming_audio import StreamingAudioBuffer
+
+    buffer = StreamingAudioBuffer(samples_per_frame=4, chunk_frames=3)
+    state = buffer.add("r")
+    state.accept_first_audio()
+    assert state.push(np.arange(8, dtype=np.float32), False) is None
+    assert state.pending_samples == 4
+    # A dtype fallback can arrive while a NumPy frame is pending.
+    chunk = state.push(torch.arange(8, 12, dtype=torch.bfloat16), True)
+    assert chunk.tolist() == list(range(4, 12))
+    assert state.pending_samples == 0
+    buffer.finish({"r"})
+    replacement = buffer.add("r")
+    assert state.push(np.ones(4, dtype=np.float32), False) is None
+    assert replacement.push(np.full(4, 99, dtype=np.float32), False).tolist() == [99] * 4
+    replacement.push(np.ones(4, dtype=np.float32), False)
+    assert replacement.pending_samples == 4
+    buffer.finish({"r"})
+    assert replacement.pending_samples == 0 and not replacement.pending
+
+
+def test_stream_audio_prefill_length_end_reaches_output(mocker) -> None:
+    from vllm.v1.worker.gpu.input_batch import InputBatch
+    from vllm.v1.worker.gpu.states import RequestState
+
+    from vllm_omni.worker_v2.model_states.eager_mtp import EagerMTPState
+    from vllm_omni.worker_v2.streaming_audio import StreamingAudioBuffer
+
+    owner = mocker.Mock()
+    owner.vllm_config.model_config.max_model_len = 4096
+    owner._stream_decode_event = None
+    eager = EagerMTPState(owner)
+    eager._audio_buffer = StreamingAudioBuffer(4, 25)
+    eager._audio_buffer.add("request")
+    eager._frame_requests = {"request"}
+    batch = InputBatch.__new__(InputBatch)
+    batch.num_reqs, batch.req_ids = 1, ["request"]
+    batch.num_computed_tokens_np = np.array([0])
+    batch.num_scheduled_tokens = np.array([3])
+    batch.idx_mapping_np = np.array([0])
+    batch.query_start_loc_np = np.array([0, 3])
+    batch.is_prefilling_np = np.array([True])
+    states = RequestState.__new__(RequestState)
+    states.max_seq_len = np.array([4])
+    mm = {
+        "model_outputs": torch.ones(3, 4),
+        "meta": {"codec_frame_valid": torch.tensor([0, 0, 1])},
+        "sr": [torch.tensor(24000)],
+    }
+    output = eager.prepare_audio_output(batch, states, mm)
+    assert output.get_output()[0]["model_outputs"].numel() == 4
+    assert "model_outputs" not in mm
+
+
+def test_stream_audio_preemption_preserves_position_rng_and_pending_pcm(mocker) -> None:
+    from vllm_omni.worker_v2.model_states.eager_mtp import EagerMTPState, TalkerInputs
+    from vllm_omni.worker_v2.streaming_audio import StreamingAudioBuffer
+
+    owner = mocker.Mock()
+    owner._stream_pos = {"request": 37}
+    owner._eager_embeds = torch.arange(18).reshape(6, 3).float()
+    owner._eager_ready = {}
+    owner.intermediate_buffer.buffers = [{} for _ in range(6)]
+    generator = torch.Generator().manual_seed(42)
+    owner._mtp_generators = {"request": generator}
+    saved = [torch.arange(6)]
+    owner.model.stream_decoder.save_slot.return_value = saved
+    eager = EagerMTPState(owner)
+    eager._talker_inputs["request"] = TalkerInputs(torch.ones(20, 3), 10)
+    eager._side_stream = mocker.Mock()
+    eager._audio_buffer = StreamingAudioBuffer(4, 25)
+    request = eager._audio_buffer.add("request")
+    request.push(torch.ones(4), False)
+    request.push(torch.ones(4), False)
+    eager.suspend_audio("request", 2)
+    owner._stream_pos.clear()
+    owner._mtp_generators.clear()
+    eager.resume_audio("request", 5)
+    assert owner._stream_pos["request"] == 37
+    assert owner._mtp_generators["request"] is generator
+    assert eager._restore_audio["request"] is saved
+    assert sum(part.numel() for part in request.pending) == 4
+    # A second preemption can happen before the restored slot is used.
+    eager.suspend_audio("request", 5)
+    owner.model.stream_decoder.save_slot.assert_called_once_with(2)
+    eager.finish_audio({"request"})
+    assert not eager._suspended_audio and not eager._restore_audio
+    assert not eager._audio_buffer.requests and not request.active
+
+
+def test_stream_audio_replay_uses_exact_inputs_and_consumes_last_frame_once(mocker):
+    from vllm_omni.worker_v2.model_states.eager_mtp import EagerMTPState, TalkerInputs
+
+    owner = mocker.Mock()
+    owner._eager_embeds = torch.tensor([[20.0, 30.0]])
+    owner.model.preprocess.return_value = (None, None, {"mtp_inputs": (None, torch.tensor([[2.0, 3.0]]))})
+    owner.intermediate_buffer.buffers = [{"req_id": "request"}]
+    eager = EagerMTPState(owner)
+    eager._restore_audio["request"] = []
+    history = torch.arange(8).reshape(4, 2).float()
+    eager._talker_inputs["request"] = TalkerInputs(history, 4)
+    first = torch.empty(2, 2)
+    assert eager.replay_inputs("request", 0, 0, torch.ones(2), first)
+    torch.testing.assert_close(first, history[:2])
+    owner.model.preprocess.assert_not_called()
+    last = torch.empty(3, 2)
+    assert eager.replay_inputs("request", 0, 2, torch.ones(3), last)
+    torch.testing.assert_close(last, torch.cat((history[2:], torch.tensor([[22.0, 33.0]]))))
+    owner.model.preprocess.assert_called_once()
+    with pytest.raises(RuntimeError, match="exceeds saved inputs"):
+        eager.replay_inputs("request", 0, 4, torch.ones(2), torch.empty(2, 2))
+
+
+def test_stream_reference_priming_groups_lengths_without_padding_or_slot_aliasing(mocker):
+    from vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_talker import Qwen3TTSTalkerForConditionalGeneration
+
+    primes = [(4, torch.full((26, 2), 4)), (1, torch.full((3, 2), 1)), (7, torch.full((26, 2), 7))]
+    calls = []
+
+    def stream(codes, slots, pos):
+        calls.append((codes.clone(), slots.clone(), pos.clone()))
+
+    decoder = mocker.Mock(side_effect=stream, device=torch.device("cpu"))
+    model = mocker.Mock(stream_decoder=decoder, stream_prime_graphs=None, stream_chunk_frames=25)
+    Qwen3TTSTalkerForConditionalGeneration.prime_stream_decoder(model, primes)
+    assert len(calls) == 3
+    for (codes, slots, pos), frames, indices, position in zip(
+        calls, [25, 1, 3], [[4, 7], [4, 7], [1]], [0, 25, 0], strict=True
+    ):
+        assert codes.shape == (len(indices), frames, 2)
+        assert codes.dtype == slots.dtype == pos.dtype == torch.int32
+        assert slots.tolist() == indices and pos.tolist() == [position] * len(indices)
+        for row, idx in enumerate(indices):
+            assert torch.all(codes[row] == idx)
+    Qwen3TTSTalkerForConditionalGeneration.prime_stream_decoder(model, [])
+    assert len(calls) == 3
+
+
+@pytest.mark.parametrize("context_frames", [25, 72])
+def test_model_owned_reference_context_respects_configured_boundary(mocker, context_frames):
+    from vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_talker import Qwen3TTSTalkerForConditionalGeneration
+
+    model = mocker.Mock(stream_ref_context_frames=context_frames, talker_config=mocker.Mock(num_code_groups=16))
+    reference = torch.arange(100 * 16).reshape(100, 16)
+    info = {"codes": {"ref": reference}}
+    got = Qwen3TTSTalkerForConditionalGeneration.get_stream_ref_context(model, info)
+    torch.testing.assert_close(got, reference[-context_frames:])
+    assert info["codes"]["ref"] is reference
+    assert Qwen3TTSTalkerForConditionalGeneration.get_stream_ref_context(model, {}) is None
+    assert (
+        Qwen3TTSTalkerForConditionalGeneration.get_stream_ref_context(model, {"codes": {"ref": torch.empty(0)}}) is None
+    )
+
+
 @pytest.mark.parametrize("streaming", [False, True])
 def test_model_owned_audio_finalizer_runs_after_copy_without_hidden(monkeypatch, streaming):
     monkeypatch.setattr(torch.cuda, "set_stream", lambda _stream: None)
@@ -434,6 +694,136 @@ def test_model_owned_audio_finalizer_runs_after_copy_without_hidden(monkeypatch,
     assert (result.pooler_output is None) == streaming
 
 
+def test_stream_audio_history_replays_exact_inputs_after_batch_reordering(mocker):
+    from vllm.v1.worker.gpu.input_batch import InputBatch
+
+    from vllm_omni.worker_v2.model_states.eager_mtp import EagerMTPState
+
+    owner = mocker.Mock()
+    owner.vllm_config.model_config.max_model_len = 12
+    eager = EagerMTPState(owner)
+    first = torch.arange(9, dtype=torch.float32).reshape(3, 3)
+    second = torch.arange(9, 18, dtype=torch.float32).reshape(3, 3)
+    expected = {"a": torch.cat((first[:2], second[2:])), "b": torch.cat((first[2:], second[:2]))}
+    batch = InputBatch.__new__(InputBatch)
+    batch.req_ids = ["a", "b"]
+    batch.query_start_loc_np = np.array([0, 2, 3])
+    batch.num_computed_tokens_np = np.array([0, 0])
+    eager.record_inputs(batch, first)
+    batch.req_ids = ["b", "a"]
+    batch.num_computed_tokens_np = np.array([1, 2])
+    eager.record_inputs(batch, second)
+    # Later input-buffer reuse must not change the conditioned replay history.
+    first.fill_(-1)
+    second.fill_(-2)
+    for request_id, wanted in expected.items():
+        eager._restore_audio[request_id] = []
+        replayed = torch.empty_like(wanted)
+        assert eager.replay_inputs(request_id, 0, 0, torch.zeros(3, dtype=torch.long), replayed)
+        torch.testing.assert_close(replayed, wanted)
+
+
+@pytest.mark.parametrize("valid", [False, True])
+def test_stream_audio_direct_first_requires_only_real_pcm_on_terminal(valid):
+    from vllm_omni.data_entry_keys import FIRST_AUDIO_REQUIRED_KEY
+    from vllm_omni.worker_v2.streaming_audio import StreamingAudioBuffer, StreamingAudioOutput
+
+    buffer = StreamingAudioBuffer(samples_per_frame=4, chunk_frames=3)
+    state = buffer.add("r")
+    state.accept_first_audio()
+    first = StreamingAudioOutput(
+        torch.arange(4).reshape(1, 4),
+        torch.tensor([valid]),
+        torch.tensor(24000),
+        None,
+        np.array([0, 1]),
+        [state],
+        np.array([True]),
+    )
+    payload = first.get_output()[0]
+    if valid:
+        assert payload == {FIRST_AUDIO_REQUIRED_KEY: torch.tensor(True)}
+    else:
+        assert payload is None  # EOS has no promised first PCM to wait for.
+    assert not state.active and not state.pending
+    buffer.finish({"r"})
+    replacement = buffer.add("r")
+    assert replacement.push(torch.full((4,), 99), False).tolist() == [99] * 4
+
+
+def test_stream_audio_direct_first_drops_one_frame_and_retains_partial_suffix():
+    from vllm_omni.worker_v2.streaming_audio import AudioRequest
+
+    state = AudioRequest(samples_per_frame=4, chunk_frames=3)
+    state.accept_first_audio()
+    assert state.push(torch.arange(4), False) is None
+    assert state.first_audio_required and state.emitted
+    assert state.push(torch.arange(4, 8), False) is None
+    chunk = state.push(torch.arange(8, 12), True)
+    assert chunk.tolist() == list(range(4, 12))
+    assert not state.active and not state.pending
+
+
+@pytest.mark.parametrize("consumer", ["pcm_only", "finalizer", "extra"])
+def test_stream_audio_snapshot_skips_discarded_code_partition(monkeypatch, mocker, consumer):
+    from vllm.v1.worker.gpu.input_batch import InputBatch
+
+    from vllm_omni.worker_v2.output_snapshot import RequestOutputSnapshot
+    from vllm_omni.worker_v2.streaming_audio import StreamingAudioBuffer, StreamingAudioOutput
+
+    monkeypatch.setattr(torch.cuda, "set_stream", lambda _stream: None)
+    monkeypatch.setattr(
+        OmniARModelRunner,
+        "_build_async_chunk_outputs_from_mm",
+        lambda *args: pytest.fail("in-stage PCM unnecessarily partitioned code outputs"),
+    )
+    batch = InputBatch.__new__(InputBatch)
+    batch.num_reqs = 3
+    batch.query_start_loc_np = np.array([0, 1, 2, 4])
+    batch.num_scheduled_tokens = np.array([1, 1, 2])
+    batch.num_tokens_after_padding = 4
+    buffer = StreamingAudioBuffer(samples_per_frame=4, chunk_frames=3)
+    states = [buffer.add(request_id) for request_id in ["frame", "eos", "prefill"]]
+    wav = torch.arange(16, dtype=torch.float32).reshape(4, 4)
+    streaming_audio = StreamingAudioOutput(
+        wav,
+        torch.tensor([True, False, True, False]),
+        torch.tensor(24000),
+        None,
+        batch.query_start_loc_np.copy(),
+        states,
+        np.array([False, False, True]),
+    )
+    copy_mm = mocker.spy(omni_ar_model_runner, "_async_copy_mm")
+
+    def finalize(payload, counts):
+        assert payload["codes"]["audio"].shape == (4, 2)
+        return RequestOutputSnapshot([None] * batch.num_reqs)
+
+    finalizer = mocker.Mock(side_effect=finalize) if consumer == "finalizer" else None
+    extra = ({"extra": torch.tensor([7])}, _FakeEvent()) if consumer == "extra" else None
+    result = _async_output(
+        req_ids=["frame", "eos", "prefill"],
+        sampler_output=SamplerOutput(torch.tensor([[1], [2], [3]]), None, None, None),
+        num_sampled_tokens=torch.tensor([1, 1, 1]),
+        multimodal_outputs={"codes": {"audio": torch.ones(4, 2)}},
+        input_batch=batch,
+        async_chunk=True,
+        streaming_audio=streaming_audio,
+        finalize_multimodal=finalizer,
+        extra_multimodal_outputs=extra,
+    ).get_output()
+    assert copy_mm.call_count == {"pcm_only": 0, "finalizer": 1, "extra": 2}[consumer]
+    if finalizer is not None:
+        finalizer.assert_called_once()
+    wav.fill_(-1)
+    assert result.inter_stage_outputs is None
+    assert result.multimodal_outputs[0]["model_outputs"].tolist() == [0, 1, 2, 3]
+    assert result.multimodal_outputs[1] == {}
+    assert result.multimodal_outputs[2]["model_outputs"].tolist() == [8, 9, 10, 11]
+    assert not states[1].active and not states[2].active
+
+
 @pytest.mark.parametrize("streaming", [False, True])
 def test_request_owned_snapshot_skips_generic_partition(monkeypatch, streaming):
     from vllm_omni.worker_v2.output_snapshot import RequestOutputSnapshot

tests/worker_v2/test_omni_model_state.py

@@ -4,13 +4,17 @@
 graph/eager paths with per-request seed independence, async snapshot ownership."""
 
 from contextlib import nullcontext
+from dataclasses import replace
 from types import SimpleNamespace
 from unittest.mock import MagicMock, patch
 
 import numpy as np
 import pytest
 import torch
+from vllm.v1.worker.gpu.input_batch import InputBatch, InputBuffers
+from vllm.v1.worker.gpu.mm.encoder_runner import EncoderRunner
 from vllm.v1.worker.gpu.model_states.default import DefaultModelState
+from vllm.v1.worker.gpu.states import RequestState
 
 from vllm_omni.model_executor.models.output_templates import OmniOutput, OwnedBatchTensor
 from vllm_omni.worker_v2.model_states.omni_model_state import OmniModelState, _make_safe_get_rope
@@ -36,6 +40,7 @@ class _DummyInputBatch:
 def _make_state(max_num_reqs=4, has_preprocess=False, has_postprocess=False, have_multimodal_outputs=False):
     state = object.__new__(OmniModelState)
     model = MagicMock()
+    model.stream_decoder = None
     model.has_preprocess = has_preprocess
     model.has_postprocess = has_postprocess
     model.have_multimodal_outputs = have_multimodal_outputs
@@ -62,6 +67,10 @@ def _make_state(max_num_reqs=4, has_preprocess=False, has_postprocess=False, hav
 
     state.intermediate_buffer = OmniIntermediateBuffer(max_num_reqs)
     state._static_inputs_embeds = None
+    from vllm_omni.worker_v2.model_states.eager_mtp import EagerMTPState
+
+    state._eager_state = EagerMTPState(state)
+    state._stream_pos = {}
     state._mtp_generators = {}
     state._mtp_runner = None
     for name in ("_mtp_input_ids", "_mtp_input_embeds", "_mtp_hidden", "_mtp_text_step", "_mtp_offsets"):
@@ -215,6 +224,61 @@ def test_static_decode_embeddings_refresh_from_input_ids():
     assert OmniModelState._preprocess_result_needs_writeback(original, original.view_as(original)) is True
 
 
+def test_moss_local_decode_runs_depth_predictor_and_routes_eos(mocker):
+    """Exercise MRV2 dispatch through the real Local hook, output and logits."""
+    from vllm.sampling_params import SamplingParams
+    from vllm.v1.worker.gpu.states import RequestState
+
+    from vllm_omni.model_executor.models.moss_tts.modeling_moss_tts_local_depth import MossTTSLocalDepthTransformer
+    from vllm_omni.model_executor.models.moss_tts.modeling_moss_tts_talker import (
+        MossTTSLocalTalkerForGeneration,
+    )
+
+    model = MossTTSLocalTalkerForGeneration.__new__(MossTTSLocalTalkerForGeneration)
+    torch.nn.Module.__init__(model)
+    model.model = torch.nn.Module()
+    model.model.embed_tokens = torch.nn.Embedding(8, 4, _weight=torch.zeros(8, 4))
+    model.hidden_size = 4
+    model.n_vq = 2
+    model.audio_pad_token_id = 8
+    model.audio_assistant_slot_token_id = 2
+    model.im_end_token_id = 3
+    model.text_vocab_size = 8
+    model.talker_mtp_output_key = ("audio_codes", "current")
+    model.talker_mtp_graph_safe = False
+    model.gpu_resident_buffer_keys = set()
+    model.audio_lm_heads = model.audio_embeddings = model.local_text_lm_head = None
+    model._audio_embed = lambda codes: codes[:, :1].expand(-1, 4).float()
+    frame = mocker.Mock(return_value=(torch.tensor([True, False]), torch.tensor([[1, 2], [3, 4]])))
+    model.local_transformer = mocker.Mock(spec=MossTTSLocalDepthTransformer, generate_frame=frame)
+
+    state = _make_state(max_num_reqs=2, has_preprocess=True, have_multimodal_outputs=True)
+    state.model = model
+    for idx in range(2):
+        state.intermediate_buffer.buffers[idx] = {
+            "req_id": f"r{idx}",
+            "audio_state": {"is_stopping": False},
+            "hidden_states": {"last": torch.full((4,), float(idx + 1))},
+        }
+    batch = _DummyInputBatch([1, 0], num_computed_tokens_cpu=[1, 1])
+    inputs = {"input_ids": torch.tensor([2, 2]), "inputs_embeds": torch.zeros(2, 4)}
+    request_state = mocker.Mock(spec=RequestState, prompt_len=np.array([1, 1]), num_computed_tokens=None)
+    state.run_preprocess(batch, inputs, request_state)
+
+    frame.assert_called_once()
+    torch.testing.assert_close(frame.call_args.args[0], torch.tensor([[2.0] * 4, [1.0] * 4]))
+    assert frame.call_args.kwargs["temperature"] == 1.7
+    assert frame.call_args.kwargs["top_k"] == 25
+    assert frame.call_args.kwargs["top_p"] == 0.8
+    torch.testing.assert_close(inputs["inputs_embeds"], torch.tensor([[1.0] * 4, [0.0] * 4]))
+    _, payload = state.postprocess_model_output(torch.zeros(2, 4), batch, request_state)
+    assert [codes.tolist() for codes in payload["codes"]["audio"]] == [[[1, 2]], [[8, 8]]]
+    assert model.compute_logits(torch.zeros(2, 4)).argmax(-1).tolist() == [2, 3]
+    # An explicit local seed must reach the request-owned MRV2 generator.
+    params = SamplingParams(extra_args={"tts_local_seed": 17}, seed=99)
+    assert state._get_mtp_generator("seeded", params, torch.device("cpu")).initial_seed() == 17
+
+
 @pytest.mark.parametrize("owned", [False, True])
 def test_batched_postprocess_gpu_snapshot_writeback(owned):
     # batch=[1, 0]: last-token indices follow the reordered batch, the scalar
@@ -264,6 +328,7 @@ def test_seed_independence_resolve_once_and_sampling_kwargs():
     # vLLM sampling seed must not produce a talker generator.
     cpu = torch.device("cpu")
     assert state._get_mtp_generator("r1", SimpleNamespace(extra_args={}, seed=42), cpu) is None
+    state._mtp_generators.clear()  # The following rows are new requests.
     # Same model-local seed reproduces identical uniforms regardless of batch makeup.
     state._mtp_sample_uniforms = torch.empty((2, 2, 4))
     assert torch.equal(
@@ -620,6 +685,12 @@ def test_publish_sampled_embeddings_for_rows_whose_sample_is_kept(monkeypatch) -
     assert [tuple(row.shape) for row in sampled] == [(1, 3), (1, 3), (0,), (1, 3)]
     assert [row.tolist() for row in sampled if row.numel()] == [[[v] * 3] for v in (11.0, 12.0, 14.0)]
 
+    batch.is_prefilling_np[:] = True
+    batch.num_computed_prefill_tokens_np[:] = 0
+    batch.num_scheduled_tokens = [1] * 4
+    extra, _done = state.publish_sampled_embeddings(batch, torch.tensor([[11], [12], [13], [14]]))
+    assert all(row.numel() == 0 for row in extra["embed"]["sampled"])
+
     # Speculative steps sample several tokens per row: not published.
     assert state.publish_sampled_embeddings(batch, torch.tensor([[11, 1], [12, 1], [13, 1], [14, 1]])) is None
 
@@ -649,3 +720,32 @@ def test_identity_preprocess_skips_only_decode_rows(prefilling):
     state.run_preprocess(batch, {"input_ids": torch.tensor([1]), "inputs_embeds": embeds})
     assert seen == (["r1"] if prefilling else [])
     assert torch.equal(embeds, torch.ones(1, 4))
+
+
+def test_mm_embeddings_exclude_zero_length_graph_padding_rows():
+    state = object.__new__(OmniModelState)
+    state.supports_mm_inputs = True
+    state.mm_pruner = None
+    state.prompt_embeds_state = None
+    state.encoder_runner = MagicMock(spec=EncoderRunner, inputs_embeds=torch.zeros(8, 2))
+    state.execute_mm_encoder = lambda _: None
+    state.gather_mm_embeddings = lambda _: ([torch.ones(1, 2)], torch.ones(5, dtype=torch.bool))
+
+    def embed(input_ids, *, query_start_loc, multimodal_embeddings, is_multimodal):
+        assert query_start_loc == [0, 2, 5]
+        return input_ids[:, None].expand(-1, 2).float()
+
+    state.model = MagicMock(supports_embed_input_ids_query_start_loc=True, embed_input_ids=embed)
+    buffers = InputBuffers(4, 8, torch.device("cpu"))
+    batch = replace(
+        InputBatch.make_dummy(2, 5, buffers),
+        num_reqs=2,
+        num_reqs_after_padding=4,
+        num_tokens=5,
+        num_tokens_after_padding=8,
+        input_ids=torch.arange(8),
+        query_start_loc_np=np.array([0, 2, 5, 5, 5], dtype=np.int32),
+    )
+    result = state.prepare_inputs_embeds({}, batch, MagicMock(spec=RequestState))
+    torch.testing.assert_close(result[:5], torch.arange(5).float()[:, None].expand(-1, 2))
+    assert result.shape == (8, 2)