Improve load balancing behavior for prefix cache-aware routing

megatron/core/inference/config.py

@@ -260,6 +260,12 @@ class InferenceConfig:
     Only applies when enable_prefix_caching is True and using a coordinator.
     """
 
+    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.
+    """
+
     prefix_caching_mamba_gb: Optional[float] = None
     """GPU memory budget (in GB) for the Mamba state cache used by prefix caching
     on hybrid models. Each cache slot stores SSM and conv states for all Mamba layers
@@ -298,7 +304,14 @@ class InferenceConfig:
     """
 
     use_synchronous_zmq_collectives: bool = False
-    """Whether to use synchronous ZMQ collectives for inference. If True, the 
-    all_reduce_max operation will be performed synchronously, which can help reduce 
+    """Whether to use synchronous ZMQ collectives for inference. If True, the
+    all_reduce_max operation will be performed synchronously, which can help reduce
     performance variability for MoEs.
     """
+
+    def __post_init__(self):
+        if not (0.0 <= self.prefix_caching_routing_alpha <= 1.0):
+            raise ValueError(
+                f"prefix_caching_routing_alpha must be in [0, 1], "
+                f"got {self.prefix_caching_routing_alpha}"
+            )

megatron/core/inference/contexts/dynamic_context.py

@@ -252,10 +252,11 @@ def __init__(self, model_config: TransformerConfig, inference_config: InferenceC
         self.enable_prefix_caching = inference_config.enable_prefix_caching
         self.prefix_caching_eviction_policy = inference_config.prefix_caching_eviction_policy
         self.prefix_caching_coordinator_policy = inference_config.prefix_caching_coordinator_policy
-        # Engine step counter (used for logging, metrics, and event tracking)
-        self.step_count = 0
 
-        # Separate monotonic clock for prefix caching LRU eviction ordering.
+        # Hyperparameter for choosing to prioritize prefix hit matches vs minimizing idle load
+        self.prefix_caching_routing_alpha = inference_config.prefix_caching_routing_alpha
+
+        # 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.
         self.prefix_cache_lru_clock = 0
@@ -264,6 +265,9 @@ def __init__(self, model_config: TransformerConfig, inference_config: InferenceC
         self.prefix_cache_hits = 0  # requests that matched at least one cached block
         self.prefix_cache_blocks_matched = 0  # total matched blocks across all requests
 
+        # Engine step counter (used for logging, metrics, and event tracking)
+        self.step_count = 0
+
         self.cache_mla_latent = (
             isinstance(model_config, MLATransformerConfig) and model_config.cache_mla_latents
         )

megatron/core/inference/data_parallel_inference_coordinator.py

@@ -11,6 +11,7 @@
 from multiprocessing import Event
 from multiprocessing.connection import Connection
 
+import numpy as np
 import torch
 
 from megatron.core.inference.config import PrefixCachingCoordinatorPolicy
@@ -87,13 +88,15 @@ def __init__(
         pipe_connection: Connection,
         data_parallel_size: int,
         tokenizer,
+        max_requests,
         inference_coordinator_port: int | None = None,
         deterministic_mode: bool = False,
         block_size_tokens: int | None = None,
         enable_prefix_caching: bool = False,
         prefix_caching_coordinator_policy: PrefixCachingCoordinatorPolicy = (
             PrefixCachingCoordinatorPolicy.FIRST_PREFIX_BLOCK
         ),
+        prefix_caching_routing_alpha: float = 0.5,
         schedule_output_path: str | None = None,
         hostname: str | None = None,
     ):
@@ -110,6 +113,10 @@ def __init__(
                 expected to connect.
             tokenizer: The tokenizer to use for prompt tokenization and detokenization.
             inference_coordinator_port (Optional[int]): The TCP port number to bind the server to.
+            prefix_caching_routing_alpha (float): Weight for prefix-aware routing score:
+                score = alpha * match + (1 - alpha) * normalized_load.
+            max_requests (int): Max concurrent requests per rank, used to
+                compute normalized_load for prefix-aware scoring.
         """
         assert HAVE_ZMQ, (
             "please install the pyzmq library to use DataParallelInferenceCoordinator\n"
@@ -179,6 +186,7 @@ def __init__(
 
         self.request_id_to_client_id = {}
         self.request_id_to_client_request_id = {}
+        self.request_id_to_rank = {}  # Maps request_id → rank identity for pending count tracking
 
         self.next_request_id = 0
         self.tokenizer = tokenizer
@@ -188,8 +196,9 @@ def __init__(
         self.block_size_tokens = block_size_tokens
         self.enable_prefix_caching = enable_prefix_caching
         self.prefix_caching_coordinator_policy = prefix_caching_coordinator_policy
-        self.hash_to_rank_info = {}  # Dict[int, Dict[bytes, int]]: hash → {rank → timestamp}
-        self._assignment_counter = 0
+        self.prefix_caching_routing_alpha = prefix_caching_routing_alpha
+        self.max_requests = max_requests
+        assert self.max_requests is not None and self.max_requests > 0
 
         # Schedule recording.
         self.schedule_output_path = schedule_output_path
@@ -201,6 +210,17 @@ def __init__(
             identity: idx for idx, identity in enumerate(sorted_identities)
         }
 
+        # Numpy arrays for vectorized scoring (indexed by rank index).
+        n_ranks = len(sorted_identities)
+        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
+
     def get_next_data_parallel_rank(self):
         """
         Selects the next data parallel rank using round-robin scheduling.
@@ -215,6 +235,26 @@ def get_next_data_parallel_rank(self):
         self._round_robin_idx = idx + 1
         return identities[idx]
 
+    def _register_rank_identity(self, identity):
+        """Register a new rank identity in the scoring data structures.
+
+        Called when a rank dynamically connects to a running coordinator
+        (e.g. in tests that spawn the coordinator with data_parallel_size=0
+        and let engines register after the fact).
+        """
+        if identity in self.identity_to_rank_index:
+            return
+        new_idx = len(self._identities_list)
+        self.identity_to_rank_index[identity] = new_idx
+        self._identities_list.append(identity)
+        self._pending_counts = np.append(self._pending_counts, np.int32(0))
+        logging.info(
+            "Coordinator: registered engine %s as rank index %d (now %d engines)",
+            identity,
+            new_idx,
+            len(self._identities_list),
+        )
+
     def _remove_engine(self, identity):
         """Remove a disconnected engine from the routing pool."""
         self.identities_of_data_parallel_ranks.remove(identity)
@@ -258,37 +298,39 @@ 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.
+        """Select the best DP rank based on prefix cache affinity and load.
 
-        Iterates request hashes in reverse order and picks the rank that cached
-        the longest matching prefix (the furthest hash found). Since hashes are
-        parent-chained, finding hash[i] in a rank guarantees hash[0..i-1] are
-        also present. Among ranks that share the longest match, the most recently
-        assigned rank (highest timestamp) is preferred. Falls back to round-robin
-        when no rank matches.
+        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).
 
         Args:
             request_hashes: List of block hashes for the request.
 
         Returns:
             bytes: The ZMQ identity of the selected data parallel rank.
         """
-        if (
-            not self.enable_prefix_caching
-            or not request_hashes
-            or self.prefix_caching_coordinator_policy == PrefixCachingCoordinatorPolicy.ROUND_ROBIN
-        ):
+        if self.prefix_caching_coordinator_policy == PrefixCachingCoordinatorPolicy.ROUND_ROBIN:
+            return self.get_next_data_parallel_rank()
+
+        if not self.enable_prefix_caching or not request_hashes:
             return self.get_next_data_parallel_rank()
 
-        # Reverse scan: first match is the longest prefix (parent-chained hashes).
-        for h in reversed(request_hashes):
-            rank_info = self.hash_to_rank_info.get(h)
-            if rank_info:
-                # Pick the most recently assigned rank.
-                best_rank = max(rank_info, key=rank_info.get)
-                return best_rank
+        match, recency = self._match_vector(request_hashes)
+
+        alpha = self.prefix_caching_routing_alpha
 
-        return self.get_next_data_parallel_rank()
+        # Vectorized score: alpha * match + (1-alpha) * free_capacity_fraction.
+        free_slots = np.maximum(0, self.max_requests - self._pending_counts).astype(np.float64)
+        scores = alpha * match + (1.0 - alpha) * (free_slots / self.max_requests)
+
+        # Tiebreak: highest score, then highest recency, then lowest rank index.
+        n_ranks = len(self._identities_list)
+        order = np.lexsort((np.arange(n_ranks), -recency, -scores))
+        best_idx = int(order[0])
+        return self._identities_list[best_idx]
 
     def _update_rank_hashes(self, rank_identity, request_hashes):
         """Record that a rank owns the given hashes.
@@ -297,10 +339,39 @@ def _update_rank_hashes(self, rank_identity, request_hashes):
             rank_identity: ZMQ identity of the target rank.
             request_hashes: List of block hashes assigned to this rank.
         """
-        self._assignment_counter += 1
-        ts = self._assignment_counter
+        rank_idx = self.identity_to_rank_index[rank_identity]
+        self._hash_assignment_counter += 1
+        ts = self._hash_assignment_counter
         for h in request_hashes:
-            self.hash_to_rank_info.setdefault(h, {})[rank_identity] = ts
+            self._hash_table.setdefault(h, {})[rank_idx] = ts
+
+    def _match_vector(self, hashes):
+        """Return ``(match, recency)`` vectors of shape ``(n_ranks,)``.
+
+        *match* is binary depth: ``(depth + 1) / len(hashes)`` for ranks that
+        have the deepest cached block, 0 otherwise.  *recency* is the raw
+        assignment timestamp for each matching rank (0 for non-matching ranks).
+
+        For ``FIRST_PREFIX_BLOCK`` the caller already truncates *hashes* to a
+        single element, so the same logic yields a binary 0/1 match score.
+        """
+        n_ranks = len(self._identities_list)
+        n = len(hashes)
+        zeros = np.zeros(n_ranks, dtype=np.float64)
+        if n == 0:
+            return zeros, zeros.copy()
+        for i in range(n - 1, -1, -1):
+            row = self._hash_table.get(hashes[i])
+            if row is None:
+                continue
+            rank_idxs = np.fromiter(row.keys(), dtype=np.intp)
+            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)
+            if present.any():
+                return present.astype(np.float64) * ((i + 1.0) / n), recency
+        return zeros, zeros.copy()
 
     def start(self):
         """
@@ -321,6 +392,7 @@ def start(self):
             if serialized_payload == b"":
                 if sender_identity not in self.identities_of_data_parallel_ranks:
                     self.identities_of_data_parallel_ranks.append(sender_identity)
+                    self._register_rank_identity(sender_identity)
                 continue
 
             deserialized_payload = msgpack.unpackb(serialized_payload, raw=False)
@@ -392,6 +464,8 @@ def start(self):
                     del self.request_id_to_client_request_id[request_id]
                     return
 
+                self.request_id_to_rank[request_id] = next_identity
+                self._pending_counts[self.identity_to_rank_index[next_identity]] += 1
                 if request_hashes:
                     self._update_rank_hashes(next_identity, request_hashes)
                 if self.schedule_records is not None:
@@ -471,6 +545,12 @@ def start(self):
                     client_request_identity = self.request_id_to_client_request_id[fid]
                     del self.request_id_to_client_id[fid]
                     del self.request_id_to_client_request_id[fid]
+                    assigned_rank = self.request_id_to_rank.pop(fid, None)
+                    if assigned_rank is not None:
+                        idx = self.identity_to_rank_index.get(assigned_rank)
+                        if idx is not None:
+                            assert self._pending_counts[idx] >= 1
+                            self._pending_counts[idx] -= 1
 
                     self.router_socket.send_multipart(
                         [
@@ -526,13 +606,15 @@ def entrypoint(
         ready_event: Event,
         data_parallel_size: int,
         tokenizer,
+        max_requests,
         inference_coordinator_port: int | None = None,
         deterministic_mode: bool = False,
         block_size_tokens: int | None = None,
         enable_prefix_caching: bool = False,
         prefix_caching_coordinator_policy: PrefixCachingCoordinatorPolicy = (
             PrefixCachingCoordinatorPolicy.FIRST_PREFIX_BLOCK
         ),
+        prefix_caching_routing_alpha: float = 0.5,
         schedule_output_path: str | None = None,
         hostname: str | None = None,
     ):
@@ -553,16 +635,20 @@ def entrypoint(
             enable_prefix_caching (bool): Whether prefix caching is enabled.
             prefix_caching_coordinator_policy (PrefixCachingCoordinatorPolicy): Routing policy.
             schedule_output_path (Optional[str]): Path to write scheduling decisions JSON.
+            prefix_caching_routing_alpha (float): Weight for prefix-aware routing score.
+            max_requests (int): Max concurrent requests per rank.
         """
         coordinator = cls(
             pipe_connection,
             data_parallel_size,
             tokenizer,
+            max_requests,
             inference_coordinator_port,
             deterministic_mode=deterministic_mode,
             block_size_tokens=block_size_tokens,
             enable_prefix_caching=enable_prefix_caching,
             prefix_caching_coordinator_policy=prefix_caching_coordinator_policy,
+            prefix_caching_routing_alpha=prefix_caching_routing_alpha,
             schedule_output_path=schedule_output_path,
             hostname=hostname,
         )

megatron/core/inference/engines/dynamic_engine.py

@@ -509,19 +509,21 @@ async def start_listening_to_data_parallel_coordinator(
             coordinator_ready_event = spawn_context.Event()
             self.inference_coordinator_process = spawn_context.Process(
                 target=DataParallelInferenceCoordinator.entrypoint,
-                args=(
-                    dp_process_pipe,
-                    coordinator_ready_event,
-                    get_pg_size(self.pg_collection.dp),
-                    self.controller.tokenizer,
-                    inference_coordinator_port,
-                    deterministic_mode,
-                    self.context.block_size_tokens,
-                    self.context.enable_prefix_caching,
-                    self.context.prefix_caching_coordinator_policy,
-                    coordinator_schedule_output_path,
-                    hostname,
-                ),
+                kwargs={
+                    "pipe_connection": dp_process_pipe,
+                    "ready_event": coordinator_ready_event,
+                    "data_parallel_size": get_pg_size(self.pg_collection.dp),
+                    "tokenizer": self.controller.tokenizer,
+                    "max_requests": self.context.max_requests,
+                    "inference_coordinator_port": inference_coordinator_port,
+                    "deterministic_mode": deterministic_mode,
+                    "block_size_tokens": self.context.block_size_tokens,
+                    "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,
+                    "schedule_output_path": coordinator_schedule_output_path,
+                    "hostname": hostname,
+                },
             )
             self.inference_coordinator_process.start()
             await await_process_call(dp_pipe.poll, self.inference_coordinator_process)

megatron/inference/utils.py

@@ -355,6 +355,7 @@ def get_inference_config_from_model_and_args(model: MegatronModule, args):
         enable_prefix_caching=args.inference_dynamic_batching_enable_prefix_caching,
         prefix_caching_eviction_policy=PrefixCachingEvictionPolicy(args.inference_dynamic_batching_prefix_caching_eviction_policy),
         prefix_caching_coordinator_policy=PrefixCachingCoordinatorPolicy(args.inference_dynamic_batching_prefix_caching_coordinator_policy),
+        prefix_caching_routing_alpha=getattr(args, 'inference_dynamic_batching_prefix_caching_routing_alpha', 0.5),
         prefix_caching_mamba_gb=getattr(args, 'inference_dynamic_batching_prefix_caching_mamba_gb', None),
         metrics_writer=metrics_writer,
         logging_step_interval=args.inference_logging_step_interval,

megatron/training/arguments.py

@@ -1881,6 +1881,13 @@ def _add_inference_args(parser):
                        'block hash only. "longest_prefix" routes to the rank with '
                        'the longest matching prefix. "round_robin" ignores prefix '
                        'affinity and cycles through ranks.')
+    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: '
+                       'score = alpha * match + (1 - alpha) * normalized_load. '
+                       'Higher alpha favors prefix cache hits; lower alpha '
+                       'favors load balance. Default: 0.5.')
     group.add_argument('--inference-dynamic-batching-prefix-caching-mamba-gb',
                        type=float, default=None,
                        dest='inference_dynamic_batching_prefix_caching_mamba_gb',