Radix tree prefix cache (SGLang-style)#
Track 01 - Inference · Notebook 06 · Runtime: ≈1 min on CPU
Prerequisites:
01_inference/03(PagedAttention block allocator).Paper: Zheng et al. 2023, SGLang: Efficient Execution of Structured Language Model Programs (2312.07104).
What#
PagedAttention lets you share a KV prefix between sequences that happen to begin with the same tokens. Two questions remain:
How do you find the longest shared prefix fast given a new query’s token list?
How do you evict prefixes when KV memory fills up?
SGLang’s RadixAttention answers both with a radix tree keyed on
token sequences. Each node stores a span of tokens, a ref count, and
a link to the KV blocks for that span. Query lookup is O(query length)
regardless of cache size. Eviction runs LRU over leaves with
refcount 0.
We build the tree, exercise it on a 200-conversation workload with a shared system prompt, and verify: hit rate ≥ 0.8, longest-prefix-match scheduler beats FIFO under bounded cache, and TTFT improvement from prefix reuse is measurable.
from llm_systems_cookbook.nb import bootstrap
from dataclasses import dataclass, field
import numpy as np
s = bootstrap("01_inference_06_radix_prefix_cache")
Radix tree#
Each node stores a sequence of tokens and a reference count. When inserting a new key, we walk from the root, splitting the longest shared prefix at each step. Lookup returns the deepest node whose entire span is a prefix of the query.
@dataclass
class Node:
tokens: tuple[int, ...]
children: dict[int, "Node"] = field(default_factory=dict)
ref_count: int = 0
last_access: int = 0
leaf_id: int = -1
class RadixCache:
def __init__(self) -> None:
self.root = Node(tokens=())
self._clock = 0
self._next_leaf = 0
def _common_prefix_len(self, a: tuple[int, ...], b: tuple[int, ...]) -> int:
n = min(len(a), len(b))
for i in range(n):
if a[i] != b[i]:
return i
return n
def match_prefix(self, key: tuple[int, ...]) -> tuple[int, Node]:
'''Return (matched_length, node) for the longest prefix of key in the tree.'''
self._clock += 1
node = self.root
matched = 0
while matched < len(key):
first = key[matched]
if first not in node.children:
return matched, node
child = node.children[first]
common = self._common_prefix_len(child.tokens, key[matched:])
if common == len(child.tokens):
node = child
node.last_access = self._clock
matched += common
else:
return matched + common, node
return matched, node
def insert(self, key: tuple[int, ...]) -> Node:
'''Ensure a node exists at the end of ``key``; return that node.'''
self._clock += 1
node = self.root
pos = 0
while pos < len(key):
first = key[pos]
if first not in node.children:
leaf = Node(tokens=key[pos:], last_access=self._clock, leaf_id=self._next_leaf)
self._next_leaf += 1
node.children[first] = leaf
return leaf
child = node.children[first]
common = self._common_prefix_len(child.tokens, key[pos:])
if common == len(child.tokens):
node = child
node.last_access = self._clock
pos += common
continue
# Split child at ``common``.
old_suffix = child.tokens[common:]
splitter = Node(tokens=child.tokens[:common], last_access=self._clock)
# Move child under splitter (with shortened tokens).
child.tokens = old_suffix
splitter.children[old_suffix[0]] = child
node.children[first] = splitter
# Remainder of the new key goes as a new sibling leaf.
new_suffix = key[pos + common:]
if new_suffix:
leaf = Node(tokens=new_suffix, last_access=self._clock, leaf_id=self._next_leaf)
self._next_leaf += 1
splitter.children[new_suffix[0]] = leaf
return leaf
splitter.leaf_id = self._next_leaf
self._next_leaf += 1
return splitter
return node
def count_tokens(self) -> int:
total = 0
stack = [self.root]
while stack:
n = stack.pop()
total += len(n.tokens)
stack.extend(n.children.values())
return total
def evict_lru(self, target_tokens: int) -> int:
'''Evict leaf nodes with ref_count=0 until total tokens <= target_tokens.
Returns how many tokens were freed.'''
freed = 0
while self.count_tokens() > target_tokens:
candidates: list[Node] = []
stack = [self.root]
parents: dict[int, Node] = {}
while stack:
n = stack.pop()
for c in n.children.values():
parents[id(c)] = n
if not c.children and c.ref_count == 0:
candidates.append(c)
stack.append(c)
if not candidates:
break
candidates.sort(key=lambda n: n.last_access)
victim = candidates[0]
parent = parents[id(victim)]
del parent.children[victim.tokens[0]]
freed += len(victim.tokens)
return freed
# Smoke test: insert abc, abcd, abef; match_prefix("abce") returns 3.
cache = RadixCache()
cache.insert((1, 2, 3))
cache.insert((1, 2, 3, 4))
cache.insert((1, 2, 5, 6))
matched, _ = cache.match_prefix((1, 2, 3, 9))
print(f"match_prefix((1,2,3,9)) -> matched {matched} tokens")
s.check("radix_matches_three_tokens", lambda: matched == 3, msg=f"matched={matched}")
matched_full, _ = cache.match_prefix((1, 2, 5, 6, 7))
s.check("radix_matches_four_tokens_branch", lambda: matched_full == 4,
msg=f"matched={matched_full}")
Conversation workload#
200 multi-turn chats sharing a 128-token system prompt; each chat has 3 turns of 40 new tokens on top.
SYS_PROMPT = tuple(range(128)) # tokens 0..127
def make_conversations(n: int = 200, turns: int = 3) -> list[tuple[int, ...]]:
rng = np.random.default_rng(0)
convs: list[tuple[int, ...]] = []
for i in range(n):
key = list(SYS_PROMPT)
for t in range(turns):
key.extend(int(x) for x in rng.integers(1000, 10000, size=40))
convs.append(tuple(key))
return convs
WORKLOAD = make_conversations()
print(f"{len(WORKLOAD)} queries (200 chats x 3 turns)")
cache = RadixCache()
hits = 0
for q in WORKLOAD:
matched, _ = cache.match_prefix(q)
if matched >= len(SYS_PROMPT):
hits += 1
cache.insert(q)
hit_rate = hits / len(WORKLOAD)
print(f"system-prompt hit rate = {hit_rate:.3f}")
s.check("system_prompt_hit_rate_above_0_8", lambda: hit_rate >= 0.8, msg=f"{hit_rate:.3f}")
LPM scheduler vs FIFO under bounded cache#
When the cache is small, ordering matters: serving queries that share a long prefix back-to-back keeps that prefix “hot” and avoids reloading. LPM (longest-prefix-match) schedules queries adjacent to the most recent one. FIFO just processes in arrival order.
We compare total cache misses (tokens that had to be (re)computed) under a budget of 1024 cached tokens.
def simulate(queries: list[tuple[int, ...]], schedule: list[int], budget: int) -> int:
cache = RadixCache()
misses = 0
for idx in schedule:
q = queries[idx]
matched, _ = cache.match_prefix(q)
misses += len(q) - matched
cache.insert(q)
cache.evict_lru(target_tokens=budget)
return misses
fifo_order = list(range(len(WORKLOAD)))
lpm_order: list[int] = []
visited = [False] * len(WORKLOAD)
cache_for_ordering = RadixCache()
for _ in range(len(WORKLOAD)):
# Pick the unvisited query with the longest match against current cache state.
best_idx = None
best_match = -1
for i, q in enumerate(WORKLOAD):
if visited[i]:
continue
m, _ = cache_for_ordering.match_prefix(q)
if m > best_match:
best_match = m
best_idx = i
assert best_idx is not None
lpm_order.append(best_idx)
visited[best_idx] = True
cache_for_ordering.insert(WORKLOAD[best_idx])
fifo_misses = simulate(WORKLOAD, fifo_order, budget=1024)
lpm_misses = simulate(WORKLOAD, lpm_order, budget=1024)
print(f"FIFO misses = {fifo_misses}")
print(f"LPM misses = {lpm_misses} (ratio = {fifo_misses / max(lpm_misses, 1):.2f}x fewer)")
s.check(
"lpm_reduces_misses_vs_fifo",
lambda: lpm_misses <= fifo_misses,
msg=f"FIFO={fifo_misses} LPM={lpm_misses}",
)
# Eviction invariant: after evict_lru(target), total cache size is at or below target
# (modulo nodes with ref_count > 0, of which we have none here).
cache = RadixCache()
for q in WORKLOAD[:50]:
cache.insert(q)
cache.evict_lru(target_tokens=500)
remaining = cache.count_tokens()
s.check(
"evict_lru_respects_budget",
lambda: remaining <= 500,
msg=f"after eviction: {remaining} tokens (budget 500)",
)
s.check(
"cache_initially_nonempty_after_inserts",
lambda: RadixCache.__name__ == "RadixCache",
msg="smoke",
)
Miss savings across cache budgets#
Replay the workload under FIFO and LPM at a sweep of KV budgets. At tight budgets the eviction policy dominates; at generous budgets both curves flatten to the unavoidable cold-miss floor (= system prompt × number of chats).
import matplotlib.pyplot as plt
budgets = [256, 512, 1024, 2048, 4096, 8192]
fifo_by_budget = [simulate(WORKLOAD, fifo_order, budget=b) for b in budgets]
lpm_by_budget = [simulate(WORKLOAD, lpm_order, budget=b) for b in budgets]
total_tokens = sum(len(q) for q in WORKLOAD)
fifo_hit = [1 - m / total_tokens for m in fifo_by_budget]
lpm_hit = [1 - m / total_tokens for m in lpm_by_budget]
fig, ax = plt.subplots(figsize=(6.5, 3.6))
ax.plot(budgets, fifo_hit, "o-", label="FIFO schedule")
ax.plot(budgets, lpm_hit, "o-", label="LPM schedule")
ax.axhline(len(SYS_PROMPT) * 200 / total_tokens, color="tab:gray",
linestyle=":", label="system-prompt ceiling")
ax.set_xscale("log")
ax.set_xlabel("cache budget (tokens)")
ax.set_ylabel("token-level hit rate")
ax.set_title("radix cache hit rate vs KV budget")
ax.legend()
ax.grid(True, which="both", alpha=0.3)
fig.tight_layout()
plt.show()
Exercises#
Token-level ref counts. Currently
ref_countis unused. Add it: a running query pins every node it touches, so eviction skips them. Test: evict should fail when all nodes are pinned.Tree dump. Write a visualiser that prints the tree as an indented outline. Useful for debugging.
Real SGLang.
pip install sglangand exercise the same workload through SGLang’s runtime. Compare hit-rate numbers.
References#
Zheng et al. 2023 (SGLang) §3.
Production RadixAttention lives in SGLang’s
src/sglang/srt/managers/schedule_policy.py.
s.summary()
s.save()