diff --git a/tests/entrypoints/test_serve.py b/tests/entrypoints/test_serve.py index 86c311a60..0915ddead 100644 --- a/tests/entrypoints/test_serve.py +++ b/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" diff --git a/vllm_omni/entrypoints/cli/serve.py b/vllm_omni/entrypoints/cli/serve.py index 50b825d7c..e4fddc830 100644 --- a/vllm_omni/entrypoints/cli/serve.py +++ b/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(