diff --git a/tests/core/sched/test_generation_scheduler_restore.py b/tests/core/sched/test_generation_scheduler_restore.py index 546084509..89e4f2847 100644 --- a/tests/core/sched/test_generation_scheduler_restore.py +++ b/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 diff --git a/tests/core/sched/test_omni_ar_scheduler_aborted_queue_sweep.py b/tests/core/sched/test_omni_ar_scheduler_aborted_queue_sweep.py index f8b9d9b4c..3b4e9c702 100644 --- a/tests/core/sched/test_omni_ar_scheduler_aborted_queue_sweep.py +++ b/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) diff --git a/tests/core/sched/test_omni_ar_scheduler_logprobs.py b/tests/core/sched/test_omni_ar_scheduler_logprobs.py index 747be2ed1..d89cfe7ed 100644 --- a/tests/core/sched/test_omni_ar_scheduler_logprobs.py +++ b/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=[[]], diff --git a/tests/core/sched/test_omni_ar_scheduler_stale_drain.py b/tests/core/sched/test_omni_ar_scheduler_stale_drain.py index 182857ea9..a4aa50982 100644 --- a/tests/core/sched/test_omni_ar_scheduler_stale_drain.py +++ b/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, diff --git a/tests/core/sched/test_omni_ar_scheduler_streaming.py b/tests/core/sched/test_omni_ar_scheduler_streaming.py index 94d796fd6..973e74c48 100644 --- a/tests/core/sched/test_omni_ar_scheduler_streaming.py +++ b/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 diff --git a/tests/core/sched/test_omni_ar_scheduler_unified_decode_deferral.py b/tests/core/sched/test_omni_ar_scheduler_unified_decode_deferral.py index b55a668dd..a2b8f82b7 100644 --- a/tests/core/sched/test_omni_ar_scheduler_unified_decode_deferral.py +++ b/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] diff --git a/tests/core/sched/test_omni_cfg_ar_scheduler.py b/tests/core/sched/test_omni_cfg_ar_scheduler.py index e9ff8ea53..b6db7a186 100644 --- a/tests/core/sched/test_omni_cfg_ar_scheduler.py +++ b/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 diff --git a/tests/core/sched/test_omni_generation_scheduler_update_session.py b/tests/core/sched/test_omni_generation_scheduler_update_session.py index 0f509c5d2..f1c099622 100644 --- a/tests/core/sched/test_omni_generation_scheduler_update_session.py +++ b/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: diff --git a/tests/core/sched/test_omni_sched_deferred_free_fence.py b/tests/core/sched/test_omni_sched_deferred_free_fence.py index 9088c6bc2..3bdbd2925 100644 --- a/tests/core/sched/test_omni_sched_deferred_free_fence.py +++ b/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) diff --git a/tests/core/sched/test_omni_scheduler_finish_requests_purge.py b/tests/core/sched/test_omni_scheduler_finish_requests_purge.py index 770c85578..9fc0a4c0e 100644 --- a/tests/core/sched/test_omni_scheduler_finish_requests_purge.py +++ b/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 diff --git a/tests/core/sched/test_omni_scheduler_input_coordinator_cleanup.py b/tests/core/sched/test_omni_scheduler_input_coordinator_cleanup.py index 008e52636..84773fc6a 100644 --- a/tests/core/sched/test_omni_scheduler_input_coordinator_cleanup.py +++ b/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"] diff --git a/tests/core/sched/test_omni_scheduler_mixin_shared.py b/tests/core/sched/test_omni_scheduler_mixin_shared.py index 82956e27a..086593593 100644 --- a/tests/core/sched/test_omni_scheduler_mixin_shared.py +++ 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), ] diff --git a/tests/core/sched/test_omni_scheduler_ready_inbox.py b/tests/core/sched/test_omni_scheduler_ready_inbox.py index 483255f21..db215eeb0 100644 --- a/tests/core/sched/test_omni_scheduler_ready_inbox.py +++ b/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 diff --git a/tests/core/sched/test_omni_scheduler_streaming_input_counter.py b/tests/core/sched/test_omni_scheduler_streaming_input_counter.py index bc392abca..d58b9f9a4 100644 --- a/tests/core/sched/test_omni_scheduler_streaming_input_counter.py +++ b/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 == [] diff --git a/tests/diffusion/cache/test_seacache.py b/tests/diffusion/cache/test_seacache.py index 5f88db5c6..a876b6b36 100644 --- a/tests/diffusion/cache/test_seacache.py +++ b/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 = [ diff --git a/tests/diffusion/diffusion_kv/test_kv_cache_utils.py b/tests/diffusion/diffusion_kv/test_kv_cache_utils.py index 543524404..ba617079f 100644 --- a/tests/diffusion/diffusion_kv/test_kv_cache_utils.py +++ b/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): diff --git a/tests/diffusion/diffusion_kv/test_mooncake_cfg_fanout.py b/tests/diffusion/diffusion_kv/test_mooncake_cfg_fanout.py index 2d3ae0a20..f1974d326 100644 --- a/tests/diffusion/diffusion_kv/test_mooncake_cfg_fanout.py +++ b/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) diff --git a/tests/diffusion/diffusion_kv/test_paged_attention_adapter.py b/tests/diffusion/diffusion_kv/test_paged_attention_adapter.py index 49aba9140..8eb60202c 100644 --- a/tests/diffusion/diffusion_kv/test_paged_attention_adapter.py +++ b/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") diff --git a/tests/diffusion/distributed/test_comm.py b/tests/diffusion/distributed/test_comm.py index 21cb18821..a8d8a032a 100644 --- a/tests/diffusion/distributed/test_comm.py +++ b/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", ) diff --git a/tests/diffusion/distributed/test_pipeline_parallel.py b/tests/diffusion/distributed/test_pipeline_parallel.py index 7cebbcc53..c64e880cb 100644 --- a/tests/diffusion/distributed/test_pipeline_parallel.py +++ b/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() diff --git a/tests/diffusion/models/boogu_image/test_pipeline_boogu_image.py b/tests/diffusion/models/boogu_image/test_pipeline_boogu_image.py index b2dea0c44..4065cc517 100644 --- a/tests/diffusion/models/boogu_image/test_pipeline_boogu_image.py +++ b/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( diff --git a/tests/diffusion/models/flux2/test_flux2_quantization.py b/tests/diffusion/models/flux2/test_flux2_quantization.py index 3c461acf5..8ba413b9e 100644 --- a/tests/diffusion/models/flux2/test_flux2_quantization.py +++ b/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, diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_quantization.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_quantization.py index e1f83fd01..0e6197c96 100644 --- a/tests/diffusion/models/minimax_h3/test_minimax_h3_quantization.py +++ b/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, diff --git a/tests/diffusion/quantization/test_fp8_config.py b/tests/diffusion/quantization/test_fp8_config.py index 9bad16792..0241c8fd1 100644 --- a/tests/diffusion/quantization/test_fp8_config.py +++ b/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" diff --git a/tests/diffusion/quantization/test_online_fp8_gpu.py b/tests/diffusion/quantization/test_online_fp8_gpu.py new file mode 100644 index 000000000..e3cccb355 --- /dev/null +++ b/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() diff --git a/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py b/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py index 8f85366e3..6936ace7c 100644 --- a/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py +++ b/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 = {} diff --git a/tests/engine/test_arg_utils_shared_fields.py b/tests/engine/test_arg_utils_shared_fields.py index 40a584fce..55f12f4ff 100644 --- a/tests/engine/test_arg_utils_shared_fields.py +++ b/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) diff --git a/tests/engine/test_async_omni_engine_input.py b/tests/engine/test_async_omni_engine_input.py index 7ad111e25..8947d109f 100644 --- a/tests/engine/test_async_omni_engine_input.py +++ b/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) diff --git a/tests/engine/test_async_omni_engine_stage_init.py b/tests/engine/test_async_omni_engine_stage_init.py index 46ac15f9b..36f43125d 100644 --- a/tests/engine/test_async_omni_engine_stage_init.py +++ b/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): diff --git a/tests/engine/test_omni_request.py b/tests/engine/test_omni_request.py index 521c92b94..433162925 100644 --- a/tests/engine/test_omni_request.py +++ b/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" diff --git a/tests/engine/test_omni_request_output.py b/tests/engine/test_omni_request_output.py index eb1bdfd41..a7c8d2764 100644 --- a/tests/engine/test_omni_request_output.py +++ b/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. diff --git a/tests/engine/test_orchestrator_error_handling.py b/tests/engine/test_orchestrator_error_handling.py index dd1f1efb5..7e1635d19 100644 --- a/tests/engine/test_orchestrator_error_handling.py +++ b/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), diff --git a/tests/engine/test_output_processor_mrv2_text.py b/tests/engine/test_output_processor_mrv2_text.py index b9dbecb1b..acfbbe5aa 100644 --- a/tests/engine/test_output_processor_mrv2_text.py +++ b/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. diff --git a/tests/engine/test_stage_engine_args.py b/tests/engine/test_stage_engine_args.py index 3e8f0ef87..9b8b70586 100644 --- a/tests/engine/test_stage_engine_args.py +++ b/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", diff --git a/tests/engine/test_stage_engine_core_client.py b/tests/engine/test_stage_engine_core_client.py index 36c95cde7..b228055b7 100644 --- a/tests/engine/test_stage_engine_core_client.py +++ b/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, + ) diff --git a/tests/engine/test_stage_pool_collective_rpc.py b/tests/engine/test_stage_pool_collective_rpc.py index 281b4b504..cede3bf16 100644 --- a/tests/engine/test_stage_pool_collective_rpc.py +++ b/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: diff --git a/tests/entrypoints/test_async_omni.py b/tests/entrypoints/test_async_omni.py index 1bcc6e08c..829e30145 100644 --- a/tests/entrypoints/test_async_omni.py +++ b/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", diff --git a/tests/entrypoints/test_async_omni_pause_sleep_routing.py b/tests/entrypoints/test_async_omni_pause_sleep_routing.py index 7da8f1396..bc4665dcc 100644 --- a/tests/entrypoints/test_async_omni_pause_sleep_routing.py +++ b/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: diff --git a/tests/helpers/cpu_collective_server.py b/tests/helpers/cpu_collective_server.py new file mode 100644 index 000000000..30cc80431 --- /dev/null +++ b/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") diff --git a/tests/helpers/distributed.py b/tests/helpers/distributed.py new file mode 100644 index 000000000..5c9c35879 --- /dev/null +++ b/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 diff --git a/tests/helpers/runtime.py b/tests/helpers/runtime.py index 5fedd17dd..2475431dd 100644 --- a/tests/helpers/runtime.py +++ b/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() # --------------------------------------------------------------------------- diff --git a/tests/helpers/tests/test_runtime_omni_runner.py b/tests/helpers/tests/test_runtime_omni_runner.py index aaada20f2..bcd29493d 100644 --- a/tests/helpers/tests/test_runtime_omni_runner.py +++ b/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 diff --git a/tests/inputs/test_mm_processor.py b/tests/inputs/test_mm_processor.py index 813d7f436..09bc5aaac 100644 --- a/tests/inputs/test_mm_processor.py +++ b/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 diff --git a/tests/inputs/test_preprocess.py b/tests/inputs/test_preprocess.py index 3fb48a7cc..c61920ba9 100644 --- a/tests/inputs/test_preprocess.py +++ b/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.""" diff --git a/tests/metrics/test_emit_calls.py b/tests/metrics/test_emit_calls.py index 02a98fc00..7a9d60100 100644 --- a/tests/metrics/test_emit_calls.py +++ b/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 diff --git a/tests/metrics/test_stat_logger.py b/tests/metrics/test_stat_logger.py index 97a18911d..86c43a500 100644 --- a/tests/metrics/test_stat_logger.py +++ b/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 diff --git a/tests/model_executor/models/omnivoice/test_processor_prompt_tokenization.py b/tests/model_executor/models/omnivoice/test_processor_prompt_tokenization.py index 9c6533c76..ab661b11a 100644 --- a/tests/model_executor/models/omnivoice/test_processor_prompt_tokenization.py +++ b/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__ diff --git a/tests/model_executor/models/qwen2_5_omni/test_per_video_audio_mask.py b/tests/model_executor/models/qwen2_5_omni/test_per_video_audio_mask.py index 699acc951..110776488 100644 --- a/tests/model_executor/models/qwen2_5_omni/test_per_video_audio_mask.py +++ b/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) diff --git a/tests/model_executor/models/test_omni_processing.py b/tests/model_executor/models/test_omni_processing.py index eaa798fba..f1493ad5e 100644 --- a/tests/model_executor/models/test_omni_processing.py +++ b/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( diff --git a/tests/model_executor/test_cfg_pairing.py b/tests/model_executor/test_cfg_pairing.py index c00313a2a..937b32b46 100644 --- a/tests/model_executor/test_cfg_pairing.py +++ b/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() diff --git a/tests/worker/test_omni_v2_profile_dispatch.py b/tests/worker/test_omni_v2_profile_dispatch.py index 4ff6fbfd4..949b2cace 100644 --- a/tests/worker/test_omni_v2_profile_dispatch.py +++ b/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 ) diff --git a/tests/worker_v2/test_capture_query_bound_gpu.py b/tests/worker_v2/test_capture_query_bound_gpu.py index 93c53f313..cc59f3218 100644 --- a/tests/worker_v2/test_capture_query_bound_gpu.py +++ b/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) diff --git a/tests/worker_v2/test_omni_ar_model_runner.py b/tests/worker_v2/test_omni_ar_model_runner.py index e425cc9c8..0257e3a33 100644 --- a/tests/worker_v2/test_omni_ar_model_runner.py +++ b/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, diff --git a/tests/worker_v2/test_omni_generation_model_runner.py b/tests/worker_v2/test_omni_generation_model_runner.py index 2b12837e4..84927b3eb 100644 --- a/tests/worker_v2/test_omni_generation_model_runner.py +++ b/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). diff --git a/vllm_omni/core/sched/omni_ar_scheduler.py b/vllm_omni/core/sched/omni_ar_scheduler.py index 0138d365e..d41fbee30 100644 --- a/vllm_omni/core/sched/omni_ar_scheduler.py +++ b/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 diff --git a/vllm_omni/core/sched/omni_cfg_ar_scheduler.py b/vllm_omni/core/sched/omni_cfg_ar_scheduler.py index f1a72ac81..7d562502e 100644 --- a/vllm_omni/core/sched/omni_cfg_ar_scheduler.py +++ b/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: diff --git a/vllm_omni/core/sched/omni_generation_scheduler.py b/vllm_omni/core/sched/omni_generation_scheduler.py index cefa2a37a..21a9c160d 100644 --- a/vllm_omni/core/sched/omni_generation_scheduler.py +++ b/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 diff --git a/vllm_omni/core/sched/omni_scheduler_mixin.py b/vllm_omni/core/sched/omni_scheduler_mixin.py index 1a1a46276..24bb5a90b 100644 --- a/vllm_omni/core/sched/omni_scheduler_mixin.py +++ b/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 diff --git a/vllm_omni/core/sched/utils.py b/vllm_omni/core/sched/utils.py index 12cadae73..3bf336dc8 100644 --- a/vllm_omni/core/sched/utils.py +++ b/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. diff --git a/vllm_omni/diffusion/models/boogu_image/pipeline_boogu_image.py b/vllm_omni/diffusion/models/boogu_image/pipeline_boogu_image.py index 09c85f9ab..b8bc67af3 100644 --- a/vllm_omni/diffusion/models/boogu_image/pipeline_boogu_image.py +++ b/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." ) diff --git a/vllm_omni/engine/__init__.py b/vllm_omni/engine/__init__.py index c36d6d986..d7cb082e4 100644 --- a/vllm_omni/engine/__init__.py +++ b/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, diff --git a/vllm_omni/engine/arg_utils.py b/vllm_omni/engine/arg_utils.py index b9ed35de0..5582d1ae4 100644 --- a/vllm_omni/engine/arg_utils.py +++ b/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"): diff --git a/vllm_omni/engine/async_omni_engine.py b/vllm_omni/engine/async_omni_engine.py index c8ab8fcec..c86b882c8 100644 --- a/vllm_omni/engine/async_omni_engine.py +++ b/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( diff --git a/vllm_omni/engine/orchestrator.py b/vllm_omni/engine/orchestrator.py index 4047a7147..d0d5a4f5b 100644 --- a/vllm_omni/engine/orchestrator.py +++ b/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(): diff --git a/vllm_omni/engine/stage_engine_core_client.py b/vllm_omni/engine/stage_engine_core_client.py index a87654cf2..b4174f057 100644 --- a/vllm_omni/engine/stage_engine_core_client.py +++ b/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 diff --git a/vllm_omni/engine/stage_engine_core_proc.py b/vllm_omni/engine/stage_engine_core_proc.py index ecc9576bc..0ae88514e 100644 --- a/vllm_omni/engine/stage_engine_core_proc.py +++ b/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 diff --git a/vllm_omni/engine/stage_pool.py b/vllm_omni/engine/stage_pool.py index 222cd3b5d..5a6e19d39 100644 --- a/vllm_omni/engine/stage_pool.py +++ b/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( diff --git a/vllm_omni/entrypoints/async_omni.py b/vllm_omni/entrypoints/async_omni.py index 34fd76f22..0f660c0cc 100644 --- a/vllm_omni/entrypoints/async_omni.py +++ b/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]: diff --git a/vllm_omni/inputs/mm_processor.py b/vllm_omni/inputs/mm_processor.py index bb02fe961..125af757b 100644 --- a/vllm_omni/inputs/mm_processor.py +++ b/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 diff --git a/vllm_omni/inputs/preprocess.py b/vllm_omni/inputs/preprocess.py index dc4105b4d..e7337d92e 100644 --- a/vllm_omni/inputs/preprocess.py +++ b/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) diff --git a/vllm_omni/model_executor/models/bagel/bagel.py b/vllm_omni/model_executor/models/bagel/bagel.py index 30dced9d0..320b3a0d1 100644 --- a/vllm_omni/model_executor/models/bagel/bagel.py +++ b/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, diff --git a/vllm_omni/model_executor/models/common/cfg_pairing.py b/vllm_omni/model_executor/models/common/cfg_pairing.py index ae4b3ac38..e19ee2679 100644 --- a/vllm_omni/model_executor/models/common/cfg_pairing.py +++ b/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: diff --git a/vllm_omni/model_executor/models/glm_image/glm_image_ar.py b/vllm_omni/model_executor/models/glm_image/glm_image_ar.py index e97a78c25..ac45c0d5e 100644 --- a/vllm_omni/model_executor/models/glm_image/glm_image_ar.py +++ b/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, diff --git a/vllm_omni/model_executor/models/mimo_audio/mimo_audio.py b/vllm_omni/model_executor/models/mimo_audio/mimo_audio.py index 4dbc9a455..9bd8246f3 100644 --- a/vllm_omni/model_executor/models/mimo_audio/mimo_audio.py +++ b/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") diff --git a/vllm_omni/model_executor/models/mimo_audio/mimo_audio_llm.py b/vllm_omni/model_executor/models/mimo_audio/mimo_audio_llm.py index ec4faadb8..9e5517faa 100644 --- a/vllm_omni/model_executor/models/mimo_audio/mimo_audio_llm.py +++ b/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, diff --git a/vllm_omni/model_executor/models/ming_flash_omni/ming_flash_omni_thinker.py b/vllm_omni/model_executor/models/ming_flash_omni/ming_flash_omni_thinker.py index 1338f2606..6887be33c 100644 --- a/vllm_omni/model_executor/models/ming_flash_omni/ming_flash_omni_thinker.py +++ b/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] = {} diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_llm.py b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_llm.py index 2781fcc30..7f1d29345 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_llm.py +++ b/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, diff --git a/vllm_omni/model_executor/models/minimax_h3/text_encoder.py b/vllm_omni/model_executor/models/minimax_h3/text_encoder.py index 9ae74036e..f233287bb 100644 --- a/vllm_omni/model_executor/models/minimax_h3/text_encoder.py +++ b/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, diff --git a/vllm_omni/model_executor/models/omnivoice/omnivoice.py b/vllm_omni/model_executor/models/omnivoice/omnivoice.py index a0b0d725e..3fc53e754 100644 --- a/vllm_omni/model_executor/models/omnivoice/omnivoice.py +++ b/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)) diff --git a/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_thinker.py b/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_thinker.py index 1f89a37c4..8289a5fa2 100644 --- a/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_thinker.py +++ b/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) diff --git a/vllm_omni/model_executor/models/step_audio2/step_audio2_thinker.py b/vllm_omni/model_executor/models/step_audio2/step_audio2_thinker.py index 6dca8b358..2122db233 100644 --- a/vllm_omni/model_executor/models/step_audio2/step_audio2_thinker.py +++ b/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: diff --git a/vllm_omni/model_executor/models/voxcpm2/scheduler.py b/vllm_omni/model_executor/models/voxcpm2/scheduler.py index 05d73653f..17111c5ad 100644 --- a/vllm_omni/model_executor/models/voxcpm2/scheduler.py +++ b/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: diff --git a/vllm_omni/model_executor/models/voxtral_tts/voxtral_tts_audio_generation.py b/vllm_omni/model_executor/models/voxtral_tts/voxtral_tts_audio_generation.py index 211159afa..5eabb4dbc 100644 --- a/vllm_omni/model_executor/models/voxtral_tts/voxtral_tts_audio_generation.py +++ b/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( diff --git a/vllm_omni/outputs/__init__.py b/vllm_omni/outputs/__init__.py index ff0a21d45..791e91f5f 100644 --- a/vllm_omni/outputs/__init__.py +++ b/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) diff --git a/vllm_omni/quantization/factory.py b/vllm_omni/quantization/factory.py index c0975f651..55a50a090 100644 --- a/vllm_omni/quantization/factory.py +++ b/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: diff --git a/vllm_omni/quantization/fp8_config.py b/vllm_omni/quantization/fp8_config.py new file mode 100644 index 000000000..4f610d8e9 --- /dev/null +++ b/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" diff --git a/vllm_omni/request.py b/vllm_omni/request.py index a6da16696..643ab020a 100644 --- a/vllm_omni/request.py +++ b/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, diff --git a/vllm_omni/worker/gpu_ar_model_runner.py b/vllm_omni/worker/gpu_ar_model_runner.py index 330355c64..c8fa03145 100644 --- a/vllm_omni/worker/gpu_ar_model_runner.py +++ b/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() diff --git a/vllm_omni/worker/gpu_generation_model_runner.py b/vllm_omni/worker/gpu_generation_model_runner.py index a05e12708..c1f4fb8eb 100644 --- a/vllm_omni/worker/gpu_generation_model_runner.py +++ b/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 diff --git a/vllm_omni/worker/gpu_model_runner.py b/vllm_omni/worker/gpu_model_runner.py index d8af6b201..931e36d1c 100644 --- a/vllm_omni/worker/gpu_model_runner.py +++ b/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. diff --git a/vllm_omni/worker_v2/omni_ar_model_runner.py b/vllm_omni/worker_v2/omni_ar_model_runner.py index 32622e8ef..757db88be 100644 --- a/vllm_omni/worker_v2/omni_ar_model_runner.py +++ b/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() diff --git a/vllm_omni/worker_v2/omni_generation_model_runner.py b/vllm_omni/worker_v2/omni_generation_model_runner.py index 605531338..2de7314fd 100644 --- a/vllm_omni/worker_v2/omni_generation_model_runner.py +++ b/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) diff --git a/vllm_omni/worker_v2/omni_model_runner.py b/vllm_omni/worker_v2/omni_model_runner.py index 10f4053cf..ff8ac96f0 100644 --- a/vllm_omni/worker_v2/omni_model_runner.py +++ b/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, )