feat(inference): prefix-affinity routing and a configurable cost policy in the DP coordinator

megatron/core/inference/config.py

@@ -114,6 +114,52 @@ class PrefixCachingCoordinatorPolicy(str, Enum):
     """Route to the rank with the fewest in-flight requests. Ignores prefix affinity."""
 
 
+class PrefixCachingCostPolicy(str, Enum):
+    """How the coordinator weighs prefix affinity against rank load.
+
+    Orthogonal to `PrefixCachingCoordinatorPolicy`, which only selects the affinity
+    signal. Both signals are normalized to the fraction of the request already cached
+    on a rank, in [0, 1], so either cost policy composes with either of
+    LONGEST_PREFIX and FIRST_PREFIX_BLOCK. Neither applies under LOAD_BALANCED, which
+    ignores affinity entirely.
+
+    RELATIVE_LOAD_WEIGHTED (default) — score = fraction - beta * relative_load, highest
+    wins, where relative_load is (load - mean) / max(1, mean). Approximates the session
+    stickiness a session-affinity router gets for free: a multi-turn request lands back
+    on the rank holding its history, with no session id to key on. Both terms are
+    normalized, so beta is dimensionless, and measuring load against the fleet mean
+    makes the penalty vanish while ranks are balanced -- at saturation this is pure
+    affinity, and load only pulls toward idle ranks as the fleet diverges. The mean is
+    floored at 1 so a near-idle fleet does not turn one in-flight request into a large
+    relative load and thrash on noise.
+
+    FREE_CAPACITY_WEIGHTED — score = alpha * fraction + (1 - alpha) * free_capacity, highest
+    wins, with alpha from `prefix_caching_routing_alpha`. Fixes the trade-off in
+    absolute terms rather than relative to how loaded the fleet actually is.
+    """
+
+    RELATIVE_LOAD_WEIGHTED = "relative_load_weighted"
+    FREE_CAPACITY_WEIGHTED = "free_capacity_weighted"
+
+
+def routes_on_prefix(policy) -> bool:
+    """Whether `policy` needs per-request block hashes to make a routing decision.
+
+    Frontends call this to decide whether hashing a prompt is worth anything: under
+    LOAD_BALANCED the coordinator discards the hashes, so computing them is pure
+    overhead on the request path. Kept beside the enum so a new prefix-aware policy
+    only has to be added in one place.
+
+    Accepts the enum, its string value, or None (no policy configured).
+    """
+    if policy is None:
+        return False
+    return PrefixCachingCoordinatorPolicy(policy) in (
+        PrefixCachingCoordinatorPolicy.LONGEST_PREFIX,
+        PrefixCachingCoordinatorPolicy.FIRST_PREFIX_BLOCK,
+    )
+
+
 class KVCacheManagementMode(str, Enum):
     """Mode for handling large tensors (KV cache, Mamba states) during suspend/resume."""
 
@@ -320,18 +366,43 @@ class InferenceConfig:
     """
 
     prefix_caching_coordinator_policy: PrefixCachingCoordinatorPolicy = (
-        PrefixCachingCoordinatorPolicy.LOAD_BALANCED
+        PrefixCachingCoordinatorPolicy.LONGEST_PREFIX
     )
     """Routing policy for the DP inference coordinator. See
     `PrefixCachingCoordinatorPolicy` for options.
 
     Only applies when enable_prefix_caching is True and using a coordinator.
     """
 
+    prefix_cache_ttl_seconds: float = 300.0
+    """How long the coordinator assumes an engine still holds a block it routed.
+
+    Only applies under `PrefixCachingEvictionPolicy.LRU`, where the engine keeps
+    released blocks and evicts them under memory pressure -- something the
+    coordinator cannot observe, so it approximates by age. Too long and it claims
+    hits on blocks already evicted, routing for affinity and paying a cold prefill
+    anyway; too short and it forgets blocks the engine still holds.
+    """
+
+    prefix_caching_cost_policy: PrefixCachingCostPolicy = (
+        PrefixCachingCostPolicy.RELATIVE_LOAD_WEIGHTED
+    )
+    """How prefix affinity is weighed against rank load. See `PrefixCachingCostPolicy`.
+
+    Only applies when enable_prefix_caching is True and using a coordinator.
+    """
+
+    prefix_caching_load_beta: float = 1.0
+    """Weight on the load penalty under `PrefixCachingCostPolicy.RELATIVE_LOAD_WEIGHTED`,
+    in units of "full cache hits per 100% above mean load". 0 disables the penalty
+    (pure affinity); 1.0 means a rank at twice the fleet mean forfeits a whole
+    prompt's worth of cache credit.
+    """
+
     prefix_caching_routing_alpha: float = 0.5
     """Weight for prefix-aware scoring: score = alpha * match + (1 - alpha) * normalized_load.
     Higher alpha favors prefix cache hits; lower alpha favors load balance.
-    Must be in [0, 1]. Only applies when enable_prefix_caching is True and using a coordinator.
+    Must be in [0, 1]. Only applies under `PrefixCachingCostPolicy.FREE_CAPACITY_WEIGHTED`.
     """
 
     prefix_caching_mamba_gb: Optional[float] = None

megatron/core/inference/contexts/dynamic_context.py

@@ -335,6 +335,11 @@ def __init__(self, model_config: TransformerConfig, inference_config: InferenceC
 
         # Hyperparameter for choosing to prioritize prefix hit matches vs minimizing idle load
         self.prefix_caching_routing_alpha = inference_config.prefix_caching_routing_alpha
+        self.prefix_caching_cost_policy = inference_config.prefix_caching_cost_policy
+        self.prefix_caching_load_beta = inference_config.prefix_caching_load_beta
+
+        # How long the coordinator's model of this engine's prefix cache survives
+        self.prefix_cache_ttl_seconds = inference_config.prefix_cache_ttl_seconds
 
         # Monotonic clock for prefix caching LRU eviction ordering.
         # Incremented each engine step but kept independent so the engine step