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

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

tests/core/sched/test_generation_batch_coalescing.py

@@ -0,0 +1,253 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Bounded readiness coalescing must not allocate or consume request state."""
+
+import queue
+from types import SimpleNamespace
+
+import pytest
+from vllm.v1.core.sched.interface import PauseState
+from vllm.v1.request import RequestStatus
+
+from vllm_omni.core.sched.omni_generation_scheduler import OmniGenerationScheduler
+from vllm_omni.outputs import OmniConnectorOutput
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def test_idle_wait_wakes_on_input_without_extending_coalescing_deadline(monkeypatch):
+    events = [notification("0"), notification("1", "2", "3")]
+    s, clock = scheduler(monkeypatch, delayed=[(0.004, events[0]), (0.002, events[1])])
+    s._generation_coalescing_policy = "idle_wait"
+    for request in s.requests.values():
+        request.status = RequestStatus.WAITING_FOR_CHUNK
+    assert s._drain_omni_connector_outputs() == events
+    assert clock[0] == pytest.approx(10.006)
+    assert s._omni_connector_output_inbox.waits == pytest.approx([0.012, 0.012])
+
+
+@pytest.mark.parametrize("pending", ["in_flight", "admission", "registration", "terminal", "empty"])
+def test_idle_wait_preserves_completion_and_control_progress(monkeypatch, pending):
+    s, clock = scheduler(monkeypatch)
+    s._generation_coalescing_policy = "idle_wait"
+    for request in s.requests.values():
+        request.status = RequestStatus.WAITING_FOR_CHUNK
+    if pending == "in_flight":
+        s.requests["0"].num_in_flight_tokens = 1
+    elif pending == "admission":
+        s.requests["0"].status = RequestStatus.WAITING
+    elif pending == "registration":
+        s.input_coordinator.pending_chunk_registrations = [object()]
+    elif pending == "terminal":
+        s._pending_data_plane_terminal_req_ids = {"0"}
+    else:
+        s.requests.clear()
+    assert s._drain_omni_connector_outputs() == []
+    assert clock[0] == 10
+    assert not s._omni_connector_output_inbox.waits
+
+
+def test_idle_wait_without_arrivals_is_bounded(monkeypatch):
+    s, clock = scheduler(monkeypatch)
+    s._generation_coalescing_policy = "idle_wait"
+    for request in s.requests.values():
+        request.status = RequestStatus.WAITING_FOR_CHUNK
+    assert s._drain_omni_connector_outputs() == []
+    assert clock[0] == pytest.approx(10.012)
+
+
+def notification(*ids, terminal=()):
+    return OmniConnectorOutput(chunk_ready_req_ids=set(ids), chunk_finished_req_ids=set(terminal))
+
+
+class Inbox:
+    def __init__(self, clock, initial=(), delayed=()):
+        self.clock = clock
+        self.initial = list(initial)
+        self.delayed = list(delayed)
+        self.waits = []
+
+    def get_nowait(self):
+        if self.initial:
+            return self.initial.pop(0)
+        raise queue.Empty
+
+    def get(self, timeout):
+        self.waits.append(timeout)
+        if not self.delayed or self.delayed[0][0] > timeout:
+            self.clock[0] += timeout
+            raise queue.Empty
+        delay, value = self.delayed.pop(0)
+        self.clock[0] += delay
+        return value
+
+
+def scheduler(monkeypatch, initial=(), delayed=()):
+    clock = [10.0]
+    monkeypatch.setattr("vllm_omni.core.sched.omni_generation_scheduler.time.monotonic", lambda: clock[0])
+    s = OmniGenerationScheduler.__new__(OmniGenerationScheduler)
+    s._generation_max_wait_s = 0.012
+    s._generation_min_batch_size = 4
+    s._pause_state = PauseState.UNPAUSED
+    s.requests = {
+        str(i): SimpleNamespace(request_id=str(i), num_in_flight_tokens=0, is_finished=lambda: False) for i in range(8)
+    }
+    s.running = list(s.requests.values())
+    s.input_coordinator = SimpleNamespace(
+        requests_with_ready_chunks=set(), finished_requests=set(), input_terminal_req_ids=set()
+    )
+    s._latest_omni_connector_output = None
+    s._omni_connector_output_inbox = Inbox(clock, initial, delayed)
+    return s, clock
+
+
+def test_arrivals_form_target_batch_in_order(monkeypatch):
+    events = [notification("0"), notification("1", "2"), notification("3")]
+    s, clock = scheduler(monkeypatch, events[:1], [(0.002, events[1]), (0.003, events[2])])
+    assert s._drain_omni_connector_outputs() == events
+    assert clock[0] == pytest.approx(10.005)
+    assert not s.input_coordinator.requests_with_ready_chunks
+    assert all(r.num_in_flight_tokens == 0 for r in s.requests.values())
+
+
+def test_duplicate_stale_and_finished_notifications_do_not_fill_batch(monkeypatch):
+    events = [notification("0", "gone", "1", "2"), notification("0")]
+    s, clock = scheduler(monkeypatch, events[:1], [(0.003, events[1])])
+    s.requests["1"].is_finished = lambda: True
+    s.requests["2"].is_finished = lambda: True
+    assert s._drain_omni_connector_outputs() == events
+    assert clock[0] == pytest.approx(10.012)
+    assert s._omni_connector_output_inbox.waits == pytest.approx([0.012, 0.009])
+
+
+@pytest.mark.parametrize("reason", ["enough", "empty", "low", "paused", "off"])
+def test_immediate_dispatch(monkeypatch, reason):
+    event = notification("0")
+    s, clock = scheduler(monkeypatch, [event])
+    if reason == "enough":
+        event.chunk_ready_req_ids.update(["1", "2", "3"])
+    elif reason == "empty":
+        event.chunk_ready_req_ids.clear()
+    elif reason == "low":
+        s.requests = {"0": s.requests["0"]}
+    elif reason == "paused":
+        s._pause_state = PauseState.PAUSED_ALL
+    elif reason == "off":
+        s._generation_max_wait_s = 0
+    assert s._drain_omni_connector_outputs() == [event]
+    assert not s._omni_connector_output_inbox.waits
+    assert clock[0] == 10
+
+
+def test_terminal_arrival_keeps_original_deadline(monkeypatch):
+    events = [notification("0"), notification("1", terminal=["1"])]
+    s, clock = scheduler(monkeypatch, events[:1], [(0.002, events[1])])
+    assert s._drain_omni_connector_outputs() == events
+    assert clock[0] == pytest.approx(10.012)
+
+
+def test_pending_ready_and_legacy_notification_preserved(monkeypatch):
+    s, _ = scheduler(monkeypatch)
+    s.input_coordinator.requests_with_ready_chunks.update(["0", "1", "2"])
+    last = notification("3")
+    s._latest_omni_connector_output = last
+    assert s._drain_omni_connector_outputs() == [last]
+    assert s._latest_omni_connector_output is None
+    assert not s._omni_connector_output_inbox.waits
+
+
+def test_deadline_is_new_for_each_batch_not_extended_by_arrivals(monkeypatch):
+    s, clock = scheduler(monkeypatch, [notification("0")], [(0.008, notification("1")), (0.008, notification("2"))])
+    assert len(s._drain_omni_connector_outputs()) == 2
+    assert clock[0] == pytest.approx(10.012)
+    s._omni_connector_output_inbox.initial.append(notification("3"))
+    assert len(s._drain_omni_connector_outputs()) == 2
+    assert clock[0] == pytest.approx(10.024)
+
+
+def test_cancelled_notification_is_filtered_before_coordinator(monkeypatch, mocker):
+    event = notification("0", "cancelled")
+    event.request_metadata = {"0": {"code_predictor_codes": [7]}, "cancelled": {"code_predictor_codes": [9]}}
+    s, _ = scheduler(monkeypatch, [event])
+    s._generation_max_wait_s = 0
+    s.waiting = []
+    s.input_coordinator._async_chunk = True
+    s.input_coordinator.update_request_metadata = mocker.Mock()
+    s.input_coordinator.process_pending_chunks = mocker.Mock()
+    s._consume_pending_connector_output("generation")
+    s.input_coordinator.update_request_metadata.assert_called_once_with(
+        s.requests, {"0": {"code_predictor_codes": [7]}}, model_mode="generation"
+    )
+    s.input_coordinator.process_pending_chunks.assert_called_once_with(s.waiting, s.running, {"0"}, set())
+
+
+@pytest.mark.parametrize("target,wait", [(0, 1), (257, 1), (4, -1), (4, float("nan")), (4, float("inf"))])
+def test_invalid_config_rejected(monkeypatch, target, wait):
+    from vllm_omni.core.sched.omni_generation_scheduler import VLLMScheduler
+
+    def initialize(s, *args, **kwargs):
+        s.max_num_running_reqs = 256
+        s.vllm_config = SimpleNamespace(
+            model_config=SimpleNamespace(
+                retains_state_across_chunks=True,
+                stage_connector_config={"extra": {"generation_min_batch_size": target, "generation_max_wait_ms": wait}},
+            )
+        )
+
+    monkeypatch.setattr(VLLMScheduler, "__init__", initialize)
+    monkeypatch.setattr(OmniGenerationScheduler, "_init_omni_io_scheduling_state", lambda s: None)
+    with pytest.raises(ValueError):
+        OmniGenerationScheduler()
+
+
+@pytest.mark.parametrize("terminal", [False, True])
+def test_first_and_terminal_chunks_have_bounded_wait(monkeypatch, terminal):
+    event = notification("0", terminal=["0"] if terminal else [])
+    s, clock = scheduler(monkeypatch, [event])
+    s.running = s.running[1:]
+    assert s._drain_omni_connector_outputs() == [event]
+    assert clock[0] == pytest.approx(10.012)
+
+
+def test_retire_in_flight_output_before_waiting_without_resetting_deadline(monkeypatch):
+    events = [notification("0"), notification("1")]
+    s, clock = scheduler(monkeypatch, events[:1])
+    s.requests["7"].num_in_flight_tokens = 1
+    assert s._drain_omni_connector_outputs() == events[:1]
+    assert s._generation_defer_batch
+    assert not s._omni_connector_output_inbox.waits
+    s.input_coordinator.requests_with_ready_chunks.add("0")
+    # Simulate EngineCore retiring the previous batch after five milliseconds.
+    clock[0] += 0.005
+    s.requests["7"].num_in_flight_tokens = 0
+    s._omni_connector_output_inbox.initial.append(events[1])
+    assert s._drain_omni_connector_outputs() == events[1:]
+    assert not s._generation_defer_batch
+    assert s._omni_connector_output_inbox.waits == pytest.approx([0.007])
+    assert clock[0] == pytest.approx(10.012)
+    assert s._generation_batch_deadline is None
+
+
+def test_expired_deadline_dispatches_even_with_other_requests_in_flight(monkeypatch):
+    s, clock = scheduler(monkeypatch, [notification("0")])
+    s.requests["7"].num_in_flight_tokens = 1
+    s._generation_batch_deadline = clock[0] - 0.001
+    assert len(s._drain_omni_connector_outputs()) == 1
+    assert not s._generation_defer_batch
+    assert not s._omni_connector_output_inbox.waits
+
+
+def test_deferred_schedule_keeps_pending_request_and_allocates_nothing(monkeypatch, mocker):
+    from tests.core.sched.test_generation_scheduler_restore import _chunk_request, _make_generation_scheduler
+
+    request = _chunk_request("pending")
+    s = _make_generation_scheduler(request)
+    s._process_pending_omni_inputs = lambda **kwargs: setattr(s, "_generation_defer_batch", True)
+    allocate = mocker.spy(s.kv_cache_manager, "allocate_slots")
+    output = s.schedule()
+    assert output.total_num_scheduled_tokens == 0
+    assert not output.num_scheduled_tokens
+    allocate.assert_not_called()
+    assert list(s.waiting) == [request]
+    assert request.num_computed_tokens == request.num_in_flight_tokens == 0
+    assert s.chunk_transfer_adapter.restore_called

tests/core/sched/test_generation_scheduler_init.py

@@ -0,0 +1,107 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Exercise generation constructor policy after actual deploy/platform resolution."""
+
+from types import SimpleNamespace
+
+import pytest
+
+from tests.helpers.stage_config import get_deploy_config_path
+from vllm_omni.config.pipeline_registry import resolve_pipeline_config
+from vllm_omni.config.stage_config import load_deploy_config, merge_pipeline_deploy
+from vllm_omni.core.sched.omni_generation_scheduler import OmniGenerationScheduler, VLLMScheduler
+from vllm_omni.platforms import current_omni_platform
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+@pytest.fixture
+def construct_scheduler(monkeypatch, mocker):
+    # Only upstream cache allocation and connector construction are omitted.
+    # Keep the actual Omni I/O initialization, native decision and constructor.
+    mocker.patch("vllm_omni.core.sched.omni_scheduler_mixin.OmniChunkTransferAdapter")
+
+    def upstream_init(self, config):
+        self.vllm_config = config
+        self.max_num_running_reqs = config.scheduler_config.max_num_seqs
+
+    monkeypatch.setattr(VLLMScheduler, "__init__", upstream_init)
+    return OmniGenerationScheduler
+
+
+def _config(*, native=True, stateful=True, tp=1, pp=1, extras=None, capacity=128):
+    return SimpleNamespace(
+        model_config=SimpleNamespace(
+            use_v2_model_runner=native,
+            supports_native_mrv2_data_plane=True,
+            retains_state_across_chunks=stateful,
+            async_chunk=True,
+            stage_id=1,
+            stage_connector_config={"extra": extras or {}},
+        ),
+        parallel_config=SimpleNamespace(tensor_parallel_size=tp, pipeline_parallel_size=pp),
+        scheduler_config=SimpleNamespace(max_num_seqs=capacity),
+    )
+
+
+@pytest.mark.parametrize("profile", ["high_concurrency", "low_latency"])
+@pytest.mark.parametrize("platform", ["cuda", "npu", "xpu", "rocm", "musa"])
+def test_moss_profile_generation_constructor_after_platform_resolution(
+    construct_scheduler, monkeypatch, mocker, profile, platform
+):
+    # merge_pipeline_deploy owns platform resolution. Mock the detected device
+    # instead of applying an override first and then resolving the worker's
+    # actual platform a second time.
+    monkeypatch.setattr(current_omni_platform, "device_name", platform)
+    deploy = load_deploy_config(get_deploy_config_path(f"moss_tts_local_mrv2_{profile}.yaml"))
+    pipeline = resolve_pipeline_config("moss_tts_local")
+    codec = merge_pipeline_deploy(pipeline, deploy)[1]
+    extras = deploy.connectors["shm"]["extra"]
+    config = _config(
+        native=codec.yaml_engine_args["use_v2_model_runner"],
+        stateful=pipeline.stages[1].retains_state_across_chunks,
+        capacity=codec.yaml_engine_args["max_num_seqs"],
+        extras=extras,
+    )
+    warning = mocker.patch("vllm_omni.core.sched.omni_generation_scheduler.logger.warning")
+    scheduler = construct_scheduler(config)
+    # The raw extras persist through fallback; the live constructor must cope.
+    assert extras["generation_min_batch_size"] == 32 and extras["generation_max_wait_ms"] == 12
+    if platform == "cuda":
+        assert scheduler._native_data_plane and scheduler.input_coordinator is not None
+        assert scheduler._generation_min_batch_size == 32
+        assert scheduler._generation_max_wait_s == 0.012
+        assert scheduler._generation_max_regular_batch == (16 if profile == "low_latency" else 0)
+        warning.assert_not_called()
+        if profile == "low_latency":
+            assert extras["codec_first_chunk_fast_path"] == extras["codec_first_chunk_gate"] == 1
+    else:
+        assert not scheduler._native_data_plane and scheduler.chunk_transfer_adapter is not None
+        assert scheduler._generation_min_batch_size == 1
+        assert scheduler._generation_max_wait_s == scheduler._generation_max_regular_batch == 0
+        warning.assert_called_once()
+
+
+@pytest.mark.parametrize("stateful,tp,pp", [(False, 1, 1), (True, 2, 1), (True, 1, 2)])
+def test_native_batch_waiting_still_rejects_unsupported_state_or_topology(construct_scheduler, stateful, tp, pp):
+    config = _config(
+        stateful=stateful, tp=tp, pp=pp, extras={"generation_min_batch_size": 32, "generation_max_wait_ms": 12}
+    )
+    with pytest.raises(ValueError, match="stateful native MRV2 TP1/PP1"):
+        construct_scheduler(config)
+
+
+@pytest.mark.parametrize(
+    "extras,message",
+    [
+        ({"generation_min_batch_size": 0}, "generation_min_batch_size"),
+        ({"generation_min_batch_size": 129}, "generation_min_batch_size"),
+        ({"generation_max_wait_ms": -1}, "generation_max_wait_ms"),
+        ({"generation_max_wait_ms": float("nan")}, "generation_max_wait_ms"),
+        ({"generation_max_regular_batch": -1}, "generation_max_regular_batch"),
+        ({"generation_coalescing_policy": "unknown"}, "generation_coalescing_policy"),
+    ],
+)
+def test_non_native_fallback_still_rejects_invalid_values(construct_scheduler, extras, message):
+    with pytest.raises(ValueError, match=message):
+        construct_scheduler(_config(native=False, extras=extras))

tests/core/sched/test_omni_ar_scheduler_kv_transfer.py

@@ -1,10 +1,16 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 from types import SimpleNamespace
 
 import pytest
 from vllm import SamplingParams
 from vllm.v1.engine import EngineCoreRequest
+from vllm.v1.request import RequestStatus
 
+from tests.helpers.omni_scheduler import bind_omits_transfer_helpers
 from vllm_omni.core.sched.omni_ar_scheduler import OmniARScheduler
+from vllm_omni.core.sched.omni_scheduler_mixin import OmniSchedulerMixin
 from vllm_omni.engine.async_engine_utils import apply_omni_final_stage_metadata
 from vllm_omni.engine.serialization import deserialize_additional_information
 
@@ -25,7 +31,7 @@ def _make_engine_request() -> EngineCoreRequest:
     )
 
 
-def _request_omits_kv_transfer(*, force_kv_transfer: bool) -> tuple[bool, dict]:
+def _omit_flags(*, force_kv_transfer: bool) -> tuple[bool, bool, dict]:
     tagged = apply_omni_final_stage_metadata(
         _make_engine_request(),
         final_stage_id=0,
@@ -37,21 +43,182 @@ def _request_omits_kv_transfer(*, force_kv_transfer: bool) -> tuple[bool, dict]:
         request_id="req",
         additional_information=tagged.additional_information,
     )
-    result = scheduler._request_omits_kv_transfer_to_next_stage(request)
+    omits_kv = scheduler._request_omits_kv_transfer_to_next_stage(request)
+    omits_chunk = scheduler._request_omits_chunk_transfer_to_next_stage(request)
     metadata = deserialize_additional_information(tagged.additional_information)
-    return result, metadata
+    return omits_kv, omits_chunk, metadata
 
 
 def test_stage_zero_request_omits_kv_transfer():
-    omits_transfer, metadata = _request_omits_kv_transfer(force_kv_transfer=False)
+    omits_kv, omits_chunk, metadata = _omit_flags(force_kv_transfer=False)
 
-    assert omits_transfer
+    assert omits_kv
+    assert omits_chunk
     assert "omni_force_kv_transfer" not in metadata
 
 
 def test_cfg_companion_forces_kv_transfer_without_downstream_payload():
-    omits_transfer, metadata = _request_omits_kv_transfer(force_kv_transfer=True)
+    omits_kv, omits_chunk, metadata = _omit_flags(force_kv_transfer=True)
 
-    assert not omits_transfer
+    assert not omits_kv
+    assert omits_chunk
     assert metadata["omni_final_stage_id"] == 0
     assert metadata["omni_force_kv_transfer"] is True
+
+
+class _ChunkRequest(SimpleNamespace):
+    def __hash__(self):
+        return hash(self.request_id)
+
+    def __eq__(self, other):
+        return isinstance(other, _ChunkRequest) and other.request_id == self.request_id
+
+
+def _make_chunk_request(final_stage_id: int, *, force_kv_transfer: bool = False):
+    tagged = apply_omni_final_stage_metadata(
+        _make_engine_request(),
+        final_stage_id=final_stage_id,
+        force_kv_transfer=force_kv_transfer,
+    )
+    request = _ChunkRequest(
+        request_id="req-shm-leak",
+        additional_information=tagged.additional_information,
+        num_in_flight_tokens=1,
+        num_stale_output_tokens=0,
+        sampling_params=SimpleNamespace(num_logprobs=None),
+        has_encoder_inputs=False,
+        pooling_params=None,
+        status=RequestStatus.RUNNING,
+        resumable=False,
+        num_output_placeholders=0,
+        spec_token_ids=[],
+        _output_token_ids=[],
+        num_computed_tokens=1,
+        stop_reason=None,
+        num_nans_in_logits=None,
+        client_index=0,
+        trace_headers=None,
+    )
+    request.is_finished = lambda: False
+    request.get_finished_reason = lambda: None
+    request.take_prefill_stats = lambda: None
+    request.take_events = lambda: None
+    return request
+
+
+def _run_finished_save_step(mocker, request, *, inter_stage_outputs=None):
+    adapter = mocker.MagicMock()
+    adapter._confirmed_num_computed_tokens.return_value = 1
+
+    sched = mocker.MagicMock()
+    sched.requests = {request.request_id: request}
+    sched.perf_metrics = None
+    sched.defer_block_free = False
+    sched.structured_output_manager.should_advance.return_value = False
+    sched._update_request_with_output.return_value = ([42], True)
+    sched._process_kv_transfer_trigger.return_value = False
+    sched._handle_stopped_request.return_value = True
+
+    def _free_request(req, delay_free_blocks=False):
+        sched._omits_kv_transfer_cache.pop(req.request_id, None)
+        return None, None
+
+    sched._free_request.side_effect = _free_request
+    sched.chunk_transfer_adapter = adapter
+    sched.running = [request]
+    sched.waiting_for_transfer_free = set()
+    sched.transfer_triggered_requests = set()
+    sched.active_kv_transfers = set()
+    sched.pending_stop_after_extraction = set()
+    sched.connector = None
+    sched.kv_cache_manager.take_events.return_value = None
+    sched.kv_cache_manager.estimate_cached_tokens.return_value = 0
+    sched.finished_req_ids_dict = {}
+    sched.make_stats.return_value = None
+    bind_omits_transfer_helpers(sched)
+
+    scheduler_output = mocker.MagicMock()
+    scheduler_output.num_scheduled_tokens = {request.request_id: 1}
+    scheduler_output.total_num_scheduled_tokens = 1
+    scheduler_output.scheduled_spec_decode_tokens = {}
+    scheduler_output.num_invalid_spec_tokens = 0
+
+    model_runner_output = mocker.MagicMock()
+    model_runner_output.sampled_token_ids = [[42]]
+    model_runner_output.logprobs = None
+    model_runner_output.prompt_logprobs_dict = {}
+    model_runner_output.pooler_output = None
+    model_runner_output.num_nans_in_logits = None
+    model_runner_output.kv_connector_output = None
+    model_runner_output.cudagraph_stats = None
+    model_runner_output.req_id_to_index = {request.request_id: 0}
+    model_runner_output.routed_experts = None
+    model_runner_output.inter_stage_outputs = inter_stage_outputs
+
+    OmniARScheduler.update_from_output(sched, scheduler_output, model_runner_output)
+    return adapter
+
+
+@pytest.mark.parametrize(
+    ("final_stage_id", "force_kv_transfer", "expect_save"),
+    [
+        pytest.param(0, False, False, id="stage-zero-final"),
+        pytest.param(1, False, True, id="downstream-final"),
+    ],
+)
+def test_finish_saves_async_only_with_downstream_consumer(mocker, final_stage_id, force_kv_transfer, expect_save):
+    request = _make_chunk_request(final_stage_id, force_kv_transfer=force_kv_transfer)
+    adapter = _run_finished_save_step(mocker, request)
+
+    assert adapter.save_async.called is expect_save
+
+
+def test_stage_zero_final_skips_inter_stage_output(mocker):
+    warn = mocker.patch("vllm_omni.core.sched.omni_ar_scheduler.logger.warning")
+    request = _make_chunk_request(0)
+    adapter = _run_finished_save_step(mocker, request, inter_stage_outputs=[{"codes": 1}])
+
+    adapter.save_async.assert_not_called()
+    assert any("inter_stage_output is present" in str(call) for call in warn.call_args_list)
+
+
+def test_streaming_session_update_invalidates_omits_kv_transfer_cache():
+    scheduler = OmniARScheduler.__new__(OmniARScheduler)
+    scheduler._omits_kv_transfer_cache = {}
+    scheduler.skipped_waiting = []
+    scheduler.log_stats = False
+    tagged = apply_omni_final_stage_metadata(_make_engine_request(), final_stage_id=0)
+    session = SimpleNamespace(
+        request_id="req",
+        additional_information=tagged.additional_information,
+        status=RequestStatus.WAITING,
+    )
+    assert scheduler._request_omits_chunk_transfer_to_next_stage(session)
+    assert "req" in scheduler._omits_kv_transfer_cache
+    update = SimpleNamespace(
+        additional_information={"omni_final_stage_id": 1},
+        model_intermediate_buffer=None,
+        arrival_time=0.0,
+        sampling_params=None,
+    )
+
+    OmniSchedulerMixin._finish_streaming_session_update(scheduler, session, update)
+
+    assert scheduler._omits_kv_transfer_cache == {}
+    assert session.additional_information == {"omni_final_stage_id": 1}
+    assert not scheduler._request_omits_chunk_transfer_to_next_stage(session)
+    assert not scheduler._request_omits_kv_transfer_to_next_stage(session)
+
+
+def test_omits_cache_misses_when_payload_object_is_replaced():
+    scheduler = OmniARScheduler.__new__(OmniARScheduler)
+    scheduler._omits_kv_transfer_cache = {}
+    tagged = apply_omni_final_stage_metadata(_make_engine_request(), final_stage_id=0)
+    request = SimpleNamespace(
+        request_id="req",
+        additional_information=tagged.additional_information,
+    )
+    assert scheduler._request_omits_chunk_transfer_to_next_stage(request)
+    request.additional_information = {"omni_final_stage_id": 1}
+
+    assert not scheduler._request_omits_chunk_transfer_to_next_stage(request)

tests/core/sched/test_omni_ar_scheduler_stale_drain.py

@@ -40,6 +40,7 @@ from vllm.v1.core.sched.output import SchedulerOutput
 from vllm.v1.outputs import ModelRunnerOutput
 from vllm.v1.request import Request, RequestStatus, StreamingUpdate
 from vllm_omni.core.sched.omni_ar_scheduler import OmniARScheduler
+from tests.helpers.omni_scheduler import bind_omits_transfer_helpers
 
 # isort: on
 
@@ -104,10 +105,11 @@ def _make_drain_sched(session: Request) -> MagicMock:
     sched.kv_cache_manager.estimate_cached_tokens.return_value = 0
     sched.finished_req_ids_dict = {}
     sched.make_stats.return_value = None
+    bind_omits_transfer_helpers(sched)
     return sched
 
 
-def _run_step(sched: MagicMock, session: Request, *, num_scheduled: int, token: int) -> bool:
+def _run_step(sched: MagicMock, session: Request, *, num_scheduled: int, token: int, spec=None) -> bool:
     """Feed one model-runner frame through update_from_output.
 
     Returns True when the frame's tokens were delivered (the append ran),
@@ -115,7 +117,7 @@ def _run_step(sched: MagicMock, session: Request, *, num_scheduled: int, token:
     """
     scheduler_output = MagicMock(spec=SchedulerOutput)
     scheduler_output.num_scheduled_tokens = {session.request_id: num_scheduled}
-    scheduler_output.scheduled_spec_decode_tokens = {}
+    scheduler_output.scheduled_spec_decode_tokens = {session.request_id: spec} if spec else {}
     scheduler_output.num_invalid_spec_tokens = 0
 
     model_runner_output = MagicMock(spec=ModelRunnerOutput)
@@ -135,6 +137,27 @@ def _run_step(sched: MagicMock, session: Request, *, num_scheduled: int, token:
     return sched._update_request_with_output.called
 
 
+def test_connector_prompt_replacement_drops_old_frame_and_delivers_new_frame() -> None:
+    from vllm_omni.core.sched.omni_scheduler_mixin import OmniSchedulerMixin
+
+    session = _make_session()
+    session.num_computed_tokens = 6
+    session.num_in_flight_tokens = 1
+    session.num_output_placeholders = 1
+    sched = _make_drain_sched(session)
+    sched.chunk_transfer_adapter.replaced_streaming_prompt_ids = {session.request_id}
+    sched.chunk_transfer_adapter.requests_with_ready_chunks = {session.request_id}
+    sched.chunk_transfer_adapter.requests_num_chunks_sent = {session.external_req_id: 2}
+
+    OmniSchedulerMixin._reset_ready_async_chunk_replacements(sched)
+    assert session.num_stale_output_tokens == 1 and session.drop_stale_output
+    assert session.num_output_placeholders == 0
+    sched._release_replaced_streaming_prompt_cache.assert_called_once_with(session)
+    assert _run_step(sched, session, num_scheduled=1, token=42) is False
+    assert session.num_stale_output_tokens == 0
+    assert _run_step(sched, session, num_scheduled=1, token=43) is True
+
+
 def test_exact_drain_delivers_new_segment_frame() -> None:
     """One in-flight decode frame at replacement: its late output drains the
     counter exactly and the new segment's first frame is delivered."""
@@ -226,3 +249,80 @@ def test_in_flight_prefill_chunk_drains_exactly_without_underflow() -> None:
     assert _run_step(sched, session, num_scheduled=5, token=42) is False  # late prefill chunk dropped
     assert session.num_stale_output_tokens == 0
     assert _run_step(sched, session, num_scheduled=1, token=43) is True  # new segment survives
+
+
+@pytest.mark.parametrize("drop", [False, True])
+def test_preemption_stale_output_follows_upstream_delivery_policy(drop):
+    request = _make_session()
+    request.status = RequestStatus.PREEMPTED
+    request.num_stale_output_tokens = 1
+    request.num_in_flight_tokens = 1
+    request.drop_stale_output = drop
+    sched = _make_drain_sched(request)
+    assert _run_step(sched, request, num_scheduled=1, token=42) is (not drop)
+    assert request.num_stale_output_tokens == 0
+    if not drop:
+        sched._update_request_with_output.assert_called_once_with(request, [42], is_stale=True)
+
+
+def test_preemption_stale_spec_rejection_does_not_roll_back_resumed_counters():
+    request = _make_session()
+    request.num_stale_output_tokens = 2
+    request.num_in_flight_tokens = 2
+    request.num_computed_tokens = 3
+    request.num_output_placeholders = 2
+    request.drop_stale_output = False
+    sched = _make_drain_sched(request)
+    assert _run_step(sched, request, num_scheduled=2, token=42, spec=[41])
+    assert request.num_computed_tokens == 3
+    assert request.num_output_placeholders == 2
+
+
+@pytest.mark.parametrize(
+    "arch,supports_reset,running,reset_running,blocked",
+    [
+        ("Qwen3TTSTalkerForConditionalGeneration", False, True, True, True),
+        ("Qwen3TTSTalkerForConditionalGeneration", False, False, True, False),
+        ("Qwen3TTSTalkerForConditionalGeneration", False, True, False, False),
+        ("Qwen3TTSTalkerForConditionalGeneration", True, True, True, False),
+        ("AnotherStatefulAudioModel", False, True, True, True),
+        ("AnotherStatefulAudioModel", True, True, True, False),
+    ],
+)
+def test_reset_running_streaming_codec_is_rejected_before_preemption(
+    mocker, arch, supports_reset, running, reset_running, blocked
+):
+    from vllm.config import VllmConfig
+    from vllm.v1.core.sched.scheduler import Scheduler
+
+    from vllm_omni.config.model import OmniModelConfig
+
+    scheduler = OmniARScheduler.__new__(OmniARScheduler)
+    scheduler.vllm_config = mocker.Mock(
+        spec=VllmConfig,
+        model_config=mocker.Mock(
+            spec=OmniModelConfig,
+            model_arch=arch,
+            supports_running_prefix_cache_reset=supports_reset,
+        ),
+    )
+    request = _make_session()
+    request.num_in_flight_tokens = 1
+    request.num_output_placeholders = 1
+    scheduler.running = [request] if running else []
+
+    def reset_cache(_self, reset_running_requests=False, reset_connector=False):
+        assert reset_running_requests is reset_running and reset_connector is True
+        return True
+
+    reset = mocker.patch.object(Scheduler, "reset_prefix_cache", autospec=True, side_effect=reset_cache)
+
+    assert scheduler.reset_prefix_cache(reset_running, reset_connector=True) is (not blocked)
+    if blocked:
+        reset.assert_not_called()
+        assert scheduler.running == [request]
+        assert request.status == RequestStatus.RUNNING
+        assert request.num_in_flight_tokens == request.num_output_placeholders == 1
+        assert not request.drop_stale_output
+    else:
+        reset.assert_called_once()

tests/core/sched/test_omni_ar_scheduler_streaming.py

@@ -26,6 +26,7 @@ from vllm_omni.core.sched.omni_ar_scheduler import OmniARAsyncScheduler, OmniARS
 from vllm_omni.distributed.omni_connectors.transfer_adapter.chunk_transfer_adapter import (
     OmniChunkTransferAdapter,
 )
+from tests.helpers.omni_scheduler import bind_omits_transfer_helpers
 
 # isort: on
 
@@ -176,6 +177,7 @@ def _run_resumable_segment_stop(
     sched.kv_cache_manager.estimate_cached_tokens.return_value = 0
     sched.finished_req_ids_dict = {}
     sched.make_stats.return_value = None
+    bind_omits_transfer_helpers(sched)
 
     scheduler_output = MagicMock(spec=SchedulerOutput)
     scheduler_output.num_scheduled_tokens = {session.request_id: 1}
@@ -313,6 +315,7 @@ def test_running_decode_step_without_inter_stage_payload_does_not_raise() -> Non
     sched.kv_cache_manager.estimate_cached_tokens.return_value = 0
     sched.finished_req_ids_dict = {}
     sched.make_stats.return_value = None
+    bind_omits_transfer_helpers(sched)
 
     scheduler_output = MagicMock(spec=SchedulerOutput)
     scheduler_output.num_scheduled_tokens = {session.request_id: 1}
@@ -401,6 +404,7 @@ def test_stale_async_frame_is_dropped_before_output_processing() -> None:
     sched.kv_cache_manager.estimate_cached_tokens.return_value = 0
     sched.finished_req_ids_dict = {}
     sched.make_stats.return_value = None
+    bind_omits_transfer_helpers(sched)
 
     scheduler_output = MagicMock(spec=SchedulerOutput)
     scheduler_output.num_scheduled_tokens = {session.request_id: 1}

tests/core/sched/test_omni_sched_deferred_free_fence.py

@@ -127,6 +127,11 @@ def _make_fast_path_sched(*, defer_block_free: bool) -> tuple[MagicMock, Request
     sched.max_num_running_reqs = 8
     sched._native_data_plane = False
     sched._retains_state_across_chunks = False
+    # Unset MagicMock attributes are truthy. Keep optional scheduling modes
+    # disabled so this fixture actually schedules its one-shot request.
+    sched._generation_defer_batch = False
+    sched._generation_first_chunk_only = False
+    sched._generation_max_regular_batch = 0
     sched._pause_state = PauseState.UNPAUSED
     sched.requests = {request.request_id: request}
     sched.running = [request]

tests/core/sched/test_omni_scheduler_input_coordinator_cleanup.py

@@ -81,7 +81,7 @@ def test_ar_free_request_cleans_input_coordinator_on_normal_free() -> None:
     scheduler = OmniARScheduler.__new__(OmniARScheduler)
     scheduler.input_coordinator = coordinator
     scheduler.chunk_transfer_adapter = None
-    scheduler._omits_kv_transfer_cache = {"req-free": True}
+    scheduler._omits_kv_transfer_cache = {"req-free": (1, True, False)}
     scheduler.encoder_cache_manager = SimpleNamespace(free=lambda request: None)
     scheduler.finished_req_ids = set()
     scheduler.finished_req_ids_dict = None