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

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

tests/core/sched/test_omni_ar_scheduler_kv_transfer.py

@@ -0,0 +1,57 @@
+from types import SimpleNamespace
+
+import pytest
+from vllm import SamplingParams
+from vllm.v1.engine import EngineCoreRequest
+
+from vllm_omni.core.sched.omni_ar_scheduler import OmniARScheduler
+from vllm_omni.engine.async_engine_utils import apply_omni_final_stage_metadata
+from vllm_omni.engine.serialization import deserialize_additional_information
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def _make_engine_request() -> EngineCoreRequest:
+    return EngineCoreRequest(
+        request_id="req",
+        prompt_token_ids=[1],
+        mm_features=None,
+        sampling_params=SamplingParams(max_tokens=1),
+        pooling_params=None,
+        arrival_time=0.0,
+        lora_request=None,
+        cache_salt=None,
+        data_parallel_rank=None,
+    )
+
+
+def _request_omits_kv_transfer(*, force_kv_transfer: bool) -> tuple[bool, dict]:
+    tagged = apply_omni_final_stage_metadata(
+        _make_engine_request(),
+        final_stage_id=0,
+        force_kv_transfer=force_kv_transfer,
+    )
+    scheduler = OmniARScheduler.__new__(OmniARScheduler)
+    scheduler._omits_kv_transfer_cache = {}
+    request = SimpleNamespace(
+        request_id="req",
+        additional_information=tagged.additional_information,
+    )
+    result = scheduler._request_omits_kv_transfer_to_next_stage(request)
+    metadata = deserialize_additional_information(tagged.additional_information)
+    return result, metadata
+
+
+def test_stage_zero_request_omits_kv_transfer():
+    omits_transfer, metadata = _request_omits_kv_transfer(force_kv_transfer=False)
+
+    assert omits_transfer
+    assert "omni_force_kv_transfer" not in metadata
+
+
+def test_cfg_companion_forces_kv_transfer_without_downstream_payload():
+    omits_transfer, metadata = _request_omits_kv_transfer(force_kv_transfer=True)
+
+    assert not omits_transfer
+    assert metadata["omni_final_stage_id"] == 0
+    assert metadata["omni_force_kv_transfer"] is True

tests/engine/test_async_omni_engine_input.py

@@ -7,7 +7,9 @@ from vllm.v1.engine import EngineCoreRequest
 from vllm_omni.distributed.omni_coordinator import ReplicaInfo, ReplicaStatus
 from vllm_omni.engine import OmniEngineCoreRequest
 from vllm_omni.engine.async_omni_engine import AsyncOmniEngine, StageRuntimeInfo
+from vllm_omni.engine.serialization import deserialize_additional_information
 from vllm_omni.engine.stage_pool import StagePool
+from vllm_omni.model_executor.stage_input_processors.bagel import ExpandedPrompt
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
@@ -107,6 +109,34 @@ def test_build_add_request_message_preserves_model_intermediate_buffer(mocker: M
     assert torch.equal(info["hidden_states"]["tts"], hidden)
 
 
+def test_cfg_companion_suppresses_payload_but_forces_kv_transfer(mocker: MockerFixture):
+    engine = object.__new__(AsyncOmniEngine)
+    params = SamplingParams(max_tokens=8)
+    engine.prompt_expand_func = lambda *_args: [
+        ExpandedPrompt(
+            prompt={"prompt": "negative"},
+            role="cfg_text",
+            request_id_suffix="__cfg_text",
+        )
+    ]
+    engine.supported_tasks = ("generate",)
+    engine.input_processor = mocker.Mock()
+    engine.input_processor.process_inputs.return_value = _make_engine_core_request("req__cfg_text")
+    engine.request_queue = mocker.Mock()
+
+    engine._enqueue_cfg_companions(
+        parent_id="req",
+        original_prompt={"prompt": "positive"},
+        stage0_params=params,
+        sampling_params_list=[params],
+    )
+
+    message = engine.request_queue.sync_q.put.call_args.args[0]
+    metadata = deserialize_additional_information(message.prompt.additional_information)
+    assert metadata["omni_final_stage_id"] == 0
+    assert metadata["omni_force_kv_transfer"] is True
+
+
 def test_build_add_request_message_with_resumable_streaming(mocker: MockerFixture):
     engine = object.__new__(AsyncOmniEngine)
     params = SamplingParams(max_tokens=8)

tests/worker/test_omni_connector_mixin.py

@@ -17,6 +17,9 @@ from unittest.mock import MagicMock, patch
 import pytest
 import torch
 
+from vllm_omni.distributed.omni_connectors.kv_transfer_manager import (
+    OmniKVTransferManager,
+)
 from vllm_omni.outputs import OmniConnectorOutput
 from vllm_omni.worker.omni_connector_model_runner_mixin import (
     OmniConnectorModelRunnerMixin,
@@ -1198,6 +1201,38 @@ class TestRankAwareKVRouting(unittest.TestCase):
         )
         host.shutdown_omni_connectors()
 
+    def test_single_rank_hooks_use_default_manager_key(self):
+        host = self._make_host(from_tp=1, to_tp=1, local_rank=0)
+        expected = ["omni_0_to_1_kv_cache_req"]
+
+        self.assertEqual(host.get_rank_aware_kv_send_keys("req", from_stage=0, to_stage=1), expected)
+        self.assertEqual(host.get_rank_aware_kv_keys("req", from_stage=0, to_stage=1), expected)
+        host.shutdown_omni_connectors()
+
+    def test_init_uses_effective_manager_tp2_topology(self):
+        model_config = _make_model_config(stage_id=0)
+        model_config.omni_kv_config = {
+            "need_send_cache": True,
+            "omni_from_stage": "0",
+            "omni_to_stage": "1",
+            "rank_mapping": {"from_tp": 2, "to_tp": 2},
+        }
+        with patch(
+            "vllm_omni.distributed.omni_connectors.kv_transfer_manager.get_local_tp_rank",
+            return_value=1,
+        ):
+            manager = OmniKVTransferManager.from_model_config(model_config)
+
+        host = MixinHost()
+        host.init_omni_connectors(model_config=model_config, kv_transfer_manager=manager)
+
+        expected = ["req_0_0_1_1"]
+        self.assertEqual(host.get_rank_aware_kv_send_keys("req", from_stage=0, to_stage=1), expected)
+        self.assertEqual(host.get_rank_aware_kv_keys("req", from_stage=0, to_stage=1), expected)
+        self.assertEqual(host.get_kv_rank_mapping()["from_tp"], 2)
+        self.assertEqual(host.get_kv_rank_mapping()["to_tp"], 2)
+        host.shutdown_omni_connectors()
+
     def test_send_keys_route_from_rank_gt_to_rank(self):
         host = self._make_host(from_tp=4, to_tp=2, local_rank=3)
         self.assertEqual(host.get_rank_aware_kv_send_keys("req", from_stage=0), ["req_0_0_3_1"])

vllm_omni/core/sched/omni_ar_scheduler.py

@@ -147,7 +147,7 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
         return None
 
     def _request_omits_kv_transfer_to_next_stage(self, request: Request) -> bool:
-        """True when orchestrator will not run stage 1+ for this request (e.g. text-only).
+        """True when this stage-zero-final request does not need downstream KV.
 
         The result is cached per request to avoid repeated deserialization of
         additional_information on every scheduler tick.
@@ -162,7 +162,7 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
             result = False
         else:
             info = deserialize_additional_information(payload)
-            result = info.get("omni_final_stage_id") == 0
+            result = info.get("omni_final_stage_id") == 0 and not bool(info.get("omni_force_kv_transfer", False))
 
         self._omits_kv_transfer_cache[rid] = result
         return result

vllm_omni/distributed/omni_connectors/kv_transfer_manager.py

@@ -403,6 +403,7 @@ class OmniKVTransferManager:
         self._prefetch_executor: ThreadPoolExecutor | None = (
             ThreadPoolExecutor(max_workers=1, thread_name_prefix="kv-prefetch") if async_prefetch else None
         )
+
         self._prefetch_futures: dict[str, Any] = {}
         self._bg_copy_stream: current_omni_platform.Stream | None = None
 
@@ -415,6 +416,11 @@ class OmniKVTransferManager:
             except Exception as e:
                 logger.warning("Failed to eagerly initialize sender connector: %s", e)
 
+    @property
+    def tp_topology(self) -> KVTPTopology:
+        """Return the effective topology after config and runtime detection."""
+        return self._tp_topo
+
     @property
     def topo_config(self) -> _TransferTopoConfig:
         """Lazily build the parallel topology snapshot on first access."""

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

vllm_omni/worker/omni_connector_model_runner_mixin.py

@@ -122,6 +122,19 @@ class OmniConnectorModelRunnerMixin:
 
         # -- heterogeneous TP rank support --
         rank_cfg = self._parse_rank_mapping(model_config)
+        if self._kv_transfer_manager is not None:
+            topology = getattr(self._kv_transfer_manager, "tp_topology", None)
+            effective_mapping = (
+                getattr(topology, "source_tp_size", None),
+                getattr(topology, "target_tp_size", None),
+                getattr(topology, "local_rank", None),
+            )
+            if all(isinstance(value, int) for value in effective_mapping):
+                rank_cfg = {
+                    "from_tp": effective_mapping[0],
+                    "to_tp": effective_mapping[1],
+                    "local_rank": effective_mapping[2],
+                }
         self._from_tp: int = rank_cfg["from_tp"]
         self._to_tp: int = rank_cfg["to_tp"]
         self._local_rank: int = rank_cfg["local_rank"]
@@ -1275,6 +1288,10 @@ class OmniConnectorModelRunnerMixin:
         For heterogeneous TP receive, the local rank is the target rank and must
         fetch one or more source-rank shards keyed as ``from_rank -> to_rank``.
         """
+        if self._from_tp <= 1 and self._to_tp <= 1:
+            resolved_to_stage = self._next_stage_id if to_stage is None else to_stage
+            return [f"omni_{from_stage}_to_{resolved_to_stage}_kv_cache_{req_id}"]
+
         remote_ranks = self.get_kv_remote_ranks()
         return [
             self.get_kv_connector_key(
@@ -1307,6 +1324,10 @@ class OmniConnectorModelRunnerMixin:
         chunk_id: int = 0,
     ) -> list[str]:
         """Build send-side connector keys for this rank's KV shard(s)."""
+        if self._from_tp <= 1 and self._to_tp <= 1:
+            resolved_to_stage = self._next_stage_id if to_stage is None else to_stage
+            return [f"omni_{from_stage}_to_{resolved_to_stage}_kv_cache_{req_id}"]
+
         target_ranks = self.get_kv_target_ranks_for_send()
         return [
             self.get_kv_connector_key(