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

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

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.