From 815842c827d3217e0e502be2cd9a3e95903d0d57 Mon Sep 17 00:00:00 2001 From: Siddharth Singh Date: Wed, 12 Aug 2026 14:15:23 -0700 Subject: [PATCH 1/2] modify coordinator wire format - split into two chunks - first chunk is metadata of constant size. This is the only thing that the coordinator needs to unpack/read and pack --- .../coordinator.py | 35 ++- .../handlers.py | 267 ++++++++++++------ .../core/inference/engines/dynamic_engine.py | 130 +++++++-- megatron/core/inference/inference_client.py | 105 +++++-- .../inference/engines/test_dynamic_engine.py | 12 +- ...est_data_parallel_inference_coordinator.py | 137 ++++++++- ...test_dynamic_prefix_caching_coordinator.py | 82 ++++-- .../inference/test_inference_client.py | 31 +- .../test_inference_client_streaming.py | 61 ++-- 9 files changed, 647 insertions(+), 213 deletions(-) diff --git a/megatron/core/inference/data_parallel_inference_coordinator/coordinator.py b/megatron/core/inference/data_parallel_inference_coordinator/coordinator.py index 67eb1bb5829..18740667d9f 100644 --- a/megatron/core/inference/data_parallel_inference_coordinator/coordinator.py +++ b/megatron/core/inference/data_parallel_inference_coordinator/coordinator.py @@ -303,14 +303,18 @@ def _remove_engine(self, identity): len(self.identities_of_data_parallel_ranks), ) - def _send_to_engine(self, identity, payload): - """Send payload to an engine, removing it from the pool if unreachable. + def _send_to_engine(self, identity, frames): + """Send a message to an engine, removing it from the pool if unreachable. + + Args: + identity: ZMQ identity of the target engine. + frames (list): Raw frames to send, metadata frame first. Returns: True if the send succeeded, False if the engine was unreachable and removed. """ try: - self.router_socket.send_multipart([identity, payload]) + self.router_socket.send_multipart([identity, *frames]) return True except zmq.error.ZMQError as e: if e.errno == zmq.EHOSTUNREACH: @@ -322,19 +326,21 @@ def _broadcast_to_engines(self, payload): """Send a deserialized payload to every connected data parallel rank.""" serialized = msgpack.packb(payload, use_bin_type=True) for data_parallel_rank_id in list(self.identities_of_data_parallel_ranks): - self._send_to_engine(data_parallel_rank_id, serialized) + self._send_to_engine(data_parallel_rank_id, [serialized]) def compute_request_hashes(self, prompt): """Compute block hashes for a prompt on CPU. + Callers decide whether hashes are wanted at all: computing them requires + the decoded prompt, and decoding it is the cost the caller is usually + trying to avoid. See handle_submit_request. + Args: prompt: Either a string (to be tokenized) or a list of token IDs. Returns: - List of integer block hashes, or empty list if prefix caching is disabled. + List of integer block hashes, one per complete block. """ - if not self.enable_prefix_caching or self.block_size_tokens is None: - return [] if isinstance(prompt, str): tokens = self.tokenizer.tokenize(prompt) else: @@ -431,20 +437,25 @@ def start(self): """ # Todo [Siddharth]: Make this more robust to handle invalid messages. while True: - sender_identity, serialized_payload = self.router_socket.recv_multipart() + # Messages are one or more frames. frames[0] is the metadata frame: + # a header plus whatever the coordinator needs to route the message. + # Any later frames are opaque payload bodies, forwarded without ever + # being decoded here -- that is what keeps this loop's cost + # independent of prompt length. + sender_identity, *frames = self.router_socket.recv_multipart() # An empty payload is a data parallel rank (re-)registering itself. - if serialized_payload == b"": + if frames[0] == b"": self._handle_rank_registration(sender_identity) continue - deserialized_payload = msgpack.unpackb(serialized_payload, raw=False) - header = Headers(deserialized_payload[0]) + metadata = msgpack.unpackb(frames[0], raw=False) + header = Headers(metadata[0]) handler = self._handlers.get(header) if handler is None: raise UnknownHeaderError(header) - if handler(self, sender_identity, deserialized_payload): + if handler(self, sender_identity, metadata, frames[1:]): break def _handle_rank_registration(self, sender_identity): diff --git a/megatron/core/inference/data_parallel_inference_coordinator/handlers.py b/megatron/core/inference/data_parallel_inference_coordinator/handlers.py index b514b5e62c7..c1d57ebfa54 100644 --- a/megatron/core/inference/data_parallel_inference_coordinator/handlers.py +++ b/megatron/core/inference/data_parallel_inference_coordinator/handlers.py @@ -8,15 +8,23 @@ is supported simply by adding a decorated function here; the coordinator's event loop never changes. -Handlers have the signature ``(coordinator, sender_identity, payload) -> bool | None`` -where ``payload`` is the already-deserialized message. Returning a truthy value -signals the coordinator's event loop to stop. +Handlers have the signature +``(coordinator, sender_identity, metadata, bodies) -> bool | None``. + +``metadata`` is the decoded contents of the message's first frame: a header +followed by whatever the coordinator needs in order to route the message. It is +small by construction and does not grow with prompt length. + +``bodies`` is the list of remaining raw frames, still packed. These carry the +prompt (inbound) or the finished request (outbound) -- the bulk of the bytes. +The coordinator forwards them as opaque frames and only decodes one when it has +to mutate it, which keeps its per-request cost flat in prompt length. + +Returning a truthy value signals the event loop to stop. """ import logging -import torch - from megatron.core.inference.config import PrefixCachingCoordinatorPolicy from megatron.core.inference.headers import Headers @@ -52,8 +60,14 @@ def decorator(fn): @message_handler(Headers.CONNECT) -def handle_connect(coordinator, sender_identity, payload): - """Handshake with a new client, replying with a CONNECT_ACK.""" +def handle_connect(coordinator, sender_identity, metadata, bodies): + """Handshake with a new client, replying with a CONNECT_ACK. + + Sent by ``InferenceClient.start``. + + ``metadata``: ``[header]``. + ``bodies``: empty. + """ if sender_identity in coordinator.known_clients: logging.info(f"Client {sender_identity} sent a duplicate connect request. Ignoring ..") return @@ -65,22 +79,27 @@ def handle_connect(coordinator, sender_identity, payload): @message_handler(Headers.SUBMIT_REQUEST) -def handle_submit_request(coordinator, sender_identity, payload): +def handle_submit_request(coordinator, sender_identity, metadata, bodies): """Route a client request to a data parallel rank. + Sent by ``InferenceClient.add_request`` / ``add_request_streaming``. + + ``metadata``: ``[header, client_request_id, sampling_params]``, where + ``sampling_params`` is the serialized dict. + ``bodies``: ``[prompt]`` -- one frame holding the packed prompt (a string or + a token id list). Forwarded to the engine untouched, and decoded here + only when the routing policy needs block hashes. + Returns True (stopping the loop) if no engines are reachable. """ - # ToDo [Siddharth]: We might want to tokenize the prompt on the - # assigned data parallel rank for this process instead - # of the coordinator. - # Message from a known client if sender_identity not in coordinator.known_clients: logging.info(f"Received message from unknown client {sender_identity}. Ignoring.") return - # this is a message from a client. - # route it to a data parallel rank - client_request_id, prompt, sampling_params = payload[1:] + + _, client_request_id, sampling_params = metadata + prompt_frame = bodies[0] + # map client request_id to server request_id # necessary because multiple clients might have the same request_id. request_id = coordinator.next_request_id @@ -89,19 +108,26 @@ def handle_submit_request(coordinator, sender_identity, payload): coordinator.request_id_to_client_request_id[request_id] = client_request_id coordinator.client_request_to_request_id[(sender_identity, client_request_id)] = request_id - # Serialize prompt. - if isinstance(prompt, (str, list)): - pass - elif isinstance(prompt, torch.Tensor): - prompt = prompt.tolist() - else: - raise Exception("specialize for <%s> prompt." % type(prompt).__name__) - - engine_payload = msgpack.packb( - [Headers.SUBMIT_REQUEST.value, request_id, prompt, sampling_params], use_bin_type=True + # Rebuilding the metadata frame is cheap: it holds no prompt tokens. + engine_metadata = msgpack.packb( + [Headers.SUBMIT_REQUEST.value, request_id, sampling_params], use_bin_type=True ) - request_hashes = coordinator.compute_request_hashes(prompt) + # Only prefix-affinity routing consults the block hashes. When nothing will + # look at them, skip hashing *and* the prompt decode it needs -- avoiding + # that decode is the reason the prompt travels in its own frame. + if ( + coordinator.enable_prefix_caching + and coordinator.block_size_tokens is not None + and coordinator.prefix_caching_coordinator_policy + != PrefixCachingCoordinatorPolicy.LOAD_BALANCED + ): + request_hashes = coordinator.compute_request_hashes( + msgpack.unpackb(prompt_frame, raw=False) + ) + else: + request_hashes = [] + if ( coordinator.prefix_caching_coordinator_policy == PrefixCachingCoordinatorPolicy.FIRST_PREFIX_BLOCK @@ -111,7 +137,7 @@ def handle_submit_request(coordinator, sender_identity, payload): # Account for the fact that some engines may have died. for _ in range(len(coordinator.identities_of_data_parallel_ranks)): next_identity = coordinator.get_best_data_parallel_rank(request_hashes) - if coordinator._send_to_engine(next_identity, engine_payload): + if coordinator._send_to_engine(next_identity, [engine_metadata, prompt_frame]): break else: # If all engines have died, we are in an abnormal state, and must exit cleanly. @@ -136,46 +162,53 @@ def handle_submit_request(coordinator, sender_identity, payload): @message_handler(Headers.SUBMIT_REQUEST_WITH_KV) -def handle_submit_request_with_kv(coordinator, sender_identity, payload): - """Route a client-supplied KV handoff to a decode engine.""" +def handle_submit_request_with_kv(coordinator, sender_identity, metadata, bodies): + """Route a client-supplied KV handoff to a decode engine. + + Sent by ``InferenceClient.add_request_with_kv_handoff`` / + ``add_request_with_kv_handoff_streaming``. + + ``metadata``: ``[header, client_request_id, sampling_params, kv_meta]``, + where ``kv_meta`` is the peer's NIXL agent/layout export. It is bounded + by TP size and num_speculative_tokens, not by prompt length, so it stays + in the metadata frame. + ``bodies``: ``[prompt, src_block_ids]``. ``src_block_ids`` names one remote + block per block_size_tokens of prompt, so it grows with the prompt and + travels as its own frame. Both bodies are forwarded to the engine + untouched, so the coordinator decodes nothing sequence-dependent. In + disaggregated serving every decode request arrives here, so this is as + hot as a plain submission. + """ if sender_identity not in coordinator.known_clients: logging.info( "Received SUBMIT_REQUEST_WITH_KV from unknown client %s; ignoring.", sender_identity ) return - if len(payload) != 6: + if len(metadata) != 4 or len(bodies) != 2: logging.error( - "Coordinator: malformed SUBMIT_REQUEST_WITH_KV payload with %d fields", len(payload) - 1 + "Coordinator: malformed SUBMIT_REQUEST_WITH_KV with %d metadata fields, %d bodies", + len(metadata) - 1, + len(bodies), ) return - client_request_id, prompt, sampling_params, kv_meta, src_block_ids = payload[1:] + _, client_request_id, sampling_params, kv_meta = metadata request_id = coordinator.next_request_id coordinator.next_request_id += 1 coordinator.request_id_to_client_id[request_id] = sender_identity coordinator.request_id_to_client_request_id[request_id] = client_request_id coordinator.client_request_to_request_id[(sender_identity, client_request_id)] = request_id - if isinstance(prompt, torch.Tensor): - prompt = prompt.tolist() - elif not isinstance(prompt, (str, list)): - raise TypeError(f"unsupported prompt type {type(prompt).__name__}") - engine_payload = msgpack.packb( - [ - Headers.SUBMIT_REQUEST_WITH_KV.value, - request_id, - prompt, - sampling_params, - kv_meta, - src_block_ids, - ], + # Rebuilding the metadata frame is cheap: it holds no prompt tokens. + engine_metadata = msgpack.packb( + [Headers.SUBMIT_REQUEST_WITH_KV.value, request_id, sampling_params, kv_meta], use_bin_type=True, ) for _ in range(len(coordinator.identities_of_data_parallel_ranks)): next_identity = coordinator.get_least_loaded_data_parallel_rank() - if coordinator._send_to_engine(next_identity, engine_payload): + if coordinator._send_to_engine(next_identity, [engine_metadata, *bodies]): break else: logging.error("Coordinator: no reachable engines for handoff request %d", request_id) @@ -189,13 +222,20 @@ def handle_submit_request_with_kv(coordinator, sender_identity, payload): @message_handler(Headers.RELEASE_KV) -def handle_release_kv(coordinator, sender_identity, payload): - """Broadcast release of prefill blocks retained for a completed handoff.""" +def handle_release_kv(coordinator, sender_identity, metadata, bodies): + """Broadcast release of prefill blocks retained for a completed handoff. + + Sent by ``InferenceClient.release_handoff``. Broadcast to every engine; + engines not holding that request id treat it as a no-op. + + ``metadata``: ``[header, client_request_id]``. + ``bodies``: empty. + """ if sender_identity not in coordinator.known_clients: logging.warning("Coordinator: ignoring RELEASE_KV from unknown client.") return - coordinator._broadcast_to_engines([Headers.RELEASE_KV.value, int(payload[1])]) + coordinator._broadcast_to_engines([Headers.RELEASE_KV.value, int(metadata[1])]) @message_handler( @@ -206,13 +246,22 @@ def handle_release_kv(coordinator, sender_identity, payload): Headers.SET_GENERATION_EPOCH, Headers.STOP, ) -def handle_control_signal(coordinator, sender_identity, payload): - """Validate a control signal against the transition table and broadcast it.""" +def handle_control_signal(coordinator, sender_identity, metadata, bodies): + """Validate a control signal against the transition table and broadcast it. + + Serves PAUSE, UNPAUSE, SUSPEND, RESUME, SET_GENERATION_EPOCH and STOP, all + sent by ``InferenceClient._send_signal_to_engines``. + + ``metadata``: ``[header, *args]``. Every signal but one carries no args; + SET_GENERATION_EPOCH carries ``[header, generation_epoch]``. The whole + list is rebroadcast verbatim so data-bearing signals keep their args. + ``bodies``: empty. + """ if sender_identity not in coordinator.known_clients: logging.warning("Coordinator: ignoring signal from unknown client.") return - header = Headers(payload[0]) + header = Headers(metadata[0]) transition = CONTROL_TRANSITIONS[header] if coordinator.state not in transition.allowed_from: # Silently ignore redundant signals; warn on genuinely invalid ones. @@ -222,9 +271,9 @@ def handle_control_signal(coordinator, sender_identity, payload): if transition.new_state is not None: coordinator.state = transition.new_state - # Broadcast the control signal. Forward the full deserialized payload so - # that data-bearing signals (e.g. SET_GENERATION_EPOCH) retain their args. - coordinator._broadcast_to_engines(payload) + # Broadcast the control signal. Forward the full metadata so that + # data-bearing signals (e.g. SET_GENERATION_EPOCH) retain their args. + coordinator._broadcast_to_engines(metadata) # STOP affects engines; reset coordinator to RUNNING to allow future engines. if header == Headers.STOP: @@ -232,21 +281,37 @@ def handle_control_signal(coordinator, sender_identity, payload): @message_handler(Headers.START_CUDA_PROFILER, Headers.STOP_CUDA_PROFILER) -def handle_cuda_profiler_signal(coordinator, sender_identity, payload): +def handle_cuda_profiler_signal(coordinator, sender_identity, metadata, bodies): """Broadcast a CUDA profiler control signal to every connected DP engine. + Serves START_CUDA_PROFILER and STOP_CUDA_PROFILER, sent by + ``InferenceClient._send_signal_to_engines``. + Profiler control is not a coordinator state transition, so there are no CoordinatorState checks — the signal is simply forwarded to all engines. + + ``metadata``: ``[header]``, rebroadcast verbatim. + ``bodies``: empty. """ if sender_identity not in coordinator.known_clients: logging.warning("Coordinator: ignoring profiler signal from unknown client.") return - coordinator._broadcast_to_engines(payload) + coordinator._broadcast_to_engines(metadata) @message_handler(Headers.ENGINE_REPLY) -def handle_engine_reply(coordinator, sender_identity, payload): - """Route completed requests from an engine back to their originating clients.""" +def handle_engine_reply(coordinator, sender_identity, metadata, bodies): + """Route completed requests from an engine back to their originating clients. + + Sent by the engine via ``_engine_reply_frames``. + + ``metadata``: ``[header, [[request_id, needs_detokenize], ...]]`` -- one + entry per finished request, in the same order as ``bodies``. + ``bodies``: one frame per entry, each a packed finished request. A finished + request echoes the prompt back, so these are the largest frames the + coordinator handles. A body is decoded only when ``needs_detokenize`` is + set; otherwise it reaches the client as the frame the engine produced. + """ # This is the output of a single engine step on some data parallel rank. if sender_identity not in coordinator.identities_of_data_parallel_ranks: # A removed engine's final replies may still be queued up. @@ -255,11 +320,8 @@ def handle_engine_reply(coordinator, sender_identity, payload): sender_identity in coordinator.removed_engine_identities ), f"ENGINE_REPLY from never-connected sender {sender_identity!r}" logging.warning("Coordinator: ENGINE_REPLY from removed engine %r", sender_identity) - finished_requests = payload[1] - for finished_request in finished_requests: - coordinator.detokenize(finished_request) - fid = finished_request["request_id"] + for (fid, needs_detokenize), body in zip(metadata[1], bodies): client_identity = coordinator.request_id_to_client_id[fid] client_request_id = coordinator.request_id_to_client_request_id[fid] del coordinator.request_id_to_client_id[fid] @@ -272,49 +334,69 @@ def handle_engine_reply(coordinator, sender_identity, payload): assert coordinator._pending_counts[idx] >= 1 coordinator._pending_counts[idx] -= 1 - coordinator.router_socket.send_multipart( - [ - client_identity, - msgpack.packb( - [Headers.ENGINE_REPLY.value, client_request_id, finished_request], - use_bin_type=True, - ), - ] + if needs_detokenize: + # Detokenizing writes generated_text into the reply, so this one has + # to be decoded and re-encoded. Clients that detokenize for + # themselves (the OpenAI frontend does) never take this path. + finished_request = msgpack.unpackb(body, raw=False) + coordinator.detokenize(finished_request) + body = msgpack.packb(finished_request, use_bin_type=True) + + reply_metadata = msgpack.packb( + [Headers.ENGINE_REPLY.value, client_request_id], use_bin_type=True ) + coordinator.router_socket.send_multipart([client_identity, reply_metadata, body]) @message_handler(Headers.ENGINE_REPLY_PARTIAL) -def handle_engine_reply_partial(coordinator, sender_identity, payload): - """Route incremental engine replies without releasing request routing state.""" +def handle_engine_reply_partial(coordinator, sender_identity, metadata, bodies): + """Route incremental engine replies without releasing request routing state. + + Sent by the engine for streaming requests at each streaming interval. + + ``metadata``: ``[header, [request_id, ...]]`` -- one id per partial, in the + same order as ``bodies``. No detokenize flag: partials are always + detokenized incrementally by the client-facing streaming layer. + ``bodies``: one frame per id, each a packed partial + (``{"request_id": int, "new_tokens": [...]}``), always forwarded + untouched. + """ if sender_identity not in coordinator.identities_of_data_parallel_ranks: assert ( sender_identity in coordinator.removed_engine_identities ), f"ENGINE_REPLY_PARTIAL from never-connected sender {sender_identity!r}" logging.warning("Coordinator: ENGINE_REPLY_PARTIAL from removed engine %r", sender_identity) return - for partial in payload[1]: - request_id = partial["request_id"] + for request_id, body in zip(metadata[1], bodies): client_identity = coordinator.request_id_to_client_id[request_id] client_request_id = coordinator.request_id_to_client_request_id[request_id] - # Partial tokens are detokenized incrementally by the client-facing streaming layer. + # Partial tokens are detokenized incrementally by the client-facing + # streaming layer, so the body is always forwarded untouched. coordinator.router_socket.send_multipart( [ client_identity, msgpack.packb( - [Headers.ENGINE_REPLY_PARTIAL.value, client_request_id, partial], - use_bin_type=True, + [Headers.ENGINE_REPLY_PARTIAL.value, client_request_id], use_bin_type=True ), + body, ] ) @message_handler(Headers.ABORT_REQUEST) -def handle_abort_request(coordinator, sender_identity, payload): - """Forward a client cancellation to the engine serving that request.""" +def handle_abort_request(coordinator, sender_identity, metadata, bodies): + """Forward a client cancellation to the engine serving that request. + + Sent by ``InferenceClient.abort_request``. Unknown or already-completed + request ids are dropped silently. + + ``metadata``: ``[header, client_request_id]``. + ``bodies``: empty. + """ if sender_identity not in coordinator.known_clients: logging.warning("Coordinator: ignoring abort from unknown client.") return - client_request_id = int(payload[1]) + client_request_id = int(metadata[1]) request_id = coordinator.client_request_to_request_id.get((sender_identity, client_request_id)) if request_id is None: return @@ -322,13 +404,19 @@ def handle_abort_request(coordinator, sender_identity, payload): if assigned_rank is not None: coordinator._send_to_engine( assigned_rank, - msgpack.packb([Headers.ABORT_REQUEST.value, request_id], use_bin_type=True), + [msgpack.packb([Headers.ABORT_REQUEST.value, request_id], use_bin_type=True)], ) @message_handler(Headers.SHUTDOWN) -def handle_shutdown(coordinator, sender_identity, payload): - """Stop the coordinator event loop on request from a known client.""" +def handle_shutdown(coordinator, sender_identity, metadata, bodies): + """Stop the coordinator event loop on request from a known client. + + Sent by ``InferenceClient.shutdown_engines``. + + ``metadata``: ``[header]``. + ``bodies``: empty. + """ if sender_identity not in coordinator.known_clients: logging.warning("Coordinator: ignoring signal from unknown client.") return @@ -336,7 +424,14 @@ def handle_shutdown(coordinator, sender_identity, payload): @message_handler(Headers.DISCONNECT) -def handle_disconnect(coordinator, sender_identity, payload): - """Remove a disconnecting engine from the routing pool.""" +def handle_disconnect(coordinator, sender_identity, metadata, bodies): + """Remove a disconnecting engine from the routing pool. + + Sent by an engine as it exits -- the only handler here whose sender is an + engine rather than a client. + + ``metadata``: ``[header]``. + ``bodies``: empty. + """ if sender_identity in coordinator.identities_of_data_parallel_ranks: coordinator._remove_engine(sender_identity) diff --git a/megatron/core/inference/engines/dynamic_engine.py b/megatron/core/inference/engines/dynamic_engine.py index 61ade21cbb3..b406e878e56 100644 --- a/megatron/core/inference/engines/dynamic_engine.py +++ b/megatron/core/inference/engines/dynamic_engine.py @@ -152,6 +152,36 @@ def format_mem_bytes(mem_bytes): return "%d bytes" % mem_bytes +def _engine_reply_frames(finished_requests: List[dict]) -> List[bytes]: + """Frame finished requests as [metadata, body, body, ...] for the coordinator. + + The metadata frame carries only what the coordinator needs to route each + reply: the request id, and whether it must detokenize into the body. Every + body stays a separate opaque frame so the coordinator can forward it without + decoding -- a finished request echoes the prompt back, so decoding it costs + more than the inbound submission did. + + Args: + finished_requests: Serialized requests, in the order their frames follow. + + Returns: + The frames to send, metadata first. + """ + metadata = [ + Headers.ENGINE_REPLY.value, + [ + [ + request["request_id"], + bool((request.get("sampling_params") or {}).get("detokenize_generations")), + ] + for request in finished_requests + ], + ] + return [msgpack.packb(metadata, use_bin_type=True)] + [ + msgpack.packb(request, use_bin_type=True) for request in finished_requests + ] + + def _get_decode_only_log_state( mode: AsyncScheduleMode, decode_only: DecodeOnly ) -> Tuple[str, Optional[bool]]: @@ -1087,11 +1117,9 @@ def _send_request_records_to_coordinator( merged.uid not in self.local_metadata_ledger ), f"finished-request ledger: duplicate uid {merged.uid!r}" self.local_metadata_ledger[merged.uid] = FinishedRequestRecord.from_request(merged) - payload = msgpack.packb( - [Headers.ENGINE_REPLY.value, [request.serialize() for request in merged_requests]], - use_bin_type=True, + self.socket_for_receiving_requests.send_multipart( + _engine_reply_frames([request.serialize() for request in merged_requests]) ) - self.socket_for_receiving_requests.send(payload) def _handle_failed_request(self, request_id: int): """Handle a failed request by sending the reply immediately. @@ -2397,9 +2425,16 @@ def _try_send_streaming_partials(self) -> None: if not partials: return - payload = msgpack.packb([Headers.ENGINE_REPLY_PARTIAL.value, partials], use_bin_type=True) nvtx_range_push("coordinator_streaming") - self.socket_for_receiving_requests.send(payload) + self.socket_for_receiving_requests.send_multipart( + [ + msgpack.packb( + [Headers.ENGINE_REPLY_PARTIAL.value, [p["request_id"] for p in partials]], + use_bin_type=True, + ) + ] + + [msgpack.packb(p, use_bin_type=True) for p in partials] + ) nvtx_range_pop("coordinator_streaming") self._partial_emit_lengths.update(emit_lengths) @@ -2786,6 +2821,52 @@ def generate( return finished_request_records_list + @staticmethod + def _pack_tp_broadcast(messages: List[List[bytes]]) -> List[bytes]: + """Flatten per-message frame lists into one TP-broadcast multipart message. + + A message is a list of frames -- metadata first, then any payload bodies. + ZMQ multipart is flat, so the frame boundaries would be lost on the wire. + They are carried instead in a manifest frame holding one frame count per + message, which lets peer ranks rebuild the grouping without any payload + being copied, decoded, or re-packed. + + Args: + messages: One frame list per message, in delivery order. + + Returns: + ``[tp_broadcast_header, manifest, *flattened frames]``. + """ + manifest = msgpack.packb([len(message) for message in messages], use_bin_type=True) + return [bytes([Headers.TP_BROADCAST.value]), manifest] + [ + frame for message in messages for frame in message + ] + + @staticmethod + def _unpack_tp_broadcast(frames: List[bytes]) -> List[List[bytes]]: + """Rebuild per-message frame lists from a TP broadcast. + + Inverse of :meth:`_pack_tp_broadcast`. + + Args: + frames: The received multipart message, header frame first. + + Returns: + One frame list per message, in the order they were packed. + """ + frame_counts = msgpack.unpackb(frames[1], raw=False) + flat = frames[2:] + messages = [] + offset = 0 + for count in frame_counts: + messages.append(flat[offset : offset + count]) + offset += count + assert offset == len(flat), ( + f"TP broadcast manifest accounts for {offset} frames but {len(flat)} were received; " + "sender and receiver disagree on message framing" + ) + return messages + def schedule_requests(self) -> int: """Drains the ZMQ socket for a batch of requests and adds them to the engine. @@ -2820,25 +2901,32 @@ def schedule_requests(self) -> int: nvtx_range_push("drain_zmq_socket") all_messages = [] if self.is_mp_coordinator: + # Locally-generated notifications are single-frame messages, so they + # are wrapped to match the frame-list shape of socket traffic. all_messages.extend( - msgpack.packb( - [Headers.KV_HANDOFF_COMPLETE.value, request_id, failed], use_bin_type=True - ) + [ + msgpack.packb( + [Headers.KV_HANDOFF_COMPLETE.value, request_id, failed], use_bin_type=True + ) + ] for request_id, failed in self._drain_handoff_completion_notifications() ) while True: try: # Receive messages in a non-blocking way. - all_messages.append(self.socket_for_receiving_requests.recv(flags=zmq.NOBLOCK)) + all_messages.append( + self.socket_for_receiving_requests.recv_multipart(flags=zmq.NOBLOCK) + ) except zmq.Again: # This exception is hit as soon as the socket is empty. break self.model_parallel_publisher_socket.send_multipart( - [bytes([Headers.TP_BROADCAST.value])] + all_messages + self._pack_tp_broadcast(all_messages) ) else: - frames = self.model_parallel_subscriber_socket.recv_multipart() - all_messages = frames[1:] + all_messages = self._unpack_tp_broadcast( + self.model_parallel_subscriber_socket.recv_multipart() + ) nvtx_range_pop("drain_zmq_socket") @@ -2846,17 +2934,23 @@ def schedule_requests(self) -> int: # Control signals are queued for the second pass. new_generation_epoch = None for message in all_messages: - data = msgpack.unpackb(message, raw=False) + data = msgpack.unpackb(message[0], raw=False) header = Headers(data[0]) if header == Headers.SUBMIT_REQUEST: - request_id, prompt, sampling_params = data[1:] + request_id, sampling_params = data[1:] + # The prompt rides in its own frame; the engine is its first + # consumer, so this is where it finally gets decoded. + prompt = msgpack.unpackb(message[1], raw=False) sampling_params = SamplingParams.deserialize(sampling_params) nvtx_range_push("add_request") self.add_request(request_id, prompt, sampling_params) nvtx_range_pop("add_request") elif header == Headers.SUBMIT_REQUEST_WITH_KV: - # Decode-side KV import. - request_id, prompt, sampling_params, kv_meta, src_block_ids = data[1:] + # Decode-side KV import. As on the plain path, the prompt rides + # in its own frame and the engine is its first consumer. + request_id, sampling_params, kv_meta = data[1:] + prompt = msgpack.unpackb(message[1], raw=False) + src_block_ids = msgpack.unpackb(message[2], raw=False) sampling_params = SamplingParams.deserialize(sampling_params) nvtx_range_push("add_request_with_kv_handoff") self.add_request_with_kv_handoff( @@ -2925,7 +3019,7 @@ def schedule_requests(self) -> int: # processes one state transition per iteration). if self._pending_signals: message = self._pending_signals.popleft() - data = msgpack.unpackb(message, raw=False) + data = msgpack.unpackb(message[0], raw=False) header = Headers(data[0]) if header == Headers.PAUSE: diff --git a/megatron/core/inference/inference_client.py b/megatron/core/inference/inference_client.py index 199018ce3a6..91783ee0906 100644 --- a/megatron/core/inference/inference_client.py +++ b/megatron/core/inference/inference_client.py @@ -6,6 +6,8 @@ import time from typing import List, Optional, Union +import torch + from megatron.core.inference.async_stream import AsyncStream from megatron.core.inference.inference_request import DynamicInferenceRequest from megatron.core.inference.sampling_params import SamplingParams @@ -111,8 +113,32 @@ def add_request( """ request_id = self.next_request_id self.next_request_id += 1 - payload = [Headers.SUBMIT_REQUEST.value, request_id, prompt, sampling_params.serialize()] - return self._submit_request(payload, request_id) + frames = [ + msgpack.packb( + [Headers.SUBMIT_REQUEST.value, request_id, sampling_params.serialize()], + use_bin_type=True, + ), + # The prompt travels in its own frame so the coordinator can route the + # request without decoding it -- at long prompts that decode dominates + # its per-request cost, and it is one serial loop shared by every rank. + self._pack_prompt(prompt), + ] + return self._submit_request(frames, request_id) + + @staticmethod + def _pack_prompt(prompt): + """Pack a prompt into its own frame. + + Coercion happens here rather than at the coordinator: the coordinator no + longer decodes the prompt, and clients are many against its one serial + loop, so this is both the only place that still sees the object and the + cheaper place to normalize it. + """ + if isinstance(prompt, torch.Tensor): + prompt = prompt.tolist() + elif not isinstance(prompt, (str, list)): + raise TypeError(f"unsupported prompt type {type(prompt).__name__}") + return msgpack.packb(prompt, use_bin_type=True) def _make_kv_handoff_request( self, @@ -121,18 +147,33 @@ def _make_kv_handoff_request( kv_meta: dict, src_block_ids: List[int], ) -> tuple[int, list]: - """Allocate an ID and build a decode request carrying remote KV metadata.""" + """Allocate an ID and build a decode request carrying remote KV metadata. + + Framed as [metadata, prompt, src_block_ids]. + + Nothing whose size follows the sequence length belongs in the metadata + frame, since that frame is the only one the coordinator decodes. + ``src_block_ids`` names one block per block_size_tokens of prompt, so it + grows with the prompt and travels as its own body. ``kv_meta`` stays in + the metadata: it is the peer's NIXL agent/layout export, bounded by TP + size and num_speculative_tokens, not by prompt length. + """ request_id = self.next_request_id self.next_request_id += 1 - payload = [ - Headers.SUBMIT_REQUEST_WITH_KV.value, - request_id, - prompt, - sampling_params.serialize(), - kv_meta, - list(src_block_ids), + frames = [ + msgpack.packb( + [ + Headers.SUBMIT_REQUEST_WITH_KV.value, + request_id, + sampling_params.serialize(), + kv_meta, + ], + use_bin_type=True, + ), + self._pack_prompt(prompt), + msgpack.packb(list(src_block_ids), use_bin_type=True), ] - return request_id, payload + return request_id, frames def add_request_with_kv_handoff( self, @@ -155,10 +196,10 @@ def add_request_with_kv_handoff( Returns: asyncio.Future: A future that resolves to the completed request. """ - request_id, payload = self._make_kv_handoff_request( + request_id, frames = self._make_kv_handoff_request( prompt, sampling_params, kv_meta, src_block_ids ) - return self._submit_request(payload, request_id) + return self._submit_request(frames, request_id) def add_request_with_kv_handoff_streaming( self, @@ -182,10 +223,10 @@ def add_request_with_kv_handoff_streaming( AsyncStream[dict]: Per-step partial and final reply frames. """ sampling_params.streaming = True - request_id, payload = self._make_kv_handoff_request( + request_id, frames = self._make_kv_handoff_request( prompt, sampling_params, kv_meta, src_block_ids ) - return self._submit_stream(payload, request_id) + return self._submit_stream(frames, request_id) def release_handoff(self, request_id: int) -> None: """Tell the coordinator to release the KV blocks pinned for `request_id`. @@ -238,21 +279,30 @@ def add_request_streaming( sampling_params.streaming = True request_id = self.next_request_id self.next_request_id += 1 - payload = [Headers.SUBMIT_REQUEST.value, request_id, prompt, sampling_params.serialize()] - return self._submit_stream(payload, request_id) + frames = [ + msgpack.packb( + [Headers.SUBMIT_REQUEST.value, request_id, sampling_params.serialize()], + use_bin_type=True, + ), + # The prompt travels in its own frame so the coordinator can route the + # request without decoding it -- at long prompts that decode dominates + # its per-request cost, and it is one serial loop shared by every rank. + self._pack_prompt(prompt), + ] + return self._submit_stream(frames, request_id) - def _submit_request(self, payload: list, request_id: int) -> asyncio.Future: + def _submit_request(self, frames: list, request_id: int) -> asyncio.Future: """Send a prepared request and register its completion future.""" - self.socket.send(msgpack.packb(payload, use_bin_type=True)) + self.socket.send_multipart(frames) assert request_id not in self.completion_futures future = asyncio.get_running_loop().create_future() self.completion_futures[request_id] = future self.request_submission_times[request_id] = time.perf_counter() return future - def _submit_stream(self, payload: list, request_id: int) -> AsyncStream[dict]: + def _submit_stream(self, frames: list, request_id: int) -> AsyncStream[dict]: """Send a prepared streaming request and register its response stream.""" - self.socket.send(msgpack.packb(payload, use_bin_type=True)) + self.socket.send_multipart(frames) stream = AsyncStream( request_id, functools.partial(self.abort_request, request_id), loop=self._loop ) @@ -275,13 +325,16 @@ async def _recv_task(self): """ while True: try: - data = msgpack.unpackb(self.socket.recv(flags=zmq.NOBLOCK), raw=False) + # frames[0] is metadata; a reply body, when present, follows it. + frames = self.socket.recv_multipart(flags=zmq.NOBLOCK) + data = msgpack.unpackb(frames[0], raw=False) header = Headers(data[0]) if header == Headers.ENGINE_REPLY: - request_id, reply = data[1:] + request_id = data[1] if request_id in self.aborted_request_ids: self.aborted_request_ids.discard(request_id) continue + reply = msgpack.unpackb(frames[1], raw=False) submitted = self.request_submission_times.pop(request_id, None) if submitted is not None: reply['latency'] = time.perf_counter() - submitted @@ -305,10 +358,10 @@ async def _recv_task(self): ) completion_future.set_result(completed_request) elif header == Headers.ENGINE_REPLY_PARTIAL: - request_id, partial = data[1:] + request_id = data[1] stream = self.streams.get(request_id) if stream is not None: - stream.put({"partial": partial}) + stream.put({"partial": msgpack.unpackb(frames[1], raw=False)}) except zmq.Again: await asyncio.sleep(0.005) continue @@ -328,7 +381,7 @@ def _connect_with_inference_coordinator(self, timeout_seconds: Optional[float] = timeout=max(0, int(timeout_seconds * 1000)) ): raise TimeoutError("Timed out connecting to the Megatron inference coordinator") - reply = msgpack.unpackb(self.socket.recv(), raw=False) + reply = msgpack.unpackb(self.socket.recv_multipart()[0], raw=False) assert Headers(reply[0]) == Headers.CONNECT_ACK def start( diff --git a/tests/unit_tests/inference/engines/test_dynamic_engine.py b/tests/unit_tests/inference/engines/test_dynamic_engine.py index 093e4062cc8..24b6e3ff6b9 100644 --- a/tests/unit_tests/inference/engines/test_dynamic_engine.py +++ b/tests/unit_tests/inference/engines/test_dynamic_engine.py @@ -35,6 +35,7 @@ ) from megatron.core.inference.engines import DynamicInferenceEngine from megatron.core.inference.engines.dynamic_engine import EngineState +from megatron.core.inference.headers import Headers from megatron.core.inference.inference_request import ( DynamicInferenceRequest, DynamicInferenceRequestRecord, @@ -772,7 +773,12 @@ def test_streaming_partials_are_sent(): engine._try_send_streaming_partials() - engine.socket_for_receiving_requests.send.assert_called_once() + # Partials go out as [metadata, body] frames: the metadata names the request + # ids so the coordinator can route without decoding the bodies. + engine.socket_for_receiving_requests.send_multipart.assert_called_once() + frames = engine.socket_for_receiving_requests.send_multipart.call_args.args[0] + assert msgpack.unpackb(frames[0], raw=False) == [Headers.ENGINE_REPLY_PARTIAL.value, [7]] + assert msgpack.unpackb(frames[1], raw=False)["new_tokens"] == [11, 12, 13] assert engine._partial_emit_lengths == {7: 3} payload = msgpack.unpackb( engine.socket_for_receiving_requests.send.call_args.args[0], raw=False @@ -798,13 +804,13 @@ def test_streaming_partials_buffer_until_token_interval(): engine._try_send_streaming_partials() - engine.socket_for_receiving_requests.send.assert_not_called() + engine.socket_for_receiving_requests.send_multipart.assert_not_called() assert engine._partial_emit_lengths == {} request.generated_tokens.append(13) engine._try_send_streaming_partials() - engine.socket_for_receiving_requests.send.assert_called_once() + engine.socket_for_receiving_requests.send_multipart.assert_called_once() assert engine._partial_emit_lengths == {7: 3} diff --git a/tests/unit_tests/inference/test_data_parallel_inference_coordinator.py b/tests/unit_tests/inference/test_data_parallel_inference_coordinator.py index 55c0f50a83d..326a780cccf 100644 --- a/tests/unit_tests/inference/test_data_parallel_inference_coordinator.py +++ b/tests/unit_tests/inference/test_data_parallel_inference_coordinator.py @@ -28,6 +28,7 @@ DynamicInferenceEngine, EngineState, RequestEntry, + _engine_reply_frames, ) from megatron.core.inference.headers import Headers from megatron.core.inference.inference_client import InferenceClient @@ -58,6 +59,115 @@ def test_coordinator_registers_client_kv_handoff_handlers(): assert Headers.RELEASE_KV in HANDLERS +class _StubCoordinator: + """Minimal stand-in exposing only what the submit handlers touch.""" + + def __init__(self, identity=b"engine-0"): + self.known_clients = {b"client-0"} + self.next_request_id = 100 + self.request_id_to_client_id = {} + self.request_id_to_client_request_id = {} + self.client_request_to_request_id = {} + self.request_id_to_rank = {} + self.identities_of_data_parallel_ranks = [identity] + self.identity_to_rank_index = {identity: 0} + self._pending_counts = np.zeros(1, dtype=np.int64) + self._identity = identity + self.sent = [] + + def get_least_loaded_data_parallel_rank(self): + return self._identity + + def _send_to_engine(self, identity, frames): + # Mirrors the real signature: a list of frames, metadata first. Passing + # raw bytes here would splay one frame per byte, so assert the shape. + assert isinstance(frames, list), f"expected a frame list, got {type(frames).__name__}" + self.sent.append((identity, frames)) + return True + + +def test_kv_handoff_round_trip_keeps_prompt_in_its_own_frame(): + """Client -> coordinator -> engine for a KV handoff, asserting the framing. + + The prompt must never be decoded by the coordinator: it is forwarded as the + opaque body frame the client packed, byte for byte. + """ + prompt_tokens = [11, 22, 33, 44] + kv_meta = {"agent": "nixl-0"} + src_block_ids = [7, 8] + + # --- client side: build the frames without needing a live socket --- + client = InferenceClient.__new__(InferenceClient) + client.next_request_id = 5 + request_id, frames = InferenceClient._make_kv_handoff_request( + client, prompt_tokens, SamplingParams(num_tokens_to_generate=4), kv_meta, src_block_ids + ) + assert request_id == 5 + assert len(frames) == 3, "KV handoff must be framed as [metadata, prompt, src_block_ids]" + + metadata = msgpack.unpackb(frames[0], raw=False) + assert metadata[0] == Headers.SUBMIT_REQUEST_WITH_KV.value + assert len(metadata) == 4, "only constant-size fields belong in the metadata frame" + assert metadata[1] == request_id + assert metadata[3] == kv_meta + assert msgpack.unpackb(frames[1], raw=False) == prompt_tokens + assert msgpack.unpackb(frames[2], raw=False) == src_block_ids + + # --- coordinator side: route it --- + coordinator = _StubCoordinator() + handler = HANDLERS[Headers.SUBMIT_REQUEST_WITH_KV] + handler(coordinator, b"client-0", metadata, frames[1:]) + + assert len(coordinator.sent) == 1 + identity, out_frames = coordinator.sent[0] + assert identity == b"engine-0" + assert len(out_frames) == 3 + # Both body frames are forwarded untouched -- not re-packed. + assert out_frames[1] is frames[1] + assert out_frames[2] is frames[2] + + engine_metadata = msgpack.unpackb(out_frames[0], raw=False) + assert engine_metadata[0] == Headers.SUBMIT_REQUEST_WITH_KV.value + server_request_id = engine_metadata[1] + assert coordinator.request_id_to_client_request_id[server_request_id] == request_id + assert coordinator.request_id_to_rank[server_request_id] == b"engine-0" + assert coordinator._pending_counts[0] == 1 + + # --- engine side: the bodies are decoded here, for the first time --- + assert msgpack.unpackb(out_frames[1], raw=False) == prompt_tokens + assert msgpack.unpackb(out_frames[2], raw=False) == src_block_ids + + +def test_kv_handoff_metadata_frame_does_not_grow_with_sequence_length(): + """The metadata frame is the only one the coordinator decodes, so its size + must not follow the prompt. Build the same request at two prompt lengths and + require the metadata frame to be byte-identical in length.""" + client = InferenceClient.__new__(InferenceClient) + sizes = [] + for n_blocks in (1, 64): + client.next_request_id = 5 + _, frames = InferenceClient._make_kv_handoff_request( + client, + list(range(n_blocks * 64)), + SamplingParams(num_tokens_to_generate=4), + {"agent": "nixl-0"}, + list(range(n_blocks)), + ) + sizes.append(len(frames[0])) + assert sizes[0] == sizes[1], ( + f"metadata frame grew with prompt length ({sizes[0]} -> {sizes[1]} bytes); " + "something sequence-dependent leaked into it" + ) + + +def test_kv_handoff_rejects_legacy_single_frame_payload(): + """A pre-split single-frame payload must be rejected, not silently mis-read.""" + coordinator = _StubCoordinator() + legacy = [Headers.SUBMIT_REQUEST_WITH_KV.value, 1, [1, 2], {}, {}, []] + HANDLERS[Headers.SUBMIT_REQUEST_WITH_KV](coordinator, b"client-0", legacy, []) + assert coordinator.sent == [] + + class DummyTokenizer: """Dummy tokenizer.""" @@ -198,13 +308,11 @@ async def async_step(self, *, verbose: Optional[bool] = False) -> Dict: finished_request_records.append(entry.record) entry.future.set_result(entry.record) to_remove.append(request_id) - # Send signal to coordinator. + # Send signal to coordinator, framed the way a real engine does. if self.is_mp_coordinator: - payload = msgpack.packb( - [Headers.ENGINE_REPLY.value, [entry.record.merge().serialize()]], - use_bin_type=True, + self.socket_for_receiving_requests.send_multipart( + _engine_reply_frames([entry.record.merge().serialize()]) ) - self.socket_for_receiving_requests.send(payload) for request_id in to_remove: del self.requests[request_id] @@ -877,10 +985,19 @@ def test_reply_routing_survives_engine_removal(self, caplog): """A removed engine's queued replies still deliver; never-connected senders assert.""" def reply(fid): - return [ - Headers.ENGINE_REPLY.value, - [{"request_id": fid, "generated_tokens": [1], "sampling_params": {}}], + """An ENGINE_REPLY as (metadata, body frames), matching the wire format. + + The metadata frame carries only the routing key and the detokenize + flag; the finished request travels as an opaque body frame. + """ + metadata = [Headers.ENGINE_REPLY.value, [[fid, False]]] + bodies = [ + msgpack.packb( + {"request_id": fid, "generated_tokens": [1], "sampling_params": {}}, + use_bin_type=True, + ) ] + return metadata, bodies coord = _make_routing_coordinator(num_ranks=2) coord.tokenizer = DummyTokenizer() @@ -892,7 +1009,7 @@ def reply(fid): # A sender that never registered is a protocol violation. with pytest.raises(AssertionError, match="never-connected"): - handle_engine_reply(coord, b"impostor", reply(11)) + handle_engine_reply(coord, b"impostor", *reply(11)) assert coord.router_socket.send_multipart.call_count == 0 assert 11 in coord.request_id_to_client_id @@ -900,7 +1017,7 @@ def reply(fid): # reply can still arrive - and must reach its client. coord._remove_engine(b"rank-0") with caplog.at_level(logging.WARNING): - handle_engine_reply(coord, b"rank-0", reply(11)) + handle_engine_reply(coord, b"rank-0", *reply(11)) assert "removed engine" in caplog.text assert coord.router_socket.send_multipart.call_args[0][0][0] == b"client-A" assert 11 not in coord.request_id_to_client_id diff --git a/tests/unit_tests/inference/test_dynamic_prefix_caching_coordinator.py b/tests/unit_tests/inference/test_dynamic_prefix_caching_coordinator.py index 1d020ee1a79..28aa82d65b8 100644 --- a/tests/unit_tests/inference/test_dynamic_prefix_caching_coordinator.py +++ b/tests/unit_tests/inference/test_dynamic_prefix_caching_coordinator.py @@ -22,7 +22,14 @@ from megatron.core.inference.data_parallel_inference_coordinator import ( DataParallelInferenceCoordinator, ) -from megatron.core.inference.engines.dynamic_engine import DynamicInferenceEngine, RequestEntry +from megatron.core.inference.data_parallel_inference_coordinator.handlers import ( + handle_submit_request, +) +from megatron.core.inference.engines.dynamic_engine import ( + DynamicInferenceEngine, + RequestEntry, + _engine_reply_frames, +) from megatron.core.inference.headers import Headers from megatron.core.inference.inference_client import InferenceClient from megatron.core.inference.inference_request import ( @@ -154,10 +161,9 @@ async def async_step(self, *, verbose: Optional[bool] = False) -> Dict: entry.future.set_result(entry.record) to_remove.append(request_id) if self.is_mp_coordinator: - payload = msgpack.packb( - [Headers.ENGINE_REPLY.value, [entry.record.serialize()]], use_bin_type=True + self.socket_for_receiving_requests.send_multipart( + _engine_reply_frames([entry.record.serialize()]) ) - self.socket_for_receiving_requests.send(payload) for request_id in to_remove: del self.requests[request_id] @@ -249,18 +255,6 @@ def test_hash_from_string_prompt(self): hashes_from_list = coordinator.compute_request_hashes([1, 2, 3, 4, 5, 6, 7, 8]) assert hashes_from_str == hashes_from_list - def test_hash_empty_when_disabled(self): - """Returns empty list when prefix caching is disabled.""" - coordinator = make_coordinator_direct(enable_prefix_caching=False) - hashes = coordinator.compute_request_hashes([1, 2, 3, 4]) - assert hashes == [] - - def test_hash_empty_when_no_block_size(self): - """Returns empty list when block_size_tokens is None.""" - coordinator = make_coordinator_direct(block_size_tokens=None) - hashes = coordinator.compute_request_hashes([1, 2, 3, 4]) - assert hashes == [] - def test_hash_partial_block_ignored(self): """Tokens that don't fill a complete block produce no hash.""" coordinator = make_coordinator_direct() @@ -288,6 +282,62 @@ def test_hash_parent_chaining(self): assert h1[1] != h2[1] +class TestSubmitDoesNotDecodePrompt: + """The coordinator must not decode the prompt when no routing needs it. + + Decoding is O(prompt length) on the coordinator's single serial loop, so + skipping it is the point of carrying the prompt in its own frame. Each test + sends a prompt frame that is deliberately *not* valid msgpack: if the + handler tried to decode it the call would raise, so completing cleanly is + proof the bytes were forwarded untouched. + """ + + UNDECODABLE_PROMPT = b"\xc1not-valid-msgpack" + + def _submit(self, coordinator): + """Drive handle_submit_request once and return the frames sent onward.""" + coordinator.known_clients = {b"client-A"} + coordinator.next_request_id = 0 + coordinator.request_id_to_client_id = {} + coordinator.request_id_to_client_request_id = {} + coordinator.client_request_to_request_id = {} + coordinator.request_id_to_rank = {} + coordinator.schedule_records = None + coordinator.router_socket = MagicMock() + + metadata = [Headers.SUBMIT_REQUEST.value, 7, {"temperature": 1.0}] + handle_submit_request(coordinator, b"client-A", metadata, [self.UNDECODABLE_PROMPT]) + return coordinator.router_socket.send_multipart.call_args.args[0] + + def test_load_balanced_forwards_prompt_verbatim(self): + """LOAD_BALANCED ignores hashes, so the prompt is never decoded.""" + coordinator = make_coordinator_direct(data_parallel_size=2) + coordinator.prefix_caching_coordinator_policy = ( + PrefixCachingCoordinatorPolicy.LOAD_BALANCED + ) + _identity, _metadata, prompt_frame = self._submit(coordinator) + assert prompt_frame is self.UNDECODABLE_PROMPT + + def test_disabled_prefix_caching_forwards_prompt_verbatim(self): + """With prefix caching off there are no hashes to compute either.""" + coordinator = make_coordinator_direct(data_parallel_size=2, enable_prefix_caching=False) + _identity, _metadata, prompt_frame = self._submit(coordinator) + assert prompt_frame is self.UNDECODABLE_PROMPT + + def test_metadata_frame_is_rewritten_with_server_request_id(self): + """The client's request id is swapped for the coordinator's own.""" + coordinator = make_coordinator_direct(data_parallel_size=2) + coordinator.prefix_caching_coordinator_policy = ( + PrefixCachingCoordinatorPolicy.LOAD_BALANCED + ) + _identity, metadata_frame, _prompt = self._submit(coordinator) + header, request_id, sampling_params = msgpack.unpackb(metadata_frame, raw=False) + assert header == Headers.SUBMIT_REQUEST.value + assert request_id == 0 # server-side id, not the client's 7 + assert sampling_params == {"temperature": 1.0} + assert coordinator.request_id_to_client_request_id[0] == 7 + + class TestCoordinatorPrefixRouting: """Test routing decisions based on prefix cache affinity.""" diff --git a/tests/unit_tests/inference/test_inference_client.py b/tests/unit_tests/inference/test_inference_client.py index 71e84d2c16d..d94f92e6004 100644 --- a/tests/unit_tests/inference/test_inference_client.py +++ b/tests/unit_tests/inference/test_inference_client.py @@ -47,12 +47,17 @@ async def test_inference_client_lifecycle(): assert client.completion_futures == {} # start(): handshake sends CONNECT, expects CONNECT_ACK, spawns listener task. - # We stage two recv() replies: the CONNECT_ACK during handshake, and an - # ENGINE_REPLY for the request we'll add below. Subsequent recvs raise - # zmq.Again so the listener loop yields back to the event loop. + # We stage two recv_multipart() replies: the CONNECT_ACK during handshake, + # and an ENGINE_REPLY for the request we'll add below. Messages arrive as + # [metadata, body] frames; the body is only present on replies that carry + # one. Subsequent recvs raise zmq.Again so the listener yields back to the + # event loop. recv_queue = [ - msgpack.packb([Headers.CONNECT_ACK.value], use_bin_type=True), - msgpack.packb([Headers.ENGINE_REPLY.value, 0, {"foo": "bar"}], use_bin_type=True), + [msgpack.packb([Headers.CONNECT_ACK.value], use_bin_type=True)], + [ + msgpack.packb([Headers.ENGINE_REPLY.value, 0], use_bin_type=True), + msgpack.packb({"foo": "bar"}, use_bin_type=True), + ], ] def fake_recv(*args, **kwargs): @@ -60,23 +65,25 @@ def fake_recv(*args, **kwargs): return recv_queue.pop(0) raise zmq.Again() - fake_socket.recv.side_effect = fake_recv + fake_socket.recv_multipart.side_effect = fake_recv client.start() assert isinstance(client.listener_task, asyncio.Task) sent_connect = fake_socket.send.call_args.args[0] assert msgpack.unpackb(sent_connect, raw=False)[0] == Headers.CONNECT.value - # add_request: SUBMIT_REQUEST payload (header, id, prompt, sampling-dict), counter increments. + # add_request frames the submission as [metadata, prompt] so the coordinator + # can route it without decoding the prompt. fut = client.add_request("hello", SamplingParams(temperature=0.5)) assert isinstance(fut, asyncio.Future) assert client.next_request_id == 1 assert 0 in client.request_submission_times - submit_payload = msgpack.unpackb(fake_socket.send.call_args.args[0], raw=False) + submit_meta, submit_prompt = fake_socket.send_multipart.call_args.args[0] + submit_payload = msgpack.unpackb(submit_meta, raw=False) assert submit_payload[0] == Headers.SUBMIT_REQUEST.value assert submit_payload[1] == 0 - assert submit_payload[2] == "hello" - assert submit_payload[3]["temperature"] == 0.5 + assert submit_payload[2]["temperature"] == 0.5 + assert msgpack.unpackb(submit_prompt, raw=False) == "hello" # Listener delivers the reply: future resolves with payload + injected latency. # Submission-time entry is popped on completion. @@ -115,7 +122,9 @@ async def test_inference_client_connect_handshake_rejects_unexpected_reply(): fatal protocol mismatch, not a recoverable error. Separated from the lifecycle test because it short-circuits before any state is established.""" client, _, fake_socket = _make_client() - fake_socket.recv.return_value = msgpack.packb([Headers.STOP.value], use_bin_type=True) + fake_socket.recv_multipart.return_value = [ + msgpack.packb([Headers.STOP.value], use_bin_type=True) + ] with pytest.raises(AssertionError): client._connect_with_inference_coordinator() diff --git a/tests/unit_tests/inference/test_inference_client_streaming.py b/tests/unit_tests/inference/test_inference_client_streaming.py index 29fdcc7497c..7615ae2918f 100644 --- a/tests/unit_tests/inference/test_inference_client_streaming.py +++ b/tests/unit_tests/inference/test_inference_client_streaming.py @@ -30,28 +30,26 @@ async def test_add_request_streaming_emits_partials_then_final(): """Two ENGINE_REPLY_PARTIAL frames followed by an ENGINE_REPLY terminate the iterator.""" client, fake_socket = _make_client() + # Partials and finals both arrive as [metadata, body] frames. recv_queue = [ - msgpack.packb([Headers.CONNECT_ACK.value], use_bin_type=True), - msgpack.packb( - [ - Headers.ENGINE_REPLY_PARTIAL.value, - 0, + [msgpack.packb([Headers.CONNECT_ACK.value], use_bin_type=True)], + [ + msgpack.packb([Headers.ENGINE_REPLY_PARTIAL.value, 0], use_bin_type=True), + msgpack.packb( {"request_id": 0, "new_tokens": [1, 2], "new_log_probs": [-0.1, -0.2]}, - ], - use_bin_type=True, - ), - msgpack.packb( - [ - Headers.ENGINE_REPLY_PARTIAL.value, - 0, - {"request_id": 0, "new_tokens": [3], "new_log_probs": [-0.3]}, - ], - use_bin_type=True, - ), - msgpack.packb( - [Headers.ENGINE_REPLY.value, 0, {"request_id": 0, "generated_tokens": [1, 2, 3]}], - use_bin_type=True, - ), + use_bin_type=True, + ), + ], + [ + msgpack.packb([Headers.ENGINE_REPLY_PARTIAL.value, 0], use_bin_type=True), + msgpack.packb( + {"request_id": 0, "new_tokens": [3], "new_log_probs": [-0.3]}, use_bin_type=True + ), + ], + [ + msgpack.packb([Headers.ENGINE_REPLY.value, 0], use_bin_type=True), + msgpack.packb({"request_id": 0, "generated_tokens": [1, 2, 3]}, use_bin_type=True), + ], ] def fake_recv(*args, **kwargs): @@ -59,7 +57,7 @@ def fake_recv(*args, **kwargs): return recv_queue.pop(0) raise zmq.Again() - fake_socket.recv.side_effect = fake_recv + fake_socket.recv_multipart.side_effect = fake_recv client.start() params = SamplingParams(temperature=0.7, return_log_probs=True) @@ -70,9 +68,10 @@ def fake_recv(*args, **kwargs): assert isinstance(iterator, AsyncStream) assert params.streaming is True assert 0 in client.streams - submit_payload = msgpack.unpackb(fake_socket.send.call_args.args[0], raw=False) + submit_meta, _ = fake_socket.send_multipart.call_args.args[0] + submit_payload = msgpack.unpackb(submit_meta, raw=False) assert submit_payload[0] == Headers.SUBMIT_REQUEST.value - assert submit_payload[3]["streaming"] is True + assert submit_payload[2]["streaming"] is True items = [] async for item in iterator: @@ -97,11 +96,11 @@ async def test_streaming_partial_for_unknown_request_is_dropped(): client, fake_socket = _make_client() recv_queue = [ - msgpack.packb([Headers.CONNECT_ACK.value], use_bin_type=True), - msgpack.packb( - [Headers.ENGINE_REPLY_PARTIAL.value, 42, {"request_id": 42, "new_tokens": [9]}], - use_bin_type=True, - ), + [msgpack.packb([Headers.CONNECT_ACK.value], use_bin_type=True)], + [ + msgpack.packb([Headers.ENGINE_REPLY_PARTIAL.value, 42], use_bin_type=True), + msgpack.packb({"request_id": 42, "new_tokens": [9]}, use_bin_type=True), + ], ] def fake_recv(*args, **kwargs): @@ -109,7 +108,7 @@ def fake_recv(*args, **kwargs): return recv_queue.pop(0) raise zmq.Again() - fake_socket.recv.side_effect = fake_recv + fake_socket.recv_multipart.side_effect = fake_recv client.start() await asyncio.sleep(0.02) @@ -122,14 +121,14 @@ async def test_client_stop_terminates_open_streams(): """stop() finishes open streams so awaiters can exit.""" client, fake_socket = _make_client() - recv_queue = [msgpack.packb([Headers.CONNECT_ACK.value], use_bin_type=True)] + recv_queue = [[msgpack.packb([Headers.CONNECT_ACK.value], use_bin_type=True)]] def fake_recv(*args, **kwargs): if recv_queue: return recv_queue.pop(0) raise zmq.Again() - fake_socket.recv.side_effect = fake_recv + fake_socket.recv_multipart.side_effect = fake_recv client.start() iterator = client.add_request_streaming("hi", SamplingParams()) From fe4e6182224e8d66d2fd55d1c8b8db47a566258f Mon Sep 17 00:00:00 2001 From: Siddharth Singh Date: Wed, 19 Aug 2026 16:36:14 -0700 Subject: [PATCH 2/2] refactor(inference): declare the coordinator wire schema in one place Addresses review feedback on #6497: the frame lists were passed around as arbitrary packed bytes, indexed positionally at each site, with no semantic meaning attached to any element. A message's layout previously existed only as an unwritten agreement between a packb site and an unpack site in another module -- `[header, request_id, sampling_params]` matched by `_, request_id, sampling_params = metadata` three files away. Nothing tied the two together, which is how a 5-field metadata frame came to be read by a handler expecting 6. Each message type is now declared once as a MessageSpec naming its metadata fields and payload frames in wire order. pack and parse are both derived from that declaration, so no frame index is written by hand and the two directions cannot drift. Adding a field means editing one tuple. The split between the two tuples is the contract that matters, and it is now structural rather than prose: metadata_fields are decoded by the coordinator and must be constant-size, payload_frames are forwarded untouched and may grow with the prompt. Putting src_block_ids in one rather than the other is now a visible review-time decision. Batched engine replies keep a dedicated BatchedMessageSpec. They carry one payload frame per finished request because the batch fans out to different clients: collapsing it into a single frame would force the coordinator to decode every reply just to find the boundaries, and a finished request echoes the prompt back. Measured on a 50-reply batch with 2000-token prompts, that split costs the coordinator 0.001 ms against 3.7 ms for a decode. pack now asserts that a batch's entry count matches its frame count, so a mismatch fails at the send rather than silently truncating the batch. No wire-format change: every layout was verified to produce byte-identical frames to the previous code. Signed-off-by: Siddharth Singh --- .../handlers.py | 74 +++--- .../core/inference/engines/dynamic_engine.py | 89 +++---- megatron/core/inference/inference_client.py | 94 ++++---- megatron/core/inference/messages.py | 224 ++++++++++++++++++ 4 files changed, 358 insertions(+), 123 deletions(-) create mode 100644 megatron/core/inference/messages.py diff --git a/megatron/core/inference/data_parallel_inference_coordinator/handlers.py b/megatron/core/inference/data_parallel_inference_coordinator/handlers.py index c1d57ebfa54..b2700b80125 100644 --- a/megatron/core/inference/data_parallel_inference_coordinator/handlers.py +++ b/megatron/core/inference/data_parallel_inference_coordinator/handlers.py @@ -27,6 +27,17 @@ from megatron.core.inference.config import PrefixCachingCoordinatorPolicy from megatron.core.inference.headers import Headers +from megatron.core.inference.messages import ( + CLIENT_REPLY, + CLIENT_REPLY_PARTIAL, + ENGINE_REPLY, + ENGINE_REPLY_PARTIAL, + SUBMIT_REQUEST, + SUBMIT_REQUEST_WITH_KV, + header_of, + pack_signal, + request_id_of, +) from .state import CONTROL_TRANSITIONS, CoordinatorState @@ -73,9 +84,7 @@ def handle_connect(coordinator, sender_identity, metadata, bodies): return coordinator.known_clients.add(sender_identity) - coordinator.router_socket.send_multipart( - [sender_identity, msgpack.packb([Headers.CONNECT_ACK.value], use_bin_type=True)] - ) + coordinator.router_socket.send_multipart([sender_identity, *pack_signal(Headers.CONNECT_ACK)]) @message_handler(Headers.SUBMIT_REQUEST) @@ -97,8 +106,8 @@ def handle_submit_request(coordinator, sender_identity, metadata, bodies): logging.info(f"Received message from unknown client {sender_identity}. Ignoring.") return - _, client_request_id, sampling_params = metadata - prompt_frame = bodies[0] + request = SUBMIT_REQUEST.parse(metadata, bodies) + client_request_id = request.request_id # map client request_id to server request_id # necessary because multiple clients might have the same request_id. @@ -109,8 +118,8 @@ def handle_submit_request(coordinator, sender_identity, metadata, bodies): coordinator.client_request_to_request_id[(sender_identity, client_request_id)] = request_id # Rebuilding the metadata frame is cheap: it holds no prompt tokens. - engine_metadata = msgpack.packb( - [Headers.SUBMIT_REQUEST.value, request_id, sampling_params], use_bin_type=True + engine_frames = SUBMIT_REQUEST.pack( + request_id=request_id, sampling_params=request.sampling_params, prompt=request.prompt ) # Only prefix-affinity routing consults the block hashes. When nothing will @@ -123,7 +132,7 @@ def handle_submit_request(coordinator, sender_identity, metadata, bodies): != PrefixCachingCoordinatorPolicy.LOAD_BALANCED ): request_hashes = coordinator.compute_request_hashes( - msgpack.unpackb(prompt_frame, raw=False) + msgpack.unpackb(request.prompt, raw=False) ) else: request_hashes = [] @@ -137,7 +146,7 @@ def handle_submit_request(coordinator, sender_identity, metadata, bodies): # Account for the fact that some engines may have died. for _ in range(len(coordinator.identities_of_data_parallel_ranks)): next_identity = coordinator.get_best_data_parallel_rank(request_hashes) - if coordinator._send_to_engine(next_identity, [engine_metadata, prompt_frame]): + if coordinator._send_to_engine(next_identity, engine_frames): break else: # If all engines have died, we are in an abnormal state, and must exit cleanly. @@ -193,7 +202,8 @@ def handle_submit_request_with_kv(coordinator, sender_identity, metadata, bodies ) return - _, client_request_id, sampling_params, kv_meta = metadata + request = SUBMIT_REQUEST_WITH_KV.parse(metadata, bodies) + client_request_id = request.request_id request_id = coordinator.next_request_id coordinator.next_request_id += 1 coordinator.request_id_to_client_id[request_id] = sender_identity @@ -201,14 +211,17 @@ def handle_submit_request_with_kv(coordinator, sender_identity, metadata, bodies coordinator.client_request_to_request_id[(sender_identity, client_request_id)] = request_id # Rebuilding the metadata frame is cheap: it holds no prompt tokens. - engine_metadata = msgpack.packb( - [Headers.SUBMIT_REQUEST_WITH_KV.value, request_id, sampling_params, kv_meta], - use_bin_type=True, + engine_frames = SUBMIT_REQUEST_WITH_KV.pack( + request_id=request_id, + sampling_params=request.sampling_params, + kv_meta=request.kv_meta, + prompt=request.prompt, + src_block_ids=request.src_block_ids, ) for _ in range(len(coordinator.identities_of_data_parallel_ranks)): next_identity = coordinator.get_least_loaded_data_parallel_rank() - if coordinator._send_to_engine(next_identity, [engine_metadata, *bodies]): + if coordinator._send_to_engine(next_identity, engine_frames): break else: logging.error("Coordinator: no reachable engines for handoff request %d", request_id) @@ -235,7 +248,7 @@ def handle_release_kv(coordinator, sender_identity, metadata, bodies): if sender_identity not in coordinator.known_clients: logging.warning("Coordinator: ignoring RELEASE_KV from unknown client.") return - coordinator._broadcast_to_engines([Headers.RELEASE_KV.value, int(metadata[1])]) + coordinator._broadcast_to_engines([Headers.RELEASE_KV.value, request_id_of(metadata)]) @message_handler( @@ -261,7 +274,7 @@ def handle_control_signal(coordinator, sender_identity, metadata, bodies): logging.warning("Coordinator: ignoring signal from unknown client.") return - header = Headers(metadata[0]) + header = header_of(metadata) transition = CONTROL_TRANSITIONS[header] if coordinator.state not in transition.allowed_from: # Silently ignore redundant signals; warn on genuinely invalid ones. @@ -321,7 +334,9 @@ def handle_engine_reply(coordinator, sender_identity, metadata, bodies): ), f"ENGINE_REPLY from never-connected sender {sender_identity!r}" logging.warning("Coordinator: ENGINE_REPLY from removed engine %r", sender_identity) - for (fid, needs_detokenize), body in zip(metadata[1], bodies): + for finished in ENGINE_REPLY.parse(metadata, bodies): + fid = finished.request_id + body = finished.reply client_identity = coordinator.request_id_to_client_id[fid] client_request_id = coordinator.request_id_to_client_request_id[fid] del coordinator.request_id_to_client_id[fid] @@ -334,7 +349,7 @@ def handle_engine_reply(coordinator, sender_identity, metadata, bodies): assert coordinator._pending_counts[idx] >= 1 coordinator._pending_counts[idx] -= 1 - if needs_detokenize: + if finished.needs_detokenize: # Detokenizing writes generated_text into the reply, so this one has # to be decoded and re-encoded. Clients that detokenize for # themselves (the OpenAI frontend does) never take this path. @@ -342,10 +357,9 @@ def handle_engine_reply(coordinator, sender_identity, metadata, bodies): coordinator.detokenize(finished_request) body = msgpack.packb(finished_request, use_bin_type=True) - reply_metadata = msgpack.packb( - [Headers.ENGINE_REPLY.value, client_request_id], use_bin_type=True + coordinator.router_socket.send_multipart( + [client_identity, *CLIENT_REPLY.pack(request_id=client_request_id, reply=body)] ) - coordinator.router_socket.send_multipart([client_identity, reply_metadata, body]) @message_handler(Headers.ENGINE_REPLY_PARTIAL) @@ -367,18 +381,15 @@ def handle_engine_reply_partial(coordinator, sender_identity, metadata, bodies): ), f"ENGINE_REPLY_PARTIAL from never-connected sender {sender_identity!r}" logging.warning("Coordinator: ENGINE_REPLY_PARTIAL from removed engine %r", sender_identity) return - for request_id, body in zip(metadata[1], bodies): - client_identity = coordinator.request_id_to_client_id[request_id] - client_request_id = coordinator.request_id_to_client_request_id[request_id] + for partial in ENGINE_REPLY_PARTIAL.parse(metadata, bodies): + client_identity = coordinator.request_id_to_client_id[partial.request_id] + client_request_id = coordinator.request_id_to_client_request_id[partial.request_id] # Partial tokens are detokenized incrementally by the client-facing # streaming layer, so the body is always forwarded untouched. coordinator.router_socket.send_multipart( [ client_identity, - msgpack.packb( - [Headers.ENGINE_REPLY_PARTIAL.value, client_request_id], use_bin_type=True - ), - body, + *CLIENT_REPLY_PARTIAL.pack(request_id=client_request_id, partial=partial.partial), ] ) @@ -396,16 +407,13 @@ def handle_abort_request(coordinator, sender_identity, metadata, bodies): if sender_identity not in coordinator.known_clients: logging.warning("Coordinator: ignoring abort from unknown client.") return - client_request_id = int(metadata[1]) + client_request_id = request_id_of(metadata) request_id = coordinator.client_request_to_request_id.get((sender_identity, client_request_id)) if request_id is None: return assigned_rank = coordinator.request_id_to_rank.get(request_id) if assigned_rank is not None: - coordinator._send_to_engine( - assigned_rank, - [msgpack.packb([Headers.ABORT_REQUEST.value, request_id], use_bin_type=True)], - ) + coordinator._send_to_engine(assigned_rank, pack_signal(Headers.ABORT_REQUEST, request_id)) @message_handler(Headers.SHUTDOWN) diff --git a/megatron/core/inference/engines/dynamic_engine.py b/megatron/core/inference/engines/dynamic_engine.py index b406e878e56..242167ca719 100644 --- a/megatron/core/inference/engines/dynamic_engine.py +++ b/megatron/core/inference/engines/dynamic_engine.py @@ -43,6 +43,18 @@ FinishedRequestRecord, Status, ) +from megatron.core.inference.messages import ( + ENGINE_REPLY, + ENGINE_REPLY_PARTIAL, + KV_HANDOFF_COMPLETE, + SEND_KV, + SET_GENERATION_EPOCH, + SUBMIT_REQUEST, + SUBMIT_REQUEST_WITH_KV, + header_of, + pack_signal, + request_id_of, +) from megatron.core.inference.sampling_params import SamplingParams from megatron.core.inference.text_generation_controllers.text_generation_controller import ( DecodeOnly, @@ -167,19 +179,16 @@ def _engine_reply_frames(finished_requests: List[dict]) -> List[bytes]: Returns: The frames to send, metadata first. """ - metadata = [ - Headers.ENGINE_REPLY.value, - [ - [ + return ENGINE_REPLY.pack( + entries=[ + ( request["request_id"], bool((request.get("sampling_params") or {}).get("detokenize_generations")), - ] + ) for request in finished_requests ], - ] - return [msgpack.packb(metadata, use_bin_type=True)] + [ - msgpack.packb(request, use_bin_type=True) for request in finished_requests - ] + payloads=[msgpack.packb(request, use_bin_type=True) for request in finished_requests], + ) def _get_decode_only_log_state( @@ -2427,13 +2436,10 @@ def _try_send_streaming_partials(self) -> None: nvtx_range_push("coordinator_streaming") self.socket_for_receiving_requests.send_multipart( - [ - msgpack.packb( - [Headers.ENGINE_REPLY_PARTIAL.value, [p["request_id"] for p in partials]], - use_bin_type=True, - ) - ] - + [msgpack.packb(p, use_bin_type=True) for p in partials] + ENGINE_REPLY_PARTIAL.pack( + entries=[p["request_id"] for p in partials], + payloads=[msgpack.packb(p, use_bin_type=True) for p in partials], + ) ) nvtx_range_pop("coordinator_streaming") @@ -2904,11 +2910,7 @@ def schedule_requests(self) -> int: # Locally-generated notifications are single-frame messages, so they # are wrapped to match the frame-list shape of socket traffic. all_messages.extend( - [ - msgpack.packb( - [Headers.KV_HANDOFF_COMPLETE.value, request_id, failed], use_bin_type=True - ) - ] + KV_HANDOFF_COMPLETE.pack(request_id=request_id, failed=failed) for request_id, failed in self._drain_handoff_completion_notifications() ) while True: @@ -2934,40 +2936,47 @@ def schedule_requests(self) -> int: # Control signals are queued for the second pass. new_generation_epoch = None for message in all_messages: - data = msgpack.unpackb(message[0], raw=False) - header = Headers(data[0]) + metadata = msgpack.unpackb(message[0], raw=False) + header = header_of(metadata) if header == Headers.SUBMIT_REQUEST: - request_id, sampling_params = data[1:] + request = SUBMIT_REQUEST.parse(metadata, message[1:]) # The prompt rides in its own frame; the engine is its first # consumer, so this is where it finally gets decoded. - prompt = msgpack.unpackb(message[1], raw=False) - sampling_params = SamplingParams.deserialize(sampling_params) nvtx_range_push("add_request") - self.add_request(request_id, prompt, sampling_params) + self.add_request( + request.request_id, + msgpack.unpackb(request.prompt, raw=False), + SamplingParams.deserialize(request.sampling_params), + ) nvtx_range_pop("add_request") elif header == Headers.SUBMIT_REQUEST_WITH_KV: # Decode-side KV import. As on the plain path, the prompt rides # in its own frame and the engine is its first consumer. - request_id, sampling_params, kv_meta = data[1:] - prompt = msgpack.unpackb(message[1], raw=False) - src_block_ids = msgpack.unpackb(message[2], raw=False) - sampling_params = SamplingParams.deserialize(sampling_params) + request = SUBMIT_REQUEST_WITH_KV.parse(metadata, message[1:]) nvtx_range_push("add_request_with_kv_handoff") self.add_request_with_kv_handoff( - request_id, prompt, sampling_params, kv_meta, src_block_ids + request.request_id, + msgpack.unpackb(request.prompt, raw=False), + SamplingParams.deserialize(request.sampling_params), + request.kv_meta, + msgpack.unpackb(request.src_block_ids, raw=False), ) nvtx_range_pop("add_request_with_kv_handoff") elif header == Headers.RELEASE_KV: # Coordinator-broadcast release. Unknown request ids are no-ops. - self.release_handoff_blocks(int(data[1])) + self.release_handoff_blocks(request_id_of(metadata)) elif header == Headers.SEND_KV: # Push transport: send a pinned hand-off's KV to the decode # instance the coordinator picked. - self.push_handoff_kv(int(data[1]), data[2]) + send_kv = SEND_KV.parse(metadata, ()) + self.push_handoff_kv(int(send_kv.request_id), send_kv.decode_metas) elif header == Headers.KV_HANDOFF_COMPLETE: - self._record_handoff_completion_notification(int(data[1]), bool(data[2])) + completion = KV_HANDOFF_COMPLETE.parse(metadata, ()) + self._record_handoff_completion_notification( + int(completion.request_id), bool(completion.failed) + ) elif header == Headers.ABORT_REQUEST: - request_id = int(data[1]) + request_id = request_id_of(metadata) entry = self.requests.get(request_id) if entry is not None: request = entry.record[-1] @@ -2983,7 +2992,7 @@ def schedule_requests(self) -> int: + self.context.request_query_lengths[idx] ) elif header == Headers.SET_GENERATION_EPOCH: - new_generation_epoch = data[1] + new_generation_epoch = SET_GENERATION_EPOCH.parse(metadata, ()).generation_epoch elif header == Headers.START_CUDA_PROFILER: # Side-effect, not a state transition: apply immediately on every # rank so an outer nsys --capture-range=cudaProfilerApi starts here. @@ -3019,8 +3028,8 @@ def schedule_requests(self) -> int: # processes one state transition per iteration). if self._pending_signals: message = self._pending_signals.popleft() - data = msgpack.unpackb(message[0], raw=False) - header = Headers(data[0]) + metadata = msgpack.unpackb(message[0], raw=False) + header = header_of(metadata) if header == Headers.PAUSE: if self.state == EngineState.RUNNING: @@ -3074,7 +3083,7 @@ async def shutdown(self): sock = getattr(self, 'socket_for_receiving_requests', None) if sock is not None and not sock.closed: try: - sock.send(msgpack.packb([Headers.DISCONNECT.value], use_bin_type=True)) + sock.send_multipart(pack_signal(Headers.DISCONNECT)) except Exception: pass for socket in getattr(self, 'zmq_sockets', []): diff --git a/megatron/core/inference/inference_client.py b/megatron/core/inference/inference_client.py index 91783ee0906..daa1d5524a0 100644 --- a/megatron/core/inference/inference_client.py +++ b/megatron/core/inference/inference_client.py @@ -14,6 +14,14 @@ from megatron.core.utils import get_asyncio_loop, trace_async_exceptions from .headers import Headers +from .messages import ( + CLIENT_REPLY, + CLIENT_REPLY_PARTIAL, + SUBMIT_REQUEST, + SUBMIT_REQUEST_WITH_KV, + header_of, + pack_signal, +) try: import zmq @@ -113,16 +121,14 @@ def add_request( """ request_id = self.next_request_id self.next_request_id += 1 - frames = [ - msgpack.packb( - [Headers.SUBMIT_REQUEST.value, request_id, sampling_params.serialize()], - use_bin_type=True, - ), - # The prompt travels in its own frame so the coordinator can route the - # request without decoding it -- at long prompts that decode dominates - # its per-request cost, and it is one serial loop shared by every rank. - self._pack_prompt(prompt), - ] + # The prompt travels in its own frame so the coordinator can route the + # request without decoding it -- at long prompts that decode dominates its + # per-request cost, and it is one serial loop shared by every rank. + frames = SUBMIT_REQUEST.pack( + request_id=request_id, + sampling_params=sampling_params.serialize(), + prompt=self._pack_prompt(prompt), + ) return self._submit_request(frames, request_id) @staticmethod @@ -160,19 +166,13 @@ def _make_kv_handoff_request( """ request_id = self.next_request_id self.next_request_id += 1 - frames = [ - msgpack.packb( - [ - Headers.SUBMIT_REQUEST_WITH_KV.value, - request_id, - sampling_params.serialize(), - kv_meta, - ], - use_bin_type=True, - ), - self._pack_prompt(prompt), - msgpack.packb(list(src_block_ids), use_bin_type=True), - ] + frames = SUBMIT_REQUEST_WITH_KV.pack( + request_id=request_id, + sampling_params=sampling_params.serialize(), + kv_meta=kv_meta, + prompt=self._pack_prompt(prompt), + src_block_ids=msgpack.packb(list(src_block_ids), use_bin_type=True), + ) return request_id, frames def add_request_with_kv_handoff( @@ -234,8 +234,7 @@ def release_handoff(self, request_id: int) -> None: Fire-and-forget. The coordinator broadcasts RELEASE_KV to every engine; engines without that request_id ignore the message. """ - payload = [Headers.RELEASE_KV.value, int(request_id)] - self.socket.send(msgpack.packb(payload, use_bin_type=True)) + self.socket.send_multipart(pack_signal(Headers.RELEASE_KV, int(request_id))) def abort_request(self, request_id: int) -> None: """Cancel an in-flight request and close its local response stream.""" @@ -248,8 +247,7 @@ def abort_request(self, request_id: int) -> None: if future is not None and not future.done(): future.cancel() self.request_submission_times.pop(request_id, None) - payload = [Headers.ABORT_REQUEST.value, request_id] - self.socket.send(msgpack.packb(payload, use_bin_type=True)) + self.socket.send_multipart(pack_signal(Headers.ABORT_REQUEST, request_id)) def add_request_streaming( self, prompt: Union[str, List[int]], sampling_params: SamplingParams @@ -279,16 +277,14 @@ def add_request_streaming( sampling_params.streaming = True request_id = self.next_request_id self.next_request_id += 1 - frames = [ - msgpack.packb( - [Headers.SUBMIT_REQUEST.value, request_id, sampling_params.serialize()], - use_bin_type=True, - ), - # The prompt travels in its own frame so the coordinator can route the - # request without decoding it -- at long prompts that decode dominates - # its per-request cost, and it is one serial loop shared by every rank. - self._pack_prompt(prompt), - ] + # The prompt travels in its own frame so the coordinator can route the + # request without decoding it -- at long prompts that decode dominates its + # per-request cost, and it is one serial loop shared by every rank. + frames = SUBMIT_REQUEST.pack( + request_id=request_id, + sampling_params=sampling_params.serialize(), + prompt=self._pack_prompt(prompt), + ) return self._submit_stream(frames, request_id) def _submit_request(self, frames: list, request_id: int) -> asyncio.Future: @@ -327,14 +323,15 @@ async def _recv_task(self): try: # frames[0] is metadata; a reply body, when present, follows it. frames = self.socket.recv_multipart(flags=zmq.NOBLOCK) - data = msgpack.unpackb(frames[0], raw=False) - header = Headers(data[0]) + metadata = msgpack.unpackb(frames[0], raw=False) + header = header_of(metadata) if header == Headers.ENGINE_REPLY: - request_id = data[1] + delivery = CLIENT_REPLY.parse(metadata, frames[1:]) + request_id = delivery.request_id if request_id in self.aborted_request_ids: self.aborted_request_ids.discard(request_id) continue - reply = msgpack.unpackb(frames[1], raw=False) + reply = msgpack.unpackb(delivery.reply, raw=False) submitted = self.request_submission_times.pop(request_id, None) if submitted is not None: reply['latency'] = time.perf_counter() - submitted @@ -358,10 +355,10 @@ async def _recv_task(self): ) completion_future.set_result(completed_request) elif header == Headers.ENGINE_REPLY_PARTIAL: - request_id = data[1] - stream = self.streams.get(request_id) + delivery = CLIENT_REPLY_PARTIAL.parse(metadata, frames[1:]) + stream = self.streams.get(delivery.request_id) if stream is not None: - stream.put({"partial": msgpack.unpackb(frames[1], raw=False)}) + stream.put({"partial": msgpack.unpackb(delivery.partial, raw=False)}) except zmq.Again: await asyncio.sleep(0.005) continue @@ -375,14 +372,13 @@ def _connect_with_inference_coordinator(self, timeout_seconds: Optional[float] = Sends a CONNECT signal and waits for a CONNECT_ACK reply to ensure the connection is established and acknowledged by the coordinator. """ - payload = [Headers.CONNECT.value] - self.socket.send(msgpack.packb(payload, use_bin_type=True)) + self.socket.send_multipart(pack_signal(Headers.CONNECT)) if timeout_seconds is not None and not self.socket.poll( timeout=max(0, int(timeout_seconds * 1000)) ): raise TimeoutError("Timed out connecting to the Megatron inference coordinator") reply = msgpack.unpackb(self.socket.recv_multipart()[0], raw=False) - assert Headers(reply[0]) == Headers.CONNECT_ACK + assert header_of(reply) == Headers.CONNECT_ACK def start( self, @@ -409,9 +405,7 @@ def _send_signal_to_engines(self, signal, *args): signal: The signal to send, typically a value from the `Headers` enum. *args: Optional extra values to include in the payload. """ - payload = [signal.value, *args] - payload_serialized = msgpack.packb(payload, use_bin_type=True) - self.socket.send(payload_serialized) + self.socket.send_multipart(pack_signal(signal, *args)) def pause_engines(self): """Sends PAUSE to all engines via coordinator. diff --git a/megatron/core/inference/messages.py b/megatron/core/inference/messages.py new file mode 100644 index 00000000000..f8881c0b58d --- /dev/null +++ b/megatron/core/inference/messages.py @@ -0,0 +1,224 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +"""Wire schema for messages exchanged with the data parallel inference coordinator. + +Every message is a list of ZMQ frames: one **metadata** frame followed by zero or +more **payload** frames. + +The metadata frame is the only frame the coordinator decodes, so it holds just +what routing needs and nothing whose size follows the sequence length. Payload +frames carry the bulk -- prompts inbound, finished requests outbound -- and are +forwarded as opaque bytes, which keeps the coordinator's per-request cost flat in +prompt length. Payload boundaries come from ZMQ multipart, so splitting a batch +costs the coordinator nothing. + +Each message type is declared once, as a :class:`MessageSpec` naming its metadata +fields and payload frames in wire order. ``pack`` and ``parse`` are both derived +from that declaration, so no frame index is written by hand and the two +directions cannot drift apart. Adding a field means editing one tuple. + +Which tuple a name goes in is the whole contract: ``metadata_fields`` are decoded +by the coordinator and so must be constant-size, while ``payload_frames`` are +never decoded in transit and may grow with the prompt. +""" + +from collections import namedtuple +from dataclasses import dataclass +from typing import Any, List, Sequence, Tuple + +from megatron.core.inference.headers import Headers + +try: + import msgpack +except ImportError: + msgpack = None + + +def header_of(metadata: Sequence) -> Headers: + """Return the header of a decoded metadata frame.""" + return Headers(metadata[0]) + + +@dataclass(frozen=True) +class MessageSpec: + """Declares one message's frame layout. + + Attributes: + name: Name of the named tuple returned by :meth:`parse`. + header: The header this layout belongs to. + metadata_fields: Names of the values following the header in the metadata + frame, in wire order. Decoded by the coordinator, so every one of + these must be constant-size. + payload_frames: Names of the opaque frames following the metadata frame, + in wire order. Never decoded in transit; these may grow with the + prompt. + """ + + name: str + header: Headers + metadata_fields: Tuple[str, ...] = () + payload_frames: Tuple[str, ...] = () + + def __post_init__(self) -> None: + object.__setattr__( + self, "_tuple", namedtuple(self.name, self.metadata_fields + self.payload_frames) + ) + + def pack(self, **values: Any) -> List[bytes]: + """Build the frames for this message from named values.""" + metadata = msgpack.packb( + [self.header.value, *(values[name] for name in self.metadata_fields)], use_bin_type=True + ) + return [metadata, *(values[name] for name in self.payload_frames)] + + def parse(self, metadata: Sequence, bodies: Sequence[bytes]) -> Any: + """Read a received message into a named tuple.""" + return self._tuple(*metadata[1:], *bodies[: len(self.payload_frames)]) + + +@dataclass(frozen=True) +class BatchedMessageSpec: + """Declares a message carrying N independent entries, one payload frame each. + + Engine replies are batched per engine step but fan out to different clients, + so each entry keeps its own frame: that is what lets the coordinator route the + batch without decoding any of it. The metadata frame holds one routing record + per entry, in the same order as the frames. + + Attributes: + name: Name of the named tuple returned per entry by :meth:`parse`. + header: The header this layout belongs to. + entry_fields: Names of the per-entry routing values held in the metadata + frame. A single field is packed as a bare scalar per entry; several + are packed as a list per entry. + payload_frame: Name of the opaque frame carried per entry. + """ + + name: str + header: Headers + entry_fields: Tuple[str, ...] + payload_frame: str + + def __post_init__(self) -> None: + object.__setattr__( + self, "_tuple", namedtuple(self.name, self.entry_fields + (self.payload_frame,)) + ) + + def pack(self, entries: Sequence[Any], payloads: Sequence[bytes]) -> List[bytes]: + """Build the frames for a batch. + + Args: + entries: One routing record per entry -- a scalar when there is a + single entry field, otherwise a sequence of that field's values. + payloads: One opaque frame per entry, in the same order. + """ + assert len(entries) == len(payloads), ( + f"{self.name}: {len(entries)} entries but {len(payloads)} payload frames; " + "the metadata frame must describe exactly the frames that follow" + ) + if len(self.entry_fields) == 1: + records = list(entries) + else: + records = [list(entry) for entry in entries] + return [msgpack.packb([self.header.value, records], use_bin_type=True), *payloads] + + def parse(self, metadata: Sequence, bodies: Sequence[bytes]) -> List[Any]: + """Read a batch, pairing each routing record with its frame.""" + records = metadata[1] + if len(self.entry_fields) == 1: + return [self._tuple(record, body) for record, body in zip(records, bodies)] + return [self._tuple(*record, body) for record, body in zip(records, bodies)] + + +# --- client -> coordinator -> engine ------------------------------------------------ + +SUBMIT_REQUEST = MessageSpec( + "SubmitRequest", + Headers.SUBMIT_REQUEST, + metadata_fields=("request_id", "sampling_params"), + payload_frames=("prompt",), +) +"""A plain inference request.""" + +SUBMIT_REQUEST_WITH_KV = MessageSpec( + "SubmitRequestWithKV", + Headers.SUBMIT_REQUEST_WITH_KV, + # kv_meta is the peer's transfer metadata, bounded by TP size and + # num_speculative_tokens. src_block_ids names one block per block_size_tokens + # of prompt, so it grows with the prompt and belongs in a frame. + metadata_fields=("request_id", "sampling_params", "kv_meta"), + payload_frames=("prompt", "src_block_ids"), +) +"""A decode request whose KV was computed by a prefill peer.""" + + +# --- engine -> coordinator ---------------------------------------------------------- + +ENGINE_REPLY = BatchedMessageSpec( + "FinishedRequest", + Headers.ENGINE_REPLY, + # needs_detokenize rides in the metadata so the coordinator can skip decoding + # the reply entirely for clients that detokenize for themselves. + entry_fields=("request_id", "needs_detokenize"), + payload_frame="reply", +) +"""A batch of finished requests, one payload frame each.""" + +ENGINE_REPLY_PARTIAL = BatchedMessageSpec( + "PartialReply", + Headers.ENGINE_REPLY_PARTIAL, + entry_fields=("request_id",), + payload_frame="partial", +) +"""A batch of incremental replies, one payload frame each.""" + + +# --- coordinator -> client ---------------------------------------------------------- + +CLIENT_REPLY = MessageSpec( + "ClientReply", Headers.ENGINE_REPLY, metadata_fields=("request_id",), payload_frames=("reply",) +) +"""One final reply, split out of an engine batch and re-addressed to its client.""" + +CLIENT_REPLY_PARTIAL = MessageSpec( + "ClientPartialReply", + Headers.ENGINE_REPLY_PARTIAL, + metadata_fields=("request_id",), + payload_frames=("partial",), +) +"""One incremental reply, re-addressed to its client.""" + + +# --- control ------------------------------------------------------------------------ + +KV_HANDOFF_COMPLETE = MessageSpec( + "KVHandoffComplete", Headers.KV_HANDOFF_COMPLETE, metadata_fields=("request_id", "failed") +) +"""A model-parallel-agreed handoff outcome, distributed over the schedule broadcast.""" + +SEND_KV = MessageSpec("SendKV", Headers.SEND_KV, metadata_fields=("request_id", "decode_metas")) +"""An instruction to push a pinned handoff's KV to a decode instance.""" + +SET_GENERATION_EPOCH = MessageSpec( + "SetGenerationEpoch", Headers.SET_GENERATION_EPOCH, metadata_fields=("generation_epoch",) +) +"""The only control signal carrying an argument.""" + + +def pack_signal(header: Headers, *args: Any) -> List[bytes]: + """Frame a control signal as a single metadata frame with no payload. + + Covers the messages whose whole content is a header plus at most a couple of + scalars: the handshake (CONNECT, CONNECT_ACK, DISCONNECT), lifecycle control + (PAUSE, UNPAUSE, SUSPEND, RESUME, STOP, SHUTDOWN), the profiler signals, and + the request-scoped controls ABORT_REQUEST and RELEASE_KV. + + The coordinator rebroadcasts decoded metadata verbatim, so args survive the + hop to the engines. + """ + return [msgpack.packb([header.value, *args], use_bin_type=True)] + + +def request_id_of(metadata: Sequence) -> int: + """Return the request id of a control signal that names exactly one request.""" + return int(metadata[1])