A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 5, 09:45 UTC
[Core] Simplify API server handshake (#39364)
Signed-off-by: Nick Hill <nickhill123@gmail.com>
Nick Hill committed
Apr 9, 2026 at 03:56 UTC
c8d98f81f676552c263f35bbde55e6edbe81b4e8
5 files changed
+53
-67
vllm/entrypoints/cli/serve.py
+8
-21
@@ -280,36 +280,23 @@ def run_multi_api_server(args: argparse.Namespace):
280
vllm_config, executor_class, log_stats, addresses, num_api_servers
281
) as (local_engine_manager, coordinator, addresses, tensor_queue):
282
# Construct common args for the APIServerProcessManager up-front.
283
- api_server_manager_kwargs = dict(
283
+ stats_update_address = None
284
+ if coordinator:
285
+ stats_update_address = coordinator.get_stats_publish_address()
286
+
287
+ # Start API servers.
288
+ api_server_manager = APIServerProcessManager(
289
listen_address=listen_address,
290
sock=sock,
291
args=args,
292
num_servers=num_api_servers,
293
input_addresses=addresses.inputs,
294
output_addresses=addresses.outputs,
290
- stats_update_address=coordinator.get_stats_publish_address()
291
- if coordinator
292
- else None,
295
+ stats_update_address=stats_update_address,
296
tensor_queue=tensor_queue,
297
)
298
296
- # For dp ranks > 0 in external/hybrid DP LB modes, we must delay the
297
- # start of the API servers until the local engine is started
298
- # (after the launcher context manager exits),
299
- # since we get the front-end stats update address from the coordinator
300
- # via the handshake with the local engine.
301
- if dp_rank == 0 or not parallel_config.local_engines_only:
302
- # Start API servers using the manager.
303
- api_server_manager = APIServerProcessManager(**api_server_manager_kwargs)
304
-
305
- # Start API servers now if they weren't already started.
306
- if api_server_manager is None:
307
- api_server_manager_kwargs["stats_update_address"] = (
308
- addresses.frontend_stats_publish_address
309
- )
310
- api_server_manager = APIServerProcessManager(**api_server_manager_kwargs)
311
-
312
- # Wait for API servers
299
+ # Wait for API servers.
300
try:
301
wait_for_completion_or_failure(
302
api_server_manager=api_server_manager,
vllm/v1/engine/__init__.py
+6
-6
@@ -4,6 +4,7 @@
4
import enum
5
import time
6
from collections.abc import Mapping
7
+from dataclasses import dataclass
8
from typing import Any, Literal
9
10
import msgspec
@@ -63,17 +64,16 @@ class FinishReason(enum.IntEnum):
64
return FINISH_REASON_STRINGS[self.value]
65
66
66
-class EngineCoreReadyResponse(
67
- msgspec.Struct,
68
- array_like=True, # type: ignore[call-arg]
69
- omit_defaults=True, # type: ignore[call-arg]
70
-):
71
- """Sent from EngineCore to the frontend during the ready handshake.
67
+@dataclass
68
+class EngineCoreReadyResponse:
69
+ """Sent from EngineCore to each frontend at the end of engine startup.
70
71
Contains post-initialization config that may differ from the original
72
values (e.g. max_model_len after KV cache auto-fitting).
73
"""
74
75
+ num_gpu_blocks: int
76
+ dp_stats_address: str | None
77
max_model_len: int | None = None
78
79
vllm/v1/engine/core.py
+9
-9
@@ -936,14 +936,20 @@ class EngineCoreProc(EngineCore):
936
vllm_config.parallel_config,
937
)
938
if client_handshake_address is None:
939
+ # We only need to handshake with one party.
940
with handshake as addresses:
941
yield addresses
942
else:
943
+ # We need to handshake with rank 0 front-end and our colocated frontend.
944
assert local_client
945
local_handshake = self._perform_handshake(
946
input_ctx, client_handshake_address, identity, True, False, vllm_config
947
)
948
with handshake as addresses, local_handshake as client_addresses:
949
+ # 1. Obtain DP Coordinator zmq address and DP process group address
950
+ # (addresses).
951
+ # 2. Add front-end input/output addresses from colocated front-end
952
+ # (client_addresses).
953
addresses.inputs = client_addresses.inputs
954
addresses.outputs = client_addresses.outputs
955
yield addresses
@@ -977,20 +983,12 @@ class EngineCoreProc(EngineCore):
983
yield addresses
984
985
# Send ready message.
980
- num_gpu_blocks = vllm_config.cache_config.num_gpu_blocks
981
- # We pass back the coordinator stats update address here for the
982
- # external LB case for our colocated front-end to use (coordinator
983
- # only runs with rank 0).
984
- dp_stats_address = self.frontend_stats_publish_address
985
-
986
- # Include config hash for DP configuration validation
986
ready_msg = {
987
"status": "READY",
988
"local": local_client,
989
"headless": headless,
991
- "num_gpu_blocks": num_gpu_blocks,
992
- "dp_stats_address": dp_stats_address,
990
}
991
+ # Include config hash for DP configuration validation
992
if vllm_config.parallel_config.data_parallel_size > 1:
993
ready_msg["parallel_config_hash"] = (
994
vllm_config.parallel_config.compute_hash()
@@ -1388,6 +1386,8 @@ class EngineCoreProc(EngineCore):
1386
poller = zmq.Poller()
1387
ready_response = EngineCoreReadyResponse(
1388
max_model_len=self.vllm_config.model_config.max_model_len,
1389
+ num_gpu_blocks=self.vllm_config.cache_config.num_gpu_blocks or 0,
1390
+ dp_stats_address=self.frontend_stats_publish_address,
1391
)
1392
ready_payload = msgspec.msgpack.encode(ready_response)
1393
for input_socket in input_sockets:
vllm/v1/engine/core_client.py
+30
-18
@@ -457,22 +457,6 @@ class ElasticScalingCache:
457
pending_notifications: dict[EEPNotificationType, set[int]]
458
459
460
-_ready_response_decoder = msgspec.msgpack.Decoder(EngineCoreReadyResponse)
461
-
462
-
463
-def _apply_ready_response(payload: bytes, vllm_config: VllmConfig) -> None:
464
- """Decode an EngineCoreReadyResponse and sync any post-initialization
465
- config changes (e.g. auto-fitted max_model_len) back to the frontend."""
466
- if not payload:
467
- return
468
- response = _ready_response_decoder.decode(payload)
469
- if response.max_model_len is not None:
470
- vllm_config.model_config.max_model_len = min(
471
- vllm_config.model_config.max_model_len,
472
- response.max_model_len,
473
- )
474
-
475
-
460
class MPClient(EngineCoreClient):
461
"""
462
MPClient: base client for multi-proc EngineCore.
@@ -608,7 +592,7 @@ class MPClient(EngineCoreClient):
592
)
593
identity, payload = sync_input_socket.recv_multipart()
594
identities.remove(identity)
611
- _apply_ready_response(payload, vllm_config)
595
+ self._apply_ready_response(payload)
596
597
self.core_engine: EngineIdentity = self.core_engines[0]
598
self.utility_results: dict[int, AnyFuture] = {}
@@ -680,6 +664,34 @@ class MPClient(EngineCoreClient):
664
target=monitor_engine_cores, daemon=True, name="MPClientEngineMonitor"
665
).start()
666
667
+ def _apply_ready_response(self, payload: bytes) -> None:
668
+ """Decode an EngineCoreReadyResponse and sync any post-initialization
669
+ config changes (e.g. auto-fitted max_model_len) back to the frontend."""
670
+ if not payload:
671
+ return
672
+ vllm_config = self.vllm_config
673
+ response = msgspec.msgpack.decode(payload, type=EngineCoreReadyResponse)
674
+ if response.max_model_len is not None:
675
+ vllm_config.model_config.max_model_len = min(
676
+ vllm_config.model_config.max_model_len,
677
+ response.max_model_len,
678
+ )
679
+
680
+ # Setup KV cache config with initialization state from
681
+ # engine core process. Sum values from all engines in DP case.
682
+ num_gpu_blocks = vllm_config.cache_config.num_gpu_blocks or 0
683
+ num_gpu_blocks += response.num_gpu_blocks
684
+ vllm_config.cache_config.num_gpu_blocks = num_gpu_blocks
685
+
686
+ # In external DP LB mode, the coordinator address that the
687
+ # front-end procs connect to is obtained by each engine via it's
688
+ # initial handshake with the rank 0 front-end.
689
+ if response.dp_stats_address is not None:
690
+ if self.stats_update_address is None:
691
+ self.stats_update_address = response.dp_stats_address
692
+ else:
693
+ assert response.dp_stats_address == self.stats_update_address
694
+
695
696
def _process_utility_output(
697
output: UtilityOutput, utility_results: dict[int, AnyFuture]
@@ -1602,7 +1614,7 @@ class DPLBAsyncMPClient(DPAsyncMPClient):
1614
)
1615
identity, payload = sync_input_socket.recv_multipart()
1616
new_engine_identities.discard(identity)
1605
- _apply_ready_response(payload, self.vllm_config)
1617
+ self._apply_ready_response(payload)
1618
1619
# NOTE(yongji): Before we schedule any requests on the new workers,
1620
# we should wait for them to switch to the new setup.
vllm/v1/engine/utils.py
-13
@@ -1212,19 +1212,6 @@ def wait_for_engine_startup(
1212
start_pending[0 if local else 1] += 1
1213
engine.state = CoreEngineState.CONNECTED
1214
elif status == "READY" and engine.state == CoreEngineState.CONNECTED:
1215
- # Setup KV cache config with initialization state from
1216
- # engine core process. Sum values from all engines in DP case.
1217
- num_gpu_blocks = cache_config.num_gpu_blocks or 0
1218
- num_gpu_blocks += msg["num_gpu_blocks"]
1219
- cache_config.num_gpu_blocks = num_gpu_blocks
1220
-
1221
- # In external DP LB mode, the coordinator address that the
1222
- # front-end procs connect to is obtained from rank 0 via
1223
- # one of the engine handshakes, and passed to the local
1224
- # front-end process in the response from the other.
1225
- if addresses.frontend_stats_publish_address is None:
1226
- addresses.frontend_stats_publish_address = msg.get("dp_stats_address")
1227
-
1215
# Validate config hash consistency across DP workers for MoE models.
1216
if coordinated_dp:
1217
worker_config_hash = msg.get("parallel_config_hash")