[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"])