Add Agentic-Aware Tail-Optimized LRU eviction to the unified radix cache (#34012)

Co-authored-by: Shuwen Wang <47200617+alphabetc1@users.noreply.github.com>

python/sglang/srt/arg_groups/choices.py

@@ -187,7 +187,7 @@ FP4_GEMM_RUNNER_BACKEND_CHOICES = [
     "marlin",
 ]
 
-RADIX_EVICTION_POLICY_CHOICES = ["lru", "lfu", "slru", "priority"]
+RADIX_EVICTION_POLICY_CHOICES = ["lru", "lfu", "slru", "priority", "tlru"]
 
 RL_ON_POLICY_TARGET_CHOICES = ["fsdp"]
 

python/sglang/srt/arg_groups/fields/memory.py

@@ -35,8 +35,11 @@ class Memory(msgspec.Struct):
             help=(
                 "The eviction policy of radix trees. 'lru' stands for Least "
                 "Recently Used, 'lfu' stands for Least Frequently Used, 'slru' "
-                "stands for Segmented Least Recently Used, and 'priority' evicts "
-                "lower-priority requests first. See "
+                "stands for Segmented Least Recently Used, 'priority' evicts "
+                "lower-priority requests first, and 'tlru' stands for "
+                "Tail-Optimized LRU (arXiv:2510.15152), which evicts the part of "
+                "a conversation that cannot affect tail TTFT before falling back "
+                "to LRU. See "
                 "https://docs.sglang.io/docs/advanced_features/radix_eviction_policy "
                 "for what each policy optimizes for."
             ),
@@ -52,10 +55,13 @@ class Memory(msgspec.Struct):
         Arg(
             help=(
                 "Tuning parameters for --radix-eviction-policy, as a json object "
-                "passed to the policy as keyword arguments. Only 'slru' takes any "
-                "today: protected_threshold (int, default 2), e.g. "
-                "'{\"protected_threshold\": 4}'. An unrecognized key fails at "
-                "startup, naming the key and the policy. See "
+                "passed to the policy as keyword arguments. 'slru' takes "
+                "protected_threshold (int, default 2), e.g. "
+                "'{\"protected_threshold\": 4}'; 'tlru' takes threshold and "
+                "next_prompt_estimate (ints, tokens), e.g. "
+                '\'{"threshold": 4096, "next_prompt_estimate": 512}\'. An '
+                "unrecognized key fails at startup, naming the key and the "
+                "policy. See "
                 "https://docs.sglang.io/docs/advanced_features/radix_eviction_policy#policy-parameters "
                 "for the full parameter list."
             ),

python/sglang/srt/arg_groups/kv_cache_hook.py

@@ -163,6 +163,24 @@ def handle_cache_compatibility(server_args: Any) -> None:
             "--disable-priority-preemption when priority scheduling is enabled."
         )
 
+    if cfg.radix_eviction_policy == "tlru":
+        tlru_config = cfg.radix_eviction_policy_config or {}
+        threshold = tlru_config.get("threshold", 0)
+        next_prompt_estimate = tlru_config.get("next_prompt_estimate", 0)
+        if threshold < 0 or next_prompt_estimate < 0:
+            raise ValueError(
+                "--radix-eviction-policy tlru requires non-negative 'threshold' and "
+                "'next_prompt_estimate' in --radix-eviction-policy-config, got "
+                f"{threshold} and {next_prompt_estimate}."
+            )
+        if threshold <= next_prompt_estimate:
+            raise ValueError(
+                "--radix-eviction-policy tlru needs 'threshold' greater than "
+                f"'next_prompt_estimate' in --radix-eviction-policy-config, got "
+                f"{threshold} <= {next_prompt_estimate}; otherwise no tokens are "
+                "ever TEL-safe and T-LRU is exactly LRU."
+            )
+
     if cfg.enable_hierarchical_cache and cfg.disable_radix_cache:
         raise ValueError(
             "The arguments enable-hierarchical-cache and disable-radix-cache are mutually exclusive "

python/sglang/srt/mem_cache/evict_policy.py

@@ -46,6 +46,41 @@ class PriorityStrategy(EvictionStrategy):
         return (node.priority, node.last_access_time)
 
 
+class TLRUStrategy(EvictionStrategy):
+    """Tail-Optimized LRU (Zhang et al., arXiv:2510.15152).
+
+    A conversation with history length L whose next prompt is expected to add
+    Q_hat tokens only has to keep L + Q_hat - threshold tokens cached to hold its
+    next prefill under the TTFT budget; tokens past that budget cannot improve
+    tail latency and are "TEL-safe", i.e. free to evict. Such nodes are reported
+    as infinitely old, which is the implementation the paper suggests: the
+    existing eviction driver then drains them before anything else (the paper's
+    phase 1) and continues in plain recency order once they run out (phase 2),
+    so neither eviction loop needs to know about T-LRU.
+
+    threshold and next_prompt_estimate are token counts, whereas the paper states
+    both in blocks; multiply the paper's values by page_size to convert.
+    """
+
+    def __init__(self, threshold: int = 0, next_prompt_estimate: int = 0):
+        self.threshold = threshold
+        self.next_prompt_estimate = next_prompt_estimate
+
+    def get_priority(self, node: TreeNode) -> Tuple[int, float]:
+        # node._tlru_history_len is the branch's high-water depth, i.e. the
+        # paper's L, and deliberately does not shrink when the tail is trimmed.
+        # Deriving L from what is still resident instead would leave the
+        # shortened conversation over budget on the next pass too, and T-LRU
+        # would walk it down to nothing rather than stopping after
+        # (threshold - Q_hat) tokens.
+        budget = max(
+            node._tlru_history_len + self.next_prompt_estimate - self.threshold, 0
+        )
+        cached_without_this_node = node._tlru_cached_prefix_len - len(node.key)
+        tel_safe = cached_without_this_node >= budget
+        return (-1 if tel_safe else 0, node.last_access_time)
+
+
 class SLRUStrategy(EvictionStrategy):
     def __init__(self, protected_threshold: int = 2):
         self.protected_threshold = protected_threshold

python/sglang/srt/mem_cache/registry.py

@@ -263,6 +263,18 @@ def create_tree_cache(ctx: TreeCacheBuildContext) -> BasePrefixCache:
             "option that selected another tree cache for this model."
         )
 
+    if get_memory().radix_eviction_policy == "tlru":
+        from sglang.srt.mem_cache.unified_radix_cache import UnifiedRadixCache
+
+        # T-LRU's per-node tail bookkeeping only exists on the unified tree;
+        # any other cache would silently fall back to LRU ordering.
+        if not isinstance(cache, UnifiedRadixCache):
+            raise ValueError(
+                "--radix-eviction-policy tlru requires UnifiedRadixCache, but "
+                f"tree_cache is {type(cache).__name__}. Drop the flag or the "
+                "option that selected another tree cache for this model."
+            )
+
     hicache_attached = cache.cache_controller is not None
     streaming_wrapped = False
     if (

python/sglang/srt/mem_cache/unified_cache/unified_tree_core.py

@@ -127,6 +127,10 @@ class UnifiedTreeNode:
         # Namespace-aware hashes used only for external KV events.
         self.event_hash_value: Optional[list[str]] = None
         self.hit_count = 0
+        # T-LRU only (0 under other policies): tokens root -> self, and the
+        # branch's high-water depth, which survives tail trimming.
+        self._tlru_cached_prefix_len = 0
+        self._tlru_history_len = 0
         self.external_cache_stored = False
         self.priority = priority
         self.lru_prev: list[UnifiedTreeNode | None] = [None] * (
@@ -182,6 +186,28 @@ class UnifiedTreeNode:
         return [value for chunk in reversed(chunks) for value in chunk]
 
 
+def _set_tlru_lens_and_raise_history(
+    node: UnifiedTreeNode, parent: UnifiedTreeNode
+) -> None:
+    """Record a new node's path depth and raise the branch high-water mark.
+
+    The walk stops early only at an ancestor whose subtree already reached this
+    depth (e.g. under a deeper sibling); a conversation that keeps setting a new
+    high-water mark walks its whole root path, so an insert costs O(path nodes)
+    in the worst case. The path length is bounded by the number of node segments
+    (roughly extend/split operations, not tokens), and the walk only runs under
+    --radix-eviction-policy tlru.
+    """
+    assert node.key is not None
+    node._tlru_cached_prefix_len = parent._tlru_cached_prefix_len + len(node.key)
+    if node._tlru_history_len < node._tlru_cached_prefix_len:
+        node._tlru_history_len = node._tlru_cached_prefix_len
+    cur = parent
+    while cur is not None and cur._tlru_history_len < node._tlru_cached_prefix_len:
+        cur._tlru_history_len = node._tlru_cached_prefix_len
+        cur = cur.parent
+
+
 class UnifiedLRUList:
     def __init__(
         self,
@@ -425,6 +451,9 @@ class UnifiedTreeCore(UnifiedTreeCoreInterface):
         self.eviction_strategy = get_eviction_strategy(
             params.eviction_policy.lower(), params.eviction_policy_config
         )
+        # The node _tlru_* lens are read only by TLRUStrategy.get_priority;
+        # every other policy skips the per-insert bookkeeping entirely.
+        self.tlru_bookkeeping = params.eviction_policy.lower() == "tlru"
 
         # ``device`` is derived from the construction-time allocator; the
         # allocator/pool themselves are owned by the cache, not the tree.
@@ -1355,6 +1384,14 @@ class UnifiedTreeCore(UnifiedTreeCoreInterface):
         new_node.hit_count = child.hit_count
         new_node.external_cache_stored = child.external_cache_stored
         new_node.creation_time = child.creation_time
+        if self.tlru_bookkeeping:
+            # A split adds no depth to the branch: the new parent sits at
+            # split_len tokens and inherits the branch's high-water mark, while
+            # child keeps its own depth because its path length is unchanged.
+            new_node._tlru_cached_prefix_len = (
+                new_node.parent._tlru_cached_prefix_len + split_len
+            )
+            new_node._tlru_history_len = child._tlru_history_len
         # Split fragments stay on the anchor's root path for the ack's walk.
         new_node.load_back_pending_id = child.load_back_pending_id
         # The rotation base is constant along a chain (position-page P keeps
@@ -1413,6 +1450,8 @@ class UnifiedTreeCore(UnifiedTreeCoreInterface):
         new_node = self._new_node(priority=priority)
         new_node.parent = parent
         new_node.key = key
+        if self.tlru_bookkeeping:
+            _set_tlru_lens_and_raise_history(new_node, parent)
         # Chain-constant under sharding: the pre-flight decline in
         # begin_insert() guarantees this tail continues the matched prefix's
         # rotation, so stamping the inserting request's base keeps every node
@@ -2165,6 +2204,8 @@ class UnifiedTreeCore(UnifiedTreeCoreInterface):
         new_node = self._new_node(priority=node.priority)
         new_node.parent = node
         new_node.key = key
+        if self.tlru_bookkeeping:
+            _set_tlru_lens_and_raise_history(new_node, node)
         new_node.hash_value = hash_value
         new_node.component_data[BASE_COMPONENT_TYPE].host_value = host_value.clone()
         node.children[child_key] = new_node

python/sglang/srt/mem_cache/utils.py

@@ -54,6 +54,7 @@ from sglang.srt.mem_cache.evict_policy import (
     MRUStrategy,
     PriorityStrategy,
     SLRUStrategy,
+    TLRUStrategy,
 )
 
 _EVICTION_POLICY_FACTORIES: dict[str, Callable[..., EvictionStrategy]] = {
@@ -64,6 +65,7 @@ _EVICTION_POLICY_FACTORIES: dict[str, Callable[..., EvictionStrategy]] = {
     "filo": FILOStrategy,
     "priority": PriorityStrategy,
     "slru": SLRUStrategy,
+    "tlru": TLRUStrategy,
 }