[Core] hunyuan-image prefetch kv (#4448)

Signed-off-by: z00806815 <zhuhonghan@huawei.com>
Signed-off-by: zhuhh97 <zhuhonghan@huawei.com>
Signed-off-by: dph97 <nenbaying@163.com>
Signed-off-by: zhh <zhh@users.noreply.github.com>
Co-authored-by: Claude Opus 4.7 <noreply@anthropic.com>
Co-authored-by: zhuhh97 <zhuhh97@users.noreply.github.com>

vllm_omni/diffusion/worker/diffusion_model_runner.py

@@ -80,9 +80,23 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
         # Cache for per-request stepwise state.
         self.state_cache: dict[str, DiffusionRequestState] = {}
 
-        # Initialize KV cache manager for connector management
+        # Initialize KV cache manager for connector management.
         self.kv_transfer_manager = OmniKVTransferManager.from_od_config(od_config)
 
+        # Prefetch covers TP / SP / CFG-Parallel / HSDP.  Disabled when a CFG
+        # companion KV collector is set (that KV is not backgrounded).
+        has_cfg_companion_kv = getattr(od_config, "cfg_kv_collect_func", None) is not None
+
+        self._kv_prefetch_enabled = (
+            bool(self.kv_transfer_manager.config.enable_kv_async_prefetch)
+            and not has_cfg_companion_kv
+            and self.kv_transfer_manager.config.need_recv_cache
+        )
+
+    @property
+    def target_device(self) -> torch.device | None:
+        return getattr(self.pipeline, "device", None)
+
     def _compile_transformer(self, attr_name: str) -> None:
         """Compile a transformer attribute on the pipeline with torch.compile."""
         model = getattr(self.pipeline, attr_name, None)
@@ -269,7 +283,7 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
             pool_overhead_gb / peak_reserved_gb * 100 if peak_reserved_gb > 0 else 0.0,
         )
 
-    def execute_model(self, req: OmniDiffusionRequest) -> DiffusionOutput:
+    def execute_model(self, req: OmniDiffusionRequest, kv_prefetch_jobs: dict | None = None) -> DiffusionOutput:
         """
         Execute a forward pass for the given requests.
 
@@ -293,12 +307,26 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
         use_hsdp = self.od_config.parallel_config.use_hsdp
         grad_context = torch.no_grad() if use_hsdp else torch.inference_mode()
         with grad_context:
-            # The manager handles the check for need_recv_cache internally
-            self.kv_transfer_manager.receive_multi_kv_cache_distributed(
-                req,
-                cfg_kv_collect_func=getattr(self.od_config, "cfg_kv_collect_func", None),
-                target_device=getattr(self.pipeline, "device", None),
-            )
+            # Receive AR KV (fetch → distribute → apply inside the entry). prefetch on:
+            # consume prior-forward payload, sync-fallback on miss; else sync receive.
+            kv_recv_t0 = time.perf_counter()
+            if self._kv_prefetch_enabled:
+                self.kv_transfer_manager.consume_and_distribute_kv_cache(
+                    req,
+                    target_device=self.target_device,
+                )
+            else:
+                self.kv_transfer_manager.receive_multi_kv_cache_distributed(
+                    req,
+                    cfg_kv_collect_func=getattr(self.od_config, "cfg_kv_collect_func", None),
+                    target_device=self.target_device,
+                )
+            kv_recv_ms = (time.perf_counter() - kv_recv_t0) * 1000
+            logger.debug("KV recv for %s %.1fms", req.request_id, kv_recv_ms)
+
+            # Kick off the next request's prefetch (+ H2D) to overlap this forward.
+            if self._kv_prefetch_enabled and kv_prefetch_jobs is not None:
+                self.kv_transfer_manager.start_prefetch(kv_prefetch_jobs, self.target_device)
 
             if req.sampling_params.generator is None and req.sampling_params.seed is not None:
                 if req.sampling_params.generator_device is not None:
@@ -404,7 +432,7 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
                 self.kv_transfer_manager.receive_multi_kv_cache_distributed(
                     state_req,
                     cfg_kv_collect_func=getattr(self.od_config, "cfg_kv_collect_func", None),
-                    target_device=getattr(self.pipeline, "device", None),
+                    target_device=self.target_device,
                 )
                 self.state_cache[request_id] = new_state
                 resolved.append(new_state)