[Feature] Bagel: Support tp+cfg parallel using mooncake transfer engine connector (#2705)

Signed-off-by: natureofnature <wzliu@connect.hku.hk>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>

tests/engine/test_async_omni_engine_stage_init.py

@@ -183,6 +183,7 @@ def test_launch_llm_stage_passes_stage_init_timeout_to_complete_stage_handshake(
     engine.model = "dummy-model"
     engine.single_stage_mode = False
     engine._omni_master_server = None
+    engine.stage_configs = []
 
     metadata = types.SimpleNamespace(stage_id=0, runtime_cfg={"devices": "0"})
     fake_vllm_config = types.SimpleNamespace()
@@ -238,6 +239,7 @@ def test_launch_llm_stage_releases_launch_lock_before_complete_stage_handshake(m
     engine.model = "dummy-model"
     engine.single_stage_mode = False
     engine._omni_master_server = None
+    engine.stage_configs = []
 
     fake_vllm_config = types.SimpleNamespace()
     fake_addresses = types.SimpleNamespace()
@@ -378,3 +380,70 @@ def test_attach_llm_stage_uses_omni_input_preprocessor(monkeypatch):
     assert input_processor is not None
     assert isinstance(input_processor.input_preprocessor, DummyOmniInputPreprocessor)
     assert input_processor.input_preprocessor.renderer is input_processor.renderer
+
+
+def test_inject_kv_stage_info_infers_sender_tp_topology():
+    from vllm_omni.engine.stage_init_utils import inject_kv_stage_info
+
+    stage0 = types.SimpleNamespace(
+        stage_id=0,
+        engine_args={
+            "tensor_parallel_size": 4,
+            "omni_kv_config": {
+                "need_send_cache": True,
+                "omni_from_stage": "0",
+                "omni_to_stage": "1",
+            },
+        },
+        engine_input_source=[],
+    )
+    stage1 = types.SimpleNamespace(
+        stage_id=1,
+        engine_args={
+            "parallel_config": {
+                "tensor_parallel_size": 2,
+                "cfg_parallel_size": 1,
+            },
+            "omni_kv_config": {"need_recv_cache": True},
+        },
+        engine_input_source=[0],
+    )
+
+    inject_kv_stage_info(stage0, 0, [stage0, stage1])
+
+    assert stage0.engine_args["omni_kv_config"]["stage_id"] == 0
+    assert stage0.engine_args["omni_kv_config"]["rank_mapping"] == {"from_tp": 4, "to_tp": 2}
+
+
+def test_inject_kv_stage_info_infers_receiver_tp_topology():
+    from vllm_omni.engine.stage_init_utils import inject_kv_stage_info
+
+    stage0 = types.SimpleNamespace(
+        stage_id=0,
+        engine_args={
+            "tensor_parallel_size": 4,
+            "omni_kv_config": {"need_send_cache": True},
+        },
+        engine_input_source=[],
+    )
+    stage1 = types.SimpleNamespace(
+        stage_id=1,
+        engine_args={
+            "parallel_config": {
+                "tensor_parallel_size": 2,
+                "cfg_parallel_size": 1,
+            },
+            "omni_kv_config": {
+                "need_recv_cache": True,
+                "omni_from_stage": "0",
+                "omni_to_stage": "1",
+            },
+        },
+        engine_input_source=[0],
+    )
+
+    inject_kv_stage_info(stage1, 1, [stage0, stage1])
+
+    assert stage1.engine_args["omni_kv_config"]["stage_id"] == 1
+    assert stage1.engine_args["omni_kv_config"]["engine_input_source"] == [0]
+    assert stage1.engine_args["omni_kv_config"]["rank_mapping"] == {"from_tp": 4, "to_tp": 2}

tests/engine/test_single_stage_mode.py

@@ -1555,6 +1555,7 @@ class TestLaunchLlmStageSingleStageMode:
         engine.single_stage_mode = True
         engine._single_stage_id_filter = 0
         engine._llm_stage_launch_lock = threading.Lock()
+        engine.stage_configs = []
         mock_oms = mocker.Mock(spec=OmniMasterServer)
         mock_oms.address = "127.0.0.1"
         mock_oms.port = 25000
@@ -1629,6 +1630,7 @@ class TestLaunchLlmStageSingleStageMode:
         engine.single_stage_mode = False
         engine._omni_master_server = None
         engine._llm_stage_launch_lock = threading.Lock()
+        engine.stage_configs = []
 
         fake_vllm_config = mocker.Mock()
         fake_executor_cls = mocker.Mock()