diff --git a/megatron/core/inference/config.py b/megatron/core/inference/config.py index 44bbf410c7a..3e95e0f4b97 100644 --- a/megatron/core/inference/config.py +++ b/megatron/core/inference/config.py @@ -114,6 +114,52 @@ class PrefixCachingCoordinatorPolicy(str, Enum): """Route to the rank with the fewest in-flight requests. Ignores prefix affinity.""" +class PrefixCachingCostPolicy(str, Enum): + """How the coordinator weighs prefix affinity against rank load. + + Orthogonal to `PrefixCachingCoordinatorPolicy`, which only selects the affinity + signal. Both signals are normalized to the fraction of the request already cached + on a rank, in [0, 1], so either cost policy composes with either of + LONGEST_PREFIX and FIRST_PREFIX_BLOCK. Neither applies under LOAD_BALANCED, which + ignores affinity entirely. + + RELATIVE_LOAD_WEIGHTED (default) — score = fraction - beta * relative_load, highest + wins, where relative_load is (load - mean) / max(1, mean). Approximates the session + stickiness a session-affinity router gets for free: a multi-turn request lands back + on the rank holding its history, with no session id to key on. Both terms are + normalized, so beta is dimensionless, and measuring load against the fleet mean + makes the penalty vanish while ranks are balanced -- at saturation this is pure + affinity, and load only pulls toward idle ranks as the fleet diverges. The mean is + floored at 1 so a near-idle fleet does not turn one in-flight request into a large + relative load and thrash on noise. + + FREE_CAPACITY_WEIGHTED — score = alpha * fraction + (1 - alpha) * free_capacity, highest + wins, with alpha from `prefix_caching_routing_alpha`. Fixes the trade-off in + absolute terms rather than relative to how loaded the fleet actually is. + """ + + RELATIVE_LOAD_WEIGHTED = "relative_load_weighted" + FREE_CAPACITY_WEIGHTED = "free_capacity_weighted" + + +def routes_on_prefix(policy) -> bool: + """Whether `policy` needs per-request block hashes to make a routing decision. + + Frontends call this to decide whether hashing a prompt is worth anything: under + LOAD_BALANCED the coordinator discards the hashes, so computing them is pure + overhead on the request path. Kept beside the enum so a new prefix-aware policy + only has to be added in one place. + + Accepts the enum, its string value, or None (no policy configured). + """ + if policy is None: + return False + return PrefixCachingCoordinatorPolicy(policy) in ( + PrefixCachingCoordinatorPolicy.LONGEST_PREFIX, + PrefixCachingCoordinatorPolicy.FIRST_PREFIX_BLOCK, + ) + + class KVCacheManagementMode(str, Enum): """Mode for handling large tensors (KV cache, Mamba states) during suspend/resume.""" @@ -320,7 +366,7 @@ class InferenceConfig: """ prefix_caching_coordinator_policy: PrefixCachingCoordinatorPolicy = ( - PrefixCachingCoordinatorPolicy.LOAD_BALANCED + PrefixCachingCoordinatorPolicy.LONGEST_PREFIX ) """Routing policy for the DP inference coordinator. See `PrefixCachingCoordinatorPolicy` for options. @@ -328,10 +374,35 @@ class InferenceConfig: Only applies when enable_prefix_caching is True and using a coordinator. """ + prefix_cache_ttl_seconds: float = 300.0 + """How long the coordinator assumes an engine still holds a block it routed. + + Only applies under `PrefixCachingEvictionPolicy.LRU`, where the engine keeps + released blocks and evicts them under memory pressure -- something the + coordinator cannot observe, so it approximates by age. Too long and it claims + hits on blocks already evicted, routing for affinity and paying a cold prefill + anyway; too short and it forgets blocks the engine still holds. + """ + + prefix_caching_cost_policy: PrefixCachingCostPolicy = ( + PrefixCachingCostPolicy.RELATIVE_LOAD_WEIGHTED + ) + """How prefix affinity is weighed against rank load. See `PrefixCachingCostPolicy`. + + Only applies when enable_prefix_caching is True and using a coordinator. + """ + + prefix_caching_load_beta: float = 1.0 + """Weight on the load penalty under `PrefixCachingCostPolicy.RELATIVE_LOAD_WEIGHTED`, + in units of "full cache hits per 100% above mean load". 0 disables the penalty + (pure affinity); 1.0 means a rank at twice the fleet mean forfeits a whole + prompt's worth of cache credit. + """ + prefix_caching_routing_alpha: float = 0.5 """Weight for prefix-aware scoring: score = alpha * match + (1 - alpha) * normalized_load. Higher alpha favors prefix cache hits; lower alpha favors load balance. - Must be in [0, 1]. Only applies when enable_prefix_caching is True and using a coordinator. + Must be in [0, 1]. Only applies under `PrefixCachingCostPolicy.FREE_CAPACITY_WEIGHTED`. """ prefix_caching_mamba_gb: Optional[float] = None diff --git a/megatron/core/inference/contexts/dynamic_context.py b/megatron/core/inference/contexts/dynamic_context.py index 0b50f50d6de..ffea55331f9 100644 --- a/megatron/core/inference/contexts/dynamic_context.py +++ b/megatron/core/inference/contexts/dynamic_context.py @@ -335,6 +335,11 @@ def __init__(self, model_config: TransformerConfig, inference_config: InferenceC # Hyperparameter for choosing to prioritize prefix hit matches vs minimizing idle load self.prefix_caching_routing_alpha = inference_config.prefix_caching_routing_alpha + self.prefix_caching_cost_policy = inference_config.prefix_caching_cost_policy + self.prefix_caching_load_beta = inference_config.prefix_caching_load_beta + + # How long the coordinator's model of this engine's prefix cache survives + self.prefix_cache_ttl_seconds = inference_config.prefix_cache_ttl_seconds # Monotonic clock for prefix caching LRU eviction ordering. # Incremented each engine step but kept independent so the engine step diff --git a/megatron/core/inference/data_parallel_inference_coordinator/coordinator.py b/megatron/core/inference/data_parallel_inference_coordinator/coordinator.py index 67eb1bb5829..bc875d5d488 100644 --- a/megatron/core/inference/data_parallel_inference_coordinator/coordinator.py +++ b/megatron/core/inference/data_parallel_inference_coordinator/coordinator.py @@ -6,6 +6,7 @@ import logging import signal import socket +import time from collections import deque from multiprocessing import Event from multiprocessing.connection import Connection @@ -13,7 +14,11 @@ import numpy as np import torch -from megatron.core.inference.config import PrefixCachingCoordinatorPolicy +from megatron.core.inference.config import ( + PrefixCachingCoordinatorPolicy, + PrefixCachingCostPolicy, + PrefixCachingEvictionPolicy, +) from megatron.core.inference.headers import Headers, UnknownHeaderError from megatron.core.inference.inference_request import compute_block_hashes_batched from megatron.core.inference.text_generation_controllers.text_generation_controller import ( @@ -96,9 +101,17 @@ def __init__( block_size_tokens: int | None = None, enable_prefix_caching: bool = False, prefix_caching_coordinator_policy: PrefixCachingCoordinatorPolicy = ( - PrefixCachingCoordinatorPolicy.FIRST_PREFIX_BLOCK + PrefixCachingCoordinatorPolicy.LONGEST_PREFIX ), prefix_caching_routing_alpha: float = 0.5, + prefix_caching_cost_policy: PrefixCachingCostPolicy = ( + PrefixCachingCostPolicy.RELATIVE_LOAD_WEIGHTED + ), + prefix_caching_load_beta: float = 1.0, + prefix_caching_eviction_policy: PrefixCachingEvictionPolicy = ( + PrefixCachingEvictionPolicy.LRU + ), + prefix_cache_ttl_seconds: float = 300.0, schedule_output_path: str | None = None, hostname: str | None = None, ): @@ -200,6 +213,8 @@ def __init__( self.enable_prefix_caching = enable_prefix_caching self.prefix_caching_coordinator_policy = prefix_caching_coordinator_policy self.prefix_caching_routing_alpha = prefix_caching_routing_alpha + self.prefix_caching_cost_policy = prefix_caching_cost_policy + self.prefix_caching_load_beta = prefix_caching_load_beta self.max_requests = max_requests assert self.max_requests is not None and self.max_requests > 0 @@ -218,11 +233,28 @@ def __init__( self._identities_list = list(sorted_identities) # rank_index → identity self._pending_counts = np.zeros(n_ranks, dtype=np.int32) - # Hash → {rank_idx: timestamp} dict for prefix cache affinity routing. - # Each key is a block hash; each value maps rank indices to assignment - # timestamps (positive int). Missing entries are implicitly zero. - self._hash_table: dict[int, dict[int, int]] = {} - self._hash_assignment_counter = 0 + # Hash → {rank_idx: [last_touch, refcount]} for prefix cache affinity + # routing. This models what each engine holds; it is a prediction, since + # the engines are never asked. Entries are retired the way the engine + # retires blocks, per prefix_caching_eviction_policy: + # REF_ZERO -- the engine deregisters at refcount 0, so drop the entry + # when the last request holding it completes. + # LRU -- the engine keeps released blocks and evicts under memory + # pressure, which the coordinator cannot observe, so + # approximate it by age: drop entries untouched for + # prefix_cache_ttl_seconds. + # Without either, the table only grows and grows steadily more optimistic, + # claiming hits on blocks the engine evicted long ago. + self._hash_table: dict[int, dict[int, list]] = {} + # (touch_time, hash, rank_idx) in insertion order, so TTL expiry is a + # sweep from the front costing only what it actually evicts. An entry + # refreshed since it was queued is skipped by comparing timestamps. + self._hash_expiry: deque = deque() + # request_id → hashes, kept only under REF_ZERO, where completion is what + # releases the blocks. + self.request_id_to_hashes: dict[int, list] = {} + self.prefix_caching_eviction_policy = prefix_caching_eviction_policy + self.prefix_cache_ttl_seconds = prefix_cache_ttl_seconds # Clients that have completed the CONNECT handshake. self.known_clients = set() @@ -303,14 +335,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 +358,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: @@ -343,13 +381,14 @@ def compute_request_hashes(self, prompt): return compute_block_hashes_batched(token_tensor, self.block_size_tokens) def get_best_data_parallel_rank(self, request_hashes): - """Select the best DP rank based on prefix cache affinity and load. + """Select the best DP rank by trading prefix-cache affinity against load. - Uses a scoring function: score = alpha * match + (1 - alpha) * normalized_load - where *match* is a policy-dependent affinity score in [0, 1] (binary for - ``first_prefix_block``, normalized prefix depth for ``longest_prefix``) - and normalized_load = free_slots / max_requests (higher means more free - capacity). + Two independent choices. `prefix_caching_coordinator_policy` picks the + affinity signal -- a binary first-block hit for ``first_prefix_block``, the + contiguous prefix depth for ``longest_prefix`` -- and both are normalized to + the fraction of the request already cached on each rank. Then + `prefix_caching_cost_policy` picks how that fraction is weighed against rank + load; see `PrefixCachingCostPolicy`. Args: request_hashes: List of block hashes for the request. @@ -365,32 +404,138 @@ def get_best_data_parallel_rank(self, request_hashes): if not self.enable_prefix_caching or not request_hashes: return self.get_least_loaded_data_parallel_rank() - match, recency = self._match_vector(request_hashes) - - alpha = self.prefix_caching_routing_alpha - - # Vectorized score: alpha * match + (1-alpha) * free_capacity_fraction. - free_slots = np.maximum(0, self.max_requests - self._pending_counts).astype(np.float64) - scores = alpha * match + (1.0 - alpha) * (free_slots / self.max_requests) - - # Tiebreak: highest score, then highest recency, then lowest rank index. n_ranks = len(self._identities_list) - order = np.lexsort((np.arange(n_ranks), -recency, -scores)) - best_idx = int(order[0]) - return self._identities_list[best_idx] + n_blocks = len(request_hashes) + + # Affinity signal, normalized to a fraction of the request in [0, 1]. + # `recency` only separates ranks under the first-block signal, where many tie + # at the same binary match; it is zero for the depth signal. + if ( + self.prefix_caching_coordinator_policy + == PrefixCachingCoordinatorPolicy.FIRST_PREFIX_BLOCK + ): + fraction, recency = self._match_vector(request_hashes) + else: + depth = self._prefix_depth_vector(request_hashes) + fraction = depth.astype(np.float64) / max(n_blocks, 1) + recency = np.zeros(n_ranks, dtype=np.float64) + + if self.prefix_caching_cost_policy == PrefixCachingCostPolicy.FREE_CAPACITY_WEIGHTED: + # Weighted sum of affinity and free capacity, both in [0, 1]. alpha fixes + # the trade-off in absolute terms, irrespective of how loaded the fleet is. + alpha = self.prefix_caching_routing_alpha + free_slots = np.maximum(0, self.max_requests - self._pending_counts).astype(np.float64) + scores = alpha * fraction + (1.0 - alpha) * (free_slots / self.max_requests) + # Tiebreak: highest score, then highest recency, then lowest rank index. + order = np.lexsort((np.arange(n_ranks), -recency, -scores)) + return self._identities_list[int(order[0])] + + # RELATIVE_LOAD_WEIGHTED: score = fraction - beta * relative_load, highest wins. + # + # Approximates session stickiness without a session id: a multi-turn request + # lands back on the rank holding its history. Both terms are normalized, so + # beta is dimensionless, and relative_load is measured against the fleet mean + # so it vanishes while ranks are balanced -- at saturation this is pure + # affinity, and load only pulls toward idle ranks as the fleet diverges. + # + # The mean is floored at 1 so a near-idle fleet does not turn one in-flight + # request into a large relative load and thrash on noise. + mean_load = float(self._pending_counts.mean()) if n_ranks else 0.0 + relative_load = (self._pending_counts - mean_load) / max(1.0, mean_load) + scores = fraction - self.prefix_caching_load_beta * relative_load + # Tiebreak: highest score, then least loaded, then lowest rank index. + order = np.lexsort((np.arange(n_ranks), self._pending_counts, -scores)) + return self._identities_list[int(order[0])] def _update_rank_hashes(self, rank_identity, request_hashes): - """Record that a rank owns the given hashes. + """Record that a rank owns the given hashes, and take a reference to them. Args: rank_identity: ZMQ identity of the target rank. request_hashes: List of block hashes assigned to this rank. """ rank_idx = self.identity_to_rank_index[rank_identity] - self._hash_assignment_counter += 1 - ts = self._hash_assignment_counter + now = time.monotonic() for h in request_hashes: - self._hash_table.setdefault(h, {})[rank_idx] = ts + entry = self._hash_table.setdefault(h, {}).get(rank_idx) + if entry is None: + self._hash_table[h][rank_idx] = [now, 1] + else: + entry[0] = now + entry[1] += 1 + # One timestamp for the whole call keeps the deque sorted. + self._hash_expiry.append((now, h, rank_idx)) + + def _release_rank_hashes(self, request_id, rank_identity): + """Drop the references a finished request held, under REF_ZERO. + + Mirrors the engine deregistering blocks once no request holds them. Blocks + shared with a still-running request keep a nonzero count and survive. + + Args: + request_id: Server-side request id that just completed. + rank_identity: ZMQ identity of the rank that served it. + """ + hashes = self.request_id_to_hashes.pop(request_id, None) + if not hashes: + return + rank_idx = self.identity_to_rank_index.get(rank_identity) + if rank_idx is None: # engine already removed; its column is gone + return + for h in hashes: + row = self._hash_table.get(h) + entry = row.get(rank_idx) if row is not None else None + if entry is None: + continue + entry[1] -= 1 + if entry[1] <= 0: + del row[rank_idx] + if not row: + del self._hash_table[h] + + def _expire_rank_hashes(self, now): + """Drop entries untouched for the TTL, under LRU. + + The deque is insertion-ordered, so expired entries form a prefix of it and + the sweep stops at the first live one -- the cost is what it evicts, not + the table size. A stale queue entry whose table entry was refreshed since + is recognized by its older timestamp and skipped. + """ + ttl = self.prefix_cache_ttl_seconds + while self._hash_expiry: + ts, h, rank_idx = self._hash_expiry[0] + if now - ts <= ttl: + break + self._hash_expiry.popleft() + row = self._hash_table.get(h) + entry = row.get(rank_idx) if row is not None else None + if entry is not None and entry[0] <= ts: + del row[rank_idx] + if not row: + del self._hash_table[h] + + def _prefix_depth_vector(self, hashes): + """Return per-rank contiguous prefix depth, in blocks. + + Prefix cache reuse requires an unbroken chain from the first block -- each + block hash chains the previous one's digest -- so a rank only benefits up + to its first miss. Walking forward and dropping ranks as they miss gives + each rank its true depth, and the loop exits as soon as no rank is left. + """ + n_ranks = len(self._identities_list) + depth = np.zeros(n_ranks, dtype=np.int64) + alive = np.ones(n_ranks, dtype=bool) + for h in hashes: + row = self._hash_table.get(h) + if row is None: + break + present = np.zeros(n_ranks, dtype=bool) + present[np.fromiter(row.keys(), dtype=np.intp, count=len(row))] = True + alive &= present + if not alive.any(): + break + depth += alive + return depth def _match_vector(self, hashes): """Return ``(match, recency)`` vectors of shape ``(n_ranks,)``. @@ -415,7 +560,10 @@ def _match_vector(self, hashes): present = np.zeros(n_ranks, dtype=bool) present[rank_idxs] = True recency = np.zeros(n_ranks, dtype=np.float64) - recency[rank_idxs] = np.fromiter(row.values(), dtype=np.float64) + # Values are [last_touch, refcount]; recency is the touch time. + recency[rank_idxs] = np.fromiter( + (e[0] for e in row.values()), dtype=np.float64, count=len(row) + ) if present.any(): return present.astype(np.float64) * ((i + 1.0) / n), recency return zeros, zeros.copy() @@ -431,20 +579,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): @@ -464,6 +617,9 @@ def detokenize(self, finished_request): finished_request (dict): The serialized merged request containing the generated tokens to be detokenized. It is modified in place. """ + if not finished_request["sampling_params"]["detokenize_generations"]: + return + detokenize_stop_sequence = (finished_request.get("sampling_params", {}) or {}).get( "detokenize_stop_sequence", False ) @@ -486,9 +642,17 @@ def entrypoint( block_size_tokens: int | None = None, enable_prefix_caching: bool = False, prefix_caching_coordinator_policy: PrefixCachingCoordinatorPolicy = ( - PrefixCachingCoordinatorPolicy.FIRST_PREFIX_BLOCK + PrefixCachingCoordinatorPolicy.LONGEST_PREFIX ), prefix_caching_routing_alpha: float = 0.5, + prefix_caching_cost_policy: PrefixCachingCostPolicy = ( + PrefixCachingCostPolicy.RELATIVE_LOAD_WEIGHTED + ), + prefix_caching_load_beta: float = 1.0, + prefix_caching_eviction_policy: PrefixCachingEvictionPolicy = ( + PrefixCachingEvictionPolicy.LRU + ), + prefix_cache_ttl_seconds: float = 300.0, schedule_output_path: str | None = None, hostname: str | None = None, ): @@ -523,6 +687,10 @@ def entrypoint( enable_prefix_caching=enable_prefix_caching, prefix_caching_coordinator_policy=prefix_caching_coordinator_policy, prefix_caching_routing_alpha=prefix_caching_routing_alpha, + prefix_caching_cost_policy=prefix_caching_cost_policy, + prefix_caching_load_beta=prefix_caching_load_beta, + prefix_caching_eviction_policy=prefix_caching_eviction_policy, + prefix_cache_ttl_seconds=prefix_cache_ttl_seconds, schedule_output_path=schedule_output_path, hostname=hostname, ) diff --git a/megatron/core/inference/data_parallel_inference_coordinator/handlers.py b/megatron/core/inference/data_parallel_inference_coordinator/handlers.py index 2c1712a1de3..1f6aee0d75c 100644 --- a/megatron/core/inference/data_parallel_inference_coordinator/handlers.py +++ b/megatron/core/inference/data_parallel_inference_coordinator/handlers.py @@ -8,16 +8,28 @@ 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 time -import torch - -from megatron.core.inference.config import PrefixCachingCoordinatorPolicy +from megatron.core.inference.config import ( + PrefixCachingCoordinatorPolicy, + PrefixCachingEvictionPolicy, +) from megatron.core.inference.headers import Headers from .state import CONTROL_TRANSITIONS, CoordinatorState @@ -52,7 +64,7 @@ def decorator(fn): @message_handler(Headers.CONNECT) -def handle_connect(coordinator, sender_identity, payload): +def handle_connect(coordinator, sender_identity, metadata, bodies): """Handshake with a new client, replying with a CONNECT_ACK.""" if sender_identity in coordinator.known_clients: logging.info(f"Client {sender_identity} sent a duplicate connect request. Ignoring ..") @@ -65,22 +77,24 @@ 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. + ``metadata`` is ``[header, client_request_id, sampling_params]``, ``bodies[0]`` + is the packed prompt, which is forwarded to the engine as-is, and ``bodies[1]`` + -- when present -- carries the prompt's block hashes, computed by the client. + 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] + hash_frame = bodies[1] if len(bodies) > 1 else None + # 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,29 +103,45 @@ 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. The client computes + # them: it already holds the tokens, there are many clients and one of these + # loops, and unpacking the prompt here to hash it would undo the frame split + # that keeps the coordinator's per-request cost flat in prompt length. + if ( + hash_frame is not None + and coordinator.enable_prefix_caching + and coordinator.prefix_caching_coordinator_policy + != PrefixCachingCoordinatorPolicy.LOAD_BALANCED + ): + request_hashes = msgpack.unpackb(hash_frame, raw=False) + else: + request_hashes = [] + if ( coordinator.prefix_caching_coordinator_policy == PrefixCachingCoordinatorPolicy.FIRST_PREFIX_BLOCK ): request_hashes = request_hashes[:1] + if ( + request_hashes + and coordinator.prefix_caching_eviction_policy == PrefixCachingEvictionPolicy.LRU + ): + # Before the routing decision below, not after: that decision reads the + # table, so an entry past its TTL has to be gone first or this request is + # placed on evidence the engine no longer holds. Sweeping here also means + # expiry needs no timer thread. + coordinator._expire_rank_hashes(time.monotonic()) + # 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. @@ -125,6 +155,9 @@ def handle_submit_request(coordinator, sender_identity, payload): coordinator._pending_counts[coordinator.identity_to_rank_index[next_identity]] += 1 if request_hashes: coordinator._update_rank_hashes(next_identity, request_hashes) + if coordinator.prefix_caching_eviction_policy == PrefixCachingEvictionPolicy.REF_ZERO: + # Completion is what releases these blocks, so remember what to drop. + coordinator.request_id_to_hashes[request_id] = request_hashes if coordinator.schedule_records is not None: coordinator.schedule_records.append( { @@ -143,13 +176,13 @@ def handle_submit_request(coordinator, sender_identity, payload): Headers.SET_GENERATION_EPOCH, Headers.STOP, ) -def handle_control_signal(coordinator, sender_identity, payload): +def handle_control_signal(coordinator, sender_identity, metadata, bodies): """Validate a control signal against the transition table and broadcast it.""" 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. @@ -159,9 +192,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: @@ -169,7 +202,7 @@ 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. Profiler control is not a coordinator state transition, so there are no @@ -178,12 +211,18 @@ def handle_cuda_profiler_signal(coordinator, sender_identity, payload): 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. + + ``metadata`` is ``[header, [[request_id, needs_detokenize], ...]]`` and + ``bodies[i]`` is the packed finished request for entry ``i``. A body is + decoded only when the coordinator has to detokenize into it; otherwise it is + handed to the client as the same 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. @@ -192,11 +231,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] @@ -208,20 +244,28 @@ def handle_engine_reply(coordinator, sender_identity, payload): if idx is not None: 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 coordinator.prefix_caching_eviction_policy == PrefixCachingEvictionPolicy.REF_ZERO: + # Under REF_ZERO the engine deregisters blocks once no request + # holds them, so mirror that here. Blocks still referenced by + # another in-flight request keep a nonzero count and survive. + coordinator._release_rank_hashes(fid, assigned_rank) + + 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): +def handle_engine_reply_partial(coordinator, sender_identity, metadata, bodies): """Route incremental engine replies without releasing request routing state.""" if sender_identity not in coordinator.identities_of_data_parallel_ranks: assert ( @@ -229,29 +273,29 @@ def handle_engine_reply_partial(coordinator, sender_identity, payload): ), 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): +def handle_abort_request(coordinator, sender_identity, metadata, bodies): """Forward a client cancellation to the engine serving that request.""" 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 @@ -259,12 +303,12 @@ 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): +def handle_shutdown(coordinator, sender_identity, metadata, bodies): """Stop the coordinator event loop on request from a known client.""" if sender_identity not in coordinator.known_clients: logging.warning("Coordinator: ignoring signal from unknown client.") @@ -273,7 +317,7 @@ def handle_shutdown(coordinator, sender_identity, payload): @message_handler(Headers.DISCONNECT) -def handle_disconnect(coordinator, sender_identity, payload): +def handle_disconnect(coordinator, sender_identity, metadata, bodies): """Remove a disconnecting engine from the routing pool.""" 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 63b4176bf37..96641eb38b1 100644 --- a/megatron/core/inference/engines/dynamic_engine.py +++ b/megatron/core/inference/engines/dynamic_engine.py @@ -148,6 +148,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]]: @@ -679,6 +709,10 @@ async def start_listening_to_data_parallel_coordinator( "enable_prefix_caching": self.context.enable_prefix_caching, "prefix_caching_coordinator_policy": self.context.prefix_caching_coordinator_policy, "prefix_caching_routing_alpha": self.context.prefix_caching_routing_alpha, + "prefix_caching_cost_policy": self.context.prefix_caching_cost_policy, + "prefix_caching_load_beta": self.context.prefix_caching_load_beta, + "prefix_caching_eviction_policy": self.context.prefix_caching_eviction_policy, + "prefix_cache_ttl_seconds": self.context.prefix_cache_ttl_seconds, "schedule_output_path": coordinator_schedule_output_path, "hostname": hostname, }, @@ -1010,11 +1044,9 @@ def _handle_failed_request(self, request_id: int): # Send the reply immediately, because it may never get a chance to be sent again. if self.use_coordinator and self.is_mp_coordinator: - payload = msgpack.packb( - [Headers.ENGINE_REPLY.value, [request_entry.record.merge().serialize()]], - use_bin_type=True, + self.socket_for_receiving_requests.send_multipart( + _engine_reply_frames([request_entry.record.merge().serialize()]) ) - self.socket_for_receiving_requests.send(payload) elif not self.use_coordinator: if request.prompt is None: request.prompt = self.controller.tokenizer.detokenize( @@ -2228,9 +2260,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) @@ -2335,11 +2374,9 @@ async def async_bookkeep( ] if records_to_send: nvtx_range_push("coordinator_communication") - payload = msgpack.packb( - [Headers.ENGINE_REPLY.value, [r.merge().serialize() for r in records_to_send]], - use_bin_type=True, + self.socket_for_receiving_requests.send_multipart( + _engine_reply_frames([r.merge().serialize() for r in records_to_send]) ) - self.socket_for_receiving_requests.send(payload) nvtx_range_pop("coordinator_communication") # Stream newly generated tokens for active requests. Finished @@ -2647,21 +2684,34 @@ def schedule_requests(self) -> int: """ nvtx_range_push("drain_zmq_socket") + # Each message is a list of frames: metadata first, then any payload + # bodies. The MP broadcast flattens them and carries a manifest of + # per-message frame counts so peer ranks can rebuild the boundaries + # without copying any payload. all_messages = [] if self.is_mp_coordinator: 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 + manifest = msgpack.packb([len(m) for m in all_messages], use_bin_type=True) self.model_parallel_publisher_socket.send_multipart( - [bytes([Headers.TP_BROADCAST.value])] + all_messages + [bytes([Headers.TP_BROADCAST.value]), manifest] + + [frame for message in all_messages for frame in message] ) else: frames = self.model_parallel_subscriber_socket.recv_multipart() - all_messages = frames[1:] + frame_counts = msgpack.unpackb(frames[1], raw=False) + flat = frames[2:] + offset = 0 + for count in frame_counts: + all_messages.append(flat[offset : offset + count]) + offset += count nvtx_range_pop("drain_zmq_socket") @@ -2669,10 +2719,13 @@ 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) @@ -2727,7 +2780,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 84cd763844b..e50968b4186 100644 --- a/megatron/core/inference/inference_client.py +++ b/megatron/core/inference/inference_client.py @@ -89,7 +89,10 @@ def __init__(self, inference_coordinator_address: str, deserialize: bool = False self.aborted_request_ids: set[int] = set() def add_request( - self, prompt: Union[str, List[int]], sampling_params: SamplingParams + self, + prompt: Union[str, List[int]], + sampling_params: SamplingParams, + block_hashes: Optional[List[int]] = None, ) -> asyncio.Future: """ Submits a new inference request to the coordinator. @@ -111,14 +114,39 @@ 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()] - payload_serialized = msgpack.packb(payload, use_bin_type=True) - self.socket.send(payload_serialized) + self.socket.send_multipart( + self._submit_frames(request_id, prompt, sampling_params, block_hashes) + ) assert request_id not in self.completion_futures self.completion_futures[request_id] = asyncio.get_running_loop().create_future() self.request_submission_times[request_id] = time.perf_counter() return self.completion_futures[request_id] + def _submit_frames(self, request_id, prompt, sampling_params, block_hashes=None): + """Frame a submission as [metadata, prompt] (+ [block_hashes]). + + The prompt travels in its own frame so the coordinator can route the + request without decoding it -- at long prompts that decode dominates the + coordinator's per-request cost, and it is a single serial loop shared by + every data parallel rank. + + Block hashes, when the caller computed them, ride in a third frame for the + same reason: prefix-affinity routing needs them, and having the client + hash the tokens it already holds keeps that work off the coordinator. The + frame is omitted entirely when there are none, so a coordinator that does + not look for it is unaffected. + """ + frames = [ + msgpack.packb( + [Headers.SUBMIT_REQUEST.value, request_id, sampling_params.serialize()], + use_bin_type=True, + ), + msgpack.packb(prompt, use_bin_type=True), + ] + if block_hashes: + frames.append(msgpack.packb(block_hashes, use_bin_type=True)) + return frames + def abort_request(self, request_id: int) -> None: """Cancel an in-flight request and close its local response stream.""" request_id = int(request_id) @@ -134,7 +162,10 @@ def abort_request(self, request_id: int) -> None: self.socket.send(msgpack.packb(payload, use_bin_type=True)) def add_request_streaming( - self, prompt: Union[str, List[int]], sampling_params: SamplingParams + self, + prompt: Union[str, List[int]], + sampling_params: SamplingParams, + block_hashes: Optional[List[int]] = None, ) -> AsyncStream[dict]: """Submit a streaming inference request. @@ -161,8 +192,9 @@ 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()] - self.socket.send(msgpack.packb(payload, use_bin_type=True)) + self.socket.send_multipart( + self._submit_frames(request_id, prompt, sampling_params, block_hashes) + ) stream = AsyncStream( request_id, functools.partial(self.abort_request, request_id), loop=self._loop ) @@ -185,13 +217,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 @@ -215,10 +250,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 @@ -238,7 +273,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/megatron/core/inference/sampling_params.py b/megatron/core/inference/sampling_params.py index 2f7d5bb4551..48a008aa694 100644 --- a/megatron/core/inference/sampling_params.py +++ b/megatron/core/inference/sampling_params.py @@ -38,6 +38,11 @@ class SamplingParams: # drops prompt_tokens before serializing the finished request, saving the ZMQ # transmission cost for long prompts. Opt in when the client needs them. return_prompt_tokens: bool = False + # When False, the DP coordinator skips detokenizing this request's output. The + # coordinator is a single process serving every DP rank, so per-request work there + # is a throughput ceiling; callers that can detokenize themselves (e.g. the HTTP + # frontend, which is replicated) should set this to False. + detokenize_generations: bool = True streaming: bool = False # Emit incremental ENGINE_REPLY_PARTIAL frames. streaming_interval: int = 1 # Minimum unsent tokens per ENGINE_REPLY_PARTIAL. diff --git a/megatron/core/inference/text_generation_server/dynamic_text_gen_server/endpoints/chat_completions.py b/megatron/core/inference/text_generation_server/dynamic_text_gen_server/endpoints/chat_completions.py index e81b5273e84..f6679a959a2 100644 --- a/megatron/core/inference/text_generation_server/dynamic_text_gen_server/endpoints/chat_completions.py +++ b/megatron/core/inference/text_generation_server/dynamic_text_gen_server/endpoints/chat_completions.py @@ -7,9 +7,19 @@ import traceback import uuid import warnings +from functools import partial -from megatron.core.inference.inference_request import unwrap_serialized_tensors +import torch + +from megatron.core.inference.inference_request import ( + compute_block_hashes_batched, + unwrap_serialized_tensors, +) +from megatron.core.inference.config import routes_on_prefix from megatron.core.inference.sampling_params import SamplingParams +from megatron.core.inference.text_generation_controllers.text_generation_controller import ( + TextGenerationController, +) from megatron.core.tokenizers.text.parsers import PARSER_MAPPING from ..incremental_detokenizer import HuggingFaceFastIncrementalDetokenizer @@ -355,6 +365,25 @@ def _replace_prefix_tokens( return previous_turn_token_ids + current_turn_additional_token_ids +def _apply_chat_template_sync( + tokenizer, messages, tools, chat_template_kwargs, add_generation_prompt=True +): + """Apply the chat template and coerce to `list[int]`, for use in a worker thread. + + The coercion runs here too: it walks every token, so leaving it on the event loop + would keep part of the stall this offload exists to remove. + """ + return _coerce_to_token_id_list( + tokenizer.apply_chat_template( + messages, + tokenize=True, + add_generation_prompt=add_generation_prompt, + tools=tools, + **chat_template_kwargs, + ) + ) + + def _coerce_to_token_id_list(result): """Convert the return value of `tokenizer.apply_chat_template` to `list[int]`. @@ -434,6 +463,8 @@ async def chat_completions(): client = current_app.config['client'] tokenizer = current_app.config['tokenizer'] parsers = current_app.config['parsers'] + block_size_tokens = current_app.config.get('block_size_tokens') + coordinator_policy = current_app.config.get('prefix_caching_coordinator_policy') req = await request.get_json() tools = req.get("tools", None) @@ -460,14 +491,15 @@ async def chat_completions(): hasattr(tokenizer, 'apply_chat_template') and getattr(tokenizer, "chat_template", None) is not None ): - prompt_tokens = _coerce_to_token_id_list( - tokenizer.apply_chat_template( + prompt_tokens = await asyncio.get_running_loop().run_in_executor( + current_app.config['tokenize_executor'], + partial( + _apply_chat_template_sync, + current_app.config['tokenize_tokenizer'], template_messages, - tokenize=True, - add_generation_prompt=True, - tools=template_tools, - **chat_template_kwargs, - ) + template_tools, + chat_template_kwargs, + ), ) if req.get("prevent_retokenization", True): @@ -508,13 +540,17 @@ async def chat_completions(): ] # Get the templated tokenization of just the previous generation - retokenized_previous_turn_token_ids = _coerce_to_token_id_list( - tokenizer.apply_chat_template( - messages_to_last_assistant_message, - tokenize=True, - add_generation_prompt=False, - tools=template_tools, - **chat_template_kwargs, + retokenized_previous_turn_token_ids = ( + await asyncio.get_running_loop().run_in_executor( + current_app.config['tokenize_executor'], + partial( + _apply_chat_template_sync, + current_app.config['tokenize_tokenizer'], + messages_to_last_assistant_message, + template_tools, + chat_template_kwargs, + add_generation_prompt=False, + ), ) ) @@ -607,11 +643,27 @@ async def chat_completions(): termination_id=-1 if ignore_eos else None, return_prompt_tokens=return_prompt_tokens, streaming_interval=int(_get_non_none(req, "streaming_interval", 1)), + # This frontend detokenizes its own output below. Keeping it off the + # coordinator matters because that is one process shared by all DP ranks. + detokenize_generations=False, ) except ValueError as e: return Response(f"Invalid sampling parameter: {e}", status=400) # --- 3. Send Requests to Engine --- + # Hash here rather than at the coordinator: the tokens are already in hand, + # frontends run many-to-one against a single serial coordinator loop, and + # hashing there would mean unpacking the prompt frame the split was + # introduced to avoid. Skipped unless the coordinator routes on prefix + # affinity, since otherwise nobody reads it and the frame is never sent. + block_hashes = ( + compute_block_hashes_batched( + torch.tensor(prompt_tokens, dtype=torch.int64), block_size_tokens + ) + if block_size_tokens and routes_on_prefix(coordinator_policy) + else None + ) + stream_requested = bool(req.get("stream", False)) if stream_requested: # Streaming currently supports only Hugging Face fast tokenizers. @@ -624,7 +676,8 @@ async def chat_completions(): return Response(str(error), status=400) streams = [ - client.add_request_streaming(prompt_tokens, sampling_params) for _ in range(n) + client.add_request_streaming(prompt_tokens, sampling_params, block_hashes) + for _ in range(n) ] include_usage = bool((req.get("stream_options") or {}).get("include_usage", False)) response = Response( @@ -641,7 +694,9 @@ async def chat_completions(): response.timeout = None return response - tasks = [client.add_request(prompt_tokens, sampling_params) for _ in range(n)] + tasks = [ + client.add_request(prompt_tokens, sampling_params, block_hashes) for _ in range(n) + ] if current_app.config['verbose']: start_time = time.perf_counter() @@ -705,7 +760,11 @@ async def chat_completions(): for result_item in batch_results: result = unwrap_serialized_tensors(result_item) - text_output = result["generated_text"] + text_output = TextGenerationController.detokenize( + tokenizer, + result["generated_tokens"], + remove_EOD=not sampling_params.detokenize_stop_sequence, + ) # The engine always reports prompt_length (for usage), but drops the # prompt_tokens tensor unless return_prompt_tokens was set. prompt_tokens_count = result.get("prompt_length") diff --git a/megatron/core/inference/text_generation_server/dynamic_text_gen_server/text_generation_server.py b/megatron/core/inference/text_generation_server/dynamic_text_gen_server/text_generation_server.py index af12bc5bbc1..68d65954ead 100644 --- a/megatron/core/inference/text_generation_server/dynamic_text_gen_server/text_generation_server.py +++ b/megatron/core/inference/text_generation_server/dynamic_text_gen_server/text_generation_server.py @@ -1,9 +1,11 @@ # Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. import asyncio +import copy import logging import multiprocessing as mp import socket +from concurrent.futures import ThreadPoolExecutor from contextlib import contextmanager from typing import List, Optional @@ -17,6 +19,7 @@ HAS_BACKEND = False import megatron.core.inference.text_generation_server.dynamic_text_gen_server.endpoints as endpoints +from megatron.core.inference.config import PrefixCachingCoordinatorPolicy from megatron.core.inference.inference_client import InferenceClient from megatron.core.utils import trace_async_exceptions @@ -24,7 +27,7 @@ # Global reference to manage the background server processes _SERVER_PROCESSES: List[mp.Process] = [] -_SHARED_SOCKET = None + @contextmanager @@ -47,8 +50,9 @@ async def _run_text_gen_server( server_port: int, parsers: Optional[List[str]] = None, verbose: bool = False, - fd: Optional[int] = None, hostname: Optional[str] = None, + block_size_tokens: Optional[int] = None, + prefix_caching_coordinator_policy: Optional[PrefixCachingCoordinatorPolicy] = None, ): """ Initializes and runs the async web server. Automatically starts and @@ -63,6 +67,10 @@ async def _run_text_gen_server( logger.info(f"Rank {rank}: InferenceClient connected.") try: + # Bind what the caller asked for -- None means every interface, which is + # not the single address gethostname() resolves to. The resolved name is + # for the log line only. + bind_host = hostname if hostname is None: try: hostname = socket.gethostname() @@ -80,6 +88,20 @@ async def _run_text_gen_server( app.config['tokenizer'] = tokenizer app.config['parsers'] = parsers app.config['verbose'] = verbose + # The frontend hashes the prompt it already holds so the coordinator does + # not have to on its single serial loop. The policy decides whether anyone + # reads the hashes; the block size is only the granularity. + app.config['block_size_tokens'] = block_size_tokens + app.config['prefix_caching_coordinator_policy'] = prefix_caching_coordinator_policy + + # Applying the chat template is synchronous and O(prompt); on the event loop it + # stalls every other request this replica owns, including delivery of responses + # that already finished. One worker is enough - the point is the yield, not + # throughput. The copy is required: HF tokenizers are not thread-safe. + app.config['tokenize_executor'] = ThreadPoolExecutor( + max_workers=1, thread_name_prefix="tokenize" + ) + app.config['tokenize_tokenizer'] = copy.deepcopy(tokenizer) # Register all blueprints from the 'endpoints' package for endpoint in endpoints.__all__: @@ -92,18 +114,20 @@ async def _run_text_gen_server( 2**14 ) # Allow many concurrent streams for HTTP/2 clients. - if fd is not None: - config.bind = [f"fd://{fd}"] - else: - config.bind = [f"{hostname}:{server_port}"] + # Held for this worker's lifetime; closing it would drop the listener. + own_socket = _bind_reuseport_socket(server_port, bind_host) + config.bind = [f"fd://{own_socket.fileno()}"] with temp_log_level(logging.INFO, logger): logger.info(f"Starting text generation server on http://{hostname}:{server_port}") logger.info(f"Using tokenizer: {type(tokenizer)}") logger.info(f"Using parsers: {parsers}") - # Quart is natively ASGI, so we can serve the app directly - await serve(app, config) + try: + # Quart is natively ASGI, so we can serve the app directly + await serve(app, config) + finally: + own_socket.close() finally: # Gracefully shut down the client when the server stops @@ -118,8 +142,9 @@ def _server_process_worker( server_port: int, parsers: Optional[List[str]] = None, verbose: bool = False, - fd: Optional[int] = None, hostname: Optional[str] = None, + block_size_tokens: Optional[int] = None, + prefix_caching_coordinator_policy: Optional[PrefixCachingCoordinatorPolicy] = None, ): """Synchronous worker function that sets up a new event loop for the separate process.""" loop = asyncio.new_event_loop() @@ -127,7 +152,15 @@ def _server_process_worker( try: loop.run_until_complete( _run_text_gen_server( - coordinator_addr, tokenizer, rank, server_port, parsers, verbose, fd, hostname + coordinator_addr, + tokenizer, + rank, + server_port, + parsers, + verbose, + hostname, + block_size_tokens, + prefix_caching_coordinator_policy, ) ) except KeyboardInterrupt: @@ -141,6 +174,38 @@ def _server_process_worker( loop.close() +def _bind_reuseport_socket(server_port: int, hostname: Optional[str]) -> socket.socket: + """Bind this worker's own socket on the shared port, with SO_REUSEPORT. + + Unlike inheriting one fd, this gives every worker its own accept queue and + lets the kernel hash each connection's 4-tuple across them. It is a large + improvement on sharing (measured 604x -> 2x spread at 32 replicas) but is + still hashing, not balancing: it cannot see that a replica is already busy, + so the spread is statistical rather than exact. + """ + sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) + # Required on every socket sharing the port; without it the second bind fails. + sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEPORT, 1) + sock.bind((hostname if hostname is not None else "0.0.0.0", server_port)) + sock.setblocking(False) + return sock + + +def _reserve_port(hostname: Optional[str]) -> int: + """Pick a free port for the replicas to bind individually. + + Replicas each bind the port themselves, so the parent cannot hold the socket + and hand out its fd; it binds only long enough to learn a free port. The gap + before the replicas bind is a small race with unrelated processes, which is + why an explicit port is preferred when one is available. + """ + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as probe: + probe.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) + probe.bind((hostname if hostname is not None else "0.0.0.0", 0)) + return probe.getsockname()[1] + + def start_text_gen_server( coordinator_addr: str, tokenizer, @@ -151,77 +216,100 @@ def start_text_gen_server( num_replicas: int = 4, hostname: Optional[str] = None, sock: Optional[socket.socket] = None, -): - """Start the text generation server.""" + block_size_tokens: Optional[int] = None, + prefix_caching_coordinator_policy: Optional[PrefixCachingCoordinatorPolicy] = None, +) -> Optional[str]: + """Start the text generation server. + + Every replica binds its own socket on ``server_port`` with SO_REUSEPORT, so + each gets its own accept queue and the kernel spreads connections across + them. Sharing one inherited socket does not balance -- replicas race to + accept from a single queue and whichever is already running keeps winning, + which concentrates most traffic on a handful of them as replica count grows. + + Call this on every rank that should host a frontend. Frontend work (chat + template, detokenize, parsers, JSON) is CPU-bound, so hosting on a single + rank confines it to that rank's CPU allocation and leaves the rest of the + job's cores unused. Each caller gets its own URL back; collecting them and + spreading requests over the result is the caller's business. + + Args: + server_port: Port to listen on. Overridden by ``sock`` when given; 0 + asks the OS to choose a free one. + sock: A socket the caller already bound, used only to fix the port. + Replicas bind that port themselves, so it is closed here rather than + shared with them. + + Returns: + The base URL this rank serves on, or None if the server was already + running. + """ global _SERVER_PROCESSES - global _SHARED_SOCKET if _SERVER_PROCESSES: logger.warning("Text gen server processes are already running.") - return + return None - # The caller may pass in a socket it has already bound ahead of time. if sock is not None: - bound_port = sock.getsockname()[1] - if bound_port == 0: + # Take the port and release the socket: replicas each bind their own with + # SO_REUSEPORT, which one shared socket cannot provide. + server_port = sock.getsockname()[1] + if server_port == 0: raise ValueError( "socket must be bound to a real port before being passed to start_text_gen_server" ) - _SHARED_SOCKET = sock - server_port = bound_port - _SHARED_SOCKET.setblocking(False) - else: - _SHARED_SOCKET = socket.socket(socket.AF_INET, socket.SOCK_STREAM) - _SHARED_SOCKET.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) - - if hasattr(socket, 'SO_REUSEPORT'): - try: - _SHARED_SOCKET.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEPORT, 1) - except OSError: - pass - - bind_address = hostname if hostname is not None else "0.0.0.0" - _SHARED_SOCKET.bind((bind_address, server_port)) - _SHARED_SOCKET.setblocking(False) - - _SHARED_SOCKET.set_inheritable(True) - fd = _SHARED_SOCKET.fileno() + sock.close() + elif server_port == 0: + server_port = _reserve_port(hostname) for i in range(num_replicas): p = mp.Process( target=_server_process_worker, - args=(coordinator_addr, tokenizer, rank, server_port, parsers, verbose, fd, hostname), + args=( + coordinator_addr, + tokenizer, + rank, + server_port, + parsers, + verbose, + hostname, + block_size_tokens, + prefix_caching_coordinator_policy, + ), daemon=True, ) p.start() _SERVER_PROCESSES.append(p) - logger.info(f"Started text gen frontend replica {i+1}/{num_replicas} (PID: {p.pid})") + logger.info( + f"Started text gen frontend replica {i+1}/{num_replicas} " + f"on port {server_port} (PID: {p.pid})" + ) + return f"http://{hostname or socket.gethostname()}:{server_port}" -def stop_text_gen_server(): - """Stop the text generation server.""" - global _SERVER_PROCESSES - global _SHARED_SOCKET - if not _SERVER_PROCESSES: +def _terminate(processes: List[mp.Process], what: str): + """Terminate a group of worker processes, escalating to kill if needed.""" + if not processes: return - - logger.info(f"Terminating {len(_SERVER_PROCESSES)} Text Gen frontend processes...") - - for p in _SERVER_PROCESSES: + logger.info(f"Terminating {len(processes)} {what} processes...") + for p in processes: if p.is_alive(): p.terminate() - - for p in _SERVER_PROCESSES: + for p in processes: p.join(timeout=3) if p.is_alive(): p.kill() p.join() - # Clean up the master socket - if _SHARED_SOCKET is not None: - _SHARED_SOCKET.close() - _SHARED_SOCKET = None +def stop_text_gen_server(): + """Stop this rank's frontend replica processes.""" + global _SERVER_PROCESSES + + if not _SERVER_PROCESSES: + return + + _terminate(_SERVER_PROCESSES, "Text Gen frontend") _SERVER_PROCESSES = [] logger.info("All text gen frontend processes terminated.") diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py index e54f8a22252..227e0593428 100644 --- a/megatron/training/arguments.py +++ b/megatron/training/arguments.py @@ -2088,22 +2088,41 @@ def _add_inference_args(parser): 'free pool when ref_count hits 0. "lru" keeps blocks ' 'cached and evicts via LRU only when space is needed.') group.add_argument('--inference-dynamic-batching-prefix-caching-coordinator-policy', - type=str, default='load_balanced', + type=str, default='longest_prefix', choices=['longest_prefix', 'first_prefix_block', 'load_balanced'], dest='inference_dynamic_batching_prefix_caching_coordinator_policy', - help='Coordinator routing policy for prefix caching. ' - '"load_balanced" (default) routes to the rank with the fewest ' - 'in-flight requests, ignoring prefix affinity. ' - '"first_prefix_block" routes based on the first block hash only. ' - '"longest_prefix" routes to the rank with the longest matching ' - 'prefix. "first_prefix_block" and "longest_prefix" both combine ' - 'prefix affinity with load balancing and fall back to ' - 'load-balanced routing when prefix caching is disabled or no ' - 'prefix match exists.') + help='Coordinator affinity signal for prefix caching; how it is ' + 'weighed against load is set by ' + '--inference-dynamic-batching-prefix-caching-cost-policy. ' + '"longest_prefix" (default) routes to the rank with the longest ' + 'matching prefix. "first_prefix_block" routes based on the first ' + 'block hash only. "load_balanced" routes to the rank with the ' + 'fewest in-flight requests, ignoring prefix affinity. All fall ' + 'back to load-balanced routing when prefix caching is disabled ' + 'or no prefix match exists.') + group.add_argument('--inference-dynamic-batching-prefix-caching-cost-policy', + type=str, default='relative_load_weighted', + choices=['relative_load_weighted', 'free_capacity_weighted'], + dest='inference_dynamic_batching_prefix_caching_cost_policy', + help='How prefix affinity is weighed against rank load. Applies ' + 'to both "longest_prefix" and "first_prefix_block". ' + '"relative_load_weighted" (default) scores ' + 'fraction - beta * (load - mean) / max(1, mean), so the load ' + 'penalty vanishes while ranks are balanced and only pulls toward ' + 'idle ranks as the fleet diverges. "free_capacity_weighted" ' + 'scores alpha * fraction + (1 - alpha) * free_capacity, fixing ' + 'the trade-off in absolute terms.') + group.add_argument('--inference-dynamic-batching-prefix-caching-load-beta', + type=float, default=1.0, + dest='inference_dynamic_batching_prefix_caching_load_beta', + help='Weight on the load penalty under the ' + '"relative_load_weighted" cost policy, in units of full cache ' + 'hits per 100%% above mean load. 0 disables the penalty (pure ' + 'affinity). Default: 1.0.') group.add_argument('--inference-dynamic-batching-prefix-caching-routing-alpha', type=float, default=0.5, dest='inference_dynamic_batching_prefix_caching_routing_alpha', - help='Weight for prefix-aware routing score: ' + help='Weight for the "free_capacity_weighted" cost policy: ' 'score = alpha * match + (1 - alpha) * normalized_load. ' 'Higher alpha favors prefix cache hits; lower alpha ' 'favors load balance. Default: 0.5.') diff --git a/megatron/training/config/inference_config.py b/megatron/training/config/inference_config.py index ee75517f41b..799f63bbd26 100644 --- a/megatron/training/config/inference_config.py +++ b/megatron/training/config/inference_config.py @@ -186,17 +186,33 @@ class InferenceSetupConfig: inference_dynamic_batching_prefix_caching_coordinator_policy: Literal[ "longest_prefix", "first_prefix_block", "load_balanced" - ] = "load_balanced" - """Coordinator routing policy for prefix caching. "load_balanced" (default) routes to the rank - with the fewest in-flight requests, ignoring prefix affinity. "first_prefix_block" routes based - on the first block hash only. "longest_prefix" routes to the rank with the longest matching - prefix. "first_prefix_block" and "longest_prefix" both combine prefix affinity with load - balancing and fall back to load-balanced routing when prefix caching is disabled or no prefix - match exists.""" + ] = "longest_prefix" + """Coordinator routing policy for prefix caching. Selects the affinity signal only; how that + signal is weighed against load is set by + --inference-dynamic-batching-prefix-caching-cost-policy. "longest_prefix" (default) routes to + the rank with the longest matching prefix. "first_prefix_block" routes based on the first block + hash only. "load_balanced" routes to the rank with the fewest in-flight requests, ignoring + prefix affinity. All fall back to load-balanced routing when prefix caching is disabled or no + prefix match exists.""" + + inference_dynamic_batching_prefix_caching_cost_policy: Literal[ + "relative_load_weighted", "free_capacity_weighted" + ] = "relative_load_weighted" + """How prefix affinity is weighed against rank load. Applies to both "longest_prefix" and + "first_prefix_block". "relative_load_weighted" (default) scores + fraction - beta * (load - mean) / max(1, mean), so the load penalty vanishes while ranks are + balanced and only pulls toward idle ranks as the fleet diverges. "free_capacity_weighted" + scores alpha * fraction + (1 - alpha) * free_capacity, fixing the trade-off in absolute + terms.""" + + inference_dynamic_batching_prefix_caching_load_beta: float = 1.0 + """Weight on the load penalty under the "relative_load_weighted" cost policy, in units of full + cache hits per 100% above mean load. 0 disables the penalty (pure affinity).""" inference_dynamic_batching_prefix_caching_routing_alpha: float = 0.5 - """Weight for prefix-aware routing score: score = alpha * match + (1 - alpha) * normalized_load. - Higher alpha favors prefix cache hits; lower alpha favors load balance.""" + """Weight for the "free_capacity_weighted" cost policy: score = alpha * match + + (1 - alpha) * normalized_load. Higher alpha favors prefix cache hits; lower alpha favors load + balance.""" inference_dynamic_batching_prefix_caching_mamba_gb: float | None = None """GPU memory budget (in GB) for the Mamba state cache used by prefix caching on hybrid models. @@ -297,6 +313,7 @@ def to_inference_config( KVCacheManagementMode, MambaInferenceStateConfig, PrefixCachingCoordinatorPolicy, + PrefixCachingCostPolicy, PrefixCachingEvictionPolicy, ) from megatron.core.utils import get_attr_wrapped_model @@ -370,6 +387,10 @@ def to_inference_config( prefix_caching_coordinator_policy=PrefixCachingCoordinatorPolicy( self.inference_dynamic_batching_prefix_caching_coordinator_policy ), + prefix_caching_cost_policy=PrefixCachingCostPolicy( + self.inference_dynamic_batching_prefix_caching_cost_policy + ), + prefix_caching_load_beta=self.inference_dynamic_batching_prefix_caching_load_beta, prefix_caching_routing_alpha=self.inference_dynamic_batching_prefix_caching_routing_alpha, prefix_caching_mamba_gb=self.inference_dynamic_batching_prefix_caching_mamba_gb, metrics_writer=metrics_writer, diff --git a/tests/unit_tests/inference/engines/test_dynamic_engine.py b/tests/unit_tests/inference/engines/test_dynamic_engine.py index af91409ef7b..b477a352f8f 100644 --- a/tests/unit_tests/inference/engines/test_dynamic_engine.py +++ b/tests/unit_tests/inference/engines/test_dynamic_engine.py @@ -13,6 +13,7 @@ from typing import Dict, List, Optional, Tuple from unittest import mock +import msgpack import pytest import torch from tqdm import tqdm @@ -34,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, @@ -765,7 +767,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} @@ -784,13 +791,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 3d95f158c32..e7914e0691e 100644 --- a/tests/unit_tests/inference/test_data_parallel_inference_coordinator.py +++ b/tests/unit_tests/inference/test_data_parallel_inference_coordinator.py @@ -25,6 +25,7 @@ DynamicInferenceEngine, EngineState, RequestEntry, + _engine_reply_frames, ) from megatron.core.inference.headers import Headers from megatron.core.inference.inference_client import InferenceClient @@ -190,13 +191,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] @@ -869,10 +868,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() @@ -884,7 +892,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 @@ -892,7 +900,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 e28bb8b8c4c..3d121ea615b 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,6 +122,8 @@ 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 5d587ffae2f..935b97c785a 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()) diff --git a/tools/run_dynamic_text_generation_server.py b/tools/run_dynamic_text_generation_server.py index a33b066e8ea..8aca26ba537 100644 --- a/tools/run_dynamic_text_generation_server.py +++ b/tools/run_dynamic_text_generation_server.py @@ -2,6 +2,7 @@ import argparse import asyncio +import logging import torch @@ -10,7 +11,7 @@ start_text_gen_server, stop_text_gen_server, ) -from megatron.core.utils import configure_nvtx_profiling, trace_async_exceptions +from megatron.core.utils import configure_nvtx_profiling, get_pg_size, trace_async_exceptions from megatron.inference.utils import add_inference_args, get_dynamic_inference_engine from megatron.post_training.arguments import add_modelopt_args from megatron.training import get_args @@ -30,6 +31,23 @@ def add_text_generation_server_args(parser: argparse.ArgumentParser): parser.add_argument( "--parsers", type=str, nargs="+", default=[], help="Parsers to use for parsing the response" ) + parser.add_argument( + "--frontend-replicas", type=int, default=-1, + help="Number of HTTP frontend processes spawned per hosting rank. " + "-1 (default) uses max(data parallel size, 4), or a flat 4 with " + "--frontend-on-all-ranks, where capacity already scales with the " + "number of ranks hosting a frontend.", + ) + parser.add_argument( + "--frontend-on-all-ranks", action="store_true", + help="Run HTTP frontends on every rank instead of only rank 0, and " + "return every rank's URL for the caller to spread requests over. " + "Frontend work (chat template, detokenize, parsers, JSON) is " + "CPU-bound and otherwise confined to the hosting rank's CPU " + "allocation, which leaves the rest of the job's cores unused. " + "Ranks still share one DP coordinator; only the HTTP tier is " + "replicated.", + ) return parser @@ -54,25 +72,62 @@ async def run_text_generation_server( hostname=hostname, ) + num_replicas = args.frontend_replicas + if num_replicas < 0: + if args.frontend_on_all_ranks: + # Capacity now scales with the number of ranks, so the per-rank + # replica count stays flat rather than tracking DP size on top of it. + num_replicas = 4 + else: + # Each replica is a single event loop, so frontend capacity has to scale with + # the number of engines it feeds. The floor of 4 preserves the previous default + # for small deployments. + num_replicas = max(get_pg_size(engine.pg_collection.dp), 4) + if rank == 0: + logging.info("Starting %d HTTP frontend replica(s) per hosting rank.", num_replicas) + + if args.frontend_on_all_ranks: + # Only the DP coordinator rank learns the coordinator's address: the + # engine broadcasts it over the DP group, which is a singleton when data + # parallel size is 1. Every rank needs it here, since every rank's + # frontend opens its own client. + address = [coordinator_addr] + torch.distributed.broadcast_object_list(address, src=0) + coordinator_addr = address[0] + assert coordinator_addr is not None, "no rank published a DP coordinator address" + try: - if rank == 0: - start_text_gen_server( + url = None + if args.frontend_on_all_ranks or rank == 0: + url = start_text_gen_server( coordinator_addr=coordinator_addr, tokenizer=engine.controller.tokenizer, parsers=args.parsers, rank=rank, - server_port=server_port, + server_port=0 if args.frontend_on_all_ranks else server_port, verbose=args.inference_text_gen_server_logging, + num_replicas=num_replicas, hostname=hostname, ) + if args.frontend_on_all_ranks: + # Unlike callers that already collect a URL per worker, this entry + # point has to gather them itself before it can report the set. + urls = [None] * torch.distributed.get_world_size() + torch.distributed.all_gather_object(urls, url) + if rank == 0: + for entry in [u for u in urls if u]: + logging.info("Frontend: %s", entry) + elif rank == 0: + logging.info("Frontend: %s", url) + # Await the engine loop directly since the server is running in a separate process await engine.engine_loop_task finally: - # Guarantee that the separate process is terminated when the engine loop stops or is interrupted - if rank == 0: - stop_text_gen_server() + # Guarantee that the separate processes are terminated when the engine loop + # stops or is interrupted. Every rank may now own frontend processes. + stop_text_gen_server() if __name__ == "__main__":