[BugFix] fix mult cli timeout with get kv (#3741)

Signed-off-by: dengyunyang <584797741@qq.com>

tests/entrypoints/test_serve.py

@@ -325,7 +325,10 @@ def test_run_headless_diffusion_registers_and_spawns_proc(mocker: MockerFixture)
 
     run_headless(_make_headless_args(stage_id=1))
 
-    mock_inject.assert_called_once_with(stage_cfg, 1)
+    mock_inject.assert_called_once()
+    assert mock_inject.call_args.args[0] is stage_cfg
+    assert mock_inject.call_args.args[1] == 1
+    assert mock_inject.call_args.args[2] == [stage_cfg]
 
     reg_kwargs = mock_register.call_args.kwargs
     assert reg_kwargs["omni_master_address"] == "127.0.0.1"

vllm_omni/entrypoints/cli/serve.py

@@ -803,7 +803,10 @@ def run_headless(args: argparse.Namespace) -> None:
         metadata = extract_stage_metadata(stage_cfg)
         if omni_conn_cfg:
             inject_omni_kv_config(stage_cfg, omni_conn_cfg, omni_from, omni_to)
-        inject_kv_stage_info(stage_cfg, stage_id)
+        # Headless single-stage launch must still infer cross-stage TP topology
+        # from the loaded deploy config so heterogeneous KV routing keys match
+        # the head process (e.g. from_tp=2, to_tp=1).
+        inject_kv_stage_info(stage_cfg, stage_id, stage_configs)
         od_config = build_diffusion_config(model, stage_cfg, metadata)
 
         logger.info(