[Feature] Transfer AR-to-diffusion payloads over NIXL (#6264)

Signed-off-by: yuanwu <yuan.wu@intel.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>

vllm_omni/diffusion/worker/diffusion_worker.py

@@ -261,6 +261,7 @@ class DiffusionWorker:
         self.init_snapshot: MemorySnapshot | None = None
         self.requested_memory: int | None = None
         self._sleep_saved_buffers: dict[str, torch.Tensor] = {}
+        self._owns_sleep_pool = False
         self.lora_manager: DiffusionLoRAManager | None = None
         # Worker-side cache of (lora_request, lora_scale) per scheduled
         # request id. Used by step mode to recover LoRA identity for cached
@@ -1010,6 +1011,7 @@ class DiffusionWorker:
             allocator = CuMemAllocator.get_instance()
             if tag == "weights":
                 assert allocator.get_current_usage() == 0, "Sleep mode can only be used for one instance per process."
+                self._owns_sleep_pool = True
             logger.info(f"[Worker {self.rank}] Activating Diffusion CuMem pool for tag: {tag}")
             return allocator.use_memory_pool(tag=tag)
         return nullcontext()
@@ -1019,25 +1021,36 @@ class DiffusionWorker:
         try:
             if self.model_runner is not None:
                 mgr = getattr(self.model_runner, "kv_transfer_manager", None)
+                if mgr is None:
+                    mgr = getattr(self.model_runner, "_kv_transfer_manager", None)
                 try:
                     offload_backend = getattr(self.model_runner, "offload_backend", None)
                     if offload_backend is not None:
                         offload_backend.disable()
                 finally:
                     if mgr is not None:
-                        mgr.shutdown_prefetch()
+                        mgr.close()
         finally:
+            self.model_runner = None
+            self.lora_manager = None
+            self._sleep_saved_buffers = {}
             try:
-                shutdown_kv_connector()
+                if getattr(self, "_owns_sleep_pool", False):
+                    gc.collect()
+                    _get_cumem_allocator_class().get_instance().release_pools()
+                    self._owns_sleep_pool = False
             finally:
                 try:
-                    a2a_permute = sys.modules.get("vllm_omni.diffusion.distributed.a2a_permute")
-                    if a2a_permute is not None:
-                        a2a_permute.clear_a2a_permute_workspaces()
-                except Exception:
-                    logger.exception("Failed to release fused Ulysses symmetric-memory workspaces")
+                    shutdown_kv_connector()
                 finally:
-                    destroy_distributed_env()
+                    try:
+                        a2a_permute = sys.modules.get("vllm_omni.diffusion.distributed.a2a_permute")
+                        if a2a_permute is not None:
+                            a2a_permute.clear_a2a_permute_workspaces()
+                    except Exception:
+                        logger.exception("Failed to release fused Ulysses symmetric-memory workspaces")
+                    finally:
+                        destroy_distributed_env()
 
 
 class CustomPipelineWorkerExtension:
@@ -1642,7 +1655,14 @@ class WorkerWrapperBase:
 
         # Re-initialize pipeline with custom pipeline if provided
         if self.uses_custom_pipeline:
-            self.worker.re_init_pipeline(self.custom_pipeline_args)
+            try:
+                self.worker.re_init_pipeline(self.custom_pipeline_args)
+            except Exception:
+                try:
+                    self.worker.shutdown()
+                except Exception:
+                    logger.exception("Failed to clean up worker after custom pipeline initialization failure")
+                raise
 
     def _prepare_worker_class(self) -> type:
         """