rebase: align vllm-omni with vLLM ac9126e58aa7

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

tests/core/sched/test_generation_scheduler_restore.py

@@ -11,7 +11,9 @@ from types import SimpleNamespace
 import pytest
 import torch
 from vllm import SamplingParams
+from vllm.distributed.aux_output_connector.connector import AuxOutputSchedulerConnector
 from vllm.v1.core.sched.interface import PauseState
+from vllm.v1.core.sched.output import CachedRequestData
 from vllm.v1.core.sched.request_queue import SchedulingPolicy, create_request_queue
 from vllm.v1.request import Request, RequestStatus
 
@@ -53,12 +55,14 @@ class FakeAdapter:
 def _make_generation_scheduler(waiting_request, *, use_v2_model_runner=False):
     scheduler = OmniGenerationScheduler.__new__(OmniGenerationScheduler)
     scheduler.max_num_scheduled_tokens = 8
+    scheduler.max_num_active_reqs = 1
     scheduler.max_num_running_reqs = 1
     scheduler._pause_state = PauseState.UNPAUSED
     scheduler.running = []
     scheduler.waiting = create_request_queue(SchedulingPolicy.FCFS)
     scheduler.waiting.add_request(waiting_request)
-    scheduler.skipped_waiting = create_request_queue(SchedulingPolicy.FCFS)
+    scheduler.kv_holding_waiting = create_request_queue(SchedulingPolicy.FCFS)
+    scheduler.deferred_waiting = set()
     scheduler.requests = {waiting_request.request_id: waiting_request}
     scheduler.policy = SchedulingPolicy.FCFS
     scheduler.chunk_transfer_adapter = FakeAdapter()
@@ -86,6 +90,7 @@ def _make_generation_scheduler(waiting_request, *, use_v2_model_runner=False):
         get_manager_metadata=lambda: None,
     )
     scheduler.connector = None
+    scheduler.aux_output_connector = None
     scheduler.ec_connector = None
     scheduler.prev_step_scheduled_req_ids = set()
     scheduler._pending_finish_reqs = []
@@ -105,6 +110,45 @@ def _make_generation_scheduler(waiting_request, *, use_v2_model_runner=False):
     return scheduler
 
 
+def test_generation_scheduler_preserves_auxiliary_output_metadata() -> None:
+    request = Request("aux-contract", [1, 2], SamplingParams(max_tokens=2), pooling_params=None)
+    request.block_hashes.append(b"block-hash")
+    scheduler = _make_generation_scheduler(request, use_v2_model_runner=True)
+    scheduler.aux_output_connector = AuxOutputSchedulerConnector()
+
+    output = scheduler._build_generation_scheduler_output(
+        new_reqs_data=[],
+        cached_reqs_data=CachedRequestData.make_empty(),
+        num_scheduled_tokens={request.request_id: 1},
+        scheduled_spec_decode_tokens={},
+        scheduled_encoder_inputs={},
+        num_common_prefix_blocks=[0],
+    )
+
+    metadata = output.aux_output_connector_metadata
+    assert metadata is not None
+    assert metadata.requests == {request.request_id: 0}
+    assert list(metadata.block_hashes[request.request_id]) == request.block_hashes
+
+
+@pytest.mark.parametrize(("runner_slots", "active_limit"), [(1, 2), (2, 1)])
+def test_generation_admits_kv_holder_before_fresh_requests(runner_slots, active_limit) -> None:
+    fresh = Request("fresh", [1, 2], SamplingParams(max_tokens=2), pooling_params=None)
+    holding = Request("holding", [3, 4], SamplingParams(max_tokens=2), pooling_params=None)
+    holding.num_computed_tokens = 1
+    scheduler = _make_generation_scheduler(fresh)
+    scheduler.max_num_running_reqs = runner_slots
+    scheduler.max_num_active_reqs = active_limit
+    scheduler.requests[holding.request_id] = holding
+    scheduler.kv_holding_waiting.add_request(holding)
+
+    output = scheduler.schedule()
+
+    assert set(output.num_scheduled_tokens) == {holding.request_id}
+    assert list(scheduler.waiting) == [fresh]
+    assert list(scheduler.kv_holding_waiting) == []
+
+
 def _chunk_request(request_id, **kwargs):
     defaults = dict(
         request_id=request_id,
@@ -123,7 +167,13 @@ def _chunk_request(request_id, **kwargs):
         record_event=lambda *args, **kwargs: None,
     )
     defaults.update(kwargs)
-    return SimpleNamespace(**defaults)
+    return _HashableChunkRequest(**defaults)
+
+
+class _HashableChunkRequest(SimpleNamespace):
+    # Deferred queues require identity hashing; SimpleNamespace types it as None.
+    __hash__ = object.__hash__  # type: ignore[assignment]
+    __eq__ = object.__eq__
 
 
 def test_chunk_lifecycle_no_resubmit_and_state_survives_requeue(monkeypatch) -> None:
@@ -138,7 +188,8 @@ def test_chunk_lifecycle_no_resubmit_and_state_survives_requeue(monkeypatch) ->
     )
     scheduler.running = []
     scheduler.waiting = create_request_queue(scheduler.policy)
-    scheduler.skipped_waiting = create_request_queue(scheduler.policy)
+    scheduler.kv_holding_waiting = create_request_queue(scheduler.policy)
+    scheduler.deferred_waiting = set()
 
     completed = Request("completed", [1, 2], SamplingParams(max_tokens=4), pooling_params=None)
     completed.status = RequestStatus.RUNNING
@@ -292,6 +343,7 @@ def _express_scheduler():
     scheduler = _make_generation_scheduler(continuation, use_v2_model_runner=True)
     scheduler.waiting.add_request(first)
     scheduler.requests[first.request_id] = first
+    scheduler.max_num_active_reqs = 4
     scheduler.max_num_running_reqs = 4
     scheduler._native_data_plane = True
     scheduler.chunk_transfer_adapter = None

tests/core/sched/test_omni_ar_scheduler_aborted_queue_sweep.py

@@ -9,7 +9,7 @@ handing control to upstream, because upstream ``Scheduler.schedule()`` raises
 admits in a finished state -- which kills the stage's engine core, not just
 the request.
 
-Regression for the sweep missing ``skipped_waiting``, the third queue
+Regression for the sweep missing ``kv_holding_waiting``, the third queue
 upstream admits from. An aborted duplex session parked there was re-selected
 by ``_select_waiting_queue_for_scheduling`` on a later tick and crashed the
 stage.
@@ -49,28 +49,30 @@ class _StubRequest:
         return (self.priority, self.arrival_time) < (other.priority, other.arrival_time)
 
 
-def _make_scheduler(policy: SchedulingPolicy, *, waiting=(), skipped_waiting=(), running=()):
+def _make_scheduler(policy: SchedulingPolicy, *, waiting=(), kv_holding_waiting=(), running=()):
     """Bare scheduler with just the surface the sweep reads."""
     scheduler = OmniARScheduler.__new__(OmniARScheduler)
     scheduler.waiting = create_request_queue(policy)
     for req in waiting:
         scheduler.waiting.add_request(req)
-    scheduler.skipped_waiting = create_request_queue(policy)
-    for req in skipped_waiting:
-        scheduler.skipped_waiting.add_request(req)
+    scheduler.kv_holding_waiting = create_request_queue(policy)
+    scheduler.deferred_waiting = set()
+    for req in kv_holding_waiting:
+        scheduler.kv_holding_waiting.add_request(req)
     scheduler.running = list(running)
     return scheduler
 
 
 @pytest.mark.parametrize("policy", _POLICIES)
-def test_schedule_sweeps_skipped_waiting_before_upstream_selection(
+def test_schedule_sweeps_kv_holding_waiting_before_upstream_selection(
     monkeypatch: pytest.MonkeyPatch,
     policy: SchedulingPolicy,
 ) -> None:
     """The public schedule path must remove aborted skipped requests first."""
     aborted = _StubRequest("req-aborted", RequestStatus.FINISHED_ABORTED)
-    scheduler = _make_scheduler(policy, skipped_waiting=[aborted])
+    scheduler = _make_scheduler(policy, kv_holding_waiting=[aborted])
     scheduler.chunk_transfer_adapter = None
+    scheduler.max_num_active_reqs = 8
     scheduler.max_num_running_reqs = 8
     scheduler.input_coordinator = None
     scheduler.num_waiting_for_streaming_input = 0
@@ -85,7 +87,7 @@ def test_schedule_sweeps_skipped_waiting_before_upstream_selection(
     sentinel = object()
 
     def upstream_schedule(_self, _throttle_prefills: bool = False):
-        assert list(scheduler.skipped_waiting) == []
+        assert list(scheduler.kv_holding_waiting) == []
         return sentinel
 
     monkeypatch.setattr(ar_sched_mod.VLLMScheduler, "schedule", upstream_schedule)

tests/core/sched/test_omni_ar_scheduler_logprobs.py

@@ -11,12 +11,14 @@ from types import MethodType, SimpleNamespace
 
 import numpy as np
 import pytest
+import torch
 from vllm.v1.engine import FinishReason
 from vllm.v1.outputs import LogprobsLists
 from vllm.v1.request import RequestStatus
 
 from vllm_omni.core.sched.omni_ar_scheduler import OmniARScheduler, _slice_sampled_logprobs
 from vllm_omni.core.sched.omni_scheduler_mixin import OmniSchedulerMixin
+from vllm_omni.outputs import OmniModelRunnerOutput
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
@@ -132,11 +134,13 @@ def _make_scheduler_stub(requests: list[_Request]) -> SimpleNamespace:
     scheduler = SimpleNamespace(
         perf_metrics=None,
         connector=None,
+        aux_output_connector=None,
         chunk_transfer_adapter=None,
         requests={request.request_id: request for request in requests},
         running=list(requests),
         waiting=_RequestQueue(),
-        skipped_waiting=_RequestQueue(),
+        kv_holding_waiting=_RequestQueue(),
+        deferred_waiting=set(),
         structured_output_manager=SimpleNamespace(accept_tokens=lambda _request, _token_ids: True),
         transfer_triggered_requests=set(),
         active_kv_transfers=set(),
@@ -292,6 +296,36 @@ def test_invalid_logprobs_finish_only_the_affected_scheduler_request() -> None:
     np.testing.assert_array_equal(output_by_id["good"].new_logprobs.logprob_token_ids[:, 0], [8])
 
 
+@pytest.mark.parametrize("scored", [False, True])
+def test_fixed_token_scores_reach_engine_output(scored) -> None:
+    request = _Request("req")
+    request.sampling_params.num_logprobs = None
+    scheduler = _make_scheduler_stub([request])
+    _bind_request_lifecycle(scheduler, update_request=lambda req, tokens: (tokens, False))
+    scheduler_output = SimpleNamespace(
+        num_scheduled_tokens={"req": 3},
+        scheduled_spec_decode_tokens={},
+        num_invalid_spec_tokens=0,
+    )
+    scores = torch.tensor([[-0.5, -1.5], [-0.2, -2.2]])
+    model_runner_output = OmniModelRunnerOutput(
+        req_ids=["req"],
+        req_id_to_index={"req": 0},
+        sampled_token_ids=[[7]],
+        prompt_logprobs_dict={},
+        prompt_token_id_logprobs_dict={"req": scores} if scored else {},
+    )
+
+    outputs = OmniARScheduler.update_from_output(scheduler, scheduler_output, model_runner_output)
+    (output,) = outputs[0].outputs
+
+    assert output.new_token_ids == [7]
+    if scored:
+        torch.testing.assert_close(output.prompt_token_id_logprobs, scores)
+    else:
+        assert output.prompt_token_id_logprobs is None
+
+
 def _pooling_model_runner_output(pooler_tensor) -> SimpleNamespace:
     return SimpleNamespace(
         sampled_token_ids=[[]],

tests/core/sched/test_omni_ar_scheduler_stale_drain.py

@@ -72,7 +72,8 @@ def _replace_streaming_session(session: Request) -> None:
     sched.num_waiting_for_streaming_input = 0
     sched.log_stats = False
     sched.chunk_transfer_adapter = None
-    sched.skipped_waiting = set()
+    sched.kv_holding_waiting = set()
+    sched.deferred_waiting = set()
     session.status = RequestStatus.WAITING_FOR_STREAMING_REQ
     update = StreamingUpdate(
         mm_features=None,

tests/core/sched/test_omni_ar_scheduler_streaming.py

@@ -41,7 +41,8 @@ def _make_scheduler(*, stage_id: int = 0, session_mode: str = "turn") -> OmniARS
     sched.num_waiting_for_streaming_input = 0
     sched.log_stats = False
     sched.chunk_transfer_adapter = None
-    sched.skipped_waiting = set()
+    sched.kv_holding_waiting = set()
+    sched.deferred_waiting = set()
     sched._free_request_blocks = MagicMock()
     sched.encoder_cache_manager = MagicMock()
     sched._inflight_prefills = set()
@@ -1165,7 +1166,8 @@ def test_chunk_segment_cleanup_keeps_requeued_resumable_receiver() -> None:
     sched.vllm_config = SimpleNamespace(model_config=SimpleNamespace(session_mode="duplex"))
     sched.running = []
     sched.waiting = queue()
-    sched.skipped_waiting = queue(session)
+    sched.kv_holding_waiting = queue(session)
+    sched.deferred_waiting = set()
     sched.num_waiting_for_streaming_input = 1
     sched.chunk_transfer_adapter = SimpleNamespace(
         receives_chunks=True,
@@ -1177,7 +1179,7 @@ def test_chunk_segment_cleanup_keeps_requeued_resumable_receiver() -> None:
 
     assert session.status == RequestStatus.WAITING
     assert session in sched.waiting.requests
-    assert session not in sched.skipped_waiting.requests
+    assert session not in sched.kv_holding_waiting.requests
     assert sched.num_waiting_for_streaming_input == 0
     assert session.request_id not in sched.chunk_transfer_adapter.segment_finished_requests
 
@@ -1204,7 +1206,8 @@ def test_chunk_segment_cleanup_keeps_explicit_update_stage_parked(
         receives_chunks=receives_chunks,
         segment_finished_requests={session.request_id},
     )
-    sched.skipped_waiting = MagicMock()
+    sched.kv_holding_waiting = MagicMock()
+    sched.deferred_waiting = set()
     sched._enqueue_waiting_request = MagicMock()
 
     sched._resume_downstream_chunk_receiver(session)
@@ -1212,7 +1215,7 @@ def test_chunk_segment_cleanup_keeps_explicit_update_stage_parked(
     assert session.status == RequestStatus.WAITING_FOR_STREAMING_REQ
     assert sched.num_waiting_for_streaming_input == 1
     assert session.request_id not in sched.chunk_transfer_adapter.segment_finished_requests
-    sched.skipped_waiting.remove_requests.assert_not_called()
+    sched.kv_holding_waiting.remove_requests.assert_not_called()
     sched._enqueue_waiting_request.assert_not_called()
 
 
@@ -1268,7 +1271,8 @@ def _park_session(sched: OmniARScheduler, session: Request) -> None:
     session.num_computed_tokens = 6
     session.num_output_placeholders = 0
     session.status = RequestStatus.WAITING_FOR_STREAMING_REQ
-    sched.skipped_waiting.add_request(session)
+    sched.kv_holding_waiting.add_request(session)
+    sched.deferred_waiting.add(session)
     sched.num_waiting_for_streaming_input = 1
 
 
@@ -1291,7 +1295,7 @@ def test_stage0_streaming_update_that_overflows_max_model_len_finishes_the_sessi
     assert session.status == RequestStatus.FINISHED_ERROR
     assert session.request_id not in sched.requests
     assert len(sched.waiting) == 0
-    assert len(sched.skipped_waiting) == 0
+    assert len(sched.kv_holding_waiting) == 0
     assert sched.num_waiting_for_streaming_input == 0
     sched._free_request_blocks.assert_called_once_with(session)
     assert sched.finished_req_ids == {session.request_id}
@@ -1399,7 +1403,8 @@ def _make_admission_scheduler(*, max_model_len: int) -> OmniARScheduler:
     sched.max_model_len = max_model_len
     sched.policy = SchedulingPolicy.FCFS
     sched.waiting = create_request_queue(sched.policy)
-    sched.skipped_waiting = create_request_queue(sched.policy)
+    sched.kv_holding_waiting = create_request_queue(sched.policy)
+    sched.deferred_waiting = set()
     sched.running = []
     sched.requests = {}
     sched._free_request = MagicMock()  # needs a KV manager; not what this covers
@@ -1428,7 +1433,7 @@ def test_queued_streaming_update_that_overflows_does_not_return_to_admission() -
 
     assert finished is True
     assert len(sched.waiting) == 0
-    assert len(sched.skipped_waiting) == 0
+    assert len(sched.kv_holding_waiting) == 0
     assert session.status == RequestStatus.FINISHED_ERROR
     assert session.resumable is False
     assert "context_length_exceeded: " in sched._streaming_context_overflow[session.request_id][1]
@@ -1447,7 +1452,8 @@ def test_queued_streaming_update_that_fits_still_resumes_the_session() -> None:
     finished = sched._handle_stopped_request(session)
 
     assert finished is False
-    assert len(sched.waiting) == 1
+    assert list(sched.kv_holding_waiting) == [session]
+    assert session not in sched.deferred_waiting
     assert session.status == RequestStatus.WAITING
     assert not getattr(sched, "_streaming_context_overflow", {})
     sched._free_request.assert_not_called()
@@ -1469,7 +1475,7 @@ def test_queued_stop_with_a_recorded_overflow_still_leaves_admission() -> None:
 
     assert finished is True
     assert len(sched.waiting) == 0
-    assert len(sched.skipped_waiting) == 0
+    assert len(sched.kv_holding_waiting) == 0
     assert session.status == RequestStatus.FINISHED_ERROR
     assert session.resumable is False
     sched._free_request.assert_not_called()
@@ -1492,7 +1498,7 @@ def test_parked_stop_with_a_recorded_overflow_keeps_the_streaming_counter_balanc
     finished = sched._handle_stopped_request(session)
 
     assert finished is True
-    assert len(sched.skipped_waiting) == 0
+    assert len(sched.kv_holding_waiting) == 0
     assert sched.num_waiting_for_streaming_input == 0
     assert session.status == RequestStatus.FINISHED_ERROR
     sched._free_request.assert_not_called()
@@ -1525,6 +1531,7 @@ def test_async_chunk_reserves_parked_slots_during_ar_admission(monkeypatch, nati
     sched.running = []
     sched._native_data_plane = native
     sched.use_v2_model_runner = native
+    sched.max_num_active_reqs = 8
     sched.max_num_running_reqs = 8
     sched.input_coordinator = (
         SimpleNamespace(_async_chunk=True, _waiting_for_chunk_running=[parked], restore_queues=lambda _w, _r: None)
@@ -1551,7 +1558,7 @@ def test_async_chunk_reserves_parked_slots_during_ar_admission(monkeypatch, nati
     observed_limits: list[int] = []
 
     def fake_schedule(self, _throttle_prefills=False):
-        observed_limits.append(self.max_num_running_reqs)
+        observed_limits.append(self.max_num_active_reqs)
         return SimpleNamespace(scheduled_new_reqs=[])
 
     monkeypatch.setattr(VLLMScheduler, "schedule", fake_schedule)
@@ -1559,4 +1566,4 @@ def test_async_chunk_reserves_parked_slots_during_ar_admission(monkeypatch, nati
     sched.schedule()
 
     assert observed_limits == [7 if native else 8]
-    assert sched.max_num_running_reqs == 8
+    assert sched.max_num_active_reqs == 8

tests/core/sched/test_omni_ar_scheduler_unified_decode_deferral.py

@@ -79,16 +79,19 @@ def _make_scheduler(
         else:
             hf_config.voxcpm2_runtime_config = runtime_config
     sched.vllm_config = SimpleNamespace(model_config=model_config)
-    # Upstream ``Scheduler.__init__`` always creates ``skipped_waiting``; the
+    # Upstream ``Scheduler.__init__`` always creates ``kv_holding_waiting``; the
     # abort sweep at the top of ``schedule()`` reads it.
-    sched.skipped_waiting = _MockQueue()
+    sched.kv_holding_waiting = _MockQueue()
+    sched.deferred_waiting = set()
     return sched
 
 
-def test_voxcpm2_unified_decode_graph_defers_waiting_when_decode_ready() -> None:
+@pytest.mark.parametrize("queue_name", ["waiting", "kv_holding_waiting"])
+def test_voxcpm2_unified_decode_graph_defers_waiting_when_decode_ready(queue_name) -> None:
     scheduler = _make_scheduler()
     scheduler.running = [_MockRequest("decode")]
-    scheduler.waiting = _MockQueue([_MockRequest("prefill", status=RequestStatus.WAITING)])
+    scheduler.waiting = _MockQueue()
+    setattr(scheduler, queue_name, _MockQueue([_MockRequest("prefill", status=RequestStatus.WAITING)]))
 
     assert scheduler._should_defer_waiting_for_unified_decode_graph()
 
@@ -144,12 +147,17 @@ def test_voxcpm2_unified_decode_graph_does_not_defer_without_cuda_graph(monkeypa
 
 def test_unified_decode_graph_deferral_restores_waiting_queue(monkeypatch) -> None:
     scheduler = _make_scheduler()
+    scheduler.max_num_active_reqs = 8
     scheduler.max_num_running_reqs = 8
     scheduler.running = [_MockRequest("decode")]
     original_waiting_req = _MockRequest("waiting", status=RequestStatus.WAITING)
     deferred_by_upstream = _MockRequest("deferred-by-upstream", status=RequestStatus.WAITING)
     original_waiting = _MockQueue([original_waiting_req])
     scheduler.waiting = original_waiting
+    kv_holder = _MockRequest("kv-holder", status=RequestStatus.WAITING)
+    original_kv_holding = _MockQueue([kv_holder])
+    scheduler.kv_holding_waiting = original_kv_holding
+    preempted_by_upstream = _MockRequest("preempted-by-upstream", status=RequestStatus.PREEMPTED)
     scheduler.policy = "fcfs"
     scheduler.chunk_transfer_adapter = None
     scheduler.input_coordinator = None
@@ -160,7 +168,10 @@ def test_unified_decode_graph_deferral_restores_waiting_queue(monkeypatch) -> No
     def fake_upstream_schedule(self, throttle_prefills: bool = False):
         assert self.waiting is not original_waiting
         assert not self.waiting
+        assert self.kv_holding_waiting is not original_kv_holding
+        assert not self.kv_holding_waiting
         self.waiting.add_request(deferred_by_upstream)
+        self.kv_holding_waiting.add_request(preempted_by_upstream)
         raise RuntimeError("stop before output wrapping")
 
     monkeypatch.setattr(scheduler_mod.VLLMScheduler, "schedule", fake_upstream_schedule)
@@ -169,3 +180,6 @@ def test_unified_decode_graph_deferral_restores_waiting_queue(monkeypatch) -> No
         scheduler.schedule()
 
     assert scheduler.waiting._items == [deferred_by_upstream, original_waiting_req]
+    assert scheduler.waiting is original_waiting
+    assert scheduler.kv_holding_waiting is original_kv_holding
+    assert scheduler.kv_holding_waiting._items == [preempted_by_upstream, kv_holder]

tests/core/sched/test_omni_cfg_ar_scheduler.py

@@ -23,10 +23,12 @@ def _scheduler(monkeypatch, waiting, running=()):
     scheduler.waiting = create_request_queue(scheduler.policy)
     for request in waiting:
         scheduler.waiting.add_request(request)
-    scheduler.skipped_waiting = create_request_queue(scheduler.policy)
+    scheduler.kv_holding_waiting = create_request_queue(scheduler.policy)
+    scheduler.deferred_waiting = set()
     scheduler.running = list(running)
     scheduler.requests = {request.request_id: request for request in [*waiting, *running]}
     scheduler.max_model_len = 128
+    scheduler.max_num_active_reqs = 2
     scheduler.max_num_running_reqs = 2
     scheduler.max_num_scheduled_tokens = 256
     scheduler.block_size = 16

tests/core/sched/test_omni_generation_scheduler_update_session.py

@@ -52,10 +52,11 @@ class _SchedulerStub(OmniGenerationScheduler):
         self.num_waiting_for_streaming_input = 0
         self.log_stats = log_stats
         self.chunk_transfer_adapter = _ChunkTransferAdapterStub()
-        self.skipped_waiting = _SkippedWaitingStub()
+        self.kv_holding_waiting = _SkippedWaitingStub()
+        self.deferred_waiting: set[Request] = set()
 
     def _enqueue_waiting_request(self, session: Request) -> None:
-        raise AssertionError("unexpected enqueue for skipped_waiting miss")
+        raise AssertionError("unexpected enqueue for kv_holding_waiting miss")
 
 
 class _AsyncChunkStopSchedulerStub(OmniGenerationScheduler):
@@ -134,7 +135,8 @@ def test_resumable_generation_stop_marks_segment_boundary() -> None:
     sched._handle_stopped_request.side_effect = rearm_resumable_request
     sched.running = [session]
     sched.waiting = MagicMock()
-    sched.skipped_waiting = MagicMock()
+    sched.kv_holding_waiting = MagicMock()
+    sched.deferred_waiting = set()
     sched.structured_output_manager.accept_tokens.return_value = True
     # Async scheduling can observe the segment boundary in the next schedule()
     # before this model output is applied and defer the same request for finish.
@@ -206,7 +208,8 @@ def test_rejected_grammar_finishes_and_frees_generation_request(prompt_complete:
         OmniSchedulerMixin._remove_stopped_requests_from_queues(sched, running, preempted)
     )
     sched.waiting = MagicMock()
-    sched.skipped_waiting = MagicMock()
+    sched.kv_holding_waiting = MagicMock()
+    sched.deferred_waiting = set()
     sched.structured_output_manager.accept_tokens.return_value = False
     sched._pending_finish_reqs = []
     sched.recompute_kv_load_failures = False
@@ -350,6 +353,8 @@ class _RealignSchedulerStub(OmniSchedulerMixin):
         self.requests = requests
         self.running = running
         self.waiting = waiting
+        self.kv_holding_waiting: list[Request] = []
+        self.deferred_waiting: set[Request] = set()
 
 
 class TestRealignRequestStatusToQueues:

tests/core/sched/test_omni_sched_deferred_free_fence.py

@@ -123,6 +123,7 @@ def _make_fast_path_sched(*, defer_block_free: bool) -> tuple[MagicMock, Request
 
     sched = MagicMock()
     sched.max_num_scheduled_tokens = 8
+    sched.max_num_active_reqs = 8
     sched.max_num_running_reqs = 8
     sched._native_data_plane = False
     sched._retains_state_across_chunks = False
@@ -130,6 +131,8 @@ def _make_fast_path_sched(*, defer_block_free: bool) -> tuple[MagicMock, Request
     sched.requests = {request.request_id: request}
     sched.running = [request]
     sched.waiting = []  # falsy: the waiting loop never runs
+    sched.kv_holding_waiting = []
+    sched.deferred_waiting = set()
     sched.policy = SchedulingPolicy.FCFS
     sched.chunk_transfer_adapter = None
     sched.scheduler_config = MagicMock(enable_chunked_prefill=True)

tests/core/sched/test_omni_scheduler_finish_requests_purge.py

@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 """Integration tests for the post-finish defensive purge of
 ``self.running`` placed in
 ``OmniARScheduler.finish_requests`` /
@@ -52,6 +55,8 @@ def _make_scheduler(scheduler_cls, *, requests, running, waiting):
     scheduler.requests = requests
     scheduler.running = running
     scheduler.waiting = waiting
+    scheduler.kv_holding_waiting = []
+    scheduler.deferred_waiting = set()
     return scheduler
 
 

tests/core/sched/test_omni_scheduler_input_coordinator_cleanup.py

@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 from __future__ import annotations
 
 from types import SimpleNamespace
@@ -49,6 +52,8 @@ def test_finish_requests_cleans_input_coordinator_for_finished_ids(
     scheduler.requests = {}
     scheduler.running = []
     scheduler.waiting = []
+    scheduler.kv_holding_waiting = []
+    scheduler.deferred_waiting = set()
 
     def fake_finish_requests(self, request_ids, finished_status):
         assert request_ids == ["req-a", "req-b"]

tests/core/sched/test_omni_scheduler_mixin_shared.py

@@ -7,7 +7,10 @@ from typing import Any
 
 import pytest
 from vllm.config import SchedulerConfig, VllmConfig
+from vllm.sampling_params import SamplingParams
+from vllm.v1.core.sched.scheduler import Scheduler
 from vllm.v1.engine import FinishReason
+from vllm.v1.request import Request
 
 from vllm_omni.config.model import OmniModelConfig
 from vllm_omni.core.sched import omni_scheduler_mixin
@@ -18,7 +21,7 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
 
 class _Scheduler(OmniSchedulerMixin):
-    pass
+    _is_blocked_waiting_status = staticmethod(Scheduler._is_blocked_waiting_status)
 
 
 def test_async_chunk_adapter_initializes_for_stage_zero_sender_and_stage_one_receiver(monkeypatch):
@@ -120,7 +123,11 @@ def test_native_data_plane_uses_the_selected_input_protocol(async_chunk, stage_i
 def test_schedule_lifecycle_helpers_process_and_restore_both_input_paths():
     calls: list[tuple[Any, ...]] = []
     scheduler = _Scheduler()
-    scheduler.waiting = ["waiting"]
+    waiting = Request("waiting", [1], SamplingParams(max_tokens=1), pooling_params=None)
+    holding = Request("holding", [1, 2], SamplingParams(max_tokens=1), pooling_params=None)
+    holding.num_computed_tokens = 1
+    scheduler.waiting = [waiting]
+    scheduler.kv_holding_waiting = [holding]
     scheduler.running = ["running"]
     scheduler.requests = {"request": object()}
     scheduler._consume_pending_connector_output = lambda mode: calls.append(("consume", mode))
@@ -137,16 +144,16 @@ def test_schedule_lifecycle_helpers_process_and_restore_both_input_paths():
     scheduler.chunk_transfer_adapter = SimpleNamespace(
         receives_chunks=True,
         process_pending_chunks=lambda waiting, running, scheduler_requests: calls.append(
-            ("process", waiting, running, scheduler_requests)
+            ("process", list(waiting), running, scheduler_requests)
         ),
         restore_queues=lambda waiting, running, scheduler_requests: calls.append(
-            ("restore-chunks", waiting, running, scheduler_requests)
+            ("restore-chunks", list(waiting), running, scheduler_requests)
         ),
         collect_timed_out_request_ids=_collect_timed_out,
         collect_failed_send_request_ids=_collect_failed_sends,
     )
     scheduler.input_coordinator = SimpleNamespace(
-        restore_queues=lambda waiting, running: calls.append(("restore-full", waiting, running))
+        restore_queues=lambda waiting, running: calls.append(("restore-full", list(waiting), running))
     )
 
     scheduler._process_pending_omni_inputs("ar")
@@ -155,13 +162,13 @@ def test_schedule_lifecycle_helpers_process_and_restore_both_input_paths():
     assert calls == [
         ("consume", "ar"),
         ("timeouts",),
-        ("process", scheduler.waiting, scheduler.running, scheduler.requests),
+        ("process", [holding, waiting], scheduler.running, scheduler.requests),
         # The chunk deadline runs after chunks are applied, so a chunk that
         # arrived this cycle resets the clock before it is measured (R1.1).
         ("chunk-timeouts", omni_scheduler_mixin.DEFAULT_INPUT_WAIT_TIMEOUT_S),
         ("failed-sends",),
-        ("restore-chunks", scheduler.waiting, scheduler.running, scheduler.requests),
-        ("restore-full", scheduler.waiting, scheduler.running),
+        ("restore-chunks", [holding, waiting], scheduler.running, scheduler.requests),
+        ("restore-full", [holding, waiting], scheduler.running),
     ]
 
 

tests/core/sched/test_omni_scheduler_ready_inbox.py

@@ -5,8 +5,13 @@
 from types import SimpleNamespace
 
 import pytest
+from vllm import SamplingParams
+from vllm.v1.core.sched.request_queue import SchedulingPolicy, create_request_queue
+from vllm.v1.request import Request, RequestStatus
 
+from vllm_omni.core.sched.omni_ar_scheduler import OmniARScheduler
 from vllm_omni.core.sched.omni_scheduler_mixin import OmniSchedulerMixin
+from vllm_omni.core.sched.omni_scheduling_coordinator import OmniSchedulingCoordinator
 from vllm_omni.outputs import OmniConnectorOutput
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
@@ -15,6 +20,8 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 def test_ready_inbox_coalesces_live_events_and_drops_cancelled(mocker):
     scheduler = OmniSchedulerMixin()
     scheduler.requests, scheduler.waiting, scheduler.running = {"r": object()}, [], []
+    scheduler.kv_holding_waiting = []
+    scheduler._is_blocked_waiting_status = OmniARScheduler._is_blocked_waiting_status
     coordinator = SimpleNamespace(
         _async_chunk=True, update_request_metadata=mocker.Mock(), process_pending_chunks=mocker.Mock()
     )
@@ -32,7 +39,92 @@ def test_ready_inbox_coalesces_live_events_and_drops_cancelled(mocker):
     coordinator.update_request_metadata.assert_called_once_with(
         scheduler.requests, {"r": {"decode_token_end": 2}}, model_mode="ar"
     )
-    coordinator.process_pending_chunks.assert_called_once_with([], [], {"r"}, {"r"})
+    coordinator.process_pending_chunks.assert_called_once_with(mocker.ANY, [], {"r"}, {"r"})
+    assert list(coordinator.process_pending_chunks.call_args.args[0]) == []
     scheduler._consume_pending_connector_output(model_mode="ar")
     assert coordinator.update_request_metadata.call_count == 1
-    coordinator.process_pending_chunks.assert_called_with([], [], set(), set())
+    coordinator.process_pending_chunks.assert_called_with(mocker.ANY, [], set(), set())
+    assert list(coordinator.process_pending_chunks.call_args.args[0]) == []
+
+
+@pytest.mark.parametrize("policy", [SchedulingPolicy.FCFS, SchedulingPolicy.PRIORITY])
+@pytest.mark.parametrize("async_chunk", [False, True])
+def test_native_input_gate_parks_and_restores_kv_holders(policy, async_chunk):
+    scheduler = OmniARScheduler.__new__(OmniARScheduler)
+    scheduler.waiting = create_request_queue(policy)
+    scheduler.kv_holding_waiting = create_request_queue(policy)
+    scheduler.deferred_waiting = set()
+    scheduler.running = []
+    scheduler.chunk_transfer_adapter = None
+    scheduler.input_coordinator = OmniSchedulingCoordinator(
+        stage_id=1, scheduler_max_num_seqs=2, async_chunk=async_chunk
+    )
+    scheduler._init_omni_connector_output_inbox()
+    holder = Request("holder", [1, 2, 3], SamplingParams(max_tokens=1), pooling_params=None)
+    holder.num_computed_tokens = 2
+    fresh = Request("fresh", [1], SamplingParams(max_tokens=1), pooling_params=None)
+    scheduler.requests = {request.request_id: request for request in (holder, fresh)}
+    scheduler._enqueue_waiting_request(fresh)
+    scheduler._enqueue_waiting_request(holder)
+
+    scheduler._consume_pending_connector_output(model_mode="ar")
+
+    assert not scheduler.waiting and not scheduler.kv_holding_waiting
+    expected_status = RequestStatus.WAITING_FOR_CHUNK if async_chunk else RequestStatus.WAITING_FOR_INPUT
+    assert holder.status == fresh.status == expected_status
+    coordinator = scheduler.input_coordinator
+    registrations = coordinator.pending_chunk_registrations if async_chunk else coordinator.pending_input_registrations
+    assert [entry.request_id for entry in registrations] == ["holder", "fresh"]
+
+    scheduler._restore_omni_wait_queues()
+    assert list(scheduler.kv_holding_waiting) == [holder]
+    assert list(scheduler.waiting) == [fresh]
+    ready_output = (
+        OmniConnectorOutput(chunk_ready_req_ids={"holder"})
+        if async_chunk
+        else OmniConnectorOutput(stage_recv_req_ids={"holder"})
+    )
+    scheduler.enqueue_omni_connector_output(ready_output)
+    scheduler._consume_pending_connector_output(model_mode="ar")
+
+    assert holder.status == RequestStatus.WAITING
+    assert fresh.status == expected_status
+    assert list(scheduler.kv_holding_waiting) == [holder]
+    assert not scheduler.waiting
+    scheduler._restore_omni_wait_queues()
+    assert list(scheduler.kv_holding_waiting) == [holder]
+    assert list(scheduler.waiting) == [fresh]
+
+
+@pytest.mark.parametrize("async_chunk", [False, True])
+@pytest.mark.parametrize(
+    "status",
+    [
+        RequestStatus.WAITING_FOR_STREAMING_REQ,
+        RequestStatus.WAITING_FOR_REMOTE_KVS,
+        RequestStatus.WAITING_FOR_STRUCTURED_OUTPUT_GRAMMAR,
+    ],
+)
+def test_native_input_gate_preserves_upstream_blocked_waits(async_chunk, status):
+    scheduler = OmniARScheduler.__new__(OmniARScheduler)
+    scheduler.waiting = create_request_queue(SchedulingPolicy.FCFS)
+    scheduler.kv_holding_waiting = create_request_queue(SchedulingPolicy.FCFS)
+    scheduler.deferred_waiting = set()
+    scheduler.running = []
+    request = Request("blocked", [1, 2, 3], SamplingParams(max_tokens=1), pooling_params=None)
+    request.status = status
+    request.num_computed_tokens = 2
+    scheduler.requests = {request.request_id: request}
+    scheduler._enqueue_waiting_request(request)
+    scheduler.input_coordinator = OmniSchedulingCoordinator(
+        stage_id=1, scheduler_max_num_seqs=1, async_chunk=async_chunk
+    )
+    scheduler._init_omni_connector_output_inbox()
+
+    scheduler._consume_pending_connector_output(model_mode="ar")
+
+    assert request.status == status
+    assert list(scheduler.kv_holding_waiting) == [request]
+    assert scheduler.deferred_waiting == {request}
+    assert not scheduler.input_coordinator.pending_chunk_registrations
+    assert not scheduler.input_coordinator.pending_input_registrations

tests/core/sched/test_omni_scheduler_streaming_input_counter.py

@@ -59,7 +59,8 @@ def _make_scheduler(scheduler_cls, *, requests, running, waiting, counter, skipp
     scheduler.requests = requests
     scheduler.running = running
     scheduler.waiting = waiting
-    scheduler.skipped_waiting = skipped if skipped is not None else []
+    scheduler.kv_holding_waiting = skipped if skipped is not None else []
+    scheduler.deferred_waiting = set()
     scheduler.num_waiting_for_streaming_input = counter
     return scheduler
 
@@ -97,8 +98,8 @@ def test_finish_requests_clears_counter_left_by_a_stomped_status(
     assert scheduler.num_waiting_for_streaming_input == 0
 
 
-def test_resync_counts_parked_requests_in_skipped_waiting() -> None:
-    """``get_num_unfinished_requests`` sums ``waiting`` and ``skipped_waiting``,
+def test_resync_counts_parked_requests_in_kv_holding_waiting() -> None:
+    """``get_num_unfinished_requests`` sums ``waiting`` and ``kv_holding_waiting``,
     so the resync has to look at both or it under-counts and reintroduces the
     spin the counter exists to prevent."""
     parked = _StubRequest("req-skipped", RequestStatus.WAITING_FOR_STREAMING_REQ)
@@ -128,9 +129,11 @@ def test_ar_schedule_resyncs_before_delegating_upstream(monkeypatch: pytest.Monk
     scheduler.requests = {parked.request_id: parked}
     scheduler.running = [parked]
     scheduler.waiting = []
-    scheduler.skipped_waiting = []
+    scheduler.kv_holding_waiting = []
+    scheduler.deferred_waiting = set()
     scheduler.num_waiting_for_streaming_input = 1  # leaked by the status rewrite
     scheduler.policy = "fcfs"
+    scheduler.max_num_active_reqs = 8
     scheduler.max_num_running_reqs = 8
 
     seen: dict[str, int] = {}
@@ -161,12 +164,13 @@ def test_generation_scheduler_sweep_removes_aborted_requests_from_all_queues() -
     scheduler = OmniGenerationScheduler.__new__(OmniGenerationScheduler)
     scheduler.waiting = create_request_queue(SchedulingPolicy.FCFS)
     scheduler.waiting.add_request(aborted)
-    scheduler.skipped_waiting = create_request_queue(SchedulingPolicy.FCFS)
-    scheduler.skipped_waiting.add_request(aborted)
+    scheduler.kv_holding_waiting = create_request_queue(SchedulingPolicy.FCFS)
+    scheduler.deferred_waiting = set()
+    scheduler.kv_holding_waiting.add_request(aborted)
     scheduler.running = [aborted]
 
     scheduler._drop_aborted_queued_requests()
 
     assert list(scheduler.waiting) == []
-    assert list(scheduler.skipped_waiting) == []
+    assert list(scheduler.kv_holding_waiting) == []
     assert scheduler.running == []

tests/diffusion/cache/test_seacache.py

@@ -329,7 +329,9 @@ def _uneven_cfg_sharded_worker(
         init_method=init_method,
         rank=rank,
         world_size=4,
-        timeout=datetime.timedelta(seconds=10),
+        # Spawned workers import torch/vLLM before joining; allow startup
+        # skew without changing the collective deadlock checks below.
+        timeout=datetime.timedelta(seconds=60),
     )
     try:
         fs_groups = [

tests/diffusion/diffusion_kv/test_kv_cache_utils.py

@@ -103,7 +103,7 @@ def test_block_hashing_uses_upstream_extra_keys_directly(monkeypatch):
     # Independent expected native extras: a crossing item has a negative
     # relative offset; a block can include the end of one item and the next.
     first, second = (feature.identifier for feature in row.mm_features)
-    extras = [((first, 2),), ((first, -2), (second, 3)), ((second, -1),)]
+    extras = [(("mm", first, 2),), (("mm", first, -2), ("mm", second, 3)), (("mm", second, -1),)]
     parent = BlockHash(hash_fn(("diffusion-prefix-root-v1", row.cache_namespace)))
     expected = []
     for start, keys in zip(range(0, 12, 4), extras, strict=True):

tests/diffusion/diffusion_kv/test_mooncake_cfg_fanout.py

@@ -31,6 +31,10 @@ def producer(monkeypatch):
     # vLLM 0.30 send_kv_to_decode also reads the PP size (ce08bb5b34,
     # #56033), the MLA/Mamba flags and the head-resharding validator.
     worker.pp_size = 1
+    worker.opaque_packed_storages = set()
+    worker._prepared_transfer_regions = {}
+    worker.region_shared_groups = []
+    worker.region_row_offsets = []
     worker.use_mla = False
     worker.kv_cache_config = SimpleNamespace(has_mamba_layers=False)
     worker._validate_head_resharding_layout = Mock(return_value=None)

tests/diffusion/diffusion_kv/test_paged_attention_adapter.py

@@ -1253,14 +1253,15 @@ def test_layer_adapter_accepts_platform_native_backend_and_uses_rank_local_heads
         speculative_config=None,
     )
 
-    def select_backend(*, backend, attn_selector_config, num_heads):
-        # Exercise the real 0.29 selector, replacing only platform resolution.
+    def select_backend(*, backend, attn_selector_config, num_heads, _run_kv_cache_dtype):
+        # Exercise the real upstream selector, replacing only platform resolution.
         # Ulysses has already gathered tokens and sharded heads before the
         # native kernel; the MoE PCP mapping must not request PCP attention.
         assert not attn_selector_config.use_pcp
         assert not attn_selector_config.use_dcp
         assert attn_selector_config.use_non_causal
         assert num_heads == 4
+        assert _run_kv_cache_dtype == config.cache_config.cache_dtype
         selected_backends.append((backend, config.attention_config.backend_per_kind))
         if selection_fails:
             raise ValueError("test backend unavailable")

tests/diffusion/distributed/test_comm.py

@@ -14,6 +14,7 @@ from typing import Literal
 import pytest
 import torch
 
+from tests.helpers.distributed import configure_cpu_collective_worker, start_collective_workers
 from tests.helpers.mark import hardware_marks
 from vllm_omni.diffusion.distributed.comm import RingComm, SeqAllToAll4D, SeqAllToAll5D
 from vllm_omni.diffusion.distributed.parallel_state import (
@@ -56,6 +57,8 @@ def _init_worker(
             "MASTER_PORT": str(master_port),
         }
     )
+    if device_kind == "cpu":
+        configure_cpu_collective_worker()
     backend = "gloo" if device_kind == "cpu" else None
     init_distributed_environment(backend=backend)
 
@@ -243,7 +246,7 @@ def _spawn_4d_identity(
     device_kind: DeviceKind,
     master_port: int,
 ) -> None:
-    torch.multiprocessing.spawn(
+    start_collective_workers(
         _run_4d_identity,
         args=(
             world_size,
@@ -257,6 +260,7 @@ def _spawn_4d_identity(
             master_port,
         ),
         nprocs=world_size,
+        device_kind=device_kind,
     )
 
 
@@ -272,7 +276,7 @@ def _spawn_5d_identity(
     device_kind: DeviceKind,
     master_port: int,
 ) -> None:
-    torch.multiprocessing.spawn(
+    start_collective_workers(
         _run_5d_identity,
         args=(
             world_size,
@@ -286,6 +290,7 @@ def _spawn_5d_identity(
             master_port,
         ),
         nprocs=world_size,
+        device_kind=device_kind,
     )
 
 
@@ -299,10 +304,11 @@ def _spawn_ring_p2p(
     device_kind: DeviceKind,
     master_port: int,
 ) -> None:
-    torch.multiprocessing.spawn(
+    start_collective_workers(
         _run_ring_p2p,
         args=(world_size, dtype, batch_size, num_heads, head_size, device_kind, master_port),
         nprocs=world_size,
+        device_kind=device_kind,
     )
 
 
@@ -499,8 +505,9 @@ def _run_fused_qkv_matches_three_4d(
 @pytest.mark.cpu
 @pytest.mark.parametrize("world_size", [2, 4])
 def test_fused_qkv_all_to_all_matches_three_4d_exchanges_cpu(world_size: int) -> None:
-    torch.multiprocessing.spawn(
+    start_collective_workers(
         _run_fused_qkv_matches_three_4d,
         args=(world_size, "cpu", 29613),
         nprocs=world_size,
+        device_kind="cpu",
     )

tests/diffusion/distributed/test_pipeline_parallel.py

@@ -19,6 +19,7 @@ from vllm.utils.network_utils import get_file_store_init_method
 from vllm.v1.worker.gpu_worker import AsyncIntermediateTensors
 
 import vllm_omni.diffusion.distributed.pipeline_parallel as pp_module
+from tests.helpers.distributed import collective_context, start_collective_workers
 from tests.helpers.mark import hardware_marks
 from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin
 from vllm_omni.diffusion.distributed.parallel_state import (
@@ -39,6 +40,25 @@ DeviceKind = Literal["cpu", "cuda"]
 _UNIT_MARKS = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
 
 
+@pytest.fixture(autouse=True)
+def _restore_collective_platforms():
+    """CPU baselines run in the parent process; retain its original platforms."""
+    import vllm.platforms
+
+    from vllm_omni.diffusion.distributed import comm, group_coordinator, parallel_state
+
+    original_platform = vllm.platforms.current_platform
+    original_omni_platforms = {
+        module: module.current_omni_platform for module in (comm, group_coordinator, parallel_state)
+    }
+    try:
+        yield
+    finally:
+        vllm.platforms.current_platform = original_platform
+        for module, platform in original_omni_platforms.items():
+            module.current_omni_platform = platform
+
+
 def _cleanup_distributed() -> None:
     """Destroy omni distributed state left dangling after a test.
 
@@ -409,6 +429,10 @@ def init_dist(
     device = _worker_device(local_rank, device_kind)
     if device_kind == "cuda":
         current_omni_platform.set_device(device)
+    else:
+        from tests.helpers.distributed import configure_cpu_collective_worker
+
+        configure_cpu_collective_worker()
     backend = "gloo" if device_kind == "cpu" else None
     init_distributed_environment(
         world_size=world_size,
@@ -500,13 +524,14 @@ def isend_irecv_worker(
 
 
 def _run_isend_irecv(pp_size: int, device_kind: DeviceKind, init_method: str) -> None:
-    mp_context = torch.multiprocessing.get_context("spawn")
+    mp_context = collective_context(device_kind)
     manager = mp_context.Manager()
     q = manager.Queue()
-    torch.multiprocessing.spawn(
+    start_collective_workers(
         isend_irecv_worker,
         args=(pp_size, init_method, device_kind, q),
         nprocs=pp_size,
+        device_kind=device_kind,
     )
     results = {label: tensor for label, tensor in [q.get(), q.get()]}
     torch.testing.assert_close(
@@ -656,15 +681,16 @@ def _run_predict_noise(
 
     baseline_out = compute_single_gpu_baseline(test_config, dtype, do_true_cfg, device_kind)
 
-    mp_context = torch.multiprocessing.get_context("spawn")
+    mp_context = collective_context(device_kind)
     manager = mp_context.Manager()
     pp_q = manager.Queue()
 
     world_size = pp_size * cfg_size
-    torch.multiprocessing.spawn(
+    start_collective_workers(
         predict_noise_worker,
         args=(world_size, init_method, pp_size, cfg_size, do_true_cfg, dtype, test_config, device_kind, pp_q),
         nprocs=world_size,
+        device_kind=device_kind,
     )
 
     pp_out = pp_q.get()
@@ -900,15 +926,16 @@ def _run_scheduler_step(
 
     baseline = compute_scheduler_step_baseline(test_config, do_true_cfg, device_kind)
 
-    mp_context = torch.multiprocessing.get_context("spawn")
+    mp_context = collective_context(device_kind)
     manager = mp_context.Manager()
     q = manager.Queue()
 
     world_size = pp_size * cfg_size
-    torch.multiprocessing.spawn(
+    start_collective_workers(
         scheduler_step_worker,
         args=(world_size, init_method, pp_size, cfg_size, do_true_cfg, test_config, device_kind, q),
         nprocs=world_size,
+        device_kind=device_kind,
     )
 
     result = q.get()

tests/diffusion/models/boogu_image/test_pipeline_boogu_image.py

@@ -239,11 +239,11 @@ def test_constructor_rejects_unsupported_mllm_quantization(mock_dependencies, mo
 )
 def test_constructor_routes_mllm_quantization(mock_dependencies, quantization_config, quantize_mllm):
     from transformers import AutoModel
-    from vllm.model_executor.layers.quantization.fp8 import Fp8Config
     from vllm.model_executor.layers.quantization.online.base import OnlineQuantizationConfig
     from vllm.model_executor.layers.quantization.utils.quant_utils import kFp8Static128BlockSym
 
     from vllm_omni.diffusion.models.boogu_image.pipeline_boogu_image import BooguImagePipeline
+    from vllm_omni.quantization.fp8_config import DiffusionFp8Config
 
     od_config = OmniDiffusionConfig(
         model="dummy-boogu",
@@ -287,15 +287,14 @@ def test_constructor_routes_mllm_quantization(mock_dependencies, quantization_co
         mock_dependencies["mllm_config_loader"].assert_not_called()
         mock_dependencies["mllm_builder"].assert_not_called()
         assert [source.subfolder for source in pipeline.weights_sources] == ["transformer"]
-    assert isinstance(mock_dependencies["transformer_cls"].call_args.kwargs["quant_config"], Fp8Config)
+    assert isinstance(mock_dependencies["transformer_cls"].call_args.kwargs["quant_config"], DiffusionFp8Config)
 
 
 def test_constructor_preserves_mllm_ignored_layers(mock_dependencies):
-    from vllm.model_executor.layers.quantization.fp8 import Fp8Config
-
     from vllm_omni.diffusion.models.boogu_image.pipeline_boogu_image import BooguImagePipeline
+    from vllm_omni.quantization.fp8_config import DiffusionFp8Config
 
-    quant_config = Fp8Config(
+    quant_config = DiffusionFp8Config(
         ignored_layers=["mllm.language_model.layers.0.self_attn.q_proj", "transformer.blocks.0.attn.to_q"]
     )
     od_config = OmniDiffusionConfig(

tests/diffusion/models/flux2/test_flux2_quantization.py

@@ -16,17 +16,13 @@ class _RecordingBlock(nn.Module):
 
 
 def test_flux2_propagates_online_quant_config_to_all_transformer_blocks(monkeypatch):
-    from vllm.model_executor.layers.quantization.fp8 import Fp8Config
-
     from vllm_omni.diffusion.models.flux2 import flux2_transformer
+    from vllm_omni.quantization import build_quant_config
 
     monkeypatch.setattr(flux2_transformer, "Flux2TransformerBlock", _RecordingBlock)
     monkeypatch.setattr(flux2_transformer, "Flux2SingleTransformerBlock", _RecordingBlock)
 
-    quant_config = Fp8Config(
-        is_checkpoint_fp8_serialized=False,
-        activation_scheme="dynamic",
-    )
+    quant_config = build_quant_config("fp8")
     model = flux2_transformer.Flux2Transformer2DModel(
         in_channels=4,
         out_channels=4,

tests/diffusion/models/minimax_h3/test_minimax_h3_quantization.py

@@ -71,12 +71,12 @@ def _small_od_config(*, class_name=None):
 
 def test_fp8_scope_and_prefix_propagation(monkeypatch):
     from vllm.model_executor.layers.linear import LinearBase, UnquantizedLinearMethod
-    from vllm.model_executor.layers.quantization.fp8 import Fp8Config
     from vllm.model_executor.layers.quantization.utils.quant_utils import (
         is_layer_skipped,
     )
 
     from vllm_omni.diffusion.models.minimax_h3 import minimax_h3_transformer as h3
+    from vllm_omni.quantization.fp8_config import DiffusionFp8Config
 
     monkeypatch.setattr(h3, "ColumnParallelLinear", _FakeLinear)
     monkeypatch.setattr(h3, "MergedColumnParallelLinear", _FakeLinear)
@@ -92,7 +92,7 @@ def test_fp8_scope_and_prefix_propagation(monkeypatch):
         "blocks.0.adaln_proj.linear",
         "final_layer.adaln_proj.linear",
     }
-    fp8_config = Fp8Config(ignored_layers=sorted(ignored_layers))
+    fp8_config = DiffusionFp8Config(ignored_layers=sorted(ignored_layers))
     model = h3.MiniMaxH3DiTModel(
         _small_od_config(),
         quant_config=fp8_config,

tests/diffusion/quantization/test_fp8_config.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 """Tests for the unified quantization framework."""
 
 import pytest
@@ -15,6 +15,74 @@ def test_build_quant_config_fp8():
     assert config is not None
     assert config.get_name() == "fp8"
     assert config.activation_scheme == "dynamic"
+    assert not config.is_checkpoint_fp8_serialized
+
+
+def test_build_quant_config_fp8_uses_native_online_quantization():
+    from vllm.model_executor.layers.quantization.online.base import OnlineQuantizationConfig
+    from vllm.model_executor.layers.quantization.utils.quant_utils import kFp8StaticTensorSym
+
+    from vllm_omni.quantization import build_quant_config
+
+    config = build_quant_config("fp8", ignored_layers=["proj_out"])
+    assert isinstance(config, OnlineQuantizationConfig)
+    assert config.args.linear.weight == kFp8StaticTensorSym
+    assert config.args.moe.weight == kFp8StaticTensorSym
+    assert config.ignored_layers == ["proj_out"]
+
+
+def test_build_quant_config_upstream_online_fp8_shorthand():
+    from vllm.model_executor.layers.quantization.online.base import OnlineQuantizationConfig
+    from vllm.model_executor.layers.quantization.utils.quant_utils import kFp8StaticTensorSym
+
+    from vllm_omni.quantization import build_quant_config
+
+    config = build_quant_config("fp8_per_tensor", ignore=["proj_out"])
+    assert isinstance(config, OnlineQuantizationConfig)
+    assert config.args.linear.weight == kFp8StaticTensorSym
+    assert config.ignored_layers == ["proj_out"]
+
+
+def test_online_fp8_rejects_static_activation():
+    from vllm_omni.quantization import build_quant_config
+
+    with pytest.raises(ValueError, match="activation_scheme='dynamic'"):
+        build_quant_config("fp8", activation_scheme="static")
+
+
+def test_build_online_quant_config_preserves_explicit_native_args():
+    from vllm.config.quantization import QuantizationConfigArgs, QuantSpec
+    from vllm.model_executor.layers.quantization.online.base import OnlineQuantizationConfig
+
+    from vllm_omni.quantization import build_quant_config
+
+    args = QuantizationConfigArgs(linear=QuantSpec(weight="fp8_per_tensor_static"), ignore=["proj_out"])
+    config = build_quant_config("online", args=args)
+    assert isinstance(config, OnlineQuantizationConfig)
+    assert config.args is args
+    assert config.ignored_layers == ["proj_out"]
+
+
+def test_build_online_quant_config_requires_args():
+    from vllm_omni.quantization import build_quant_config
+
+    with pytest.raises(ValueError, match="requires quantization config arguments"):
+        build_quant_config("online")
+
+
+def test_serialized_checkpoint_replaces_online_fp8_config():
+    from vllm.model_executor.layers.quantization.fp8 import Fp8Config
+
+    from vllm_omni.quantization import build_quant_config
+    from vllm_omni.quantization.factory import resolve_quant_config_from_disk
+
+    config = resolve_quant_config_from_disk(
+        build_quant_config("fp8"),
+        {"quant_method": "fp8", "is_checkpoint_fp8_serialized": True, "activation_scheme": "static"},
+    )
+    assert isinstance(config, Fp8Config)
+    assert config.is_checkpoint_fp8_serialized
+    assert config.activation_scheme == "static"
 
 
 def test_build_quant_config_none():
@@ -39,7 +107,7 @@ def test_build_quant_config_invalid():
 def test_build_quant_config_dict():
     from vllm_omni.quantization import build_quant_config
 
-    config = build_quant_config({"method": "fp8", "activation_scheme": "static"})
+    config = build_quant_config({"method": "fp8", "is_checkpoint_fp8_serialized": True, "activation_scheme": "static"})
     assert config is not None
     assert config.get_name() == "fp8"
     assert config.activation_scheme == "static"
@@ -48,7 +116,7 @@ def test_build_quant_config_dict():
 def test_build_quant_config_dict_not_mutated():
     from vllm_omni.quantization import build_quant_config
 
-    original = {"method": "fp8", "activation_scheme": "static"}
+    original = {"method": "fp8", "is_checkpoint_fp8_serialized": True, "activation_scheme": "static"}
     copy = original.copy()
     build_quant_config(original)
     assert original == copy
@@ -99,7 +167,7 @@ def test_build_quant_config_per_component_inner_dict_not_mutated():
     """Inner component dicts should not be mutated by build_quant_config."""
     from vllm_omni.quantization import build_quant_config
 
-    inner = {"method": "fp8", "activation_scheme": "static"}
+    inner = {"method": "fp8", "is_checkpoint_fp8_serialized": True, "activation_scheme": "static"}
     original = inner.copy()
     build_quant_config({"transformer": inner, "vae": None})
     assert inner == original
@@ -115,20 +183,18 @@ def test_flat_dict_not_misdetected_as_per_component():
 
 
 def test_build_quant_config_passthrough():
-    from vllm.model_executor.layers.quantization.fp8 import Fp8Config
-
     from vllm_omni.quantization import build_quant_config
+    from vllm_omni.quantization.fp8_config import DiffusionFp8Config
 
-    fp8 = Fp8Config(is_checkpoint_fp8_serialized=False, activation_scheme="dynamic")
+    fp8 = DiffusionFp8Config()
     assert build_quant_config(fp8) is fp8
 
 
 def test_component_config_routing():
-    from vllm.model_executor.layers.quantization.fp8 import Fp8Config
-
     from vllm_omni.quantization import ComponentQuantizationConfig
+    from vllm_omni.quantization.fp8_config import DiffusionFp8Config
 
-    fp8 = Fp8Config(is_checkpoint_fp8_serialized=False, activation_scheme="dynamic")
+    fp8 = DiffusionFp8Config()
     config = ComponentQuantizationConfig(
         component_configs={"transformer": fp8, "vae": None},
     )
@@ -140,11 +206,10 @@ def test_component_config_routing():
 
 
 def test_component_config_with_default():
-    from vllm.model_executor.layers.quantization.fp8 import Fp8Config
-
     from vllm_omni.quantization import ComponentQuantizationConfig
+    from vllm_omni.quantization.fp8_config import DiffusionFp8Config
 
-    fp8 = Fp8Config(is_checkpoint_fp8_serialized=False, activation_scheme="dynamic")
+    fp8 = DiffusionFp8Config()
     config = ComponentQuantizationConfig(
         component_configs={"vae": None},
         default_config=fp8,
@@ -167,7 +232,7 @@ def test_integration_dict():
 
     config = OmniDiffusionConfig(
         model="test",
-        quantization_config={"method": "fp8", "activation_scheme": "static"},
+        quantization_config={"method": "fp8", "is_checkpoint_fp8_serialized": True, "activation_scheme": "static"},
     )
     assert config.quantization_config is not None
     assert config.quantization_config.get_name() == "fp8"

tests/diffusion/quantization/test_online_fp8_gpu.py

@@ -0,0 +1,72 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Exercise online FP8 loading and inference through the native CUDA path."""
+
+import pytest
+import torch
+from torch.nn import functional as F
+from vllm.config import set_current_vllm_config
+from vllm.distributed import (
+    destroy_distributed_environment,
+    destroy_model_parallel,
+    init_distributed_environment,
+    initialize_model_parallel,
+)
+from vllm.model_executor.layers.linear import ReplicatedLinear, UnquantizedLinearMethod
+from vllm.model_executor.layers.quantization.online.fp8 import Fp8PerTensorOnlineLinearMethod
+from vllm.utils.network_utils import get_open_port
+from vllm.utils.torch_utils import set_default_torch_dtype
+
+from tests.helpers.mark import hardware_test
+from vllm_omni.diffusion.data import OmniDiffusionConfig
+from vllm_omni.diffusion.vllm_config import create_diffusion_vllm_config
+from vllm_omni.quantization import build_quant_config
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion]
+
+
+@pytest.mark.parametrize("method", ["fp8", "fp8_per_tensor"])
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+def test_online_fp8_loads_bf16_weights_and_runs_native_kernel(method):
+    ignored_key = "ignored_layers" if method == "fp8" else "ignore"
+    config = build_quant_config(method, **{ignored_key: ["skip"]})
+    od_config = OmniDiffusionConfig(model="test", dtype=torch.bfloat16, quantization_config=config)
+    vllm_config = create_diffusion_vllm_config(torch.device("cuda:0"), od_config)
+
+    try:
+        init_distributed_environment(
+            world_size=1,
+            rank=0,
+            local_rank=0,
+            distributed_init_method=f"tcp://127.0.0.1:{get_open_port()}",
+            backend="nccl",
+        )
+        with set_current_vllm_config(vllm_config), set_default_torch_dtype(torch.bfloat16), torch.device("cuda:0"):
+            initialize_model_parallel(tensor_model_parallel_size=1, pipeline_model_parallel_size=1)
+            layer = ReplicatedLinear(
+                128, 128, bias=False, params_dtype=torch.bfloat16, quant_config=config, prefix="linear", disable_tp=True
+            )
+            assert isinstance(layer.quant_method, Fp8PerTensorOnlineLinearMethod)
+            assert layer.weight.is_meta
+
+            generator = torch.Generator(device="cuda").manual_seed(42)
+            weights = torch.randn(128, 128, generator=generator, dtype=torch.bfloat16) / 128**0.5
+            inputs = torch.randn(32, 128, generator=generator, dtype=torch.bfloat16)
+            expected = F.linear(inputs, weights).float()
+
+            # This is the loader installed by the native online quantizer.
+            # It materializes and quantizes the BF16 checkpoint in one pass.
+            layer.weight.weight_loader(layer.weight, weights)
+            assert layer.weight.dtype == torch.float8_e4m3fn
+            actual = layer(inputs)[0].float()
+            relative_mse = (actual - expected).square().mean() / expected.square().mean()
+            assert relative_mse.item() < 0.002
+
+            ignored = ReplicatedLinear(
+                128, 128, bias=False, params_dtype=torch.bfloat16, quant_config=config, prefix="skip", disable_tp=True
+            )
+            assert isinstance(ignored.quant_method, UnquantizedLinearMethod)
+            assert ignored.weight.dtype == torch.bfloat16
+    finally:
+        destroy_model_parallel()
+        destroy_distributed_environment()

tests/distributed/omni_connectors/test_chunk_transfer_adapter.py

@@ -11,6 +11,7 @@ import pytest
 import torch
 from pytest_mock import MockerFixture
 from vllm.sampling_params import SamplingParams
+from vllm.v1.core.sched.request_queue import SchedulingPolicy, create_request_queue
 from vllm.v1.core.sched.scheduler import Scheduler as VLLMScheduler
 from vllm.v1.metrics.stats import PrefillStats, PromptTokenStats
 from vllm.v1.request import Request, RequestStatus
@@ -610,6 +611,45 @@ def build_adapter(monkeypatch, mocker: MockerFixture):
     return _build
 
 
+@pytest.mark.parametrize("active_stream_window", [0, 2])
+def test_scheduler_chunk_gate_preserves_kv_holding_queue(build_adapter, active_stream_window):
+    adapter, _ = build_adapter(active_stream_window=active_stream_window)
+    scheduler = OmniARScheduler.__new__(OmniARScheduler)
+    scheduler.waiting = create_request_queue(SchedulingPolicy.FCFS)
+    scheduler.kv_holding_waiting = create_request_queue(SchedulingPolicy.FCFS)
+    scheduler.deferred_waiting = set()
+    scheduler.running = []
+    scheduler.chunk_transfer_adapter = adapter
+    scheduler.input_coordinator = None
+    holder = Request("holder", [1, 2, 3], SamplingParams(max_tokens=1), pooling_params=None)
+    holder.num_computed_tokens = 2
+    fresh = Request("fresh", [1], SamplingParams(max_tokens=1), pooling_params=None)
+    scheduler.requests = {request.request_id: request for request in (holder, fresh)}
+    scheduler._enqueue_waiting_request(fresh)
+    scheduler._enqueue_waiting_request(holder)
+
+    scheduler._process_pending_omni_inputs(model_mode="ar")
+
+    assert not scheduler.waiting and not scheduler.kv_holding_waiting
+    assert list(adapter.waiting_for_chunk_waiting_requests) == [holder, fresh]
+    assert holder.status == fresh.status == RequestStatus.WAITING_FOR_CHUNK
+    scheduler._restore_omni_wait_queues()
+    assert list(scheduler.kv_holding_waiting) == [holder]
+    assert list(scheduler.waiting) == [fresh]
+
+    adapter._finished_load_reqs.add(holder.request_id)
+    scheduler._process_pending_omni_inputs(model_mode="ar")
+
+    assert holder.status == RequestStatus.WAITING
+    assert fresh.status == RequestStatus.WAITING_FOR_CHUNK
+    assert list(scheduler.kv_holding_waiting) == [holder]
+    assert not scheduler.waiting
+    assert adapter.requests_with_ready_chunks == {holder.request_id}
+    scheduler._restore_omni_wait_queues()
+    assert list(scheduler.kv_holding_waiting) == [holder]
+    assert list(scheduler.waiting) == [fresh]
+
+
 def _dequeue_load_entry(adapter, request):
     """Register and take the entry as the background recv loop would."""
     adapter.load_async(request)
@@ -2420,7 +2460,8 @@ def test_finish_requests_reclaims_resumable_segment_and_reuses_capacity(
     scheduler.requests = {request.request_id: request}
     scheduler.running = [request] if placement == "running" else []
     scheduler.waiting = DummyWaitingQueue([request] if placement == "waiting" else [])
-    scheduler.skipped_waiting = DummyWaitingQueue([request] if placement.startswith("skipped") else [])
+    scheduler.kv_holding_waiting = DummyWaitingQueue([request] if placement.startswith("skipped") else [])
+    scheduler.deferred_waiting = set()
     scheduler.num_waiting_for_streaming_input = int(streaming_counter_owned)
 
     freed = []
@@ -2444,7 +2485,7 @@ def test_finish_requests_reclaims_resumable_segment_and_reuses_capacity(
     assert not adapter.waiting_for_chunk_running_requests
     assert scheduler.running == []
     assert list(scheduler.waiting) == []
-    assert list(scheduler.skipped_waiting) == []
+    assert list(scheduler.kv_holding_waiting) == []
     assert scheduler.num_waiting_for_streaming_input == 0
 
     fresh_a = _req("req-fresh-a", RequestStatus.WAITING)
@@ -2467,7 +2508,8 @@ def test_finish_requests_does_not_reopen_off_queue_resumable_segment(build_adapt
     scheduler.requests = {request.request_id: request}
     scheduler.running = []
     scheduler.waiting = DummyWaitingQueue()
-    scheduler.skipped_waiting = DummyWaitingQueue()
+    scheduler.kv_holding_waiting = DummyWaitingQueue()
+    scheduler.deferred_waiting = set()
     scheduler._free_request = lambda *args, **kwargs: pytest.fail("off-queue request was freed")
 
     assert (
@@ -2541,6 +2583,7 @@ def test_generation_scheduler_calls_cleanup_on_finished(monkeypatch, mocker: Moc
 
     scheduler = mocker.MagicMock()
     scheduler._first_chunk_express = False
+    scheduler.aux_output_connector = None
     scheduler._express_min_slack_s = 0.0
     scheduler.chunk_transfer_adapter = adapter_mock
     scheduler.connector = None
@@ -2623,6 +2666,7 @@ def test_ar_scheduler_defers_cleanup_and_queues_save_on_finished(mocker: MockerF
     from vllm_omni.core.sched.omni_ar_scheduler import OmniARScheduler
 
     scheduler = mocker.MagicMock()
+    scheduler.aux_output_connector = None
     scheduler.chunk_transfer_adapter = adapter_mock
     scheduler.connector = None
     scheduler.perf_metrics = None
@@ -2720,6 +2764,8 @@ def test_omni_ar_scheduler_finish_requests(mocker: MockerFixture):
     sched.requests = {}
     sched.running = []
     sched.waiting = []
+    sched.kv_holding_waiting = []
+    sched.deferred_waiting = set()
 
     with patch.object(VLLMScheduler, "finish_requests", _super_finish):
         OmniARScheduler.finish_requests(sched, ["r1"], RequestStatus.FINISHED_ABORTED)
@@ -2735,7 +2781,8 @@ def _parked_sender_scheduler(adapter, session):
     scheduler.requests = {session.request_id: session}
     scheduler.running = []
     scheduler.waiting = DummyWaitingQueue()
-    scheduler.skipped_waiting = DummyWaitingQueue([session])
+    scheduler.kv_holding_waiting = DummyWaitingQueue([session])
+    scheduler.deferred_waiting = set()
     scheduler.num_waiting_for_streaming_input = 1
     scheduler.finished_req_ids_dict = {}
 

tests/engine/test_arg_utils_shared_fields.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 """Tests for vllm_omni.engine.arg_utils — invariants that must
 hold for the orchestrator/engine/server CLI flag partition."""
 
@@ -8,6 +8,7 @@ from __future__ import annotations
 from dataclasses import fields
 
 import pytest
+from vllm.config import AuxOutputConfig
 
 from vllm_omni.engine.arg_utils import (
     SHARED_FIELDS,
@@ -49,3 +50,24 @@ def test_internal_blacklist_keys_derived_from_orchestrator():
     """
     blacklist = internal_blacklist_keys()
     assert blacklist == orchestrator_field_names() - SHARED_FIELDS
+
+
+@pytest.mark.parametrize(
+    "aux_args",
+    [
+        {"enable_return_routed_experts": True},
+        {"aux_output_config": {"enable_return_routed_experts": True}},
+        {"aux_output_config": AuxOutputConfig(enable_return_routed_experts=True)},
+    ],
+)
+@pytest.mark.parametrize("use_v2_model_runner", [False, True])
+def test_aux_outputs_require_the_actual_omni_stage_v2_runner(aux_args, use_v2_model_runner, monkeypatch):
+    # Upstream runner selection reads this environment variable; Omni workers
+    # select their runner from the separate per-stage flag instead.
+    monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "1")
+    if use_v2_model_runner:
+        args = OmniEngineArgs(model="unused", use_v2_model_runner=True, **aux_args)
+        assert args.aux_output_config.enabled
+    else:
+        with pytest.raises(ValueError, match="use_v2_model_runner=True for this Omni stage"):
+            OmniEngineArgs(model="unused", **aux_args)

tests/engine/test_async_omni_engine_input.py

@@ -168,6 +168,37 @@ def test_build_add_request_message_preserves_additional_information(mocker: Mock
     output_processor.add_request.assert_not_called()
 
 
+def test_build_add_request_message_preserves_kv_hints_after_omni_payload_upgrade(mocker):
+    from vllm.v1.engine.input_processor import InputProcessor
+    from vllm.v1.kv_hints import KvHintsEnvelope
+
+    engine = object.__new__(AsyncOmniEngine)
+    params = SamplingParams(max_tokens=8)
+    engine.default_sampling_params_list = [params]
+    engine.stage_metadata = [StageRuntimeInfo(final_output=False, final_output_type=None, stage_type="llm")]
+    engine.supported_tasks = ("generate",)
+    processor = mocker.create_autospec(InputProcessor, instance=True)
+    hints = KvHintsEnvelope("1", "hint-message", [])
+
+    def process(**kwargs):
+        request = _make_engine_core_request()
+        request.kv_hints = kwargs["kv_hints"]
+        return request
+
+    processor.process_inputs.side_effect = process
+    engine.input_processor = processor
+    msg = engine._build_add_request_message(
+        request_id="req-1",
+        prompt={"prompt_token_ids": [1], "additional_information": {"speaker": "voice"}},
+        sampling_params_list=[params],
+        kv_hints=hints,
+    )
+    assert processor.process_inputs.call_args.kwargs["kv_hints"] is hints
+    assert isinstance(msg.prompt, OmniEngineCoreRequest)
+    assert msg.prompt.kv_hints is hints
+    assert msg.prompt.additional_information is not None
+
+
 def test_build_add_request_message_injects_global_id_before_prompt_transform(mocker: MockerFixture):
     engine = object.__new__(AsyncOmniEngine)
     params = SamplingParams(max_tokens=8)

tests/engine/test_async_omni_engine_stage_init.py

@@ -268,11 +268,21 @@ def test_stage_engine_core_client_module_reload_keeps_forward_refs_deferred():
     """Regression test for forward references in make_async_mp_client."""
     import vllm_omni.engine.stage_engine_core_client as client_mod
 
-    importlib.reload(client_mod)
+    # Snapshot the module namespace so a reload cannot strand other test
+    # modules that captured the pre-reload class objects (isinstance checks
+    # against the stale class would then fail). ``importlib.reload`` mutates
+    # the module dict in place; restoring it leaves the runtime exactly as it
+    # was before this test ran.
+    saved = dict(client_mod.__dict__)
+    try:
+        importlib.reload(client_mod)
 
-    assert client_mod.StageEngineCoreClientBase.make_async_mp_client.__annotations__["return"] == (
-        "StageEngineCoreClient | DPLBStageEngineCoreClient"
-    )
+        assert client_mod.StageEngineCoreClientBase.make_async_mp_client.__annotations__["return"] == (
+            "StageEngineCoreClient | DPLBStageEngineCoreClient"
+        )
+    finally:
+        client_mod.__dict__.clear()
+        client_mod.__dict__.update(saved)
 
 
 def test_async_omni_engine_initialize_stages_passes_log_stats_and_client_config_to_runtime(monkeypatch):

tests/engine/test_omni_request.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 """Regression tests for OmniRequest substitutability with the base Request.
 
 vllm-omni rebinds ``vllm.v1.request.Request`` to ``OmniRequest`` at import
@@ -13,13 +13,16 @@ substitutable with the base class — including base-style calls that pass
 
 import inspect
 
+import msgspec
 import numpy as np
 import pytest
 import torch
 from vllm.sampling_params import SamplingParams
+from vllm.v1.engine import EngineCoreRequest
+from vllm.v1.kv_hints import KvHintAction, KvHintsEnvelope
 from vllm.v1.request import Request
 
-from vllm_omni.engine import PromptEmbedsPayload
+from vllm_omni.engine import OmniEngineCoreRequest, PromptEmbedsPayload
 from vllm_omni.request import OmniRequest, OmniStreamingUpdate
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
@@ -111,3 +114,26 @@ def test_model_intermediate_buffer_round_trips_to_streaming_update():
     assert update is not None
     assert request.model_intermediate_buffer is buffer
     assert update.model_intermediate_buffer is buffer
+
+
+def test_kv_hints_survive_omni_upgrade_wire_transport_and_scheduler_conversion():
+    hints = KvHintsEnvelope("1", "hint-message", [KvHintAction("a", "prefetch", "1", {"blocks": [7]})])
+    core = EngineCoreRequest(
+        request_id="hint-request",
+        prompt_token_ids=[1, 2],
+        mm_features=None,
+        sampling_params=SamplingParams(max_tokens=2),
+        pooling_params=None,
+        arrival_time=0.0,
+        lora_request=None,
+        cache_salt=None,
+        data_parallel_rank=None,
+        session_id="session",
+        kv_hints=hints,
+    )
+    upgraded = OmniEngineCoreRequest.from_request(core)
+    decoded = msgspec.msgpack.decode(msgspec.msgpack.encode(upgraded), type=OmniEngineCoreRequest)
+    scheduled = OmniRequest.from_engine_core_request(decoded, block_hasher=None)
+    assert decoded.session_id == "session"
+    assert scheduled.kv_hints == hints
+    assert scheduled.session_id == "session"

tests/engine/test_omni_request_output.py

@@ -1,7 +1,8 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 """Tests for OmniRequestOutput class."""
 
+import numpy as np
 import pytest
 from PIL import Image
 from vllm.outputs import CompletionOutput, RequestOutput
@@ -157,6 +158,7 @@ class TestMsgpackRoundTrip:
         )
 
         source = _make_text_request_output(text="round trip")
+        source.prompt_token_id_logprobs = np.array([[-0.5, -1.5], [-0.2, -2.2]], dtype=np.float32)
         out = OmniRequestOutput.from_stage_output(
             source,
             request_id="rt-1",
@@ -178,6 +180,7 @@ class TestMsgpackRoundTrip:
         assert decoded.outputs[0].text == "round trip"
         assert decoded.prompt == "test prompt"
         assert decoded.prompt_token_ids == [1, 2, 3]
+        np.testing.assert_array_equal(decoded.prompt_token_id_logprobs, source.prompt_token_id_logprobs)
 
     def test_round_trip_diffusion_output(self):
         """A diffusion output round-trips without data loss."""
@@ -302,6 +305,16 @@ class TestRequestOutputFieldParity:
         assert omni.num_cached_tokens is None
         assert omni.num_cache_creation_tokens is None
 
+    def test_fixed_token_scores_survive_stage_output_wrapping(self) -> None:
+        source = _make_text_request_output()
+        source.prompt_token_id_logprobs = np.array([[-0.5, -1.5], [-0.2, -2.2]], dtype=np.float32)
+
+        omni = OmniRequestOutput.from_stage_output(source, stage_id=0)
+        forwarded = OmniRequestOutput.from_stage_output(omni, stage_id=1)
+
+        np.testing.assert_array_equal(forwarded.prompt_token_id_logprobs, source.prompt_token_id_logprobs)
+        assert OmniRequestOutput(request_id="unscored").prompt_token_id_logprobs is None
+
     def test_transfer_param_fields_are_copied_from_the_stage_output(self) -> None:
         """Declaring the fields is not enough — they must also be copied.
 

tests/engine/test_orchestrator_error_handling.py

@@ -1282,6 +1282,7 @@ async def test_cache_reset_rpc_failure_is_nonfatal_and_visits_all_replicas(metho
     [
         ("reset_mm_cache", {}, True),
         ("sleep", {}, True),
+        ("release_kv_cache_memory", {}, True),
         ("pause_scheduler", {"clear_cache": True}, True),
         ("pause_scheduler", {"clear_cache": False}, False),
         ("reset_prefix_cache", {}, False),

tests/engine/test_output_processor_mrv2_text.py

@@ -7,8 +7,9 @@ from unittest.mock import MagicMock
 
 import pytest
 import torch
-from vllm.sampling_params import RequestOutputKind
-from vllm.v1.engine import FinishReason
+from vllm.sampling_params import RequestOutputKind, SamplingParams
+from vllm.v1.engine import EngineCoreRequest, FinishReason
+from vllm.v1.engine.logprobs import LogprobsProcessor
 
 from vllm_omni.engine import OmniEngineCoreOutput
 from vllm_omni.engine.output_modality import OutputModality
@@ -16,6 +17,7 @@ from vllm_omni.engine.output_processor import (
     MultimodalOutputProcessor,
     OmniRequestState,
 )
+from vllm_omni.outputs import OmniRequestOutput
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
@@ -109,6 +111,37 @@ def test_text_tokens_are_detokenized_when_mrv2_ar_output_has_pooling_payload():
     assert not state.mm_accumulated.is_empty
 
 
+def test_fixed_token_scores_survive_output_processing_and_stage_wrapping():
+    processor, state = _make_processor(OutputModality.TEXT, _Detokenizer())
+    request = EngineCoreRequest(
+        request_id="r",
+        prompt_token_ids=[1, 2, 3],
+        mm_features=None,
+        sampling_params=SamplingParams(prompt_logprob_token_ids=[4, 5]),
+        pooling_params=None,
+        arrival_time=0.0,
+        lora_request=None,
+        cache_salt=None,
+        data_parallel_rank=None,
+    )
+    state.logprobs_processor = LogprobsProcessor.from_new_request(None, request)
+    scores = torch.tensor([[-0.5, -1.5], [-0.2, -2.2]])
+    processed = processor.process_outputs(
+        [
+            OmniEngineCoreOutput(
+                request_id="r",
+                new_token_ids=[42],
+                prompt_token_id_logprobs=scores,
+                finish_reason=FinishReason.STOP,
+            )
+        ]
+    )
+
+    (raw_output,) = processed.request_outputs
+    final_output = OmniRequestOutput.from_stage_output(raw_output, stage_id=0)
+    torch.testing.assert_close(torch.from_numpy(final_output.prompt_token_id_logprobs), scores)
+
+
 @pytest.mark.parametrize("streaming", [False, True])
 def test_generation_stage_mm_delivered_only_with_terminal(streaming):
     # Non-terminal chunks accumulate; the terminal output carries the audio.

tests/engine/test_stage_engine_args.py

@@ -97,7 +97,6 @@ _OMNI_ONLY_LLM_STAGE_ENGINE_FIELDS = frozenset(
         "enable_multithread_weight_load",
         "env",
         "has_sampling_extra_args",
-        "log_level",
         "log_stats",
         "model_arch",
         "model_subdir",

tests/engine/test_stage_engine_core_client.py

@@ -1,15 +1,23 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 """Tests for StageEngineCoreClient.check_health()."""
 
 from __future__ import annotations
 
 from types import SimpleNamespace
+from unittest.mock import patch
 
 import pytest
+from vllm import SamplingParams
+from vllm.v1.engine.core_client import AsyncMPClient, DPLBAsyncMPClient
 from vllm.v1.engine.exceptions import EngineDeadError
 
-from vllm_omni.engine.stage_engine_core_client import StageEngineCoreClient
+from vllm_omni.engine.stage_engine_core_client import (
+    DPLBStageEngineCoreClient,
+    StageEngineCoreClient,
+    StageEngineCoreClientBase,
+)
+from vllm_omni.engine.stage_init_utils import StageMetadata
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
@@ -30,3 +38,46 @@ def test_check_health_raises_when_resources_engine_dead():
     client = _make_client(engine_dead=True)
     with pytest.raises(EngineDeadError, match="engine core is dead"):
         client.check_health()
+
+
+@pytest.mark.parametrize("dp_size, external_lb", [(1, False), (2, False), (2, True)])
+@pytest.mark.parametrize("renderer", [None, object()])
+def test_stage_client_forwards_renderer_to_the_actual_upstream_constructor(dp_size, external_lb, renderer):
+    config = SimpleNamespace(
+        model_config=None,
+        parallel_config=SimpleNamespace(data_parallel_size=dp_size, data_parallel_external_lb=external_lb),
+    )
+    metadata = StageMetadata(
+        stage_id=0,
+        stage_type="llm",
+        engine_output_type="text",
+        is_comprehension=True,
+        requires_multimodal_data=False,
+        engine_input_source=[],
+        final_output=True,
+        final_output_type="text",
+        default_sampling_params=SamplingParams(),
+        custom_process_input_func=None,
+        model_stage=None,
+        runtime_cfg=None,
+    )
+    internal_lb = dp_size > 1 and not external_lb
+    parent = DPLBAsyncMPClient if internal_lb else AsyncMPClient
+    # Keep Omni construction/factory routing real; intercept only upstream's
+    # subprocess/socket boundary, using its current signature via autospec.
+    with patch.object(parent, "__init__", autospec=True, return_value=None) as init:
+        addresses = {"input_address": "tcp://192.0.2.1:1"}
+        client = StageEngineCoreClientBase.make_async_mp_client(
+            config, object, metadata=metadata, client_addresses=addresses, renderer=renderer
+        )
+        assert isinstance(client, DPLBStageEngineCoreClient if internal_lb else StageEngineCoreClient)
+        init.assert_called_once_with(
+            client,
+            config,
+            object,
+            log_stats=False,
+            client_addresses=addresses,
+            client_count=1,
+            client_index=0,
+            renderer=renderer,
+        )

tests/engine/test_stage_pool_collective_rpc.py

@@ -83,6 +83,27 @@ def test_collective_rpc_control_method_reraises_worker_error():
     asyncio.run(run())
 
 
+@pytest.mark.cpu
+def test_kv_memory_release_uses_engine_core_and_propagates_pause_errors():
+    async def run() -> None:
+        release = AsyncMock(return_value=None)
+        pool, client = _make_pool(release_kv_cache_memory_async=release)
+        await pool.collective_rpc(0, "release_kv_cache_memory", timeout=2.0)
+        release.assert_awaited_once_with()
+        client.collective_rpc_async.assert_not_awaited()
+
+        release.side_effect = RuntimeError("requires a completed pause first")
+        with pytest.raises(RuntimeError, match="completed pause"):
+            await pool.collective_rpc(0, "release_kv_cache_memory")
+
+        del client.release_kv_cache_memory_async
+        result = await pool.collective_rpc(0, "release_kv_cache_memory")
+        assert result["supported"] is False
+        client.collective_rpc_async.assert_not_awaited()
+
+    asyncio.run(run())
+
+
 @pytest.mark.cpu
 def test_collective_rpc_reset_caches_use_control_helpers():
     async def run() -> None:

tests/entrypoints/test_async_omni.py

@@ -10,6 +10,7 @@ import pytest
 from vllm.engine.protocol import StreamingInput
 from vllm.lora.request import LoRARequest
 from vllm.sampling_params import RequestOutputKind, SamplingParams
+from vllm.v1.kv_hints import KvHintsEnvelope
 
 from tests.helpers.mark import hardware_test
 from tests.helpers.stage_config import get_deploy_config_path
@@ -138,6 +139,39 @@ def test_generate_forwards_lora_request_to_engine():
     asyncio.run(run())
 
 
+@pytest.mark.cpu
+@pytest.mark.parametrize("streaming", [None, "chunk", "empty"])
+def test_generate_forwards_kv_hints_on_initial_stage_submission(streaming):
+    async def run():
+        submissions = []
+
+        async def add(**kwargs):
+            submissions.append(kwargs)
+
+        omni = get_async_omni_instance(fake_add_request=add)
+        omni.engine.add_streaming_update_async = _noop
+        omni.engine.stage_configs = [SimpleNamespace(is_comprehension=True, stage_type="llm")]
+        omni.engine.stage_vllm_configs = [SimpleNamespace(model_config=SimpleNamespace(is_encoder_decoder=False))]
+        hints = KvHintsEnvelope("1", "message", [])
+        params = SamplingParams(max_tokens=2, output_kind=RequestOutputKind.DELTA)
+
+        async def chunks():
+            if streaming == "chunk":
+                yield StreamingInput(prompt={"prompt_token_ids": [1]})
+
+        async for _ in omni.generate(
+            prompt=chunks() if streaming else {"prompt_token_ids": [1]},
+            request_id="hints",
+            sampling_params_list=[params],
+            kv_hints=hints,
+        ):
+            pass
+        assert len(submissions) == 1
+        assert submissions[0]["kv_hints"] is hints
+
+    asyncio.run(run())
+
+
 @pytest.mark.cpu
 @pytest.mark.parametrize(
     "req_ids,cancel_prefix,expected_cancel_count",

tests/entrypoints/test_async_omni_pause_sleep_routing.py

@@ -194,6 +194,56 @@ def test_reset_encoder_cache_forwards_to_ar_stages():
     asyncio.run(run())
 
 
+@pytest.mark.cpu
+def test_release_kv_memory_requires_pause_and_wakes_only_ar_kv_allocations():
+    async def run() -> None:
+        omni = _make_omni(stage_types=["llm", "diffusion", "llm"])
+        with pytest.raises(RuntimeError, match="completed pause"):
+            await omni.release_kv_cache_memory()
+        omni.collective_rpc.assert_not_awaited()
+
+        omni._paused = True
+        await omni.release_kv_cache_memory()
+        omni._clear_frontend_mm_cache.assert_awaited_once_with()
+        omni.collective_rpc.assert_awaited_once_with(
+            method="release_kv_cache_memory",
+            args=(),
+            kwargs=None,
+            stage_ids=[0, 2],
+            timeout=CACHE_RESET_TIMEOUT_S,
+        )
+        assert omni._stage_sleeping_tags == {0: {"kv_cache"}, 2: {"kv_cache"}}
+        assert not omni._level2_sleeping
+
+        omni.collective_rpc.reset_mock()
+        await omni.wake_up(stage_ids=[0, 2], tags=["kv_cache"])
+        omni.collective_rpc.assert_awaited_once_with(
+            method="wake_up", args=(), kwargs={"tags": ["kv_cache"]}, stage_ids=[0, 2]
+        )
+        assert not omni._sleeping_tags
+        assert omni._paused
+
+    asyncio.run(run())
+
+
+@pytest.mark.cpu
+def test_release_kv_memory_skips_diffusion_and_reports_engine_failures():
+    async def run() -> None:
+        diffusion = _make_omni(stage_types=["diffusion"])
+        await diffusion.release_kv_cache_memory()
+        diffusion.collective_rpc.assert_not_awaited()
+        assert not diffusion._paused
+
+        omni = _make_omni(stage_types=["llm"])
+        omni._paused = True
+        omni.collective_rpc.return_value = [{"supported": False, "error": "requires all executor memory resident"}]
+        with pytest.raises(RuntimeError, match="all executor memory resident"):
+            await omni.release_kv_cache_memory()
+        assert not omni._sleeping_tags
+
+    asyncio.run(run())
+
+
 @pytest.mark.cpu
 def test_reset_mm_cache_clears_frontend_then_forwards_to_ar_stages():
     async def run() -> None:

tests/helpers/cpu_collective_server.py

@@ -0,0 +1,19 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Import preload for the CPU-only collective forkserver.
+
+This module is imported in the forkserver process, before it starts workers.
+Its CUDA visibility does not alter the pytest process or CUDA parity workers.
+"""
+
+import os
+
+os.environ["CUDA_VISIBLE_DEVICES"] = ""
+
+import torch  # noqa: E402
+
+# Retain the expensive worker-module imports across CPU process groups.
+from tests.diffusion.distributed import test_comm, test_pipeline_parallel  # noqa: E402, F401
+
+if torch.cuda.is_initialized():
+    raise RuntimeError("CPU collective forkserver must not initialize CUDA")

tests/helpers/distributed.py

@@ -0,0 +1,62 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Process and platform setup for isolated collective test workers."""
+
+from collections.abc import Callable
+from multiprocessing.context import BaseContext
+
+import torch
+
+
+def collective_context(device_kind: str) -> BaseContext:
+    """Reuse imports in a clean CPU forkserver; CUDA workers use spawn."""
+    if device_kind == "cpu" and "forkserver" in torch.multiprocessing.get_all_start_methods():
+        torch.multiprocessing.set_forkserver_preload(["tests.helpers.cpu_collective_server"])
+        return torch.multiprocessing.get_context("forkserver")
+    return torch.multiprocessing.get_context("spawn")
+
+
+def start_collective_workers(
+    worker: Callable[..., None], *, args: tuple[object, ...], nprocs: int, device_kind: str
+) -> None:
+    context = collective_context(device_kind)
+    if context.get_start_method() == "spawn":
+        torch.multiprocessing.spawn(worker, args=args, nprocs=nprocs)
+    else:
+        torch.multiprocessing.start_processes(worker, args=args, nprocs=nprocs, start_method=context.get_start_method())
+
+
+def configure_cpu_collective_worker() -> None:
+    """Select CPU platforms inside a spawned worker that only uses Gloo.
+
+    Upstream coordinators now resolve local ranks through the selected
+    platform's visible devices. CPU ranks can outnumber visible GPUs, so they
+    must use the CPU platform even when the test host has CUDA installed.
+    The worker still creates real process groups and runs real collectives.
+    """
+    import vllm.platforms
+    from vllm.platforms.cpu import CpuPlatform
+
+    from vllm_omni.diffusion.distributed import comm, group_coordinator, parallel_state
+    from vllm_omni.platforms.interface import UnspecifiedOmniPlatform
+
+    class CpuCollectivePlatform(UnspecifiedOmniPlatform):
+        dist_backend = "gloo"
+
+        @classmethod
+        def get_device_count(cls) -> int:
+            # CPU tensors share one logical device regardless of process rank.
+            return 1
+
+        @classmethod
+        def set_device(cls, device: torch.device) -> None:
+            assert device.type == "cpu"
+
+        @classmethod
+        def synchronize(cls) -> None:
+            pass
+
+    vllm.platforms.current_platform = CpuPlatform()
+    platform = CpuCollectivePlatform()
+    for module in (comm, group_coordinator, parallel_state):
+        module.current_omni_platform = platform

tests/helpers/runtime.py

@@ -659,6 +659,14 @@ class OmniRunner:
         self.seed = seed
         self._prompt_len_estimate_cache: dict[str, Any] = {}
         self.omni: Any = None
+        self._runner_process = psutil.Process()
+        self._initial_child_ids: set[tuple[int, float]] = set()
+        self._owned_engine_processes: dict[int, float] = {}
+        for child in self._runner_process.children(recursive=True):
+            try:
+                self._initial_child_ids.add((child.pid, child.create_time()))
+            except psutil.Error:
+                pass
         try:
             from vllm_omni.entrypoints.omni import Omni
 
@@ -670,12 +678,14 @@ class OmniRunner:
                 deploy_config=deploy_config,
                 **kwargs,
             )
+            self._remember_engine_processes(startup=True)
             startup_s = time.perf_counter() - startup_t0
             if log_stats:
                 print(f"OmniRunner startup took {startup_s:.3f}s (model={model_name})", flush=True)
         except BaseException:
             # ``with OmniRunner(...)`` never reaches ``__enter__``/``__exit__``
             # when construction fails after worker processes have started.
+            self._remember_engine_processes(startup=True)
             self.__exit__(None, None, None)
             raise
 
@@ -899,18 +909,51 @@ class OmniRunner:
     def stop_profile(self, stages: list[int] | None = None) -> list[Any]:
         return self.omni.stop_profile(stages=stages)
 
+    def _remember_engine_processes(self, *, startup: bool = False) -> None:
+        """Track startup workers, then only descendants of those workers.
+
+        Existing children may belong to another runner or a shared transcriber.
+        Record identities before close() can reparent surviving workers.
+        """
+        try:
+            if startup:
+                candidates = self._runner_process.children(recursive=True)
+            else:
+                candidates = []
+                for pid, created in list(self._owned_engine_processes.items()):
+                    try:
+                        proc = psutil.Process(pid)
+                        if proc.create_time() == created:
+                            candidates.extend(proc.children(recursive=True))
+                    except psutil.Error:
+                        pass
+            for proc in candidates:
+                try:
+                    identity = (proc.pid, proc.create_time())
+                    if identity in self._initial_child_ids:
+                        continue
+                    if startup and any(
+                        (parent.pid, parent.create_time()) in self._initial_child_ids for parent in proc.parents()
+                    ):
+                        continue
+                    cmdline = " ".join(proc.cmdline()).lower()
+                    if "enginecore" in cmdline or "enginecore" in proc.name().lower():
+                        self._owned_engine_processes[identity[0]] = identity[1]
+                except psutil.Error:
+                    pass
+        except psutil.Error as error:
+            print(f"Warning: could not inspect runner workers: {error}")
+
     def _cleanup_process(self):
         try:
-            keywords = ["enginecore"]
             matched = []
-            for proc in psutil.process_iter(["pid", "name", "cmdline", "username"]):
+            for pid, created in self._owned_engine_processes.items():
                 try:
-                    cmdline = " ".join(proc.cmdline()).lower() if proc.cmdline() else ""
-                    name = proc.name().lower()
-                    if any(k in cmdline for k in keywords) or any(k in name for k in keywords):
-                        print(f"Found vllm process: PID={proc.pid}, cmd={cmdline[:100]}")
+                    proc = psutil.Process(pid)
+                    if proc.create_time() == created:
+                        print(f"Found owned vllm process: PID={proc.pid}")
                         matched.append(proc)
-                except (psutil.NoSuchProcess, psutil.AccessDenied):
+                except psutil.Error:
                     pass
             for proc in matched:
                 try:
@@ -938,11 +981,14 @@ class OmniRunner:
         return self
 
     def __exit__(self, exc_type, exc_val, exc_tb):
-        omni = getattr(self, "omni", None)
-        if omni is not None and hasattr(omni, "close"):
-            omni.close()
-        self._cleanup_process()
-        cleanup_test_environment()
+        self._remember_engine_processes()
+        try:
+            omni = getattr(self, "omni", None)
+            if omni is not None and hasattr(omni, "close"):
+                omni.close()
+        finally:
+            self._cleanup_process()
+            cleanup_test_environment()
 
 
 # ---------------------------------------------------------------------------

tests/helpers/tests/test_runtime_omni_runner.py

@@ -1,10 +1,15 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
-"""Tests for OmniRunner startup rollback."""
+"""Tests for OmniRunner startup rollback and worker ownership."""
 
 from __future__ import annotations
 
+import subprocess
+import sys
+from collections.abc import Generator
+
+import psutil
 import pytest
 
 from tests.helpers.runtime import OmniRunner
@@ -38,3 +43,114 @@ def test_runner_init_rolls_back_on_omni_startup_failure(monkeypatch: pytest.Monk
     assert cleaned.count("process") == 1
     # Once at the start of ``__init__``, once from the constructor rollback.
     assert cleaned.count("env") == 2
+
+
+@pytest.fixture
+def live_workers() -> Generator[list[subprocess.Popen[bytes]], None, None]:
+    workers: list[subprocess.Popen[bytes]] = []
+    try:
+        yield workers
+    finally:
+        for worker in workers:
+            if worker.poll() is None:
+                worker.kill()
+            worker.wait(timeout=10)
+            if worker.stdout is not None:
+                worker.stdout.close()
+
+
+def _start_engine_worker(
+    workers: list[subprocess.Popen[bytes]], *, ignore_terminate: bool = False
+) -> subprocess.Popen[bytes]:
+    script = "import signal, time; "
+    if ignore_terminate:
+        script += "signal.signal(signal.SIGTERM, signal.SIG_IGN); "
+    script += "print('ready', flush=True); time.sleep(90)"
+    worker = subprocess.Popen([sys.executable, "-c", script, "enginecore"], stdout=subprocess.PIPE)
+    workers.append(worker)
+    assert worker.stdout is not None
+    assert worker.stdout.readline() == b"ready\n"
+    return worker
+
+
+@pytest.mark.parametrize("close_raises", [False, True])
+def test_runner_cleanup_preserves_unrelated_engine_workers(
+    monkeypatch: pytest.MonkeyPatch,
+    live_workers: list[subprocess.Popen[bytes]],
+    close_raises: bool,
+) -> None:
+    existing = _start_engine_worker(live_workers)
+    owned: list[subprocess.Popen[bytes]] = []
+    late_unrelated: list[subprocess.Popen[bytes]] = []
+    monkeypatch.setattr("tests.helpers.runtime.cleanup_test_environment", lambda: None)
+
+    class _OwnedOmni:
+        def __init__(self, *args: object, **kwargs: object) -> None:
+            owned.append(_start_engine_worker(live_workers))
+
+        def close(self) -> None:
+            late_unrelated.append(_start_engine_worker(live_workers))
+            if close_raises:
+                raise RuntimeError("close failed")
+
+    monkeypatch.setattr("vllm_omni.entrypoints.omni.Omni", _OwnedOmni)
+    if close_raises:
+        with pytest.raises(RuntimeError, match="close failed"), OmniRunner("fake-model"):
+            pass
+    else:
+        with OmniRunner("fake-model"):
+            pass
+
+    assert owned[0].poll() is not None
+    assert existing.poll() is None
+    assert late_unrelated[0].poll() is None
+
+
+def test_runner_startup_failure_cleans_only_new_workers(
+    monkeypatch: pytest.MonkeyPatch, live_workers: list[subprocess.Popen[bytes]]
+) -> None:
+    unrelated = _start_engine_worker(live_workers)
+    owned: list[subprocess.Popen[bytes]] = []
+    monkeypatch.setattr("tests.helpers.runtime.cleanup_test_environment", lambda: None)
+
+    class _FailingOmni:
+        def __init__(self, *args: object, **kwargs: object) -> None:
+            owned.append(_start_engine_worker(live_workers))
+            raise RuntimeError("startup failed")
+
+    monkeypatch.setattr("vllm_omni.entrypoints.omni.Omni", _FailingOmni)
+    with pytest.raises(RuntimeError, match="startup failed"):
+        OmniRunner("fake-model")
+
+    assert owned[0].poll() is not None
+    assert unrelated.poll() is None
+
+
+def test_runner_cleanup_rejects_stale_process_identity(live_workers: list[subprocess.Popen[bytes]]) -> None:
+    unrelated = _start_engine_worker(live_workers)
+    runner = object.__new__(OmniRunner)
+    runner._owned_engine_processes = {unrelated.pid: psutil.Process(unrelated.pid).create_time() + 1}
+
+    runner._cleanup_process()
+
+    assert unrelated.poll() is None
+
+
+def test_runner_cleanup_kills_owned_worker_ignoring_terminate(
+    monkeypatch: pytest.MonkeyPatch, live_workers: list[subprocess.Popen[bytes]]
+) -> None:
+    owned: list[subprocess.Popen[bytes]] = []
+    monkeypatch.setattr("tests.helpers.runtime.cleanup_test_environment", lambda: None)
+
+    class _StubbornOmni:
+        def __init__(self, *args: object, **kwargs: object) -> None:
+            owned.append(_start_engine_worker(live_workers, ignore_terminate=True))
+
+        def close(self) -> None:
+            pass
+
+    monkeypatch.setattr("vllm_omni.entrypoints.omni.Omni", _StubbornOmni)
+    with OmniRunner("fake-model"):
+        assert owned[0].poll() is None
+
+    assert owned[0].poll() is not None

tests/inputs/test_mm_processor.py

@@ -6,6 +6,7 @@ from unittest.mock import Mock
 
 import pytest
 from transformers import BatchFeature
+from vllm.multimodal.parse import AudioProcessorItems, ImageProcessorItems, MultiModalDataItems
 
 from vllm_omni.inputs.mm_processor import OmniMultiModalProcessor
 
@@ -24,10 +25,14 @@ def test_custom_processor_bridge_preserves_prompt_kwargs_and_passthrough(counts,
 
     processor = object.__new__(Processor)
     processor.dummy_inputs = SimpleNamespace(get_dummy_text=Mock(return_value="dummy"))
-    processor._get_hf_mm_data = Mock(return_value=({"audios": ["wave"]}, {"embedding": "preserve"}))
     processor._call_hf_processor = Mock(return_value=BatchFeature({"features": "processed"}))
-    items = Mock()
-    items.get_all_counts.return_value = counts
+
+    class AudioItems(AudioProcessorItems):
+        def get_passthrough_data(self):
+            return {"embedding": "preserve"}
+
+    items = MultiModalDataItems({"audio": AudioItems(["wave"]), "image": ImageProcessorItems([])} if counts else {})
+    items.select = Mock(wraps=items.select)
     kwargs = {"sampling_rate": 16000}
     if original_text is not None:
         kwargs[processor._OMNI_PROMPT_TEXT_KEY] = original_text
@@ -36,9 +41,9 @@ def test_custom_processor_bridge_preserves_prompt_kwargs_and_passthrough(counts,
     items.select.assert_called_once_with({k for k, v in counts.items() if v > 0})
     processor._call_hf_processor.assert_called_once_with(
         "dummy" if original_text is None else original_text,
-        {"audios": ["wave"]},
+        {"audios": ["wave"]} if counts else {},
         {"sampling_rate": 16000},
         {},
     )
-    assert dict(result) == {"features": "processed", "embedding": "preserve"}
+    assert dict(result) == ({"features": "processed", "embedding": "preserve"} if counts else {"features": "processed"})
     assert kwargs == before

tests/inputs/test_preprocess.py

@@ -8,7 +8,9 @@ from unittest.mock import AsyncMock, Mock
 
 import pytest
 import torch
-from vllm.inputs import tokens_input
+from vllm.exceptions import VLLMValidationError
+from vllm.inputs import mm_input, tokens_input
+from vllm.multimodal.inputs import PlaceholderRange
 from vllm.pooling_params import PoolingParams
 from vllm.renderers import BaseRenderer
 from vllm.renderers.params import TokenizeParams
@@ -115,6 +117,35 @@ def test_omni_renderer_is_a_real_subclass_created_once():
     assert not hasattr(instance, "_renderer")
 
 
+@pytest.mark.parametrize("async_mode", [False, True])
+@pytest.mark.parametrize("no_media_kwargs", [False, True])
+@pytest.mark.parametrize("limit", [2, 4])
+def test_renderer_applies_expanded_prompt_bound_without_splitting_media(renderer, async_mode, no_media_kwargs, limit):
+    renderer._process_multimodal.return_value = mm_input(
+        list(range(6)), {}, {"image": ["hash"]}, {"image": [PlaceholderRange(offset=1, length=2)]}
+    )
+    prompt = {"prompt_token_ids": [1, 2, 3], "additional_information": {"speaker": 1}}
+    if no_media_kwargs:
+        prompt["mm_processor_kwargs"] = {"target_h": 512}
+    else:
+        prompt["multi_modal_data"] = {"image": "image"}
+    params = TokenizeParams(max_total_tokens=128, truncate_prompt_tokens=limit, truncation_side="right")
+
+    def render():
+        if async_mode:
+            return asyncio.run(renderer.render_cmpl_async([prompt], params))
+        return renderer.render_cmpl([prompt], params)
+
+    if limit == 2:
+        with pytest.raises(VLLMValidationError, match="would split or drop"):
+            render()
+    else:
+        (result,) = render()
+        assert result["prompt_token_ids"] == [0, 1, 2, 3]
+        assert result["mm_placeholders"]["image"] == [PlaceholderRange(offset=1, length=2)]
+        assert result["additional_information"] == {"speaker": 1}
+
+
 def test_concrete_renderer_overrides_are_not_suppressed():
     """The old proxy bound BaseRenderer.render_cmpl itself and hid subclass overrides."""
 

tests/metrics/test_emit_calls.py

@@ -426,6 +426,8 @@ def _make_scheduler_shell() -> OmniARScheduler:
     obj.requests = {}
     obj.running = []
     obj.waiting = []
+    obj.kv_holding_waiting = []
+    obj.deferred_waiting = set()
     obj._kv_wait_start_ts = {}
     obj._omni_kv_config = None
     return obj

tests/metrics/test_stat_logger.py

@@ -477,6 +477,7 @@ def stats_config():
         observability_config=ObservabilityConfig(kv_cache_metrics=True),
         speculative_config=Mock(spec=SpeculativeConfig, num_speculative_tokens=3),
         kv_transfer_config=None,
+        ec_transfer_config=None,
         lora_config=None,
     )
     yield config

tests/model_executor/models/omnivoice/test_processor_prompt_tokenization.py

@@ -42,10 +42,26 @@ def _make_processor(monkeypatch: pytest.MonkeyPatch) -> OmniVoiceMultiModalProce
     return processor
 
 
+class _FakeProcessorItem:
+    def __init__(self, processor_data: dict[str, object]) -> None:
+        self._processor_data = processor_data
+
+    def get_processor_data(self) -> dict[str, object]:
+        return dict(self._processor_data)
+
+    def get_passthrough_data(self) -> dict[str, object]:
+        return {}
+
+
 def _items(counts: dict[str, int], mm_data: dict[str, object]) -> SimpleNamespace:
+    class _Selected:
+        @staticmethod
+        def values():
+            return [_FakeProcessorItem(mm_data)] if mm_data else []
+
     return SimpleNamespace(
         get_all_counts=lambda: counts,
-        select=lambda _keys: mm_data,
+        select=lambda _keys: _Selected(),
     )
 
 
@@ -75,9 +91,8 @@ def test_apply_tokenizes_the_prompt_once(monkeypatch: pytest.MonkeyPatch) -> Non
 def test_media_path_encodes_audio_without_retokenizing(monkeypatch: pytest.MonkeyPatch) -> None:
     processor = _make_processor(monkeypatch)
     audio = np.zeros(2400, dtype=np.float32)
-    monkeypatch.setattr(processor, "_get_hf_mm_data", lambda _items: ({"audios": [audio]}, {}))
 
-    out = processor._apply_hf_processor_main(_items({"audio": 1}, {}), {"ref_text": "reference"})
+    out = processor._apply_hf_processor_main(_items({"audio": 1}, {"audios": [audio]}), {"ref_text": "reference"})
 
     assert tuple(out["ref_audio_tokens"].shape) == (8, 4)
     assert list(out["ref_audio_len"]) == [4]
@@ -89,7 +104,6 @@ def test_media_path_encodes_audio_without_retokenizing(monkeypatch: pytest.Monke
 def test_media_path_with_no_audio_touches_no_tokenizer(monkeypatch: pytest.MonkeyPatch) -> None:
     """Text-only requests and fully cached audio both reach this path with empty items."""
     processor = _make_processor(monkeypatch)
-    monkeypatch.setattr(processor, "_get_hf_mm_data", lambda _items: ({}, {}))
 
     out = processor._apply_hf_processor_main(_items({}, {}), {"ref_text": "reference"})
 
@@ -99,4 +113,6 @@ def test_media_path_with_no_audio_touches_no_tokenizer(monkeypatch: pytest.Monke
 
 
 def test_legacy_call_hf_processor_is_gone() -> None:
-    assert not hasattr(OmniVoiceMultiModalProcessor, "_call_hf_processor")
+    # Upstream reintroduced a 2-arg `_call_hf_processor` on the base class; the
+    # omni-specific 4-arg hook must not be defined on the omni class itself.
+    assert "_call_hf_processor" not in OmniVoiceMultiModalProcessor.__dict__

tests/model_executor/models/qwen2_5_omni/test_per_video_audio_mask.py

@@ -262,6 +262,7 @@ def test_cached_apply_hf_processor_expands_pair_miss_for_audio_hit_video_miss():
         "audio": 1,
     }[modality]
     mm_data_items.__getitem__ = MagicMock(side_effect=lambda modality: [f"{modality}-0"])
+    mm_data_items.values = MagicMock(return_value=[])
 
     inputs = MagicMock()
     inputs.mm_data_items = mm_data_items
@@ -275,6 +276,7 @@ def test_cached_apply_hf_processor_expands_pair_miss_for_audio_hit_video_miss():
     cache = MagicMock()
     # Video miss, audio hit under pair-aware keys — expansion must force audio.
     cache.is_cached.side_effect = lambda keys: [False] if keys[0].startswith("v") else [True]
+    inputs.cache = cache
 
     captured = {}
 
@@ -309,6 +311,8 @@ def test_cached_apply_hf_processor_expands_pair_miss_for_audio_hit_video_miss():
     fake_self._get_mm_fields_config = MagicMock(return_value={})
     fake_self._get_mm_prompt_updates = MagicMock(return_value={"video": [[object()]], "audio": [[object()]]})
     fake_self._merge_mm_kwargs = capture_merge
+    fake_self._postprocess_prompt = lambda prompt: prompt
+    inputs.prompt = [0, 1, 2]
 
     with patch(
         "vllm_omni.model_executor.models.qwen2_5_omni.qwen2_5_omni_thinker.MultiModalKwargsItems.from_hf_inputs",
@@ -549,24 +553,6 @@ def test_qwen2_5_prompt_updates_validate_mask_against_request_video_count():
         )
 
 
-def test_qwen2_5_mm_only_dummy_counts_subtract_only_videos_using_audio():
-    captured_counts = {}
-    fake_self = SimpleNamespace(
-        dummy_inputs=SimpleNamespace(get_dummy_text=lambda counts: captured_counts.update(counts) or "dummy"),
-        _apply_hf_processor_text_mm=lambda **_kwargs: ([], {}, False),
-    )
-    mm_items = SimpleNamespace(get_all_counts=lambda: {"video": 2, "audio": 1})
-
-    Qwen2_5OmniThinkerMultiModalProcessor._apply_hf_processor_mm_only(
-        fake_self,
-        mm_items,
-        {"use_audio_in_video": [True, False]},
-        {},
-    )
-
-    assert captured_counts == {"video": 2, "audio": 0}
-
-
 def test_get_video_use_audio_in_video_falls_back_to_prompt_updates_for_cache_hits():
     fake_self = _fake_processor()
     update_with_audio = SimpleNamespace(content=SimpleNamespace(full=[VIDEO_TOKEN_ID, AUDIO_TOKEN_ID]))
@@ -661,7 +647,12 @@ def test_qwen3_processor_inherits_vllm_omni_per_video_helpers():
         Qwen3OmniMoeThinkerMultiModalProcessor,
         Qwen2_5OmniThinkerMultiModalProcessor,
     )
-    assert Qwen3OmniMoeThinkerMultiModalProcessor._get_hf_mm_data is UpstreamQwen3Processor._get_hf_mm_data
+    # Keep Qwen3's native HF preprocessing while inheriting Omni pair handling.
+    assert Qwen3OmniMoeThinkerMultiModalProcessor._get_hf_mm_inputs is UpstreamQwen3Processor._get_hf_mm_inputs
+    assert (
+        Qwen3OmniMoeThinkerMultiModalProcessor._get_audio_in_video_pairs
+        is Qwen2_5OmniThinkerMultiModalProcessor._get_audio_in_video_pairs
+    )
     mro = Qwen3OmniMoeThinkerMultiModalProcessor.__mro__
     assert mro.index(Qwen2_5OmniThinkerMultiModalProcessor) < mro.index(UpstreamQwen3Processor)
 

tests/model_executor/models/test_omni_processing.py

@@ -79,31 +79,16 @@ def get_text_token_prompts(
     """Return ``(text_prompt, token_prompt)`` for the given multimodal data."""
     dummy_inputs = processor.dummy_inputs
     tokenizer: TokenizerLike = processor.info.get_tokenizer()
-    model_config = processor.info.ctx.model_config
 
     parsed_data = processor.info.parse_mm_data(mm_data)
     mm_counts = {k: len(vs) for k, vs in parsed_data.items()}
 
-    inputs = dummy_inputs.get_dummy_processor_inputs(
-        model_config.max_model_len,
-        mm_counts,
-        mm_options={},
+    text_prompt = dummy_inputs.get_dummy_text(mm_counts)
+    token_prompt = tokenizer.encode(
+        text_prompt,
+        **processor.info.get_default_tok_params().get_encode_kwargs(),
     )
 
-    text_prompt: str | None
-    token_prompt: list[int]
-    if isinstance(inputs.prompt, list):
-        text_prompt = None
-        token_prompt = inputs.prompt
-    elif isinstance(inputs.prompt, str):
-        text_prompt = inputs.prompt
-        token_prompt = tokenizer.encode(
-            text_prompt,
-            **processor.info.get_default_tok_params().get_encode_kwargs(),
-        )
-    else:
-        raise TypeError(type(inputs.prompt))
-
     return text_prompt, token_prompt
 
 
@@ -199,8 +184,8 @@ def _test_processing_correctness(
         modality: _to_dummy_options(modality, count) for modality, count in limit_mm_per_prompt_ints.items()
     }
 
-    baseline_processor = factories.build_processor(ctx, cache=None)
-    cached_processor = factories.build_processor(ctx, cache=cache)
+    baseline_processor = factories.build_processor(ctx)
+    cached_processor = factories.build_processor(ctx)
 
     rng = np.random.RandomState(0)
 
@@ -238,6 +223,7 @@ def _test_processing_correctness(
             baseline_processor,
             cached_processor,
             batch_idx,
+            cache,
         )
 
 
@@ -247,6 +233,7 @@ def _test_processing_correctness_one(
     baseline_processor: BaseMultiModalProcessor,
     cached_processor: BaseMultiModalProcessor,
     batch_idx: int,
+    cache: MultiModalProcessorOnlyCache,
 ):
     model_type = model_config.hf_config.model_type
 
@@ -258,12 +245,14 @@ def _test_processing_correctness_one(
         token_prompt,
         mm_items=mm_items,
         hf_processor_mm_kwargs={},
+        cache=None,
     )
 
     cached_tokenized_result = cached_processor(
         token_prompt,
         mm_items=mm_items,
         hf_processor_mm_kwargs={},
+        cache=cache,
     )
 
     _assert_inputs_equal(
@@ -278,11 +267,13 @@ def _test_processing_correctness_one(
             text_prompt,
             mm_items=mm_items,
             hf_processor_mm_kwargs={},
+            cache=None,
         )
         cached_text_result = cached_processor(
             text_prompt,
             mm_items=mm_items,
             hf_processor_mm_kwargs={},
+            cache=cache,
         )
 
         _assert_inputs_equal(

tests/model_executor/test_cfg_pairing.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 """Unit tests for the model-neutral CFG request-pairing scheduler patches.
 
 The patches keep a cond/uncond request pair step-locked inside one engine:
@@ -66,6 +66,8 @@ def _make_scheduler_cls() -> type:
     class FakeScheduler:
         def __init__(self) -> None:
             self.waiting = _FakeWaitingQueue()
+            self.kv_holding_waiting = _FakeWaitingQueue()
+            self.deferred_waiting: set[_FakeRequest] = set()
             self.running: list[_FakeRequest] = []
             self.requests: dict[str, _FakeRequest] = {}
             self.finished: list[set[str]] = []
@@ -196,6 +198,44 @@ def test_lone_member_is_held_then_released(scheduler: Any) -> None:
     assert scheduler.running == []
 
 
+def test_lone_kv_holding_member_is_held(scheduler: Any) -> None:
+    request = _FakeRequest("req-1", {"cfg_role": "cond"})
+    scheduler.add_request(request)
+    scheduler.waiting.remove_requests([request])
+    scheduler.kv_holding_waiting.append(request)
+
+    held = cfg_pairing._hold_incomplete_pairs(scheduler)
+
+    assert list(scheduler.kv_holding_waiting) == []
+    cfg_pairing._release_held(held)
+    assert list(scheduler.kv_holding_waiting) == [request]
+
+
+def test_deferred_kv_partner_holds_the_fresh_member(scheduler: Any) -> None:
+    cond, uncond = _add_pair(scheduler)
+    scheduler.waiting.remove_requests([uncond])
+    scheduler.kv_holding_waiting.append(uncond)
+    scheduler.deferred_waiting.add(uncond)
+
+    held = cfg_pairing._hold_incomplete_pairs(scheduler)
+
+    assert list(scheduler.waiting) == []
+    assert list(scheduler.kv_holding_waiting) == [uncond]
+    cfg_pairing._release_held(held)
+    assert list(scheduler.waiting) == [cond]
+
+
+def test_kv_holding_queue_keeps_pair_partners_adjacent(scheduler: Any) -> None:
+    cond, uncond = _add_pair(scheduler)
+    other = _FakeRequest("other")
+    scheduler.kv_holding_waiting.extend([cond, other, uncond])
+    scheduler.waiting.clear()
+
+    cfg_pairing._reorder_waiting_for_cfg(scheduler)
+
+    assert list(scheduler.kv_holding_waiting) == [cond, uncond, other]
+
+
 def test_complete_pair_is_admitted_together(scheduler: Any) -> None:
     scheduler.add_request(_FakeRequest("req-1", {"cfg_role": "cond"}))
     scheduler.schedule()

tests/worker/test_omni_v2_profile_dispatch.py

@@ -15,8 +15,7 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 def test_dispatch_eager_and_manager_paths(monkeypatch, dp_size):
     runner = object.__new__(OmniGPUModelRunner)
     runner.cudagraph_manager = None
-    runner.dp_size = dp_size
-    runner.dp_rank = 1
+    runner.parallel_config = SimpleNamespace(data_parallel_size=dp_size, data_parallel_rank=1)
 
     import vllm.v1.worker.gpu.dp_utils as dp_utils
 
@@ -52,7 +51,7 @@ def test_dispatch_eager_and_manager_paths(monkeypatch, dp_size):
         return expected
 
     runner.cudagraph_manager = SimpleNamespace(dispatch=dispatch)
-    runner.dp_size = 1
+    runner.parallel_config.data_parallel_size = 1
     batch_desc, _ = runner._dispatch_batch_descriptor(
         num_reqs=1, num_toks=8, uniform_tok_count=8, num_active_loras=0, use_eager=False, max_query_len=8
     )

tests/worker_v2/test_capture_query_bound_gpu.py

@@ -44,7 +44,7 @@ def _check_prefill():
     spec = FullAttentionSpec(block_size=16, num_kv_heads=8, head_size=128, dtype=torch.bfloat16)
     kvconfig = KVCacheConfig(num_blocks=1024, kv_cache_tensors=[], kv_cache_groups=[KVCacheGroupSpec(["attn"], spec)])
     state = OmniModelState.__new__(OmniModelState)
-    state.max_model_len = 256
+    state.model_config = SimpleNamespace(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_omni_ar_model_runner.py

@@ -11,10 +11,13 @@ from unittest.mock import MagicMock
 import numpy as np
 import pytest
 import torch
-from vllm.v1.outputs import RoutedExpertsTensors
+from vllm.distributed.aux_output_connector.connector import AuxRequestOutput
+from vllm.sampling_params import SamplingParams
 from vllm.v1.worker.gpu.sample.output import SamplerOutput, SamplingMaskTensors
+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.worker_v2.omni_ar_model_runner import OmniARModelRunner, OmniAsyncOutput
 from vllm_omni.worker_v2.output_snapshot import pack_output_snapshot
 
@@ -37,7 +40,7 @@ class _FakeEvent:
 def _async_output(req_ids=("req-0",), **overrides) -> OmniAsyncOutput:
     rid2idx = {rid: i for i, rid in enumerate(req_ids)}
     mro = omni_ar_model_runner.OmniModelRunnerOutput(list(req_ids), rid2idx, None, prompt_logprobs_dict={})
-    sampler_output = SamplerOutput(torch.tensor([[123]]), None, None, None)
+    sampler_output = SamplerOutput(torch.tensor([[123]]), None, None, torch.tensor([1]), torch.tensor([0]))
     kwargs = dict(model_runner_output=mro, sampler_output=sampler_output)
     kwargs.update(num_sampled_tokens=torch.tensor([1] * len(req_ids)), copy_event=_FakeEvent())
     kwargs.update(main_stream=_FakeStream(), copy_stream=_FakeStream())
@@ -60,23 +63,38 @@ def test_async_output_blocking_event_and_routing_masks(monkeypatch, compact_widt
         counts=torch.tensor([2, 0]),
         vocab_size=4,
     )
-    sampler_output = SamplerOutput(torch.tensor([[2], [0]]), None, None, torch.tensor([1, 0]), None, masks)
-    routed = RoutedExpertsTensors(torch.tensor([[[2, 3]], [[4, 5]]]), torch.tensor([7, 9]))
+    sampler_output = SamplerOutput(
+        torch.tensor([[2], [0]]), None, None, torch.tensor([1, 0]), torch.tensor([0, 0]), masks
+    )
+    copied: list[tuple[np.ndarray, np.ndarray]] = []
 
+    class _FakePendingAuxOutput:
+        def enqueue_cpu_copy(self, *, num_sampled, num_rejected) -> None:
+            copied.append((num_sampled, num_rejected))
+
+        def process_output(self) -> dict[str, AuxRequestOutput]:
+            return {
+                "decode": AuxRequestOutput(token_start=0, rows=np.array([[2, 3]], dtype=np.uint8)),
+                "prefill": AuxRequestOutput(token_start=0, rows=np.array([[4, 5]], dtype=np.uint8)),
+            }
+
+    pending = _FakePendingAuxOutput()
     output = _async_output(
         req_ids=["decode", "prefill"],
         sampler_output=sampler_output,
         num_sampled_tokens=torch.tensor([1, 0]),
         copy_event=None,  # None → constructor builds the default blocking event
-        routed_experts=routed,
+        pending_aux_output=pending,
     ).get_output()
     assert event_kwargs == [{"blocking": True}]  # blocking event by default
     assert output.sampled_token_ids == [[2], []]
     assert output.sampling_masks.to_nested_list() == [[0, 2], []]
     np.testing.assert_array_equal(output.sampling_masks.offsets, [0, 2, 2])
-    np.testing.assert_array_equal(output.routed_experts.routing_data, [[[2, 3]], [[4, 5]]])
-    np.testing.assert_array_equal(output.routed_experts.slot_mapping, [7, 9])
+    np.testing.assert_array_equal(output.aux_output_connector_output["decode"].rows, [[2, 3]])
+    np.testing.assert_array_equal(output.aux_output_connector_output["prefill"].rows, [[4, 5]])
+    assert output.aux_output_connector_output["decode"].token_start == 0
     np.testing.assert_array_equal(output.sampling_masks.token_ids, [0, 2])
+    assert len(copied) == 1 and copied[0][0].tolist() == [1, 0]
 
 
 @pytest.mark.parametrize("needs_history", [False, True])
@@ -84,19 +102,26 @@ def test_last_pp_rank_orchestration_and_kv_resolver(monkeypatch, needs_history)
     runner = OmniARModelRunner.__new__(OmniARModelRunner)
     input_batch = SimpleNamespace(req_ids=["req"], num_reqs=1, seq_lens=torch.tensor([3]))
     input_batch.idx_mapping, input_batch.query_start_loc = torch.tensor([0]), torch.tensor([0, 1])
+    input_batch.idx_mapping_np = np.array([0])
+    input_batch.num_computed_prefill_tokens_np = np.array([0])
+    input_batch.num_scheduled_tokens = np.array([3])
+    input_batch.prefill_len_np = np.array([3])
+    input_batch.query_start_loc_np = np.array([0, 3])
     state = SimpleNamespace(input_batch=input_batch, hidden_states=torch.zeros(1, 2))
-    state.finished_req_ids, state.ec_connector_output, state.routed_experts = {"finished"}, None, None
+    state.finished_req_ids, state.ec_connector_output = {"finished"}, None
     runner.execute_model_state = state
     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.aux_output_connector = None
     runner.model_config = SimpleNamespace(async_chunk=False)
     runner.vllm_config = SimpleNamespace(model_config=SimpleNamespace(engine_output_type="text"))
-    runner.model_state = SimpleNamespace(postprocess_model_output=MagicMock(return_value=(torch.zeros(1, 2), None)))
+    text_hidden = torch.tensor([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]])
+    runner.model_state = SimpleNamespace(postprocess_model_output=MagicMock(return_value=(text_hidden, None)))
     runner.model_state.intermediate_buffer = SimpleNamespace(buffers={0: {"global_request_id": "global-req"}})
     runner.req_states = SimpleNamespace(req_id_to_index={"req": 0})
     runner.req_states.all_token_ids = SimpleNamespace(gpu=torch.tensor([[1]]))
     runner.req_states.num_computed_tokens = SimpleNamespace(gpu=torch.tensor([0]))
-    runner.req_states.prompt_len = SimpleNamespace(np=np.array([1]), gpu=torch.tensor([1]))
+    runner.req_states.prompt_len = SimpleNamespace(np=np.array([3]), gpu=torch.tensor([3]))
     runner.main_stream = runner.output_copy_stream = MagicMock()
     runner.eplb = runner._finalize_native_data_plane_output = runner._reserve_native_data_plane_outputs = MagicMock()
     sampler_out = (SimpleNamespace(sampled_token_ids=torch.tensor([[2]])), MagicMock(), MagicMock())
@@ -110,13 +135,17 @@ def test_last_pp_rank_orchestration_and_kv_resolver(monkeypatch, needs_history)
         yield
         sampling_active = False
 
-    runner.model = SimpleNamespace(compute_logits=None, logitsprocs_need_output_token_ids=needs_history)
+    runner.model = SimpleNamespace(
+        compute_logits=lambda hidden: hidden, logitsprocs_need_output_token_ids=needs_history
+    )
     runner.model.mrv2_sampling_context = sampling_context
     runner.sampler = None
     runner.sample = MagicMock(return_value=sampler_out)
     runner.sample.side_effect = lambda *_: sampler_out if sampling_active is needs_history else pytest.fail()
     logprobs_mock = MagicMock(side_effect=lambda *_: pytest.fail("inside ctx") if sampling_active else {})
-    runner.prompt_logprobs_worker = SimpleNamespace(compute_prompt_logprobs=logprobs_mock)
+    runner.prompt_logprobs_worker = PromptLogprobsWorker(1, torch.device("cpu"), logprobs_mode="raw_logits")
+    runner.prompt_logprobs_worker.add_request("req", 0, SamplingParams(prompt_logprob_token_ids=[1, 0]))
+    runner.prompt_logprobs_worker.compute_prompt_logprobs = logprobs_mock
     runner.postprocess_sampled, connector_output = MagicMock(), object()
 
     def post_forward(finished_req_ids):
@@ -131,10 +160,67 @@ def test_last_pp_rank_orchestration_and_kv_resolver(monkeypatch, needs_history)
     assert runner.sample_tokens(None) is mock_out
     built = omni_ar_model_runner.OmniAsyncOutput.call_args.kwargs["model_runner_output"]
     assert built.kv_connector_output is connector_output
+    torch.testing.assert_close(built.prompt_token_id_logprobs_dict["req"], torch.tensor([[2.0, 1.0], [4.0, 3.0]]))
     assert runner._resolve_global_request_id("req") == "global-req"  # from the intermediate buffer
     assert runner._resolve_global_request_id("unknown") == "unknown"  # fallback to the local id
 
 
+def test_async_output_preserves_fixed_token_scores_on_host(monkeypatch) -> None:
+    monkeypatch.setattr(torch.cuda, "set_stream", lambda _stream: None)
+    scores = torch.tensor([[-0.5, -1.5], [-0.2, -2.2]])
+    mro = omni_ar_model_runner.OmniModelRunnerOutput(
+        req_ids=["req-0"],
+        req_id_to_index={"req-0": 0},
+        sampled_token_ids=None,
+        prompt_logprobs_dict={},
+        prompt_token_id_logprobs_dict={"req-0": scores},
+    )
+    output = _async_output(model_runner_output=mro).get_output()
+    assert output.prompt_token_id_logprobs_dict["req-0"].device.type == "cpu"
+    torch.testing.assert_close(output.prompt_token_id_logprobs_dict["req-0"], scores)
+
+
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+def test_chunked_fixed_token_scores_are_copied_from_cuda() -> None:
+    device = torch.device("cuda")
+    hidden = torch.tensor([[1.0, 2.0, 3.0], [3.0, 1.0, 2.0], [2.0, 3.0, 1.0], [4.0, 5.0, 6.0]], device=device)
+    worker = PromptLogprobsWorker(1, device)
+    worker.add_request("req", 0, SamplingParams(prompt_logprob_token_ids=[2, 0], prompt_logprob_start=1))
+    batch = SimpleNamespace(
+        req_ids=["req"],
+        idx_mapping_np=np.array([0]),
+        num_computed_prefill_tokens_np=np.array([0]),
+        num_scheduled_tokens=np.array([2]),
+        prefill_len_np=np.array([4]),
+        query_start_loc_np=np.array([0, 2]),
+    )
+    assert worker.compute_prompt_token_id_logprobs(lambda hidden: hidden, hidden[:2], batch, np.array([4])) == {}
+    batch.num_computed_prefill_tokens_np = np.array([2])
+    scores = worker.compute_prompt_token_id_logprobs(lambda hidden: hidden, hidden[2:], batch, np.array([4]))
+    mro = omni_ar_model_runner.OmniModelRunnerOutput(
+        req_ids=["req"],
+        req_id_to_index={"req": 0},
+        sampled_token_ids=None,
+        prompt_logprobs_dict={},
+        prompt_token_id_logprobs_dict=scores,
+    )
+    num_sampled = torch.tensor([1], device=device, dtype=torch.int32)
+    sampler_output = SamplerOutput(
+        torch.tensor([[1]], device=device), None, None, num_sampled, torch.zeros_like(num_sampled)
+    )
+    output = OmniAsyncOutput(
+        model_runner_output=mro,
+        sampler_output=sampler_output,
+        num_sampled_tokens=num_sampled,
+        main_stream=torch.cuda.current_stream(),
+        copy_stream=torch.cuda.Stream(),
+    ).get_output()
+
+    expected = hidden[1:3].log_softmax(dim=-1)[:, [2, 0]].cpu()
+    assert output.prompt_token_id_logprobs_dict["req"].device.type == "cpu"
+    torch.testing.assert_close(output.prompt_token_id_logprobs_dict["req"], expected)
+
+
 def test_async_mm_snapshot_owns_output_until_copy_finishes() -> None:
     runner = OmniARModelRunner.__new__(OmniARModelRunner)
     runner.model_config = SimpleNamespace(async_chunk=True)
@@ -249,7 +335,13 @@ def test_async_output_slices_request_payloads_with_graph_padding(
     batch.num_tokens_after_padding = padded_total
     output = _async_output(
         req_ids=[f"req-{i}" for i in range(len(lengths))],
-        sampler_output=SamplerOutput(torch.ones(len(lengths), 1, dtype=torch.long), None, None, None),
+        sampler_output=SamplerOutput(
+            torch.ones(len(lengths), 1, dtype=torch.long),
+            None,
+            None,
+            torch.ones(len(lengths), dtype=torch.long),
+            torch.zeros(len(lengths), dtype=torch.long),
+        ),
         text_hidden=hidden,
         multimodal_outputs={"codes": {"audio": codes, "ref": refs}},
         input_batch=batch,
@@ -325,7 +417,9 @@ def test_model_owned_audio_finalizer_runs_after_copy_without_hidden(monkeypatch,
 
     result = _async_output(
         req_ids=["audio", "partial"],
-        sampler_output=SamplerOutput(torch.tensor([[99], [100]]), None, None, None),
+        sampler_output=SamplerOutput(
+            torch.tensor([[99], [100]]), None, None, torch.tensor([1, 1]), torch.tensor([0, 0])
+        ),
         num_sampled_tokens=torch.tensor([1, 0]),
         multimodal_outputs={"snapshot": snapshot},
         input_batch=batch,
@@ -363,7 +457,9 @@ def test_request_owned_snapshot_skips_generic_partition(monkeypatch, streaming):
 
     result = _async_output(
         req_ids=["audio", "partial"],
-        sampler_output=SamplerOutput(torch.tensor([[99], [100]]), None, None, None),
+        sampler_output=SamplerOutput(
+            torch.tensor([[99], [100]]), None, None, torch.tensor([1, 1]), torch.tensor([0, 0])
+        ),
         num_sampled_tokens=torch.tensor([1, 0]),
         multimodal_outputs={"snapshot": snapshot},
         input_batch=batch,

tests/worker_v2/test_omni_generation_model_runner.py

@@ -4,12 +4,14 @@
 """OmniGenerationModelRunner contracts: empty-step lifecycle, output partition,
 CPU-sync vs CUDA-async dispatch, and async-chunk slot recycling."""
 
+from contextlib import nullcontext
 from types import SimpleNamespace
 from unittest.mock import MagicMock
 
 import numpy as np
 import pytest
 import torch
+from vllm.distributed.aux_output_connector.connector import AuxRequestOutput
 
 from vllm_omni.model_executor.models.output_templates import OmniOutput
 from vllm_omni.outputs import OmniModelRunnerOutput
@@ -45,6 +47,7 @@ def _make_runner(model_output, num_reqs=1, prompt_len=10):
     runner.execute_model_state = SimpleNamespace(finished_req_ids={"finished"}, ec_connector_output=None)
     runner.kv_connector = SimpleNamespace(post_forward=MagicMock(return_value=None))
     runner.check_ep_fault = False
+    runner.aux_output_connector = None
     runner.req_states = MagicMock(
         prompt_len=SimpleNamespace(np=np.full(num_reqs, prompt_len, dtype=np.int32)),
         num_computed_tokens=_FakeStagedField(np.zeros(num_reqs, dtype=np.int32)),
@@ -52,7 +55,8 @@ def _make_runner(model_output, num_reqs=1, prompt_len=10):
     return runner
 
 
-def test_control_only_step_keeps_lifecycle_and_skips_input_construction():
+@pytest.mark.parametrize("aux_enabled", [False, True])
+def test_control_only_step_keeps_lifecycle_and_skips_input_construction(aux_enabled):
     runner = object.__new__(OmniGenerationModelRunner)
     order = []
     for name in (
@@ -68,8 +72,13 @@ def test_control_only_step_keeps_lifecycle_and_skips_input_construction():
     runner.kv_connector = SimpleNamespace(no_forward=lambda _s: output)
     runner._merge_ec_connector_no_forward = lambda _s, value: value
     runner._attach_native_data_plane_signals = lambda value: value
+    metadata = object()
+    runner.aux_output_connector = SimpleNamespace(begin_step=MagicMock()) if aux_enabled else None
     scheduler_output = SimpleNamespace(
-        total_num_scheduled_tokens=0, scheduled_new_reqs=[], scheduled_cached_reqs=SimpleNamespace(req_ids=[])
+        total_num_scheduled_tokens=0,
+        scheduled_new_reqs=[],
+        scheduled_cached_reqs=SimpleNamespace(req_ids=[]),
+        aux_output_connector_metadata=metadata,
     )
     assert runner.execute_model(scheduler_output) is output
     assert order == [
@@ -78,6 +87,9 @@ def test_control_only_step_keeps_lifecycle_and_skips_input_construction():
         "free_states",
         "_apply_block_table_staged_writes_if_available",
     ]
+    if aux_enabled:
+        assert runner.aux_output_connector is not None
+        runner.aux_output_connector.begin_step.assert_called_once_with(metadata)
 
 
 def test_released_chunk_reuses_scheduler_output_and_slot_recycle_clears_state():
@@ -184,6 +196,72 @@ def test_sample_tokens_uses_async_output_for_cuda_and_snapshots_req_ids(monkeypa
     assert output_req_ids == ["req-0"]
 
 
+def test_generation_aux_output_copies_and_materializes_without_sampling(monkeypatch):
+    from vllm_omni.worker_v2 import omni_generation_model_runner as module
+
+    events = []
+
+    class PendingOutput:
+        def enqueue_cpu_copy(self, *, num_sampled, num_rejected):
+            np.testing.assert_array_equal(num_sampled, [0])
+            np.testing.assert_array_equal(num_rejected, [0])
+            events.append("copy")
+
+        def process_output(self):
+            assert events == ["copy", "record", "synchronize"]
+            return {"req-0": AuxRequestOutput(0, np.array([[2, 3]], dtype=np.uint8))}
+
+    class CopyEvent:
+        def record(self, stream):
+            events.append("record")
+
+        def synchronize(self):
+            events.append("synchronize")
+
+    monkeypatch.setattr(torch.cuda, "Event", lambda **kw: CopyEvent())
+    monkeypatch.setattr(torch.cuda, "stream", lambda stream: nullcontext())
+    monkeypatch.setattr(module, "_async_copy_mm", lambda *a, **kw: {})
+    output = OmniModelRunnerOutput(["req-0"], {"req-0": 0}, [[]])
+    result = module.OmniGenerationAsyncOutput(
+        model_runner_output=output,
+        multimodal_outputs={},
+        num_reqs=1,
+        main_stream=object(),
+        copy_stream=SimpleNamespace(wait_stream=lambda stream: None),
+        pending_aux_output=PendingOutput(),
+    ).get_output()
+
+    assert result.sampled_token_ids == [[]]
+    np.testing.assert_array_equal(result.aux_output_connector_output["req-0"].rows, [[2, 3]])
+
+
+def test_pending_aux_output_forces_async_delivery_without_cuda_payload(monkeypatch):
+    from vllm_omni.worker_v2 import omni_generation_model_runner as module
+
+    runner = _make_runner(OmniOutput(text_hidden_states=torch.empty(0), multimodal_outputs={}))
+    runner.device = SimpleNamespace(type="cuda")
+    runner.main_stream = object()
+    runner.output_copy_stream = object()
+    pending = object()
+    runner.aux_output_connector = SimpleNamespace(prepare_output=MagicMock(return_value=pending))
+    runner._reserve_native_data_plane_outputs = MagicMock()
+    runner._release_generation_slots = MagicMock()
+    runner._finalize_native_data_plane_output = MagicMock()
+    captured = {}
+
+    class AsyncOutput:
+        def __init__(self, **kwargs):
+            captured.update(kwargs)
+
+    monkeypatch.setattr(module, "OmniGenerationAsyncOutput", AsyncOutput)
+    input_batch = runner._gen_input_batch
+    result = runner.sample_tokens()
+
+    assert isinstance(result, AsyncOutput)
+    runner.aux_output_connector.prepare_output.assert_called_once_with(input_batch)
+    assert captured["pending_aux_output"] is pending
+
+
 @pytest.mark.parametrize("exact", [True, False])
 def test_exact_input_shape_models_reject_padded_batches(exact):
     # Qwen3-Omni Code2Wav splits input_ids by seq_token_counts (#6712).

vllm_omni/core/sched/omni_ar_scheduler.py

@@ -25,7 +25,6 @@ from vllm.v1.spec_decode.metrics import SpecDecodingStats
 from vllm_omni.core.sched.omni_scheduler_mixin import OmniSchedulerMixin
 from vllm_omni.core.sched.utils import (
     free_kv_blocks_in_physical_order,
-    omni_routed_experts_for_request,
 )
 from vllm_omni.engine import OmniEngineCoreOutput
 from vllm_omni.engine.serialization import deserialize_additional_information
@@ -99,7 +98,7 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
     core scheduling logic.
     """
 
-    max_num_running_reqs: int
+    max_num_active_reqs: int
 
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
@@ -346,32 +345,32 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
         # them. Upstream vllm raises RuntimeError on this status; omni allows
         # async abort (e.g. client disconnect during TTS streaming) to leave
         # requests in the waiting/running queues temporarily.
-        waiting = getattr(self, "waiting")
         self._drop_aborted_queued_requests()
         self._process_pending_omni_inputs(model_mode="ar")
         self._drop_aborted_queued_requests()
         self._resync_streaming_input_counter()
-        original_waiting = None
+        original_wait_queues = None
         if self._should_defer_waiting_admission():
-            original_waiting = waiting
+            original_wait_queues = self.waiting, self.kv_holding_waiting
             self.waiting = create_request_queue(self.policy)
+            self.kv_holding_waiting = create_request_queue(self.policy)
 
-        original_max_num_running_reqs = self.max_num_running_reqs
+        original_max_num_active_reqs = self.max_num_active_reqs
         async_chunk_transport = self._async_chunk_transport_enabled()
         reserved_running_slots = (
             self._get_async_chunk_reserved_running_slots() if async_chunk_transport and self.use_v2_model_runner else 0
         )
         if reserved_running_slots:
-            self.max_num_running_reqs = max(0, original_max_num_running_reqs - reserved_running_slots)
+            self.max_num_active_reqs = max(0, original_max_num_active_reqs - reserved_running_slots)
         try:
             scheduler_output = super().schedule(throttle_prefills)
         finally:
-            self.max_num_running_reqs = original_max_num_running_reqs
-            if original_waiting is not None:
-                deferred_waiting = list(self.waiting)
-                if deferred_waiting:
-                    original_waiting.prepend_requests(deferred_waiting)
-                self.waiting = original_waiting
+            self.max_num_active_reqs = original_max_num_active_reqs
+            if original_wait_queues is not None:
+                original_waiting, original_kv_holding = original_wait_queues
+                original_waiting.prepend_requests(self.waiting)
+                original_kv_holding.prepend_requests(self.kv_holding_waiting)
+                self.waiting, self.kv_holding_waiting = original_wait_queues
             self._restore_omni_wait_queues()
 
         self._postprocess_omni_schedule_output(
@@ -394,6 +393,7 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
         sampled_token_ids = model_runner_output.sampled_token_ids
         logprobs = model_runner_output.logprobs
         prompt_logprobs_dict = model_runner_output.prompt_logprobs_dict
+        prompt_token_id_logprobs_dict = getattr(model_runner_output, "prompt_token_id_logprobs_dict", {})
         num_scheduled_tokens = scheduler_output.num_scheduled_tokens
         pooler_outputs = model_runner_output.pooler_output
         mm_outputs = getattr(model_runner_output, "multimodal_outputs", None)
@@ -656,8 +656,11 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
             if stopped:
                 if self.chunk_transfer_adapter is not None:
                     confirmed_num_computed_tokens = self.chunk_transfer_adapter._confirmed_num_computed_tokens(request)
-                if model_runner_output.routed_experts is not None:
-                    routed_experts = omni_routed_experts_for_request(model_runner_output.routed_experts, request)
+                if (
+                    self.aux_output_connector is not None
+                    and (aux_output := model_runner_output.aux_output_connector_output) is not None
+                ):
+                    routed_experts = self.aux_output_connector.take_output(request, aux_output)
 
                 # Capture finish_reason BEFORE _handle_stopped_request, which may
                 # reset the status to WAITING for streaming requests that continue.
@@ -729,6 +732,7 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
 
             # Get prompt logprobs for this request.
             prompt_logprobs_tensors = prompt_logprobs_dict.get(req_id)
+            prompt_token_id_logprobs = prompt_token_id_logprobs_dict.get(req_id)
             has_stage_output = (
                 bool(new_token_ids)
                 or mm_output is not None
@@ -749,6 +753,7 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
                     finish_reason=finish_reason,
                     new_logprobs=new_logprobs,
                     new_prompt_logprobs_tensors=prompt_logprobs_tensors,
+                    prompt_token_id_logprobs=prompt_token_id_logprobs,
                     pooling_output=pooling_output_payload,
                     multimodal_output=mm_output,
                     stop_reason=request.stop_reason,
@@ -764,6 +769,7 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
             else:
                 # Invariant: EngineCore returns no partial prefill outputs.
                 assert not prompt_logprobs_tensors
+                assert prompt_token_id_logprobs is None
 
             if self.chunk_transfer_adapter is not None and (
                 inter_stage_output is not None or is_segment_finished or finished
@@ -913,8 +919,9 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
                 if session.status == RequestStatus.WAITING_FOR_STREAMING_REQ:
                     self.num_waiting_for_streaming_input -= 1
                 session.status = RequestStatus.WAITING
-                if session in self.skipped_waiting:
-                    self.skipped_waiting.remove_requests((session,))
+                self.deferred_waiting.discard(session)
+                if session in self.kv_holding_waiting:
+                    self.kv_holding_waiting.remove_requests((session,))
                     self._enqueue_waiting_request(session)
 
                 if self.log_stats:
@@ -1338,7 +1345,8 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
         ``_free_blocks``) and skips the caller's input-coordinator cleanup.
         """
         self.waiting.remove_requests((request,))
-        self.skipped_waiting.remove_requests((request,))
+        self.kv_holding_waiting.remove_requests((request,))
+        self.deferred_waiting.discard(request)
         if request.status == RequestStatus.WAITING_FOR_STREAMING_REQ:
             self.num_waiting_for_streaming_input -= 1
         request.status = RequestStatus.FINISHED_ERROR

vllm_omni/core/sched/omni_cfg_ar_scheduler.py

@@ -32,7 +32,7 @@ class OmniCFGARScheduler(OmniARScheduler):
             raise ValueError("Atomic CFG reservations require one full-attention KV group")
         self._capacity = self.kv_cache_manager.block_pool.num_gpu_blocks - 1 - self.kv_cache_manager.watermark_blocks
         pair_pages = 2 * -(-self.max_model_len // self.block_size)
-        if self._capacity < pair_pages or self.max_num_running_reqs < 2:
+        if self._capacity < pair_pages or min(self.max_num_running_reqs, self.max_num_active_reqs) < 2:
             raise ValueError("Atomic CFG requires KV capacity and max_num_seqs for two maximum-length branches")
         if self.max_num_scheduled_tokens < 2 * self.max_model_len:
             raise ValueError("Atomic CFG requires max_num_batched_tokens >= 2 * max_model_len")
@@ -81,7 +81,7 @@ class OmniCFGARScheduler(OmniARScheduler):
         admitted = create_request_queue(self.policy)
         pending = {request.request_id: request for request in original_waiting}
         pages = self._capacity - sum(self._reserved_pages(request) for request in self.running)
-        slots = self.max_num_running_reqs - len(self.running)
+        slots = self.max_num_active_reqs - len(self.running)
         tokens = self.max_num_scheduled_tokens - len(self.running)
         chosen: set[str] = set()
         for request in original_waiting:

vllm_omni/core/sched/omni_generation_scheduler.py

@@ -24,7 +24,6 @@ from vllm.v1.spec_decode.metrics import SpecDecodingStats
 
 from vllm_omni.core.sched.omni_scheduler_mixin import OmniSchedulerMixin
 from vllm_omni.core.sched.output import OmniCachedRequestData, OmniNewRequestData
-from vllm_omni.core.sched.utils import omni_routed_experts_for_request
 from vllm_omni.engine.serialization import deserialize_additional_information
 from vllm_omni.outputs import OmniModelRunnerOutput
 
@@ -136,6 +135,10 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
         if self.connector is not None:
             meta = self._build_kv_connector_meta(self.connector, scheduler_output)
             scheduler_output.kv_connector_metadata = meta
+        if self.aux_output_connector is not None:
+            scheduler_output.aux_output_connector_metadata = self.aux_output_connector.build_connector_meta(
+                scheduler_output, self.requests
+            )
         # EC Connector: package metadata
         if self.ec_connector is not None:
             ec_meta = self.ec_connector.build_connector_meta(scheduler_output)
@@ -191,7 +194,7 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
         ):
             # Downstream async-chunk stages receive the next segment from the
             # connector, not from an API StreamingUpdate. Enqueue them as
-            # schedulable before the base class can park them in skipped_waiting.
+            # schedulable before the base class can defer a streaming-input wait.
             request.status = RequestStatus.WAITING
             self._enqueue_waiting_request(request)
             return False
@@ -203,7 +206,7 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
         A generation batch slot belongs to one chunk, not the entire stream.
         Keep requests with outstanding output in ``running`` so their state
         cannot be reused before completion. Other live requests join the tail
-        of ``waiting`` without freeing their decoder state or KV blocks.
+        of the appropriate waiting queue without freeing decoder state or KV blocks.
         Preserve WAITING_FOR_CHUNK until the coordinator observes fresh input.
         Stateful codecs retain lifetime admission until their decoder state is
         released, even when no chunk is currently executing.
@@ -221,7 +224,7 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
                 continue
             if request.status == RequestStatus.RUNNING:
                 request.status = RequestStatus.WAITING
-            self.waiting.add_request(request)
+            self._enqueue_waiting_request(request)
         self.running = in_flight
 
     def schedule(self, throttle_prefills: bool = False) -> SchedulerOutput:
@@ -252,6 +255,7 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
 
         # Temporary queue: preserve waiting order while requests await input.
         skipped_waiting_requests = create_request_queue(self.policy)
+        skipped_kv_holding_requests = create_request_queue(self.policy)
         req_index = 0
         self._drop_aborted_queued_requests()
         self._requeue_completed_native_chunks()
@@ -358,7 +362,7 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
         # Fast path selection and scheduling (treat all as diffusion requests,
         # independent of pooling_params)
         while (
-            self.waiting
+            (self.kv_holding_waiting or self.waiting)
             and token_budget > 0
             and len(num_scheduled_tokens) < execution_batch_size
             and (
@@ -366,38 +370,47 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
                 if native_chunks and not self._retains_state_across_chunks
                 else len(self.running) + reserved_running_slots
             )
-            < self.max_num_running_reqs
+            < self.max_num_active_reqs
             and self._pause_state == PauseState.UNPAUSED
         ):
-            request = self.waiting.peek_request()
+            request_queue = self.kv_holding_waiting or self.waiting
+            skipped_requests = (
+                skipped_kv_holding_requests if request_queue is self.kv_holding_waiting else skipped_waiting_requests
+            )
+            request = request_queue.peek_request()
             if express and request.request_id in self._chunk_started:
-                self.waiting.pop_request()
-                skipped_waiting_requests.add_request(request)
+                request_queue.pop_request()
+                self.deferred_waiting.add(request)
+                skipped_requests.add_request(request)
                 continue
             if native_chunks and request.num_in_flight_tokens > 0:
                 # A restored waiting entry can still own an outstanding batch.
                 # Do not let it block other ready streams or execute it twice.
-                self.waiting.pop_request()
-                skipped_waiting_requests.add_request(request)
+                request_queue.pop_request()
+                self.deferred_waiting.add(request)
+                skipped_requests.add_request(request)
                 continue
             # OMNI: Skip requests that are not in self.requests
             if request.request_id not in self.requests or (
                 self.chunk_transfer_adapter is None and request.status == RequestStatus.FINISHED_STOPPED
             ):
                 # Pop the finished request from waiting queue and don't schedule it
-                self.waiting.pop_request()
+                request_queue.pop_request()
+                self.deferred_waiting.discard(request)
                 continue
 
             # async_chunk: wait for the first upstream chunk (don't start with placeholders).
             if async_chunk_transport and len(request.prompt_token_ids) == 0:
                 if self._is_done_receiving_chunks(request.request_id):
                     if not _has_async_chunk_payload_to_run(request):
-                        self.waiting.pop_request()
+                        request_queue.pop_request()
+                        self.deferred_waiting.discard(request)
                         self._pending_finish_reqs.append(request)
                         continue
                 else:
-                    self.waiting.pop_request()
-                    skipped_waiting_requests.prepend_request(request)
+                    request_queue.pop_request()
+                    self.deferred_waiting.add(request)
+                    skipped_requests.prepend_request(request)
                     continue
 
             # Allocate all input tokens for the request in one shot
@@ -413,11 +426,13 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
                     ):
                         required_tokens = 1
                     else:
-                        self.waiting.pop_request()
+                        request_queue.pop_request()
                         if self._is_done_receiving_chunks(request.request_id):
+                            self.deferred_waiting.discard(request)
                             self._pending_finish_reqs.append(request)
                         else:
-                            skipped_waiting_requests.add_request(request)
+                            self.deferred_waiting.add(request)
+                            skipped_requests.add_request(request)
                         continue
             num_new_tokens = min(required_tokens, token_budget)
             new_blocks = self.kv_cache_manager.allocate_slots(
@@ -433,7 +448,8 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
                 break
 
             # Officially schedule this request
-            request = self.waiting.pop_request()
+            request = request_queue.pop_request()
+            self.deferred_waiting.discard(request)
             self.running.append(request)
             if self.log_stats:
                 request.record_event(EngineCoreEventType.SCHEDULED, scheduled_timestamp)
@@ -448,7 +464,9 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
         if first_chunk_express:
             self._chunk_started.update(num_scheduled_tokens)
 
-        # Return skipped waiting requests
+        # Return skipped requests to their original queue.
+        if skipped_kv_holding_requests:
+            self.kv_holding_waiting.prepend_requests(skipped_kv_holding_requests)
         if skipped_waiting_requests:
             self.waiting.prepend_requests(skipped_waiting_requests)
 
@@ -792,8 +810,11 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
                     prefill_stats.finalize(self.kv_cache_manager.estimate_cached_tokens(request))
 
             if stopped:
-                if model_runner_output.routed_experts is not None:
-                    routed_experts = omni_routed_experts_for_request(model_runner_output.routed_experts, request)
+                if (
+                    self.aux_output_connector is not None
+                    and (aux_output := model_runner_output.aux_output_connector_output) is not None
+                ):
+                    routed_experts = self.aux_output_connector.take_output(request, aux_output)
                 finish_reason = request.get_finished_reason()
                 finished = self._handle_stopped_request(request)
                 is_segment_finished = not finished

vllm_omni/core/sched/omni_scheduler_mixin.py

@@ -111,6 +111,56 @@ else:
     _SchedulerMixinBase = object
 
 
+class _OmniWaitingQueue:
+    """Let one input-gating pass visit both upstream waiting queues.
+
+    Requests holding KV blocks retain upstream admission priority when input
+    gates park, restore or requeue them. Connector state is consumed once per
+    scheduler tick, including when both physical queues contain requests.
+    """
+
+    def __init__(self, scheduler: Any) -> None:
+        self.scheduler = scheduler
+
+    def __iter__(self) -> Iterator[Request]:
+        for request_queue in (self.scheduler.kv_holding_waiting, self.scheduler.waiting):
+            for request in request_queue:
+                # Upstream owns grammar, remote-KV and streaming-update waits.
+                # Omni input readiness must not overwrite those statuses.
+                if not self.scheduler._is_blocked_waiting_status(request.status):
+                    yield request
+
+    def add_request(self, request: Request) -> None:
+        self.scheduler._enqueue_waiting_request(request)
+
+    def prepend_requests(self, requests: Iterable[Request]) -> None:
+        holding: list[Request] = []
+        fresh: list[Request] = []
+        for request in requests:
+            if self.scheduler._is_blocked_waiting_status(request.status):
+                self.scheduler.deferred_waiting.add(request)
+            (holding if self.scheduler._holds_kv_blocks(request) else fresh).append(request)
+        self.scheduler.kv_holding_waiting.prepend_requests(holding)
+        self.scheduler.waiting.prepend_requests(fresh)
+
+    def remove_request(self, request: Request) -> None:
+        # Input metadata can change the computed frontier before removal, so
+        # choose the queue by membership rather than the request's new state.
+        if request in self.scheduler.kv_holding_waiting:
+            self.scheduler.kv_holding_waiting.remove_request(request)
+        else:
+            self.scheduler.waiting.remove_request(request)
+        self.scheduler.deferred_waiting.discard(request)
+
+    remove = remove_request
+
+    def remove_requests(self, requests: Iterable[Request]) -> None:
+        requests = set(requests)
+        self.scheduler.kv_holding_waiting.remove_requests(requests)
+        self.scheduler.waiting.remove_requests(requests)
+        self.scheduler.deferred_waiting.difference_update(requests)
+
+
 class OmniSchedulerMixin(_SchedulerMixinBase):
     """Shared scheduler helpers for omni-specific request handling."""
 
@@ -118,6 +168,8 @@ class OmniSchedulerMixin(_SchedulerMixinBase):
     vllm_config: VllmConfig
     requests: dict[str, Request]
     waiting: RequestQueue
+    kv_holding_waiting: RequestQueue
+    deferred_waiting: set[Request]
     running: list[Request]
     _pending_input_timeout_outputs: dict[str, tuple[int, OmniEngineCoreOutput]]
 
@@ -239,8 +291,9 @@ class OmniSchedulerMixin(_SchedulerMixinBase):
         if session.status == RequestStatus.WAITING_FOR_STREAMING_REQ:
             self.num_waiting_for_streaming_input -= 1
         session.status = RequestStatus.WAITING
-        if session in self.skipped_waiting:
-            self.skipped_waiting.remove_requests((session,))
+        self.deferred_waiting.discard(session)
+        if session in self.kv_holding_waiting:
+            self.kv_holding_waiting.remove_requests((session,))
             self._enqueue_waiting_request(session)
         if self.log_stats:
             session.record_event(EngineCoreEventType.QUEUED)
@@ -372,14 +425,14 @@ class OmniSchedulerMixin(_SchedulerMixinBase):
             )
         if input_coordinator._async_chunk:
             input_coordinator.process_pending_chunks(
-                self.waiting,
+                _OmniWaitingQueue(self),
                 self.running,
                 chunk_ready_req_ids,
                 chunk_finished_req_ids,
             )
         else:
             input_coordinator.process_pending_full_payload_inputs(
-                self.waiting,
+                _OmniWaitingQueue(self),
                 stage_recv_req_ids,
             )
 
@@ -389,7 +442,7 @@ class OmniSchedulerMixin(_SchedulerMixinBase):
         self._process_pending_input_timeouts()
         if self.chunk_transfer_adapter:
             self.chunk_transfer_adapter.process_pending_chunks(
-                self.waiting,
+                _OmniWaitingQueue(self),
                 self.running,
                 scheduler_requests=self.requests,
             )
@@ -421,12 +474,12 @@ class OmniSchedulerMixin(_SchedulerMixinBase):
         """Restore requests temporarily parked by Omni input gates."""
         if self.chunk_transfer_adapter:
             self.chunk_transfer_adapter.restore_queues(
-                self.waiting,
+                _OmniWaitingQueue(self),
                 self.running,
                 scheduler_requests=self.requests,
             )
         if self.input_coordinator:
-            self.input_coordinator.restore_queues(self.waiting, self.running)
+            self.input_coordinator.restore_queues(_OmniWaitingQueue(self), self.running)
 
     def _process_pending_input_timeouts(self) -> None:
         """Force-fail requests waiting on the full-payload coordinator too long.
@@ -664,6 +717,7 @@ class OmniSchedulerMixin(_SchedulerMixinBase):
         finish_reason: Any = None,
         new_logprobs: Any = None,
         new_prompt_logprobs_tensors: Any = None,
+        prompt_token_id_logprobs: Any = None,
         pooling_output: Any = None,
         multimodal_output: Any = None,
         stop_reason: Any = None,
@@ -688,6 +742,7 @@ class OmniSchedulerMixin(_SchedulerMixinBase):
             finish_reason=finish_reason,
             new_logprobs=new_logprobs,
             new_prompt_logprobs_tensors=new_prompt_logprobs_tensors,
+            prompt_token_id_logprobs=prompt_token_id_logprobs,
             pooling_output=pooling_output,
             pooling_output_payload=pooling_output_payload,
             multimodal_output=multimodal_output,
@@ -829,7 +884,8 @@ class OmniSchedulerMixin(_SchedulerMixinBase):
                 )
             }
             self.waiting.remove_requests(removable)
-            self.skipped_waiting.remove_requests(removable)
+            self.kv_holding_waiting.remove_requests(removable)
+            self.deferred_waiting.difference_update(removable)
 
     def _resume_downstream_chunk_receiver(self, request: Request) -> None:
         """Resume duplex connector polling without an external update."""
@@ -845,7 +901,8 @@ class OmniSchedulerMixin(_SchedulerMixinBase):
             return
         self.num_waiting_for_streaming_input -= 1
         request.status = RequestStatus.WAITING
-        self.skipped_waiting.remove_requests((request,))
+        self.deferred_waiting.discard(request)
+        self.kv_holding_waiting.remove_requests((request,))
         self._enqueue_waiting_request(request)
 
     def _aggregate_kv_connector_stats(
@@ -968,10 +1025,10 @@ class OmniSchedulerMixin(_SchedulerMixinBase):
                 request_ids = tuple(request_ids)
             target_request_ids = set(request_ids)
 
-        skipped_waiting_ids = {r.request_id for r in getattr(self, "skipped_waiting", ())}
+        deferred_waiting_ids = {r.request_id for r in self.deferred_waiting}
         pre_adapter_streaming_wait_ids = {
             rid
-            for rid in target_request_ids & skipped_waiting_ids
+            for rid in target_request_ids & deferred_waiting_ids
             if (req := self.requests.get(rid)) is not None
             and (
                 req.status == RequestStatus.WAITING_FOR_STREAMING_REQ
@@ -1033,7 +1090,7 @@ class OmniSchedulerMixin(_SchedulerMixinBase):
         status is not ``RUNNING``, set it to ``RUNNING``; symmetrically
         flip ``RUNNING → WAITING`` when the request is actually in
         ``self.waiting``. A resumable segment stop still held in
-        ``self.skipped_waiting`` is restored to
+        ``self.kv_holding_waiting`` is restored to
         ``WAITING_FOR_STREAMING_REQ`` so upstream also balances its
         paused-session counter. Because adapter cleanup may restore a
         connector-owned request's prior status first, ``finish_requests``
@@ -1070,8 +1127,8 @@ class OmniSchedulerMixin(_SchedulerMixinBase):
             return
 
         running_ids = {r.request_id for r in self.running}
-        waiting_ids = {r.request_id for r in self.waiting}
-        skipped_waiting_ids = {r.request_id for r in getattr(self, "skipped_waiting", ())}
+        waiting_ids = {r.request_id for queue in (self.waiting, self.kv_holding_waiting) for r in queue}
+        deferred_waiting_ids = {r.request_id for r in self.deferred_waiting}
 
         for rid in ids_to_align:
             req = self.requests.get(rid)
@@ -1093,7 +1150,7 @@ class OmniSchedulerMixin(_SchedulerMixinBase):
             )
             if req.is_finished() and not (resumable_segment_stop or streaming_wait):
                 continue
-            if rid in skipped_waiting_ids and streaming_wait:
+            if rid in deferred_waiting_ids and streaming_wait:
                 req.status = RequestStatus.WAITING_FOR_STREAMING_REQ
             elif rid in running_ids and req.status != RequestStatus.RUNNING:
                 req.status = RequestStatus.RUNNING
@@ -1156,10 +1213,11 @@ class OmniSchedulerMixin(_SchedulerMixinBase):
         self.running[:] = [req for req in self.running if keep_running(req)]
 
     def _drop_aborted_queued_requests(self) -> None:
-        for request_queue in (self.waiting, self.skipped_waiting):
+        for request_queue in (self.waiting, self.kv_holding_waiting):
             aborted = [req for req in request_queue if req.status == RequestStatus.FINISHED_ABORTED]
             if aborted:
                 request_queue.remove_requests(aborted)
+                self.deferred_waiting.difference_update(aborted)
         self.running[:] = [req for req in self.running if req.status != RequestStatus.FINISHED_ABORTED]
 
     def _resync_streaming_input_counter(self) -> None:
@@ -1168,7 +1226,7 @@ class OmniSchedulerMixin(_SchedulerMixinBase):
             return
         parked = sum(
             1
-            for queue in (getattr(self, "waiting", None), getattr(self, "skipped_waiting", None))
+            for queue in (getattr(self, "waiting", None), getattr(self, "kv_holding_waiting", None))
             if queue
             for request in queue
             if request.status == RequestStatus.WAITING_FOR_STREAMING_REQ

vllm_omni/core/sched/utils.py

@@ -1,30 +1,11 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 """Shared utilities for omni schedulers."""
 
-import numpy as np
 from vllm.v1.core.kv_cache_manager import KVCacheManager
-from vllm.v1.outputs import RoutedExpertsLists
 from vllm.v1.request import Request
 
 
-def omni_routed_experts_for_request(routed_experts: RoutedExpertsLists, request) -> np.ndarray | None:
-    """Extract per-request routed experts from RoutedExpertsLists using slot_mapping.
-
-    Matches upstream RoutedExpertsManager.get() pattern — filters routing_data
-    rows whose slot_mapping entries belong to this request's block_table.
-    """
-    if routed_experts is None:
-        return None
-    slots = getattr(request, "block_table", None)
-    if slots is None:
-        return None
-    slot_set = set(slots)
-    mask = np.isin(routed_experts.slot_mapping, list(slot_set))
-    data = routed_experts.routing_data[mask]
-    return data if data.size > 0 else None
-
-
 def free_kv_blocks_in_physical_order(manager: KVCacheManager, request: Request) -> None:
     """Return non-cached request blocks in increasing physical order.
 

vllm_omni/diffusion/models/boogu_image/pipeline_boogu_image.py

@@ -64,6 +64,7 @@ from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch, spli
 from vllm_omni.model_executor.model_loader.weight_utils import download_weights_from_hf_specific
 from vllm_omni.platforms import current_omni_platform
 from vllm_omni.quantization.component_config import resolve_component_quant_config
+from vllm_omni.quantization.fp8_config import DiffusionFp8Config
 
 logger = init_logger(__name__)
 
@@ -322,7 +323,7 @@ class BooguImagePipeline(CFGParallelMixin, nn.Module, ProgressBarMixin, Supports
         """Load MLLM through HF, or prepare it for deferred online FP8 loading."""
         config = None
         if quant_config is not None:
-            if not isinstance(quant_config, Fp8Config):
+            if not isinstance(quant_config, (Fp8Config, DiffusionFp8Config)):
                 raise ValueError(
                     "Boogu MLLM only supports FP8 quantization. Set mllm to null to disable online quantization."
                 )

vllm_omni/engine/__init__.py

@@ -123,6 +123,8 @@ class OmniEngineCoreRequest(EngineCoreRequest):
             trace_headers=request.trace_headers,
             resumable=request.resumable,
             external_req_id=request.external_req_id,
+            session_id=request.session_id,
+            kv_hints=request.kv_hints,
             reasoning_ended=request.reasoning_ended,
             reasoning_parser_kwargs=request.reasoning_parser_kwargs,
             abort_immediately=request.abort_immediately,

vllm_omni/engine/arg_utils.py

@@ -269,6 +269,8 @@ class OmniEngineArgs(EngineArgs):
         )
         validate_worker_omni_connector(self.worker_cls, needs_connector)
         super().__post_init__()
+        if self.aux_output_config.enabled and not self.use_v2_model_runner:
+            raise ValueError("Auxiliary outputs require use_v2_model_runner=True for this Omni stage.")
 
     def _ensure_omni_models_registered(self):
         if hasattr(self, "_omni_models_registered"):

vllm_omni/engine/async_omni_engine.py

@@ -14,6 +14,7 @@ from typing import Any, Literal, cast
 from vllm.inputs import PromptType
 from vllm.logger import init_logger
 from vllm.v1.engine import EngineCoreRequest
+from vllm.v1.kv_hints import KvHintsEnvelope
 
 from vllm_omni.data_entry_keys import REQUEST_ARTIFACT_DIRS_KEY, TRANSFORM_OWNED_META_KEYS
 from vllm_omni.engine import OmniEngineCoreRequest
@@ -259,6 +260,7 @@ class AsyncOmniEngine(OmniEngineBase):
         data_parallel_rank: int | None = None,
         reasoning_ended: bool | None = None,
         *,
+        kv_hints: KvHintsEnvelope | None = None,
         resumable: bool = False,
         message_type: Literal["add_request", "streaming_update"] = "add_request",
     ) -> StageSubmissionMessage:
@@ -289,6 +291,8 @@ class AsyncOmniEngine(OmniEngineBase):
         request_artifact_dirs: list[str] = []
 
         stage_type = self.stage_metadata[0].stage_type
+        if stage_type == "diffusion" and kv_hints is not None:
+            raise ValueError("kv_hints require an AR/LLM stage with a KV cache")
         output_prompt_text: Any = None
         _preprocess_ms = 0.0
         if stage_type != "diffusion" and not isinstance(prompt, EngineCoreRequest):
@@ -343,6 +347,7 @@ class AsyncOmniEngine(OmniEngineBase):
                     priority=priority,
                     data_parallel_rank=data_parallel_rank,
                     resumable=resumable,
+                    **({"kv_hints": kv_hints} if kv_hints is not None else {}),
                 )
             except Exception:
                 if preselected_stage0_replica is not None and self.stage_pools:
@@ -518,6 +523,7 @@ class AsyncOmniEngine(OmniEngineBase):
         data_parallel_rank: int | None = None,
         reasoning_ended: bool | None = None,
         *,
+        kv_hints: KvHintsEnvelope | None = None,
         resumable: bool = False,
     ) -> None:
         """Process stage-0 input locally, then send to the Orchestrator.
@@ -543,6 +549,7 @@ class AsyncOmniEngine(OmniEngineBase):
                 data_parallel_rank=data_parallel_rank,
                 reasoning_ended=reasoning_ended,
                 resumable=resumable,
+                **({"kv_hints": kv_hints} if kv_hints is not None else {}),
             )
         except BaseException:
             if isinstance(prompt, dict):
@@ -599,6 +606,7 @@ class AsyncOmniEngine(OmniEngineBase):
         data_parallel_rank: int | None = None,
         reasoning_ended: bool | None = None,
         *,
+        kv_hints: KvHintsEnvelope | None = None,
         resumable: bool = False,
     ) -> None:
         """Async add_request API."""
@@ -617,6 +625,7 @@ class AsyncOmniEngine(OmniEngineBase):
             data_parallel_rank=data_parallel_rank,
             reasoning_ended=reasoning_ended,
             resumable=resumable,
+            **({"kv_hints": kv_hints} if kv_hints is not None else {}),
         )
 
     def add_streaming_update(

vllm_omni/engine/orchestrator.py

@@ -718,7 +718,7 @@ class OrchestratorBase:
         stage_ids: list[int] = []
         for pool in target_pools:
             clear_sender = pool.stage_type != "diffusion" and (
-                method in ("reset_mm_cache", "sleep")
+                method in ("reset_mm_cache", "sleep", "release_kv_cache_memory")
                 or (method == "pause_scheduler" and kwargs.get("clear_cache", args[1] if len(args) > 1 else True))
             )
             for replica_id in pool.live_replica_ids():

vllm_omni/engine/stage_engine_core_client.py

@@ -34,6 +34,7 @@ from vllm_omni.engine.stage_client import StageClientBase
 from vllm_omni.engine.stage_init_utils import StageMetadata
 
 if TYPE_CHECKING:
+    from vllm.renderers import BaseRenderer
     from vllm.v1.engine import EngineCoreOutput
 
     from vllm_omni.inputs.data import OmniTokensPrompt
@@ -90,6 +91,7 @@ class StageEngineCoreClientBase(StageClientBase):
         coordinator: Any = None,
         client_count: int = 1,
         client_index: int = 0,
+        renderer: BaseRenderer | None = None,
     ) -> StageEngineCoreClient | DPLBStageEngineCoreClient:
         """Create the appropriate stage async client for the DP mode."""
         parallel_config = vllm_config.parallel_config
@@ -103,6 +105,7 @@ class StageEngineCoreClientBase(StageClientBase):
             coordinator=coordinator,
             client_count=client_count,
             client_index=client_index,
+            renderer=renderer,
         )
 
         if parallel_config.data_parallel_size > 1 and not parallel_config.data_parallel_external_lb:
@@ -118,6 +121,7 @@ class StageEngineCoreClientBase(StageClientBase):
         client_addresses: dict[str, str] | None = None,
         client_count: int = 1,
         client_index: int = 0,
+        renderer: BaseRenderer | None = None,
         *,
         metadata: StageMetadata | None = None,
         engine_manager: Any = None,
@@ -183,13 +187,16 @@ class StageEngineCoreClientBase(StageClientBase):
         )
 
         try:
-            super().__init__(
+            # Concrete clients place AsyncMPClient next in the cooperative MRO;
+            # the mixin's declared base exposes only object.__init__ to mypy.
+            super().__init__(  # type: ignore[call-arg]
                 vllm_config,
                 executor_class,
                 log_stats=log_stats,
                 client_addresses=client_addresses,
                 client_count=client_count,
                 client_index=client_index,
+                renderer=renderer,
             )
             if engine_manager is not None:
                 self.resources.engine_manager = engine_manager

vllm_omni/engine/stage_engine_core_proc.py

@@ -16,7 +16,8 @@ import signal
 from typing import Any
 
 import vllm.v1.engine.core as _vllm_engine_core_module
-from vllm.logger import init_logger
+from vllm.config import VllmConfig
+from vllm.logger import configure_logging, init_logger
 from vllm.transformers_utils.config import (
     maybe_register_config_serialize_by_value,
 )
@@ -163,6 +164,10 @@ class StageEngineCoreProc(EngineCoreProc):
             logging / metrics only.
         """
         signal_callback: SignalCallback | None = None
+        vllm_config: VllmConfig = kwargs["vllm_config"]
+        if logging_config := getattr(vllm_config, "logging_config", None):
+            configure_logging(logging_config)
+
         maybe_register_config_serialize_by_value()
 
         # Register vllm-omni reasoning parsers (e.g. step_audio) in this

vllm_omni/engine/stage_pool.py

@@ -107,6 +107,7 @@ class StagePool:
             "pause_scheduler",
             "resume_scheduler",
             "sleep",
+            "release_kv_cache_memory",
             "wake_up",
             "reset_prefix_cache",
             "reset_encoder_cache",
@@ -1334,7 +1335,7 @@ class StagePool:
                     if timeout is not None:
                         return await asyncio.wait_for(result, timeout=timeout)
                     return await result
-                if method in self._CACHE_RESET_METHODS:
+                if method in self._CACHE_RESET_METHODS or method == "release_kv_cache_memory":
                     return {"supported": False, "error": f"EngineCore helper {method}_async is unavailable"}
 
             return await client.collective_rpc_async(

vllm_omni/entrypoints/async_omni.py

@@ -41,6 +41,7 @@ if TYPE_CHECKING:
     from vllm.tokenizers import TokenizerLike
     from vllm.v1.engine import PauseMode
     from vllm.v1.engine.input_processor import InputProcessor
+    from vllm.v1.kv_hints import KvHintsEnvelope
 
     from vllm_omni.inputs.data import OmniInteractionPrompt, OmniPromptType
 
@@ -125,6 +126,7 @@ class AsyncOmni(AsyncOmniBase, EngineClient):
         priority: int = 0,
         data_parallel_rank: int | None = None,
         session_id: str | None = None,
+        kv_hints: KvHintsEnvelope | None = None,
         reasoning_ended: bool | None = None,
         reasoning_parser_kwargs: dict[str, Any] | None = None,
         arrival_time: float | None = None,
@@ -277,6 +279,7 @@ class AsyncOmni(AsyncOmniBase, EngineClient):
                     arrival_time=wall_start_ts,
                     lora_request=lora_request,
                     first_chunk_submitted=first_chunk_submitted,
+                    **({"kv_hints": kv_hints} if kv_hints is not None else {}),
                 )
                 await first_chunk_submitted
             else:
@@ -288,6 +291,7 @@ class AsyncOmni(AsyncOmniBase, EngineClient):
                     final_output_stage_ids=final_output_stage_ids,
                     arrival_time=wall_start_ts,
                     lora_request=lora_request,
+                    **({"kv_hints": kv_hints} if kv_hints is not None else {}),
                 )
             submit_ts = time.time()
             stage_first_ts = cast(list[float | None], req_state.metrics.stage_first_ts)
@@ -361,6 +365,7 @@ class AsyncOmni(AsyncOmniBase, EngineClient):
         arrival_time: float,
         lora_request: Any = None,
         first_chunk_submitted: asyncio.Future[None] | None = None,
+        kv_hints: KvHintsEnvelope | None = None,
     ) -> asyncio.Task:
         """Submit a streaming input generator as incremental stage-0 updates."""
         if not sampling_params_list:
@@ -422,6 +427,7 @@ class AsyncOmni(AsyncOmniBase, EngineClient):
                                 arrival_time=arrival_time,
                                 lora_request=lora_request,
                                 resumable=True,
+                                **({"kv_hints": kv_hints} if kv_hints is not None else {}),
                             )
                         )
                         has_submitted_first_chunk = True
@@ -479,6 +485,7 @@ class AsyncOmni(AsyncOmniBase, EngineClient):
                                     arrival_time=arrival_time,
                                     lora_request=lora_request,
                                     resumable=False,
+                                    **({"kv_hints": kv_hints} if kv_hints is not None else {}),
                                 )
                             )
                             has_submitted_first_chunk = True
@@ -903,6 +910,28 @@ class AsyncOmni(AsyncOmniBase, EngineClient):
         )
         return all(self._coerce_stage_bool(result) for result in results)
 
+    async def release_kv_cache_memory(
+        self, *, stage_ids: list[int] | None = None, timeout: float = CACHE_RESET_TIMEOUT_S
+    ) -> None:
+        """Discard AR-stage KV memory after a completed generation pause.
+
+        EngineCore checks that the scheduler is paused and all executor
+        memory is resident. Diffusion stages have no KV cache to release.
+        Restore KV memory with ``wake_up(tags=["kv_cache"])`` before resuming.
+        """
+        ar_stage_ids, _ = self._split_stage_ids_by_type(stage_ids)
+        if not ar_stage_ids:
+            return
+        async with self._pause_cond:
+            if not self._paused:
+                raise RuntimeError("release_kv_cache_memory() requires a completed pause first")
+            await self._pause_cond.wait_for(lambda: getattr(self, "_admitting", 0) == 0)
+        if 0 in ar_stage_ids:
+            await asyncio.wait_for(self._clear_frontend_mm_cache(), timeout=timeout)
+        await self._engine_core_rpc("release_kv_cache_memory", stage_ids=ar_stage_ids, timeout=timeout)
+        self._record_stage_sleep(ar_stage_ids, [CuMemTag.KV_CACHE.value])
+        self._hold_admission_until_resume = True
+
     async def sleep(
         self, stage_ids: list[int] | None = None, level: int = 2, mode: PauseMode = "abort"
     ) -> list[OmniACK]:

vllm_omni/inputs/mm_processor.py

@@ -26,14 +26,24 @@ class OmniMultiModalProcessor(BaseMultiModalProcessor[_I]):
     def _apply_hf_processor_main(
         self,
         mm_items: MultiModalDataItems,
-        hf_processor_mm_kwargs: Mapping[str, object],
+        hf_kwargs: Mapping[str, object],
     ) -> BatchFeature:
-        kwargs = dict(hf_processor_mm_kwargs)
+        kwargs = dict(hf_kwargs)
         prompt = kwargs.pop(self._OMNI_PROMPT_TEXT_KEY, None)
         if prompt is None:
             prompt = self.dummy_inputs.get_dummy_text(mm_items.get_all_counts())
         valid_items = mm_items.select({key for key, count in mm_items.get_all_counts().items() if count > 0})
-        mm_data, passthrough = self._get_hf_mm_data(valid_items)
+        # Raw processor/passthrough data (upstream removed `_get_hf_mm_data` in
+        # favour of `_get_hf_mm_inputs`, which also renames `audios` -> `audio`
+        # and injects a dummy text key). Omni's custom `_call_hf_processor`
+        # subclasses still expect the un-renamed keys, so extract them here.
+        mm_data: dict[str, object] = {}
+        passthrough: dict[str, object] = {}
+        for items in valid_items.values():
+            if not items:
+                continue
+            mm_data.update(items.get_processor_data())
+            passthrough.update(items.get_passthrough_data())
         result = self._call_hf_processor(str(prompt), dict(mm_data), kwargs, {})
         result.update(passthrough)
         return result

vllm_omni/inputs/preprocess.py

@@ -52,7 +52,7 @@ class OmniRenderer:
     def _routes_no_media_kwargs(prompt) -> bool:
         return "prompt_embeds" not in prompt and "mm_processor_kwargs" in prompt and not prompt.get("multi_modal_data")
 
-    def _process_singleton(self, prompt, *, skip_mm_cache: bool = False):
+    def _process_singleton(self, prompt, *, skip_mm_cache: bool = False, tok_params=None):
         renderer = cast(BaseRenderer, self)
         if self._routes_no_media_kwargs(prompt):
             inputs = renderer._process_multimodal(
@@ -62,11 +62,14 @@ class OmniRenderer:
                 mm_uuids=prompt.get("multi_modal_uuids"),
                 skip_mm_cache=skip_mm_cache,
             )
+            renderer._truncate_expanded_prompt(inputs, tok_params)
         else:
-            inputs = cast(BaseRenderer, super())._process_singleton(prompt, skip_mm_cache=skip_mm_cache)
+            inputs = cast(BaseRenderer, super())._process_singleton(
+                prompt, skip_mm_cache=skip_mm_cache, tok_params=tok_params
+            )
         return self._with_omni_extras(inputs, prompt)
 
-    async def _process_singleton_async(self, prompt, *, skip_mm_cache: bool = False):
+    async def _process_singleton_async(self, prompt, *, skip_mm_cache: bool = False, tok_params=None):
         renderer = cast(BaseRenderer, self)
         if self._routes_no_media_kwargs(prompt):
             inputs = await renderer._process_multimodal_async(
@@ -76,8 +79,11 @@ class OmniRenderer:
                 mm_uuids=prompt.get("multi_modal_uuids"),
                 skip_mm_cache=skip_mm_cache,
             )
+            renderer._truncate_expanded_prompt(inputs, tok_params)
         else:
-            inputs = await cast(BaseRenderer, super())._process_singleton_async(prompt, skip_mm_cache=skip_mm_cache)
+            inputs = await cast(BaseRenderer, super())._process_singleton_async(
+                prompt, skip_mm_cache=skip_mm_cache, tok_params=tok_params
+            )
         return self._with_omni_extras(inputs, prompt)
 
 

vllm_omni/model_executor/models/bagel/bagel.py

@@ -259,19 +259,25 @@ class OmniBagelMultiModalProcessor(BaseMultiModalProcessor[OmniBagelProcessingIn
     def _apply_hf_processor_main(
         self,
         mm_items: MultiModalDataItems,
-        hf_processor_mm_kwargs: Mapping[str, object],
+        hf_kwargs: Mapping[str, object],
     ) -> BatchFeature:
         valid_mm_items = mm_items.select({key for key, count in mm_items.get_all_counts().items() if count > 0})
-        mm_data, passthrough_data = self._get_hf_mm_data(valid_mm_items)
+        mm_data: dict[str, object] = {}
+        passthrough_data: dict[str, object] = {}
+        for items in valid_mm_items.values():
+            if not items:
+                continue
+            mm_data.update(items.get_processor_data())
+            passthrough_data.update(items.get_passthrough_data())
         prompt_text = self.dummy_inputs.get_dummy_text(mm_items.get_all_counts())
-        processor = self.info.get_hf_processor(**hf_processor_mm_kwargs)
+        processor = self.info.get_hf_processor(**hf_kwargs)
 
         has_image = "images" in mm_data
         has_img2img = "pixel_values_img2img" in mm_data
         processed_data = BatchFeature()
 
         if has_image:
-            image_kwargs = {**hf_processor_mm_kwargs, "is_img2img": False}
+            image_kwargs = {**hf_kwargs, "is_img2img": False}
             image_outputs = self.info.ctx.call_hf_processor(
                 processor,
                 {"text": prompt_text, "images": mm_data["images"]},
@@ -280,7 +286,7 @@ class OmniBagelMultiModalProcessor(BaseMultiModalProcessor[OmniBagelProcessingIn
             processed_data.update(image_outputs)
 
         if has_img2img:
-            img2img_kwargs = self._mm_kwargs_for_bagel_img2img_hf(hf_processor_mm_kwargs)
+            img2img_kwargs = self._mm_kwargs_for_bagel_img2img_hf(hf_kwargs)
             img2img_kwargs["is_img2img"] = True
             img2img_outputs = self.info.ctx.call_hf_processor(
                 processor,

vllm_omni/model_executor/models/common/cfg_pairing.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 #
 # Adapted from nvidia/Nemotron-Labs-Audex-2B (Apache-2.0):
 #   inference_scripts_vllm/audiogen_scripts/vllm_cfg_patch.py
@@ -89,14 +89,12 @@ def _cfg_extra(request: Any, key: str) -> Any:
 
 
 def _wait_queues(scheduler: Any) -> list[Any]:
-    # Divergence from the official (vLLM 0.20-era) file: vLLM 0.24 parks
-    # structurally not-ready requests in a second ``skipped_waiting`` queue;
-    # pair-hold must cover both queues or a lone member parked there would
-    # be admitted without its partner.
+    # Pair holding must cover requests that retain KV blocks as well as
+    # fresh requests, or a lone member can be admitted without its partner.
     queues = [scheduler.waiting]
-    skipped = getattr(scheduler, "skipped_waiting", None)
-    if skipped is not None:
-        queues.append(skipped)
+    kv_holding = getattr(scheduler, "kv_holding_waiting", None)
+    if kv_holding is not None:
+        queues.append(kv_holding)
     return queues
 
 
@@ -108,7 +106,7 @@ def _pair_complete(scheduler: Any, request_id: str, blocked_ids: set[str] | None
     if len(roles) != 2 or not all(rid in scheduler.requests for rid in roles.values()):
         return False
     if blocked_ids:
-        # A partner parked in skipped_waiting (e.g. waiting for remote KVs)
+        # A deferred partner (e.g. waiting for remote KVs)
         # may fail promotion this step while this member gets scheduled from
         # waiting; hold this member until the partner is schedulable.
         for rid in roles.values():
@@ -149,14 +147,14 @@ def _hold_incomplete_pairs(scheduler: Any) -> list[tuple[Any, Any]]:
     hold_counts: dict[str, int] = getattr(scheduler, "_cfg_hold_counts", None) or {}
     scheduler._cfg_hold_counts = hold_counts
 
-    skipped = getattr(scheduler, "skipped_waiting", None)
-    blocked_ids = {req.request_id for req in list(skipped)} if skipped is not None else set()
+    kv_holding = getattr(scheduler, "kv_holding_waiting", None)
+    blocked_ids = {req.request_id for req in getattr(scheduler, "deferred_waiting", ())}
 
     held: list[tuple[Any, Any]] = []
     for queue in _wait_queues(scheduler):
         # A member is only "blocked" for its partner's sake when it sits in
         # the OTHER queue; members of the queue being scanned move together.
-        partner_blockers = blocked_ids if queue is not skipped else set()
+        partner_blockers = blocked_ids if queue is not kv_holding else set()
         held_here = []
         for request in list(queue):
             if _pair_complete(scheduler, request.request_id, partner_blockers):
@@ -181,8 +179,12 @@ def _release_held(held: list[tuple[Any, Any]]) -> None:
 
 
 def _reorder_waiting_for_cfg(scheduler: Any) -> None:
-    """Move CFG pair partners adjacent in the FCFS waiting queue."""
-    waiting = scheduler.waiting
+    """Move CFG pair partners adjacent in each FCFS waiting queue."""
+    for waiting in _wait_queues(scheduler):
+        _reorder_cfg_queue(scheduler, waiting)
+
+
+def _reorder_cfg_queue(scheduler: Any, waiting: Any) -> None:
     # Only the FCFS queue (a deque) has caller-controlled ordering; priority
     # queues order themselves and pair sync then relies on hold + equalize.
     if not isinstance(waiting, deque) or len(waiting) < 2:

vllm_omni/model_executor/models/glm_image/glm_image_ar.py

@@ -459,16 +459,22 @@ class GlmImageMultiModalProcessor(BaseMultiModalProcessor[GlmImageProcessingInfo
     def _apply_hf_processor_main(
         self,
         mm_items: MultiModalDataItems,
-        hf_processor_mm_kwargs: Mapping[str, object],
+        hf_kwargs: Mapping[str, object],
     ) -> BatchFeature:
         """Process source images under vLLM's token-only prompt contract."""
         mm_counts = mm_items.get_all_counts()
         num_images = mm_counts.get("image", 0)
         if num_images == 0:
-            return BatchFeature({"mrope_image_grid_thw": self._build_generation_grids(hf_processor_mm_kwargs)})
+            return BatchFeature({"mrope_image_grid_thw": self._build_generation_grids(hf_kwargs)})
 
         valid_mm_items = mm_items.select({key for key, count in mm_counts.items() if count > 0})
-        processor_data, passthrough_data = self._get_hf_mm_data(valid_mm_items)
+        processor_data: dict[str, object] = {}
+        passthrough_data: dict[str, object] = {}
+        for items in valid_mm_items.values():
+            if not items:
+                continue
+            processor_data.update(items.get_processor_data())
+            passthrough_data.update(items.get_passthrough_data())
         images = processor_data.get("images")
         if not images:
             return BatchFeature(dict(passthrough_data))
@@ -483,7 +489,7 @@ class GlmImageMultiModalProcessor(BaseMultiModalProcessor[GlmImageProcessingInfo
             source_grid_thw = image_grid_thw[:num_images]
             image_inputs["image_grid_thw"] = source_grid_thw
 
-            target_grid = self._build_generation_grids(hf_processor_mm_kwargs)[:1].to(dtype=source_grid_thw.dtype)
+            target_grid = self._build_generation_grids(hf_kwargs)[:1].to(dtype=source_grid_thw.dtype)
             image_inputs["mrope_image_grid_thw"] = torch.cat(
                 [source_grid_thw, target_grid],
                 dim=0,

vllm_omni/model_executor/models/mimo_audio/mimo_audio.py

@@ -372,7 +372,7 @@ class MiMoAudioLLMMultiModalProcessor(OmniMultiModalProcessor[MiMoAudioLLMProces
 
         # Text-only input not supported in composite processor
         if not mm_data.get("audio", []):
-            prompt_ids = self._apply_hf_processor_tokens_only(prompt_ids)
+            prompt_ids = self._postprocess_prompt(prompt_ids)
 
             return BatchFeature(dict(input_ids=[prompt_ids]), tensor_type="pt")
 

vllm_omni/model_executor/models/mimo_audio/mimo_audio_llm.py

@@ -396,25 +396,21 @@ class MimoAudioDummyInputsBuilder(BaseDummyInputsBuilder[MimoAudioProcessingInfo
 
 
 class MimoAudioMultiModalProcessor(BaseMultiModalProcessor[MimoAudioProcessingInfo]):
-    def _preprocess_hf_mm_data(
+    def _get_hf_mm_inputs(
         self,
-        mm_data: Mapping[str, object],
-        hf_processor_mm_kwargs: Mapping[str, object],
-    ) -> tuple[Mapping[str, object], Mapping[str, object]]:
-        mm_data, hf_processor_mm_kwargs = super()._preprocess_hf_mm_data(
-            mm_data,
-            hf_processor_mm_kwargs,
-        )
-        mm_data = dict(mm_data)
-        audios = mm_data.pop("audios", [])
-        if audios:
-            mm_data["audio"] = audios
-            feature_extractor = self.info.get_feature_extractor(**hf_processor_mm_kwargs)
-            hf_processor_mm_kwargs = {
-                **hf_processor_mm_kwargs,
-                "sampling_rate": feature_extractor.sampling_rate,
-            }
-        return mm_data, hf_processor_mm_kwargs
+        mm_items: MultiModalDataItems,
+        hf_kwargs: Mapping[str, object],
+    ):
+        hf_inputs = super()._get_hf_mm_inputs(mm_items, hf_kwargs)
+        if hf_inputs.hf_data.get("audio"):
+            feature_extractor = self.info.get_feature_extractor(**hf_kwargs)
+            hf_inputs = hf_inputs._replace(
+                hf_kwargs={
+                    **hf_kwargs,
+                    "sampling_rate": feature_extractor.sampling_rate,
+                },
+            )
+        return hf_inputs
 
     def _get_mm_fields_config(
         self,

vllm_omni/model_executor/models/ming_flash_omni/ming_flash_omni_thinker.py

@@ -516,7 +516,7 @@ class MingFlashOmniThinkerMultiModalProcessor(BaseMultiModalProcessor[MingFlashO
     def _apply_hf_processor_main(
         self,
         mm_items: MultiModalDataItems,
-        hf_processor_mm_kwargs: Mapping[str, object],
+        hf_kwargs: Mapping[str, object],
     ) -> BatchFeature:
         """Call Ming's modality processors without expanding prompt text.
 
@@ -525,7 +525,13 @@ class MingFlashOmniThinkerMultiModalProcessor(BaseMultiModalProcessor[MingFlashO
         placeholder tokens remain under vLLM's prompt-update control.
         """
         valid_mm_items = mm_items.select({key for key, count in mm_items.get_all_counts().items() if count > 0})
-        mm_data, passthrough_data = self._get_hf_mm_data(valid_mm_items)
+        mm_data: dict[str, object] = {}
+        passthrough_data: dict[str, object] = {}
+        for items in valid_mm_items.values():
+            if not items:
+                continue
+            mm_data.update(items.get_processor_data())
+            passthrough_data.update(items.get_passthrough_data())
         hf_processor = self.info.get_hf_processor()
 
         data: dict[str, object] = {}

vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_llm.py

@@ -3473,19 +3473,25 @@ class MiniCPMO45OmniLLMMultiModalProcessor(BaseMultiModalProcessor[MiniCPMO45Omn
     def _apply_hf_processor_main(
         self,
         mm_items: MultiModalDataItems,
-        hf_processor_mm_kwargs: Mapping[str, object],
+        hf_kwargs: Mapping[str, object],
     ) -> BatchFeature:
         """
         Process each modality independently because the MiniCPM processor
         asserts that image tags and image sizes have matching lengths.
         """
         valid_mm_items = mm_items.select({key for key, count in mm_items.get_all_counts().items() if count > 0})
-        mm_data, passthrough_data = self._get_hf_mm_data(valid_mm_items)
+        mm_data: dict[str, object] = {}
+        passthrough_data: dict[str, object] = {}
+        for items in valid_mm_items.values():
+            if not items:
+                continue
+            mm_data.update(items.get_processor_data())
+            passthrough_data.update(items.get_passthrough_data())
 
         tokenizer = self.info.get_tokenizer()
         prompt_text = self.dummy_inputs.get_dummy_text(mm_items.get_all_counts())
         input_ids = torch.tensor([tokenizer.encode(prompt_text)])
-        mm_inputs = self.process_mm_inputs(mm_data, hf_processor_mm_kwargs)
+        mm_inputs = self.process_mm_inputs(mm_data, hf_kwargs)
         processed_data = BatchFeature(
             {
                 "input_ids": input_ids,

vllm_omni/model_executor/models/minimax_h3/text_encoder.py

@@ -94,9 +94,9 @@ class MiniMaxH3MultiModalProcessor(Qwen3VLMultiModalProcessor):
     def _apply_hf_processor_main(
         self,
         mm_items: Any,
-        hf_processor_mm_kwargs: Mapping[str, object],
+        hf_kwargs: Mapping[str, object],
     ):
-        base_kwargs = self._base_processor_kwargs(hf_processor_mm_kwargs)
+        base_kwargs = self._base_processor_kwargs(hf_kwargs)
         return super()._apply_hf_processor_main(
             mm_items,
             base_kwargs,

vllm_omni/model_executor/models/omnivoice/omnivoice.py

@@ -117,7 +117,7 @@ class OmniVoiceMultiModalProcessor(OmniMultiModalProcessor[OmniVoiceMultiModalPr
     def _apply_hf_processor_main(
         self,
         mm_items: MultiModalDataItems,
-        hf_processor_mm_kwargs: Mapping[str, object],
+        hf_kwargs: Mapping[str, object],
     ) -> BatchFeature:
         """Encode the reference audio, if any; the prompt was tokenized in apply().
 
@@ -126,7 +126,13 @@ class OmniVoiceMultiModalProcessor(OmniMultiModalProcessor[OmniVoiceMultiModalPr
         touch the text tokenizer.
         """
         valid_items = mm_items.select({key for key, count in mm_items.get_all_counts().items() if count > 0})
-        mm_data, passthrough = self._get_hf_mm_data(valid_items)
+        mm_data: dict[str, object] = {}
+        passthrough: dict[str, object] = {}
+        for items in valid_items.values():
+            if not items:
+                continue
+            mm_data.update(items.get_processor_data())
+            passthrough.update(items.get_passthrough_data())
         audio = self._get_reference_audio(mm_data)
         if audio is None:
             return BatchFeature(dict(passthrough))

vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_thinker.py

@@ -78,7 +78,7 @@ from vllm.multimodal.parse import AudioProcessorItems, MultiModalDataItems, Vide
 from vllm.multimodal.processing.context import TimingContext
 from vllm.multimodal.processing.inputs import ProcessorInputs
 from vllm.multimodal.processing.processor import (
-    MultiModalProcessingInfo,
+    MultiModalProcessingResult,
     MultiModalPromptUpdates,
     PlaceholderFeaturesInfo,
     PromptReplacement,
@@ -367,7 +367,7 @@ class Qwen2_5OmniThinkerMultiModalProcessor(
     def _apply_hf_processor_main(
         self,
         mm_items: MultiModalDataItems,
-        hf_processor_mm_kwargs: Mapping[str, object],
+        hf_kwargs: Mapping[str, object],
     ) -> BatchFeature:
         """Apply omni-specific HF kwargs before the upstream Qwen2.5-Omni
         processor runs.
@@ -378,23 +378,30 @@ class Qwen2_5OmniThinkerMultiModalProcessor(
         per-video ``use_audio_in_video`` mask handling have to run here.
         """
         valid_mm_items = mm_items.select({k for k, c in mm_items.get_all_counts().items() if c > 0})
-        processor_data, _ = self._get_hf_mm_data(valid_mm_items)
+        # Upstream removed `_get_hf_mm_data`; extract the raw processor data
+        # (still keyed as `videos`/`video_metadata`/`audios`) for the omni
+        # kwargs helpers below, instead of `_get_hf_mm_inputs` which renames
+        # and injects a dummy text key.
+        processor_data: dict[str, object] = {}
+        for items in valid_mm_items.values():
+            if items:
+                processor_data.update(items.get_processor_data())
 
-        hf_processor_mm_kwargs = _presampled_videos_hf_kwargs(
+        hf_kwargs = _presampled_videos_hf_kwargs(
             processor_data,
-            hf_processor_mm_kwargs,
+            hf_kwargs,
         )
-        hf_processor_mm_kwargs = _coerce_use_audio_in_video_for_hf_processor(
+        hf_kwargs = _coerce_use_audio_in_video_for_hf_processor(
             processor_data,
-            hf_processor_mm_kwargs,
+            hf_kwargs,
         )
 
         processed_data = super()._apply_hf_processor_main(
             mm_items,
-            hf_processor_mm_kwargs,
+            hf_kwargs,
         )
 
-        per_video_mask = hf_processor_mm_kwargs.get(_PER_VIDEO_USE_AUDIO_IN_VIDEO_KEY)
+        per_video_mask = hf_kwargs.get(_PER_VIDEO_USE_AUDIO_IN_VIDEO_KEY)
         if per_video_mask is not None:
             processed_data["use_audio_in_video"] = torch.tensor(per_video_mask)
         return processed_data
@@ -564,32 +571,6 @@ class Qwen2_5OmniThinkerMultiModalProcessor(
             ),
         ]
 
-    def _apply_hf_processor_mm_only(
-        self,
-        mm_items: MultiModalDataItems,
-        hf_processor_mm_kwargs: Mapping[str, object],
-        tokenization_kwargs: Mapping[str, object],
-    ):
-        mm_counts = mm_items.get_all_counts()
-
-        if "video" in mm_counts:
-            video_use_audio_in_video = _get_request_video_use_audio_in_video(
-                hf_processor_mm_kwargs,
-                mm_counts["video"],
-            )
-            if any(video_use_audio_in_video):
-                assert "audio" in mm_counts
-                mm_counts["audio"] -= sum(video_use_audio_in_video)
-
-        _, mm_processed_data, _ = self._apply_hf_processor_text_mm(
-            prompt_text=self.dummy_inputs.get_dummy_text(mm_counts),
-            mm_items=mm_items,
-            hf_processor_mm_kwargs=hf_processor_mm_kwargs,
-            tokenization_kwargs=tokenization_kwargs,
-        )
-
-        return mm_processed_data
-
     def _get_audio_in_video_pairs(
         self,
         mm_data_items: MultiModalDataItems,
@@ -639,14 +620,14 @@ class Qwen2_5OmniThinkerMultiModalProcessor(
         self,
         inputs: ProcessorInputs,
         timing_ctx: TimingContext,
-    ) -> MultiModalProcessingInfo:
+    ) -> MultiModalProcessingResult:
         """Cache-aware: processes video/audio pairs as units,
         using pair-specific cache keys. Falls back to standard cache if no pairs."""
 
-        cache = self.cache
+        cache = inputs.cache
 
-        _, passthrough_data = self._get_hf_mm_data(inputs.mm_data_items)
-        if cache is None or passthrough_data:
+        has_passthrough_data = any(len(items.get_passthrough_data()) > 0 for items in inputs.mm_data_items.values())
+        if cache is None or has_passthrough_data:
             return self._apply_hf_processor(inputs, timing_ctx)
 
         aiv_pairs = self._get_audio_in_video_pairs(
@@ -694,7 +675,7 @@ class Qwen2_5OmniThinkerMultiModalProcessor(
         with timing_ctx.record("apply_hf_processor"):
             mm_missing_processed_data = self._apply_hf_processor_main(
                 mm_items=mm_missing_data_items,
-                hf_processor_mm_kwargs=hf_processor_mm_kwargs,
+                hf_kwargs=hf_processor_mm_kwargs,
             )
 
         mm_missing_kwargs = MultiModalKwargsItems.from_hf_inputs(
@@ -720,7 +701,8 @@ class Qwen2_5OmniThinkerMultiModalProcessor(
                 mm_missing_prompt_updates=mm_missing_prompt_updates,
             )
 
-        mm_info = MultiModalProcessingInfo(
+        mm_info = MultiModalProcessingResult(
+            prompt_ids=self._postprocess_prompt(inputs.prompt),
             kwargs=mm_kwargs,
             hashes=mm_hashes,
             prompt_updates=mm_prompt_updates,
@@ -819,11 +801,12 @@ class Qwen2_5OmniThinkerMultiModalProcessor(
     def _maybe_apply_prompt_updates(
         self,
         mm_items: MultiModalDataItems,
-        prompt_ids: list[int],
-        mm_kwargs: MultiModalKwargsItems,
-        mm_prompt_updates: MultiModalPromptUpdates,
+        mm_res: MultiModalProcessingResult,
     ) -> tuple[list[int], Mapping[str, list[PlaceholderFeaturesInfo]]]:
         mm_item_counts = mm_items.get_all_counts()
+        prompt_ids = mm_res.prompt_ids
+        mm_kwargs = mm_res.kwargs
+        mm_prompt_updates = mm_res.prompt_updates
         self._validate_mm_kwargs(mm_kwargs, mm_item_counts)
         self._validate_mm_updates(mm_prompt_updates, mm_item_counts)
 

vllm_omni/model_executor/models/step_audio2/step_audio2_thinker.py

@@ -561,35 +561,15 @@ class StepAudio2MultiModalProcessor(BaseMultiModalProcessor[StepAudio2Processing
             )
         ]
 
-    def _preprocess_hf_mm_data(
-        self,
-        mm_data: Mapping[str, object],
-        hf_processor_mm_kwargs: Mapping[str, object],
-    ) -> tuple[Mapping[str, object], Mapping[str, object]]:
-        mm_data, hf_processor_mm_kwargs = super()._preprocess_hf_mm_data(
-            mm_data,
-            hf_processor_mm_kwargs,
-        )
-        mm_data = dict(mm_data)
-        audios = mm_data.pop("audios", [])
-        if audios:
-            if isinstance(audios, str):
-                mm_data["audio"] = [audios]
-            elif isinstance(audios, (list, tuple)):
-                mm_data["audio"] = audios
-            else:
-                mm_data["audio"] = [audios]
-        return mm_data, hf_processor_mm_kwargs
-
     def _postprocess_hf_mm_data(
         self,
-        mm_data: Mapping[str, object],
-        hf_processor_mm_kwargs: Mapping[str, object],
+        hf_data: Mapping[str, object],
+        hf_kwargs: Mapping[str, object],
         processed_data: BatchFeature,
     ) -> BatchFeature:
         processed_data = super()._postprocess_hf_mm_data(
-            mm_data,
-            hf_processor_mm_kwargs,
+            hf_data,
+            hf_kwargs,
             processed_data,
         )
         if "audio_mels" not in processed_data:

vllm_omni/model_executor/models/voxcpm2/scheduler.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 from __future__ import annotations
 
@@ -30,7 +30,7 @@ class VoxCPM2OmniARAsyncScheduler(OmniARAsyncScheduler):
     def _should_defer_waiting_for_unified_decode_graph(self) -> bool:
         if not self._unified_decode_graph_enabled():
             return False
-        if not self.waiting or not self.running:
+        if not (self.waiting or self.kv_holding_waiting) or not self.running:
             return False
 
         for request in self.running:

vllm_omni/model_executor/models/voxtral_tts/voxtral_tts_audio_generation.py

@@ -48,6 +48,8 @@ from vllm.multimodal.parse import AudioProcessorItems, MultiModalDataItems, Mult
 from vllm.multimodal.processing import BaseDummyInputsBuilder, BaseMultiModalProcessor
 from vllm.multimodal.processing.processor import (
     BaseProcessingInfo,
+    MultiModalProcessingResult,
+    PlaceholderFeaturesInfo,
     ProcessorInputs,
     PromptReplacement,
     PromptUpdate,
@@ -874,39 +876,17 @@ class VoxtralTTSMultiModalProcessor(BaseMultiModalProcessor[VoxtralTTSProcessing
             ),
         ]
 
-    def _apply_hf_processor_mm_only(
-        self,
-        mm_items: MultiModalDataItems,
-        hf_processor_mm_kwargs: Mapping[str, object],
-        tokenization_kwargs: Mapping[str, object],
-    ) -> BatchFeature:
-        """
-        Apply the HF processor on the multi-modal data only.
-
-        Issue: Voxtral TTS use Mistral Tokenizer with custom audio encoder. It doesn't
-        inherit Transformers ProcessorMixin and can't use call_hf_processor_mm_only.
-
-        Solution: Override this method to call _apply_hf_processor_text_mm directly.
-        """
-        mm_counts = mm_items.get_all_counts()
-        _, mm_processed_data, _ = self._apply_hf_processor_text_mm(
-            prompt_text=self.dummy_inputs.get_dummy_text(mm_counts),
-            mm_items=mm_items,
-            hf_processor_mm_kwargs=hf_processor_mm_kwargs,
-            tokenization_kwargs=tokenization_kwargs,
-        )
-        return mm_processed_data
-
     def _maybe_apply_prompt_updates(
         self,
         mm_items: MultiModalDataItems,
-        prompt_ids: list[int],
-        mm_kwargs: MultiModalKwargsItems,
-        mm_prompt_updates,
-    ):
+        mm_res: MultiModalProcessingResult,
+    ) -> tuple[list[int], Mapping[str, list[PlaceholderFeaturesInfo]]]:
         # Voxtral's Mistral chat template has already inserted the complete
         # audio-token run, so locate it without applying the replacement again.
         mm_item_counts = mm_items.get_all_counts()
+        prompt_ids = mm_res.prompt_ids
+        mm_kwargs = mm_res.kwargs
+        mm_prompt_updates = mm_res.prompt_updates
         self._validate_mm_kwargs(mm_kwargs, mm_item_counts)
         self._validate_mm_updates(mm_prompt_updates, mm_item_counts)
         mm_placeholders = self._find_mm_placeholders(

vllm_omni/outputs/__init__.py

@@ -4,6 +4,7 @@
 from dataclasses import dataclass, field
 from typing import Any
 
+import numpy as np
 import torch
 from PIL import Image
 from vllm.outputs import CompletionOutput, RequestOutput
@@ -86,6 +87,7 @@ _REQUEST_OUTPUT_CONTENT_ATTRS = (
     "prompt",
     "prompt_token_ids",
     "prompt_logprobs",
+    "prompt_token_id_logprobs",
     "outputs",
     "finished",
     "lora_request",
@@ -149,6 +151,7 @@ class OmniRequestOutput(RequestOutput):
     prompt: OmniPromptType | None = None
     prompt_token_ids: list[int] | None = None
     prompt_logprobs: Any = None
+    prompt_token_id_logprobs: np.ndarray | None = None
     outputs: list[CompletionOutput] = field(default_factory=list)
     finished: bool = True
     metrics: Any = field(default_factory=dict)

vllm_omni/quantization/factory.py

@@ -81,6 +81,16 @@ from .component_config import ComponentQuantizationConfig  # noqa: E402
 logger = init_logger(__name__)
 
 
+def _build_fp8(**kw: Any) -> QuantizationConfig:
+    if kw.get("is_checkpoint_fp8_serialized", False):
+        from vllm.model_executor.layers.quantization.fp8 import Fp8Config
+
+        return Fp8Config(**kw)
+    from .fp8_config import DiffusionFp8Config
+
+    return DiffusionFp8Config(**kw)
+
+
 def _build_int8(**kw: Any) -> QuantizationConfig:
     """Lazy import for Int8 diffusion config (supports CUDA + NPU)."""
     from .int8_config import DiffusionInt8Config
@@ -168,6 +178,7 @@ def _build_torchao_float8_weight_only(**kw: Any) -> QuantizationConfig:
 
 
 _OVERRIDES: dict[str, Callable[..., QuantizationConfig]] = {
+    "fp8": _build_fp8,
     "int8": _build_int8,
     "bitsandbytes": _build_bitsandbytes,
     "mxfp8": _build_mxfp8,
@@ -330,6 +341,17 @@ def _build_single(method: str, **kwargs: Any) -> QuantizationConfig:
 
     config_cls = get_quantization_config(method)
 
+    from vllm.config.quantization import resolve_quantization_config
+    from vllm.model_executor.layers.quantization.online.base import OnlineQuantizationConfig
+
+    if config_cls is OnlineQuantizationConfig:
+        if "args" in kwargs:
+            return config_cls(**kwargs)
+        args = resolve_quantization_config(method, kwargs or None)
+        if args is None:
+            raise ValueError("Online quantization requires quantization config arguments")
+        return config_cls(args)
+
     try:
         return config_cls(**kwargs)
     except TypeError:

vllm_omni/quantization/fp8_config.py

@@ -0,0 +1,47 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Omni's online FP8 option backed by vLLM's online quantization API."""
+
+from vllm.config.quantization import QuantizationConfigArgs, QuantSpec
+from vllm.model_executor.layers.quantization.online.base import OnlineQuantizationConfig
+
+
+class DiffusionFp8Config(OnlineQuantizationConfig):
+    """Keep the diffusion ``fp8`` option for quantizing BF16/FP16 weights.
+
+    Upstream Fp8Config now describes serialized checkpoints only. Keeping the
+    Omni method name also lets checkpoint metadata replace this online config
+    through the existing serialized-checkpoint reconciliation path.
+    """
+
+    def __init__(
+        self,
+        is_checkpoint_fp8_serialized: bool = False,
+        activation_scheme: str = "dynamic",
+        ignored_layers: list[str] | None = None,
+        weight_block_size: list[int] | None = None,
+        store_dtype: str | None = None,
+    ) -> None:
+        if is_checkpoint_fp8_serialized:
+            raise ValueError("Serialized FP8 checkpoints require upstream Fp8Config")
+        if activation_scheme != "dynamic":
+            raise ValueError("Online FP8 quantization requires activation_scheme='dynamic'")
+        if weight_block_size is not None:
+            raise ValueError("Block-wise FP8 quantization requires a serialized checkpoint or fp8_per_block")
+        if store_dtype is not None:
+            raise ValueError("Online FP8 quantization does not support store_dtype")
+        super().__init__(
+            QuantizationConfigArgs(
+                linear=QuantSpec(weight="fp8_per_tensor_static"),
+                moe=QuantSpec(weight="fp8_per_tensor_static"),
+                ignore=list(ignored_layers or []),
+            )
+        )
+        self.is_checkpoint_fp8_serialized = False
+        self.activation_scheme = activation_scheme
+        self.weight_block_size = None
+        self.store_dtype = None
+
+    @classmethod
+    def get_name(cls) -> str:
+        return "fp8"

vllm_omni/request.py

@@ -107,6 +107,7 @@ class OmniRequest(Request):
             model_intermediate_buffer=getattr(request, "model_intermediate_buffer", None),
             resumable=request.resumable,
             session_id=request.session_id,
+            kv_hints=request.kv_hints,
             reasoning_ended=request.reasoning_ended,
             reasoning_parser_kwargs=request.reasoning_parser_kwargs,
             abort_immediately=request.abort_immediately,

vllm_omni/worker/gpu_ar_model_runner.py

@@ -206,7 +206,6 @@ class OmniAsyncGPUModelRunnerOutput(AsyncGPUModelRunnerOutput):
         invalid_req_indices = kwargs.pop("invalid_req_indices")
         async_output_copy_stream = kwargs.pop("async_output_copy_stream")
         vocab_size = kwargs.pop("vocab_size")
-        routed_experts = kwargs.pop("routed_experts", None)
         num_nans = kwargs.pop("num_nans", None)
         # Upstream AsyncGPUModelRunnerOutput added check_ep_fault / _has_fault
         # for EP all2all fault tolerance (PR #43637). Omni doesn't use this
@@ -223,7 +222,6 @@ class OmniAsyncGPUModelRunnerOutput(AsyncGPUModelRunnerOutput):
         self._sampled_token_ids = sampled_token_ids
         self.vocab_size = vocab_size
         self._logprobs_tensors = logprobs_tensors
-        self._routed_experts = routed_experts
         self._has_fault: torch.Tensor | None = None
         # Upstream b1e12d142d (PR #51304) added device-side NaN-in-logits
         # counts (num_nans) to AsyncGPUModelRunnerOutput. Omni keeps the
@@ -240,9 +238,6 @@ class OmniAsyncGPUModelRunnerOutput(AsyncGPUModelRunnerOutput):
             # output asynchronously.
             self.sampled_token_ids_cpu = self._sampled_token_ids.to("cpu", non_blocking=True)
             self._logprobs_tensors_cpu = self._logprobs_tensors.to_cpu_nonblocking() if self._logprobs_tensors else None
-            self._routed_experts_cpu = (
-                self._routed_experts.to_cpu_nonblocking() if self._routed_experts is not None else None
-            )
             self._num_nans_cpu = self._num_nans.to("cpu", non_blocking=True) if self._num_nans is not None else None
             self.async_copy_ready_event.record()
 
@@ -902,9 +897,6 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin, Duplex
         if self.execute_model_state is not None:
             raise RuntimeError("State error: sample_tokens() must be called after execute_model() returns None.")
 
-        if self.routed_experts_initialized:
-            self.routed_experts_capturer.clear_buffer()
-
         if not getattr(self, "_warmup_state_cleared", False):
             self._warmup_state_cleared = True
             if hasattr(self.model, "_clear_warmup_state"):
@@ -936,11 +928,6 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin, Duplex
                 if flush_ids:
                     self.flush_full_payload_outputs(flush_ids)
 
-        if self.routed_experts_initialized:
-            capturer = self.routed_experts_capturer
-            if capturer is not None and hasattr(capturer, "finalize_pending_copy"):
-                capturer.finalize_pending_copy()
-
         # If ngram_gpu is used, we need to copy the scheduler_output to avoid
         # the modification has influence on the scheduler_output in engine core process.
         # The replace is much faster than deepcopy.
@@ -1332,9 +1319,6 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin, Duplex
         if deferred_state_corrections_fn:
             deferred_state_corrections_fn()
 
-        if self._should_return_omni_routed_experts() and hasattr(self, "_positions_cpu"):
-            self._omni_routed_experts_d2h(scheduler_output)
-
         return None
 
     def _sample(
@@ -1587,14 +1571,6 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin, Duplex
         except TypeError:
             return scheduler_output
 
-    def _should_return_omni_routed_experts(self) -> bool:
-        model_config = getattr(self, "model_config", None)
-        if model_config is None:
-            model_config = getattr(getattr(self, "vllm_config", None), "model_config", None)
-        return bool(getattr(model_config, "enable_return_routed_experts", False)) and bool(
-            getattr(self, "routed_experts_initialized", False)
-        )
-
     @staticmethod
     def _model_omni_flag(model: Any, name: str, default: bool = False) -> bool:
         return bool(getattr(model, name, default)) if model is not None else default
@@ -1659,8 +1635,6 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin, Duplex
             getattr(self, "model", None), "supports_async_whole_payload"
         ):
             return False
-        if bool(getattr(model_config, "enable_return_routed_experts", False)):
-            return False
 
         model = getattr(self, "model", None)
         if not self._model_omni_flag(model, "use_async_omni_output"):
@@ -1945,9 +1919,6 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin, Duplex
             )
 
         with record_function_or_nullcontext("gpu_model_runner: ModelRunnerOutput"):
-            routed_experts_lists = None
-            if self._should_return_omni_routed_experts():
-                routed_experts_lists = self._omni_extract_routed_experts(scheduler_output)
             output = OmniModelRunnerOutput(
                 req_ids=req_ids_output_copy,
                 req_id_to_index=req_id_to_index_output_copy,
@@ -1963,7 +1934,6 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin, Duplex
                 cudagraph_stats=cudagraph_stats,
             )
             output.kv_extracted_req_ids = kv_extracted_req_ids
-            output.routed_experts = routed_experts_lists
         return output
 
     @torch.inference_mode()

vllm_omni/worker/gpu_generation_model_runner.py

@@ -137,9 +137,6 @@ class GPUGenerationModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin
         if self.execute_model_state is not None:
             raise RuntimeError("State error: sample_tokens() must be called after execute_model() returns None.")
 
-        if self.routed_experts_initialized:
-            self.routed_experts_capturer.clear_buffer()
-
         if hasattr(self, "_omni_connector"):
             for request in getattr(scheduler_output, "pending_input_registrations", []):
                 self.register_chunk_recv(request)
@@ -150,11 +147,6 @@ class GPUGenerationModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin
                 if flush_ids:
                     self.flush_full_payload_outputs(flush_ids)
 
-        if self.routed_experts_initialized:
-            capturer = self.routed_experts_capturer
-            if capturer is not None and hasattr(capturer, "finalize_pending_copy"):
-                capturer.finalize_pending_copy()
-
         # If ngram_gpu is used, we need to copy the scheduler_output to avoid
         # the modification has influence on the scheduler_output in engine core process.
         # The replace is much faster than deepcopy.
@@ -401,9 +393,6 @@ class GPUGenerationModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin
         if deferred_state_corrections_fn:
             deferred_state_corrections_fn()
 
-        if self.routed_experts_initialized and hasattr(self, "_positions_cpu"):
-            self._omni_routed_experts_d2h(scheduler_output)
-
         return None
 
     @torch.inference_mode()
@@ -516,9 +505,6 @@ class GPUGenerationModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin
         # [Omni] Copy req_id mappings to avoid async scheduling mutation.
         req_ids_output_copy = self.input_batch.req_ids.copy()
         req_id_to_index_output_copy = self.input_batch.req_id_to_index.copy()
-        routed_experts_lists = None
-        if self.routed_experts_initialized:
-            routed_experts_lists = self._omni_extract_routed_experts(scheduler_output)
         if inter_stage_outputs and self._should_accumulate_full_payload_output():
             for i, rid in enumerate(req_ids_output_copy):
                 req_state = self.requests.get(rid)
@@ -541,7 +527,6 @@ class GPUGenerationModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin
             ec_connector_output=ec_connector_output if self.supports_mm_inputs else None,
         )
         output.omni_connector_output = self.get_omni_connector_output()
-        output.routed_experts = routed_experts_lists
 
         if not self.use_async_scheduling:
             return output

vllm_omni/worker/gpu_model_runner.py

@@ -45,7 +45,6 @@ from vllm_omni.utils.device_copy import index_to_device
 
 if TYPE_CHECKING:
     from vllm.v1.core.sched.output import SchedulerOutput
-    from vllm.v1.outputs import RoutedExpertsLists
 else:
     xgr = LazyLoader("xgr", globals(), "xgrammar")
     xgr_torch_compile = LazyLoader(
@@ -108,52 +107,6 @@ class OmniGPUModelRunner(PrefixCacheRunnerMixin, GPUModelRunner):
                 return sampled
         return super()._to_list(sampled_token_ids)
 
-    def _omni_routed_experts_d2h(self, scheduler_output) -> None:
-        """Issue routed-experts D2H copy matching upstream GPUModelRunner pattern.
-
-        Upstream does this inline in ``execute_model``:
-            buf = self.routed_experts_capturer.get_device_buffer()
-            total = scheduler_output.total_num_scheduled_tokens
-            self.routed_experts_cpu[:total].copy_(buf[:total], non_blocking=True)
-            self.routed_experts_slot_mapping_cpu[:total].copy_(
-                self.routed_experts_slot_mapping_device[:total], non_blocking=True)
-        """
-        if not self.routed_experts_initialized:
-            return
-        buf = self.routed_experts_capturer.get_device_buffer()
-        total = scheduler_output.total_num_scheduled_tokens
-        self.routed_experts_cpu[:total].copy_(buf[:total], non_blocking=True)
-        if hasattr(self, "routed_experts_slot_mapping_device"):
-            self.routed_experts_slot_mapping_cpu[:total].copy_(
-                self.routed_experts_slot_mapping_device[:total],
-                non_blocking=True,
-            )
-
-    def _omni_extract_routed_experts(self, scheduler_output) -> "RoutedExpertsLists | None":
-        """Extract routed experts matching upstream GPUModelRunner pattern.
-
-        Upstream (sync path, sample_tokens):
-            total = scheduler_output.total_num_scheduled_tokens
-            output.routed_experts = RoutedExpertsLists(
-                routing_data=self.routed_experts_cpu[:total].numpy(),
-                slot_mapping=self.routed_experts_slot_mapping_cpu[:total].numpy(),
-            )
-
-        Returns RoutedExpertsLists (batch-level, with slot_mapping) so that
-        downstream schedulers can use slot_mapping to map back to requests.
-        """
-        from vllm.v1.outputs import RoutedExpertsLists
-
-        if not self.routed_experts_initialized:
-            return None
-        total = scheduler_output.total_num_scheduled_tokens
-        if total <= 0:
-            return None
-        return RoutedExpertsLists(
-            routing_data=self.routed_experts_cpu[:total].numpy(),
-            slot_mapping=self.routed_experts_slot_mapping_cpu[:total].numpy(),
-        )
-
     def initialize_metadata_builders(self, kv_cache_config, kernel_block_sizes):
         """Initialize metadata builders and keep FA3 graph metadata buffers sized.
 

vllm_omni/worker_v2/omni_ar_model_runner.py

@@ -24,7 +24,7 @@ from vllm.model_executor.layers.fused_moe.all2all_utils import (
 )
 from vllm.utils.torch_utils import PIN_MEMORY
 from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput
-from vllm.v1.outputs import AsyncModelRunnerOutput, ModelRunnerOutput, RoutedExpertsTensors
+from vllm.v1.outputs import AsyncModelRunnerOutput, ModelRunnerOutput
 from vllm.v1.worker.gpu.eplb_utils import step_eplb_after
 
 from vllm_omni.data_entry_keys import OmniPayload, flatten_payload
@@ -151,7 +151,9 @@ class OmniARModelRunner(OmniGPUModelRunner):
         hidden_states = self.execute_model_state.hidden_states
         finished_req_ids = self.execute_model_state.finished_req_ids
         ec_connector_output = self.execute_model_state.ec_connector_output
-        routed_experts = self.execute_model_state.routed_experts
+        pending_aux_output = None
+        if self.aux_output_connector is not None:
+            pending_aux_output = self.aux_output_connector.prepare_output(input_batch)
         self.execute_model_state = None
 
         # --- Omni: reconstruct raw model output and post-process ---
@@ -234,6 +236,12 @@ class OmniARModelRunner(OmniGPUModelRunner):
             self.req_states.num_computed_tokens.gpu,
             self.req_states.prompt_len.np,
         )
+        prompt_token_id_logprobs_dict = self.prompt_logprobs_worker.compute_prompt_token_id_logprobs(
+            self.model.compute_logits,
+            text_hidden,
+            input_batch,
+            self.req_states.prompt_len.np,
+        )
 
         # --- Omni: pooler_output ---
         engine_output_type = getattr(self.vllm_config.model_config, "engine_output_type", "text")
@@ -245,6 +253,7 @@ class OmniARModelRunner(OmniGPUModelRunner):
             req_id_to_index={rid: i for i, rid in enumerate(input_batch.req_ids)},
             sampled_token_ids=None,  # type: ignore[arg-type]
             prompt_logprobs_dict=prompt_logprobs_dict,
+            prompt_token_id_logprobs_dict=prompt_token_id_logprobs_dict,
             kv_connector_output=None,
         )
         model_runner_output.kv_extracted_req_ids = kv_extracted
@@ -265,7 +274,7 @@ class OmniARModelRunner(OmniGPUModelRunner):
             finalize_output=(None if materialize_native else self._finalize_native_data_plane_output),
             finalize_multimodal=sampling_output.finalize_multimodal,
             check_ep_fault=self.check_ep_fault,
-            routed_experts=routed_experts,
+            pending_aux_output=pending_aux_output,
             extra_multimodal_outputs=extra_outputs,
         )
         self._release_multimodal_snapshot(snapshot_slot, async_output.copy_event)
@@ -678,13 +687,13 @@ class OmniAsyncOutput(AsyncModelRunnerOutput):
         finalize_output: Any | None = None,
         finalize_multimodal: Any | None = None,
         check_ep_fault: bool = False,
-        routed_experts: RoutedExpertsTensors | None = None,
+        pending_aux_output: Any | None = None,
         extra_multimodal_outputs: tuple[dict[str, Any], torch.cuda.Event] | None = None,
     ):
         self.model_runner_output = model_runner_output
         self.sampler_output = sampler_output
         self.num_sampled_tokens = num_sampled_tokens
-        self.routed_experts = routed_experts
+        self.pending_aux_output = pending_aux_output
         self.copy_event = copy_event if copy_event is not None else torch.cuda.Event(blocking=True)
         self._async_chunk = bool(async_chunk)
         self._finalize_output = finalize_output
@@ -739,9 +748,15 @@ class OmniAsyncOutput(AsyncModelRunnerOutput):
             self.sampling_mask_tensors = None
             if sampler_output.sampling_mask_tensors is not None:
                 self.sampling_mask_tensors = sampler_output.sampling_mask_tensors.to_cpu_nonblocking()
-            self.routed_experts_cpu = None
-            if routed_experts is not None:
-                self.routed_experts_cpu = routed_experts.to_cpu_nonblocking()
+            if pending_aux_output is not None:
+                pending_aux_output.enqueue_cpu_copy(
+                    num_sampled=self.num_sampled_tokens_np,
+                    num_rejected=_async_copy_to_np(
+                        sampler_output.num_rejected,
+                        copy_stream=copy_stream,
+                        pin_memory=pin_memory,
+                    ),
+                )
 
             # Logprobs
             self.logprobs_tensors = None
@@ -760,6 +775,10 @@ class OmniAsyncOutput(AsyncModelRunnerOutput):
                 k: v.to_cpu_nonblocking() if v is not None else None
                 for k, v in self.model_runner_output.prompt_logprobs_dict.items()
             }
+            self.prompt_token_id_logprobs_dict = {
+                k: _async_copy_tensor(v, copy_stream=copy_stream, pin_memory=pin_memory)
+                for k, v in self.model_runner_output.prompt_token_id_logprobs_dict.items()
+            }
             if check_ep_fault:
                 has_fault = get_ep_all2all_manager().query_fault()
                 self._has_fault = has_fault.to("cpu", non_blocking=True)
@@ -820,8 +839,8 @@ class OmniAsyncOutput(AsyncModelRunnerOutput):
         self.model_runner_output.sampled_token_ids = sampled_token_ids
         if self.sampling_mask_tensors is not None:
             self.model_runner_output.sampling_masks = self.sampling_mask_tensors.tolists()
-        if self.routed_experts_cpu is not None:
-            self.model_runner_output.routed_experts = self.routed_experts_cpu.tolists()
+        if self.pending_aux_output is not None:
+            self.model_runner_output.aux_output_connector_output = self.pending_aux_output.process_output()
         self.model_runner_output.sampled_token_ids_materialized = True
 
         # Logprobs
@@ -832,6 +851,7 @@ class OmniAsyncOutput(AsyncModelRunnerOutput):
         if self.logprobs_tensors is not None:
             self.model_runner_output.logprobs = self.logprobs_tensors.tolists()
         self.model_runner_output.prompt_logprobs_dict = self.prompt_logprobs_dict
+        self.model_runner_output.prompt_token_id_logprobs_dict = self.prompt_token_id_logprobs_dict
 
         if self._has_fault is not None and self._has_fault.item():
             mask = get_ep_all2all_manager().query_active_mask()

vllm_omni/worker_v2/omni_generation_model_runner.py

@@ -14,6 +14,7 @@ from __future__ import annotations
 from dataclasses import replace
 from typing import Any
 
+import numpy as np
 import torch
 from vllm.config.compilation import CUDAGraphMode
 from vllm.forward_context import set_forward_context
@@ -123,10 +124,12 @@ class OmniGenerationAsyncOutput(AsyncModelRunnerOutput):
         copy_stream: torch.cuda.Stream,
         finalize_output: Any | None = None,
         check_ep_fault: bool = False,
+        pending_aux_output: Any | None = None,
     ) -> None:
         self.model_runner_output = model_runner_output
         self.num_reqs = num_reqs
         self.finalize_output = finalize_output
+        self.pending_aux_output = pending_aux_output
         self.copy_event = torch.cuda.Event(blocking=True)
         self._has_fault: torch.Tensor | None = None
 
@@ -138,6 +141,10 @@ class OmniGenerationAsyncOutput(AsyncModelRunnerOutput):
                 copy_stream=copy_stream,
                 pin_memory=PIN_MEMORY,
             )
+            if pending_aux_output is not None:
+                # Generation stages emit no sampled/rejected token IDs.
+                counts = np.zeros(num_reqs, dtype=np.int32)
+                pending_aux_output.enqueue_cpu_copy(num_sampled=counts, num_rejected=counts)
             if check_ep_fault:
                 has_fault = get_ep_all2all_manager().query_fault()
                 self._has_fault = has_fault.to("cpu", non_blocking=True)
@@ -145,6 +152,8 @@ class OmniGenerationAsyncOutput(AsyncModelRunnerOutput):
 
     def get_output(self) -> OmniModelRunnerOutput:
         self.copy_event.synchronize()
+        if self.pending_aux_output is not None:
+            self.model_runner_output.aux_output_connector_output = self.pending_aux_output.process_output()
         if self._has_fault is not None and self._has_fault.item():
             mask = get_ep_all2all_manager().query_active_mask()
             raise RuntimeError(
@@ -321,6 +330,8 @@ class OmniGenerationModelRunner(OmniGPUModelRunner):
                 # Registration/terminal/abort work has been applied above.
                 # There is no new or cached request state to add or update.
                 self._apply_block_table_staged_writes_if_available()
+                if self.aux_output_connector is not None:
+                    self.aux_output_connector.begin_step(scheduler_output.aux_output_connector_metadata)
                 empty_output = self.kv_connector.no_forward(scheduler_output)
                 return self._attach_native_data_plane_signals(
                     self._merge_ec_connector_no_forward(scheduler_output, empty_output)
@@ -333,6 +344,8 @@ class OmniGenerationModelRunner(OmniGPUModelRunner):
             self.update_requests(scheduler_output)
             self._sync_native_data_plane_payloads(scheduler_output)
             self._apply_block_table_staged_writes_if_available()
+            if self.aux_output_connector is not None:
+                self.aux_output_connector.begin_step(scheduler_output.aux_output_connector_metadata)
             if scheduler_output.total_num_scheduled_tokens == 0:
                 empty_output = self.kv_connector.no_forward(scheduler_output)
                 return self._attach_native_data_plane_signals(
@@ -366,7 +379,7 @@ class OmniGenerationModelRunner(OmniGPUModelRunner):
 
         if not dummy_run:
             assert batch_req_state is not None
-            input_batch = self.prepare_inputs(scheduler_output, batch_req_state, batch_desc)
+            input_batch = self.prepare_inputs(scheduler_output, batch_req_state, batch_desc, 0)
         else:
             from vllm.v1.worker.gpu.input_batch import InputBatch
 
@@ -451,7 +464,6 @@ class OmniGenerationModelRunner(OmniGPUModelRunner):
             finished_req_ids=scheduler_output.finished_req_ids,
             dp_sync=dp_sync,
             ec_connector_output=ec_connector_output,
-            routed_experts=None,
             cudagraph_stats=None,
         )
         return None
@@ -476,6 +488,9 @@ class OmniGenerationModelRunner(OmniGPUModelRunner):
             return None
 
         num_reqs = input_batch.num_reqs
+        pending_aux_output = None
+        if self.aux_output_connector is not None:
+            pending_aux_output = self.aux_output_connector.prepare_output(input_batch)
 
         # Mark all scheduled tokens as computed so the scheduler does
         # not re-schedule them.  Unlike AR stages we do NOT call
@@ -517,7 +532,7 @@ class OmniGenerationModelRunner(OmniGPUModelRunner):
 
         raw_multimodal_outputs = model_output.multimodal_outputs
         can_copy_async = getattr(getattr(self, "device", None), "type", "cpu") == "cuda" and (
-            self.check_ep_fault or _contains_cuda_tensor(raw_multimodal_outputs)
+            self.check_ep_fault or pending_aux_output is not None or _contains_cuda_tensor(raw_multimodal_outputs)
         )
         if can_copy_async:
             async_output = OmniGenerationAsyncOutput(
@@ -528,6 +543,7 @@ class OmniGenerationModelRunner(OmniGPUModelRunner):
                 copy_stream=self.output_copy_stream,
                 finalize_output=self._finalize_native_data_plane_output,
                 check_ep_fault=self.check_ep_fault,
+                pending_aux_output=pending_aux_output,
             )
             self._reserve_native_data_plane_outputs(list(req_ids))
             self._release_generation_slots(input_batch)

vllm_omni/worker_v2/omni_model_runner.py

@@ -499,6 +499,8 @@ class OmniGPUModelRunner(GPUModelRunner):
             self.update_requests(scheduler_output)
             self._sync_native_data_plane_payloads(scheduler_output)
             self.block_tables.apply_staged_writes()
+            if self.aux_output_connector is not None:
+                self.aux_output_connector.begin_step(scheduler_output.aux_output_connector_metadata)
             if scheduler_output.total_num_scheduled_tokens == 0:
                 empty_output = self.kv_connector.no_forward(scheduler_output)
                 return self._attach_native_data_plane_signals(
@@ -540,7 +542,7 @@ class OmniGPUModelRunner(GPUModelRunner):
 
         if not dummy_run:
             assert batch_req_state is not None
-            input_batch = self.prepare_inputs(scheduler_output, batch_req_state, batch_desc)
+            input_batch = self.prepare_inputs(scheduler_output, batch_req_state, batch_desc, num_active_loras)
             block_tables, slot_mappings = self.prepare_attn(input_batch)
             self.model_state.preprocess_state(
                 input_batch,
@@ -681,11 +683,6 @@ class OmniGPUModelRunner(GPUModelRunner):
         if not dummy_run and isinstance(hidden_states, torch.Tensor):
             self.model_state.run_postprocess(hidden_states, input_batch)
 
-        routed_experts = None
-        if not dummy_run and (capturer := self.routed_experts_capturer) is not None:
-            assert slot_mappings is not None
-            routed_experts = capturer.get_routed_experts(slot_mappings, num_toks)
-
         self.execute_model_state = ExecuteModelState(
             input_batch=input_batch,
             attn_metadata=attn_metadata,
@@ -695,7 +692,6 @@ class OmniGPUModelRunner(GPUModelRunner):
             finished_req_ids=scheduler_output.finished_req_ids,
             dp_sync=dp_sync,
             ec_connector_output=ec_connector_output,
-            routed_experts=routed_experts,
             cudagraph_stats=None,
         )