From ab15dce3c887b8934ac4717a5ca0a3a98c63412f Mon Sep 17 00:00:00 2001
From: Siddharth Singh <sidsingh@nvidia.com>
Date: Mon, 3 Aug 2026 13:42:10 -0700
Subject: [PATCH 1/7] move tokenizer to a separate thread, so that it does not
 block the asyncio event loop

---
 .../endpoints/chat_completions.py             | 53 ++++++++++++++-----
 .../text_generation_server.py                 | 11 ++++
 tools/run_dynamic_text_generation_server.py   | 18 +++++++
 3 files changed, 68 insertions(+), 14 deletions(-)

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..a039c6cf3de 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,6 +7,7 @@
 import traceback
 import uuid
 import warnings
+from functools import partial
 
 from megatron.core.inference.inference_request import unwrap_serialized_tensors
 from megatron.core.inference.sampling_params import SamplingParams
@@ -355,6 +356,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]`.
 
@@ -460,14 +480,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 +529,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,
+                                ),
                             )
                         )
 
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..9e7a8da716e 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
 
@@ -81,6 +83,15 @@ async def _run_text_gen_server(
         app.config['parsers'] = parsers
         app.config['verbose'] = verbose
 
+        # 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__:
             app.register_blueprint(endpoint)
diff --git a/tools/run_dynamic_text_generation_server.py b/tools/run_dynamic_text_generation_server.py
index a33b066e8ea..a07ff444d5c 100644
--- a/tools/run_dynamic_text_generation_server.py
+++ b/tools/run_dynamic_text_generation_server.py
@@ -2,9 +2,11 @@
 
 import argparse
 import asyncio
+import logging
 
 import torch
 
+from megatron.core import parallel_state
 from megatron.core.inference.engines import DynamicInferenceEngine
 from megatron.core.inference.text_generation_server.dynamic_text_gen_server import (
     start_text_gen_server,
@@ -30,6 +32,12 @@ 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 on the primary rank. "
+             "-1 (default) uses max(data parallel size, 4), so frontend capacity "
+             "scales with the number of engines.",
+    )
     return parser
 
 
@@ -54,6 +62,15 @@ async def run_text_generation_server(
         hostname=hostname,
     )
 
+    num_replicas = args.frontend_replicas
+    if num_replicas < 0:
+        # 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(parallel_state.get_data_parallel_world_size(), 4)
+    if rank == 0:
+        logging.info("Starting %d HTTP frontend replica(s).", num_replicas)
+
     try:
         if rank == 0:
             start_text_gen_server(
@@ -63,6 +80,7 @@ async def run_text_generation_server(
                 rank=rank,
                 server_port=server_port,
                 verbose=args.inference_text_gen_server_logging,
+                num_replicas=num_replicas,
                 hostname=hostname,
             )
 

From 77c62ff6f23d61d1d8d3a3c9fc4b05d27dccfa5b Mon Sep 17 00:00:00 2001
From: Siddharth Singh <sidsingh@nvidia.com>
Date: Mon, 3 Aug 2026 14:24:09 -0700
Subject: [PATCH 2/7] move detokenization from the coordinator to the HTTP
 server

---
 .../coordinator.py                                   |  3 +++
 megatron/core/inference/sampling_params.py           |  5 +++++
 .../endpoints/chat_completions.py                    | 12 +++++++++++-
 3 files changed, 19 insertions(+), 1 deletion(-)

diff --git a/megatron/core/inference/data_parallel_inference_coordinator/coordinator.py b/megatron/core/inference/data_parallel_inference_coordinator/coordinator.py
index 67eb1bb5829..ac590bdb8ca 100644
--- a/megatron/core/inference/data_parallel_inference_coordinator/coordinator.py
+++ b/megatron/core/inference/data_parallel_inference_coordinator/coordinator.py
@@ -464,6 +464,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
         )
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 a039c6cf3de..35ae357d00d 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
@@ -11,6 +11,9 @@
 
 from megatron.core.inference.inference_request import unwrap_serialized_tensors
 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
@@ -632,6 +635,9 @@ 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)
@@ -730,7 +736,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")

From 52094e5e29673deb2bbb3a3f83735f88a4e1a2d1 Mon Sep 17 00:00:00 2001
From: Siddharth Singh <sidsingh@nvidia.com>
Date: Thu, 6 Aug 2026 13:27:15 -0700
Subject: [PATCH 3/7] take DP size from engine

---
 tools/run_dynamic_text_generation_server.py | 5 ++---
 1 file changed, 2 insertions(+), 3 deletions(-)

diff --git a/tools/run_dynamic_text_generation_server.py b/tools/run_dynamic_text_generation_server.py
index a07ff444d5c..bd169818ef3 100644
--- a/tools/run_dynamic_text_generation_server.py
+++ b/tools/run_dynamic_text_generation_server.py
@@ -6,13 +6,12 @@
 
 import torch
 
-from megatron.core import parallel_state
 from megatron.core.inference.engines import DynamicInferenceEngine
 from megatron.core.inference.text_generation_server.dynamic_text_gen_server import (
     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
@@ -67,7 +66,7 @@ async def run_text_generation_server(
         # 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(parallel_state.get_data_parallel_world_size(), 4)
+        num_replicas = max(get_pg_size(engine.pg_collection.dp), 4)
     if rank == 0:
         logging.info("Starting %d HTTP frontend replica(s).", num_replicas)
 

From 929bf14d80b45c80025114157ec41317968ac7bf Mon Sep 17 00:00:00 2001
From: Siddharth Singh <sidsingh@aws-cmh-slurm-1-login-01.cm.cluster>
Date: Thu, 13 Aug 2026 23:23:54 -0700
Subject: [PATCH 4/7] Give each frontend replica its own SO_REUSEPORT socket

The replicas shared one listening socket: the parent bound it and passed the
same fd to every forked worker, so all of them accepted from a single queue.
That does not balance. Whichever worker is already running tends to win the
wakeup, and it keeps winning, because an event loop with work in flight polls
more often than one blocked in accept. SO_REUSEPORT was set on that socket but
was inert -- the kernel only load-balances when several sockets are bound to
the port and it can hash a connection's 4-tuple to choose between them.

Measured with 32 replicas and a fresh connection per request, ~90% of traffic
landed on 5 of them, 20 replicas served exactly one request each, and
throughput was 3.7x lower than the same server under a pooled client that
opened its connections up front. Load made it worse rather than averaging it
out: at 2048 requests the busiest replica took 604x the quietest.

Each replica now binds its own socket on the shared port, so every one gets its
own accept queue. Spread became max/min 1.5x with all replicas serving, and
throughput 2.5x on the fresh-connection path. A pooled client is unaffected in
steady state, which is the point: how well the frontend spreads no longer
depends on connection behaviour the server cannot observe.

start_text_gen_server now returns the base URL it is serving on. Callers that
start a frontend on more than one rank need the addresses to spread requests
over; previously they had to reconstruct them. The signature is otherwise
unchanged, including sock, which still fixes the port -- it is closed rather
than shared, since replicas bind their own.

tools/run_dynamic_text_generation_server.py gains --frontend-on-all-ranks,
which hosts a frontend on every rank and gathers the URLs. Frontend work is
CPU-bound and otherwise confined to one rank's CPU allocation while the rest of
the job's cores go unused.

Signed-off-by: Siddharth Singh <sidsingh@nvidia.com>
(cherry picked from commit 1001ca96e8e11c1faf79c53ed5a1b63aef6ecc25)
---
 .../text_generation_server.py                 | 153 ++++++++++++------
 tools/run_dynamic_text_generation_server.py   |  66 ++++++--
 2 files changed, 152 insertions(+), 67 deletions(-)

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 9e7a8da716e..65b916f60b3 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
@@ -26,7 +26,7 @@
 
 # Global reference to manage the background server processes
 _SERVER_PROCESSES: List[mp.Process] = []
-_SHARED_SOCKET = None
+
 
 
 @contextmanager
@@ -49,7 +49,6 @@ 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,
 ):
     """
@@ -65,6 +64,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()
@@ -103,18 +106,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
@@ -129,7 +134,6 @@ def _server_process_worker(
     server_port: int,
     parsers: Optional[List[str]] = None,
     verbose: bool = False,
-    fd: Optional[int] = None,
     hostname: Optional[str] = None,
 ):
     """Synchronous worker function that sets up a new event loop for the separate process."""
@@ -138,7 +142,7 @@ 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
             )
         )
     except KeyboardInterrupt:
@@ -152,6 +156,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,
@@ -162,77 +198,88 @@ def start_text_gen_server(
     num_replicas: int = 4,
     hostname: Optional[str] = None,
     sock: Optional[socket.socket] = None,
-):
-    """Start the text generation server."""
+) -> 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),
             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/tools/run_dynamic_text_generation_server.py b/tools/run_dynamic_text_generation_server.py
index bd169818ef3..8aca26ba537 100644
--- a/tools/run_dynamic_text_generation_server.py
+++ b/tools/run_dynamic_text_generation_server.py
@@ -33,9 +33,20 @@ def add_text_generation_server_args(parser: argparse.ArgumentParser):
     )
     parser.add_argument(
         "--frontend-replicas", type=int, default=-1,
-        help="Number of HTTP frontend processes spawned on the primary rank. "
-             "-1 (default) uses max(data parallel size, 4), so frontend capacity "
-             "scales with the number of engines.",
+        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
 
@@ -63,33 +74,60 @@ async def run_text_generation_server(
 
     num_replicas = args.frontend_replicas
     if num_replicas < 0:
-        # 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 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).", num_replicas)
+        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__":

From 82caf5fe96356b2a0c4b40733b10cc93f4cdb23f Mon Sep 17 00:00:00 2001
From: Siddharth Singh <sidsingh@aws-cmh-slurm-1-login-01.cm.cluster>
Date: Wed, 12 Aug 2026 14:15:23 -0700
Subject: [PATCH 5/7] modify coordinator wire format - split into two chunks -
 first chunk is metadata of constant size. This is the only thing that the
 coordinator needs to unpack/read and pack

---
 .../coordinator.py                            |  35 +++--
 .../handlers.py                               | 138 ++++++++++--------
 .../core/inference/engines/dynamic_engine.py  |  81 ++++++++--
 megatron/core/inference/inference_client.py   |  36 +++--
 .../inference/engines/test_dynamic_engine.py  |  13 +-
 ...est_data_parallel_inference_coordinator.py |  28 ++--
 ...test_dynamic_prefix_caching_coordinator.py |  82 +++++++++--
 .../inference/test_inference_client.py        |  31 ++--
 .../test_inference_client_streaming.py        |  61 ++++----
 9 files changed, 338 insertions(+), 167 deletions(-)

diff --git a/megatron/core/inference/data_parallel_inference_coordinator/coordinator.py b/megatron/core/inference/data_parallel_inference_coordinator/coordinator.py
index ac590bdb8ca..88a2fe8015e 100644
--- a/megatron/core/inference/data_parallel_inference_coordinator/coordinator.py
+++ b/megatron/core/inference/data_parallel_inference_coordinator/coordinator.py
@@ -303,14 +303,18 @@ def _remove_engine(self, identity):
             len(self.identities_of_data_parallel_ranks),
         )
 
-    def _send_to_engine(self, identity, payload):
-        """Send payload to an engine, removing it from the pool if unreachable.
+    def _send_to_engine(self, identity, frames):
+        """Send a message to an engine, removing it from the pool if unreachable.
+
+        Args:
+            identity: ZMQ identity of the target engine.
+            frames (list): Raw frames to send, metadata frame first.
 
         Returns:
             True if the send succeeded, False if the engine was unreachable and removed.
         """
         try:
-            self.router_socket.send_multipart([identity, payload])
+            self.router_socket.send_multipart([identity, *frames])
             return True
         except zmq.error.ZMQError as e:
             if e.errno == zmq.EHOSTUNREACH:
@@ -322,19 +326,21 @@ def _broadcast_to_engines(self, payload):
         """Send a deserialized payload to every connected data parallel rank."""
         serialized = msgpack.packb(payload, use_bin_type=True)
         for data_parallel_rank_id in list(self.identities_of_data_parallel_ranks):
-            self._send_to_engine(data_parallel_rank_id, serialized)
+            self._send_to_engine(data_parallel_rank_id, [serialized])
 
     def compute_request_hashes(self, prompt):
         """Compute block hashes for a prompt on CPU.
 
+        Callers decide whether hashes are wanted at all: computing them requires
+        the decoded prompt, and decoding it is the cost the caller is usually
+        trying to avoid. See handle_submit_request.
+
         Args:
             prompt: Either a string (to be tokenized) or a list of token IDs.
 
         Returns:
-            List of integer block hashes, or empty list if prefix caching is disabled.
+            List of integer block hashes, one per complete block.
         """
-        if not self.enable_prefix_caching or self.block_size_tokens is None:
-            return []
         if isinstance(prompt, str):
             tokens = self.tokenizer.tokenize(prompt)
         else:
@@ -431,20 +437,25 @@ def start(self):
         """
         # Todo [Siddharth]: Make this more robust to handle invalid messages.
         while True:
-            sender_identity, serialized_payload = self.router_socket.recv_multipart()
+            # Messages are one or more frames. frames[0] is the metadata frame:
+            # a header plus whatever the coordinator needs to route the message.
+            # Any later frames are opaque payload bodies, forwarded without ever
+            # being decoded here -- that is what keeps this loop's cost
+            # independent of prompt length.
+            sender_identity, *frames = self.router_socket.recv_multipart()
 
             # An empty payload is a data parallel rank (re-)registering itself.
-            if serialized_payload == b"":
+            if frames[0] == b"":
                 self._handle_rank_registration(sender_identity)
                 continue
 
-            deserialized_payload = msgpack.unpackb(serialized_payload, raw=False)
-            header = Headers(deserialized_payload[0])
+            metadata = msgpack.unpackb(frames[0], raw=False)
+            header = Headers(metadata[0])
 
             handler = self._handlers.get(header)
             if handler is None:
                 raise UnknownHeaderError(header)
-            if handler(self, sender_identity, deserialized_payload):
+            if handler(self, sender_identity, metadata, frames[1:]):
                 break
 
     def _handle_rank_registration(self, sender_identity):
diff --git a/megatron/core/inference/data_parallel_inference_coordinator/handlers.py b/megatron/core/inference/data_parallel_inference_coordinator/handlers.py
index 2c1712a1de3..804e63d0641 100644
--- a/megatron/core/inference/data_parallel_inference_coordinator/handlers.py
+++ b/megatron/core/inference/data_parallel_inference_coordinator/handlers.py
@@ -8,15 +8,23 @@
 is supported simply by adding a decorated function here; the coordinator's event
 loop never changes.
 
-Handlers have the signature ``(coordinator, sender_identity, payload) -> bool | None``
-where ``payload`` is the already-deserialized message. Returning a truthy value
-signals the coordinator's event loop to stop.
+Handlers have the signature
+``(coordinator, sender_identity, metadata, bodies) -> bool | None``.
+
+``metadata`` is the decoded contents of the message's first frame: a header
+followed by whatever the coordinator needs in order to route the message. It is
+small by construction and does not grow with prompt length.
+
+``bodies`` is the list of remaining raw frames, still packed. These carry the
+prompt (inbound) or the finished request (outbound) -- the bulk of the bytes.
+The coordinator forwards them as opaque frames and only decodes one when it has
+to mutate it, which keeps its per-request cost flat in prompt length.
+
+Returning a truthy value signals the event loop to stop.
 """
 
 import logging
 
-import torch
-
 from megatron.core.inference.config import PrefixCachingCoordinatorPolicy
 from megatron.core.inference.headers import Headers
 
@@ -52,7 +60,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 +73,23 @@ 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]`` and
+    ``bodies[0]`` is the packed prompt, which is forwarded to the engine as-is
+    unless the routing policy needs to hash it.
+
     Returns True (stopping the loop) if no engines are reachable.
     """
-    # ToDo [Siddharth]: We might want to tokenize the prompt on the
-    # assigned data parallel rank for this process instead
-    # of the coordinator.
-
     # Message from a known client
     if sender_identity not in coordinator.known_clients:
         logging.info(f"Received message from unknown client {sender_identity}. Ignoring.")
         return
-    # this is a message from a client.
-    # route it to a data parallel rank
-    client_request_id, prompt, sampling_params = payload[1:]
+
+    _, client_request_id, sampling_params = metadata
+    prompt_frame = bodies[0]
+
     # map client request_id to server request_id
     # necessary because multiple clients might have the same request_id.
     request_id = coordinator.next_request_id
@@ -89,19 +98,26 @@ def handle_submit_request(coordinator, sender_identity, payload):
     coordinator.request_id_to_client_request_id[request_id] = client_request_id
     coordinator.client_request_to_request_id[(sender_identity, client_request_id)] = request_id
 
-    # Serialize prompt.
-    if isinstance(prompt, (str, list)):
-        pass
-    elif isinstance(prompt, torch.Tensor):
-        prompt = prompt.tolist()
-    else:
-        raise Exception("specialize for <%s> prompt." % type(prompt).__name__)
-
-    engine_payload = msgpack.packb(
-        [Headers.SUBMIT_REQUEST.value, request_id, prompt, sampling_params], use_bin_type=True
+    # Rebuilding the metadata frame is cheap: it holds no prompt tokens.
+    engine_metadata = msgpack.packb(
+        [Headers.SUBMIT_REQUEST.value, request_id, sampling_params], use_bin_type=True
     )
 
-    request_hashes = coordinator.compute_request_hashes(prompt)
+    # Only prefix-affinity routing consults the block hashes. When nothing will
+    # look at them, skip hashing *and* the prompt decode it needs -- avoiding
+    # that decode is the reason the prompt travels in its own frame.
+    if (
+        coordinator.enable_prefix_caching
+        and coordinator.block_size_tokens is not None
+        and coordinator.prefix_caching_coordinator_policy
+        != PrefixCachingCoordinatorPolicy.LOAD_BALANCED
+    ):
+        request_hashes = coordinator.compute_request_hashes(
+            msgpack.unpackb(prompt_frame, raw=False)
+        )
+    else:
+        request_hashes = []
+
     if (
         coordinator.prefix_caching_coordinator_policy
         == PrefixCachingCoordinatorPolicy.FIRST_PREFIX_BLOCK
@@ -111,7 +127,7 @@ def handle_submit_request(coordinator, sender_identity, payload):
     # Account for the fact that some engines may have died.
     for _ in range(len(coordinator.identities_of_data_parallel_ranks)):
         next_identity = coordinator.get_best_data_parallel_rank(request_hashes)
-        if coordinator._send_to_engine(next_identity, engine_payload):
+        if coordinator._send_to_engine(next_identity, [engine_metadata, prompt_frame]):
             break
     else:
         # If all engines have died, we are in an abnormal state, and must exit cleanly.
@@ -143,13 +159,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 +175,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 +185,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 +194,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 +214,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]
@@ -209,19 +228,22 @@ def handle_engine_reply(coordinator, sender_identity, payload):
                 assert coordinator._pending_counts[idx] >= 1
                 coordinator._pending_counts[idx] -= 1
 
-        coordinator.router_socket.send_multipart(
-            [
-                client_identity,
-                msgpack.packb(
-                    [Headers.ENGINE_REPLY.value, client_request_id, finished_request],
-                    use_bin_type=True,
-                ),
-            ]
+        if needs_detokenize:
+            # Detokenizing writes generated_text into the reply, so this one has
+            # to be decoded and re-encoded. Clients that detokenize for
+            # themselves (the OpenAI frontend does) never take this path.
+            finished_request = msgpack.unpackb(body, raw=False)
+            coordinator.detokenize(finished_request)
+            body = msgpack.packb(finished_request, use_bin_type=True)
+
+        reply_metadata = msgpack.packb(
+            [Headers.ENGINE_REPLY.value, client_request_id], use_bin_type=True
         )
+        coordinator.router_socket.send_multipart([client_identity, reply_metadata, body])
 
 
 @message_handler(Headers.ENGINE_REPLY_PARTIAL)
-def handle_engine_reply_partial(coordinator, sender_identity, payload):
+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 +251,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 +281,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 +295,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..346147a84ba 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]]:
@@ -1010,11 +1040,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 +2256,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 +2370,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 +2680,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 +2715,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 +2776,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..c6aac68ae2f 100644
--- a/megatron/core/inference/inference_client.py
+++ b/megatron/core/inference/inference_client.py
@@ -111,14 +111,28 @@ 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))
         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):
+        """Frame a submission as [metadata, prompt].
+
+        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.
+        """
+        return [
+            msgpack.packb(
+                [Headers.SUBMIT_REQUEST.value, request_id, sampling_params.serialize()],
+                use_bin_type=True,
+            ),
+            msgpack.packb(prompt, use_bin_type=True),
+        ]
+
     def abort_request(self, request_id: int) -> None:
         """Cancel an in-flight request and close its local response stream."""
         request_id = int(request_id)
@@ -161,8 +175,7 @@ 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))
         stream = AsyncStream(
             request_id, functools.partial(self.abort_request, request_id), loop=self._loop
         )
@@ -185,13 +198,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 +231,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 +254,7 @@ def _connect_with_inference_coordinator(self, timeout_seconds: Optional[float] =
             timeout=max(0, int(timeout_seconds * 1000))
         ):
             raise TimeoutError("Timed out connecting to the Megatron inference coordinator")
-        reply = msgpack.unpackb(self.socket.recv(), raw=False)
+        reply = msgpack.unpackb(self.socket.recv_multipart()[0], raw=False)
         assert Headers(reply[0]) == Headers.CONNECT_ACK
 
     def start(
diff --git a/tests/unit_tests/inference/engines/test_dynamic_engine.py b/tests/unit_tests/inference/engines/test_dynamic_engine.py
index 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())

From e8ca92c2048974e3f49f6bf40fd39ccd66a5bf8c Mon Sep 17 00:00:00 2001
From: Siddharth Singh <sidsingh@nvidia.com>
Date: Mon, 17 Aug 2026 11:14:54 -0700
Subject: [PATCH 6/7] feat(inference): prefix-affinity routing in the DP
 coordinator

Route each request to the rank that already holds the most of its prompt
rather than to the least loaded one. Frontends compute per-block prompt
hashes and ship them alongside the request; the coordinator scores ranks by
(prefill blocks still to compute) x (1 + load), so idle ranks fill first and
a rank holding the prefix wins thereafter.

Hashing happens on the frontend, not the coordinator: the tokens are already
in hand there, frontends run many-to-one against a single serial coordinator
loop, and hashing at the coordinator would mean unpacking the prompt frame
the request/prompt frame split exists to avoid.

Whether to hash follows from the coordinator's routing policy, which is the
only component that knows if anyone will read the hashes. `routes_on_prefix`
lives beside the policy enum so a new prefix-aware policy is a one-line
change. `block_size_tokens` is granularity only and is always passed; it must
match the engine's KV block size or the hashes name blocks the engine never
cached.

Coordinator-side cache tracking assumes an engine still holds a block for
`prefix_cache_ttl_seconds` under the LRU eviction policy, since engine-side
eviction is not observable from the coordinator.

On a 16-engine nanoV3.5 SWE-RL run this moved the prefill skip rate from
59-63% to 97.7% and the prefill share of step time from 42-62% to 10.3%.

Signed-off-by: Siddharth Singh <sidsingh@nvidia.com>
(cherry picked from commit fcb23e1ee84045775eb068d0e30c4b299adff0d3)
---
 megatron/core/inference/config.py             |  28 ++++
 .../inference/contexts/dynamic_context.py     |   3 +
 .../coordinator.py                            | 153 ++++++++++++++++--
 .../handlers.py                               |  46 ++++--
 .../core/inference/engines/dynamic_engine.py  |   2 +
 megatron/core/inference/inference_client.py   |  33 +++-
 .../endpoints/chat_completions.py             |  30 +++-
 .../text_generation_server.py                 |  34 +++-
 8 files changed, 294 insertions(+), 35 deletions(-)

diff --git a/megatron/core/inference/config.py b/megatron/core/inference/config.py
index 44bbf410c7a..3b8887a2f40 100644
--- a/megatron/core/inference/config.py
+++ b/megatron/core/inference/config.py
@@ -114,6 +114,24 @@ class PrefixCachingCoordinatorPolicy(str, Enum):
     """Route to the rank with the fewest in-flight requests. Ignores prefix affinity."""
 
 
+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."""
 
@@ -328,6 +346,16 @@ 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_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.
diff --git a/megatron/core/inference/contexts/dynamic_context.py b/megatron/core/inference/contexts/dynamic_context.py
index 0b50f50d6de..739a2a0daab 100644
--- a/megatron/core/inference/contexts/dynamic_context.py
+++ b/megatron/core/inference/contexts/dynamic_context.py
@@ -336,6 +336,9 @@ 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
 
+        # 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
         # counter is not overloaded with cache-eviction semantics.
diff --git a/megatron/core/inference/data_parallel_inference_coordinator/coordinator.py b/megatron/core/inference/data_parallel_inference_coordinator/coordinator.py
index 88a2fe8015e..c2792b64611 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,10 @@
 import numpy as np
 import torch
 
-from megatron.core.inference.config import PrefixCachingCoordinatorPolicy
+from megatron.core.inference.config import (
+    PrefixCachingCoordinatorPolicy,
+    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 (
@@ -99,6 +103,10 @@ def __init__(
             PrefixCachingCoordinatorPolicy.FIRST_PREFIX_BLOCK
         ),
         prefix_caching_routing_alpha: float = 0.5,
+        prefix_caching_eviction_policy: PrefixCachingEvictionPolicy = (
+            PrefixCachingEvictionPolicy.LRU
+        ),
+        prefix_cache_ttl_seconds: float = 300.0,
         schedule_output_path: str | None = None,
         hostname: str | None = None,
     ):
@@ -218,11 +226,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()
@@ -371,6 +396,26 @@ 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()
 
+        if self.prefix_caching_coordinator_policy == PrefixCachingCoordinatorPolicy.LONGEST_PREFIX:
+            # Cost of sending here: the prefill this rank would still have to
+            # compute, scaled by how contended it is. Multiplicative rather than a
+            # weighted sum because prefill saved on a rank that cannot start it
+            # promptly is not saved.
+            #
+            # 1 + load, so an idle rank is charged for the prefill it would still
+            # have to do rather than scoring zero on any prefix. Reading it as
+            # queue position: this request is the (load + 1)-th on that rank, so
+            # the product approximates when its prefill would finish. Affinity
+            # therefore decides from the first request, including while every rank
+            # is idle.
+            depth = self._prefix_depth_vector(request_hashes)
+            remaining_blocks = np.maximum(0, len(request_hashes) - depth)
+            cost = remaining_blocks.astype(np.float64) * (1.0 + self._pending_counts)
+            # Tiebreak: lowest cost, then least loaded, then lowest rank index.
+            n_ranks = len(self._identities_list)
+            order = np.lexsort((np.arange(n_ranks), self._pending_counts, cost))
+            return self._identities_list[int(order[0])]
+
         match, recency = self._match_vector(request_hashes)
 
         alpha = self.prefix_caching_routing_alpha
@@ -386,17 +431,94 @@ def get_best_data_parallel_rank(self, request_hashes):
         return self._identities_list[best_idx]
 
     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,)``.
@@ -421,7 +543,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()
@@ -503,6 +628,10 @@ def entrypoint(
             PrefixCachingCoordinatorPolicy.FIRST_PREFIX_BLOCK
         ),
         prefix_caching_routing_alpha: float = 0.5,
+        prefix_caching_eviction_policy: PrefixCachingEvictionPolicy = (
+            PrefixCachingEvictionPolicy.LRU
+        ),
+        prefix_cache_ttl_seconds: float = 300.0,
         schedule_output_path: str | None = None,
         hostname: str | None = None,
     ):
@@ -537,6 +666,8 @@ 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_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 804e63d0641..1f6aee0d75c 100644
--- a/megatron/core/inference/data_parallel_inference_coordinator/handlers.py
+++ b/megatron/core/inference/data_parallel_inference_coordinator/handlers.py
@@ -24,8 +24,12 @@
 """
 
 import logging
+import time
 
-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
@@ -76,9 +80,9 @@ def handle_connect(coordinator, sender_identity, metadata, bodies):
 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]`` and
-    ``bodies[0]`` is the packed prompt, which is forwarded to the engine as-is
-    unless the routing policy needs to hash it.
+    ``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.
     """
@@ -89,6 +93,7 @@ def handle_submit_request(coordinator, sender_identity, metadata, bodies):
 
     _, 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.
@@ -103,18 +108,17 @@ def handle_submit_request(coordinator, sender_identity, metadata, bodies):
         [Headers.SUBMIT_REQUEST.value, request_id, sampling_params], use_bin_type=True
     )
 
-    # Only prefix-affinity routing consults the block hashes. When nothing will
-    # look at them, skip hashing *and* the prompt decode it needs -- avoiding
-    # that decode is the reason the prompt travels in its own frame.
+    # 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 (
-        coordinator.enable_prefix_caching
-        and coordinator.block_size_tokens is not None
+        hash_frame is not None
+        and coordinator.enable_prefix_caching
         and coordinator.prefix_caching_coordinator_policy
         != PrefixCachingCoordinatorPolicy.LOAD_BALANCED
     ):
-        request_hashes = coordinator.compute_request_hashes(
-            msgpack.unpackb(prompt_frame, raw=False)
-        )
+        request_hashes = msgpack.unpackb(hash_frame, raw=False)
     else:
         request_hashes = []
 
@@ -124,6 +128,16 @@ def handle_submit_request(coordinator, sender_identity, metadata, bodies):
     ):
         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)
@@ -141,6 +155,9 @@ def handle_submit_request(coordinator, sender_identity, metadata, bodies):
     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(
             {
@@ -227,6 +244,11 @@ def handle_engine_reply(coordinator, sender_identity, metadata, bodies):
             if idx is not None:
                 assert coordinator._pending_counts[idx] >= 1
                 coordinator._pending_counts[idx] -= 1
+            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
diff --git a/megatron/core/inference/engines/dynamic_engine.py b/megatron/core/inference/engines/dynamic_engine.py
index 346147a84ba..b6109165a3b 100644
--- a/megatron/core/inference/engines/dynamic_engine.py
+++ b/megatron/core/inference/engines/dynamic_engine.py
@@ -709,6 +709,8 @@ 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_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,
                 },
diff --git a/megatron/core/inference/inference_client.py b/megatron/core/inference/inference_client.py
index c6aac68ae2f..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,27 +114,38 @@ def add_request(
         """
         request_id = self.next_request_id
         self.next_request_id += 1
-        self.socket.send_multipart(self._submit_frames(request_id, prompt, sampling_params))
+        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):
-        """Frame a submission as [metadata, prompt].
+    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.
         """
-        return [
+        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."""
@@ -148,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.
 
@@ -175,7 +192,9 @@ def add_request_streaming(
         sampling_params.streaming = True
         request_id = self.next_request_id
         self.next_request_id += 1
-        self.socket.send_multipart(self._submit_frames(request_id, prompt, sampling_params))
+        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
         )
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 35ae357d00d..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
@@ -9,7 +9,13 @@
 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,
@@ -457,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)
@@ -643,6 +651,19 @@ async def chat_completions():
             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.
@@ -655,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(
@@ -672,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()
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 65b916f60b3..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
@@ -19,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
 
@@ -50,6 +51,8 @@ async def _run_text_gen_server(
     parsers: Optional[List[str]] = None,
     verbose: bool = False,
     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
@@ -85,6 +88,11 @@ 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
@@ -135,6 +143,8 @@ def _server_process_worker(
     parsers: Optional[List[str]] = None,
     verbose: bool = False,
     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()
@@ -142,7 +152,15 @@ def _server_process_worker(
     try:
         loop.run_until_complete(
             _run_text_gen_server(
-                coordinator_addr, tokenizer, rank, server_port, parsers, verbose, hostname
+                coordinator_addr,
+                tokenizer,
+                rank,
+                server_port,
+                parsers,
+                verbose,
+                hostname,
+                block_size_tokens,
+                prefix_caching_coordinator_policy,
             )
         )
     except KeyboardInterrupt:
@@ -198,6 +216,8 @@ def start_text_gen_server(
     num_replicas: int = 4,
     hostname: Optional[str] = None,
     sock: Optional[socket.socket] = None,
+    block_size_tokens: Optional[int] = None,
+    prefix_caching_coordinator_policy: Optional[PrefixCachingCoordinatorPolicy] = None,
 ) -> Optional[str]:
     """Start the text generation server.
 
@@ -245,7 +265,17 @@ def start_text_gen_server(
     for i in range(num_replicas):
         p = mp.Process(
             target=_server_process_worker,
-            args=(coordinator_addr, tokenizer, rank, server_port, parsers, verbose, hostname),
+            args=(
+                coordinator_addr,
+                tokenizer,
+                rank,
+                server_port,
+                parsers,
+                verbose,
+                hostname,
+                block_size_tokens,
+                prefix_caching_coordinator_policy,
+            ),
             daemon=True,
         )
         p.start()

From 457d3b0a919688c2a62c38e6654ccc51f4e0dbf9 Mon Sep 17 00:00:00 2001
From: Siddharth Singh <sidsingh@nvidia.com>
Date: Wed, 19 Aug 2026 10:58:50 -0700
Subject: [PATCH 7/7] feat(inference): make the prefix-caching cost policy
 configurable

The coordinator conflated two decisions: which affinity signal to read
(first-block hit vs contiguous prefix depth) and how that affinity is weighed
against rank load. Split them so the cost function is selectable independently
of the signal.

`PrefixCachingCoordinatorPolicy` keeps its meaning -- it picks the signal --
and both signals are now normalized to the fraction of the request already
cached on a rank, in [0, 1]. The new `PrefixCachingCostPolicy` picks how that
fraction is scored, and composes with either signal:

  RELATIVE_LOAD_WEIGHTED (new default)
      score = fraction - beta * (load - mean) / max(1, mean)
      Load is measured against the fleet mean, so the penalty vanishes while
      ranks are balanced: at saturation this is pure affinity, and load only
      pulls toward idle ranks as the fleet diverges. Approximates the session
      stickiness a session-affinity router gets for free, with no session id.
      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 (the previous first_prefix_block behaviour)
      score = alpha * fraction + (1 - alpha) * free_slots / max_requests
      Fixes the trade-off in absolute terms rather than relative to load.

Also flips the coordinator policy default from LOAD_BALANCED to LONGEST_PREFIX,
so prefix affinity is used by default once prefix caching is enabled.

Measured on a 16-replica SWE rollout workload, RELATIVE_LOAD_WEIGHTED cuts
drain-phase imbalance sharply versus the previous multiplicative cost:
busy-phase load CV 0.434 -> 0.272, tail CV 1.275 -> 0.555,
idle-while-others-busy 15.7% -> 6.7%, max/median replica load 2.13 -> 1.44.

Signed-off-by: Siddharth Singh <sidsingh@nvidia.com>
---
 megatron/core/inference/config.py             | 47 ++++++++-
 .../inference/contexts/dynamic_context.py     |  2 +
 .../coordinator.py                            | 99 ++++++++++++-------
 .../core/inference/engines/dynamic_engine.py  |  2 +
 megatron/training/arguments.py                | 41 +++++---
 megatron/training/config/inference_config.py  | 39 ++++++--
 6 files changed, 170 insertions(+), 60 deletions(-)

diff --git a/megatron/core/inference/config.py b/megatron/core/inference/config.py
index 3b8887a2f40..3e95e0f4b97 100644
--- a/megatron/core/inference/config.py
+++ b/megatron/core/inference/config.py
@@ -114,6 +114,34 @@ 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.
 
@@ -338,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.
@@ -356,10 +384,25 @@ class InferenceConfig:
     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 739a2a0daab..ffea55331f9 100644
--- a/megatron/core/inference/contexts/dynamic_context.py
+++ b/megatron/core/inference/contexts/dynamic_context.py
@@ -335,6 +335,8 @@ 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
diff --git a/megatron/core/inference/data_parallel_inference_coordinator/coordinator.py b/megatron/core/inference/data_parallel_inference_coordinator/coordinator.py
index c2792b64611..bc875d5d488 100644
--- a/megatron/core/inference/data_parallel_inference_coordinator/coordinator.py
+++ b/megatron/core/inference/data_parallel_inference_coordinator/coordinator.py
@@ -16,6 +16,7 @@
 
 from megatron.core.inference.config import (
     PrefixCachingCoordinatorPolicy,
+    PrefixCachingCostPolicy,
     PrefixCachingEvictionPolicy,
 )
 from megatron.core.inference.headers import Headers, UnknownHeaderError
@@ -100,9 +101,13 @@ 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
         ),
@@ -208,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
 
@@ -374,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.
@@ -396,39 +404,48 @@ 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()
 
-        if self.prefix_caching_coordinator_policy == PrefixCachingCoordinatorPolicy.LONGEST_PREFIX:
-            # Cost of sending here: the prefill this rank would still have to
-            # compute, scaled by how contended it is. Multiplicative rather than a
-            # weighted sum because prefill saved on a rank that cannot start it
-            # promptly is not saved.
-            #
-            # 1 + load, so an idle rank is charged for the prefill it would still
-            # have to do rather than scoring zero on any prefix. Reading it as
-            # queue position: this request is the (load + 1)-th on that rank, so
-            # the product approximates when its prefill would finish. Affinity
-            # therefore decides from the first request, including while every rank
-            # is idle.
+        n_ranks = len(self._identities_list)
+        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)
-            remaining_blocks = np.maximum(0, len(request_hashes) - depth)
-            cost = remaining_blocks.astype(np.float64) * (1.0 + self._pending_counts)
-            # Tiebreak: lowest cost, then least loaded, then lowest rank index.
-            n_ranks = len(self._identities_list)
-            order = np.lexsort((np.arange(n_ranks), self._pending_counts, cost))
-            return self._identities_list[int(order[0])]
-
-        match, recency = self._match_vector(request_hashes)
-
-        alpha = self.prefix_caching_routing_alpha
+            fraction = depth.astype(np.float64) / max(n_blocks, 1)
+            recency = np.zeros(n_ranks, dtype=np.float64)
 
-        # 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)
+        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])]
 
-        # 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]
+        # 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, and take a reference to them.
@@ -625,9 +642,13 @@ 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
         ),
@@ -666,6 +687,8 @@ 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,
diff --git a/megatron/core/inference/engines/dynamic_engine.py b/megatron/core/inference/engines/dynamic_engine.py
index b6109165a3b..96641eb38b1 100644
--- a/megatron/core/inference/engines/dynamic_engine.py
+++ b/megatron/core/inference/engines/dynamic_engine.py
@@ -709,6 +709,8 @@ 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,
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,