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..b2700b80125 100644 --- a/megatron/core/inference/data_parallel_inference_coordinator/handlers.py +++ b/megatron/core/inference/data_parallel_inference_coordinator/handlers.py @@ -8,17 +8,36 @@ 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 +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 @@ -52,35 +71,44 @@ 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 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) -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:] + + 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. request_id = coordinator.next_request_id @@ -89,19 +117,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_frames = SUBMIT_REQUEST.pack( + request_id=request_id, sampling_params=request.sampling_params, prompt=request.prompt ) - 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(request.prompt, raw=False) + ) + else: + request_hashes = [] + if ( coordinator.prefix_caching_coordinator_policy == PrefixCachingCoordinatorPolicy.FIRST_PREFIX_BLOCK @@ -111,7 +146,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_frames): break else: # If all engines have died, we are in an abnormal state, and must exit cleanly. @@ -136,46 +171,57 @@ 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:] + 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 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, - ], - use_bin_type=True, + # Rebuilding the metadata frame is cheap: it holds no prompt tokens. + 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_payload): + if coordinator._send_to_engine(next_identity, engine_frames): break else: logging.error("Coordinator: no reachable engines for handoff request %d", request_id) @@ -189,13 +235,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, request_id_of(metadata)]) @message_handler( @@ -206,13 +259,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 = header_of(metadata) transition = CONTROL_TRANSITIONS[header] if coordinator.state not in transition.allowed_from: # Silently ignore redundant signals; warn on genuinely invalid ones. @@ -222,9 +284,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 +294,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 +333,10 @@ 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 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] @@ -272,63 +349,82 @@ def handle_engine_reply(coordinator, sender_identity, payload): assert coordinator._pending_counts[idx] >= 1 coordinator._pending_counts[idx] -= 1 + 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. + finished_request = msgpack.unpackb(body, raw=False) + coordinator.detokenize(finished_request) + body = msgpack.packb(finished_request, use_bin_type=True) + coordinator.router_socket.send_multipart( - [ - client_identity, - msgpack.packb( - [Headers.ENGINE_REPLY.value, client_request_id, finished_request], - use_bin_type=True, - ), - ] + [client_identity, *CLIENT_REPLY.pack(request_id=client_request_id, reply=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"] - 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. + 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, partial], - use_bin_type=True, - ), + *CLIENT_REPLY_PARTIAL.pack(request_id=client_request_id, partial=partial.partial), ] ) @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 = 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) -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 +432,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..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, @@ -152,6 +164,33 @@ 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. + """ + return ENGINE_REPLY.pack( + entries=[ + ( + request["request_id"], + bool((request.get("sampling_params") or {}).get("detokenize_generations")), + ) + for request in finished_requests + ], + payloads=[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 +1126,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 +2434,13 @@ 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( + 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") self._partial_emit_lengths.update(emit_lengths) @@ -2786,6 +2827,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 +2907,28 @@ 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 - ) + KV_HANDOFF_COMPLETE.pack(request_id=request_id, failed=failed) 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,34 +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, raw=False) - header = Headers(data[0]) + metadata = msgpack.unpackb(message[0], raw=False) + header = header_of(metadata) if header == Headers.SUBMIT_REQUEST: - request_id, prompt, sampling_params = data[1:] - sampling_params = SamplingParams.deserialize(sampling_params) + 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. 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. - request_id, prompt, sampling_params, kv_meta, src_block_ids = data[1:] - sampling_params = SamplingParams.deserialize(sampling_params) + # Decode-side KV import. As on the plain path, the prompt rides + # in its own frame and the engine is its first consumer. + 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] @@ -2889,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. @@ -2925,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, 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: @@ -2980,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 199018ce3a6..daa1d5524a0 100644 --- a/megatron/core/inference/inference_client.py +++ b/megatron/core/inference/inference_client.py @@ -6,12 +6,22 @@ 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 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 @@ -111,8 +121,30 @@ 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) + # 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 + 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 +153,27 @@ 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), - ] - return request_id, payload + 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( 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`. @@ -193,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.""" @@ -207,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 @@ -238,21 +277,28 @@ 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) + # 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, 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 +321,17 @@ async def _recv_task(self): """ while True: try: - data = msgpack.unpackb(self.socket.recv(flags=zmq.NOBLOCK), raw=False) - header = Headers(data[0]) + # frames[0] is metadata; a reply body, when present, follows it. + frames = self.socket.recv_multipart(flags=zmq.NOBLOCK) + metadata = msgpack.unpackb(frames[0], raw=False) + header = header_of(metadata) if header == Headers.ENGINE_REPLY: - request_id, reply = 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(delivery.reply, raw=False) submitted = self.request_submission_times.pop(request_id, None) if submitted is not None: reply['latency'] = time.perf_counter() - submitted @@ -305,10 +355,10 @@ async def _recv_task(self): ) completion_future.set_result(completed_request) elif header == Headers.ENGINE_REPLY_PARTIAL: - request_id, partial = 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": partial}) + stream.put({"partial": msgpack.unpackb(delivery.partial, raw=False)}) except zmq.Again: await asyncio.sleep(0.005) continue @@ -322,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(), raw=False) - assert Headers(reply[0]) == Headers.CONNECT_ACK + reply = msgpack.unpackb(self.socket.recv_multipart()[0], raw=False) + assert header_of(reply) == Headers.CONNECT_ACK def start( self, @@ -356,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]) 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())