diff --git a/tests/diffusion/diffusion_kv/test_native_prefetch.py b/tests/diffusion/diffusion_kv/test_native_prefetch.py new file mode 100644 index 000000000..05c34abda --- /dev/null +++ b/tests/diffusion/diffusion_kv/test_native_prefetch.py @@ -0,0 +1,552 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from types import SimpleNamespace +from unittest.mock import Mock + +import pytest +import torch +import vllm.v1.core.single_type_kv_cache_manager as native_kv_managers +from vllm.config import KVTransferConfig +from vllm.v1.kv_cache_interface import FullAttentionSpec, KVCacheConfig, KVCacheGroupSpec, KVCacheTensor +from vllm.v1.outputs import KVConnectorOutput + +from vllm_omni.diffusion.diffusion_engine import DiffusionEngine +from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode +from vllm_omni.diffusion.diffusion_kv.kv_connector import KVReceiveProgress, native_prefetch_enabled +from vllm_omni.diffusion.diffusion_kv.request import DiffusionKVRequest +from vllm_omni.diffusion.executor.abstract import DiffusionExecutor +from vllm_omni.diffusion.request import OmniDiffusionRequest +from vllm_omni.diffusion.sched import DiffusionRequestStatus, RequestScheduler +from vllm_omni.diffusion.sched.interface import CachedRequestData, DiffusionSchedulerOutput +from vllm_omni.inputs.data import OmniDiffusionSamplingParams + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion] + + +def _output(*, submit=(), required=(), retire=(), metadata=None): + return DiffusionSchedulerOutput( + step_id=0, + scheduled_new_reqs=[], + scheduled_cached_reqs=CachedRequestData.make_empty(), + finished_req_ids=set(), + num_running_reqs=0, + num_waiting_reqs=0, + kv_connector_metadata=metadata, + kv_transfer_request_ids=set(submit), + kv_required_request_ids=set(required), + kv_finished_request_ids=set(retire), + kv_poll_only=metadata is None, + ) + + +def _completed(*ids, invalid=()): + return KVConnectorOutput(finished_recving=set(ids), invalid_block_ids=set(invalid)) + + +def test_prefetch_survives_post_forward_and_is_not_submitted_twice(): + progress = KVReceiveProgress() + active = Mock() + # B finishes before A, in a different get_finished/post_forward call. + active.post_forward.side_effect = [_completed("B0"), _completed("A0"), _completed(), _completed()] + first = _output(submit=["A0", "B0"], required=["A0"], metadata=object()) + assert progress.prepare(active, first, 1).finished_recving == {"A0", "B0"} + second = _output(required=["B0"], retire=["A0"]) + assert progress.prepare(active, second, 1).finished_recving == {"B0"} + active.pre_forward.assert_called_once_with(first) + progress.prepare(active, _output(retire=["B0"]), 1) + assert not progress.submitted + assert not progress.received + assert not progress.deadlines + + +def test_submit_b_does_not_wait_for_b(): + progress = KVReceiveProgress() + active = Mock() + active.post_forward.return_value = _completed("A0") + result = progress.prepare(active, _output(submit=["A0", "B0"], required=["A0"], metadata=object()), 10) + assert result.finished_recving == {"A0"} + assert progress.submitted == {"A0", "B0"} + active.post_forward.assert_called_once() + active.post_forward.side_effect = [_completed(), _completed("B0")] + progress.prepare(active, _output(required=["B0"]), 10) + assert progress.received == {"A0", "B0"} + assert active.pre_forward.call_count == 1 + + +def test_duplicate_submission_and_unknown_required_fail_before_transfer(): + progress = KVReceiveProgress(submitted={"B0"}) + active = Mock() + with pytest.raises(RuntimeError, match="Duplicate"): + progress.prepare(active, _output(submit=["B0"], metadata=object()), 1) + with pytest.raises(RuntimeError, match="never submitted"): + progress.prepare(active, _output(required=["C0"]), 1) + active.pre_forward.assert_not_called() + + +def test_poll_checks_timeout_without_resetting_it(monkeypatch): + now = [1.0] + monkeypatch.setattr("vllm_omni.diffusion.diffusion_kv.kv_connector.time.monotonic", lambda: now[0]) + progress = KVReceiveProgress() + active = Mock() + active.post_forward.side_effect = lambda *_: _completed() + progress.prepare(active, _output(submit=["B0"], metadata=object()), 5) + now[0] = 7 + with pytest.raises(TimeoutError, match="B0"): + progress.prepare(active, _output(), 5) + assert progress.submitted == {"B0"} + + +def test_poll_checks_load_errors_and_refuses_early_retirement(): + progress = KVReceiveProgress(submitted={"B0"}) + active = Mock() + active.post_forward.return_value = _completed(invalid=[42]) + with pytest.raises(RuntimeError, match="invalid remote pages"): + progress.prepare(active, _output(), 1) + active.post_forward.return_value = _completed() + with pytest.raises(RuntimeError, match="unfinished"): + progress.prepare(active, _output(retire=["B0"]), 1) + assert progress.submitted == {"B0"} + + +def test_late_unsubmitted_cancellation_ack_is_not_retained(): + progress = KVReceiveProgress() + active = Mock() + active.post_forward.return_value = _completed("cancelled-before-admission") + result = progress.prepare(active, _output(), 1) + assert not result.finished_recving + assert not progress.received + + +def test_executor_intersects_rank_completions_and_waits_only_required(): + executor = SimpleNamespace(od_config=SimpleNamespace(num_gpus=2), collective_rpc=Mock()) + output = _output(submit=["A0", "B0"], required=["A0"], metadata=object()) + executor.collective_rpc.return_value = [_completed("A0", "B0"), _completed("A0")] + result = DiffusionExecutor.prepare_kv_for_forward(executor, output) + assert result.finished_recving == {"A0"} + assert executor.collective_rpc.call_args.kwargs["exec_all_ranks"] is True + executor.collective_rpc.return_value = [_completed("A0", "B0"), _completed("A0", "B0")] + assert DiffusionExecutor.prepare_kv_for_forward(executor, _output()).finished_recving == {"A0", "B0"} + executor.collective_rpc.return_value = [_completed("B0"), _completed()] + # Background polling can return early, but compute B requires every rank. + assert not DiffusionExecutor.prepare_kv_for_forward(executor, _output()).finished_recving + with pytest.raises(RuntimeError, match="every rank"): + DiffusionExecutor.prepare_kv_for_forward(executor, _output(required=["B0"])) + + +def test_executor_submits_prefetch_after_current_rank_barrier(): + calls = [] + + def rpc(method, **kwargs): + output = kwargs["args"][0] + calls.append(output) + if len(calls) == 1: + assert output.kv_transfer_request_ids == output.kv_required_request_ids == {"A0"} + assert output.has_sync_kv_loads + else: + assert output.kv_transfer_request_ids == {"B0"} + assert not output.kv_required_request_ids + assert not output.has_sync_kv_loads + assert output.kv_connector_metadata == "B metadata" + # B is still in flight on both ranks. This must not block compute A. + return [_completed("A0"), _completed("A0")] + + executor = SimpleNamespace(od_config=SimpleNamespace(num_gpus=2), collective_rpc=rpc) + executor.prepare_kv_for_forward = lambda out: DiffusionExecutor.prepare_kv_for_forward(executor, out) + output = _output(submit=["A0", "B0"], required=["A0"], metadata="A metadata") + output.kv_prefetch_connector_metadata = "B metadata" + output.kv_prefetch_request_ids = {"B0"} + assert executor.prepare_kv_for_forward(output).finished_recving == {"A0"} + assert len(calls) == 2 + + +def _config(): + return SimpleNamespace( + kv_transfer_config=KVTransferConfig( + kv_connector="MooncakeConnector", + kv_role="kv_consumer", + engine_id="prefetch-test", + kv_connector_extra_config={"mooncake_protocol": "tcp", "enable_kv_async_prefetch": True}, + ), + diffusion_kv_mode=DiffusionKVCacheMode.PAGED_SCHEDULER, + max_num_seqs=1, + model_class_name="TestDiffusionPipeline", + ) + + +def test_mock_config_does_not_enable_native_prefetch(): + assert not native_prefetch_enabled(Mock()) + + +@pytest.mark.parametrize("enabled", ["true", "false", 0, 1, None]) +def test_native_prefetch_rejects_non_boolean_flag(enabled): + config = _config() + config.kv_transfer_config.kv_connector_extra_config["enable_kv_async_prefetch"] = enabled + with pytest.raises(ValueError, match="must be a boolean"): + native_prefetch_enabled(config) + + +def test_opt_in_and_platform_validation(monkeypatch): + monkeypatch.setattr("vllm_omni.platforms.current_omni_platform.is_cuda", lambda: True) + config = _config() + assert native_prefetch_enabled(config) + config.model_class_name = "AnotherDiffusionPipeline" + assert native_prefetch_enabled(config) + config.max_num_seqs = 2 + with pytest.raises(ValueError, match="max_num_seqs=1"): + native_prefetch_enabled(config) + config.max_num_seqs = 1 + monkeypatch.setattr("vllm_omni.platforms.current_omni_platform.is_cuda", lambda: False) + with pytest.raises(ValueError, match="CUDA"): + native_prefetch_enabled(config) + config.kv_transfer_config.kv_connector_extra_config["enable_kv_async_prefetch"] = False + assert not native_prefetch_enabled(config) + assert not native_prefetch_enabled(SimpleNamespace()) + + +class _Connector: + def __init__(self): + self.commits = [] + self.metas = [] + + def get_num_new_matched_tokens(self, request, computed): + return 4, True + + def update_state_after_alloc(self, request, blocks, tokens): + self.commits.append(request.request_id) + self.metas.append(request.request_id) + # Upstream Mooncake mutates this flag after staging the receive. + request.kv_transfer_params["do_remote_prefill"] = False + + def build_connector_meta(self, output): + meta, self.metas = tuple(self.metas), [] + return meta + + def update_connector_output(self, output): + pass + + def request_finished(self, request, blocks): + return False, None + + +def _scheduler(monkeypatch, num_blocks=24): + monkeypatch.setattr("vllm_omni.platforms.current_omni_platform.is_cuda", lambda: True) + connector = _Connector() + monkeypatch.setattr( + "vllm_omni.diffusion.diffusion_kv.kv_connector.create_scheduler_kv_connector", lambda *args: connector + ) + native_kv_managers.register_all_kvcache_specs(None) + spec = FullAttentionSpec(block_size=4, num_kv_heads=2, head_size=8, dtype=torch.bfloat16) + config = KVCacheConfig( + num_blocks=num_blocks, + kv_cache_tensors=[ + KVCacheTensor( + size=spec.page_size_bytes * num_blocks, + layers=["layer0"], + layer_stride=spec.page_size_bytes, + block_stride=spec.page_size_bytes, + ) + ], + kv_cache_groups=[KVCacheGroupSpec(layer_names=["layer0"], kv_cache_spec=spec)], + ) + scheduler = RequestScheduler() + scheduler.initialize( + _config(), + kv_cache_config=config, + scheduler_block_size=4, + hash_block_size=4, + kv_vllm_config=SimpleNamespace(model_config=SimpleNamespace(max_model_len=64), max_in_flight_tokens=64), + ) + return scheduler, connector + + +def _request(rid): + return OmniDiffusionRequest( + prompt=rid, + request_id=rid, + sampling_params=OmniDiffusionSamplingParams(num_inference_steps=1), + diffusion_kv_requests=tuple( + DiffusionKVRequest( + f"{rid}/{row}", sequence_id=row, prefix_len=4, target_len=4, seq_len=8, prompt_token_ids=[1] * 4 + ) + for row in range(2) + ), + kv_transfer_params={ + "do_remote_prefill": True, + "transfer_id": f"xfer-{rid}", + "remote_engine_id": "ar", + "remote_bootstrap_addr": "localhost:8998", + "num_transfer_tokens": 4, + }, + ) + + +def test_scheduler_prefetches_one_waiting_request_and_reuses_allocation(monkeypatch): + scheduler, connector = _scheduler(monkeypatch) + for rid in ("A", "B", "C"): + scheduler.add_request(_request(rid)) + output = scheduler.schedule() + assert output.scheduled_request_ids == ["A"] + assert output.kv_transfer_request_ids == {"A/0", "A/1", "B/0", "B/1"} + assert output.kv_required_request_ids == {"A/0", "A/1"} + assert output.kv_connector_metadata == ("A/0", "A/1") + assert output.kv_prefetch_connector_metadata == ("B/0", "B/1") + assert scheduler.get_request_state("B").status == DiffusionRequestStatus.WAITING + manager = scheduler._diffusion_kv_manager + b_meta = manager.get_metadata("B") + assert not manager.has_request("C") + a_blocks = {b for row in manager.get_metadata("A").sequences for group in row.block_ids for b in group} + b_blocks = {b for row in b_meta.sequences for group in row.block_ids for b in group} + assert a_blocks.isdisjoint(b_blocks) + scheduler.update_kv_connector_output(_completed("A/0", "A/1", "B/0", "B/1")) + scheduler.finish_requests("A", DiffusionRequestStatus.FINISHED_COMPLETED) + output = scheduler.schedule() + assert output.scheduled_request_ids == ["B"] + assert manager.get_metadata("B") is b_meta + assert output.kv_transfer_request_ids == {"C/0", "C/1"} + assert not output.kv_required_request_ids + assert output.kv_finished_request_ids == {"A/0", "A/1"} + assert connector.commits.count("B/0") == connector.commits.count("B/1") == 1 + + +def test_prefetch_capacity_failure_rolls_back_all_cfg_rows(monkeypatch): + # One sentinel + A's four pages + only one of B's two CFG rows. + scheduler, connector = _scheduler(monkeypatch, num_blocks=7) + for rid in ("A", "B"): + scheduler.add_request(_request(rid)) + output = scheduler.schedule() + assert output.kv_transfer_request_ids == {"A/0", "A/1"} + assert not scheduler._diffusion_kv_manager.has_request("B") + assert scheduler._native_prefetch_request_id is None + assert connector.commits == ["A/0", "A/1"] + + +def test_abort_drains_b_before_releasing_pages(monkeypatch): + scheduler, _ = _scheduler(monkeypatch) + for rid in ("A", "B"): + scheduler.add_request(_request(rid)) + scheduler.schedule() + scheduler.update_kv_connector_output(_completed("A/0", "A/1")) + scheduler.finish_requests("B", DiffusionRequestStatus.FINISHED_ABORTED) + assert "B" in scheduler._kv_draining_requests + assert scheduler._diffusion_kv_manager.has_request("B") + engine = object.__new__(DiffusionEngine) + engine.scheduler = scheduler + engine._remove_diffusion_kv_requests = Mock() + engine.abort_queue = SimpleNamespace(empty=lambda: True) + + def drain(output): + assert output.kv_required_request_ids == {"B/0", "B/1"} + assert scheduler._diffusion_kv_manager.has_request("B") + engine._remove_diffusion_kv_requests.assert_not_called() + return _completed("A/0", "A/1", "B/0", "B/1") + + engine.executor = SimpleNamespace(prepare_kv_for_forward=drain) + engine._abort_requests(["B"]) + assert not scheduler._diffusion_kv_manager.has_request("B") + assert scheduler.get_request_state("B").status == DiffusionRequestStatus.FINISHED_ABORTED + assert scheduler._native_prefetch_request_id is None + + +def test_pending_prefetch_reuses_transfer_when_admitted(monkeypatch): + scheduler, connector = _scheduler(monkeypatch) + for rid in ("A", "B"): + scheduler.add_request(_request(rid)) + scheduler.schedule() + scheduler.update_kv_connector_output(_completed("A/0", "A/1", "B/0")) + assert "B" in scheduler._kv_loading_request_ids + scheduler.finish_requests("A", DiffusionRequestStatus.FINISHED_COMPLETED) + output = scheduler.schedule() + assert output.scheduled_request_ids == ["B"] + assert output.kv_required_request_ids == {"B/0", "B/1"} + assert not output.kv_transfer_request_ids + assert len(connector.commits) == 4 + + +def test_missing_source_skips_prefetch_without_allocating(monkeypatch): + scheduler, connector = _scheduler(monkeypatch) + scheduler.add_request(_request("A")) + b = _request("B") + del b.kv_transfer_params["remote_bootstrap_addr"] + scheduler.add_request(b) + scheduler.schedule() + assert not scheduler._diffusion_kv_manager.has_request("B") + assert connector.commits == ["A/0", "A/1"] + + +def test_prefetch_reservation_error_does_not_fail_current_request(monkeypatch): + scheduler, connector = _scheduler(monkeypatch) + for rid in ("A", "B"): + scheduler.add_request(_request(rid)) + manager = scheduler._diffusion_kv_manager + reserve = manager.reserve_request + + def fail_b(request_id, requests): + if request_id == "B": + raise ValueError("invalid B layout") + return reserve(request_id, requests) + + monkeypatch.setattr(manager, "reserve_request", fail_b) + output = scheduler.schedule() + assert output.scheduled_request_ids == ["A"] + assert connector.commits == ["A/0", "A/1"] + assert not manager.has_request("B") + scheduler.update_kv_connector_output(_completed("A/0", "A/1")) + scheduler.finish_requests("A", DiffusionRequestStatus.FINISHED_COMPLETED) + output = scheduler.schedule() + assert scheduler.get_request_state("B").status == DiffusionRequestStatus.FINISHED_ERROR + assert "B" in output.finished_req_ids + + +def test_invalid_prefetch_boundary_is_rejected_before_reservation(monkeypatch): + scheduler, connector = _scheduler(monkeypatch) + for rid in ("A", "B"): + scheduler.add_request(_request(rid)) + b_rows = scheduler.get_request_state("B").diffusion_kv_requests + b_rows[1].kv_transfer_params["num_transfer_tokens"] = b_rows[1].num_tokens + 1 + manager = scheduler._diffusion_kv_manager + reserve = manager.reserve_request + + def fail_if_b_reserved(request_id, requests): + if request_id == "B": + raise AssertionError("invalid B must be rejected before reservation") + return reserve(request_id, requests) + + monkeypatch.setattr(manager, "reserve_request", fail_if_b_reserved) + output = scheduler.schedule() + assert output.scheduled_request_ids == ["A"] + assert scheduler.get_request_state("B").status == DiffusionRequestStatus.WAITING + assert not manager.has_request("B") + + +def test_abort_timeout_keeps_target_pages(monkeypatch): + scheduler, _ = _scheduler(monkeypatch) + for rid in ("A", "B"): + scheduler.add_request(_request(rid)) + scheduler.schedule() + engine = object.__new__(DiffusionEngine) + engine.scheduler = scheduler + engine.executor = SimpleNamespace(prepare_kv_for_forward=Mock(side_effect=TimeoutError("receive"))) + engine._fail_engine = Mock() + engine._remove_diffusion_kv_requests = Mock() + with pytest.raises(TimeoutError): + engine._abort_requests(["B"]) + engine._fail_engine.assert_called_once() + engine._remove_diffusion_kv_requests.assert_not_called() + assert scheduler._diffusion_kv_manager.has_request("B") + + +def test_sleep_rejects_live_prefetch_before_touching_allocator(): + from vllm_omni.diffusion.worker.diffusion_worker import DiffusionWorker + + worker = SimpleNamespace(model_runner=SimpleNamespace(_kv_receive_progress=KVReceiveProgress(submitted={"B/0"}))) + with pytest.raises(RuntimeError, match="Cannot sleep"): + DiffusionWorker.sleep(worker) + + +def test_registration_failure_retires_b_and_flushes_cancellation(monkeypatch): + from vllm_omni.diffusion.diffusion_kv import kv_connector + + scheduler, connector = _scheduler(monkeypatch) + for rid in ("A", "B"): + scheduler.add_request(_request(rid)) + commit = kv_connector.commit_kv_load + + def fail_b(conn, manager, requests, matched): + if requests[0].request_id.startswith("B/"): + conn.metas.append("B cancellation") + raise kv_connector.KVTransferRegistrationError("registration rollback") + return commit(conn, manager, requests, matched) + + monkeypatch.setattr("vllm_omni.diffusion.sched.base_scheduler.commit_kv_load", fail_b) + output = scheduler.schedule() + assert output.scheduled_request_ids == ["A"] + assert output.finished_req_ids == {"B"} + assert output.kv_finished_request_ids == {"B/0", "B/1"} + assert output.kv_prefetch_connector_metadata == ("B cancellation",) + assert not output.kv_prefetch_request_ids + assert output.num_waiting_reqs == 0 + assert scheduler.get_request_state("B").status == DiffusionRequestStatus.FINISHED_ERROR + assert not scheduler._diffusion_kv_manager.has_request("B") + assert scheduler._diffusion_kv_manager.has_request("A") + + +def test_partial_cfg_registration_rolls_back_every_row(monkeypatch): + from vllm_omni.diffusion.diffusion_kv import kv_connector + + class FakeMooncakeConnector: + def __init__(self): + self.staged = [] + self.cancelled = [] + + def update_state_after_alloc(self, request, blocks, tokens): + self.staged.append(request.request_id) + if request.request_id.endswith("/1"): + raise RuntimeError("second CFG row failed") + + def request_finished(self, request, blocks): + self.cancelled.append(request.request_id) + return False, None + + monkeypatch.setattr( + "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.mooncake_connector.MooncakeConnector", + FakeMooncakeConnector, + ) + scheduler, _ = _scheduler(monkeypatch) + scheduler.add_request(_request("B")) + state = scheduler.get_request_state("B") + rows = state.diffusion_kv_requests + for row in rows: + row.kv_transfer_params = dict(state.req.kv_transfer_params) + allocation = scheduler._diffusion_kv_manager.reserve_request("B", rows) + connector = FakeMooncakeConnector() + + with pytest.raises(kv_connector.KVTransferRegistrationError, match="Could not register"): + kv_connector.commit_kv_load( + connector, + scheduler._diffusion_kv_manager.native_manager, + rows, + [4, 4], + ) + assert allocation is not None + assert connector.staged == ["B/0", "B/1"] + assert connector.cancelled == ["B/0", "B/1"] + + +def test_drain_after_terminal_request_state_was_popped(monkeypatch): + scheduler, _ = _scheduler(monkeypatch) + for rid in ("A", "B"): + scheduler.add_request(_request(rid)) + scheduler.schedule() + scheduler.finish_requests("B", DiffusionRequestStatus.FINISHED_ABORTED) + scheduler.pop_request_state("B") + output = scheduler.native_kv_poll_output(drain_request_ids=["B"]) + assert output.kv_required_request_ids == {"B/0", "B/1"} + assert scheduler._diffusion_kv_manager.has_request("B") + + +@pytest.mark.parametrize("prefetch", [False, True]) +def test_real_active_connector_starts_required_load_before_polling(prefetch): + from vllm.v1.worker.gpu.kv_connector import ActiveKVConnector + + from vllm_omni.diffusion.diffusion_kv.kv_connector import wait_for_kv_load + + active = object.__new__(ActiveKVConnector) + active._disabled = False + active._pending_load_start = False + active.kv_connector = Mock() + active.kv_connector.get_block_ids_with_load_errors.return_value = set() + active._start_load_kv = Mock() + + def finished(_): + active._start_load_kv.assert_called_once() + return set(), {"A0"} + + active.kv_connector.get_finished.side_effect = finished + output = _output(submit=["A0"], required=["A0"], metadata=object()) + if prefetch: + result = KVReceiveProgress().prepare(active, output, 1) + else: + output.kv_required_request_ids = None + result = wait_for_kv_load(active, output, 1) + assert result.finished_recving == {"A0"} diff --git a/tests/diffusion/test_diffusion_engine_cleanup.py b/tests/diffusion/test_diffusion_engine_cleanup.py index ba7d7d5d2..eea749571 100644 --- a/tests/diffusion/test_diffusion_engine_cleanup.py +++ b/tests/diffusion/test_diffusion_engine_cleanup.py @@ -38,6 +38,7 @@ def _make_engine() -> DiffusionEngine: engine._cv = threading.Condition(engine._rpc_lock) engine._out_streams = {} engine._closed = False + engine._shutting_down = False engine._shutdown_complete = False engine.abort_queue = queue.Queue() engine._loop_started = False @@ -379,3 +380,22 @@ def test_close_defers_resource_shutdown_until_worker_thread_stops() -> None: engine.executor.shutdown.assert_called_once() assert engine._shutdown_complete is True assert engine._loop_started is False + + +def test_fail_engine_does_not_reenter_when_executor_shutdown_fails() -> None: + engine = _make_engine() + shutdown_error = RuntimeError("worker shutdown failed") + engine.executor.shutdown.side_effect = shutdown_error + engine.scheduler.close = Mock() + engine._fail_pending_rpcs = Mock() + + with pytest.raises(RuntimeError, match="worker shutdown failed"): + engine._fail_engine(RuntimeError("engine failed")) + + engine._fail_engine(RuntimeError("follow-up failure")) + + engine.executor.shutdown.assert_called_once_with() + engine.scheduler.close.assert_not_called() + assert engine._closed is True + assert engine._shutting_down is True + assert engine._shutdown_complete is False diff --git a/tests/diffusion/test_diffusion_engine_dummy_run.py b/tests/diffusion/test_diffusion_engine_dummy_run.py index 1f1e8cbab..6b8a71792 100644 --- a/tests/diffusion/test_diffusion_engine_dummy_run.py +++ b/tests/diffusion/test_diffusion_engine_dummy_run.py @@ -220,6 +220,7 @@ def test_paged_kv_profile_requests_match_per_rank_batch( width=1024, guidance_scale=5.0, num_image_inputs=3, + num_frames=1, ) engine._prepare_request_for_admission.assert_called_once_with(request) diff --git a/tests/diffusion/test_diffusion_worker.py b/tests/diffusion/test_diffusion_worker.py index a27cd2cc8..fd6260f41 100644 --- a/tests/diffusion/test_diffusion_worker.py +++ b/tests/diffusion/test_diffusion_worker.py @@ -48,6 +48,7 @@ def mock_gpu_worker(mocker: MockerFixture, mock_od_config): worker = DiffusionWorker(local_rank=0, rank=0, od_config=mock_od_config) # Mock the model_runner with pipeline worker.model_runner = mocker.Mock() + worker.model_runner._kv_receive_progress = None worker.model_runner.pipeline = mocker.Mock() worker.device = torch.device("cuda", 0) worker._sleep_saved_buffers = {} diff --git a/vllm_omni/diffusion/diffusion_engine.py b/vllm_omni/diffusion/diffusion_engine.py index 6440c2fe8..ae2a2f86d 100644 --- a/vllm_omni/diffusion/diffusion_engine.py +++ b/vllm_omni/diffusion/diffusion_engine.py @@ -385,6 +385,7 @@ class DiffusionEngine: self._cv = threading.Condition(self._rpc_lock) self._out_streams: dict[str, asyncio.Queue[DiffusionOutput]] = {} self._closed = False + self._shutting_down = False self._shutdown_complete = False self.abort_queue: queue.Queue[str] = queue.Queue() self._rpc_queue: queue.Queue[_RpcTask] = queue.Queue() @@ -635,6 +636,7 @@ class DiffusionEngine: if sched_output.scheduled_request_ids else BatchRunnerOutput.from_list([]) ) + self._poll_native_kv() worker_execution_completed = True except Exception as exc: if self._closed: @@ -845,7 +847,12 @@ class DiffusionEngine: raise def _prepare_kv_for_forward(self, sched_output: DiffusionSchedulerOutput) -> None: - if getattr(sched_output, "kv_connector_metadata", None) is None: + if ( + getattr(sched_output, "kv_connector_metadata", None) is None + and getattr(sched_output, "kv_prefetch_connector_metadata", None) is None + and not getattr(sched_output, "kv_required_request_ids", None) + and not getattr(sched_output, "kv_poll_only", False) + ): return try: output = self.executor.prepare_kv_for_forward(sched_output) @@ -853,7 +860,10 @@ class DiffusionEngine: self.scheduler.update_kv_connector_output(output) if not self.abort_queue.empty(): self._process_aborts_queue() - incomplete = sched_output.kv_transfer_request_ids - (output.finished_recving or set()) + required_ids = sched_output.kv_required_request_ids + if required_ids is None: + required_ids = sched_output.kv_transfer_request_ids + incomplete = required_ids - (output.finished_recving or set()) sched_output.finished_req_ids.update(self.scheduler.fail_incomplete_kv_loads(incomplete)) # Timed-out/cancelled requests must never reach model execution. terminal = { @@ -879,26 +889,37 @@ class DiffusionEngine: self._fail_engine(exc) raise - def _fail_engine(self, exc: Exception) -> None: - if getattr(self, "_shutdown_complete", False): + def _poll_native_kv(self, *, drain_request_ids: list[str] | None = None) -> None: + poll = getattr(self.scheduler, "native_kv_poll_output", None) + if poll is None: return - logger.error("Diffusion engine failed; stopping workers before releasing KV pages", exc_info=exc) + output = poll(drain_request_ids=drain_request_ids) + if output is not None: + self._prepare_kv_for_forward(output) + + def _fail_engine(self, exc: Exception) -> None: with self._cv: + if getattr(self, "_shutdown_complete", False) or getattr(self, "_shutting_down", False): + return + # Mark shutdown before invoking any component callbacks. A + # failing shutdown can re-enter this method from another error + # path, and repeating executor.shutdown() is unsafe. + self._shutting_down = True self._closed = True if self.stop_event is not None: self.stop_event.set() streams = list(self._out_streams.values()) self._cv.notify_all() + + logger.error("Diffusion engine failed; stopping workers before releasing KV pages", exc_info=exc) for stream in streams: self._put_queue_output(stream, DiffusionOutput.from_exception(exc)) self._fail_pending_rpcs(exc) - try: - self.executor.shutdown() - finally: - try: - self.scheduler.close() - finally: - self._shutdown_complete = True + # If Worker shutdown fails, retain the Scheduler reservations. A + # remote producer may still be writing into those allocations. + self.executor.shutdown() + self.scheduler.close() + self._shutdown_complete = True def _emit_finished_outputs( self, @@ -1150,6 +1171,7 @@ class DiffusionEngine: if sched_output.scheduled_request_ids else BatchRunnerOutput.from_list([]) ) + self._poll_native_kv() except EngineDeadError: raise except Exception as exc: @@ -1225,6 +1247,7 @@ class DiffusionEngine: guidance_scale: float, num_image_inputs: int = 1, num_inference_steps: int = 1, + num_frames: int | None = None, ) -> OmniDiffusionRequest | None: """Build a minimal model request for startup profiling or warmup.""" prompt = OmniTextPrompt(prompt="dummy run") @@ -1241,7 +1264,8 @@ class DiffusionEngine: audio_sr = 16000 prompt.setdefault("multi_modal_data", {})["audio"] = np.random.randn(audio_sr * 2).astype(np.float32) - num_frames = get_dummy_run_num_frames(model_class_name, supports_audio_input) + if num_frames is None: + num_frames = get_dummy_run_num_frames(model_class_name, supports_audio_input) if num_frames <= 0: return None return OmniDiffusionRequest( @@ -1293,6 +1317,9 @@ class DiffusionEngine: width=profile_width, guidance_scale=5.0, num_image_inputs=get_dummy_run_num_image_inputs(model_class_name), + # Hunyuan skips generic warmup, but paged KV still needs its + # prepared image request to measure the startup memory envelope. + num_frames=1, ) if request is None: raise RuntimeError("paged_scheduler requires a runnable Diffusion KV memory profile request") @@ -1495,8 +1522,15 @@ class DiffusionEngine: else: self._loop_started = False - self.scheduler.close() - self.executor.shutdown() + if getattr(self.scheduler, "_native_prefetch_enabled", False): + # No new schedules after the busy loop exits. Finish outstanding + # writes before shutting down consumers and releasing reservations. + self._poll_native_kv(drain_request_ids=list(self.scheduler._kv_loading_request_ids)) + self.executor.shutdown() + self.scheduler.close() + else: + self.scheduler.close() + self.executor.shutdown() self._shutdown_complete = True def abort(self, request_id: str | Iterable[str]) -> None: @@ -1532,6 +1566,8 @@ class DiffusionEngine: request_ids = [request_ids] if isinstance(request_ids, str) else list(request_ids) request_ids = list(dict.fromkeys(request_ids)) + self._poll_native_kv(drain_request_ids=request_ids) + for request_id in request_ids: if self.scheduler.get_request_state(request_id) is not None: self.scheduler.finish_requests(request_id, DiffusionRequestStatus.FINISHED_ABORTED) diff --git a/vllm_omni/diffusion/diffusion_kv/kv_connector.py b/vllm_omni/diffusion/diffusion_kv/kv_connector.py index eaa478300..b3a326c56 100644 --- a/vllm_omni/diffusion/diffusion_kv/kv_connector.py +++ b/vllm_omni/diffusion/diffusion_kv/kv_connector.py @@ -8,6 +8,7 @@ import math import time from collections import Counter from collections.abc import Mapping +from dataclasses import dataclass, field from typing import TYPE_CHECKING, Any from vllm.config import KVTransferConfig, VllmConfig @@ -35,6 +36,96 @@ class KVTransferRegistrationError(ValueError): """A failed handoff whose destination pages have not been dispatched.""" +def native_prefetch_enabled(od_config: OmniDiffusionConfig) -> bool: + """Validate the opt-in GPU consumer implementation before allocating pages.""" + config = getattr(od_config, "kv_transfer_config", None) + if not isinstance(config, KVTransferConfig): + return False + extra = config.kv_connector_extra_config or {} + enabled = extra.get("enable_kv_async_prefetch", False) + if not isinstance(enabled, bool): + raise ValueError("enable_kv_async_prefetch must be a boolean") + if not enabled: + return False + from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode + from vllm_omni.platforms import current_omni_platform + + if ( + not current_omni_platform.is_cuda() + or config.kv_connector != "MooncakeConnector" + or config.kv_role != "kv_consumer" + or extra.get("mooncake_protocol", "rdma") != "tcp" + or od_config.diffusion_kv_mode is not DiffusionKVCacheMode.PAGED_SCHEDULER + or od_config.max_num_seqs != 1 + or getattr(od_config, "cfg_kv_collect_func", None) is not None + ): + raise ValueError( + "Native KV prefetch requires CUDA, MooncakeConnector TCP consumer, " + "paged_scheduler, max_num_seqs=1 and no CFG companion collector" + ) + return True + + +@dataclass +class KVReceiveProgress: + """Rank-local state survives consumption of Mooncake completion events.""" + + submitted: set[str] = field(default_factory=set) + received: set[str] = field(default_factory=set) + sent: set[str] = field(default_factory=set) + deadlines: dict[str, float] = field(default_factory=dict) + + def prepare( + self, active_connector: ActiveKVConnector, output: DiffusionSchedulerOutput, timeout: float + ) -> KVConnectorOutput: + new_ids = output.kv_transfer_request_ids + if new_ids & self.submitted: + raise RuntimeError("Duplicate native KV receive submission") + required = output.kv_required_request_ids or set() + if not required.issubset(self.submitted | new_ids): + raise RuntimeError("Required native KV receive was never submitted") + if output.kv_connector_metadata is not None: + active_connector.pre_forward(output) + elif new_ids: + raise RuntimeError("Native KV submission requires connector metadata") + self.submitted.update(new_ids) + self.deadlines.update(dict.fromkeys(new_ids, time.monotonic() + timeout)) + + while True: + # post_forward also consumes get_finished(); always retain its + # events, including completions for a request not required yet. + result = active_connector.post_forward(output.kv_finished_request_ids) + # Cancelling a request before admission can send a zero-page + # notification to release its producer. Its late acknowledgement + # has no receive reservation and must not become a leaked record. + self.received.update(self.submitted.intersection(result.finished_recving or ())) + self.sent.update(self.submitted.intersection(result.finished_sending or ())) + if result.invalid_block_ids: + raise RuntimeError("Diffusion KV connector reported invalid remote pages") + expired = { + rid + for rid, deadline in self.deadlines.items() + if rid not in self.received and time.monotonic() >= deadline + } + if expired: + raise TimeoutError(f"Timed out receiving diffusion KV for {sorted(expired)}") + if required.issubset(self.received): + break + time.sleep(0.001) + + retired = output.kv_finished_request_ids + if (retired & self.submitted) - self.received: + raise RuntimeError("Cannot retire an unfinished native KV receive") + self.submitted.difference_update(retired) + self.received.difference_update(retired) + self.sent.difference_update(retired) + for rid in retired: + self.deadlines.pop(rid, None) + result.finished_recving = set(self.received) + result.finished_sending = set(self.sent) + return result + + def mint_transfer_id(request_id: str) -> str: """Return the stable ticket shared by one AR -> Diffusion handoff.""" @@ -235,6 +326,21 @@ def install_mooncake_cfg_fanout(connector: KVConnectorBase_V1) -> None: setattr(worker, "_omni_cfg_fanout_installed", True) +def validate_kv_transfer_boundaries(requests: tuple[DiffusionKVRequest, ...], matched_tokens: list[int]) -> None: + """Validate every CFG row before reserving or registering destination pages.""" + for request, num_tokens in zip(requests, matched_tokens, strict=True): + params = request.kv_transfer_params + if num_tokens <= 0 or params is None: + continue + transfer_tokens = params.get("num_transfer_tokens") + if type(transfer_tokens) is not int or not num_tokens <= transfer_tokens <= request.num_tokens: + raise KVTransferRegistrationError( + "Diffusion KV transfer boundary must cover the reusable prefix " + f"without exceeding the allocated sequence: reusable={num_tokens}, " + f"transfer={transfer_tokens!r}, allocated={request.num_tokens}" + ) + + def commit_kv_load( connector: KVConnectorBase_V1, manager: KVCacheManager, @@ -246,6 +352,7 @@ def commit_kv_load( # All rows sharing a transfer_id must reach the same connector metadata: # the producer uses this complete per-rank row set for fan-out accounting. # Validate every CFG row before mutating any connector state. + validate_kv_transfer_boundaries(requests, matched_tokens) for request, num_tokens in zip(requests, matched_tokens, strict=True): blocks = manager.get_blocks(request.request_id) # Mooncake's producer advertises complete physical blocks. Keep every @@ -256,12 +363,6 @@ def commit_kv_load( transfer_tokens = num_tokens if num_tokens > 0 and request.kv_transfer_params is not None: transfer_tokens = request.kv_transfer_params["num_transfer_tokens"] - if type(transfer_tokens) is not int or not num_tokens <= transfer_tokens <= request.num_tokens: - raise KVTransferRegistrationError( - "Diffusion KV transfer boundary must cover the reusable prefix " - f"without exceeding the allocated sequence: reusable={num_tokens}, " - f"transfer={transfer_tokens!r}, allocated={request.num_tokens}" - ) prefix_blocks = KVCacheBlocks( tuple( group[: (transfer_tokens + spec.kv_cache_spec.block_size - 1) // spec.kv_cache_spec.block_size] diff --git a/vllm_omni/diffusion/executor/abstract.py b/vllm_omni/diffusion/executor/abstract.py index 9b0f71e23..0f184552f 100644 --- a/vllm_omni/diffusion/executor/abstract.py +++ b/vllm_omni/diffusion/executor/abstract.py @@ -177,7 +177,31 @@ class DiffusionExecutor(ABC): ) def prepare_kv_for_forward(self, scheduler_output: DiffusionSchedulerOutput) -> KVConnectorOutput | None: - if scheduler_output.kv_connector_metadata is None: + if scheduler_output.kv_prefetch_connector_metadata is not None: + current = replace( + scheduler_output, + kv_transfer_request_ids=scheduler_output.kv_transfer_request_ids + - scheduler_output.kv_prefetch_request_ids, + kv_prefetch_connector_metadata=None, + kv_prefetch_request_ids=set(), + ) + # First complete the current request on every rank. Only then + # submit B, preventing Mooncake from coalescing its bytes with A. + self.prepare_kv_for_forward(current) + return self.prepare_kv_for_forward( + replace( + current, + kv_connector_metadata=scheduler_output.kv_prefetch_connector_metadata, + kv_transfer_request_ids=scheduler_output.kv_prefetch_request_ids, + kv_required_request_ids=set(), + kv_finished_request_ids=set(), + ) + ) + if ( + scheduler_output.kv_connector_metadata is None + and not scheduler_output.kv_required_request_ids + and not scheduler_output.kv_poll_only + ): return transfer_output = replace( scheduler_output, @@ -194,6 +218,14 @@ class DiffusionExecutor(ABC): if len(outputs) != self.od_config.num_gpus or any(output.invalid_block_ids for output in outputs): # Missing ranks / invalid pages cannot establish safe ownership. raise RuntimeError("Diffusion KV receive failed on one or more ranks") + if scheduler_output.kv_required_request_ids is not None: + # Prefetch workers retain completion events until retirement. + # Their cumulative snapshots can be intersected directly. + finished = set.intersection(*(set(output.finished_recving or ()) for output in outputs)) + if not scheduler_output.kv_required_request_ids.issubset(finished): + raise RuntimeError("Required diffusion KV receive did not complete on every rank") + outputs[0].finished_recving = finished + return outputs[0] completed = getattr(self, "_kv_receive_completed_ranks", {}) for rank, output in enumerate(outputs): for request_id in output.finished_recving or (): diff --git a/vllm_omni/diffusion/sched/base_scheduler.py b/vllm_omni/diffusion/sched/base_scheduler.py index f0ca6c3ac..bc897b0b0 100644 --- a/vllm_omni/diffusion/sched/base_scheduler.py +++ b/vllm_omni/diffusion/sched/base_scheduler.py @@ -6,6 +6,7 @@ from __future__ import annotations import time from abc import ABC, abstractmethod from collections import deque +from collections.abc import Iterable from dataclasses import fields from vllm.config import VllmConfig @@ -19,7 +20,9 @@ from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode from vllm_omni.diffusion.diffusion_kv.kv_connector import ( KVTransferRegistrationError, commit_kv_load, + native_prefetch_enabled, prepare_kv_requests, + validate_kv_transfer_boundaries, ) from vllm_omni.diffusion.diffusion_kv.manager import DiffusionKVCacheManager from vllm_omni.diffusion.diffusion_kv.metadata import DiffusionKVMetadata @@ -71,6 +74,8 @@ class BaseScheduler(ABC): self._kv_request_generations: dict[str, int] = {} self._kv_draining_requests: dict[str, SchedulerRequestState] = {} self._kv_received_request_ids: set[str] = set() + self._native_prefetch_enabled = False + self._native_prefetch_request_id: str | None = None def initialize( self, @@ -82,6 +87,8 @@ class BaseScheduler(ABC): kv_vllm_config: VllmConfig | None = None, ) -> None: self.od_config = od_config + self._native_prefetch_enabled = native_prefetch_enabled(od_config) + self._native_prefetch_request_id = None self._request_states.clear() self._step_id = 0 self._waiting.clear() @@ -227,6 +234,8 @@ class BaseScheduler(ABC): sequence.num_computed_tokens = request.num_computed_tokens self._waiting.popleft() + if request_id == self._native_prefetch_request_id: + self._native_prefetch_request_id = None was_new_request = state.status == DiffusionRequestStatus.WAITING if not self._running: self._running_sampling_params_key = state.sampling_params_key @@ -277,6 +286,23 @@ class BaseScheduler(ABC): scheduler_output.kv_connector_metadata = self._kv_connector.build_connector_meta(scheduler_output) scheduler_output.kv_transfer_request_ids = self._kv_transfer_request_ids self._kv_transfer_request_ids = set() + if self._native_prefetch_enabled: + scheduler_output.kv_required_request_ids = self._loading_sequence_ids(self._running) + # Flush current-request metadata before staging B. Mooncake + # batches ready requests from the same producer into one write; + # combining A+B would make A's completion depend on B's bytes. + self._try_native_prefetch() + prefetch_finished_ids = self._kv_finished_request_ids - scheduler_output.kv_finished_request_ids + if self._kv_transfer_request_ids or prefetch_finished_ids: + scheduler_output.kv_prefetch_connector_metadata = self._kv_connector.build_connector_meta( + scheduler_output + ) + scheduler_output.kv_prefetch_request_ids = self._kv_transfer_request_ids + scheduler_output.kv_transfer_request_ids |= self._kv_transfer_request_ids + self._kv_transfer_request_ids = set() + scheduler_output.finished_req_ids.update(self._finished_req_ids) + scheduler_output.kv_finished_request_ids.update(prefetch_finished_ids) + scheduler_output.num_waiting_reqs = len(self._waiting) # update after schedule self._step_id += 1 @@ -284,6 +310,92 @@ class BaseScheduler(ABC): self._kv_finished_request_ids.clear() return scheduler_output + def _try_native_prefetch(self) -> None: + if not self._running or not self._waiting or self._native_prefetch_request_id is not None: + return + request_id = self._waiting[0] + state = self._request_states.get(request_id) + if state is None or state.is_finished() or not state.diffusion_kv_requests: + return + manager = self._diffusion_kv_manager + if manager.has_request(request_id): + return + matched_tokens = [] + for request in state.diffusion_kv_requests: + params = request.kv_transfer_params or {} + if ( + not params.get("do_remote_prefill") + or not all(params.get(key) for key in ("remote_engine_id", "remote_bootstrap_addr", "transfer_id")) + or type(params.get("num_transfer_tokens")) is not int + or params["num_transfer_tokens"] <= 0 + ): + return + num_tokens, _ = self._kv_connector.get_num_new_matched_tokens(request, 0) + if num_tokens is None: + return + matched_tokens.append(num_tokens) + try: + validate_kv_transfer_boundaries(state.diffusion_kv_requests, matched_tokens) + except KVTransferRegistrationError as exc: + # Do not reserve B's pages for a transfer that cannot be registered. + # Normal admission will report the request-scoped error later. + logger.debug("Native KV prefetch skipped for %s: invalid boundary (%s)", request_id, exc) + return + # The manager atomically reserves all CFG rows, including capacity + # needed when B eventually executes. B remains in WAITING. + try: + allocation = manager.reserve_request(request_id, state.diffusion_kv_requests) + except Exception as exc: + # reserve_request rolls back atomically. Normal admission will + # surface B's error later; speculative allocation must not fail A. + logger.debug("Native KV prefetch skipped for %s: reservation failed (%s)", request_id, exc) + return + if allocation is None: + logger.debug("Native KV prefetch skipped for %s: insufficient blocks", request_id) + return + self._kv_request_generations[request_id] = allocation.allocation_generation + try: + transfer_ids = commit_kv_load( + self._kv_connector, manager.native_manager, state.diffusion_kv_requests, matched_tokens + ) + except KVTransferRegistrationError as exc: + # Rollback may already notify the producer to release its pages. + # B cannot retry that transfer ticket, but A can still proceed. + self._finish_requests({request_id: DiffusionRequestStatus.FINISHED_ERROR}, {request_id: str(exc)}) + return + self._kv_transfer_request_ids.update(transfer_ids) + if transfer_ids: + self._kv_loading_request_ids.add(request_id) + for sequence, request in zip(allocation.sequences, state.diffusion_kv_requests, strict=True): + sequence.num_computed_tokens = request.num_computed_tokens + self._native_prefetch_request_id = request_id + logger.debug("Native KV prefetch submitted for %s: %s", request_id, sorted(transfer_ids)) + + def _loading_sequence_ids(self, request_ids: Iterable[str]) -> set[str]: + sequence_ids: set[str] = set() + for request_id in request_ids: + if request_id not in self._kv_loading_request_ids: + continue + state = self._request_states.get(request_id) or self._kv_draining_requests.get(request_id) + if state is not None: + sequence_ids.update(request.request_id for request in state.diffusion_kv_requests) + return sequence_ids + + def native_kv_poll_output(self, *, drain_request_ids: list[str] | None = None) -> DiffusionSchedulerOutput | None: + """Poll after compute; cancellation/close may wait for selected loads.""" + if not self._native_prefetch_enabled or not self._kv_loading_request_ids: + return None + return DiffusionSchedulerOutput( + step_id=self._step_id, + scheduled_new_reqs=[], + scheduled_cached_reqs=CachedRequestData.make_empty(), + finished_req_ids=set(), + num_running_reqs=len(self._running), + num_waiting_reqs=len(self._waiting), + kv_required_request_ids=self._loading_sequence_ids(drain_request_ids or []), + kv_poll_only=True, + ) + @abstractmethod def update_from_output(self, sched_output: DiffusionSchedulerOutput, output: RunnerOutput) -> set[str]: pass @@ -426,6 +538,7 @@ class BaseScheduler(ABC): self._kv_draining_requests.clear() self._kv_received_request_ids.clear() self._kv_request_generations.clear() + self._native_prefetch_request_id = None def _shutdown_diffusion_kv(self) -> None: """Stop transfers before releasing their cache and request state.""" @@ -471,6 +584,8 @@ class BaseScheduler(ABC): for request_id in finished_req_ids: state = self._request_states[request_id] status = statuses[request_id] + if request_id == self._native_prefetch_request_id: + self._native_prefetch_request_id = None if request_id in self._kv_loading_request_ids: self._kv_draining_requests[request_id] = state else: diff --git a/vllm_omni/diffusion/sched/interface.py b/vllm_omni/diffusion/sched/interface.py index fd2b47c63..12beef9d3 100644 --- a/vllm_omni/diffusion/sched/interface.py +++ b/vllm_omni/diffusion/sched/interface.py @@ -256,9 +256,23 @@ class DiffusionSchedulerOutput: kv_prefetch_job: KVPrefetchJob | None = None kv_connector_metadata: KVConnectorMetadata | None = None kv_transfer_request_ids: set[str] = field(default_factory=set) + # None preserves the synchronous path; an empty set means submit/poll + # without waiting. Transfer IDs above are submitted only once. + kv_required_request_ids: set[str] | None = None + kv_poll_only: bool = False + # Keep background transfers out of the current request's Mooncake batch. + kv_prefetch_connector_metadata: KVConnectorMetadata | None = None + kv_prefetch_request_ids: set[str] = field(default_factory=set) # Connector lifecycle uses per-sequence IDs, not public request IDs. kv_finished_request_ids: set[str] = field(default_factory=set) + @property + def has_sync_kv_loads(self) -> bool: + # vLLM defers async submission until post_forward. The synchronous + # receive loop must start its loads in pre_forward before polling. + required = self.kv_required_request_ids + return bool(self.kv_transfer_request_ids if required is None else required) + @cached_property def scheduled_request_ids(self) -> list[str]: """ @@ -275,4 +289,10 @@ class DiffusionSchedulerOutput: @property def is_empty(self) -> bool: - return self.num_scheduled_reqs == 0 and self.kv_connector_metadata is None + return ( + self.num_scheduled_reqs == 0 + and self.kv_connector_metadata is None + and self.kv_prefetch_connector_metadata is None + and not self.kv_required_request_ids + and not self.kv_poll_only + ) diff --git a/vllm_omni/diffusion/worker/diffusion_model_runner.py b/vllm_omni/diffusion/worker/diffusion_model_runner.py index d4d457bdc..fe5c1a21f 100644 --- a/vllm_omni/diffusion/worker/diffusion_model_runner.py +++ b/vllm_omni/diffusion/worker/diffusion_model_runner.py @@ -200,6 +200,9 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin): else None ) self._kv_connector = None + from vllm_omni.diffusion.diffusion_kv.kv_connector import KVReceiveProgress, native_prefetch_enabled + + self._kv_receive_progress = KVReceiveProgress() if native_prefetch_enabled(od_config) else None # Prefetch covers TP / SP / CFG-Parallel / HSDP. Disabled when a CFG # companion KV collector is set (that KV is not backgrounded). @@ -487,6 +490,8 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin): from vllm_omni.diffusion.diffusion_kv.kv_connector import wait_for_kv_load timeout = self.od_config.kv_transfer_config.kv_connector_extra_config.get("transfer_timeout", 60.0) + if self._kv_receive_progress is not None: + return self._kv_receive_progress.prepare(self._kv_connector, scheduler_output, timeout) return wait_for_kv_load(self._kv_connector, scheduler_output, timeout) def install_diffusion_kv_metadata(self, metadata: DiffusionKVMetadata) -> bool: diff --git a/vllm_omni/diffusion/worker/diffusion_worker.py b/vllm_omni/diffusion/worker/diffusion_worker.py index f4574b35e..0684a9387 100644 --- a/vllm_omni/diffusion/worker/diffusion_worker.py +++ b/vllm_omni/diffusion/worker/diffusion_worker.py @@ -801,6 +801,9 @@ class DiffusionWorker: Args: level: Sleep level. Level 1 offloads weights, level 2 also saves buffers. """ + progress = getattr(getattr(self, "model_runner", None), "_kv_receive_progress", None) + if progress is not None and progress.submitted: + raise RuntimeError("Cannot sleep with live native KV prefetch reservations; finish requests first") # The config validator rejects sleep for the native paged path. Keep # this worker-side guard precise as well: test doubles and legacy # configs may expose arbitrary attributes through Mock/getattr.