diff --git a/docs/configuration/stage_configs.md b/docs/configuration/stage_configs.md index 305c68360..1085b4e64 100644 --- a/docs/configuration/stage_configs.md +++ b/docs/configuration/stage_configs.md @@ -80,6 +80,19 @@ Note: for the diffusion path, an omitted `distributed_executor_backend` selects segments) and `mp` when `num_gpus > 1`. Set `mp` explicitly to keep a worker subprocess on one GPU. `ray` / `external_launcher` are not fully supported yet. +### Stage-level runner selection + +`model_runner: v1` or `v2` at the deploy level sets the default runner. +A `model_runner` on an individual stage overrides that default, allowing +one stage to migrate or roll back independently. This does not change +`PipelineConfig` topology, input adapters, or the connector payload contract. +An omitted stage override preserves the deploy-level selection. + +MRv2 native downstream receivers currently support turn-based requests only; +streaming sessions and resumable input require the V1 prompt-replacement path. +Selecting a runner does not add the session capabilities it lacks. Platform +fallbacks and stage overrides are validated after configuration resolution. + ### Stage fields Each entry under `stages:` accepts any `StageDeployConfig` field directly (no nested `engine_args:`). Only fields whose value legitimately varies across stages live here; pipeline-wide settings (trust_remote_code, distributed_executor_backend, dtype, quantization, prefix/chunked prefill, DP/PP sizes) are declared at the top level and applied to every stage. Unknown keys fall through to `engine_extras:` and are forwarded to the engine. Frequently used fields are listed below; the source-of-truth schema is `StageDeployConfig` in `vllm_omni/config/stage_config.py`. diff --git a/tests/config/test_environment_variables.py b/tests/config/test_environment_variables.py index e0434a358..14613b1cb 100644 --- a/tests/config/test_environment_variables.py +++ b/tests/config/test_environment_variables.py @@ -169,7 +169,7 @@ def test_inventory_matches_reviewed_snapshot_counts(): EnvironmentVariableCategory.PLATFORM_EXTERNAL: 29, EnvironmentVariableCategory.MODEL_SPECIFIC: 66, EnvironmentVariableCategory.BENCHMARK_TRANSITIONAL: 21, - EnvironmentVariableCategory.INTERNAL: 2, + EnvironmentVariableCategory.INTERNAL: 6, } disposition_counts = Counter( diff --git a/tests/config/test_omni_config.py b/tests/config/test_omni_config.py index da69ec5f1..822cd2e51 100644 --- a/tests/config/test_omni_config.py +++ b/tests/config/test_omni_config.py @@ -822,6 +822,62 @@ def test_mrv2_fails_fast_on_platforms_without_native_workers(platform: str): _apply_platform_overrides(DeployConfig(model_runner="v2"), platform=platform) +def _qwen3_omni_pipeline() -> PipelineConfig: + from dataclasses import replace + + pipeline = _resolve_pipeline_or_skip("qwen3_omni_moe", Qwen3OmniMoeConfig()) + return replace(pipeline, stages=tuple(replace(s, supports_native_mrv2_data_plane=True) for s in pipeline.stages)) + + +def test_stage_model_runner_overrides_deploy_runner(tmp_path: Path): + deploy_path = tmp_path / "qwen3_omni_mixed_runner.yaml" + deploy_path.write_text( + """\ +model_runner: v1 +async_chunk: true +stages: + - stage_id: 0 + - stage_id: 1 + model_runner: v2 + - stage_id: 2 + model_runner: v2 +""" + ) + deploy = load_deploy_config(deploy_path) + assert [stage.model_runner for stage in deploy.stages] == [None, "v2", "v2"] + + pipeline = _qwen3_omni_pipeline() + legacy = merge_pipeline_deploy(pipeline, deploy) + assert [stage.yaml_engine_args["use_v2_model_runner"] for stage in legacy] == [False, True, True] + # The test pipeline declares native support; only v2 stages use it. + assert [stage.yaml_engine_args["supports_native_mrv2_data_plane"] for stage in legacy] == [True, True, True] + + structured = VllmOmniConfig.from_pipeline_config(pipeline, deploy_config_path=str(deploy_path)) + assert [stage.model_config.use_v2_model_runner for stage in structured.stage_configs] == [False, True, True] + + +def test_platform_v1_fallback_overrides_stage_v2(): + deploy = DeployConfig( + stages=[StageDeployConfig(stage_id=0), StageDeployConfig(stage_id=1, model_runner="v2")], + platforms={"rocm": {"model_runner": "v1"}}, + ) + deploy = _apply_platform_overrides(deploy, platform="rocm") + assert deploy.model_runner == "v1" + assert deploy.stages[1].model_runner is None + + +def test_downstream_mrv2_stage_requires_turn_sessions(): + stages = [StageDeployConfig(stage_id=0), StageDeployConfig(stage_id=1, model_runner="v2")] + with pytest.raises(ValueError, match="session_mode 'turn' only"): + merge_pipeline_deploy(_qwen3_omni_pipeline(), DeployConfig(stages=stages, session_mode="duplex")) + + +def test_v1_downstream_stage_keeps_duplex_sessions(): + deploy = DeployConfig(stages=[StageDeployConfig(stage_id=0)], session_mode="duplex") + stages = merge_pipeline_deploy(_qwen3_omni_pipeline(), deploy) + assert not any(stage.yaml_engine_args["use_v2_model_runner"] for stage in stages) + + def test_qwen3_tts_high_concurrency_mrv2_profile_is_explicit_opt_in(): assert load_deploy_config(_DEPLOY_DIR / "qwen3_tts_high_concurrency.yaml").model_runner == "v1" assert load_deploy_config(_DEPLOY_DIR / "qwen3_tts_high_concurrency_mrv2.yaml").model_runner == "v2" @@ -2220,3 +2276,12 @@ def test_mps_stays_in_runtime_instead_of_engine_arguments(): assert stage.runtime_config.cuda_mps args = build_engine_args_dict_from_omni_stage_config(stage, model="test-model") assert "cuda_mps" not in args + + +def test_platform_stage_overlay_rejects_invalid_model_runner(): + deploy = DeployConfig( + stages=[StageDeployConfig(stage_id=0)], + platforms={"cuda": {"stages": [{"stage_id": 0, "model_runner": "v3"}]}}, + ) + with pytest.raises(ValueError, match="model_runner must be 'v1' or 'v2'"): + _apply_platform_overrides(deploy, platform="cuda") diff --git a/tests/core/sched/test_generation_scheduler_restore.py b/tests/core/sched/test_generation_scheduler_restore.py index ab7c65aa7..546084509 100644 --- a/tests/core/sched/test_generation_scheduler_restore.py +++ b/tests/core/sched/test_generation_scheduler_restore.py @@ -25,8 +25,8 @@ class FakeAdapter: """Minimal mock of OmniChunkTransferAdapter tracking restore calls.""" def __init__(self): - self.waiting_for_chunk_waiting_requests = deque() - self.waiting_for_chunk_running_requests = deque() + self.waiting_for_chunk_waiting_requests: deque = deque() + self.waiting_for_chunk_running_requests: deque = deque() self.restore_called = False self.done_request_ids = set() @@ -249,3 +249,103 @@ class TestRestoreQueuesOnError: assert adapter.restore_called is True assert "req-B" in running + + +def test_first_chunk_express_slack_guard_tracks_emitted_audio(monkeypatch): + """Express steps wait while a ready later chunk's stream is close to underrun.""" + import vllm_omni.core.sched.omni_generation_scheduler as module + + scheduler = OmniGenerationScheduler.__new__(OmniGenerationScheduler) + scheduler._first_chunk_express = True + scheduler._express_min_slack_s = 0.5 + scheduler._express_skipped_for_slack = 0 + scheduler._stream_audio = {} + scheduler._chunk_started = {"started", "idle"} + ready = SimpleNamespace( + request_id="started", num_in_flight_tokens=0, prompt_token_ids=[0] * 28, num_computed_tokens=0 + ) + # Started, but no chunk to decode now: never blocks an express step. + idle = SimpleNamespace(request_id="idle", num_in_flight_tokens=0, prompt_token_ids=[0] * 28, num_computed_tokens=28) + scheduler.running = [ready, idle] + scheduler.waiting = [] + clock = [100.0] + monkeypatch.setattr(module.time, "monotonic", lambda: clock[0]) + + # No measured audio credit: do not delay the ready stream. + assert not scheduler._continuations_have_slack() + # 1 s of audio at t=100; at t=100.2 it holds 0.8 s, at t=100.7 only 0.3 s. + scheduler._record_stream_audio("started", {"model_outputs": torch.zeros(24000), "sr": torch.tensor(24000)}) + scheduler._record_stream_audio("started", {"model_outputs": torch.zeros(0), "sr": torch.tensor(24000)}) + clock[0] = 100.2 + assert scheduler._continuations_have_slack() + clock[0] = 100.7 + assert not scheduler._continuations_have_slack() + # Another second emitted: slack is back to 1.3 s. + scheduler._record_stream_audio("started", {"model_outputs": torch.zeros(24000), "sr": torch.tensor(24000)}) + assert scheduler._continuations_have_slack() + assert scheduler._stream_audio["started"] == [100.0, 2.0] + + +def _express_scheduler(): + continuation = Request("continuation", [1, 2], SamplingParams(max_tokens=4), pooling_params=None) + first = Request("first", [1], SamplingParams(max_tokens=4), pooling_params=None) + scheduler = _make_generation_scheduler(continuation, use_v2_model_runner=True) + scheduler.waiting.add_request(first) + scheduler.requests[first.request_id] = first + scheduler.max_num_running_reqs = 4 + scheduler._native_data_plane = True + scheduler.chunk_transfer_adapter = None + scheduler.input_coordinator = SimpleNamespace( + _async_chunk=True, finished_requests=set(), restore_queues=lambda *args, **kwargs: None + ) + scheduler._first_chunk_express = True + scheduler._last_step_express = False + scheduler._chunk_started = {"continuation"} + scheduler._express_min_slack_s = 0 + scheduler._stream_audio = {} + return scheduler, continuation, first + + +def test_express_schedules_only_first_chunks_then_allows_continuations(): + scheduler, continuation, first = _express_scheduler() + output = scheduler.schedule() + assert output.num_scheduled_tokens == {"first": 1} + assert scheduler._last_step_express + assert continuation in list(scheduler.waiting) + # Another first chunk arrives before the continuation runs. + first.num_computed_tokens = len(first.prompt_token_ids) + second = Request("second", [1], SamplingParams(max_tokens=4), pooling_params=None) + scheduler.waiting.add_request(second) + scheduler.requests[second.request_id] = second + output = scheduler.schedule() + assert not scheduler._last_step_express + assert "continuation" in output.num_scheduled_tokens + + +def test_express_does_not_delay_continuation_for_unready_first_chunk(): + scheduler, continuation, first = _express_scheduler() + first.num_in_flight_tokens = 1 + output = scheduler.schedule() + assert not scheduler._last_step_express + assert output.num_scheduled_tokens == {"continuation": 2} + + +def test_express_slack_guard_is_used_by_schedule(): + scheduler, continuation, first = _express_scheduler() + scheduler._express_min_slack_s = 0.5 + scheduler._express_skipped_for_slack = 0 + output = scheduler.schedule() + assert not scheduler._last_step_express + assert "continuation" in output.num_scheduled_tokens + + +def test_express_cancel_cleans_playback_state(monkeypatch): + from vllm.v1.core.sched.scheduler import Scheduler + + scheduler, continuation, _ = _express_scheduler() + scheduler.input_coordinator = None + scheduler._stream_audio["continuation"] = [1.0, 2.0] + monkeypatch.setattr(Scheduler, "_free_request", lambda *args: (None, None)) + scheduler._free_request(continuation) + assert "continuation" not in scheduler._chunk_started + assert "continuation" not in scheduler._stream_audio 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 fc165c968..0f509c5d2 100644 --- a/tests/core/sched/test_omni_generation_scheduler_update_session.py +++ b/tests/core/sched/test_omni_generation_scheduler_update_session.py @@ -118,6 +118,8 @@ def test_resumable_generation_stop_marks_segment_boundary() -> None: session.num_in_flight_tokens = 1 sched = MagicMock() + sched._first_chunk_express = False + sched._express_min_slack_s = 0.0 sched.requests = {session.request_id: session} sched.perf_metrics = None sched.chunk_transfer_adapter = SimpleNamespace( @@ -185,6 +187,8 @@ def test_rejected_grammar_finishes_and_frees_generation_request(prompt_complete: session.num_in_flight_tokens = 1 sched = MagicMock() + sched._first_chunk_express = False + sched._express_min_slack_s = 0.0 sched.requests = {session.request_id: session} sched.perf_metrics = None sched.chunk_transfer_adapter = SimpleNamespace( diff --git a/tests/core/sched/test_omni_sched_ec_request_finish.py b/tests/core/sched/test_omni_sched_ec_request_finish.py index 8316a4308..f1b747adc 100644 --- a/tests/core/sched/test_omni_sched_ec_request_finish.py +++ b/tests/core/sched/test_omni_sched_ec_request_finish.py @@ -120,6 +120,8 @@ def _make_finish_sched(request: Request) -> MagicMock: """MagicMock scheduler driving the real ``update_from_output`` through a finishing request; ``_free_request`` reports EC transfer params.""" sched = MagicMock() + sched._first_chunk_express = False + sched._express_min_slack_s = 0.0 sched.requests = {request.request_id: request} sched.perf_metrics = None sched.defer_block_free = False diff --git a/tests/core/sched/test_omni_sched_prefill_stats_finalize.py b/tests/core/sched/test_omni_sched_prefill_stats_finalize.py index f512b5a6e..3438466b7 100644 --- a/tests/core/sched/test_omni_sched_prefill_stats_finalize.py +++ b/tests/core/sched/test_omni_sched_prefill_stats_finalize.py @@ -53,6 +53,8 @@ def _make_request() -> Request: def _make_sched(request: Request) -> MagicMock: sched = MagicMock() + sched._first_chunk_express = False + sched._express_min_slack_s = 0.0 sched.requests = {request.request_id: request} sched.perf_metrics = None sched.defer_block_free = False diff --git a/tests/core/sched/test_omni_scheduler_mixin_shared.py b/tests/core/sched/test_omni_scheduler_mixin_shared.py index ed18de0b9..82956e27a 100644 --- a/tests/core/sched/test_omni_scheduler_mixin_shared.py +++ b/tests/core/sched/test_omni_scheduler_mixin_shared.py @@ -175,7 +175,7 @@ def test_finished_request_attachment_keeps_ar_abort_policy_explicit( ): scheduler = _Scheduler() scheduler.finished_req_ids_dict = defaultdict(set, {2: {"req-finished"}}) - outputs = {} + outputs: dict = {} scheduler._attach_finished_request_sets( outputs, @@ -207,3 +207,26 @@ def test_output_helper_preserves_required_nan_counter_default(): output = scheduler._make_omni_engine_output(request, new_token_ids=[]) assert output.num_nans_in_logits == 0 + + +@pytest.mark.parametrize(("role", "coordinated"), [("sender", False), ("receiver", True), (None, True)]) +def test_native_downstream_sender_stage_does_not_wait_for_chunks(role, coordinated): + """An orchestrator-fed downstream stage (MiniCPM-o's Talker) owns only an + outgoing connector; the native plane must not park it for input chunks.""" + scheduler = _Scheduler() + model_config = object.__new__(OmniModelConfig) + model_config.stage_id = 1 + model_config.async_chunk = True + model_config.requires_full_payload_input = False + model_config.supports_native_mrv2_data_plane = True + model_config.use_v2_model_runner = True + model_config.stage_connector_config = {"name": "SharedMemoryConnector", "extra": {"role": role} if role else {}} + scheduler.vllm_config = object.__new__(VllmConfig) + scheduler.vllm_config.model_config = model_config + scheduler.vllm_config.scheduler_config = object.__new__(SchedulerConfig) + scheduler.vllm_config.scheduler_config.max_num_seqs = 1 + scheduler._init_omni_io_scheduling_state() + + assert scheduler._native_data_plane + assert (scheduler.input_coordinator is not None) is coordinated + assert scheduler._async_chunk_transport_enabled() is coordinated diff --git a/tests/diffusion/models/qwen_image/test_qwen_image_pipeline_device.py b/tests/diffusion/models/qwen_image/test_qwen_image_pipeline_device.py index 7f0ce51c8..f36baad0d 100644 --- a/tests/diffusion/models/qwen_image/test_qwen_image_pipeline_device.py +++ b/tests/diffusion/models/qwen_image/test_qwen_image_pipeline_device.py @@ -64,7 +64,7 @@ class _MoveableModule(nn.Module): return [True, True, True] -def _run_pipeline_init(monkeypatch, *, enable_cpu_offload, loader_device): +def _run_pipeline_init(monkeypatch, *, enable_cpu_offload, loader_device, diffusion_offload_config=None): import vllm_omni.diffusion.models.qwen_image.pipeline_qwen_image as pipe_mod from vllm_omni.diffusion.models.qwen_image.pipeline_qwen_image import QwenImagePipeline @@ -110,14 +110,21 @@ def _run_pipeline_init(monkeypatch, *, enable_cpu_offload, loader_device): tf_model_config={}, quantization_config=None, enable_cpu_offload=enable_cpu_offload, + diffusion_offload_config=diffusion_offload_config, enable_diffusion_pipeline_profiler=False, ) return QwenImagePipeline(od_config=od_config) -def test_pipeline_cpu_offload_parks_encoder_vae_on_cpu(monkeypatch): +@pytest.mark.parametrize("compact", [False, True]) +def test_pipeline_cpu_offload_parks_encoder_vae_on_cpu(monkeypatch, compact): """#7555: cpu_offload keeps BF16 encoder/VAE off the GPU during DiT init.""" - pipe = _run_pipeline_init(monkeypatch, enable_cpu_offload=True, loader_device="cuda") + pipe = _run_pipeline_init( + monkeypatch, + enable_cpu_offload=not compact, + loader_device="cuda", + diffusion_offload_config={"mode": "module", "components": ["dit"]} if compact else None, + ) assert pipe.text_encoder.placed_device.type == "cpu" assert pipe.vae.placed_device.type == "cpu" assert pipe.transformer.probe_device_type == "cuda" diff --git a/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py b/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py index bffa8a026..8f85366e3 100644 --- a/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py +++ b/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py @@ -224,7 +224,7 @@ def test_streaming_payload_can_replace_placeholder_prompt(mocker: MockerFixture) def test_turn_start_replacement_ignores_accumulated_prompt_capacity(mocker: MockerFixture) -> None: request = _streaming_request(mocker, num_computed_tokens=4064) - payload = { + payload: dict = { "ids": {"prompt": [1]}, "meta": { "replace_streaming_prompt": True, @@ -336,7 +336,7 @@ def test_streaming_window_builds_one_chunk_recompute_recipe(mocker: MockerFixtur num_output_placeholders=1, update_block_hashes=mocker.Mock(), ) - payload = { + payload: dict = { "ids": {"prompt": [1]}, "meta": { "next_stage_prompt_len": 10, @@ -406,7 +406,7 @@ def test_streaming_window_appends_until_capacity_then_recomputes(mocker: MockerF num_output_placeholders=0, update_block_hashes=mocker.Mock(), ) - second = { + second: dict = { "ids": {"prompt": [1]}, "meta": { "next_stage_prompt_len": 12, @@ -432,7 +432,7 @@ def test_streaming_window_appends_until_capacity_then_recomputes(mocker: MockerF request._all_token_ids.extend([201, 202, 203]) request._output_token_ids.extend([201, 202, 203]) request.num_computed_tokens = 4017 - third = { + third: dict = { "ids": {"prompt": [2]}, "meta": { "next_stage_prompt_len": 60, @@ -458,7 +458,7 @@ def test_streaming_window_appends_until_capacity_then_recomputes(mocker: MockerF request._all_token_ids.extend([301]) request._output_token_ids.extend([301]) request.num_computed_tokens = 76 - fourth = { + fourth: dict = { "ids": {"prompt": [3]}, "meta": { "next_stage_prompt_len": 8, @@ -1806,7 +1806,7 @@ def test_sender_only_adapter_does_not_park_or_clear_requests(build_adapter): request = _req("req-1", RequestStatus.WAITING) request.additional_information = {"tts_token_ids": torch.tensor([1])} waiting_queue = DummyWaitingQueue([request]) - running_queue = [] + running_queue: list = [] adapter.load_async(request) adapter.process_pending_chunks( @@ -1846,7 +1846,7 @@ def test_fifo_promotion(build_adapter): adapter, _ = build_adapter(stage_id=1, model_mode="generation", max_num_seqs=2, active_stream_window=2) reqs = [_req(f"req-{idx}", RequestStatus.WAITING) for idx in range(1, 5)] waiting_queue = DummyWaitingQueue(reqs) - running_queue = [] + running_queue: list = [] adapter.process_pending_chunks(waiting_queue, running_queue) @@ -1876,7 +1876,7 @@ def test_non_active_waiting_request_is_held_off_scheduler(build_adapter): active = _req("req-active", RequestStatus.WAITING) non_active = _req("req-non-active", RequestStatus.WAITING) waiting_queue = DummyWaitingQueue([active, non_active]) - running_queue = [] + running_queue: list = [] adapter.process_pending_chunks(waiting_queue, running_queue) @@ -1911,7 +1911,7 @@ def test_finished_releases_slot(build_adapter): req_1 = _req("req-1", RequestStatus.WAITING) req_2 = _req("req-2", RequestStatus.WAITING) waiting_queue = DummyWaitingQueue([req_1, req_2]) - running_queue = [] + running_queue: list = [] adapter.process_pending_chunks(waiting_queue, running_queue) assert list(adapter._active_streams) == ["req-1"] @@ -2004,7 +2004,7 @@ def test_cleanup_receiver_releases_multiple_slots_in_sequence(build_adapter): adapter, _ = build_adapter(stage_id=1, model_mode="generation", max_num_seqs=2, active_stream_window=2) reqs = [_req(f"req-{idx}", RequestStatus.WAITING) for idx in range(1, 5)] waiting_queue = DummyWaitingQueue(reqs) - running_queue = [] + running_queue: list = [] adapter.process_pending_chunks(waiting_queue, running_queue) assert list(adapter._active_streams) == ["req-1", "req-2"] @@ -2364,7 +2364,7 @@ def test_finish_requests_releases_active_stream_slot(build_adapter): aborted = _req("req-aborted", RequestStatus.RUNNING) waiting = _req("req-waiting", RequestStatus.WAITING) waiting_queue = DummyWaitingQueue([waiting]) - running_queue = [] + running_queue: list = [] adapter._active_streams[aborted.request_id] = aborted adapter._held_non_active.append(aborted) @@ -2488,7 +2488,7 @@ def test_restore_queues_skips_requests_missing_from_scheduler_requests(build_ada zombie = _req("req-zombie", RequestStatus.WAITING_FOR_CHUNK) live = _req("req-live", RequestStatus.WAITING_FOR_CHUNK) waiting_queue = DummyWaitingQueue() - running_queue = [] + running_queue: list = [] adapter.waiting_for_chunk_waiting_requests = deque([zombie, live]) adapter.waiting_for_chunk_running_requests = deque([zombie, live]) @@ -2522,7 +2522,7 @@ class _HashableRequest(SimpleNamespace): # double needs it too. num_stale_output_tokens = 0 - def __hash__(self): + def __hash__(self): # type: ignore[override] # Test requests are hashable by request ID. return hash(self.request_id) def __eq__(self, other): @@ -2540,6 +2540,8 @@ def test_generation_scheduler_calls_cleanup_on_finished(monkeypatch, mocker: Moc from vllm_omni.core.sched.omni_generation_scheduler import OmniGenerationScheduler scheduler = mocker.MagicMock() + scheduler._first_chunk_express = False + scheduler._express_min_slack_s = 0.0 scheduler.chunk_transfer_adapter = adapter_mock scheduler.connector = None scheduler.ec_connector = None diff --git a/tests/distributed/omni_connectors/test_shm_connector.py b/tests/distributed/omni_connectors/test_shm_connector.py index 8d9ad0270..39a0e8327 100644 --- a/tests/distributed/omni_connectors/test_shm_connector.py +++ b/tests/distributed/omni_connectors/test_shm_connector.py @@ -294,3 +294,280 @@ def test_consumed_bad_payload_removes_lock_file(connector, monkeypatch): assert not os.path.exists(f"/dev/shm/shm_{key}_lockfile.lock") connector.reap_consumed() assert key not in connector._pending_keys + + +# ── Arrival wakeups ─────────────────────────────────────────────────── + + +def _stage_connector(stage_id): + return SharedMemoryConnector({"stage_id": stage_id}) + + +def test_wakeup_discovery_only_scans_deployment_directory(monkeypatch, tmp_path): + from vllm_omni.distributed.omni_connectors.connectors import shm_connector as module + + directory = tmp_path / "deployment" + monkeypatch.setattr(module, "_wakeup_directory", lambda: str(directory)) + original_scandir = os.scandir + scanned = [] + + def scoped_scandir(path): + # A parent-directory scan would recreate the performance regression. + assert os.fspath(path) == str(directory) + scanned.append(path) + return original_scandir(path) + + monkeypatch.setattr(module.glob.os, "scandir", scoped_scandir) + sender, receiver = _stage_connector(86), _stage_connector(87) + try: + generation = receiver.get_wakeup_generation() + assert generation is not None + assert os.path.dirname(receiver._wake_path) == str(directory) + sender._wake_receiver(87) + assert receiver.wait_for_change(generation, timeout=0.5) + assert scanned + finally: + sender.close() + receiver.close() + + +def test_missing_wakeup_directory_keeps_payload_transport_working(monkeypatch, tmp_path): + from vllm_omni.distributed.omni_connectors.connectors import shm_connector as module + + directory = tmp_path / "missing" + monkeypatch.setattr(module, "_wakeup_directory", lambda: str(directory)) + sender, receiver = _stage_connector(88), _stage_connector(89) + key = f"wake_{uuid.uuid4().hex}" + try: + assert sender.put("88", "89", key, {"value": 1})[0] + assert not directory.exists() + assert receiver.get("88", "89", key)[0] == {"value": 1} + finally: + sender.close() + receiver.close() + + +def test_put_wakes_waiting_receiver(): + sender, receiver = _stage_connector(70), _stage_connector(71) + try: + generation = receiver.get_wakeup_generation() + assert generation is not None + start = time.monotonic() + assert receiver.wait_for_change(generation, timeout=0.05) is False + assert time.monotonic() - start >= 0.04 # no arrival: blocks for the timeout + + key = f"wake_{uuid.uuid4().hex}" + sender.put("70", "71", key, {"value": 1}) + start = time.monotonic() + assert receiver.wait_for_change(generation, timeout=5) is True + assert time.monotonic() - start < 0.5 + assert receiver.get("70", "71", key)[0] == {"value": 1} + # Drained: the next wait blocks again. + assert receiver.wait_for_change(receiver.get_wakeup_generation(), timeout=0.02) is False + finally: + sender.close() + receiver.close() + + +def test_receiver_does_not_spin_after_sender_closes(): + sender, receiver = _stage_connector(72), _stage_connector(73) + try: + receiver.get_wakeup_generation() + sender.put("72", "73", f"wake_{uuid.uuid4().hex}", {"value": 1}) + sender.close() # the only external writer goes away + generation = receiver.get_wakeup_generation() + receiver.wait_for_change(generation, timeout=1) + generation = receiver.get_wakeup_generation() + start = time.monotonic() + assert receiver.wait_for_change(generation, timeout=0.05) is False + assert time.monotonic() - start >= 0.04 + finally: + receiver.close() + + +def test_wakeups_can_be_disabled_and_close_unlinks_fifo(monkeypatch): + receiver = _stage_connector(74) + assert receiver.get_wakeup_generation() is not None + path = receiver._wake_path + assert os.path.exists(path) + receiver.close() + assert not os.path.exists(path) + + monkeypatch.setenv("VLLM_OMNI_SHM_WAKEUP", "0") + sender, receiver = _stage_connector(75), _stage_connector(76) + try: + assert receiver.get_wakeup_generation() is None # caller keeps its timed poll + key = f"wake_{uuid.uuid4().hex}" + assert sender.put("75", "76", key, {"value": 2})[0] + assert receiver.get("75", "76", key)[0] == {"value": 2} + finally: + sender.close() + receiver.close() + + +def _sibling_receiver(stage_id, ready, result): + receiver = _stage_connector(stage_id) + generation = receiver.get_wakeup_generation() + ready.set() + start = time.monotonic() + woke = receiver.wait_for_change(generation, timeout=5) + result.put((woke, time.monotonic() - start)) + receiver.close() + + +def _sibling_sender(stage_id, key): + sender = _stage_connector(stage_id - 1) + sender.put(str(stage_id - 1), str(stage_id), key, {"value": 3}) + + +def test_sibling_stage_processes_share_wakeups(): + # Stage engine processes are siblings under the launching process. + import multiprocessing + + ctx = multiprocessing.get_context("fork") + ready, result = ctx.Event(), ctx.Queue() + key = f"wake_{uuid.uuid4().hex}" + receiver = ctx.Process(target=_sibling_receiver, args=(78, ready, result)) + receiver.start() + assert ready.wait(5) + time.sleep(0.05) + sender = ctx.Process(target=_sibling_sender, args=(78, key)) + sender.start() + sender.join(5) + woke, waited = result.get(timeout=5) + receiver.join(5) + assert woke and waited < 1.0 + cleanup = SharedMemoryConnector({}) + cleanup.cleanup(key) + + +def test_same_stage_receivers_are_independent_across_close_and_restart(): + sender = _stage_connector(80) + first, second = _stage_connector(81), _stage_connector(81) + replacement = None + try: + first_generation = first.get_wakeup_generation() + second_generation = second.get_wakeup_generation() + assert first._wake_path != second._wake_path + sender._wake_receiver(81) + assert first.wait_for_change(first_generation, timeout=0.5) + assert second.wait_for_change(second_generation, timeout=0.5) + second_generation = second.get_wakeup_generation() + first.close() + assert first.get_wakeup_generation() is None + assert os.path.exists(second._wake_path) + replacement = _stage_connector(81) + replacement_generation = replacement.get_wakeup_generation() + sender._wake_receiver(81) + assert second.wait_for_change(second_generation, timeout=0.5) + assert replacement.wait_for_change(replacement_generation, timeout=0.5) + finally: + sender.close() + first.close() + second.close() + if replacement is not None: + replacement.close() + + +def test_partial_receiver_open_closes_fd_and_removes_only_own_fifo(monkeypatch, tmp_path): + import errno + + from vllm_omni.distributed.omni_connectors.connectors import shm_connector as module + + monkeypatch.setattr(module, "_wakeup_directory", lambda: str(tmp_path)) + receiver = _stage_connector(82) + peer_path = tmp_path / "stage_82_peer" + os.mkfifo(peer_path, 0o600) + original_open = os.open + opened = [] + + def fail_hold_open(path, flags, *args, **kwargs): + if flags & os.O_WRONLY: + raise OSError(errno.EMFILE, "injected fd exhaustion") + fd = original_open(path, flags, *args, **kwargs) + opened.append(fd) + return fd + + monkeypatch.setattr(module.os, "open", fail_hold_open) + try: + assert receiver.get_wakeup_generation() is None + assert list(tmp_path.iterdir()) == [peer_path] + assert len(opened) == 1 + with pytest.raises(OSError, match="Bad file descriptor"): + os.fstat(opened[0]) + finally: + receiver.close() + + +def test_wakeup_ignores_stale_fifo_and_non_fifo_paths(monkeypatch, tmp_path): + from vllm_omni.distributed.omni_connectors.connectors import shm_connector as module + + monkeypatch.setattr(module, "_wakeup_directory", lambda: str(tmp_path)) + sender, receiver = _stage_connector(83), _stage_connector(84) + plain = tmp_path / "84_plain" + plain.write_text("unchanged") + (tmp_path / "84_link").symlink_to(plain) + os.mkfifo(tmp_path / "84_stale", 0o600) + try: + generation = receiver.get_wakeup_generation() + sender._wake_receiver(84) + assert receiver.wait_for_change(generation, timeout=0.5) + assert plain.read_text() == "unchanged" + assert (tmp_path / "84_stale").exists() + finally: + sender.close() + receiver.close() + + +def test_waiter_does_not_read_fd_reused_after_close(monkeypatch, tmp_path): + from vllm_omni.distributed.omni_connectors.connectors import shm_connector as module + + receiver = _stage_connector(85) + generation = receiver.get_wakeup_generation() + read_fd = receiver._wake_read_fd + path = tmp_path / "unrelated_ipc" + path.write_bytes(b"untouched") + reused = [] + + def close_during_select(*args): + receiver.close() + fd = os.open(path, os.O_RDONLY) + if fd != read_fd: + os.dup2(fd, read_fd) + os.close(fd) + reused.append(read_fd) + return [read_fd], [], [] + + monkeypatch.setattr(module.select, "select", close_during_select) + try: + assert not receiver.wait_for_change(generation, timeout=0.1) + assert os.lseek(read_fd, 0, os.SEEK_CUR) == 0 + finally: + for fd in reused: + os.close(fd) + receiver.close() + + +def test_deployment_scope_survives_different_worker_parents(monkeypatch): + config = {"stage_id": 1, "extra": {"wakeup_scope": uuid.uuid4().hex}} + monkeypatch.setattr(os, "getppid", lambda: 111) + receiver = SharedMemoryConnector(config) + monkeypatch.setattr(os, "getppid", lambda: 222) + sender = SharedMemoryConnector(config) + try: + generation = receiver.get_wakeup_generation() + sender._wake_receiver(1) + assert receiver.wait_for_change(generation, timeout=0.1) + directory = receiver._wake_directory + finally: + receiver.close() + sender.close() + assert not os.path.exists(directory) + + +@pytest.mark.parametrize("value,expected", [("bad", 0.005), ("nan", 0.005), ("-1", 0.005), ("20", 0.02)]) +def test_receive_poll_interval_validation(monkeypatch, value, expected): + from vllm_omni.distributed.omni_connectors.model_runner.omni_connector_payload_transport import _recv_poll_seconds + + monkeypatch.setenv("VLLM_OMNI_CONNECTOR_RECV_POLL_MS", value) + assert _recv_poll_seconds() == expected diff --git a/tests/engine/test_orchestrator_error_handling.py b/tests/engine/test_orchestrator_error_handling.py index f05c5f132..dd1f1efb5 100644 --- a/tests/engine/test_orchestrator_error_handling.py +++ b/tests/engine/test_orchestrator_error_handling.py @@ -519,6 +519,40 @@ async def test_async_chunk_prewarm_without_prompt_token_ids_fails_only_that_requ orchestrator_fixture.thread.join(timeout=5) +@pytest.mark.asyncio +async def test_streaming_input_to_mrv2_downstream_stage_fails_with_client_error(orchestrator_factory) -> None: + """A downstream MRv2 native-data-plane stage supports turn-based requests + only: a resumable (realtime) request must fail up front, not hang.""" + stage0 = FakeStageClient(stage_type="llm", final_output=False) + stage1 = FakeStageClient(stage_type="llm", final_output=True) + v1 = SimpleNamespace(model_config=SimpleNamespace(max_model_len=64)) + mrv2 = SimpleNamespace( + model_config=SimpleNamespace(max_model_len=64, use_v2_model_runner=True, supports_native_mrv2_data_plane=True) + ) + fixture = orchestrator_factory([stage0, stage1], stage_vllm_configs=[v1, mrv2], async_chunk=True) + + try: + await _enqueue_add_request( + fixture, + request_id="req-stream", + prompt=SimpleNamespace(request_id="req-stream", prompt_token_ids=[1, 2], resumable=True), + original_prompt={"prompt": "stream"}, + sampling_params_list=[_sampling_params(), _sampling_params()], + final_stage_id=1, + ) + + error_msg = await _wait_for_error_message(fixture, request_id="req-stream") + assert error_msg.fatal is False + assert error_msg.status_code == 400 + assert "stage 1 runs on model_runner v2" in error_msg.error + assert stage0.add_request_calls == [] and stage1.add_request_calls == [] + assert fixture.thread.is_alive() + finally: + if fixture.thread.is_alive(): + fixture.request_sync_q.put_nowait(ShutdownRequestMessage()) + fixture.thread.join(timeout=5) + + # ───────── Direct unit tests for the fault-isolation helpers ───────── diff --git a/tests/worker/test_device_copy.py b/tests/worker/test_device_copy.py new file mode 100644 index 000000000..9ae8935b1 --- /dev/null +++ b/tests/worker/test_device_copy.py @@ -0,0 +1,28 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +import pytest +import torch + +from tests.helpers.mark import hardware_test +from vllm_omni.utils.device_copy import to_device_nonblocking + +pytestmark = pytest.mark.core_model + + +@pytest.mark.cpu +def test_tensor_copy_cpu_preserves_view_and_dtype(): + source = torch.arange(12, dtype=torch.float64).reshape(3, 4).T + result = to_device_nonblocking(source, "cpu") + assert result is source + assert result.dtype == torch.float64 + + +@hardware_test(res={"cuda": "L4"}, num_cards=1) +def test_tensor_copy_cuda_handles_pageable_and_pinned_sources(): + if not torch.cuda.is_available(): + pytest.skip("CUDA required") + source = torch.arange(12, dtype=torch.float32).reshape(3, 4).T + for host in (source, source.pin_memory()): + result = to_device_nonblocking(host, "cuda") + torch.testing.assert_close(result.cpu(), host) + assert to_device_nonblocking(result, result.device) is result diff --git a/tests/worker/test_omni_gpu_model_runner.py b/tests/worker/test_omni_gpu_model_runner.py index 7e6258960..0a8997f97 100644 --- a/tests/worker/test_omni_gpu_model_runner.py +++ b/tests/worker/test_omni_gpu_model_runner.py @@ -607,6 +607,39 @@ def test_update_additional_information_deserializes_new_request_payload(): ) +@pytest.mark.parametrize("replace_payload", [False, True]) +def test_resumed_streaming_payload_converts_to_typed_model_input(replace_payload): + from vllm_omni.core.sched.output import OmniNewRequestData + from vllm_omni.data_entry_keys import to_struct + + runner = _make_runner(req_ids=("r1",), hidden_size=4) + runner.model.replace_runtime_additional_information = replace_payload + payload = { + "codes": {"audio": torch.tensor([1, 2])}, + "meta": {"req_id": ["r1"], "stream_finished": torch.tensor(False)}, + } + new_req = OmniNewRequestData( + req_id="r1", + prompt_token_ids=[1, 2], + mm_features=[], + sampling_params=None, + pooling_params=None, + block_ids=([],), + num_computed_tokens=0, + lora_request=None, + model_intermediate_buffer=payload, + ) + + OmniGPUModelRunner._update_streaming_input_additional_info(runner, new_req, "r1") + # CosyVoice and other typed-payload consumers perform this conversion + # after the runner adds its resume metadata. + result = to_struct(runner.model_intermediate_buffer["r1"]) + assert result.meta.resumable is True + assert result.meta.num_processed_tokens == 0 + assert result.meta.req_id == ["r1"] + torch.testing.assert_close(result.codes.audio, payload["codes"]["audio"]) + + def test_streaming_new_request_marker_replaces_terminal_chunk_snapshot(): from vllm_omni.engine.serialization import serialize_additional_information @@ -1222,3 +1255,30 @@ def test_decode_batch_without_mtp_rejects_incomplete_outputs(monkeypatch, invali with pytest.raises(ValueError, match="decode preprocessing"): runner._preprocess(scheduled, 2) assert torch.equal(runner.inputs_embeds.gpu, before) + + +@pytest.mark.parametrize("supports_boundaries", [False, True]) +def test_multimodal_embeddings_receive_opt_in_host_request_boundaries(supports_boundaries): + runner = object.__new__(OmniGPUModelRunner) + received = {} + + def embed(ids, **kwargs): + received.update(kwargs) + return ids + + runner.model = SimpleNamespace(supports_embed_input_ids_query_start_loc=supports_boundaries, embed_input_ids=embed) + runner.input_batch = SimpleNamespace(num_reqs=3) + runner.query_start_loc = SimpleNamespace(cpu=torch.tensor([0, 6, 7, 11, 999])) + runner.input_ids = SimpleNamespace(gpu=torch.arange(16)) + mm = [torch.ones(2, 4), torch.ones(1, 4)] + mask = torch.tensor([True] * 4 + [False] * 3 + [True] * 3 + [False]) + + result = runner._embed_multimodal_input_ids(11, mm, mask) + + torch.testing.assert_close(result, torch.arange(11)) + assert received["multimodal_embeddings"] is mm + assert received["is_multimodal"] is mask + if supports_boundaries: + assert received["query_start_loc"] == [0, 6, 7, 11] + else: + assert "query_start_loc" not in received diff --git a/tests/worker_v2/test_intermediate_buffer.py b/tests/worker_v2/test_intermediate_buffer.py index a84fef711..6ddd9738a 100644 --- a/tests/worker_v2/test_intermediate_buffer.py +++ b/tests/worker_v2/test_intermediate_buffer.py @@ -94,3 +94,25 @@ def test_update_merge_semantics(monkeypatch): assert stored is not tensor assert torch.equal(stored, tensor) assert "kv" in buf.buffers[0] + + +def test_add_request_merges_orchestrator_model_intermediate_buffer(): + """Stage bridges such as MiniCPM-o's llm2tts hand over ``model_intermediate_buffer``.""" + buf = OmniIntermediateBuffer(max_num_reqs=2) + hidden = torch.randn(3, 4) + data = _make_new_req_data( + "talker-1", + additional_information={"meta": {"a": 1}}, + model_intermediate_buffer={ + "ids": {"tts": [1, 2, 3]}, + "hidden_states": {"tts": hidden}, + "meta": {"next_stage_prompt_len": 5}, + "request_id": "thinker-1", + }, + ) + buf.add_request(0, data) + info = buf.buffers[0] + assert info["ids"]["tts"] == [1, 2, 3] + assert torch.equal(info["hidden_states"]["tts"], hidden) + assert info["meta"] == {"a": 1, "next_stage_prompt_len": 5} + assert info["request_id"] == "thinker-1" and info["req_id"] == "talker-1" diff --git a/tests/worker_v2/test_omni_data_plane.py b/tests/worker_v2/test_omni_data_plane.py index ccb51c460..978346070 100644 --- a/tests/worker_v2/test_omni_data_plane.py +++ b/tests/worker_v2/test_omni_data_plane.py @@ -185,6 +185,7 @@ def _start_save_thread(plane, connector): _send_side_request_payload={}, _code_prompt_token_ids=defaultdict(list), _work_available=threading.Event(), + _save_work_available=threading.Event(), _stop_event=threading.Event(), _MAX_SEND_RETRIES=0, _can_send=True, @@ -199,7 +200,7 @@ def _start_save_thread(plane, connector): def _stop_save_thread(plane): plane._stop_event.set() - plane._work_available.set() + plane._save_work_available.set() plane._save_thread.join(timeout=2) assert not plane._save_thread.is_alive() @@ -551,3 +552,71 @@ def test_abort_before_first_chunk_cleans_receiver_state(): assert plane.abort_requests({"r"}) == 0 assert plane.abort_requests({"r"}) == 0 assert not plane._pending_load_reqs and not plane._request_ids_mapping and not plane._local_stage_payload_cache + + +def test_ar_receiver_accumulates_rows_and_replaces_chunk_metadata(): + # A V1 upstream sender stamps scalar flags on every chunk; they describe + # the latest chunk and must not be concatenated like decode rows. + plane = _bare_plane() + plane._send_side_request_payload = {} + + def chunk(rows, finished): + return { + "embed": {"decode": torch.ones((rows, 2))}, + "meta": {"finished": torch.tensor(finished), "is_segment_finished": torch.tensor(False)}, + } + + plane._accumulate_payload("ext", chunk(1, False)) + merged = plane._accumulate_payload("ext", chunk(2, True)) + + assert merged["embed"]["decode"].shape == (3, 2) + assert merged["meta"]["finished"].item() is True + assert merged["meta"]["is_segment_finished"].ndim == 0 + + +def test_terminal_without_processor_payload_still_sends_finish_marker(raw_plane, monkeypatch): + # Qwen3-Omni's Talker processor returns None on the terminal call when the + # last frames ended exactly on a chunk boundary; Code2Wav must still learn + # that the stream finished, as with the V1 chunk adapter. + raw_plane._async_chunk = True + raw_plane._omni_connector = object() + raw_plane._request_ids_mapping = {} + raw_plane._custom_process_func = lambda transfer_manager, multimodal_output, request, is_finished=False: None + raw_plane._custom_process_batch_func = None + monkeypatch.setattr(raw_plane, "is_data_transfer_rank", lambda: True) + enqueued = [] + + def enqueue(request, payload, **_kw): + enqueued.append((request.request_id, payload)) + return True, None + + monkeypatch.setattr(raw_plane, "_enqueue_chunk_payload", enqueue) + raw_plane.register_request(_new_request()) + + _complete(raw_plane, [{"codes": {"audio": torch.ones((1, 16))}}]) + assert enqueued == [] # a non-terminal None payload sends nothing + assert raw_plane.request_terminal({"internal"}) == 1 + [(req_id, payload)] = enqueued + assert req_id == "internal" + assert payload["meta"]["finished"].item() is True + + +def test_payload_builders_see_stage_model_config(monkeypatch): + # Producer-side builders (Qwen3-Omni's Thinker reads + # hf_config.talker_config.accept_hidden_layer) resolve the stage config + # through _get_model_config(), as they do on the V1 runner mixin. + from vllm_omni.model_executor.stage_input_processors.qwen3_omni import _get_accept_hidden_layer_index + + def init_connectors(self, *, model_config): + self._async_chunk = True + self._custom_process_func = None + + monkeypatch.setattr(OmniRunnerDataPlane, "init_omni_connectors", init_connectors) + monkeypatch.setattr(OmniRunnerDataPlane, "_start_output_worker", lambda self, **_kw: None) + model_config = SimpleNamespace(hf_config=SimpleNamespace(talker_config=SimpleNamespace(accept_hidden_layer=24))) + vllm_config = SimpleNamespace(model_config=model_config) + + plane = OmniRunnerDataPlane(vllm_config, model_config) + + assert plane._get_model_config() is model_config + assert _get_accept_hidden_layer_index(plane) == 24 diff --git a/tests/worker_v2/test_omni_generation_model_runner.py b/tests/worker_v2/test_omni_generation_model_runner.py index 63c480291..2b12837e4 100644 --- a/tests/worker_v2/test_omni_generation_model_runner.py +++ b/tests/worker_v2/test_omni_generation_model_runner.py @@ -13,7 +13,7 @@ import torch from vllm_omni.model_executor.models.output_templates import OmniOutput from vllm_omni.outputs import OmniModelRunnerOutput -from vllm_omni.worker_v2.omni_generation_model_runner import OmniGenerationModelRunner +from vllm_omni.worker_v2.omni_generation_model_runner import OmniGenerationModelRunner, check_exact_input_shape pytestmark = [pytest.mark.core_model, pytest.mark.cpu] @@ -182,3 +182,16 @@ def test_sample_tokens_uses_async_output_for_cuda_and_snapshots_req_ids(monkeypa output_req_ids = captured["model_runner_output"].req_ids input_batch.req_ids[0] = "reused" # snapshot must insulate the published output assert output_req_ids == ["req-0"] + + +@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). + model = SimpleNamespace(requires_exact_input_shape=exact) + check_exact_input_shape(model, torch.zeros(32, dtype=torch.long), 32) + check_exact_input_shape(model, None, 16) + if exact: + with pytest.raises(RuntimeError, match="exact input shape"): + check_exact_input_shape(model, torch.zeros(48, dtype=torch.long), 32) + else: + check_exact_input_shape(model, torch.zeros(48, dtype=torch.long), 32) diff --git a/tests/worker_v2/test_omni_gpu_model_runner.py b/tests/worker_v2/test_omni_gpu_model_runner.py index 835f2d5f5..dd4c760c6 100644 --- a/tests/worker_v2/test_omni_gpu_model_runner.py +++ b/tests/worker_v2/test_omni_gpu_model_runner.py @@ -181,6 +181,84 @@ def test_capture_contract_uses_model_declaration(stage, declared): runner._configure_cudagraph_output_contract() assert runner._model_returns_tuple is declared assert runner._exclude_full_graph is declared + assert runner._full_graph_aux_outputs is False + + +def _aux_output(num_tokens): + hidden = torch.arange(num_tokens * 2, dtype=torch.float32).reshape(num_tokens, 2) + return hidden, {"hidden_states": {"layers": {0: hidden + 1, 24: hidden + 2}}} + + +def test_full_graph_aux_contract_keeps_full_and_round_trips_leaves(): + runner = object.__new__(OmniGPUModelRunner) + runner.model = SimpleNamespace(_returns_tuple=True, supports_mrv2_full_graph_aux_outputs=True) + runner._configure_cudagraph_output_contract() + assert runner._full_graph_aux_outputs and not runner._exclude_full_graph + + calls = [] + + def forward(num_tokens=3): + calls.append(num_tokens) + return _aux_output(num_tokens) + + runner.model.forward = forward + original_forward = runner.model.forward + runner.use_aux_hidden_state_outputs = False + full = SimpleNamespace(cg_mode=CUDAGraphMode.FULL) + runner.cudagraph_manager = SimpleNamespace(_capture_descs={CUDAGraphMode.FULL: [full]}, _candidates={}) + runner.model_state = SimpleNamespace() + captured = {} + + def capture(_self): + # The graph manager stores (hidden, leaves) in its aux buffers. + assert runner.use_aux_hidden_state_outputs is True + hidden, leaves = runner.model.forward(4) + captured["hidden"], captured["leaves"] = hidden, leaves + assert [tuple(leaf.shape) for leaf in leaves] == [(4, 2), (4, 2)] + runner.model.forward(2) # smaller capture, same structure + return 1 + + with patch.object(GPUModelRunner, "capture_model", capture): + assert runner.capture_model() == 1 + assert runner.model.forward is original_forward and runner.use_aux_hidden_state_outputs is False + assert runner.cudagraph_manager._capture_descs == {CUDAGraphMode.FULL: [full]} # FULL kept + + hidden, aux = runner._split_fullgraph_output((captured["hidden"], captured["leaves"])) + reference_hidden, reference_aux = _aux_output(4) + assert torch.equal(hidden, reference_hidden) + assert set(aux["hidden_states"]["layers"]) == {0, 24} + assert torch.equal(aux["hidden_states"]["layers"][24], reference_aux["hidden_states"]["layers"][24]) + + +@pytest.mark.parametrize( + "aux", + [ + {}, # no leaves + {"layers": {0: torch.zeros(5, 2)}}, # leaf not on the token axis + {"layers": {0: 1.0}}, # non-tensor leaf + ], +) +def test_full_graph_aux_contract_rejects_invalid_outputs(aux): + runner = object.__new__(OmniGPUModelRunner) + runner.model = SimpleNamespace(_returns_tuple=True, supports_mrv2_full_graph_aux_outputs=True) + runner._configure_cudagraph_output_contract() + with pytest.raises(RuntimeError, match="supports_mrv2_full_graph_aux_outputs"): + runner._flatten_capture_aux_output((torch.zeros(3, 2), aux)) + + +def test_full_graph_aux_contract_rejects_structure_change(): + runner = object.__new__(OmniGPUModelRunner) + runner.model = SimpleNamespace(_returns_tuple=True, supports_mrv2_full_graph_aux_outputs=True) + runner._configure_cudagraph_output_contract() + runner._flatten_capture_aux_output((torch.zeros(3, 2), {"layers": {0: torch.zeros(3, 2)}})) + with pytest.raises(RuntimeError, match="structure changed"): + runner._flatten_capture_aux_output((torch.zeros(3, 2), {"layers": {1: torch.zeros(3, 2)}})) + + +def test_full_graph_output_without_contract_is_hidden_only(): + runner = object.__new__(OmniGPUModelRunner) + hidden = torch.ones(2, 2) + assert runner._split_fullgraph_output(hidden) == (hidden, None) @pytest.mark.parametrize("runner_kind", ["gpu", "ar", "generation"]) @@ -202,6 +280,7 @@ def test_dummy_forward_uses_upstream_execution_state(runner_kind, monkeypatch): else hidden ) runner._dummy_hidden = hidden + runner.model.requires_request_ids = False runner.model_config = SimpleNamespace() runner.vllm_config = SimpleNamespace() runner.req_states = SimpleNamespace() diff --git a/tests/worker_v2/test_omni_model_state.py b/tests/worker_v2/test_omni_model_state.py index 3b585527a..b0178e9a5 100644 --- a/tests/worker_v2/test_omni_model_state.py +++ b/tests/worker_v2/test_omni_model_state.py @@ -19,6 +19,8 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu] class _DummyInputBatch: + is_prefilling_np: np.ndarray + input_ids: SimpleNamespace query_start_loc: torch.Tensor @@ -592,3 +594,58 @@ def test_first_audio_marker_requires_accepted_delivery(monkeypatch, accepted): # A missing route must leave the normal codec path responsible for frame 0; # otherwise it skips the frame and the orchestrator waits forever for it. assert outputs["meta"]["first_audio"].tolist() == ["r0" in accepted, "r1" in accepted] + + +class _FakeEvent: + def record(self, stream=None) -> None: + self.stream = stream + + +def test_publish_sampled_embeddings_for_rows_whose_sample_is_kept(monkeypatch) -> None: + monkeypatch.setattr(torch.cuda, "Event", _FakeEvent) + state = _make_state() + state.model.publishes_sampled_embeddings = True + state.model.embed_input_ids = lambda ids: ids.to(torch.float32).reshape(-1, 1).repeat(1, 3) + # Rows: final prefill chunk, decode, non-final prefill chunk, one-chunk prefill. + batch = SimpleNamespace( + num_reqs=4, + num_scheduled_tokens=[4, 1, 2, 6], + num_computed_prefill_tokens_np=np.array([2, 9, 0, 0], dtype=np.int32), + prefill_len_np=np.array([6, 9, 5, 6], dtype=np.int32), + is_prefilling_np=np.array([True, False, True, True]), + ) + extra, _done = state.publish_sampled_embeddings(batch, torch.tensor([[11], [12], [13], [14]])) + sampled = extra["embed"]["sampled"] + # A non-final prefill chunk's sample is discarded, so it publishes nothing. + assert [tuple(row.shape) for row in sampled] == [(1, 3), (1, 3), (0,), (1, 3)] + assert [row.tolist() for row in sampled if row.numel()] == [[[v] * 3] for v in (11.0, 12.0, 14.0)] + + # Speculative steps sample several tokens per row: not published. + assert state.publish_sampled_embeddings(batch, torch.tensor([[11, 1], [12, 1], [13, 1], [14, 1]])) is None + + +def test_publish_sampled_embeddings_is_opt_in() -> None: + state = _make_state() + state.model.publishes_sampled_embeddings = False + assert state.publish_sampled_embeddings(SimpleNamespace(num_reqs=1), torch.tensor([[1]])) is None + + +@pytest.mark.parametrize("prefilling", [False, True]) +def test_identity_preprocess_skips_only_decode_rows(prefilling): + state = _make_state(has_preprocess=True) + state._decode_preprocess_is_identity = True + state._decode_preprocess = None + state.intermediate_buffer.buffers[0] = {"req_id": "r1"} + seen = [] + + def preprocess(input_ids, input_embeds, **info): + seen.append(info["req_id"]) + return input_ids, input_embeds, {} + + state.model.preprocess = preprocess + batch = _DummyInputBatch([0]) + batch.is_prefilling_np = np.array([prefilling]) + embeds = torch.ones(1, 4) + state.run_preprocess(batch, {"input_ids": torch.tensor([1]), "inputs_embeds": embeds}) + assert seen == (["r1"] if prefilling else []) + assert torch.equal(embeds, torch.ones(1, 4)) diff --git a/tests/worker_v2/test_shared_stage_contracts.py b/tests/worker_v2/test_shared_stage_contracts.py new file mode 100644 index 000000000..58e3f68a1 --- /dev/null +++ b/tests/worker_v2/test_shared_stage_contracts.py @@ -0,0 +1,79 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Generic model opt-in contracts; no model checkpoints required.""" + +from types import SimpleNamespace + +import pytest +import torch + +from vllm_omni.model_executor.models.output_templates import OmniOutput +from vllm_omni.worker_v2.model_states.omni_model_state import OmniModelState +from vllm_omni.worker_v2.omni_ar_model_runner import _merge_payload_trees + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def test_device_output_hook_receives_live_batch_and_request_state(): + state = object.__new__(OmniModelState) + state.have_multimodal_outputs = True + batch, requests, buffers = object(), object(), [{"req_id": "r"}] + hidden = torch.ones(1, 2) + seen = [] + + def make_output(value, **kwargs): + seen.append((value, kwargs)) + return OmniOutput(text_hidden_states=value, multimodal_outputs={"codes": {"audio": value}}) + + state.model = SimpleNamespace(make_omni_output_mrv2=make_output) + state.intermediate_buffer = SimpleNamespace(gather=lambda value: buffers) + result_hidden, payload = state.postprocess_model_output(hidden, batch, requests) + assert result_hidden is hidden + assert payload["codes"]["audio"] is hidden + assert seen == [(hidden, {"input_batch": batch, "req_states": requests, "model_intermediate_buffer": buffers})] + + +def test_sampled_embeddings_merge_preserves_sibling_fields_and_input(): + prefill, sampled = torch.ones(2, 3), torch.zeros(1, 3) + base = {"embed": {"prefill": prefill}, "meta": {"finished": False}} + merged = _merge_payload_trees(base, {"embed": {"sampled": [sampled]}}) + assert merged["embed"]["prefill"] is prefill + assert merged["embed"]["sampled"] == [sampled] + assert "sampled" not in base["embed"] + assert merged["meta"] == {"finished": False} + + +def test_custom_sampler_falls_back_to_parent(monkeypatch): + from vllm.v1.worker.gpu.model_states.default import DefaultModelState + + state = object.__new__(OmniModelState) + state.model = SimpleNamespace() + sampler, marker = object(), object() + monkeypatch.setattr(DefaultModelState, "custom_sampler", lambda self, original: (original, marker)) + assert state.custom_sampler(sampler) == (sampler, marker) + + +def test_logits_vocab_is_adopted_before_parent_sampler_initialization(monkeypatch): + import vllm.v1.worker.gpu.model_runner as upstream + + from vllm_omni.worker_v2 import omni_model_runner as module + + runner = object.__new__(module.OmniGPUModelRunner) + runner.vocab_size = 100 + runner.req_states = SimpleNamespace(vocab_size=100) + runner.model_config = SimpleNamespace(model_arch_config=SimpleNamespace(vocab_size=100)) + model = SimpleNamespace(logits_vocab_size=32) + original_initializer = upstream.init_model_state + sentinel = object() + monkeypatch.setattr(module, "init_omni_model_state", lambda *args: sentinel) + + def load(parent): + assert upstream.init_model_state(None, model, None, None) is sentinel + assert parent.vocab_size == parent.req_states.vocab_size == 32 + assert parent.model_config.model_arch_config.vocab_size == 32 + raise RuntimeError("stop before sampler allocation") + + monkeypatch.setattr(upstream.GPUModelRunner, "load_model", load) + with pytest.raises(RuntimeError, match="stop before sampler allocation"): + runner.load_model() + assert upstream.init_model_state is original_initializer diff --git a/tests/worker_v2/test_shared_stage_cuda.py b/tests/worker_v2/test_shared_stage_cuda.py new file mode 100644 index 000000000..5a8bb72f3 --- /dev/null +++ b/tests/worker_v2/test_shared_stage_cuda.py @@ -0,0 +1,45 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Small CUDA correctness checks for shared stage output contracts.""" + +from types import SimpleNamespace + +import pytest +import torch + +from vllm_omni.worker.gpu_generation_model_runner import _HostCopyBatch +from vllm_omni.worker_v2.omni_model_runner import OmniGPUModelRunner + +pytestmark = [pytest.mark.core_model, pytest.mark.cuda] + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_host_copy_batch_preserves_noncontiguous_outputs_on_current_stream(): + stream = torch.cuda.Stream() + with torch.cuda.stream(stream): + source = torch.arange(24, device="cuda", dtype=torch.float32).reshape(4, 6) + copies = _HostCopyBatch(pin_memory=True) + first = copies.copy(source.t()) + second = copies.copy(source + 10) + copies.wait() + assert first.is_contiguous() and first.is_pinned() + assert torch.equal(first, torch.arange(24, dtype=torch.float32).reshape(4, 6).t()) + assert torch.equal(second, torch.arange(24, dtype=torch.float32).reshape(4, 6) + 10) + copies.wait() # no outstanding copies + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_aux_pytree_tracks_replayed_graph_values(): + runner = object.__new__(OmniGPUModelRunner) + runner.model = SimpleNamespace(_returns_tuple=True, supports_mrv2_full_graph_aux_outputs=True) + runner._configure_cudagraph_output_contract() + source = torch.ones((4, 3), device="cuda") + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + captured = runner._flatten_capture_aux_output((source * 2, {"layers": [source + 1, source + 2]})) + source.fill_(7) + graph.replay() + hidden, aux = runner._split_fullgraph_output(captured) + assert torch.equal(hidden, torch.full_like(source, 14)) + assert torch.equal(aux["layers"][0], torch.full_like(source, 8)) + assert torch.equal(aux["layers"][1], torch.full_like(source, 9)) diff --git a/vllm_omni/config/environment_variable_inventory.py b/vllm_omni/config/environment_variable_inventory.py index 4ae58c00f..7a25034a1 100644 --- a/vllm_omni/config/environment_variable_inventory.py +++ b/vllm_omni/config/environment_variable_inventory.py @@ -251,6 +251,10 @@ _BENCHMARK_TRANSITIONAL = ( ) _INTERNAL = ( + "VLLM_OMNI_CODEC_FIRST_CHUNK_EXPRESS", + "VLLM_OMNI_CODEC_FIRST_CHUNK_EXPRESS_SLACK_S", + "VLLM_OMNI_CONNECTOR_RECV_POLL_MS", + "VLLM_OMNI_SHM_WAKEUP", "VLLM_OMNI_DLO_DP_WAVE_TIMEOUT", "VLLM_OMNI_REPLICA_ID", ) diff --git a/vllm_omni/config/omni_config.py b/vllm_omni/config/omni_config.py index 3a650dd03..fb2695856 100644 --- a/vllm_omni/config/omni_config.py +++ b/vllm_omni/config/omni_config.py @@ -57,6 +57,8 @@ from vllm_omni.config.stage_config import ( normalize_pipeline_cli_overrides, reconcile_diffusion_attention_overrides, resolve_stage_async_chunk, + resolve_stage_model_runner, + validate_native_mrv2_session, validate_stage_async_chunk_edges, ) from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode @@ -1988,7 +1990,9 @@ def _build_model_config( kwargs["final_output"] = topology.final_output if "custom_voice_dir" not in kwargs and deploy.custom_voice_dir is not None: kwargs["custom_voice_dir"] = _copy_value(deploy.custom_voice_dir) - kwargs.setdefault("use_v2_model_runner", deploy.model_runner == "v2") + stage_runner = resolve_stage_model_runner(deploy, stage_deploy) + validate_native_mrv2_session(deploy, topology, stage_runner) + kwargs.setdefault("use_v2_model_runner", stage_runner == "v2") kwargs.setdefault( "supports_native_mrv2_data_plane", topology.supports_native_mrv2_data_plane, diff --git a/vllm_omni/config/stage_config.py b/vllm_omni/config/stage_config.py index 1721a338c..08f588a83 100644 --- a/vllm_omni/config/stage_config.py +++ b/vllm_omni/config/stage_config.py @@ -403,6 +403,9 @@ class StageDeployConfig: # False opts this stage out of pipeline-wide async chunking. async_chunk: bool | None = None + # Overrides the deploy-level ``model_runner`` for this stage, so a pipeline + # can run e.g. its LLM stage on v1 and its codec stages on MRv2. + model_runner: Literal["v1", "v2"] | None = None # Inter-stage connector wiring and request defaults. output_connectors: dict[str, str] | None = None @@ -609,6 +612,7 @@ class DeployConfig: _STAGE_RESERVED_KEYS = frozenset( { "async_chunk", + "model_runner", "stage_id", "devices", "num_replicas", @@ -657,6 +661,9 @@ def _parse_stage_deploy(stage_data: dict[str, Any]) -> StageDeployConfig: else: flat_args[k] = v + stage_runner = stage_data.get("model_runner") + if stage_runner is not None and stage_runner not in ("v1", "v2"): + raise ValueError(f"stage {stage_data['stage_id']}: model_runner must be 'v1' or 'v2', got {stage_runner!r}") kwargs: dict[str, Any] = { "stage_id": stage_data["stage_id"], "devices": devices, @@ -668,6 +675,7 @@ def _parse_stage_deploy(stage_data: dict[str, Any]) -> StageDeployConfig: kwargs[name] = flat_args.pop(name) kwargs["async_chunk"] = stage_data.get("async_chunk") + kwargs["model_runner"] = stage_runner kwargs["output_connectors"] = stage_data.get("output_connectors") kwargs["input_connectors"] = stage_data.get("input_connectors") kwargs["default_sampling_params"] = stage_data.get("default_sampling_params") @@ -889,17 +897,12 @@ def _apply_platform_overrides( if model_runner not in ("v1", "v2"): raise ValueError(f"platform model_runner must be one of ('v1', 'v2'), got {model_runner!r}") deploy.model_runner = model_runner - if deploy.model_runner == "v2" and platform in {"npu", "xpu"}: - raise NotImplementedError( - f"Model Runner V2 is not supported on {platform.upper()}: " - "the platform worker still uses the legacy chunk-transfer data plane." - ) - if platform is None or deploy.platforms is None: - return deploy - if platform_section is None: - return deploy - - platform_stages = platform_section.get("stages", []) + if model_runner == "v1": + # A platform's V1 fallback also covers stages that opt into V2. + for stage in deploy.stages: + if stage.model_runner == "v2": + stage.model_runner = None + platform_stages = platform_section.get("stages", []) if platform_section else [] base_by_id = {s.stage_id: s for s in deploy.stages} for ps in platform_stages: @@ -935,6 +938,18 @@ def _apply_platform_overrides( else: base.engine_extras[key] = val + # Validate the final values, including stage entries from the platform + # overlay. A global V2 default may still apply to omitted pipeline stages. + for stage in deploy.stages: + if stage.model_runner is not None and stage.model_runner not in ("v1", "v2"): + raise ValueError(f"stage {stage.stage_id}: model_runner must be 'v1' or 'v2', got {stage.model_runner!r}") + uses_v2 = deploy.model_runner == "v2" or any(stage.model_runner == "v2" for stage in deploy.stages) + if uses_v2 and platform in {"npu", "xpu"}: + raise NotImplementedError( + f"Model Runner V2 is not supported on {platform.upper()}: " + "the platform worker still uses the legacy chunk-transfer data plane." + ) + return deploy @@ -952,6 +967,28 @@ def _resolve_execution_mode( return _EXECUTION_TYPE_TO_STAGE_WORKER.get(execution_type, (StageType.LLM, None)) +def resolve_stage_model_runner(deploy: DeployConfig, stage: StageDeployConfig | None) -> str: + """The stage's own ``model_runner`` if set, else the deploy-level one.""" + runner = getattr(stage, "model_runner", None) if stage is not None else None + return runner or deploy.model_runner + + +def validate_native_mrv2_session(deploy: DeployConfig, ps: StagePipelineConfig, stage_runner: str) -> None: + """Reject session modes a downstream MRv2 native-data-plane stage lacks. + + Streaming-session prompt replacement exists only in the V1 chunk adapter, + so a stage that receives from an upstream stage on MRv2 supports + turn-based sessions only. + """ + if stage_runner != "v2" or not ps.supports_native_mrv2_data_plane or not ps.input_sources: + return + if deploy.session_mode != "turn": + raise ValueError( + f"stage {ps.stage_id}: model_runner v2 supports session_mode 'turn' only, got " + f"{deploy.session_mode!r}; run this stage with model_runner: v1 for streaming sessions." + ) + + def resolve_stage_async_chunk(deploy: DeployConfig, stage: StageDeployConfig | None) -> bool: if not isinstance(deploy.async_chunk, bool): raise ValueError("async_chunk must be a boolean") @@ -1076,9 +1113,11 @@ def _build_engine_args( "it is derived from the deploy-level `model_runner` field and the " "pipeline's `supports_native_mrv2_data_plane` declaration." ) - engine_args["use_v2_model_runner"] = deploy.model_runner == "v2" + stage_runner = resolve_stage_model_runner(deploy, ds) + validate_native_mrv2_session(deploy, ps, stage_runner) + engine_args["use_v2_model_runner"] = stage_runner == "v2" engine_args["supports_native_mrv2_data_plane"] = bool(ps.supports_native_mrv2_data_plane) - if deploy.model_runner == "v2" and not ps.supports_native_mrv2_data_plane: + if stage_runner == "v2" and not ps.supports_native_mrv2_data_plane: logger.warning( "Stage %s (%s) selects model_runner=v2 without declaring " "supports_native_mrv2_data_plane. It will use the legacy transport path; " diff --git a/vllm_omni/core/sched/omni_generation_scheduler.py b/vllm_omni/core/sched/omni_generation_scheduler.py index 179781c90..cefa2a37a 100644 --- a/vllm_omni/core/sched/omni_generation_scheduler.py +++ b/vllm_omni/core/sched/omni_generation_scheduler.py @@ -3,10 +3,12 @@ from __future__ import annotations +import os import time from collections import defaultdict from typing import TYPE_CHECKING, Any +import torch from vllm.compilation.cuda_graph import CUDAGraphStat from vllm.logger import init_logger from vllm.v1.core.kv_cache_manager import KVCacheBlocks @@ -56,6 +58,23 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler): self._init_omni_io_scheduling_state() self._retains_state_across_chunks = bool(getattr(model_config, "retains_state_across_chunks", False)) self._pending_finish_reqs: list[Request] = [] + self._first_chunk_express = os.environ.get("VLLM_OMNI_CODEC_FIRST_CHUNK_EXPRESS") == "1" + self._last_step_express = False + # Streams that already had a chunk scheduled; the rest await their first. + self._chunk_started: set[str] = set() + # Optional guard: take an express step only while every ready later + # chunk's stream still holds at least this much unplayed audio, so a + # loaded stage does not trade first packets for playback underruns. + # Playback is assumed to start with a stream's first emitted audio. + self._express_min_slack_s = float(os.environ.get("VLLM_OMNI_CODEC_FIRST_CHUNK_EXPRESS_SLACK_S", "0") or 0) + # request id -> [monotonic time of first emitted audio, emitted seconds] + self._stream_audio: dict[str, list[float]] = {} + self._express_skipped_for_slack = 0 + if self._first_chunk_express: + logger.info( + "Generation scheduler decodes first chunks in express steps (min continuation slack %.2f s)", + self._express_min_slack_s, + ) def _build_generation_scheduler_output( self, @@ -250,7 +269,31 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler): # OMNI: Track requests that are already finished (e.g., marked by connector) # These should be removed from running and not scheduled already_finished_reqs: set[Request] = set() - while req_index < len(self.running) and token_budget > 0 and len(num_scheduled_tokens) < execution_batch_size: + # An express step isolates first chunks from ready continuations. + # Two express steps never run back to back; the optional slack guard + # requires playback headroom before delaying a continuation. + first_chunk_express = getattr(self, "_first_chunk_express", False) + express = ( + first_chunk_express + and native_chunks + and not self._last_step_express + and any( + r.request_id not in self._chunk_started + and r.num_in_flight_tokens == 0 + and len(r.prompt_token_ids) > r.num_computed_tokens + for r in self.waiting + ) + ) + if express and getattr(self, "_express_min_slack_s", 0) > 0 and not self._continuations_have_slack(): + express = False + if first_chunk_express: + self._last_step_express = express + while ( + not express + and req_index < len(self.running) + and token_budget > 0 + and len(num_scheduled_tokens) < execution_batch_size + ): request = self.running[req_index] # OMNI: Skip requests that are not in self.requests if request.request_id not in self.requests or ( @@ -327,6 +370,10 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler): and self._pause_state == PauseState.UNPAUSED ): request = self.waiting.peek_request() + if express and request.request_id in self._chunk_started: + self.waiting.pop_request() + skipped_waiting_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. @@ -398,6 +445,9 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler): token_budget -= num_new_tokens scheduled_new_reqs.append(request) + if first_chunk_express: + self._chunk_started.update(num_scheduled_tokens) + # Return skipped waiting requests if skipped_waiting_requests: self.waiting.prepend_requests(skipped_waiting_requests) @@ -493,9 +543,52 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler): return self._wrap_omni_scheduler_output(scheduler_output) + def _continuations_have_slack(self) -> bool: + """Whether every started stream with a chunk ready to decode can wait one express step.""" + now = time.monotonic() + for request in (*self.running, *self.waiting): + if request.request_id not in self._chunk_started or request.num_in_flight_tokens > 0: + continue + if len(request.prompt_token_ids) <= request.num_computed_tokens: + continue + audio = self._stream_audio.get(request.request_id) + if audio is None: + # No known playback credit: do not delay this continuation. + return False + if audio[0] + audio[1] - now < self._express_min_slack_s: + self._express_skipped_for_slack += 1 + if self._express_skipped_for_slack in (1, 100, 10000): + logger.info( + "First-chunk express deferred: a stream holds under %.2f s of audio (%d times)", + self._express_min_slack_s, + self._express_skipped_for_slack, + ) + return False + return True + + def _record_stream_audio(self, request_id: str, mm_output: Any) -> None: + """Accumulate the seconds of audio each stream has emitted (``model_outputs`` at ``sr``).""" + if not isinstance(mm_output, dict): + return + audio = mm_output.get("model_outputs") + rate = mm_output.get("sr") + # Only mono sample vectors have an unambiguous duration here. + # Unknown/multichannel layouts earn no slack credit. + if not isinstance(audio, torch.Tensor) or audio.ndim != 1 or audio.numel() == 0: + return + if isinstance(rate, torch.Tensor): + rate = int(rate.reshape(-1)[0]) if rate.numel() else 0 + if not isinstance(rate, int) or rate <= 0: + return + entry = self._stream_audio.setdefault(request_id, [time.monotonic(), 0.0]) + entry[1] += audio.numel() / rate + def _free_request( self, request: Request, delay_free_blocks: bool = False ) -> tuple[dict[str, Any] | None, dict[str, Any] | None]: + if getattr(self, "_first_chunk_express", False): + self._chunk_started.discard(request.request_id) + getattr(self, "_stream_audio", {}).pop(request.request_id, None) if self.input_coordinator is None: return super()._free_request(request, delay_free_blocks) @@ -641,6 +734,8 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler): ec_transfer_params = None pooler_output = pooler_outputs[req_index] if pooler_outputs else None mm_output = mm_outputs[req_index] if mm_outputs else None + if getattr(self, "_express_min_slack_s", 0) > 0 and self._first_chunk_express: + self._record_stream_audio(req_id, mm_output) status_before_stop = request.status finish_reason = None is_segment_finished = False diff --git a/vllm_omni/core/sched/omni_scheduler_mixin.py b/vllm_omni/core/sched/omni_scheduler_mixin.py index 0f17fb54d..1a1a46276 100644 --- a/vllm_omni/core/sched/omni_scheduler_mixin.py +++ b/vllm_omni/core/sched/omni_scheduler_mixin.py @@ -8,7 +8,7 @@ import os import queue import time from collections.abc import Iterable, Iterator -from typing import Any +from typing import TYPE_CHECKING, Any import torch from vllm.compilation.cuda_graph import CUDAGraphStat @@ -44,6 +44,7 @@ from vllm_omni.core.sched.output import ( from vllm_omni.distributed.omni_connectors.transfer_adapter.chunk_transfer_adapter import ( OmniChunkTransferAdapter, ) +from vllm_omni.distributed.omni_connectors.utils.config import stage_receives_chunks from vllm_omni.engine import OmniEngineCoreOutput from vllm_omni.engine.serialization import serialize_additional_information from vllm_omni.outputs import OmniConnectorOutput @@ -104,7 +105,13 @@ elif DEFAULT_INPUT_WAIT_TIMEOUT_S == 0: ) -class OmniSchedulerMixin: +if TYPE_CHECKING: + from vllm.v1.core.sched.scheduler import Scheduler as _SchedulerMixinBase +else: + _SchedulerMixinBase = object + + +class OmniSchedulerMixin(_SchedulerMixinBase): """Shared scheduler helpers for omni-specific request handling.""" # Provided by the concrete vLLM scheduler. @@ -127,17 +134,12 @@ class OmniSchedulerMixin: def _init_omni_io_scheduling_state(self) -> None: """Initialize scheduler state shared by AR and generation stages.""" - # Every omni worker forces the v1 model runner (the v2 runner carries no - # omni hooks), so the scheduler has to agree or the two disagree about - # what the SchedulerOutput carries. vLLM 0.29 defaults this to v2 where - # 0.28 did not, so the scheduler took its v2 fast path -- the one that - # deliberately omits resumed-request token ids -- while the v1 runner - # still read them, and resuming a request raised ``KeyError: `` - # in ``_update_states``. - if getattr(self, "use_v2_model_runner", False): - logger.warning("OMNI scheduler forces v1 model runner for omni hooks.") - self.use_v2_model_runner = False - + # The scheduler must agree with the stage's worker about which runner + # consumes its SchedulerOutput: vLLM's own default (v2 since 0.29) would + # take the v2 fast path, which omits resumed-request token ids, while a + # v1 omni runner still reads them (``KeyError: `` in + # ``_update_states``). The stage's model config carries the choice + # (deploy ``model_runner``, optionally per stage). model_config = self.vllm_config.model_config self.use_v2_model_runner = bool(getattr(model_config, "use_v2_model_runner", False)) self._native_data_plane = uses_native_mrv2_data_plane( @@ -150,7 +152,14 @@ class OmniSchedulerMixin: else None ) self.input_coordinator: OmniSchedulingCoordinator | None = None - if self._native_data_plane and getattr(model_config, "async_chunk", False): + if ( + self._native_data_plane + and getattr(model_config, "async_chunk", False) + # A downstream sender-only stage (e.g. MiniCPM-o's Talker, fed by + # the orchestrator) has no connector input to wait for; parking + # its requests for chunks would never release them. + and (getattr(model_config, "stage_id", 0) == 0 or stage_receives_chunks(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), @@ -825,6 +834,8 @@ class OmniSchedulerMixin: def _resume_downstream_chunk_receiver(self, request: Request) -> None: """Resume duplex connector polling without an external update.""" adapter = self.chunk_transfer_adapter + if adapter is None: + return adapter.segment_finished_requests.discard(request.request_id) if ( not adapter.receives_chunks diff --git a/vllm_omni/data_entry_keys.py b/vllm_omni/data_entry_keys.py index e7c27fa70..e18e3c9cb 100644 --- a/vllm_omni/data_entry_keys.py +++ b/vllm_omni/data_entry_keys.py @@ -41,6 +41,8 @@ class Embeddings(TypedDict, total=False): prepared_prefill: torch.Tensor prefill: torch.Tensor decode: torch.Tensor + # [1, H] embedding of the token sampled by a row's final prefill step (MRv2 producers). + sampled: torch.Tensor decode_token_start: int decode_token_end: int cached_decode: torch.Tensor @@ -88,6 +90,7 @@ class OmniPayloadMeta(TypedDict, total=False): right_holdback_size: int override_keys: list[tuple[str, str]] num_processed_tokens: int + resumable: bool next_stage_prompt_len: int next_stage_generation_tokens: int replace_streaming_prompt: bool @@ -202,6 +205,8 @@ class MetaStruct(_StructBase): right_holdback_size: int | None = None override_keys: list[tuple[str, str]] | None = None num_processed_tokens: int | None = None + # The model runner sets this when a streaming request is resumed. + resumable: bool | None = None next_stage_prompt_len: int | None = None next_stage_generation_tokens: int | None = None replace_streaming_prompt: bool | None = None diff --git a/vllm_omni/distributed/omni_connectors/connectors/shm_connector.py b/vllm_omni/distributed/omni_connectors/connectors/shm_connector.py index 3ad69e661..09854d22d 100644 --- a/vllm_omni/distributed/omni_connectors/connectors/shm_connector.py +++ b/vllm_omni/distributed/omni_connectors/connectors/shm_connector.py @@ -2,9 +2,14 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project import fcntl +import glob +import hashlib import os +import select +import stat import threading import time +import uuid from collections import OrderedDict from multiprocessing import shared_memory as shm_pkg from typing import Any @@ -17,6 +22,15 @@ from .base import OmniConnectorBase logger = get_connector_logger(__name__) +def _wakeup_enabled() -> bool: + return os.environ.get("VLLM_OMNI_SHM_WAKEUP", "1") == "1" + + +def _wakeup_directory() -> str: + # Stage engine processes of one deployment share their launching parent. + return f"/dev/shm/omni_shm_wake_{os.getuid()}_{os.getppid()}" + + class SharedMemoryConnector(OmniConnectorBase): """Key-addressed local shared-memory connector. @@ -25,6 +39,12 @@ class SharedMemoryConnector(OmniConnectorBase): remote-transport metadata such as ``source_host`` / ``source_port`` (that is the RDMA connector's job). When such metadata is passed in, the connector silently falls back to key-based lookup. + + Arrival wakeups: each receiving connector owns a unique named FIFO. A + ``put`` broadcasts one byte to every FIFO for the destination stage, so the receive loop blocks in + ``wait_for_change`` until data arrives instead of re-polling on a fixed + interval. Wakeups are hints: a stage that cannot reach the FIFO (another + deployment layout, or ``VLLM_OMNI_SHM_WAKEUP=0``) keeps the timed poll. """ def get_with_deadline(self, from_stage, to_stage, get_key, metadata=None, *, deadline): @@ -35,6 +55,12 @@ class SharedMemoryConnector(OmniConnectorBase): def __init__(self, config: dict[str, Any]): self.config = config self.stage_id = config.get("stage_id", -1) + scope = config.get("extra", {}).get("wakeup_scope") + self._wake_directory = ( + f"/dev/shm/omni_shm_wake_{os.getuid()}_{hashlib.sha256(str(scope).encode()).hexdigest()[:24]}" + if scope is not None + else _wakeup_directory() + ) self._pending_keys: OrderedDict[str, None] = OrderedDict() self._pending_keys_lock = threading.Lock() self._metrics = { @@ -42,6 +68,129 @@ class SharedMemoryConnector(OmniConnectorBase): "gets": 0, "bytes_transferred": 0, } + # Receiver side: FIFO read end (plus a write end of our own, so the + # FIFO never reports EOF when no sender has it open) and a counter of + # drained wakeups. Each receiver owns its path, including across restarts. + self._wake_lock = threading.Lock() + self._wake_read_fd: int | None = None + self._wake_hold_fd: int | None = None + self._wake_path: str | None = None + self._wake_generation = 0 + self._wake_closed = False + + def _open_wakeup_receiver(self) -> bool: + if self._wake_closed: + return False + if self._wake_read_fd is not None: + return True + if not _wakeup_enabled(): + return False + try: + if int(self.stage_id) < 0: + return False + path = f"{self._wake_directory}/{int(self.stage_id)}_{uuid.uuid4().hex}" + except (TypeError, ValueError): + return False + read_fd = hold_fd = None + created = False + try: + os.makedirs(os.path.dirname(path), mode=0o700, exist_ok=True) + os.mkfifo(path, 0o600) + created = True + read_fd = os.open(path, os.O_RDONLY | os.O_NONBLOCK) + hold_fd = os.open(path, os.O_WRONLY | os.O_NONBLOCK) + except OSError as e: + for fd in (read_fd, hold_fd): + if fd is not None: + os.close(fd) + if created: + try: + os.unlink(path) + except FileNotFoundError: + pass + logger.debug("SHM wakeup FIFO unavailable at %s: %s", path, e) + return False + self._wake_read_fd, self._wake_hold_fd, self._wake_path = read_fd, hold_fd, path + return True + + def get_wakeup_generation(self) -> int | None: + """Snapshot before polling; a later arrival changes it (see ``wait_for_change``). + + ``None`` when this stage has no wakeup FIFO: the caller keeps its timed poll. + """ + with self._wake_lock: + return self._wake_generation if self._open_wakeup_receiver() else None + + def wait_for_change(self, generation: int, *, timeout: float) -> bool: + """Block until a ``put`` to this stage since ``generation``, or ``timeout``.""" + with self._wake_lock: + read_fd = self._wake_read_fd + if self._wake_closed or read_fd is None: + return False + if self._wake_generation == generation: + try: + readable, _, _ = select.select([read_fd], [], [], timeout) + except (OSError, ValueError): + return False + if not readable: + return False + with self._wake_lock: + if self._wake_closed or self._wake_read_fd != read_fd: + return False + try: + while os.read(read_fd, 4096): + pass + except BlockingIOError: + pass + except OSError: + return False + self._wake_generation += 1 + return self._wake_generation != generation + + def _wake_receiver(self, to_stage: Any) -> None: + if not _wakeup_enabled(): + return + try: + pattern = f"{self._wake_directory}/{int(to_stage)}_*" + except (TypeError, ValueError): + return # non-numeric stage names use the ordinary polling path + with self._wake_lock: + if self._wake_closed: + return + # Discover every replica, including newly restarted receivers. Do + # not cache writers across receiver lifetimes or unlink peers' paths. + for path in glob.iglob(pattern): + fd = None + try: + fd = os.open(path, os.O_WRONLY | os.O_NONBLOCK | os.O_NOFOLLOW) + if stat.S_ISFIFO(os.fstat(fd).st_mode): + os.write(fd, b"\0") + except OSError: + # Missing reader/file or full pipe: timed polling remains + # the correctness fallback. A full FIFO already has hints. + pass + finally: + if fd is not None: + os.close(fd) + + def _close_wakeups(self) -> None: + with self._wake_lock: + self._wake_closed = True + for fd in (self._wake_read_fd, self._wake_hold_fd): + if fd is not None: + os.close(fd) + self._wake_read_fd = self._wake_hold_fd = None + if self._wake_path is not None: + try: + os.unlink(self._wake_path) + except FileNotFoundError: + pass + self._wake_path = None + try: + os.rmdir(self._wake_directory) + except OSError: + # Other receivers may still own FIFOs in this deployment. + pass def put( self, @@ -65,6 +214,7 @@ class SharedMemoryConnector(OmniConnectorBase): metadata = {"shm": meta, "size": size} with self._pending_keys_lock: self._pending_keys[put_key] = None + self._wake_receiver(to_stage) self._metrics["puts"] += 1 self._metrics["bytes_transferred"] += size @@ -184,6 +334,7 @@ class SharedMemoryConnector(OmniConnectorBase): keys = list(self._pending_keys) for key in keys: self.cleanup(key) + self._close_wakeups() def reap_consumed(self) -> None: """Bounded round-robin sweep; receivers unlink SHM in another process.""" diff --git a/vllm_omni/distributed/omni_connectors/model_runner/omni_connector_payload_transport.py b/vllm_omni/distributed/omni_connectors/model_runner/omni_connector_payload_transport.py index 265a67dca..76ca7eaa5 100644 --- a/vllm_omni/distributed/omni_connectors/model_runner/omni_connector_payload_transport.py +++ b/vllm_omni/distributed/omni_connectors/model_runner/omni_connector_payload_transport.py @@ -6,6 +6,8 @@ from __future__ import annotations import importlib import inspect +import math +import os import time from collections import deque from collections.abc import Callable @@ -22,6 +24,23 @@ from vllm_omni.distributed.omni_connectors.model_runner.omni_connector_runtime i ) from vllm_omni.outputs import OmniConnectorOutput + +# No-progress recheck for connectors without a change notification (SHM polls). +def _recv_poll_seconds() -> float: + """Resolve a finite positive poll interval, falling back for invalid input.""" + raw = os.environ.get("VLLM_OMNI_CONNECTOR_RECV_POLL_MS", "5") + try: + value = float(raw) / 1000 + if math.isfinite(value) and value > 0: + return value + except ValueError: + pass + logger.warning("Invalid VLLM_OMNI_CONNECTOR_RECV_POLL_MS=%r; using 5 ms", raw) + return 0.005 + + +_RECV_POLL_S = _recv_poll_seconds() + if TYPE_CHECKING: from vllm_omni.distributed.omni_connectors.connectors.base import ( OmniConnectorBase, @@ -189,10 +208,12 @@ class _OmniConnectorPayloadTransportMixin(_OmniConnectorRuntimeMixin): connector: OmniConnectorBase, from_stage: str, to_stage: str, - connector_get_key: str, + connector_get_key: str | None, metadata: dict[str, Any] | None = None, ) -> Any: """Receive one ordinary non-KV stage payload on the local leader rank only.""" + if connector_get_key is None: + return None if not self._stage_payload_broadcast_groups(): if metadata is None: return connector.get(from_stage, to_stage, connector_get_key) @@ -208,7 +229,7 @@ class _OmniConnectorPayloadTransportMixin(_OmniConnectorRuntimeMixin): connector: OmniConnectorBase, from_stage: str, to_stage: str, - connector_get_key: str, + connector_get_key: str | None, metadata: dict[str, Any] | None = None, ) -> Any: """Receive one full-payload transfer on the local leader rank only.""" @@ -225,7 +246,7 @@ class _OmniConnectorPayloadTransportMixin(_OmniConnectorRuntimeMixin): connector: OmniConnectorBase, from_stage: str, to_stage: str, - connector_get_key: str, + connector_get_key: str | None, metadata: dict[str, Any] | None = None, ) -> Any: """Receive one ordinary async chunk on the local leader rank only.""" @@ -803,7 +824,7 @@ class _OmniConnectorPayloadTransportMixin(_OmniConnectorRuntimeMixin): self._pending_save_counts[req_id] += 1 sent_ids.append(req_id) if sent_ids: - self._work_available.set() + self._save_work_available.set() return sent_ids # ------------------------------------------------------------------ # @@ -930,6 +951,8 @@ class _OmniConnectorPayloadTransportMixin(_OmniConnectorRuntimeMixin): """Background thread: poll connector for incoming data.""" _recv_poll_count = 0 while not self._stop_event.is_set(): + snapshot = getattr(self._omni_connector, "get_wakeup_generation", None) + generation = snapshot() if callable(snapshot) else None with self._lock: pending_ids = list(self._pending_load_reqs.keys()) @@ -951,8 +974,16 @@ class _OmniConnectorPayloadTransportMixin(_OmniConnectorRuntimeMixin): made_progress = self._poll_pending_requests_once(pending_ids) if not made_progress and not self._stop_event.is_set(): - self._work_available.wait(timeout=0.005) - self._work_available.clear() + # Wait for a *new* arrival: existing future chunks may be + # blocked by scheduler ownership of the previous payload. + # A key-presence predicate would spin and compete for the + # GIL. Honor the configured readiness recheck interval. + wait_for_change = getattr(self._omni_connector, "wait_for_change", None) + if generation is not None and callable(wait_for_change): + wait_for_change(generation, timeout=_RECV_POLL_S) + else: + self._work_available.wait(timeout=_RECV_POLL_S) + self._work_available.clear() _MAX_SEND_RETRIES = 3 @@ -986,8 +1017,8 @@ class _OmniConnectorPayloadTransportMixin(_OmniConnectorRuntimeMixin): self._requeue_or_drop_failed_send(task, error=send_error) continue - self._work_available.wait(timeout=0.01) - self._work_available.clear() + self._save_work_available.wait(timeout=0.01) + self._save_work_available.clear() def _requeue_or_drop_failed_send( self, @@ -1547,7 +1578,7 @@ class _OmniConnectorPayloadTransportMixin(_OmniConnectorRuntimeMixin): with self._lock: self._pending_save_reqs.setdefault(request_id, deque()).append(task) self._pending_save_counts[request_id] += 1 - self._work_available.set() + self._save_work_available.set() return True, completion def _poll_pending_requests_once(self, pending_ids: list[str]) -> bool: @@ -1564,8 +1595,25 @@ class _OmniConnectorPayloadTransportMixin(_OmniConnectorRuntimeMixin): self.publish_omni_connector_output_to_sink() return made_progress + @staticmethod + def _finish_marker_payload() -> OmniPayload: + return {"meta": {"finished": torch.tensor(True, dtype=torch.bool)}} + + @staticmethod + def _request_is_finished(request: Any) -> bool: + is_finished = getattr(request, "is_finished", None) + return bool(is_finished()) if callable(is_finished) else False + def _publish_chunk_cohort(self, entries: list[tuple[Any, Any]], *, wait_for_delivery: bool) -> int: - entries = [(request, payload) for request, payload in entries if payload is not None] + # A request's last chunk must carry its finish marker even when the + # processor has nothing left to send (e.g. its frames ended exactly on + # a chunk boundary), as the V1 chunk adapter guarantees; otherwise the + # receiver waits for input until its deadline. + entries = [ + (request, payload if payload is not None else self._finish_marker_payload()) + for request, payload in entries + if payload is not None or self._request_is_finished(request) + ] if not entries: return 0 emitted = 0 diff --git a/vllm_omni/distributed/omni_connectors/model_runner/omni_connector_runtime.py b/vllm_omni/distributed/omni_connectors/model_runner/omni_connector_runtime.py index 3349c2166..6667c0e4b 100644 --- a/vllm_omni/distributed/omni_connectors/model_runner/omni_connector_runtime.py +++ b/vllm_omni/distributed/omni_connectors/model_runner/omni_connector_runtime.py @@ -103,6 +103,10 @@ class _SendCompletion: class _OmniConnectorRuntimeMixin: """Own connector lifecycle, shared state, and KV transfer delegation.""" + if TYPE_CHECKING: + + def _stage_payload_broadcast_groups(self) -> tuple[Any, ...]: ... + _omni_connector: Any _kv_transfer_manager: Any _async_chunk: bool @@ -145,6 +149,7 @@ class _OmniConnectorRuntimeMixin: _lock: Any _stop_event: Any _work_available: Any + _save_work_available: Any _recv_thread: Any _save_thread: Any _omni_connector_initialized: bool @@ -313,6 +318,10 @@ class _OmniConnectorRuntimeMixin: self._omni_connector_output_drain_lock = threading.Lock() self._stop_event = threading.Event() self._work_available = threading.Event() + # The save loop has its own event: both loops clear theirs after a + # wait, so a shared event let the receiver swallow a send wakeup and + # left the payload waiting for the saver's 10 ms timeout. + self._save_work_available = threading.Event() # Start background threads only when there's a connector self._recv_thread: threading.Thread | None = None @@ -343,6 +352,8 @@ class _OmniConnectorRuntimeMixin: if getattr(self, "_synchronous_payload_transport", False): return self._stop_event.set() + self._work_available.set() + self._save_work_available.set() if self._recv_thread is not None: self._recv_thread.join(timeout=5) if self._save_thread is not None: diff --git a/vllm_omni/distributed/omni_connectors/utils/initialization.py b/vllm_omni/distributed/omni_connectors/utils/initialization.py index d0139b59a..7c12ab6a7 100644 --- a/vllm_omni/distributed/omni_connectors/utils/initialization.py +++ b/vllm_omni/distributed/omni_connectors/utils/initialization.py @@ -1,10 +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 """Utilities for OmniConnector configuration and validation.""" import json import os +import uuid import warnings from pathlib import Path from typing import TYPE_CHECKING, Any, Literal @@ -195,7 +196,9 @@ def create_connectors_from_config( replica_id = max(int(os.environ.get("VLLM_OMNI_REPLICA_ID", "0")), 0) resolved_spec = resolve_connector_spec( ConnectorSpec(name=connector_spec.name, extra=extra), - stage_id=int(caller_stage_id) if str(caller_stage_id).isdigit() else from_stage_id, + stage_id=int(caller_stage_id) + if caller_stage_id is not None and str(caller_stage_id).isdigit() + else from_stage_id, role=role, purpose=resolved_purpose, local_rank=local_rank, @@ -308,7 +311,7 @@ def load_omni_transfer_config( config_dict = normalized # Parse connectors - connectors = {} + connectors: dict[tuple[str, str], ConnectorSpec] = {} runtime_config = config_dict.get("runtime", {}) # Parse global connectors (from runtime.connectors) @@ -429,6 +432,12 @@ def load_omni_transfer_config( f"{missing_str}. Define connectors or allow auto SHM creation for these edges." ) + # Resolve once in the launcher; serialized connector extras carry the same + # namespace through EngineCore and multiprocess/TP workers. + wakeup_scope = uuid.uuid4().hex + for spec in connectors.values(): + if spec.name == "SharedMemoryConnector": + spec.extra.setdefault("wakeup_scope", wakeup_scope) config = OmniTransferConfig(connectors=connectors) logger.info(f"Loaded OmniTransferConfig with {len(connectors)} connector configurations") diff --git a/vllm_omni/engine/orchestrator.py b/vllm_omni/engine/orchestrator.py index 64eee1cff..4047a7147 100644 --- a/vllm_omni/engine/orchestrator.py +++ b/vllm_omni/engine/orchestrator.py @@ -34,6 +34,7 @@ from vllm.v1.engine import EngineCoreOutputs, FinishReason from vllm.v1.engine.exceptions import EngineDeadError from vllm.v1.metrics.stats import IterationStats +from vllm_omni.core.sched.omni_scheduling_coordinator import uses_native_mrv2_data_plane from vllm_omni.data_entry_keys import FIRST_AUDIO_KEY, FIRST_AUDIO_REQUIRED_KEY from vllm_omni.diffusion.data import is_diffusion_request_started_output from vllm_omni.distributed.omni_connectors.utils.config import stage_receives_chunks @@ -2907,6 +2908,18 @@ class Orchestrator(OrchestratorBase): return False return True + def _native_mrv2_receiver_stage(self, final_stage_id: int) -> int | None: + """First downstream stage up to ``final_stage_id`` that receives on MRv2's native data plane.""" + for stage_id in range(1, min(final_stage_id, len(self.stage_pools) - 1) + 1): + vllm_config = getattr(self.stage_pools[stage_id], "stage_vllm_config", None) + model_config = getattr(vllm_config, "model_config", None) + if uses_native_mrv2_data_plane( + model_config, + use_v2_model_runner=bool(getattr(model_config, "use_v2_model_runner", False)), + ): + return stage_id + return None + async def _handle_add_request(self, msg: StageSubmissionMessage) -> None: """Handle an add_request message from the main thread.""" stage_id = 0 @@ -2926,6 +2939,20 @@ class Orchestrator(OrchestratorBase): await self._fail_request_dead_stage(request_id, stage_id) return + if getattr(prompt, "resumable", False): + mrv2_stage_id = self._native_mrv2_receiver_stage(final_stage_id) + if mrv2_stage_id is not None: + # Streaming-input sessions replace downstream prompts through + # the V1 chunk adapter, which MRv2's native data plane lacks. + await self._fail_request_client_error( + request_id, + stage_id, + f"streaming (resumable) input is not supported: stage {mrv2_stage_id} runs on " + "model_runner v2, which supports turn-based requests only; use model_runner: v1 " + "for that stage", + ) + return + logger.debug( "[Orchestrator] _handle_add_request: stage=%s req=%s " "prompt_type=%s original_prompt_type=%s final_stage=%s " diff --git a/vllm_omni/utils/device_copy.py b/vllm_omni/utils/device_copy.py index 2b271940f..973055547 100644 --- a/vllm_omni/utils/device_copy.py +++ b/vllm_omni/utils/device_copy.py @@ -20,6 +20,14 @@ from collections.abc import Sequence import torch +def to_device_nonblocking(tensor: torch.Tensor, device: torch.device | str) -> torch.Tensor: + """``tensor.to(device)`` without a host sync when ``tensor`` is on the CPU and ``device`` is CUDA.""" + device = torch.device(device) + if tensor.device.type != "cpu" or device.type != "cuda": + return tensor.to(device) + return tensor.pin_memory().to(device, non_blocking=True) + + def index_to_device(values: Sequence[int], device: torch.device | str, dtype: torch.dtype = torch.long) -> torch.Tensor: """A host index list as a device tensor, without a host sync.""" if torch.device(device).type != "cuda": diff --git a/vllm_omni/worker/gpu_ar_model_runner.py b/vllm_omni/worker/gpu_ar_model_runner.py index 28af89a6f..ad3e50d52 100644 --- a/vllm_omni/worker/gpu_ar_model_runner.py +++ b/vllm_omni/worker/gpu_ar_model_runner.py @@ -309,6 +309,10 @@ def _ensure_tensor_values(payload: dict[str, object]) -> dict[str, torch.Tensor] """ result: dict[str, torch.Tensor] = {} for key, val in payload.items(): + # Sparse per-request conditioning uses None for requests with no + # update. Absence is expected, not an unsupported wire value. + if val is None: + continue if isinstance(val, torch.Tensor): result[key] = val elif isinstance(val, (int, float, bool)): @@ -1817,6 +1821,7 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin, Duplex if not needs_pooler_payload and prefix_cache_step_id is not None: # No consumer for this step's merge: consume the step context by # id (exactly-once contract). The cache write still lands. + assert self.omni_prefix_cache is not None self.omni_prefix_cache.discard_step(prefix_cache_step_id) prefix_cache_step_id = None if self.omni_prefix_cache is None and needs_scheduled_hidden_payload and not audio_sparse_output: diff --git a/vllm_omni/worker/gpu_generation_model_runner.py b/vllm_omni/worker/gpu_generation_model_runner.py index c82cf71e2..a05e12708 100644 --- a/vllm_omni/worker/gpu_generation_model_runner.py +++ b/vllm_omni/worker/gpu_generation_model_runner.py @@ -13,6 +13,7 @@ import gc import logging from collections.abc import Mapping from dataclasses import replace +from typing import cast import numpy as np import torch @@ -23,6 +24,7 @@ from vllm.distributed.parallel_state import get_pp_group from vllm.forward_context import set_forward_context from vllm.sequence import IntermediateTensors from vllm.utils.math_utils import cdiv +from vllm.utils.platform_utils import is_pin_memory_available from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput from vllm.v1.outputs import EMPTY_MODEL_RUNNER_OUTPUT, AsyncModelRunnerOutput, make_empty_encoder_model_runner_output from vllm.v1.spec_decode.dflash import DFlashProposer @@ -51,6 +53,36 @@ from vllm_omni.worker.omni_connector_model_runner_mixin import ( logger = logging.getLogger(__name__) +class _HostCopyBatch: + """Device-to-host copies of one step's outputs behind a single host sync. + + ``tensor.to("cpu")`` blocks the host once per tensor, so a step returning + one waveform per request paid one sync per request. Here each CUDA tensor + is copied into pinned host memory without blocking and ``wait`` blocks + once, after which every returned tensor holds its data, exactly like the + per-tensor ``.detach().to("cpu").contiguous()`` it replaces. Without + pinned memory, or off CUDA, the per-tensor blocking copy is kept. + """ + + def __init__(self, pin_memory: bool) -> None: + self._pin_memory = bool(pin_memory) + self._pending = False + + def copy(self, tensor: torch.Tensor) -> torch.Tensor: + tensor = tensor.detach() + if tensor.device.type != "cuda" or not self._pin_memory: + return tensor.to("cpu").contiguous() + host = torch.empty(tensor.shape, dtype=tensor.dtype, device="cpu", pin_memory=True) + host.copy_(tensor, non_blocking=True) + self._pending = True + return host + + def wait(self) -> None: + if self._pending: + torch.cuda.current_stream().synchronize() + self._pending = False + + class GPUGenerationModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin): """Generation model runner for vLLM-Omni (non-autoregressive). @@ -59,6 +91,8 @@ class GPUGenerationModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin - Executes generation process and returns tensors via `pooler_output`. """ + execute_model_state: ExecuteModelState | None + def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) self._async_chunk = getattr(self.model_config, "async_chunk", False) @@ -419,6 +453,8 @@ class GPUGenerationModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin # Build per-request multimodal_outputs list (dedicated channel). # pooler_output is no longer used for multimodal data. per_req_payloads: list[dict[str, object]] = [] + # One host sync for the whole step instead of one per request tensor. + to_host = _HostCopyBatch(is_pin_memory_available()) if isinstance(multimodal_outputs_raw, torch.Tensor): # One row per request. The old asserts (`shape[0] == 1` AND # `shape[0] == num_reqs`) jointly forced num_reqs == 1, silently @@ -431,7 +467,7 @@ class GPUGenerationModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin "to return multiple tensors per request, use a dict)." ) for i in range(num_reqs): - per_req_payloads.append({"model_outputs": multimodal_outputs_raw[i].detach().to("cpu").contiguous()}) + per_req_payloads.append({"model_outputs": to_host.copy(multimodal_outputs_raw[i])}) elif isinstance(multimodal_outputs_raw, list): # One entry per request. The old `len == 1` assert did not check # num_reqs, so a batched step built a length-1 payload list that @@ -444,9 +480,7 @@ class GPUGenerationModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin "to return multiple lists per request, use a dict)." ) for out in multimodal_outputs_raw: - per_req_payloads.append( - {"model_outputs": out.detach().to("cpu").contiguous() if out is not None else None} - ) + per_req_payloads.append({"model_outputs": to_host.copy(out) if out is not None else None}) elif isinstance(multimodal_outputs_raw, Mapping): num_reqs = self.input_batch.num_reqs for i in range(num_reqs): @@ -458,21 +492,26 @@ class GPUGenerationModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin f"Multimodal output list for key '{key}' has length {len(out)} " f"but expected {num_reqs} (one entry per request)." ) - mm_payload[key] = out[i].detach().to("cpu").contiguous() + mm_payload[key] = to_host.copy(out[i]) elif isinstance(out, torch.Tensor): - mm_payload[key] = out.detach().to("cpu").contiguous() + mm_payload[key] = to_host.copy(out) else: logger.warning(f"Unsupported multimodal output type for key '{key}': {type(out)}") per_req_payloads.append(_ensure_tensor_values(mm_payload)) else: raise RuntimeError("Unsupported diffusion output type") + # Sync mode hands these payloads out as finished host tensors. + to_host.wait() + inter_stage_outputs: list[dict[str, object] | None] | None + multimodal_outputs: list[dict[str, object] | None] | None if self._async_chunk: inter_stage_outputs, multimodal_outputs = partition_payload_list(per_req_payloads) else: # See gpu_ar_model_runner: non-async-chunk ships the full payload to the next # stage; #4527's (None, per_req_payloads) starved the downstream stage. (PR #4792) - inter_stage_outputs, multimodal_outputs = per_req_payloads, per_req_payloads + inter_stage_outputs = cast(list[dict[str, object] | None], per_req_payloads) + multimodal_outputs = inter_stage_outputs # [Omni] Copy req_id mappings to avoid async scheduling mutation. req_ids_output_copy = self.input_batch.req_ids.copy() @@ -493,7 +532,8 @@ class GPUGenerationModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin logprobs=None, prompt_logprobs_dict={}, pooler_output=None, - multimodal_outputs=multimodal_outputs, + # Preserve absent per-row client payloads; the scheduler distinguishes None from {}. + multimodal_outputs=multimodal_outputs, # type: ignore[arg-type] inter_stage_outputs=inter_stage_outputs, kv_connector_output=kv_connector_output, num_nans_in_logits={}, diff --git a/vllm_omni/worker/gpu_model_runner.py b/vllm_omni/worker/gpu_model_runner.py index 210550f2c..9cd81fb1a 100644 --- a/vllm_omni/worker/gpu_model_runner.py +++ b/vllm_omni/worker/gpu_model_runner.py @@ -41,6 +41,7 @@ from vllm_omni.model_executor.layers.rotary_embedding.mrope import OmniMRotaryEm from vllm_omni.model_executor.models.model_local_kv import collect_model_local_kv_specs from vllm_omni.model_executor.models.output_templates import OmniOutput from vllm_omni.platforms import current_omni_platform +from vllm_omni.utils.device_copy import index_to_device if TYPE_CHECKING: from vllm.v1.core.sched.output import SchedulerOutput @@ -210,6 +211,32 @@ class OmniGPUModelRunner(PrefixCacheRunnerMixin, GPUModelRunner): self._init_talker_mtp() self._prewarm_attention_capture_workspaces() self._report_model_local_kv() + self._warn_unexposed_stage_hooks(model) + + # Read on the model this runner holds. A multi-stage wrapper that builds its + # stage module as a child must re-export them, or the runner silently takes + # the per-row, host-synchronizing default paths. + _STAGE_HOOKS = ("gpu_resident_buffer_keys", "preprocess_decode_batch", "use_async_omni_output") + + @classmethod + def _warn_unexposed_stage_hooks(cls, model: Any) -> None: + if model is None or not hasattr(model, "named_children"): + return + for child_name, child in model.named_children(): + missing = [ + hook + for hook in cls._STAGE_HOOKS + if getattr(child, hook, None) not in (None, False) and getattr(model, hook, None) in (None, False) + ] + if missing: + logger.warning( + "%s.%s defines %s but %s does not expose them; the runner reads these on %s.", + type(model).__name__, + child_name, + ", ".join(missing), + type(model).__name__, + type(model).__name__, + ) def _report_model_local_kv(self) -> None: """Log attention KV this model holds outside the paged manager. @@ -1608,6 +1635,21 @@ class OmniGPUModelRunner(PrefixCacheRunnerMixin, GPUModelRunner): device=device, ) + def _embed_multimodal_input_ids(self, num_scheduled_tokens, mm_embeds, is_mm_embed): + embedding_kwargs = {} + if mm_embeds and getattr(self.model, "supports_embed_input_ids_query_start_loc", False): + # The batch has already been reordered. Read host-owned boundaries + # so model prompt rearrangement does not assume a request order or + # introduce a GPU-to-host transfer on every decode step. + num_reqs = self.input_batch.num_reqs + embedding_kwargs["query_start_loc"] = self.query_start_loc.cpu[: num_reqs + 1].tolist() + return self.model.embed_input_ids( + self.input_ids.gpu[:num_scheduled_tokens], + multimodal_embeddings=mm_embeds, + is_multimodal=is_mm_embed, + **embedding_kwargs, + ) + def _preprocess( self, scheduler_output: "SchedulerOutput", @@ -1644,11 +1686,7 @@ class OmniGPUModelRunner(PrefixCacheRunnerMixin, GPUModelRunner): # NOTE(woosuk): To unify token ids and soft tokens (vision # embeddings), we always use embeddings (rather than token ids) # as input to the multimodal model, even when the input is text. - inputs_embeds_scheduled = self.model.embed_input_ids( - self.input_ids.gpu[:num_scheduled_tokens], - multimodal_embeddings=mm_embeds, - is_multimodal=is_mm_embed, - ) + inputs_embeds_scheduled = self._embed_multimodal_input_ids(num_scheduled_tokens, mm_embeds, is_mm_embed) # TODO(woosuk): Avoid the copy. Optimize. self.inputs_embeds.gpu[:num_scheduled_tokens].copy_(inputs_embeds_scheduled) @@ -1821,7 +1859,8 @@ class OmniGPUModelRunner(PrefixCacheRunnerMixin, GPUModelRunner): dtype=req_embeds.dtype, ) - offsets_t = torch.tensor(start_offsets_b, device=req_embeds.device, dtype=torch.long) + # A pageable H2D here would sync the host every decode step. + offsets_t = index_to_device(start_offsets_b, req_embeds.device) inputs_embeds.index_copy_(0, offsets_t, req_embeds) preprocess_input_ids.index_copy_( 0, diff --git a/vllm_omni/worker_v2/model_states/intermediate_buffer.py b/vllm_omni/worker_v2/model_states/intermediate_buffer.py index b9614a729..539755579 100644 --- a/vllm_omni/worker_v2/model_states/intermediate_buffer.py +++ b/vllm_omni/worker_v2/model_states/intermediate_buffer.py @@ -79,6 +79,22 @@ class OmniIntermediateBuffer: if ai is not None: info.update(_resolve_additional_information(ai)) + # Orchestrator stage bridges (e.g. MiniCPM-o's llm2tts) hand the + # request's runner payload over directly, like the V1 runner's + # ``_update_intermediate_buffer``: nested sections merge key by key. + mib = getattr(new_req_data, "model_intermediate_buffer", None) + if isinstance(mib, dict) and mib: + for key, value in mib.items(): + if isinstance(value, dict): + section = info.get(key) + if not isinstance(section, dict): + section = {} + info[key] = section + for qualifier, item in value.items(): + self._store_value(section, qualifier, item, set()) + else: + self._store_value(info, key, value, set()) + if new_req_data.mm_features: info["mm_features"] = new_req_data.mm_features diff --git a/vllm_omni/worker_v2/model_states/omni_model_state.py b/vllm_omni/worker_v2/model_states/omni_model_state.py index ca1e76cfc..1e5963cc2 100644 --- a/vllm_omni/worker_v2/model_states/omni_model_state.py +++ b/vllm_omni/worker_v2/model_states/omni_model_state.py @@ -19,6 +19,7 @@ import types from collections.abc import Callable from typing import Any, cast +import numpy as np import torch import torch.nn as nn from vllm.config import VllmConfig @@ -32,6 +33,7 @@ from vllm.v1.worker.gpu.states import RequestState from vllm_omni.model_executor.models.output_templates import OmniOutput, OwnedBatchTensor from vllm_omni.platforms import current_omni_platform +from vllm_omni.utils.device_copy import index_to_device from vllm_omni.worker_v2.model_states.eager_mtp import EagerMTPState from vllm_omni.worker_v2.model_states.intermediate_buffer import ( OmniIntermediateBuffer, @@ -87,6 +89,7 @@ class OmniModelState(DefaultModelState): """ # Eager-frame defaults for states built without __init__ (tests, adapters). + _decode_preprocess_is_identity = False _eager_mtp = False _eager_fastpath = False _eager_rows: tuple[InputBatch, list[tuple[int, int, str, bool]], torch.Tensor] | None = None @@ -139,6 +142,7 @@ class OmniModelState(DefaultModelState): self.has_postprocess: bool = getattr(model, "has_postprocess", False) self.have_multimodal_outputs: bool = getattr(model, "have_multimodal_outputs", False) self._decode_preprocess = self._resolve_decode_preprocess(model) + self._decode_preprocess_is_identity = bool(getattr(model, "mrv2_decode_preprocess_is_identity", False)) self._mtp_generators: dict[str, torch.Generator] = {} # Talker's codec_embedding dim may differ from hf_text_config.hidden_size; probe real dim. self._embed_dim = self._get_embed_dim(model, device) if self.has_preprocess else 0 @@ -211,6 +215,13 @@ class OmniModelState(DefaultModelState): input_batch, text_hidden, sampled_token_ids, multimodal_outputs, mtp_batch_descriptor_dispatcher ) + def custom_sampler(self, sampler: Any) -> tuple[Any, Any] | None: + """Let the model replace or wrap the MRv2 sampler (``mrv2_custom_sampler``).""" + hook = getattr(self.model, "mrv2_custom_sampler", None) + if callable(hook): + return hook(sampler) + return super().custom_sampler(sampler) + @staticmethod def _resolve_decode_preprocess(model: nn.Module) -> Callable | None: """Adapt existing V1 hooks once, outside the per-step hot path. @@ -603,7 +614,12 @@ class OmniModelState(DefaultModelState): preprocess_entries: list[tuple[int, int, int, int, dict[str, Any], bool]] = [] settled = self._eager_settled if self._eager_fastpath else None settled_rows: list[tuple[int, int, int, str]] = [] + # Declared identity decode has no per-row updates or hook work. + skip_decode_rows = self._decode_preprocess_is_identity and not self._eager_mtp + is_prefilling_np = getattr(input_batch, "is_prefilling_np", None) if skip_decode_rows else None for i, req_idx in enumerate(req_indices): + if is_prefilling_np is not None and not is_prefilling_np[i]: + continue buf = self.intermediate_buffer.buffers[req_idx] if not buf or "req_id" not in buf: continue @@ -791,6 +807,44 @@ class OmniModelState(DefaultModelState): prepacked_mtp_inputs=prepacked_mtp_inputs, ) + def publish_sampled_embeddings( + self, input_batch: InputBatch, sampled_token_ids: torch.Tensor + ) -> tuple[dict[str, Any], torch.cuda.Event] | None: + """Embedding of the token each row sampled this step. + + A producer whose consumer is fed the embeddings of its generated tokens + (``publishes_sampled_embeddings`` on the model) can then hand each token + over with the step that sampled it instead of one step later, when that + token's forward would capture the same embedding. Published beside the + packed snapshot as ``embed.sampled``: one ``[1, H]`` row per row whose + sample is kept, an empty tensor for rows still inside their prefill. + """ + if not getattr(self.model, "publishes_sampled_embeddings", False): + return None + num_reqs = input_batch.num_reqs + if not num_reqs: + return None + if sampled_token_ids.reshape(num_reqs, -1).shape[1] != 1: + return None # speculative steps sample several tokens per row + computed = input_batch.num_computed_prefill_tokens_np[:num_reqs] + scheduled = np.asarray(input_batch.num_scheduled_tokens[:num_reqs]) + finishing = computed + scheduled >= input_batch.prefill_len_np[:num_reqs] + kept = ~input_batch.is_prefilling_np[:num_reqs] | finishing + rows = np.flatnonzero(kept).tolist() + if not rows: + return None + device = sampled_token_ids.device + first = sampled_token_ids.reshape(num_reqs, -1)[:, 0].index_select(0, index_to_device(rows, device)) + embeds = self.model.embed_input_ids(first.long()) + empty = torch.empty(0, dtype=embeds.dtype) + sampled: list[torch.Tensor] = [empty] * num_reqs + for k, row in enumerate(rows): + sampled[row] = embeds[k : k + 1] + done = torch.cuda.Event() + done.record() + extra: dict[str, Any] = {"embed": {"sampled": sampled}} + return extra, done + def _pack_mtp_batch( self, mtp_batches: list[tuple[int, int, tuple[torch.Tensor, torch.Tensor]]], @@ -1288,6 +1342,16 @@ class OmniModelState(DefaultModelState): Handles ``OmniOutput`` unwrapping and ``make_omni_output`` conversion. """ + make_output_mrv2 = getattr(self.model, "make_omni_output_mrv2", None) + if not isinstance(model_output, OmniOutput) and callable(make_output_mrv2) and self.have_multimodal_outputs: + # Device-side output: the hook reads this step's batch and request + # state directly instead of per-request host views. + model_output = make_output_mrv2( + model_output, + input_batch=input_batch, + req_states=req_states, + model_intermediate_buffer=self.intermediate_buffer.gather(input_batch), + ) if not isinstance(model_output, OmniOutput) and hasattr(self.model, "make_omni_output"): if isinstance(model_output, (list, tuple)) or self.have_multimodal_outputs: buffer_list = self.intermediate_buffer.gather(input_batch) diff --git a/vllm_omni/worker_v2/omni_ar_model_runner.py b/vllm_omni/worker_v2/omni_ar_model_runner.py index 0f020c47b..28f363a1b 100644 --- a/vllm_omni/worker_v2/omni_ar_model_runner.py +++ b/vllm_omni/worker_v2/omni_ar_model_runner.py @@ -201,6 +201,17 @@ class OmniARModelRunner(OmniGPUModelRunner): multimodal_outputs, self._dispatch_mtp_batch_descriptor, ) + publish_sampled = getattr(self.model_state, "publish_sampled_embeddings", None) + extra_outputs = ( + publish_sampled(input_batch, sampler_output.sampled_token_ids) + if ( + multimodal_outputs + and callable(publish_sampled) + and bool(getattr(self.model_config, "async_chunk", False)) + and getattr(self.model_config, "engine_output_type", "text") != "text" + ) + else None + ) if self.pp_handler is not None: self.pp_handler.broadcast( sampler_output.sampled_token_ids, @@ -251,6 +262,7 @@ class OmniARModelRunner(OmniGPUModelRunner): finalize_output=(None if materialize_native else self._finalize_native_data_plane_output), check_ep_fault=self.check_ep_fault, routed_experts=routed_experts, + extra_multimodal_outputs=extra_outputs, ) self._release_multimodal_snapshot(snapshot_slot, async_output.copy_event) _guard_graph_replay_for_pooler_copy( @@ -543,6 +555,18 @@ def _async_copy_mm( } +def _merge_payload_trees(base: dict[str, Any], extra: dict[str, Any]) -> dict[str, Any]: + """``base`` with ``extra`` merged in; nested mappings merge key by key.""" + merged = dict(base) + for key, value in extra.items(): + current = merged.get(key) + if isinstance(current, dict) and isinstance(value, dict): + merged[key] = _merge_payload_trees(current, value) + else: + merged[key] = value + return merged + + def _slice_pooler_value( value: Any, *, @@ -649,6 +673,7 @@ class OmniAsyncOutput(AsyncModelRunnerOutput): finalize_output: Any | None = None, check_ep_fault: bool = False, routed_experts: RoutedExpertsTensors | 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 @@ -745,6 +770,20 @@ class OmniAsyncOutput(AsyncModelRunnerOutput): copy_stream=copy_stream, pin_memory=pin_memory, ) + if extra_multimodal_outputs: + # Produced after sampling on the producer stream; copy only + # once its completion event has been observed. + extra_outputs, extra_ready = extra_multimodal_outputs + copy_stream.wait_event(extra_ready) + self._mm_snapshot = _merge_payload_trees( + self._mm_snapshot, + _async_copy_mm( + extra_outputs, + self._total_tokens, + copy_stream=copy_stream, + pin_memory=pin_memory, + ), + ) elif self._need_pooler and text_hidden is not None: self._hidden_cpu = _async_copy_tensor( text_hidden, diff --git a/vllm_omni/worker_v2/omni_data_plane.py b/vllm_omni/worker_v2/omni_data_plane.py index 50a449b2d..d1ace2a9a 100644 --- a/vllm_omni/worker_v2/omni_data_plane.py +++ b/vllm_omni/worker_v2/omni_data_plane.py @@ -51,6 +51,7 @@ class _NativeRequestState: output_token_ids=output, all_token_ids=prompt + output, output_token_count=len(self.output_token_ids), + last_output_token_id=self.output_token_ids[-1] if self.output_token_ids else None, additional_information=self.additional_information, sampling_params=self.sampling_params, num_computed_tokens=self.num_computed_tokens, @@ -83,6 +84,11 @@ class OmniRunnerDataPlane(OmniConnectorModelRunnerMixin): self._native_send_lock = threading.Lock() self._native_outputs_in_flight: dict[str, int] = defaultdict(int) self._native_terminal_pending: set[str] = set() + # Stage payload builders read the stage config through + # ``_get_model_config()`` (e.g. Qwen3-Omni's Thinker reads + # ``hf_config.talker_config.accept_hidden_layer``). + self.vllm_config = vllm_config + self.model_config = model_config self.init_omni_connectors(model_config=model_config) self._delivery_manager = OmniDeliveryManager( delivery_timeout_s=self._connector_delivery_timeout(), @@ -394,7 +400,9 @@ class OmniRunnerDataPlane(OmniConnectorModelRunnerMixin): for qual, qval in value.items(): if qual in span_handled: continue - if key == "meta" and qual == "finished": + if key == "meta": + # Per-chunk state (``finished``, the V1 sender's + # ``is_segment_finished`` flag, ...), not rows. merged_sub[qual] = qval continue if (key, qual) in override_keys: @@ -563,6 +571,7 @@ class OmniRunnerDataPlane(OmniConnectorModelRunnerMixin): self._delivery_manager.shutdown(RuntimeError("MRv2 connector shutdown")) self._stop_event.set() self._work_available.set() + self._save_work_available.set() close_errors: list[BaseException] = [] connector = getattr(self, "_omni_connector", None) diff --git a/vllm_omni/worker_v2/omni_generation_model_runner.py b/vllm_omni/worker_v2/omni_generation_model_runner.py index 68e16f42e..605531338 100644 --- a/vllm_omni/worker_v2/omni_generation_model_runner.py +++ b/vllm_omni/worker_v2/omni_generation_model_runner.py @@ -56,6 +56,23 @@ def _contains_cuda_tensor(value: Any) -> bool: return False +def check_exact_input_shape(model: Any, input_ids: torch.Tensor | None, num_tokens: int) -> None: + """Enforce ``requires_exact_input_shape`` (#6712). + + Such models split the flat ``input_ids`` by per-request + ``seq_token_counts``, so padding rows would misalign every request. V1 + trims cudagraph padding for them (``maybe_unpad_input_ids``); this runner + dispatches eagerly and must already hand over the exact batch. + """ + if input_ids is None or not getattr(model, "requires_exact_input_shape", False): + return + if int(input_ids.shape[0]) != int(num_tokens): + raise RuntimeError( + f"{type(model).__name__} requires an exact input shape: got {int(input_ids.shape[0])} " + f"input ids for {int(num_tokens)} scheduled tokens" + ) + + def _materialize_generation_value( value: Any, req_index: int, @@ -368,6 +385,8 @@ class OmniGenerationModelRunner(OmniGPUModelRunner): input_batch, dummy_run=dummy_run, ) + if not dummy_run: + check_exact_input_shape(self.model, input_ids, input_batch.num_tokens) model_inputs = { "input_ids": input_ids, @@ -376,6 +395,10 @@ class OmniGenerationModelRunner(OmniGPUModelRunner): "intermediate_tensors": intermediate_tensors, **self.model_state.prepare_inputs(input_batch, self.req_states), } + if getattr(self.model, "requires_request_ids", False): + # Request-owned model state (e.g. MiniCPM-o's vocoder caches) is + # keyed by the scheduler request id, in batch order. + model_inputs["request_ids"] = list(input_batch.req_ids) self._add_legacy_forward_inputs(model_inputs, input_batch) eplb = getattr(self, "eplb", None) diff --git a/vllm_omni/worker_v2/omni_model_runner.py b/vllm_omni/worker_v2/omni_model_runner.py index 4eed3f84c..10f4053cf 100644 --- a/vllm_omni/worker_v2/omni_model_runner.py +++ b/vllm_omni/worker_v2/omni_model_runner.py @@ -9,6 +9,7 @@ import threading from typing import TYPE_CHECKING, Any, cast import torch +from torch.utils._pytree import tree_flatten, tree_unflatten from vllm.config.compilation import CUDAGraphMode from vllm.forward_context import ( set_forward_context, @@ -55,6 +56,16 @@ def _needs_capture_tensor_unwrap(model: Any) -> bool: return bool(getattr(model, "_returns_tuple", False)) +def _supports_full_graph_aux_outputs(model: Any) -> bool: + """Tuple-returning model whose auxiliary output FULL graphs may replay. + + The model declares ``supports_mrv2_full_graph_aux_outputs`` when its + auxiliary output is a tensor-only pytree with the same structure for every + batch shape and each leaf's leading dimension is the token axis. + """ + return _needs_capture_tensor_unwrap(model) and bool(getattr(model, "supports_mrv2_full_graph_aux_outputs", False)) + + class OmniGPUModelRunner(GPUModelRunner): """Thin layer over v2 ``GPUModelRunner`` for Omni lifecycle hooks.""" @@ -62,14 +73,48 @@ class OmniGPUModelRunner(GPUModelRunner): encoder_cache: EncoderCache | None _native_output_materializer: NativeOutputWorker | None model_state: OmniModelState + vocab_size: int + use_aux_hidden_state_outputs: bool _last_aux_output: Any _last_multimodal_outputs: dict[str, Any] | None _model_returns_tuple: bool + # FULL-graph aux-output contract (see _configure_cudagraph_output_contract). + _full_graph_aux_outputs = False + _aux_output_spec: Any = None def _configure_cudagraph_output_contract(self) -> None: """Select the CUDA graph output contract declared by the model.""" self._model_returns_tuple = _needs_capture_tensor_unwrap(self.model) - self._exclude_full_graph = self._model_returns_tuple or hasattr(self.model, "_last_captured_layers") + # FULL replay returns what capture stored: the hidden states and, for a + # declared tuple model, its auxiliary pytree flattened to leaves. + self._full_graph_aux_outputs = _supports_full_graph_aux_outputs(self.model) + self._aux_output_spec = None + self._exclude_full_graph = (self._model_returns_tuple and not self._full_graph_aux_outputs) or hasattr( + self.model, "_last_captured_layers" + ) + + def _split_fullgraph_output(self, replayed: Any) -> tuple[Any, Any]: + """FULL replay output -> ``(hidden_states, aux pytree or None)``.""" + if self._full_graph_aux_outputs and isinstance(replayed, tuple): + hidden, leaves = replayed + return hidden, tree_unflatten(list(leaves), self._aux_output_spec) + return replayed, None + + def _flatten_capture_aux_output(self, model_output: Any) -> tuple[torch.Tensor, list[torch.Tensor]]: + """``(hidden, aux pytree)`` -> ``(hidden, leaves)`` for the graph manager's aux buffers.""" + hidden, aux = model_output + leaves, spec = tree_flatten(aux) + num_tokens = hidden.shape[0] + if not leaves or any(not isinstance(leaf, torch.Tensor) or leaf.shape[:1] != (num_tokens,) for leaf in leaves): + raise RuntimeError( + f"{type(self.model).__name__} declares supports_mrv2_full_graph_aux_outputs, but its auxiliary " + "output is not a non-empty tensor pytree on the token axis" + ) + if self._aux_output_spec is None: + self._aux_output_spec = spec + elif spec != self._aux_output_spec: + raise RuntimeError(f"{type(self.model).__name__} auxiliary output structure changed between captures") + return hidden, leaves @staticmethod def _prepare_cudagraph_capture_output(model_output: Any) -> Any: @@ -171,10 +216,16 @@ class OmniGPUModelRunner(GPUModelRunner): # With no new requests, its UVA snapshots remain valid and immutable. # Keep the upstream path for custom/speculative samplers, whose staged # writes may have a different lifecycle. - if not scheduler_output.scheduled_new_reqs and type(getattr(self, "sampler", None)) is Sampler: + sampler = getattr(self, "sampler", None) + if not scheduler_output.scheduled_new_reqs and ( + type(sampler) is Sampler or getattr(sampler, "omni_static_staged_writes", False) + ): return logits_processor = getattr(self.model, "logits_processor", None) logits_vocab = getattr(logits_processor, "vocab_size", None) + if logits_vocab is None: + # Narrow model heads without a vLLM LogitsProcessor (codec heads). + logits_vocab = getattr(self.model, "logits_vocab_size", None) if isinstance(logits_vocab, int) and logits_vocab > 0: for request_data in scheduler_output.scheduled_new_reqs: sampling_params = request_data.sampling_params @@ -289,9 +340,17 @@ class OmniGPUModelRunner(GPUModelRunner): def load_model(self, *args: Any, **kwargs: Any) -> None: import vllm.v1.worker.gpu.model_runner as _mr_module + def _init_model_state(vllm_config: Any, model: Any, encoder_cache: Any, device: Any) -> Any: + # Runs after the model is built and before the sampler is: a model + # whose logits head is narrower than its HF config reports (the + # MiniCPM-o codec Talker's config has no vocab_size at all) sizes + # the MRv2 sampler and request state by that head. + self._adopt_model_logits_vocab(model) + return init_omni_model_state(vllm_config, model, encoder_cache, device) + with _model_state_patch_lock: _orig = _mr_module.init_model_state - _mr_module.init_model_state = init_omni_model_state + _mr_module.init_model_state = _init_model_state try: super().load_model(*args, **kwargs) finally: @@ -305,6 +364,18 @@ class OmniGPUModelRunner(GPUModelRunner): self.supports_mm_inputs = False self.encoder_cache = None + def _adopt_model_logits_vocab(self, model: Any) -> None: + vocab = getattr(model, "logits_vocab_size", None) + if not isinstance(vocab, int) or vocab <= 0 or vocab == self.vocab_size: + return + logger.info("Sizing the MRv2 sampler by the model's logits head: vocab %s -> %d", self.vocab_size, vocab) + self.vocab_size = vocab + self.req_states.vocab_size = vocab + # Kernel warmup sizes its structured-output bitmask from the config. + arch_config = getattr(self.model_config, "model_arch_config", None) + if arch_config is not None and getattr(arch_config, "vocab_size", vocab) != vocab: + arch_config.vocab_size = vocab + # ------------------------------------------------------------------ # CUDA Graph: conditionally exclude FULL mode # ------------------------------------------------------------------ @@ -336,19 +407,27 @@ class OmniGPUModelRunner(GPUModelRunner): self._exclude_unsupported_full_graphs() # Wrap model forward during capture so tuple returns don't crash - # torch.empty_like() in the PIECEWISE warmup pass. + # torch.empty_like() in the PIECEWISE warmup pass. A model with the + # FULL aux-output contract instead hands the graph manager its leaves, + # which it keeps in persistent buffers and returns on replay. if self._model_returns_tuple: original_forward = self.model.forward + aux_outputs = self._full_graph_aux_outputs def _capture_forward(*args: Any, **kwargs: Any) -> Any: output = original_forward(*args, **kwargs) + if aux_outputs: + return self._flatten_capture_aux_output(output) return self._prepare_cudagraph_capture_output(output) self.model.forward = _capture_forward # type: ignore[assignment] + use_aux = self.use_aux_hidden_state_outputs + self.use_aux_hidden_state_outputs = use_aux or aux_outputs try: result = super().capture_model() finally: self.model.forward = original_forward # type: ignore[assignment] + self.use_aux_hidden_state_outputs = use_aux else: result = super().capture_model() @@ -554,9 +633,9 @@ class OmniGPUModelRunner(GPUModelRunner): # inputs_embeds buffer above. assert self.cudagraph_manager is not None self.kv_connector.pre_forward(scheduler_output) - hidden_states = self.cudagraph_manager.run_fullgraph(batch_desc) - self._last_aux_output = None + replayed = self.cudagraph_manager.run_fullgraph(batch_desc) self._last_multimodal_outputs = None + hidden_states, self._last_aux_output = self._split_fullgraph_output(replayed) if hasattr(self.model, "_last_captured_layers"): self.model._last_captured_layers = self._last_aux_output else: