[Feat] Add prefetch for Mooncake cross-stage paged KV transfer in HunyuanImage3 (#7637)

Signed-off-by: Acerak01-fy <wfy2003324@163.com>

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"}

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

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)
 

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 = {}

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)

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]

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 ():

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:

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
+        )

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:

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.