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

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

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.