diff --git a/tests/core/sched/test_omni_scheduler_input_coordinator_cleanup.py b/tests/core/sched/test_omni_scheduler_input_coordinator_cleanup.py new file mode 100644 index 000000000..4088cbf78 --- /dev/null +++ b/tests/core/sched/test_omni_scheduler_input_coordinator_cleanup.py @@ -0,0 +1,94 @@ +from __future__ import annotations + +from types import SimpleNamespace + +import pytest +from vllm.v1.request import RequestStatus + +import vllm_omni.core.sched.omni_ar_scheduler as ar_sched_mod +import vllm_omni.core.sched.omni_generation_scheduler as gen_sched_mod +from vllm_omni.core.sched.omni_ar_scheduler import OmniARScheduler +from vllm_omni.core.sched.omni_generation_scheduler import OmniGenerationScheduler + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +class FakeInputCoordinator: + def __init__(self) -> None: + self.freed: list[str] = [] + + def free_finished_request(self, request_id: str) -> None: + self.freed.append(request_id) + + +class DummyRequest: + request_id = "req-free" + client_index = 0 + additional_information = None + + def is_finished(self) -> bool: + return True + + +@pytest.mark.parametrize( + ("scheduler_cls", "scheduler_mod"), + [ + (OmniARScheduler, ar_sched_mod), + (OmniGenerationScheduler, gen_sched_mod), + ], +) +def test_finish_requests_cleans_input_coordinator_for_finished_ids( + monkeypatch: pytest.MonkeyPatch, + scheduler_cls, + scheduler_mod, +) -> None: + coordinator = FakeInputCoordinator() + scheduler = scheduler_cls.__new__(scheduler_cls) + scheduler.chunk_transfer_adapter = None + scheduler.input_coordinator = coordinator + + def fake_finish_requests(self, request_ids, finished_status): + assert request_ids == ["req-a", "req-b"] + assert finished_status == RequestStatus.FINISHED_STOPPED + return [("req-a", 0), ("req-b", 1)] + + monkeypatch.setattr(scheduler_mod.VLLMScheduler, "finish_requests", fake_finish_requests) + + result = scheduler_cls.finish_requests( + scheduler, + ["req-a", "req-b"], + RequestStatus.FINISHED_STOPPED, + ) + + assert result == [("req-a", 0), ("req-b", 1)] + assert coordinator.freed == ["req-a", "req-b"] + + +def test_ar_free_request_cleans_input_coordinator_on_normal_free() -> None: + coordinator = FakeInputCoordinator() + scheduler = OmniARScheduler.__new__(OmniARScheduler) + scheduler.input_coordinator = coordinator + scheduler._omits_kv_transfer_cache = {"req-free": True} + scheduler.encoder_cache_manager = SimpleNamespace(free=lambda request: None) + scheduler.finished_req_ids = set() + scheduler.finished_req_ids_dict = None + scheduler._new_prompt_len_snapshot = {"req-free": 3} + scheduler._connector_finished = lambda request: (False, None) + scheduler._should_transfer_kv_for_request = lambda request_id: False + scheduler._free_blocks = lambda request: None + + result = OmniARScheduler._free_request(scheduler, DummyRequest()) + + assert result is None + assert coordinator.freed == ["req-free"] + assert scheduler._omits_kv_transfer_cache == {} + assert scheduler._new_prompt_len_snapshot == {} + + +def test_generation_free_request_cleans_input_coordinator(monkeypatch: pytest.MonkeyPatch) -> None: + coordinator = FakeInputCoordinator() + scheduler = OmniGenerationScheduler.__new__(OmniGenerationScheduler) + scheduler.input_coordinator = coordinator + + def fake_free_request(self, request, delay_free_blocks=False): + assert delay_free_blocks is True diff --git a/tests/core/sched/test_chunk_scheduling_coordinator.py b/tests/core/sched/test_omni_scheduling_coordinator.py similarity index 80% rename from tests/core/sched/test_chunk_scheduling_coordinator.py rename to tests/core/sched/test_omni_scheduling_coordinator.py index c9459d8ec..8c7cc20b0 100644 --- a/tests/core/sched/test_chunk_scheduling_coordinator.py +++ b/tests/core/sched/test_omni_scheduling_coordinator.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Unit tests for OmniSchedulingCoordinator (formerly ChunkSchedulingCoordinator). +"""Unit tests for OmniSchedulingCoordinator. These tests use mock request objects and mock queues. They do not require GPU, vLLM runtime, or any connector. @@ -11,10 +11,12 @@ from __future__ import annotations import unittest from types import SimpleNamespace +import torch + import vllm_omni.core.sched.omni_scheduling_coordinator as coord_mod from vllm_omni.core.sched.omni_scheduling_coordinator import ( - ChunkSchedulingCoordinator, OmniSchedulingCoordinator, + uses_qwen3_omni_full_payload_input_coordinator, ) # ------------------------------------------------------------------ # @@ -89,11 +91,55 @@ class MockQueue: # ------------------------------------------------------------------ # +class TestFullPayloadCoordinatorSelection(unittest.TestCase): + def test_qwen3_omni_talker_and_code2wav_use_full_payload_input_coordinator(self): + for model_stage in ("talker", "code2wav"): + model_config = SimpleNamespace( + stage_id=1, + async_chunk=False, + model_arch="Qwen3OmniMoeForConditionalGeneration", + model_stage=model_stage, + ) + + self.assertTrue(uses_qwen3_omni_full_payload_input_coordinator(model_config)) + + def test_async_chunk_and_non_qwen3_omni_do_not_use_full_payload_input_coordinator(self): + cases = [ + SimpleNamespace( + stage_id=1, + async_chunk=True, + model_arch="Qwen3OmniMoeForConditionalGeneration", + model_stage="talker", + ), + SimpleNamespace( + stage_id=1, + async_chunk=False, + model_arch="Qwen3TTSForConditionalGeneration", + model_stage="code2wav", + ), + SimpleNamespace( + stage_id=1, + async_chunk=False, + model_arch="Qwen2_5OmniForConditionalGeneration", + model_stage="talker", + ), + SimpleNamespace( + stage_id=0, + async_chunk=False, + model_arch="Qwen3OmniMoeForConditionalGeneration", + model_stage="thinker", + ), + ] + + for model_config in cases: + self.assertFalse(uses_qwen3_omni_full_payload_input_coordinator(model_config)) + + class TestChunkCoordinatorStateTransition(unittest.TestCase): """Test 5: process_pending_chunks transitions WAITING_FOR_CHUNK → target.""" def test_ready_request_transitions_to_waiting(self): - coord = ChunkSchedulingCoordinator(scheduler_max_num_seqs=10, stage_id=1, async_chunk=True) + coord = OmniSchedulingCoordinator(scheduler_max_num_seqs=10, stage_id=1, async_chunk=True) req = _make_request("r1", status=RequestStatus.WAITING_FOR_CHUNK) waiting = MockQueue([req]) @@ -110,7 +156,7 @@ class TestChunkCoordinatorStateTransition(unittest.TestCase): self.assertIn("r1", coord.requests_with_ready_chunks) def test_non_ready_stays_waiting_for_chunk(self): - coord = ChunkSchedulingCoordinator(scheduler_max_num_seqs=10, stage_id=1, async_chunk=True) + coord = OmniSchedulingCoordinator(scheduler_max_num_seqs=10, stage_id=1, async_chunk=True) req = _make_request("r1", status=RequestStatus.WAITING_FOR_CHUNK) waiting = MockQueue([req]) @@ -126,7 +172,7 @@ class TestChunkCoordinatorStateTransition(unittest.TestCase): self.assertEqual(req.status, RequestStatus.WAITING_FOR_CHUNK) def test_stage_0_is_noop(self): - coord = ChunkSchedulingCoordinator(scheduler_max_num_seqs=10, stage_id=0) + coord = OmniSchedulingCoordinator(scheduler_max_num_seqs=10, stage_id=0) req = _make_request("r1") waiting = MockQueue([req]) running: list = [] @@ -144,7 +190,7 @@ class TestChunkCoordinatorRestoreQueues(unittest.TestCase): """Test 6: restore_queues returns waiting-for-chunk requests.""" def test_restore(self): - coord = ChunkSchedulingCoordinator(scheduler_max_num_seqs=10, stage_id=1) + coord = OmniSchedulingCoordinator(scheduler_max_num_seqs=10, stage_id=1) r1 = _make_request("r1") r2 = _make_request("r2") @@ -166,7 +212,7 @@ class TestChunkCoordinatorFinishedSignal(unittest.TestCase): """Test 8: chunk_finished_req_ids → finished_requests.""" def test_finished_signal(self): - coord = ChunkSchedulingCoordinator(scheduler_max_num_seqs=10, stage_id=1, async_chunk=True) + coord = OmniSchedulingCoordinator(scheduler_max_num_seqs=10, stage_id=1, async_chunk=True) req = _make_request("r1", status=RequestStatus.WAITING_FOR_CHUNK) waiting = MockQueue([req]) @@ -187,7 +233,7 @@ class TestChunkCoordinatorUpdateRequestMetadata(unittest.TestCase): def test_ar_mode_no_longer_sets_additional_information(self): """AR mode only processes scheduling metadata, not full payloads.""" - coord = ChunkSchedulingCoordinator(scheduler_max_num_seqs=10, stage_id=1) + coord = OmniSchedulingCoordinator(scheduler_max_num_seqs=10, stage_id=1) req = _make_request("r1") requests = {"r1": req} @@ -204,10 +250,14 @@ class TestChunkCoordinatorUpdateRequestMetadata(unittest.TestCase): self.assertIsNone(getattr(req, "additional_information", None)) def test_generation_mode(self): - coord = ChunkSchedulingCoordinator(scheduler_max_num_seqs=10, stage_id=1) + coord = OmniSchedulingCoordinator(scheduler_max_num_seqs=10, stage_id=1) req = _make_request("r1") req.prompt_token_ids = [0, 0, 0] + req.num_prompt_tokens = 3 + req.num_computed_tokens = 3 + req._all_token_ids = [0, 0, 0, 99] + req._output_token_ids = [99] requests = {"r1": req} request_metadata = { @@ -220,16 +270,61 @@ class TestChunkCoordinatorUpdateRequestMetadata(unittest.TestCase): coord.update_request_metadata(requests, request_metadata, model_mode="generation") self.assertEqual(req.prompt_token_ids, [10, 20, 30]) + self.assertEqual(req.num_prompt_tokens, 3) self.assertEqual(req.num_computed_tokens, 0) + self.assertEqual(req._all_token_ids, [10, 20, 30]) + self.assertEqual(req._output_token_ids, []) self.assertIsNone(req.additional_information) self.assertEqual(req._omni_initial_model_buffer, {"meta": {"left_context_size": 25}}) + def test_generation_mode_flattens_tensor_code_predictor_codes(self): + coord = OmniSchedulingCoordinator(scheduler_max_num_seqs=10, stage_id=1) + + req = _make_request("r1") + req.prompt_token_ids = [9] + req.num_prompt_tokens = 1 + req._all_token_ids = [9, 8] + req._output_token_ids = [8] + requests = {"r1": req} + + coord.update_request_metadata( + requests, + {"r1": {"code_predictor_codes": torch.tensor([[1, 2, 3]], dtype=torch.long)}}, + model_mode="generation", + ) + + self.assertEqual(req.prompt_token_ids, [1, 2, 3]) + self.assertEqual(req.num_prompt_tokens, 3) + self.assertEqual(req._all_token_ids, [1, 2, 3]) + self.assertEqual(req._output_token_ids, []) + + def test_generation_mode_flattens_nested_code_predictor_codes(self): + coord = OmniSchedulingCoordinator(scheduler_max_num_seqs=10, stage_id=1) + + req = _make_request("r1") + req.prompt_token_ids = [9] + req.num_prompt_tokens = 1 + req._all_token_ids = [9, 8] + req._output_token_ids = [8] + requests = {"r1": req} + + coord.update_request_metadata( + requests, + {"r1": {"code_predictor_codes": [[1, 2], [3, 4]]}}, + model_mode="generation", + ) + + self.assertEqual(req.prompt_token_ids, [1, 2, 3, 4]) + self.assertEqual(req.num_prompt_tokens, 4) + self.assertEqual(req._all_token_ids, [1, 2, 3, 4]) + self.assertEqual(req._output_token_ids, []) + class TestChunkCoordinatorPostprocess(unittest.TestCase): """Test postprocess_scheduler_output clears ready chunks.""" def test_clear_ready(self): - coord = ChunkSchedulingCoordinator(scheduler_max_num_seqs=10, stage_id=1) + coord = OmniSchedulingCoordinator(scheduler_max_num_seqs=10, stage_id=1) coord.requests_with_ready_chunks = {"r1", "r2"} new_req = SimpleNamespace(req_id="r1") @@ -366,6 +461,33 @@ class TestWaitingForInputTransition(unittest.TestCase): self.assertEqual(len(coord.pending_input_registrations), 1) self.assertEqual(coord.pending_input_registrations[0].request_id, "r1") + def test_idle_cycles_retain_received_marker_before_request_appears(self): + coord = OmniSchedulingCoordinator( + scheduler_max_num_seqs=10, + stage_id=1, + async_chunk=False, + ) + coord._full_payload_input_received.add("late") + coord.finished_requests.add("late") + + waiting = MockQueue() + running: list = [] + + coord.process_pending_full_payload_inputs(waiting, running, stage_recv_req_ids=set()) + + self.assertIn("late", coord._full_payload_input_received) + self.assertIn("late", coord.finished_requests) + + late_req = _make_request("late", status=RequestStatus.WAITING) + waiting.add_request(late_req) + + coord.process_pending_full_payload_inputs(waiting, running, stage_recv_req_ids=set()) + + self.assertEqual(late_req.status, RequestStatus.WAITING) + self.assertEqual(coord.pending_input_registrations, []) + self.assertIn("late", coord._full_payload_input_received) + self.assertIn("late", coord.finished_requests) + class TestTimeoutDetection(unittest.TestCase): """Regression tests for orphaned pending-recv timeout detection. @@ -508,14 +630,18 @@ class TestTimeoutDetection(unittest.TestCase): self.assertEqual(len(coord._waiting_for_input), 0) def test_free_finished_request_clears_waiting_since(self): - """free_finished_request clears _waiting_since.""" + """free_finished_request clears coordinator lifecycle markers.""" coord = OmniSchedulingCoordinator( scheduler_max_num_seqs=10, stage_id=1, ) coord._waiting_since["r1"] = 0.0 + coord._full_payload_input_received.add("r1") + coord.finished_requests.add("r1") coord.free_finished_request("r1") self.assertNotIn("r1", coord._waiting_since) + self.assertNotIn("r1", coord._full_payload_input_received) + self.assertNotIn("r1", coord.finished_requests) def test_timeout_from_running_queue_full_lifecycle(self): """End-to-end: request from running → WAITING_FOR_CHUNK → restore → diff --git a/tests/model_executor/stage_input_processors/test_qwen3_omni_streaming_helpers.py b/tests/model_executor/stage_input_processors/test_qwen3_omni_streaming_helpers.py index 18972c91d..f11a4654e 100644 --- a/tests/model_executor/stage_input_processors/test_qwen3_omni_streaming_helpers.py +++ b/tests/model_executor/stage_input_processors/test_qwen3_omni_streaming_helpers.py @@ -7,6 +7,7 @@ from __future__ import annotations from types import SimpleNamespace import pytest +import torch import vllm_omni.model_executor.stage_input_processors.qwen3_omni as q3 @@ -79,3 +80,110 @@ def test_get_streaming_codec_delta_len_increments_and_finishes(_streaming_contex assert d3 == 1 state = q3._get_qwen3_streaming_state("c1", _streaming_context) assert state.talker2code2wav_last_seq_len == 0 + + +def test_talker2code2wav_full_payload_filters_by_output_token_ids() -> None: + request = SimpleNamespace( + request_id="codec", + output_token_ids=[4197, 1, 2, 4198, -1, 2048], + ) + rows = torch.tensor( + [ + [100, 101, 102], + [10, 11, 12], + [20, 21, 22], + [30, 31, 32], + [40, 41, 42], + [50, 51, 52], + ], + dtype=torch.long, + ) + + payload = q3.talker2code2wav_full_payload(None, {"codes.audio": rows}, request) + + assert payload is not None + assert payload["codes"]["audio"] == [10, 20, 11, 21, 12, 22] + assert payload["code_predictor_codes"] == payload["codes"]["audio"] + + +def test_talker2code2wav_full_payload_drops_count_matched_terminal_row() -> None: + request = SimpleNamespace( + request_id="codec_terminal_row", + output_token_ids=[0, 4198], + ) + rows = torch.tensor( + [ + [10, 11, 12], + ], + dtype=torch.long, + ) + + payload = q3.talker2code2wav_full_payload(None, {"codes.audio": rows}, request) + + assert payload is None + + +def test_talker2code2wav_full_payload_drops_rows_aligned_to_non_codec_ids() -> None: + request = SimpleNamespace( + request_id="codec_invalid_ids", + output_token_ids=[4197, 0, 4198, 4196, -1, 2048], + ) + rows = torch.tensor( + [ + [91, 92, 93], + [0, 0, 0], + [81, 82, 83], + [71, 72, 73], + [61, 62, 63], + [51, 52, 53], + ], + dtype=torch.long, + ) + + payload = q3.talker2code2wav_full_payload(None, {"codes.audio": rows}, request) + + assert payload is not None + assert payload["codes"]["audio"] == [0, 0, 0] + assert payload["code_predictor_codes"] == payload["codes"]["audio"] + + +def test_talker2code2wav_full_payload_keeps_all_zero_codec_rows() -> None: + request = SimpleNamespace( + request_id="codec_zero", + output_token_ids=[0, 1], + ) + rows = torch.tensor( + [ + [0, 0, 0], + [7, 8, 9], + ], + dtype=torch.long, + ) + + payload = q3.talker2code2wav_full_payload(None, {"codes.audio": rows}, request) + + assert payload is not None + assert payload["codes"]["audio"] == [0, 7, 0, 8, 0, 9] + assert payload["code_predictor_codes"] == payload["codes"]["audio"] + + +def test_thinker2talker_full_payload_packs_complete_tensors() -> None: + request = SimpleNamespace( + request_id="thinker", + prompt_token_ids=[151644, 872], + output_token_ids=[3], + all_token_ids=[151644, 872, 3], + ) + pooling_output = { + "hidden_states.layer_0": torch.ones(3, 2), + "hidden_states.layer_24": torch.full((3, 2), 2.0), + "embed.tts_bos": torch.zeros(1, 2), + } + + payload = q3.thinker2talker_full_payload(None, pooling_output, request) + + assert payload is not None + assert payload["ids"]["all"] == [151644, 872, 3] + assert payload["embed"]["prefill"].device.type == "cpu" + assert payload["hidden_states"]["output"].device.type == "cpu" + assert payload["next_stage_prompt_len"] > 0 diff --git a/tests/test_config_factory.py b/tests/test_config_factory.py index 57313fe3b..37566088d 100644 --- a/tests/test_config_factory.py +++ b/tests/test_config_factory.py @@ -1504,7 +1504,7 @@ class TestSentinelDefaultPrecedence: def test_none_value_skipped_yaml_wins(self): stages = self._stages({"max_num_seqs": None}) assert stages[2].runtime_overrides.get("max_num_seqs") is None - assert "max_num_seqs" not in stages[2].yaml_engine_args + assert stages[2].yaml_engine_args["max_num_seqs"] == 64 def test_empty_kwargs_yaml_only(self): stages = self._stages({}) @@ -1598,6 +1598,41 @@ class TestSentinelDefaultPrecedence: assert sync_stages[1].custom_process_input_func is not None assert sync_stages[1].custom_process_input_func.endswith("talker2code2wav") + def test_async_chunk_dispatches_qwen3_omni_processors(self): + import runpy + from pathlib import Path + + from vllm_omni.config.stage_config import DeployConfig, merge_pipeline_deploy + + pipeline_path = ( + Path(__file__).parent.parent / "vllm_omni" / "model_executor" / "models" / "qwen3_omni" / "pipeline.py" + ) + pipeline = runpy.run_path(str(pipeline_path))["QWEN3_OMNI_PIPELINE"] + + async_stages = merge_pipeline_deploy(pipeline, DeployConfig(async_chunk=True)) + assert ( + async_stages[0] + .yaml_engine_args["custom_process_next_stage_input_func"] + .endswith("thinker2talker_async_chunk") + ) + assert ( + async_stages[1] + .yaml_engine_args["custom_process_next_stage_input_func"] + .endswith("talker2code2wav_async_chunk") + ) + + sync_stages = merge_pipeline_deploy(pipeline, DeployConfig(async_chunk=False)) + assert ( + sync_stages[0] + .yaml_engine_args["custom_process_next_stage_input_func"] + .endswith("thinker2talker_full_payload") + ) + assert ( + sync_stages[1] + .yaml_engine_args["custom_process_next_stage_input_func"] + .endswith("talker2code2wav_full_payload") + ) + class TestSamplingConstraintsPrecedence: """Test that pipeline sampling_constraints override deploy defaults.""" diff --git a/tests/worker/test_omni_gpu_model_runner.py b/tests/worker/test_omni_gpu_model_runner.py index c2f3f3622..683f2d0a5 100644 --- a/tests/worker/test_omni_gpu_model_runner.py +++ b/tests/worker/test_omni_gpu_model_runner.py @@ -5,6 +5,7 @@ import pytest import torch from vllm_omni.worker.gpu_model_runner import OmniGPUModelRunner +from vllm_omni.worker.omni_connector_model_runner_mixin import OmniConnectorModelRunnerMixin pytestmark = [pytest.mark.core_model, pytest.mark.cpu] @@ -330,6 +331,117 @@ def test_update_intermediate_buffer_skips_unknown_req_id(): assert "unknown_req" not in runner.model_intermediate_buffer +def _make_full_payload_accumulation_runner( + model_arch="Qwen3OmniMoeForConditionalGeneration", + model_stage="talker", + async_chunk=False, +): + runner = object.__new__(OmniConnectorModelRunnerMixin) + runner.model_config = SimpleNamespace( + model_arch=model_arch, + model_stage=model_stage, + async_chunk=async_chunk, + ) + runner._custom_process_func = object() + runner._pending_full_payload_send = {} + runner._stage_id = 1 + return runner + + +def test_accumulate_full_payload_output_preserves_aligned_all_zero_qwen3_omni_codec_rows(): + runner = _make_full_payload_accumulation_runner() + request = SimpleNamespace(output_token_ids=[0, 1]) + codes = torch.zeros((2, 3), dtype=torch.long) + + OmniConnectorModelRunnerMixin.accumulate_full_payload_output(runner, "r1", {"codes.audio": codes}, request) + + stored, _ = OmniConnectorModelRunnerMixin._materialize_full_payload_entry(runner._pending_full_payload_send["r1"]) + assert torch.equal(stored["codes.audio"], codes) + + +def test_accumulate_full_payload_output_keeps_misaligned_all_zero_qwen3_omni_codec_rows(): + # After removing the sender-side zero filter, the accumulator keeps every + # codec row including misaligned all-zero rows. The downstream consumer + # (_extract_qwen3_full_payload_codec_rows) is the authoritative crop and + # filters by output_token_ids. + runner = _make_full_payload_accumulation_runner() + request = SimpleNamespace(output_token_ids=[0, 1]) + codes = torch.zeros((1, 3), dtype=torch.long) + + OmniConnectorModelRunnerMixin.accumulate_full_payload_output(runner, "r1", {"codes.audio": codes}, request) + + stored, _ = OmniConnectorModelRunnerMixin._materialize_full_payload_entry(runner._pending_full_payload_send["r1"]) + assert "codes.audio" in stored + assert torch.equal(stored["codes.audio"], codes) + + +def test_accumulate_full_payload_output_preserves_incremental_aligned_all_zero_qwen3_omni_codec_rows(): + runner = _make_full_payload_accumulation_runner() + request = SimpleNamespace(output_token_ids=[0, 1]) + runner._pending_full_payload_send["r1"] = ( + {"codes.audio": torch.ones((1, 3), dtype=torch.long)}, + request, + ) + codes = torch.zeros((1, 3), dtype=torch.long) + + OmniConnectorModelRunnerMixin.accumulate_full_payload_output(runner, "r1", {"codes.audio": codes}, request) + + stored, _ = OmniConnectorModelRunnerMixin._materialize_full_payload_entry(runner._pending_full_payload_send["r1"]) + assert stored["codes.audio"].shape == (2, 3) + assert torch.equal(stored["codes.audio"][1], torch.zeros(3, dtype=torch.long)) + + +def test_accumulate_full_payload_output_keeps_all_zero_qwen3_omni_prefill_placeholder(): + # Prefill placeholder rows (output_token_ids empty) are no longer dropped + # at the sender. The consumer-side crop trims them off using + # output_token_ids, so the end-to-end semantics are unchanged. + runner = _make_full_payload_accumulation_runner() + request = SimpleNamespace(output_token_ids=[]) + codes = torch.zeros((2, 3), dtype=torch.long) + + OmniConnectorModelRunnerMixin.accumulate_full_payload_output(runner, "r1", {"codes.audio": codes}, request) + + stored, _ = OmniConnectorModelRunnerMixin._materialize_full_payload_entry(runner._pending_full_payload_send["r1"]) + assert "codes.audio" in stored + assert torch.equal(stored["codes.audio"], codes) + + +def test_full_payload_output_accumulation_hook_matrix(): + assert _make_full_payload_accumulation_runner(model_stage="thinker")._should_accumulate_full_payload_output() + assert _make_full_payload_accumulation_runner(model_stage="talker")._should_accumulate_full_payload_output() + assert not _make_full_payload_accumulation_runner(model_stage="code2wav")._should_accumulate_full_payload_output() + assert not _make_full_payload_accumulation_runner( + model_stage="talker", async_chunk=True + )._should_accumulate_full_payload_output() + assert not _make_full_payload_accumulation_runner( + model_arch="Qwen3TTSForConditionalGeneration" + )._should_accumulate_full_payload_output() + assert not _make_full_payload_accumulation_runner( + model_arch="Qwen2_5OmniForConditionalGeneration" + )._should_accumulate_full_payload_output() + + +def test_sync_local_stage_payloads_retains_payload_until_request_is_active(): + runner = object.__new__(OmniGPUModelRunner) + payload = {"codes": {"audio": [1, 2, 3]}} + runner._local_stage_payload_cache = {"late": payload} + runner._full_payload_pending_broadcast_req_ids = set() + runner.requests = {} + runner.model_intermediate_buffer = {} + + OmniGPUModelRunner._sync_local_stage_payloads(runner) + + assert runner._local_stage_payload_cache == {"late": payload} + assert runner.model_intermediate_buffer == {} + + runner.requests = {"late": DummyReqState()} + OmniGPUModelRunner._sync_local_stage_payloads(runner) + + assert runner._local_stage_payload_cache == {} + assert runner.model_intermediate_buffer["late"] == payload + assert runner.requests["late"].additional_information_cpu == payload + + def test_maybe_attach_mimo_audio_req_infos_enriches_dict(): runner = _make_runner_for_mimo() req_id = "r_mimo" diff --git a/vllm_omni/core/sched/omni_ar_scheduler.py b/vllm_omni/core/sched/omni_ar_scheduler.py index 297584812..0a7fd431c 100644 --- a/vllm_omni/core/sched/omni_ar_scheduler.py +++ b/vllm_omni/core/sched/omni_ar_scheduler.py @@ -20,11 +20,16 @@ from vllm.v1.request import Request, RequestStatus, StreamingUpdate from vllm.v1.spec_decode.metrics import SpecDecodingStats from vllm_omni.core.sched.omni_scheduler_mixin import OmniSchedulerMixin +from vllm_omni.core.sched.omni_scheduling_coordinator import ( + OmniSchedulingCoordinator, + uses_qwen3_omni_full_payload_input_coordinator, +) from vllm_omni.core.sched.output import OmniSchedulerOutput from vllm_omni.distributed.omni_connectors.transfer_adapter.chunk_transfer_adapter import ( OmniChunkTransferAdapter, ) from vllm_omni.engine.serialization import deserialize_additional_information +from vllm_omni.outputs import OmniConnectorOutput logger = init_logger(__name__) @@ -75,6 +80,14 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler): self.chunk_transfer_adapter = None if getattr(model_config, "async_chunk", False): self.chunk_transfer_adapter = OmniChunkTransferAdapter(self.vllm_config) + self.input_coordinator: OmniSchedulingCoordinator | None = None + if uses_qwen3_omni_full_payload_input_coordinator(model_config): + self.input_coordinator = OmniSchedulingCoordinator( + scheduler_max_num_seqs=self.vllm_config.scheduler_config.max_num_seqs, + stage_id=getattr(model_config, "stage_id", 0), + async_chunk=False, + ) + self._latest_omni_connector_output: OmniConnectorOutput | None = None # Snapshot prompt length for each streaming input update self._new_prompt_len_snapshot: dict[str, int] = {} @@ -190,6 +203,19 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler): return False def schedule(self) -> SchedulerOutput: # type: ignore[override] + connector_output = self._latest_omni_connector_output + self._latest_omni_connector_output = None + if self.input_coordinator: + if connector_output and connector_output.request_metadata: + self.input_coordinator.update_request_metadata( + self.requests, connector_output.request_metadata, model_mode="ar" + ) + self.input_coordinator.process_pending_full_payload_inputs( + self.waiting, + self.running, + connector_output.stage_recv_req_ids if connector_output else set(), + ) + if self.chunk_transfer_adapter: self.chunk_transfer_adapter.process_pending_chunks(self.waiting, self.running) @@ -199,6 +225,8 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler): if self.chunk_transfer_adapter: # Add request waiting for chunk to the waiting and running queue self.chunk_transfer_adapter.restore_queues(self.waiting, self.running) + if self.input_coordinator: + self.input_coordinator.restore_queues(self.waiting, self.running) try: # Late import to avoid circulars in some launch modes from .output import OmniNewRequestData @@ -239,9 +267,11 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler): # Wrap in omni scheduler output to carry transfer metadata. base_fields = SchedulerOutput.__dataclass_fields__.keys() base_data = {name: getattr(scheduler_output, name) for name in base_fields} + input_regs = self.input_coordinator.pending_input_registrations if self.input_coordinator else [] return OmniSchedulerOutput( **base_data, finished_requests_needing_kv_transfer=finished_reqs, + pending_input_registrations=input_regs, ) def update_from_output( @@ -534,6 +564,16 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler): engine_core_outputs[0] = eco = EngineCoreOutputs() eco.scheduler_stats = stats + omni_output = getattr(model_runner_output, "omni_connector_output", None) + if omni_output is not None: + self._latest_omni_connector_output = omni_output + if self.input_coordinator and omni_output.request_metadata: + self.input_coordinator.update_request_metadata( + self.requests, + omni_output.request_metadata, + model_mode="ar", + ) + # Free blocks that were held for transfer (kv_ready and # active_kv_transfers updates already done before the per-request loop). if kv_extracted_ids: @@ -572,7 +612,12 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler): if self.chunk_transfer_adapter: self.chunk_transfer_adapter.finish_requests(request_ids, finished_status, self.requests) - return super().finish_requests(request_ids, finished_status) + finished = super().finish_requests(request_ids, finished_status) + input_coordinator = getattr(self, "input_coordinator", None) + if input_coordinator is not None: + for request_id, _ in finished: + self._free_input_coordinator_request(request_id) + return finished def _update_request_as_session(self, session: Request, update: StreamingUpdate) -> None: """ @@ -616,10 +661,14 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler): # It triggered but hasn't finished yet. We MUST wait. logger.debug(f"[Omni] Request {request_id} finished but transfer is still ACTIVE. Waiting.") self.waiting_for_transfer_free.add(request_id) + if self.input_coordinator is not None: + self._free_input_coordinator_request(request_id) kv_xfer_params = None return kv_xfer_params elif request_id in self.waiting_for_transfer_free: # Blocks held until KV extraction completes in a future step. + if self.input_coordinator is not None: + self._free_input_coordinator_request(request_id) return None else: logger.debug( @@ -654,10 +703,14 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler): if isinstance(add_info, dict): add_info.update(kv_xfer_params) + if self.input_coordinator is not None: + self._free_input_coordinator_request(request_id) return kv_xfer_params # 3. Standard Freeing delay_free_blocks |= connector_delay_free_blocks + if self.input_coordinator is not None: + self._free_input_coordinator_request(request_id) if not delay_free_blocks: self._free_blocks(request) diff --git a/vllm_omni/core/sched/omni_generation_scheduler.py b/vllm_omni/core/sched/omni_generation_scheduler.py index f01672b12..5fc20d2da 100644 --- a/vllm_omni/core/sched/omni_generation_scheduler.py +++ b/vllm_omni/core/sched/omni_generation_scheduler.py @@ -2,6 +2,7 @@ from __future__ import annotations import time from collections import defaultdict +from typing import Any from vllm.compilation.cuda_graph import CUDAGraphStat from vllm.distributed.kv_events import KVEventBatch @@ -23,11 +24,15 @@ from vllm.v1.request import Request, RequestStatus, StreamingUpdate 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.omni_scheduling_coordinator import ( + OmniSchedulingCoordinator, + uses_qwen3_omni_full_payload_input_coordinator, +) +from vllm_omni.core.sched.output import OmniCachedRequestData, OmniNewRequestData, OmniSchedulerOutput from vllm_omni.distributed.omni_connectors.transfer_adapter.chunk_transfer_adapter import ( OmniChunkTransferAdapter, ) -from vllm_omni.outputs import OmniModelRunnerOutput +from vllm_omni.outputs import OmniConnectorOutput, OmniModelRunnerOutput logger = init_logger(__name__) @@ -39,6 +44,14 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler): self.chunk_transfer_adapter = None if getattr(model_config, "async_chunk", False): self.chunk_transfer_adapter = OmniChunkTransferAdapter(self.vllm_config) + self.input_coordinator: OmniSchedulingCoordinator | None = None + if uses_qwen3_omni_full_payload_input_coordinator(model_config): + self.input_coordinator = OmniSchedulingCoordinator( + scheduler_max_num_seqs=self.vllm_config.scheduler_config.max_num_seqs, + stage_id=getattr(model_config, "stage_id", 0), + async_chunk=False, + ) + self._latest_omni_connector_output: OmniConnectorOutput | None = None def schedule(self) -> SchedulerOutput: """Diffusion fast path: @@ -68,6 +81,18 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler): # Temporary queue: preserve waiting order, do not disturb non-diffusion requests skipped_waiting_requests = create_request_queue(self.policy) req_index = 0 + connector_output = self._latest_omni_connector_output + self._latest_omni_connector_output = None + if self.input_coordinator: + if connector_output and connector_output.request_metadata: + self.input_coordinator.update_request_metadata( + self.requests, connector_output.request_metadata, model_mode="generation" + ) + self.input_coordinator.process_pending_full_payload_inputs( + self.waiting, + self.running, + connector_output.stage_recv_req_ids if connector_output else set(), + ) if self.chunk_transfer_adapter: self.chunk_transfer_adapter.process_pending_chunks(self.waiting, self.running) @@ -201,7 +226,16 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler): self.chunk_transfer_adapter.restore_queues(self.waiting, self.running) else: res = super().schedule() - return res + if self.input_coordinator: + self.input_coordinator.restore_queues(self.waiting, self.running) + base_fields = SchedulerOutput.__dataclass_fields__.keys() + base_data = {name: getattr(res, name) for name in base_fields} + return OmniSchedulerOutput( + **base_data, + pending_input_registrations=( + self.input_coordinator.pending_input_registrations if self.input_coordinator else [] + ), + ) # Compute common prefix blocks (aligned with v1) num_common_prefix_blocks = [0] * len(self.kv_cache_config.kv_cache_groups) @@ -321,8 +355,17 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler): # internal deques and never scheduled again. if self.chunk_transfer_adapter: self.chunk_transfer_adapter.restore_queues(self.waiting, self.running) + if self.input_coordinator: + self.input_coordinator.restore_queues(self.waiting, self.running) - return scheduler_output + base_fields = SchedulerOutput.__dataclass_fields__.keys() + base_data = {name: getattr(scheduler_output, name) for name in base_fields} + return OmniSchedulerOutput( + **base_data, + pending_input_registrations=( + self.input_coordinator.pending_input_registrations if self.input_coordinator else [] + ), + ) def finish_requests(self, request_ids, finished_status: RequestStatus) -> list[tuple[str, int]]: """Handles the finish signal from outside the scheduler. @@ -340,7 +383,20 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler): if self.chunk_transfer_adapter: self.chunk_transfer_adapter.finish_requests(request_ids, finished_status, self.requests) - return super().finish_requests(request_ids, finished_status) + finished = super().finish_requests(request_ids, finished_status) + if self.input_coordinator is not None: + for request_id, _ in finished: + self._free_input_coordinator_request(request_id) + return finished + + def _free_request(self, request: Request, delay_free_blocks: bool = False) -> dict[str, Any] | None: + if self.input_coordinator is None: + return super()._free_request(request, delay_free_blocks) + + try: + return super()._free_request(request, delay_free_blocks) + finally: + self._free_input_coordinator_request(request.request_id) """ Scheduler for the diffusion model. @@ -589,6 +645,16 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler): engine_core_outputs[0] = eco = EngineCoreOutputs() eco.scheduler_stats = stats + omni_output = getattr(model_runner_output, "omni_connector_output", None) + if omni_output is not None: + self._latest_omni_connector_output = omni_output + if self.input_coordinator and omni_output.request_metadata: + self.input_coordinator.update_request_metadata( + self.requests, + omni_output.request_metadata, + model_mode="generation", + ) + return engine_core_outputs def _update_request_as_session(self, session: Request, update: StreamingUpdate) -> None: diff --git a/vllm_omni/core/sched/omni_scheduler_mixin.py b/vllm_omni/core/sched/omni_scheduler_mixin.py index 36080e63a..606739e90 100644 --- a/vllm_omni/core/sched/omni_scheduler_mixin.py +++ b/vllm_omni/core/sched/omni_scheduler_mixin.py @@ -7,6 +7,12 @@ from vllm.v1.request import Request, RequestStatus, StreamingUpdate class OmniSchedulerMixin: """Shared scheduler helpers for omni-specific request handling.""" + def _free_input_coordinator_request(self, request_id: str) -> None: + """Prune full-payload coordinator state for a completed request.""" + input_coordinator = getattr(self, "input_coordinator", None) + if input_coordinator is not None: + input_coordinator.free_finished_request(request_id) + def _replace_session_with_streaming_update( self, session: Request, diff --git a/vllm_omni/core/sched/omni_scheduling_coordinator.py b/vllm_omni/core/sched/omni_scheduling_coordinator.py index 3b6b4892e..6c32ed4cd 100644 --- a/vllm_omni/core/sched/omni_scheduling_coordinator.py +++ b/vllm_omni/core/sched/omni_scheduling_coordinator.py @@ -22,6 +22,15 @@ from vllm.v1.request import Request, RequestStatus logger = init_logger(__name__) +def uses_qwen3_omni_full_payload_input_coordinator(model_config: Any) -> bool: + return ( + getattr(model_config, "stage_id", 0) > 0 + and not getattr(model_config, "async_chunk", False) + and getattr(model_config, "model_arch", None) == "Qwen3OmniMoeForConditionalGeneration" + and getattr(model_config, "model_stage", None) in {"talker", "code2wav"} + ) + + class OmniSchedulingCoordinator: """Pure-scheduling coordinator for chunk and full_payload input waiting. @@ -166,8 +175,10 @@ class OmniSchedulingCoordinator: to_remove.append(request) self._waiting_for_input.append(request) self.pending_input_registrations.append(request) - for request in to_remove: - waiting_queue.remove(request) + if to_remove: + # Use the bulk-remove helper: one O(N) sweep instead of N + # repeated O(N) removes from a list-backed queue. + waiting_queue.remove_requests(to_remove) def process_pending_full_payload_inputs_legacy( self, @@ -255,6 +266,30 @@ class OmniSchedulingCoordinator: waiting_queue.add_request(request) self._waiting_for_input = deque() + @staticmethod + def _flatten_prompt_token_ids(value: Any) -> list[int]: + """Normalize connector metadata into flat prompt token ids.""" + if value is None: + return [] + if hasattr(value, "detach") and hasattr(value, "cpu") and hasattr(value, "tolist"): + value = value.detach().cpu().tolist() + elif hasattr(value, "tolist") and not isinstance(value, (list, tuple)): + value = value.tolist() + + if isinstance(value, (list, tuple)): + flattened: list[int] = [] + for item in value: + if hasattr(item, "detach") and hasattr(item, "cpu") and hasattr(item, "tolist"): + item = item.detach().cpu().tolist() + elif hasattr(item, "tolist") and not isinstance(item, (list, tuple)): + item = item.tolist() + if isinstance(item, (list, tuple)): + flattened.extend(int(token_id) for token_id in item) + else: + flattened.append(int(item)) + return flattened + return [int(value)] + def update_request_metadata( self, requests: dict[str, Request], @@ -310,7 +345,7 @@ class OmniSchedulingCoordinator: ) if model_mode != "ar": - new_ids = metadata.get("code_predictor_codes", []) + new_ids = self._flatten_prompt_token_ids(metadata.get("code_predictor_codes")) runtime_seed = None if "left_context_size" in metadata: runtime_seed = { @@ -319,6 +354,10 @@ class OmniSchedulingCoordinator: request._omni_initial_model_buffer = runtime_seed if new_ids: request.prompt_token_ids = new_ids + request.num_prompt_tokens = len(new_ids) + request._all_token_ids.clear() + request._all_token_ids.extend(new_ids) + request._output_token_ids.clear() request.num_computed_tokens = 0 def postprocess_scheduler_output( @@ -374,7 +413,3 @@ class OmniSchedulingCoordinator: if scheduler_output.scheduled_cached_reqs: for req_id in scheduler_output.scheduled_cached_reqs.req_ids: self.requests_with_ready_chunks.discard(req_id) - - -# Backward-compatible alias -ChunkSchedulingCoordinator = OmniSchedulingCoordinator diff --git a/vllm_omni/core/sched/output.py b/vllm_omni/core/sched/output.py index c7a8c07c5..c016cbd81 100644 --- a/vllm_omni/core/sched/output.py +++ b/vllm_omni/core/sched/output.py @@ -1,4 +1,5 @@ from dataclasses import dataclass, field +from typing import Any from vllm.v1.core.sched.output import CachedRequestData, NewRequestData, SchedulerOutput from vllm.v1.request import Request @@ -75,3 +76,4 @@ class OmniSchedulerOutput(SchedulerOutput): """Scheduler output with omni-specific transfer metadata.""" finished_requests_needing_kv_transfer: dict[str, dict] = field(default_factory=dict) + pending_input_registrations: list[Any] = field(default_factory=list) diff --git a/vllm_omni/distributed/omni_connectors/utils/initialization.py b/vllm_omni/distributed/omni_connectors/utils/initialization.py index f012af3c9..e6d68f9bc 100644 --- a/vllm_omni/distributed/omni_connectors/utils/initialization.py +++ b/vllm_omni/distributed/omni_connectors/utils/initialization.py @@ -5,6 +5,7 @@ import json import sys +import warnings from pathlib import Path from typing import TYPE_CHECKING, Any @@ -403,7 +404,16 @@ def build_stage_connectors( connectors_config: dict[str, Any], purpose: str = "request_forwarding", ) -> dict[tuple[str, str], Any] | None: - """Instantiate OmniConnectors for a stage based on config.""" + """Instantiate OmniConnectors for a stage based on config. + + Deprecated: prefer ``get_stage_connector_config`` plus the unified + connector factory. Kept as a thin shim so legacy callers keep working. + """ + warnings.warn( + "build_stage_connectors is deprecated; use get_stage_connector_config and OmniConnectorFactory instead.", + DeprecationWarning, + stacklevel=2, + ) if not connectors_config: return {} diff --git a/vllm_omni/entrypoints/openai/realtime_connection.py b/vllm_omni/entrypoints/openai/realtime_connection.py index 9fc2a1ee3..f2ce75df2 100644 --- a/vllm_omni/entrypoints/openai/realtime_connection.py +++ b/vllm_omni/entrypoints/openai/realtime_connection.py @@ -120,6 +120,7 @@ class RealtimeConnection(VllmRealtimeConnection): full_text = "" prompt_token_ids_len = 0 completion_tokens_len = 0 + last_prompt_token_ids_len = 0 # detect Stage-0 segment rollover self._realtime_audio_ref = None # Coerce cumulative outputs to delta outputs; this ensures @@ -144,6 +145,18 @@ class RealtimeConnection(VllmRealtimeConnection): new_token_ids = list(first_output.token_ids) new_tokens_len = len(new_token_ids) + cur_prompt_token_ids_len = len(output.prompt_token_ids or []) + # Stage-0 segment rollover: buffer_realtime_audio may yield + # multiple TokensPrompt segments and the second segment's + # prompt_token_ids include the first segment's decoded output. + # Clear accumulated full_text at the boundary so the + # carried-over prefix is not duplicated in the final + # TranscriptionDone.text. + stage_id = getattr(output, "stage_id", None) + if stage_id == 0 and cur_prompt_token_ids_len > last_prompt_token_ids_len > 0: + full_text = "" + last_prompt_token_ids_len = cur_prompt_token_ids_len + if not prompt_token_ids_len and output.prompt_token_ids: prompt_token_ids_len = len(output.prompt_token_ids) diff --git a/vllm_omni/model_executor/models/qwen3_omni/pipeline.py b/vllm_omni/model_executor/models/qwen3_omni/pipeline.py index 1c69ec795..1f3df4206 100644 --- a/vllm_omni/model_executor/models/qwen3_omni/pipeline.py +++ b/vllm_omni/model_executor/models/qwen3_omni/pipeline.py @@ -30,7 +30,8 @@ QWEN3_OMNI_PIPELINE = PipelineConfig( requires_multimodal_data=True, hf_config_name="thinker_config", engine_output_type="latent", - custom_process_next_stage_input_func=(f"{_PROC}.thinker2talker_async_chunk"), + custom_process_next_stage_input_func=(f"{_PROC}.thinker2talker_full_payload"), + async_chunk_process_next_stage_input_func=(f"{_PROC}.thinker2talker_async_chunk"), sampling_constraints={"detokenize": True}, ), StagePipelineConfig( @@ -41,7 +42,9 @@ QWEN3_OMNI_PIPELINE = PipelineConfig( hf_config_name="talker_config", engine_output_type="latent", custom_process_input_func=f"{_PROC}.thinker2talker", - custom_process_next_stage_input_func=(f"{_PROC}.talker2code2wav_async_chunk"), + sync_process_input_func=f"{_PROC}.thinker2talker_token_only", + custom_process_next_stage_input_func=(f"{_PROC}.talker2code2wav_full_payload"), + async_chunk_process_next_stage_input_func=(f"{_PROC}.talker2code2wav_async_chunk"), sampling_constraints={ "detokenize": False, "stop_token_ids": [2150], diff --git a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py index d77650265..0bdca2fd2 100644 --- a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py +++ b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py @@ -237,14 +237,22 @@ class Qwen3OmniMoeForConditionalGeneration( prompt_token_ids = tokenizer.encode(prompt_template) + # In non-async-chunk (full-payload) mode the engine treats each + # streaming TokensPrompt as a fresh decode, so mid-stream segment + # yields cause the thinker to emit duplicate responses and the + # talker to emit duplicate audio. Defer all audio to the final + # flush so the thinker sees one complete prompt. + async_chunk = getattr(model_config, "async_chunk", False) + async for audio_chunk in audio_stream: buffer.write_audio(audio_chunk) - while (segment := buffer.read_audio()) is not None: - yield TokensPrompt( - prompt_token_ids=prompt_token_ids, - multi_modal_data={"audio": segment}, - ) + if async_chunk: + while (segment := buffer.read_audio()) is not None: + yield TokensPrompt( + prompt_token_ids=prompt_token_ids, + multi_modal_data={"audio": segment}, + ) remaining = buffer.flush() if remaining is not None and len(remaining) > 0: diff --git a/vllm_omni/model_executor/stage_input_processors/qwen3_omni.py b/vllm_omni/model_executor/stage_input_processors/qwen3_omni.py index 63403619e..b1672612b 100644 --- a/vllm_omni/model_executor/stage_input_processors/qwen3_omni.py +++ b/vllm_omni/model_executor/stage_input_processors/qwen3_omni.py @@ -35,6 +35,10 @@ logger = logging.getLogger(__name__) # Pooling output layer keys: "0" = word embedding, "24" = accept_hidden_layer _EMBED_LAYER_KEY = "0" _HIDDEN_LAYER_KEY = "24" +_QWEN3_CODEC_CODEBOOK_SIZE = 2048 +_QWEN3_CODEC_PAD_TOKEN_ID = 4196 +_QWEN3_CODEC_BOS_TOKEN_ID = 4197 +_QWEN3_CODEC_EOS_TOKEN_ID = 4198 def _layer_tensor(layers: dict[Any, Any], key: str) -> torch.Tensor | None: @@ -99,6 +103,84 @@ def _ensure_list(x): return list(x) +def _as_tensor_or_none(value: Any) -> torch.Tensor | None: + if isinstance(value, torch.Tensor): + return value.detach().cpu() + if isinstance(value, list) and value and isinstance(value[0], torch.Tensor): + return value[0].detach().cpu() + return None + + +def _is_valid_qwen3_codec_token_id(token_id: Any) -> bool: + try: + token_id = int(token_id) + except (TypeError, ValueError): + return False + return 0 <= token_id < _QWEN3_CODEC_CODEBOOK_SIZE + + +def should_accumulate_qwen3_omni_full_payload_output( + model_config: Any, + custom_process_func: Any, +) -> bool: + """Return whether Qwen3-Omni should accumulate full-payload outputs.""" + return ( + custom_process_func is not None + and not getattr(model_config, "async_chunk", False) + and getattr(model_config, "model_arch", None) == "Qwen3OmniMoeForConditionalGeneration" + and getattr(model_config, "model_stage", None) in {"thinker", "talker"} + ) + + +def _extract_qwen3_full_payload_codec_rows( + code_predictor_codes: torch.Tensor, + output_token_ids: list[int], +) -> tuple[torch.Tensor, dict[str, int]]: + """Filter full-payload codec rows by the authoritative output ids.""" + if code_predictor_codes.ndim != 2 or code_predictor_codes.numel() == 0: + return code_predictor_codes, { + "raw_rows": int(code_predictor_codes.shape[0]) if code_predictor_codes.ndim > 0 else 0, + "aligned_rows": 0, + "valid_rows": 0, + "trailing_placeholder_count": 0, + } + + trailing_placeholder_count = 0 + while ( + trailing_placeholder_count < len(output_token_ids) and output_token_ids[-1 - trailing_placeholder_count] == -1 + ): + trailing_placeholder_count += 1 + + aligned_len = min(int(code_predictor_codes.shape[0]), len(output_token_ids)) + if aligned_len <= 0: + return code_predictor_codes[:0], { + "raw_rows": int(code_predictor_codes.shape[0]), + "aligned_rows": 0, + "valid_rows": 0, + "trailing_placeholder_count": trailing_placeholder_count, + } + + aligned_rows = code_predictor_codes[-aligned_len:] + aligned_token_ids = output_token_ids[-aligned_len:] + aligned_token_mask = torch.tensor( + [_is_valid_qwen3_codec_token_id(token_id) for token_id in aligned_token_ids], + dtype=torch.bool, + device=aligned_rows.device, + ) + row_valid_mask = (aligned_rows.max(dim=1).values < _QWEN3_CODEC_CODEBOOK_SIZE) & ( + aligned_rows.min(dim=1).values >= 0 + ) + filtered_rows = aligned_rows[aligned_token_mask & row_valid_mask] + if filtered_rows.numel() == 0: + filtered_rows = aligned_rows[:0] + return filtered_rows, { + "raw_rows": int(code_predictor_codes.shape[0]), + "aligned_rows": aligned_len, + "valid_rows": int(filtered_rows.shape[0]) if filtered_rows.ndim > 0 else 0, + "trailing_placeholder_count": trailing_placeholder_count, + } + + # ========================= # PD disaggregation helpers # ========================= @@ -373,6 +455,79 @@ def thinker2talker_async_chunk( return payload +def thinker2talker_full_payload( + transfer_manager: Any, + pooling_output: dict[str, Any], + request: OmniEngineCoreRequest, +) -> dict[str, Any] | None: + """Pack complete thinker output for the non-async connector path.""" + if not isinstance(pooling_output, dict): + return None + + layers = { + 0: pooling_output.get("hidden_states.layer_0"), + 24: pooling_output.get("hidden_states.layer_24"), + } + thinker_emb = _layer_tensor(layers, _EMBED_LAYER_KEY) + thinker_hid = _layer_tensor(layers, _HIDDEN_LAYER_KEY) + if thinker_emb is None: + hidden = pooling_output.get("hidden") + thinker_emb = hidden if isinstance(hidden, torch.Tensor) else None + if thinker_emb is None or thinker_hid is None: + logger.debug( + "thinker2talker_full_payload: missing thinker tensors for req=%s (embed=%s hidden=%s)", + getattr(request, "request_id", None), + thinker_emb is not None, + thinker_hid is not None, + ) + return None + + prompt_token_ids = _ensure_list(getattr(request, "prompt_token_ids", []) or []) + all_token_ids = _ensure_list(getattr(request, "all_token_ids", None) or []) + if not all_token_ids: + output_token_ids = _ensure_list(getattr(request, "output_token_ids", []) or []) + all_token_ids = list(prompt_token_ids) + list(output_token_ids) + + # Trim the trailing stop-token row from the accumulated thinker output. + # The accumulator captures one hidden-state row per executed thinker + # forward (prefill + every decode step including the one that emitted + # the stop_token), so for a finished request thinker_emb has exactly one + # row more than the rows the talker should consume. async_chunk's + # chunk-0 path naturally captures only the prefill / non-stop portion, + # which is why the [async_chunk] parametrization passes while [default] + # over-generates one codec frame on short outputs (e.g. + # test_one_word_prompt_001[default]: audio extends "London" with + # spurious phonemes). + if isinstance(thinker_emb, torch.Tensor) and thinker_emb.shape[0] > 0: + thinker_emb_prefill = thinker_emb[:-1] + else: + thinker_emb_prefill = thinker_emb + if isinstance(thinker_hid, torch.Tensor) and thinker_hid.shape[0] > 0: + thinker_hid_prefill = thinker_hid[:-1] + else: + thinker_hid_prefill = thinker_hid + + payload: OmniPayload = { + "embed": { + "prefill": thinker_emb_prefill.detach().cpu(), + "tts_bos": _as_tensor_or_none(pooling_output.get("embed.tts_bos")), + "tts_eos": _as_tensor_or_none(pooling_output.get("embed.tts_eos")), + "tts_pad": _as_tensor_or_none(pooling_output.get("embed.tts_pad")), + }, + "hidden_states": {"output": thinker_hid_prefill.detach().cpu()}, + "ids": {"all": list(all_token_ids), "prompt": list(prompt_token_ids)}, + "meta": {"finished": torch.tensor(True, dtype=torch.bool)}, + } + payload["next_stage_prompt_len"] = _compute_talker_prompt_ids_length(payload, device="cpu") + speaker = extract_speaker_from_request(request) + if speaker is not None: + payload["speaker"] = speaker + language = extract_language_from_request(request) + if language is not None: + payload["language"] = language + return payload + + def thinker2talker( source_outputs: list[Any], prompt: OmniTokensPrompt | TextPrompt | None = None, @@ -482,6 +637,85 @@ def thinker2talker( return talker_inputs +def thinker2talker_token_only( + source_outputs: list[Any], + prompt: OmniTokensPrompt | TextPrompt | None = None, + requires_multimodal_data: bool = False, + streaming_context: Any | None = None, +) -> list[OmniTokensPrompt]: + """Non-async-chunk Stage-1 input builder for the connector data plane. + + The worker connector (Stage-0 ``thinker2talker_full_payload`` → + ``_sync_local_stage_payloads``) supplies the bulk talker conditioning + tensors (embed / hidden_states / ids) via ``model_intermediate_buffer``. + The orchestrator only needs to ship a placeholder prefill prompt of the + correct length so the scheduler can allocate KV-cache slots. + + Small per-request voice metadata (``speaker`` / ``language``) is forwarded + here from the user prompt so the worker's line-408 buffer seed picks it + up. The connector-side ``extract_speaker_from_request`` reads the + strongly-typed ``request.additional_information.entries["speaker"]`` which + currently does not always round-trip the user-supplied voice; until that + plumbing is normalized, providing the small fields directly preserves + voice selection (regression discovered on Buildkite 9668: + ``test_speaker_002[default]`` lost the preset voice). + """ + talker_inputs: list[OmniTokensPrompt] = [] + for i, thinker_output in enumerate(source_outputs): + output = thinker_output.outputs[0] + req_id = str(getattr(thinker_output, "request_id", f"idx-{i}")) + # Skip-on-missing parity with thinker2talker_full_payload: if the + # connector builder would drop this request (no MM dict or missing + # hidden-state layers), do the same here so the worker buffer + # presence agrees with the orchestrator's scheduling decision. + thinker_mm_raw = getattr(output, "multimodal_output", None) + if not isinstance(thinker_mm_raw, dict): + logger.debug("thinker2talker_token_only: skip req=%s due to empty multimodal_output", req_id) + continue + mm_hs = thinker_mm_raw.get("hidden_states", {}) + mm_layers = mm_hs.get("layers", {}) if isinstance(mm_hs, dict) else {} + if _layer_tensor(mm_layers, _EMBED_LAYER_KEY) is None or _layer_tensor(mm_layers, _HIDDEN_LAYER_KEY) is None: + logger.debug("thinker2talker_token_only: skip req=%s due to missing hidden-state layers", req_id) + continue + prompt_token_ids = _ensure_list(thinker_output.prompt_token_ids) + output_ids = _ensure_list(output.cumulative_token_ids) + is_streaming_session = bool(getattr(streaming_context, "enabled", False)) + if is_streaming_session: + prompt_token_ids, output_ids, thinker_sequences, thinker_input_ids = _get_streaming_talker_tokens( + req_id, + prompt_token_ids, + output_ids, + getattr(streaming_context, "new_prompt_len_snapshot", None), + streaming_context, + clear_state=bool(getattr(thinker_output, "finished", False)), + ) + else: + thinker_sequences = prompt_token_ids + output_ids + thinker_input_ids = prompt_token_ids + info_for_len = {"ids": {"all": thinker_sequences, "prompt": thinker_input_ids}} + prompt_len = _compute_talker_prompt_ids_length(info_for_len, device="cpu") + + # Forward only small voice metadata; bulk tensors come from the + # connector path via _sync_local_stage_payloads. + small_info: dict[str, Any] = {} + speaker = extract_speaker_from_prompt(prompt, index=i) + if speaker is not None: + small_info["speaker"] = speaker + language = extract_language_from_prompt(prompt, index=i) + if language is not None: + small_info["language"] = language + + talker_inputs.append( + OmniTokensPrompt( + prompt_token_ids=[0] * prompt_len, + additional_information=(small_info if small_info else None), + multi_modal_data=None, + mm_processor_kwargs=None, + ) + ) + return talker_inputs + + # ========================= # Talker -> Code2Wav # ========================= @@ -556,6 +790,56 @@ def talker2code2wav_async_chunk( ) +def talker2code2wav_full_payload( + transfer_manager: Any, + pooling_output: dict[str, Any], + request: OmniEngineCoreRequest, +) -> dict[str, Any] | None: + """Pack complete talker codec output for the non-async connector path.""" + if not isinstance(pooling_output, dict): + return None + code_predictor_codes = pooling_output.get("codes.audio") + if code_predictor_codes is None: + codes = pooling_output.get("codes") + if isinstance(codes, dict): + code_predictor_codes = codes.get("audio") + if code_predictor_codes is None: + return None + if not isinstance(code_predictor_codes, torch.Tensor): + code_predictor_codes = torch.as_tensor(code_predictor_codes) + if code_predictor_codes.numel() == 0: + return None + + output_token_ids = _ensure_list(getattr(request, "output_token_ids", []) or []) + raw_shape = tuple(code_predictor_codes.shape) + code_predictor_codes, codec_stats = _extract_qwen3_full_payload_codec_rows( + code_predictor_codes.to(torch.long), + list(output_token_ids), + ) + if code_predictor_codes.numel() == 0: + return None + + codec_codes = code_predictor_codes.transpose(0, 1).cpu().reshape(-1).tolist() + logger.debug( + "talker2code2wav_full_payload: raw_shape=%s output_ids_len=%s aligned_rows=%s " + "valid_rows=%s placeholders=%s flattened_len=%s pad4196=%s bos4197=%s eos4198=%s", + raw_shape, + len(output_token_ids), + codec_stats["aligned_rows"], + codec_stats["valid_rows"], + codec_stats["trailing_placeholder_count"], + len(codec_codes), + sum(1 for tid in output_token_ids if tid == _QWEN3_CODEC_PAD_TOKEN_ID), + sum(1 for tid in output_token_ids if tid == _QWEN3_CODEC_BOS_TOKEN_ID), + sum(1 for tid in output_token_ids if tid == _QWEN3_CODEC_EOS_TOKEN_ID), + ) + return { + "codes": {"audio": codec_codes}, + "code_predictor_codes": codec_codes, + "meta": {"finished": torch.tensor(True, dtype=torch.bool)}, + } + + def talker2code2wav( source_outputs: list[Any], _prompt: OmniTokensPrompt | TextPrompt | None = None, diff --git a/vllm_omni/patch.py b/vllm_omni/patch.py index 74a3d8a76..b78323c21 100644 --- a/vllm_omni/patch.py +++ b/vllm_omni/patch.py @@ -102,6 +102,11 @@ if not hasattr(RequestStatus, "WAITING_FOR_CHUNK"): # as a non-finished state and remains compatible with existing comparisons. extend_enum(RequestStatus, "WAITING_FOR_CHUNK", -1) +if not hasattr(RequestStatus, "WAITING_FOR_INPUT"): + # Full-payload stage handoff uses a distinct waiting state so the + # scheduler can restore the request once non-stage-0 inputs arrive. + extend_enum(RequestStatus, "WAITING_FOR_INPUT", -2) + # Snapshot sys.modules: `hasattr` below can trigger lazy submodule imports # (e.g. transformers' `_LazyModule.__getattr__`), which mutate sys.modules # during iteration and raise `dictionary changed size during iteration`. diff --git a/vllm_omni/worker/gpu_ar_model_runner.py b/vllm_omni/worker/gpu_ar_model_runner.py index f87529106..b36072bec 100644 --- a/vllm_omni/worker/gpu_ar_model_runner.py +++ b/vllm_omni/worker/gpu_ar_model_runner.py @@ -81,6 +81,16 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin): self.inputs_embeds = self._make_buffer(self.max_num_tokens, self.hidden_size, dtype=self.dtype, numpy=False) # Initialize KV cache manager (preserve vllm_config fallback behavior) self.kv_transfer_manager = OmniKVTransferManager.from_vllm_config(self.vllm_config, self.model_config) + # Only Qwen3-Omni currently consumes the connector-based full-payload + # handoff added in this PR. Other model architectures (e.g. Bagel + # diffusion) retain their pre-existing runner behavior so this PR + # does not perturb them. + if getattr(self.model_config, "model_arch", None) == "Qwen3OmniMoeForConditionalGeneration": + self.init_omni_connectors( + vllm_config=self.vllm_config, + model_config=self.model_config, + kv_transfer_manager=self.kv_transfer_manager, + ) self._downstream_payload_cache: dict[str, bool] = {} def _make_buffer(self, *size, dtype, numpy=True): @@ -337,6 +347,16 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin): request_id_resolver=self._resolve_global_request_id, ) + if hasattr(self, "_omni_connector"): + for request in getattr(scheduler_output, "pending_input_registrations", []): + self.register_chunk_recv(request) + self.recv_full_payload_inputs(scheduler_output) + if self._pending_full_payload_send: + flush_ids = set(getattr(scheduler_output, "finished_req_ids", set())) + flush_ids.update({rid for rid in self._pending_full_payload_send if rid not in self.requests}) + if flush_ids: + self.flush_full_payload_outputs(flush_ids) + if self.routed_experts_initialized: capturer = RoutedExpertsCapturer.get_instance() if capturer is not None: @@ -375,7 +395,7 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin): if kv_ids: output = copy(output) output.kv_extracted_req_ids = kv_ids - return output + return self.attach_omni_connector_output(output) if not num_scheduled_tokens: if ( @@ -399,7 +419,7 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin): output = copy(output) output.kv_extracted_req_ids = kv_ids - return output + return self.attach_omni_connector_output(output) if self.cache_config.kv_sharing_fast_prefill: assert not self.num_prompt_logprobs, ( @@ -735,11 +755,11 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin): # In case of PP with kv transfer, we need to pass through the # kv_connector_output if kv_connector_output.is_empty(): - return EMPTY_MODEL_RUNNER_OUTPUT + return self.attach_omni_connector_output(EMPTY_MODEL_RUNNER_OUTPUT) output = copy(EMPTY_MODEL_RUNNER_OUTPUT) output.kv_connector_output = kv_connector_output - return output + return self.attach_omni_connector_output(output) # Unpack ephemeral state. ( @@ -985,6 +1005,13 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin): # Flatten nested dicts to dotted keys so pooling_output # stays dict[str, torch.Tensor] for msgspec serialization. pooler_output.append(flatten_payload(payload)) + + if pooler_output and self._should_accumulate_full_payload_output(): + for i, rid in enumerate(req_ids_output_copy): + req_state = self.requests.get(rid) + if req_state is not None and pooler_output[i]: + self.accumulate_full_payload_output(rid, pooler_output[i], req_state) + with record_function_or_nullcontext("gpu_model_runner: ModelRunnerOutput"): if self.routed_experts_initialized: capturer = RoutedExpertsCapturer.get_instance() @@ -1005,6 +1032,7 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin): cudagraph_stats=cudagraph_stats, ) output.kv_extracted_req_ids = kv_extracted_req_ids + output.omni_connector_output = self.get_omni_connector_output() if not self.use_async_scheduling: return output diff --git a/vllm_omni/worker/gpu_generation_model_runner.py b/vllm_omni/worker/gpu_generation_model_runner.py index 96db0bd74..3a7411459 100644 --- a/vllm_omni/worker/gpu_generation_model_runner.py +++ b/vllm_omni/worker/gpu_generation_model_runner.py @@ -53,6 +53,17 @@ class GPUGenerationModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin - Executes generation process and returns tensors via `pooler_output`. """ + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + # Scope full-payload connector init to Qwen3-Omni: other generation + # models (e.g. Bagel DiT) retain their pre-existing runner setup + # so this refactor does not perturb them. + if getattr(self.model_config, "model_arch", None) == "Qwen3OmniMoeForConditionalGeneration": + self.init_omni_connectors( + vllm_config=self.vllm_config, + model_config=self.model_config, + ) + def _update_request_states(self, scheduler_output: SchedulerOutput): # remove requests for req_id in scheduler_output.finished_req_ids: @@ -89,6 +100,16 @@ 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 hasattr(self, "_omni_connector"): + for request in getattr(scheduler_output, "pending_input_registrations", []): + self.register_chunk_recv(request) + self.recv_full_payload_inputs(scheduler_output) + if self._pending_full_payload_send: + flush_ids = set(getattr(scheduler_output, "finished_req_ids", set())) + flush_ids.update({rid for rid in self._pending_full_payload_send if rid not in self.requests}) + if flush_ids: + self.flush_full_payload_outputs(flush_ids) + if self.routed_experts_initialized: capturer = RoutedExpertsCapturer.get_instance() if capturer is not None: @@ -110,7 +131,7 @@ class GPUGenerationModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin self._update_request_states(scheduler_output) deferred_state_corrections_fn = self._update_states(scheduler_output) if not scheduler_output.total_num_scheduled_tokens: - return EMPTY_MODEL_RUNNER_OUTPUT + return self.attach_omni_connector_output(EMPTY_MODEL_RUNNER_OUTPUT) if has_ec_transfer() and not get_ec_transfer().is_consumer: with self.maybe_get_ec_connector_output( @@ -118,7 +139,7 @@ class GPUGenerationModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin encoder_cache=self.encoder_cache, ) as ec_connector_output: self._execute_mm_encoder(scheduler_output) - return make_empty_encoder_model_runner_output(scheduler_output) + return self.attach_omni_connector_output(make_empty_encoder_model_runner_output(scheduler_output)) if not num_scheduled_tokens: if ( @@ -134,9 +155,10 @@ class GPUGenerationModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin self._dummy_run(1) if not has_kv_transfer_group(): # Return empty ModelRunnerOutput if no work to do. - return EMPTY_MODEL_RUNNER_OUTPUT + return self.attach_omni_connector_output(EMPTY_MODEL_RUNNER_OUTPUT) - return self.kv_connector_no_forward(scheduler_output, self.vllm_config) + result = self.kv_connector_no_forward(scheduler_output, self.vllm_config) + return self.attach_omni_connector_output(result) if self.cache_config.kv_sharing_fast_prefill: assert not self.num_prompt_logprobs, ( @@ -338,11 +360,11 @@ class GPUGenerationModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin # In case of PP with kv transfer, we need to pass through the # kv_connector_output if kv_connector_output.is_empty(): - return EMPTY_MODEL_RUNNER_OUTPUT + return self.attach_omni_connector_output(EMPTY_MODEL_RUNNER_OUTPUT) output = copy(EMPTY_MODEL_RUNNER_OUTPUT) output.kv_connector_output = kv_connector_output - return output + return self.attach_omni_connector_output(output) # Unpack ephemeral state. ( @@ -403,6 +425,12 @@ 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() + if self._should_accumulate_full_payload_output(): + for i, rid in enumerate(req_ids_output_copy): + req_state = self.requests.get(rid) + if req_state is not None and pooler_output[i]: + self.accumulate_full_payload_output(rid, pooler_output[i], req_state) + output = OmniModelRunnerOutput( req_ids=req_ids_output_copy, req_id_to_index=req_id_to_index_output_copy, @@ -415,6 +443,7 @@ class GPUGenerationModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin cudagraph_stats=cudagraph_stats, ec_connector_output=ec_connector_output if self.supports_mm_inputs else None, ) + output.omni_connector_output = self.get_omni_connector_output() 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 f4eadedc0..b3f746eb4 100644 --- a/vllm_omni/worker/gpu_model_runner.py +++ b/vllm_omni/worker/gpu_model_runner.py @@ -1,3 +1,4 @@ +import contextlib from typing import TYPE_CHECKING, Any, cast import numpy as np @@ -264,6 +265,14 @@ class OmniGPUModelRunner(GPUModelRunner): self.omni_prefix_cache.reset_prefix_cached_new_req_ids() # Remove finished requests from the cached states. + # cleanup_finished_request lives on OmniConnectorModelRunnerMixin and + # is only safe to call once init_omni_connectors() has populated the + # mixin state. Archs that inherit the method via MRO without running + # that init must be skipped, so probe a mixin-owned attribute as the + # "state initialized" gate. + cleanup_finished_request = ( + getattr(self, "cleanup_finished_request", None) if hasattr(self, "_request_ids_mapping") else None + ) for req_id in scheduler_output.finished_req_ids: self.requests.pop(req_id, None) self.model_intermediate_buffer.pop(req_id, None) @@ -272,6 +281,8 @@ class OmniGPUModelRunner(GPUModelRunner): self._downstream_payload_cache.pop(req_id, None) if hasattr(self, "_talker_mtp_generators"): self._talker_mtp_generators.pop(req_id, None) + if cleanup_finished_request is not None: + cleanup_finished_request(req_id) if hasattr(self, "late_interaction_runner"): self.late_interaction_runner.on_requests_finished(scheduler_output.finished_req_ids) @@ -1033,8 +1044,31 @@ class OmniGPUModelRunner(GPUModelRunner): req_token_spans.append((start_offset, start_offset + sched_tokens)) return req_token_spans + def _sync_local_stage_payloads(self) -> None: + """Move received full-payload stage inputs into model_intermediate_buffer.""" + cache = getattr(self, "_local_stage_payload_cache", None) + if not cache: + return + lock = getattr(self, "_lock", None) + ctx = lock if lock is not None else contextlib.nullcontext() + with ctx: + if not cache: + return + active_req_ids = set(getattr(self, "requests", {})) + pending = set(getattr(self, "_full_payload_pending_broadcast_req_ids", set())) + staged = { + req_id: payload + for req_id, payload in cache.items() + if req_id not in pending and req_id in active_req_ids and isinstance(payload, dict) + } + for req_id in staged: + cache.pop(req_id, None) + for req_id, payload in staged.items(): + self._update_intermediate_buffer(req_id, payload) + def _build_model_kwargs_extra(self) -> dict: """Build extra keyword arguments passed to the model for this step.""" + self._sync_local_stage_payloads() model_kwargs_extra: dict[str, object] = {} try: buffer_map = self._gather_runtime_additional_information() @@ -1299,6 +1333,13 @@ class OmniGPUModelRunner(GPUModelRunner): if hasattr(self.model, "has_preprocess") or hasattr(self.model, "enable_update_additional_information"): if self.vllm_config.model_config.async_chunk: self._update_additional_information(scheduler_output) + else: + # In full-payload (non-async-chunk) mode, connector-delivered + # stage payloads must override any earlier engine-level + # additional_information written by the legacy + # custom_process_input_func codec, so talker_preprocess reads + # the full thinker payload. + self._sync_local_stage_payloads() if hasattr(self.model, "has_preprocess") and self.model.has_preprocess: # Overlay custom prompt_embeds per request for the prompt portion; diff --git a/vllm_omni/worker/omni_connector_model_runner_mixin.py b/vllm_omni/worker/omni_connector_model_runner_mixin.py index 6955aebd8..eb5b47b53 100644 --- a/vllm_omni/worker/omni_connector_model_runner_mixin.py +++ b/vllm_omni/worker/omni_connector_model_runner_mixin.py @@ -153,7 +153,7 @@ class OmniConnectorModelRunnerMixin: # -- full_payload_mode: accumulate latest pooler_output per request, # send only when the request finishes (next-cycle flush) -- - self._pending_full_payload_send: dict[str, tuple[Any, Any]] = {} + self._pending_full_payload_send: dict[str, tuple[Any, ...]] = {} # -- KV sent accumulator -- self._kv_sent_req_ids: list[str] = [] @@ -206,21 +206,33 @@ class OmniConnectorModelRunnerMixin: """Clean up per-request state after a request is fully finished. Call this when a request is freed from the model runner to prevent - memory leaks in the mixin's tracking dicts/sets. The external - request ID is resolved before cleaning up ``_put_req_chunk`` which - is keyed by external ID. + memory leaks in the mixin's tracking dicts/sets. + + Two senders use different keys: ``send_chunk`` keys per-request + state under the EXTERNAL id (after mapping resolution), while + ``send_full_payload_outputs`` keys under the INTERNAL id. To cover + both modes (and forward compat with id-rename scenarios) we attempt + cleanup against both keys; the entry that doesn't exist for the + active mode is a no-op pop. Only the key that actually has pending + saves is added to ``_deferred_send_cleanup`` so the bg save's + decrement path drains it without leaving orphans. """ ext_id = self._request_ids_mapping.pop(req_id, None) - send_req_id = ext_id if ext_id is not None else req_id + keys_to_clean: list[str] = [req_id] + if ext_id is not None and ext_id != req_id: + keys_to_clean.append(ext_id) with self._lock: - if self._pending_save_counts.get(send_req_id, 0): - self._deferred_send_cleanup.add(send_req_id) - else: - self._put_req_chunk.pop(send_req_id, None) - self._send_side_request_payload.pop(send_req_id, None) - self._code_prompt_token_ids.pop(send_req_id, None) - self._cached_ic.pop(send_req_id, None) + keys_pending = [k for k in keys_to_clean if self._pending_save_counts.get(k, 0)] + for k in keys_pending: + self._deferred_send_cleanup.add(k) + for k in keys_to_clean: + if k in keys_pending: + continue + self._put_req_chunk.pop(k, None) + self._send_side_request_payload.pop(k, None) + self._code_prompt_token_ids.pop(k, None) + self._cached_ic.pop(k, None) self._kv_pending_transfers.pop(req_id, None) self._kv_active_transfers.discard(req_id) self._kv_completed_transfers.discard(req_id) @@ -628,6 +640,16 @@ class OmniConnectorModelRunnerMixin: that has arrived, or ``None`` if nothing is ready. Stores full payloads in the local cache and extracts scheduling metadata. """ + # Fast path: when TP is trivial (no peer ranks waiting on a broadcast) + # and the bg recv thread has not staged anything, skip the lock + TP + # broadcast cycle entirely. _broadcast_tp_payload_packet already + # returns its input unchanged under the same world_size<=1 condition, + # so the original code path was a no-op here on every empty step. + tp_group = self._get_local_tp_group() + if ( + tp_group is None or getattr(tp_group, "world_size", 1) <= 1 + ) and not self._full_payload_pending_broadcast_req_ids: + return None with self._lock: results = self._collect_full_payload_results_locked() if self.is_data_transfer_rank() else None results = self._broadcast_tp_payload_packet(results) @@ -649,10 +671,63 @@ class OmniConnectorModelRunnerMixin: ) return results + def _get_model_config(self) -> Any: + model_config = getattr(self, "model_config", None) + if model_config is not None: + return model_config + return getattr(getattr(self, "vllm_config", None), "model_config", None) + + def _should_accumulate_full_payload_output(self) -> bool: + """Gate send-side full-payload output accumulation only. + + Cached per instance: the result depends only on model_config / + _custom_process_func, both of which are set at init time. Avoid + the per-step dynamic import inside the model decode loop. + """ + cached = getattr(self, "_should_accumulate_full_payload_output_cached", None) + if cached is not None: + return cached + model_config = self._get_model_config() + if model_config is None: + self._should_accumulate_full_payload_output_cached = False + return False + if getattr(model_config, "model_arch", None) == "Qwen3OmniMoeForConditionalGeneration": + from vllm_omni.model_executor.stage_input_processors.qwen3_omni import ( + should_accumulate_qwen3_omni_full_payload_output, + ) + + result = should_accumulate_qwen3_omni_full_payload_output( + model_config, + getattr(self, "_custom_process_func", None), + ) + self._should_accumulate_full_payload_output_cached = result + return result + self._should_accumulate_full_payload_output_cached = False + return False + @staticmethod - def _is_all_zero_tensor(t: Any) -> bool: - """Return True if *t* is a torch.Tensor whose elements are all zero.""" - return isinstance(t, torch.Tensor) and t.numel() > 0 and not t.any() + def _new_full_payload_accumulator(output: dict[str, Any]): + chunks: dict[str, list[torch.Tensor]] = {} + latest: dict[str, Any] = {} + rows: dict[str, int] = {} + for k, v in output.items(): + if isinstance(v, torch.Tensor) and v.dim() >= 2: + chunks[k] = [v] + rows[k] = int(v.shape[0]) + else: + latest[k] = v + return chunks, latest, rows + + @staticmethod + def _materialize_full_payload_entry(entry): + if len(entry) == 2: + return entry + chunks, latest, _rows, request = entry + output = dict(latest) + for k, tensors in chunks.items(): + if tensors: + output[k] = tensors[0] if len(tensors) == 1 else torch.cat(tensors, dim=0) + return output, request def accumulate_full_payload_output( self, @@ -666,55 +741,45 @@ class OmniConnectorModelRunnerMixin: along dim-0. Scalar / global tensors (1-D or 0-D) are replaced with the latest value. - All-zero tensors (e.g. ``code_predictor_codes`` emitted during - prefill) are dropped so that they do not pollute downstream stages - with garbage / noise frames. + Note: codec rows are NOT filtered for zero placeholders here. The + downstream consumer ``_extract_qwen3_full_payload_codec_rows`` crops + codec rows using ``output_token_ids`` as the authoritative source, + which makes any sender-side zero filtering redundant. Skipping the + sender-side ``t.any()`` scan also avoids a per-tensor GPU->CPU device + sync that stalled the decode pipeline. The data is actually sent when ``flush_full_payload_outputs`` is called with the finished request IDs from the next scheduler cycle. """ - # ---- Filter out all-zero tensors from the incoming pooler_output ---- - filtered: dict[str, Any] = {} - dropped_zero_keys: list[tuple[str, tuple[int, ...]]] = [] - for k, v in pooler_output.items(): - if self._is_all_zero_tensor(v): - dropped_zero_keys.append((k, tuple(v.shape))) - continue # skip prefill zero-filled placeholders - filtered[k] = v - if dropped_zero_keys: - logger.info( - "[Stage-%s] accumulate_full_payload_output: req=%s dropped_zero_keys=%s", - self._stage_id, - req_id, - dropped_zero_keys, - ) - pooler_output = filtered - existing = self._pending_full_payload_send.get(req_id) + if existing is None: - self._pending_full_payload_send[req_id] = (pooler_output, request) + chunks, latest, rows = self._new_full_payload_accumulator(pooler_output) + self._pending_full_payload_send[req_id] = (chunks, latest, rows, request) return - prev_output, _ = existing - merged: dict[str, Any] = {} - for k in set(prev_output) | set(pooler_output): - v_new = pooler_output.get(k) - v_old = prev_output.get(k) - if v_new is None: - merged[k] = v_old - elif v_old is None: - merged[k] = v_new - elif ( - isinstance(v_new, torch.Tensor) - and isinstance(v_old, torch.Tensor) - and v_new.dim() >= 2 - and v_old.dim() >= 2 - and v_new.shape[1:] == v_old.shape[1:] - ): - merged[k] = torch.cat([v_old, v_new], dim=0) + if len(existing) == 2: + chunks, latest, rows = self._new_full_payload_accumulator(existing[0]) + else: + chunks, latest, rows, _ = existing + + for k, v in pooler_output.items(): + if v is None: + continue + if isinstance(v, torch.Tensor) and v.dim() >= 2: + if k in chunks and chunks[k] and v.shape[1:] == chunks[k][0].shape[1:]: + chunks[k].append(v) + rows[k] += int(v.shape[0]) + else: + latest.pop(k, None) + chunks[k] = [v] + rows[k] = int(v.shape[0]) else: - merged[k] = v_new - self._pending_full_payload_send[req_id] = (merged, request) + chunks.pop(k, None) + rows.pop(k, None) + latest[k] = v + + self._pending_full_payload_send[req_id] = (chunks, latest, rows, request) def flush_full_payload_outputs(self, finished_req_ids: set[str]) -> None: """Send accumulated full_payload outputs for requests that just finished.""" @@ -728,7 +793,7 @@ class OmniConnectorModelRunnerMixin: for req_id in finished_req_ids: entry = self._pending_full_payload_send.pop(req_id, None) if entry is not None: - to_send[req_id] = entry + to_send[req_id] = self._materialize_full_payload_entry(entry) logger.info("[Stage-%s] flush_full_payload_outputs: to_send=%s", self._stage_id, list(to_send.keys())) if to_send: self.send_full_payload_outputs(scheduler_output=None, outputs=to_send) @@ -1476,7 +1541,7 @@ class OmniConnectorModelRunnerMixin: logger.warning("Error receiving data for %s", req_id, exc_info=True) if not made_progress and not self._stop_event.is_set(): - self._work_available.wait(timeout=0.001) + self._work_available.wait(timeout=0.005) self._work_available.clear() _MAX_SEND_RETRIES = 3