[Bugfix] Restore BAGEL CFG KV transfers (#5620)

Signed-off-by: lishunyang12 <lishunyang03@gmail.com>
Co-authored-by: lishunyang12 <lishunyang03@gmail.com>

vllm_omni/engine/async_engine_utils.py

@@ -77,12 +77,17 @@ def upgrade_to_omni_request(
 def apply_omni_final_stage_metadata(
     request: EngineCoreRequest,
     final_stage_id: int,
+    *,
+    force_kv_transfer: bool = False,
 ) -> EngineCoreRequest:
-    """Tag a request so the AR scheduler can skip DiT KV for stage zero."""
+    """Tag a request with its final stage and optional KV-transfer override."""
     merged: dict[str, Any] = {}
     if isinstance(request, OmniEngineCoreRequest) and request.additional_information is not None:
         merged = deserialize_additional_information(request.additional_information)
     merged["omni_final_stage_id"] = final_stage_id
+    merged.pop("omni_force_kv_transfer", None)
+    if force_kv_transfer:
+        merged["omni_force_kv_transfer"] = True
     payload = serialize_additional_information(merged)
     return OmniEngineCoreRequest.from_request(
         request,

vllm_omni/engine/async_omni_engine.py

@@ -798,11 +798,9 @@ class AsyncOmniEngine:
                 supported_tasks=self.supported_tasks,
             )
             request.external_req_id = cid
-            # Companions are stage-0-final: without this worker-visible tag
-            # the AR runner would still emit downstream connector payloads
-            # for them (the orchestrator-side final_stage_id=0 registration
-            # alone does not reach the worker).
-            request = apply_omni_final_stage_metadata(request, 0)
+            # Companions are stage-0-final for ordinary downstream payloads,
+            # but diffusion still needs their CFG KV caches.
+            request = apply_omni_final_stage_metadata(request, 0, force_kv_transfer=True)
 
             # Registration of this companion on stage-0's output processor is
             # deferred to Orchestrator._handle_add_companion, which routes