[Feature]: Support cfg kv-cache transfer in multi-stage (#1422)

Signed-off-by: princepride <wangzhipeng628@gmail.com>
Signed-off-by: 汪志鹏 <wangzhipeng628@gmail.com>
Co-authored-by: Cursor <cursoragent@cursor.com>

vllm_omni/entrypoints/async_omni_diffusion.py

@@ -60,6 +60,7 @@ class AsyncOmniDiffusion:
         # Capture stage info from kwargs before they might be filtered out
         stage_id = kwargs.get("stage_id")
         engine_input_source = kwargs.get("engine_input_source")
+        cfg_kv_collect_func = kwargs.pop("cfg_kv_collect_func", None)
 
         # Build config
         if od_config is None:
@@ -115,6 +116,9 @@ class AsyncOmniDiffusion:
             else:
                 raise
 
+        if cfg_kv_collect_func is not None:
+            od_config.cfg_kv_collect_func = cfg_kv_collect_func
+
         # Initialize engine
         self.engine: DiffusionEngine = DiffusionEngine.make_engine(od_config)