diff --git a/docs/docs/advanced_features/radix_eviction_policy.mdx b/docs/docs/advanced_features/radix_eviction_policy.mdx index 368f2b2e4..787430be3 100644 --- a/docs/docs/advanced_features/radix_eviction_policy.mdx +++ b/docs/docs/advanced_features/radix_eviction_policy.mdx @@ -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. | | `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. | +| `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: @@ -55,7 +56,7 @@ TypeError: SLRUStrategy.__init__() got an unexpected keyword argument 'protected ### 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 @@ -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. +#### `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 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. diff --git a/python/sglang/srt/arg_groups/choices.py b/python/sglang/srt/arg_groups/choices.py index 8ce2c00ec..1ac5b914c 100644 --- a/python/sglang/srt/arg_groups/choices.py +++ b/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"] diff --git a/python/sglang/srt/arg_groups/fields/memory.py b/python/sglang/srt/arg_groups/fields/memory.py index d00108b31..7de6d3a40 100644 --- a/python/sglang/srt/arg_groups/fields/memory.py +++ b/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." ), diff --git a/python/sglang/srt/arg_groups/kv_cache_hook.py b/python/sglang/srt/arg_groups/kv_cache_hook.py index 4d473c79e..8c246858b 100644 --- a/python/sglang/srt/arg_groups/kv_cache_hook.py +++ b/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 " diff --git a/python/sglang/srt/mem_cache/evict_policy.py b/python/sglang/srt/mem_cache/evict_policy.py index dad1625fe..7a6e20981 100644 --- a/python/sglang/srt/mem_cache/evict_policy.py +++ b/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 diff --git a/python/sglang/srt/mem_cache/registry.py b/python/sglang/srt/mem_cache/registry.py index 75afef76b..de35ad014 100644 --- a/python/sglang/srt/mem_cache/registry.py +++ b/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 ( diff --git a/python/sglang/srt/mem_cache/unified_cache/unified_tree_core.py b/python/sglang/srt/mem_cache/unified_cache/unified_tree_core.py index 2b279540c..b105e4f57 100644 --- a/python/sglang/srt/mem_cache/unified_cache/unified_tree_core.py +++ b/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 diff --git a/python/sglang/srt/mem_cache/utils.py b/python/sglang/srt/mem_cache/utils.py index 1dd70fbaf..e44d48f02 100644 --- a/python/sglang/srt/mem_cache/utils.py +++ b/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, } diff --git a/test/registered/unit/mem_cache/test_tlru_eviction_policy.py b/test/registered/unit/mem_cache/test_tlru_eviction_policy.py new file mode 100644 index 000000000..0e56fdf24 --- /dev/null +++ b/test/registered/unit/mem_cache/test_tlru_eviction_policy.py @@ -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)