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

Co-authored-by: Shuwen Wang <47200617+alphabetc1@users.noreply.github.com>
This commit is contained in:
yl3469
2026-09-17 17:52:59 +08:00
committed by GitHub
co-authored by Shuwen Wang
parent 575759d90a
commit 1a90ae6727
9 changed files with 349 additions and 8 deletions
@@ -24,6 +24,7 @@ Select one with `--radix-eviction-policy`. All of them fall back to least-recent
| `lfu` | The prefix with the fewest cache hits, then the least recent. | A small set of prompts is reused far more than the rest, and you want those to survive bursts of one-off traffic. | | `lfu` | The prefix with the fewest cache hits, then the least recent. | A small set of prompts is reused far more than the rest, and you want those to survive bursts of one-off traffic. |
| `slru` | Prefixes still in the probationary segment, then the least recent within a segment. | Like `lfu`, but you want a hard floor on how much a one-off prefix can displace a proven one. See [`slru` parameters](#slru-parameters). | | `slru` | Prefixes still in the probationary segment, then the least recent within a segment. | Like `lfu`, but you want a hard floor on how much a one-off prefix can displace a proven one. See [`slru` parameters](#slru-parameters). |
| `priority` | The prefix belonging to the lowest-priority request, then the least recent. | You run [priority scheduling](/docs/advanced_features/server_arguments) and want cache retention to follow the same ranking as admission. | | `priority` | The prefix belonging to the lowest-priority request, then the least recent. | You run [priority scheduling](/docs/advanced_features/server_arguments) and want cache retention to follow the same ranking as admission. |
| `tlru` | The tail of a conversation beyond what its next prefill needs to meet the TTFT budget ("TEL-safe" tokens), then the least recent. | Agentic / multi-turn workloads where tail TTFT matters more than raw hit rate. See [`tlru` parameters](#tlru-parameters). Requires the unified radix cache (the default tree). |
Notes on the scoring inputs: Notes on the scoring inputs:
@@ -55,7 +56,7 @@ TypeError: SLRUStrategy.__init__() got an unexpected keyword argument 'protected
### Policy parameters ### Policy parameters
Only `slru` currently takes a parameter. `lru`, `lfu`, and `priority` take none, so `--radix-eviction-policy-config` has no effect with them and any key is an error. `slru` and `tlru` take parameters. `lru`, `lfu`, and `priority` take none, so `--radix-eviction-policy-config` has no effect with them and any key is an error.
#### `slru` parameters #### `slru` parameters
@@ -67,6 +68,24 @@ Only `slru` currently takes a parameter. `lru`, `lfu`, and `priority` take none,
Raising it makes promotion harder, so the protected set stays small and closer to your genuinely hot prefixes; a prefix hit three times stays probationary at `4` but is protected at the default `2`. Lowering it to `1` promotes any prefix that is reused even once, which approaches `lru` with a one-hit grace period. Raising it makes promotion harder, so the protected set stays small and closer to your genuinely hot prefixes; a prefix hit three times stays probationary at `4` but is protected at the default `2`. Lowering it to `1` promotes any prefix that is reused even once, which approaches `lru` with a one-hit grace period.
#### `tlru` parameters
`tlru` (Tail-Optimized LRU, [arXiv:2510.15152](https://arxiv.org/abs/2510.15152)) protects only the cached history a conversation's next prefill needs to stay under a TTFT budget; the rest of its tail is evicted first, and eviction continues in plain recency order once the TEL-safe tokens run out.
| Key | Type | Meaning |
|---|---|---|
| `threshold` | int | Tail-latency threshold ξ, in tokens: a conversation only keeps enough cache to hold its next prefill under ξ uncached tokens. Convert from a TTFT target by dividing it by the measured ms per uncached token. The paper states ξ in blocks, so multiply its values by `--page-size`. |
| `next_prompt_estimate` | int | Estimated tokens the next turn of a conversation will add; use the trace's empirical mean. |
Both keys are required in practice: `threshold` must be greater than `next_prompt_estimate` (only the difference affects behaviour), and values at or above it reduce T-LRU to plain LRU, which startup rejects.
```bash Command
python3 -m sglang.launch_server \
--model-path MODEL_PATH \
--radix-eviction-policy tlru \
--radix-eviction-policy-config '{"threshold": 4096, "next_prompt_estimate": 512}'
```
## Which policy to pick ## Which policy to pick
Start with `lru` and change it only against a measured cache hit rate — the counters are exposed under `--enable-metrics`. `lfu` and `slru` help when your traffic has a stable hot set that a recency-only policy keeps flushing; they hurt when prefix popularity shifts over time, because a prefix that earned a high hit count keeps its advantage after it stops being useful. Start with `lru` and change it only against a measured cache hit rate — the counters are exposed under `--enable-metrics`. `lfu` and `slru` help when your traffic has a stable hot set that a recency-only policy keeps flushing; they hurt when prefix popularity shifts over time, because a prefix that earned a high hit count keeps its advantage after it stops being useful.
+1 -1
View File
@@ -187,7 +187,7 @@ FP4_GEMM_RUNNER_BACKEND_CHOICES = [
"marlin", "marlin",
] ]
RADIX_EVICTION_POLICY_CHOICES = ["lru", "lfu", "slru", "priority"] RADIX_EVICTION_POLICY_CHOICES = ["lru", "lfu", "slru", "priority", "tlru"]
RL_ON_POLICY_TARGET_CHOICES = ["fsdp"] RL_ON_POLICY_TARGET_CHOICES = ["fsdp"]
+12 -6
View File
@@ -35,8 +35,11 @@ class Memory(msgspec.Struct):
help=( help=(
"The eviction policy of radix trees. 'lru' stands for Least " "The eviction policy of radix trees. 'lru' stands for Least "
"Recently Used, 'lfu' stands for Least Frequently Used, 'slru' " "Recently Used, 'lfu' stands for Least Frequently Used, 'slru' "
"stands for Segmented Least Recently Used, and 'priority' evicts " "stands for Segmented Least Recently Used, 'priority' evicts "
"lower-priority requests first. See " "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 " "https://docs.sglang.io/docs/advanced_features/radix_eviction_policy "
"for what each policy optimizes for." "for what each policy optimizes for."
), ),
@@ -52,10 +55,13 @@ class Memory(msgspec.Struct):
Arg( Arg(
help=( help=(
"Tuning parameters for --radix-eviction-policy, as a json object " "Tuning parameters for --radix-eviction-policy, as a json object "
"passed to the policy as keyword arguments. Only 'slru' takes any " "passed to the policy as keyword arguments. 'slru' takes "
"today: protected_threshold (int, default 2), e.g. " "protected_threshold (int, default 2), e.g. "
"'{\"protected_threshold\": 4}'. An unrecognized key fails at " "'{\"protected_threshold\": 4}'; 'tlru' takes threshold and "
"startup, naming the key and the policy. See " "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 " "https://docs.sglang.io/docs/advanced_features/radix_eviction_policy#policy-parameters "
"for the full parameter list." "for the full parameter list."
), ),
@@ -163,6 +163,24 @@ def handle_cache_compatibility(server_args: Any) -> None:
"--disable-priority-preemption when priority scheduling is enabled." "--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: if cfg.enable_hierarchical_cache and cfg.disable_radix_cache:
raise ValueError( raise ValueError(
"The arguments enable-hierarchical-cache and disable-radix-cache are mutually exclusive " "The arguments enable-hierarchical-cache and disable-radix-cache are mutually exclusive "
@@ -46,6 +46,41 @@ class PriorityStrategy(EvictionStrategy):
return (node.priority, node.last_access_time) 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): class SLRUStrategy(EvictionStrategy):
def __init__(self, protected_threshold: int = 2): def __init__(self, protected_threshold: int = 2):
self.protected_threshold = protected_threshold self.protected_threshold = protected_threshold
+12
View File
@@ -263,6 +263,18 @@ def create_tree_cache(ctx: TreeCacheBuildContext) -> BasePrefixCache:
"option that selected another tree cache for this model." "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 hicache_attached = cache.cache_controller is not None
streaming_wrapped = False streaming_wrapped = False
if ( if (
@@ -127,6 +127,10 @@ class UnifiedTreeNode:
# Namespace-aware hashes used only for external KV events. # Namespace-aware hashes used only for external KV events.
self.event_hash_value: Optional[list[str]] = None self.event_hash_value: Optional[list[str]] = None
self.hit_count = 0 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.external_cache_stored = False
self.priority = priority self.priority = priority
self.lru_prev: list[UnifiedTreeNode | None] = [None] * ( self.lru_prev: list[UnifiedTreeNode | None] = [None] * (
@@ -182,6 +186,28 @@ class UnifiedTreeNode:
return [value for chunk in reversed(chunks) for value in chunk] 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: class UnifiedLRUList:
def __init__( def __init__(
self, self,
@@ -425,6 +451,9 @@ class UnifiedTreeCore(UnifiedTreeCoreInterface):
self.eviction_strategy = get_eviction_strategy( self.eviction_strategy = get_eviction_strategy(
params.eviction_policy.lower(), params.eviction_policy_config 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 # ``device`` is derived from the construction-time allocator; the
# allocator/pool themselves are owned by the cache, not the tree. # 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.hit_count = child.hit_count
new_node.external_cache_stored = child.external_cache_stored new_node.external_cache_stored = child.external_cache_stored
new_node.creation_time = child.creation_time 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. # 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 new_node.load_back_pending_id = child.load_back_pending_id
# The rotation base is constant along a chain (position-page P keeps # 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 = self._new_node(priority=priority)
new_node.parent = parent new_node.parent = parent
new_node.key = key 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 # Chain-constant under sharding: the pre-flight decline in
# begin_insert() guarantees this tail continues the matched prefix's # begin_insert() guarantees this tail continues the matched prefix's
# rotation, so stamping the inserting request's base keeps every node # 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 = self._new_node(priority=node.priority)
new_node.parent = node new_node.parent = node
new_node.key = key 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.hash_value = hash_value
new_node.component_data[BASE_COMPONENT_TYPE].host_value = host_value.clone() new_node.component_data[BASE_COMPONENT_TYPE].host_value = host_value.clone()
node.children[child_key] = new_node node.children[child_key] = new_node
+2
View File
@@ -54,6 +54,7 @@ from sglang.srt.mem_cache.evict_policy import (
MRUStrategy, MRUStrategy,
PriorityStrategy, PriorityStrategy,
SLRUStrategy, SLRUStrategy,
TLRUStrategy,
) )
_EVICTION_POLICY_FACTORIES: dict[str, Callable[..., EvictionStrategy]] = { _EVICTION_POLICY_FACTORIES: dict[str, Callable[..., EvictionStrategy]] = {
@@ -64,6 +65,7 @@ _EVICTION_POLICY_FACTORIES: dict[str, Callable[..., EvictionStrategy]] = {
"filo": FILOStrategy, "filo": FILOStrategy,
"priority": PriorityStrategy, "priority": PriorityStrategy,
"slru": SLRUStrategy, "slru": SLRUStrategy,
"tlru": TLRUStrategy,
} }
@@ -0,0 +1,208 @@
"""Unit tests for the Tail-Optimized LRU eviction strategy (arXiv:2510.15152).
T-LRU reports a node as infinitely old once the conversation holding it is above
its TEL-safe budget, so the ordinary eviction driver drains those nodes first
(the paper's phase 1) and then continues in recency order (phase 2). These tests
exercise that ordering against stub nodes, so they need no GPU, model or CUDA
build and run in well under a second.
evict_policy is loaded straight from its file: importing sglang as a package
pulls in the engine's runtime dependencies, which would make a pure-logic test
require a full install.
"""
import importlib.util
import math
import os
from dataclasses import dataclass, field
try:
from sglang.test.ci.ci_register import register_cpu_ci
except ImportError: # standalone run without an sglang install; CI parses the
# registration below from the AST, so the stub changes nothing for CI.
def register_cpu_ci(**kwargs):
pass
register_cpu_ci(est_time=1, suite="base-a-test-cpu")
_HERE = os.path.dirname(os.path.abspath(__file__))
_EVICT_POLICY = os.path.normpath(
os.path.join(_HERE, "../../../../python/sglang/srt/mem_cache/evict_policy.py")
)
_spec = importlib.util.spec_from_file_location("evict_policy", _EVICT_POLICY)
_mod = importlib.util.module_from_spec(_spec)
_spec.loader.exec_module(_mod)
LRUStrategy, TLRUStrategy = _mod.LRUStrategy, _mod.TLRUStrategy
PAGE = 256
XI = 4096
Q_HAT = 1024
DELTA = XI - Q_HAT # tokens of tail the policy may free
@dataclass
class FakeKey:
n: int
def __len__(self):
return self.n
@dataclass
class FakeNode:
"""A radix node on one conversation's path; only the fields T-LRU reads."""
_tlru_cached_prefix_len: int
key_len: int
_tlru_history_len: int
last_access_time: float = 0.0
key: FakeKey = field(init=False)
def __post_init__(self):
self.key = FakeKey(self.key_len)
def chain(node_lens, convo_length=None, t0=0.0):
"""Build a root->leaf chain, deepest last.
convo_length defaults to the full depth, which is the state right after the
conversation's newest turn was inserted.
"""
convo = sum(node_lens) if convo_length is None else convo_length
nodes, depth = [], 0
for i, n in enumerate(node_lens):
depth += n
nodes.append(
FakeNode(
_tlru_cached_prefix_len=depth,
key_len=n,
_tlru_history_len=convo,
last_access_time=t0 + i,
)
)
return nodes
def strategy(threshold=XI, next_prompt_estimate=Q_HAT):
return TLRUStrategy(threshold=threshold, next_prompt_estimate=next_prompt_estimate)
def is_tel_safe(s, node):
return s.get_priority(node)[0] < 0
def freed_tokens(s, nodes):
return sum(n.key_len for n in nodes if is_tel_safe(s, n))
def test_fresh_conversation_frees_exactly_the_tail_budget():
s = strategy()
assert freed_tokens(s, chain([PAGE] * 40)) == DELTA
def test_trimming_stops_after_the_budget():
"""The survivors of a trim must be protected.
This is the regression guard for deriving the history length from what is
still resident: that would leave the shortened conversation over budget on
every subsequent pass and walk it down to nothing.
"""
s = strategy()
nodes = chain([PAGE] * 40)
survivors = [n for n in nodes if not is_tel_safe(s, n)]
assert freed_tokens(s, survivors) == 0
def test_conversation_under_threshold_is_entirely_free():
s = strategy()
short = chain([PAGE] * 4) # 1024 + Q_hat <= xi, so no caching is needed
assert freed_tokens(s, short) == sum(n.key_len for n in short)
def test_phase_one_spreads_across_conversations_then_falls_back_to_lru():
s = strategy()
old = chain([PAGE] * 40, t0=0.0)
new = chain([PAGE] * 40, t0=100.0)
order = sorted(old + new, key=s.get_priority)
n_safe = 2 * (DELTA // PAGE)
assert all(is_tel_safe(s, n) for n in order[:n_safe])
# Both conversations donate their tail, which is what the paper's
# per-conversation loop exists to produce.
assert any(n in old for n in order[:n_safe])
assert any(n in new for n in order[:n_safe])
# Phase 2 is plain recency.
assert order[n_safe].last_access_time == old[0].last_access_time
safe_times = [n.last_access_time for n in order[:n_safe]]
assert safe_times == sorted(safe_times)
def test_degenerates_to_lru_when_estimate_reaches_threshold():
pool = chain([PAGE] * 40, t0=0.0) + chain([PAGE] * 40, t0=100.0)
degenerate = strategy(next_prompt_estimate=XI)
lru = LRUStrategy()
assert [id(n) for n in sorted(pool, key=degenerate.get_priority)] == [
id(n) for n in sorted(pool, key=lru.get_priority)
]
assert freed_tokens(degenerate, chain([PAGE] * 40)) == 0
def test_oversized_tail_node_is_protected_rather_than_partially_freed():
"""Node granularity under-trims instead of over-trimming.
The paper trims one block at a time; a radix tree can only drop whole leaves,
so a turn larger than the budget stays put and phase 2 decides its fate.
"""
s = strategy()
assert freed_tokens(s, chain([PAGE * 40])) == 0
def test_compacted_branch_keeps_shared_prefix_protected():
"""Context compaction shortens a conversation instead of extending it.
It occurs in about 2% of turns in the agentic traces we benchmark, and forks a
shallow branch off a shared ancestor that still carries the deeper branch's
high-water mark. The ancestor is then measured against a history longer than
what hangs below it, which must stay conservative: over-protect the shared
prefix, never free it early.
"""
s = strategy()
assert freed_tokens(s, chain([PAGE] * 8, convo_length=200_000)) == 0
def test_budget_clamps_at_zero():
"""A budget below zero must mean nothing needs caching, not wrap around."""
s = strategy()
node = FakeNode(_tlru_cached_prefix_len=PAGE, key_len=PAGE, _tlru_history_len=0)
assert is_tel_safe(s, node)
def test_priority_is_finite_and_orderable():
"""The driver pushes (priority, node) onto a heap, so the keys must compare
without falling through to comparing nodes."""
s = strategy()
a, b = chain([PAGE] * 4)[:2]
for node in (a, b):
flag, when = s.get_priority(node)
assert flag in (-1, 0)
assert math.isfinite(when)
assert (s.get_priority(a) < s.get_priority(b)) or (
s.get_priority(b) < s.get_priority(a)
)
if __name__ == "__main__":
failures = 0
for name, fn in sorted(globals().items()):
if name.startswith("test_") and callable(fn):
try:
fn()
print(f"OK {name}")
except AssertionError as e:
failures += 1
print(f"FAIL {name}: {e}")
print(
"\n" + ("TLRU_TESTS_OK" if not failures else f"TLRU_TESTS_FAILED ({failures})")
)
raise SystemExit(1 if failures else 0)