[Core] Don't issue a blocking collective RPC during engine handshake (#57226)
Signed-off-by: okorzh <okorzh@amd.com> Co-authored-by: Andreas Karatzas <akaratza@amd.com>
Oxana Korzh committed
Sep 16, 2026 at 20:30 UTC
2307a1d16cdd0077cd882bb9fc41b453f6f1999e
2 files changed
+40
tests/v1/engine/test_engine_core_client.py
+38
@@ -1452,3 +1452,41 @@ def test_engine_core_proc_instantiation_cuda_empty(monkeypatch: pytest.MonkeyPat
1452
)
1453
1454
engine_core_proc.shutdown()
1455
+
1456
+
1457
+def _ready_response_with_weight_transfer(weight_transfer_config):
1458
+ """Build a ready response against a mock executor."""
1459
+ from vllm.v1.engine.core import EngineCoreProc
1460
+
1461
+ executor = MagicMock()
1462
+ executor.supports_draft_weight_updates.return_value = True
1463
+
1464
+ proc = MagicMock()
1465
+ proc.model_executor = executor
1466
+ proc.vllm_config = MagicMock()
1467
+ proc.vllm_config.weight_transfer_config = weight_transfer_config
1468
+ proc.vllm_config.lora_config = None
1469
+ proc.vllm_config.model_config.enable_sleep_mode = False
1470
+ proc.scheduler = MagicMock()
1471
+
1472
+ response = EngineCoreProc._make_ready_response(proc)
1473
+ return response, executor
1474
+
1475
+
1476
+def test_ready_response_skips_draft_rpc_without_weight_transfer():
1477
+ """No weight-transfer config: the executor must not be consulted."""
1478
+ response, executor = _ready_response_with_weight_transfer(None)
1479
+
1480
+ executor.supports_draft_weight_updates.assert_not_called()
1481
+ assert response.supports_draft_weight_updates is False
1482
+
1483
+
1484
+def test_ready_response_queries_executor_with_weight_transfer():
1485
+ """Weight transfer configured: the executor is still consulted."""
1486
+ weight_transfer_config = MagicMock()
1487
+ weight_transfer_config.backend = "nixl"
1488
+
1489
+ response, executor = _ready_response_with_weight_transfer(weight_transfer_config)
1490
+
1491
+ executor.supports_draft_weight_updates.assert_called_once()
1492
+ assert response.supports_draft_weight_updates is True
vllm/v1/engine/core.py
+2
@@ -1691,6 +1691,8 @@ class EngineCoreProc(EngineCore):
1691
enable_sleep_mode=self.vllm_config.model_config.enable_sleep_mode,
1692
supports_draft_weight_updates=(
1693
self.model_executor.supports_draft_weight_updates()
1694
+ if self.vllm_config.weight_transfer_config is not None
1695
+ else False
1696
),
1697
)
1698