[Diffusion] Implement paged AR↔DiT KV connector ("v1", reusing vllm's native mooncake) (#6310)

Signed-off-by: asukaqaq-s <1311722138@qq.com>
Signed-off-by: Huang, Zeyu <11222265+fhfuih@users.noreply.github.com>
Co-authored-by: asukaqaq-s <1311722138@qq.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>

vllm_omni/diffusion/worker/diffusion_worker.py

@@ -46,6 +46,10 @@ from vllm_omni.diffusion.data import (
     OmniWakeTask,
 )
 from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode
+from vllm_omni.diffusion.diffusion_kv.kv_connector import (
+    init_worker_kv_connector,
+    shutdown_kv_connector,
+)
 from vllm_omni.diffusion.diffusion_kv.metadata import DiffusionKVMetadata
 from vllm_omni.diffusion.distributed.parallel_state import (
     destroy_distributed_env,
@@ -522,6 +526,7 @@ class DiffusionWorker:
         self.vllm_config.model_config.max_model_len = resolved_max_model_len
         kv_cache_config = kv_cache_configs[self.rank]
         self.vllm_config.cache_config.num_gpu_blocks = kv_cache_config.num_blocks
+        init_worker_kv_connector(self.vllm_config, kv_cache_config)
         with self._maybe_get_memory_pool_context("kv_cache"):
             self.model_runner.set_kv_cache_config(kv_cache_config)
 
@@ -978,13 +983,16 @@ class DiffusionWorker:
                         mgr.shutdown_prefetch()
         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: