Skip to content
mini-sglang from Scratch Stage 2 · Scheduling it well: many requests at once

The radix cache: reusing prefixes across requests

In a multi-turn chat the N-th prompt contains everything from the first N−1 turns; many applications put the same system prompt in front of every request; an agent carries an ever longer history at every step. These requests share a long opening, and so does their KV. SGLang's RadixAttention remembers every prefix it has computed in a radix tree, and a new request first looks for the longest match and reuses whatever it hits. It is SGLang's most characteristic design, and mini-sglang implements it in 237 lines.

Self-test: if you can answer these, skip the chapter

  1. How does a radix tree differ from an ordinary trie?
  2. When can a node be evicted? Why only leaves?
  3. Does matching a prefix modify the tree?
  4. With a page size of 16, how much of a 30-token prompt can be cached?
Answers (try it yourself first, then expand)
  1. A trie stores one token per node; a radix tree compresses a run of tokens with no branching into a single node that holds a stretch of tokens and their KV locations, so matching compares a whole stretch at once and there are far fewer nodes.
  2. When it is not locked by a running request (ref_count == 0) and it is a leaf. An internal node is a prefix of other nodes, and evicting it would make its children unreachable; evicting only leaves keeps the tree a valid prefix structure (and a parent may become the new leaf once its child is gone).
  3. Yes: when the match ends in the middle of a node, that node is split in two so the handle points exactly at where the match ended.
  4. 16 tokens: insertion only caches whole pages, so only the first page of the 30 tokens goes into the cache.

Files you will write: kvcache/radix_cache.py, fast_compare_key in kernel/torch_ops.py, and the radix registration in kvcache/__init__.py.

Files so far: 62 (1 new, 2 modified, 0 brought in early)
minisgl/
├── __init__.py
├── core.py
├── env.py
├── attention/  (4 files, unchanged)
├── distributed/  (3 files, unchanged)
├── engine/  (5 files, unchanged)
├── kernel/
│   ├── __init__.py
│   ├── torch_ops.py  ← modified in this chapter
├── kvcache/
│   ├── __init__.py  ← modified in this chapter
│   ├── base.py
│   ├── mha_pool.py
│   ├── naive_cache.py
│   ├── radix_cache.py  ← new in this chapter
├── layers/  (8 files, unchanged)
├── llm/  (2 files, unchanged)
├── message/  (4 files, unchanged)
├── models/  (7 files, unchanged)
├── moe/  (2 files, unchanged)
├── scheduler/  (9 files, unchanged)
├── utils/  (8 files, unchanged)
This step's main: examples/ch09_radix.py — it uses only the files above; python tools/steps.py check rebuilds this tree chapter by chapter and runs it.

Animation: insertion, splitting, matching, locking and eviction in the radix cache (about 2 minutes, Chinese narration and subtitles)

The shape of the tree

Each node of the tree holds a stretch of tokens (key) and where those tokens' KV lives in the pool (value, token-level locations like the page table). Walking from the root to a node and concatenating the keys gives a cached prefix, and concatenating the values gives its KV locations.

Unlike a trie that creates a node per token, a radix tree compresses "a run with no branching" into one node. A 4000-token system prompt is 4000 nodes in a trie and one node in a radix tree, compared in a single stretch.

minisgl/kvcache/radix_cache.py · RadixTreeNode
class RadixTreeNode:
    counter: int = 0

    def __init__(self, key_fn: KeyFn, tic: int | None = None) -> None:
        self.key_fn = key_fn
        self.children: Dict[Any, RadixTreeNode] = {}
        self._parent: RadixTreeNode | None = None
        self.ref_count = 0
        self.uuid = RadixTreeNode.counter
        RadixTreeNode.counter += 1
        self.timestamp = tic or time.monotonic_ns()
        self._key: torch.Tensor
        self._value: torch.Tensor
        self._length: int

    def set_key_value(self, key: torch.Tensor, value: torch.Tensor) -> None:
        assert len(key) == len(value)
        self._key, self._value, self._length = key, value, len(key)

    def set_parent(self, parent: RadixTreeNode) -> None:
        self._parent = parent
        parent.children[self.key_fn(self._key)] = self  # children are indexed by the first page's tokens

    @property
    def length(self) -> int:
        return self._length

    @property
    def parent(self) -> RadixTreeNode:
        assert self._parent is not None
        return self._parent

    @property
    def value(self) -> torch.Tensor:
        return self._value

    def is_root(self) -> bool:
        return self._parent is None

    def is_leaf(self) -> bool:
        return len(self.children) == 0

    def get_match_len(self, input_ids: torch.Tensor) -> int:
        return fast_compare_key(self._key, input_ids)

    def split_at(self, pos: int) -> RadixTreeNode:
        """Cut this node in two at pos: the new first half becomes this node's parent, and is returned."""
        assert 0 < pos < self.length
        parent = self.parent
        new_node = RadixTreeNode(self.key_fn, self.timestamp)
        new_node.set_key_value(self._key[:pos], self._value[:pos])
        new_node.set_parent(parent)  # replace the entry in parent.children that pointed at this node
        new_node.ref_count = self.ref_count
        self.set_key_value(self._key[pos:], self._value[pos:])
        self.set_parent(new_node)
        return new_node

    def __lt__(self, other: RadixTreeNode) -> bool:  # for the heap: the older the timestamp, the sooner it is evicted
        return self.timestamp < other.timestamp
- children is indexed by the first token of a child's key (when page size is 1) or the tuple of the first page's tokens: under one node no two children can start the same way, so a single dictionary lookup finds the only child that could match; - ref_count: how many running requests depend on this node. It cannot be evicted while this is above 0; - timestamp: when it was last visited, for LRU eviction; - split_at(pos): cuts the node in two at pos, inserting the new first half between the parent and itself.

Walking the tree

Matching and insertion both start from the same function:

minisgl/kvcache/radix_cache.py · RadixPrefixCache._tree_walk
def _tree_walk(self, input_ids: torch.Tensor) -> Tuple[RadixTreeNode, int]:
    prefix_len, total = 0, len(input_ids)
    node = self.root_node
    tic = time.monotonic_ns()
    while prefix_len < total:
        child = node.children.get(self.key_fn(input_ids[prefix_len:]))
        if child is None:
            return node, prefix_len
        node = child
        # reaching a child means at least its first page matched; the real match length is then rounded down to a page
        match_len = align_down(node.get_match_len(input_ids[prefix_len:]), self.page_size)
        prefix_len += match_len
        if match_len != node.length:  # only part of the node matched: split it there
            node = node.split_at(match_len)
            node.timestamp = tic
            return node, prefix_len
        node.timestamp = tic  # a visited node refreshes its timestamp (LRU)
    return node, prefix_len
Starting at the root, look up a child by the start of what is left of the input; if one is found, compare how long a prefix that child's key shares with the input (fast_compare_key, rounded down to a page). A complete match walks on; a partial match splits the node at the matching point and returns the first half.

So matching does change the shape of the tree. It adds and removes nothing from the cache, but it may split one node into two so the returned handle points exactly at "where the match ended". That matters for the locking that follows, because what gets locked is precisely what was hit, no more and no less.

fast_compare_key finds the position of the first difference; upstream implements it in C++ with std::mismatch (bound through tvm-ffi), and we use PyTorch:

minisgl/kernel/torch_ops.py · fast_compare_key
def fast_compare_key(x: torch.Tensor, y: torch.Tensor) -> int:
    """How many leading elements two one-dimensional integer tensors share (the index of the first difference)."""
    n = min(len(x), len(y))
    diff = (x[:n] != y[:n]).nonzero()
    return int(diff[0, 0]) if len(diff) > 0 else n

Matching, inserting, locking

minisgl/kvcache/radix_cache.py · RadixPrefixCache.match_prefix
def match_prefix(self, input_ids: torch.Tensor) -> MatchResult:
    node, prefix_len = self._tree_walk(input_ids)
    return MatchResult(RadixCacheHandle(prefix_len, node))
minisgl/kvcache/radix_cache.py · RadixPrefixCache.insert_prefix
def insert_prefix(self, input_ids: torch.Tensor, indices: torch.Tensor) -> InsertResult:
    insert_len = align_down(len(input_ids), self.page_size)  # only whole pages are cached
    input_ids, indices = input_ids[:insert_len], indices[:insert_len]
    node, prefix_len = self._tree_walk(input_ids)
    if prefix_len != insert_len:  # there is still an uncached part: hang on a new leaf
        new_node = RadixTreeNode(self.key_fn)
        new_node.set_key_value(input_ids[prefix_len:], indices[prefix_len:].clone())
        new_node.set_parent(node)
        self.evictable_size += new_node.length
        node = new_node
    return InsertResult(prefix_len, RadixCacheHandle(insert_len, node))
Insertion only caches whole pages (align_down): with a page size of 16, only the first 16 tokens of a 30-token prompt enter the cache. After walking to the end of the existing prefix, the rest is hung on as a new leaf. The cached_len in the return value is "how much was already in the tree before this insertion", which the caller (CacheManager.cache_req) uses to free its own duplicate pages.

minisgl/kvcache/radix_cache.py · RadixPrefixCache.lock_handle
def lock_handle(self, handle: BaseCacheHandle, unlock: bool = False) -> None:
    assert isinstance(handle, RadixCacheHandle)
    node = handle.node
    while not node.is_root():  # walk from the node that was hit up to the root: the request depends on the whole path
        if unlock:
            node.ref_count -= 1
            assert node.ref_count >= 0
            if node.ref_count == 0:
                self.evictable_size += node.length
                self.protected_size -= node.length
        else:
            if node.ref_count == 0:
                self.evictable_size -= node.length
                self.protected_size += node.length
            node.ref_count += 1
        node = node.parent
Locking walks from the handle's node all the way to the root, adding one to each node's ref_count along the way, because a request depends on the whole prefix and not just its last stretch. When a node's ref_count goes from 0 to 1, its length moves from "evictable" to "protected". The two counts together are the total number of tokens in the cache, and CacheManager.available_size is exactly "free pages plus evictable".

Eviction

minisgl/kvcache/radix_cache.py · RadixPrefixCache.evict
def evict(self, size: int) -> torch.Tensor:
    if size == 0:
        return self.empty_tensor
    assert size <= self.evictable_size, f"Cannot evict {size}, only {self.evictable_size}"
    leaves = self._collect_evictable_leaves()
    heapq.heapify(leaves)
    evicted: List[torch.Tensor] = []
    evicted_size = 0
    while evicted_size < size:
        assert leaves, f"Cannot evict enough cache: need {size}, got {evicted_size}"
        node = heapq.heappop(leaves)
        evicted_size += node.length
        evicted.append(node.value)
        self.evictable_size -= node.length
        parent = node.parent
        del parent.children[self.key_fn(node._key)]
        if parent.is_leaf() and parent.ref_count == 0:  # the parent has become an evictable leaf
            heapq.heappush(leaves, parent)
    return torch.cat(evicted)
Only a leaf with ref_count 0 can be evicted: evicting an internal node would strip its children of their prefix and make them unreachable. Every evictable leaf goes into a min-heap by timestamp, and the least recently visited is popped each time; if its parent thereby becomes a leaf and is unlocked, that goes into the heap too. Eviction works a node at a time, so it may free more than was asked for.

An example

insert [1,2,3,4], [1,2,5,6], [7,8] root ref=1 [1,2] ref=0 [3,4] ref=0 [5,6] ref=0 [7,8] ref=0 6 evictable · 0 protected match [1,2,3,9] and lock it root ref=1 [1,2] ref=1 [3] ref=1 [4] ref=0 [5,6] ref=0 [7,8] ref=0 [3,4] splits ref+1 along the path evict 2, then 1 more root ref=1 [1,2] ref=1 [3] ref=1 LRU: [4] and [5,6] go first, then [7,8] protected (ref>0) evictable leaf
insertion, matching (splitting and locking), LRU eviction

examples/ch09_radix.py
"""Chapter 9: the radix cache's tree structure, and what prefix reuse buys end to end."""

import torch
from minisgl.core import Context, SamplingParams, set_global_ctx
from minisgl.env import ENV
from minisgl.kvcache.radix_cache import RadixPrefixCache


def show(node, depth=0):
    for child in sorted(node.children.values(), key=lambda n: n._key.tolist()):
        print(f"{'    ' * depth}└─ key={child._key.tolist()} value={child.value.tolist()} ref={child.ref_count}")
        show(child, depth + 1)


def t(*xs):
    return torch.tensor(xs, dtype=torch.int32)


set_global_ctx(Context(page_size=1))
cache = RadixPrefixCache(torch.device("cpu"))
cache.insert_prefix(t(1, 2, 3, 4), t(10, 11, 12, 13))
cache.insert_prefix(t(1, 2, 5, 6), t(10, 11, 20, 21))
cache.insert_prefix(t(7, 8), t(30, 31))
print("after inserting [1,2,3,4], [1,2,5,6], [7,8]:")
show(cache.root_node)
handle = cache.match_prefix(t(1, 2, 3, 9)).cuda_handle
cache.lock_handle(handle)
print(f"match [1,2,3,9]: {handle.cached_len} hit at {handle.get_matched_indices().tolist()}; after locking:")
show(cache.root_node)
print("evictable / protected:", tuple(cache.size_info))
print("evict 2 tokens, locations freed:", cache.evict(2).tolist())
print("evict 1 more token, locations freed:", cache.evict(1).tolist())
print("the tree that is left:")
show(cache.root_node)
cache.check_integrity()

# end to end: one long "system prompt" plus different questions
from minisgl.core import reset_global_ctx
from minisgl.llm import LLM

reset_global_ctx()  # the global context was set by hand above, and the engine sets it again

ENV.DISABLE_OVERLAP_SCHEDULING.value = True
system = "You are a helpful assistant. Answer concisely and accurately. " * 4
questions = ["What is 2 + 2?", "Name a primary color.", "What is the capital of Japan?", "Spell cat."]
for cache_type in ("naive", "radix"):
    llm = LLM("models/Qwen3-0.6B", dtype=torch.float32, max_running_req=8, num_page_override=1024,
              max_seq_len_override=256, cache_type=cache_type)
    computed = []
    orig = llm._forward
    llm._forward = lambda fi: (computed.append(sum(r.extend_len for r in fi.batch.reqs))
                               if fi.batch.is_prefill else None, orig(fi))[1]
    first = llm.generate([system + questions[0]], SamplingParams(max_tokens=4))  # one request first
    rest = llm.generate([system + q for q in questions[1:]], SamplingParams(max_tokens=4))
    total = sum(len(llm.tokenizer(system + q).input_ids) for q in questions)
    print(f"{cache_type:<5}: {total} prompt tokens in all, prefill actually computed {sum(computed)}")
    llm.shutdown()
What it prints
after inserting [1,2,3,4], [1,2,5,6], [7,8]:
└─ key=[1, 2] value=[10, 11] ref=0
    └─ key=[3, 4] value=[12, 13] ref=0
    └─ key=[5, 6] value=[20, 21] ref=0
└─ key=[7, 8] value=[30, 31] ref=0
match [1,2,3,9]: 3 hit at [10, 11, 12]; after locking:
└─ key=[1, 2] value=[10, 11] ref=1
    └─ key=[3] value=[12] ref=1
        └─ key=[4] value=[13] ref=0
    └─ key=[5, 6] value=[20, 21] ref=0
└─ key=[7, 8] value=[30, 31] ref=0
evictable / protected: (5, 3)
evict 2 tokens, locations freed: [13, 20, 21]
evict 1 more token, locations freed: [30, 31]
the tree that is left:
└─ key=[1, 2] value=[10, 11] ref=1
    └─ key=[3] value=[12] ref=1
naive: 231 prompt tokens in all, prefill actually computed 231
radix: 231 prompt tokens in all, prefill actually computed 73
- After inserting [1,2,3,4] and [1,2,5,6], the common prefix [1,2] becomes a node with two branches below it; - Matching [1,2,3,9] hits 3 tokens: [3,4] splits into [3] and [4], and the handle points at [3]. After locking, [1,2] and [3] have ref_count 1, so 3 tokens are protected and 5 are evictable; - Asked to evict 2 tokens: the evictable leaves are [4], [5,6] and [7,8]. The match only refreshed the timestamps of [1,2] and [3], so the split-off [4] kept its original insertion time and goes first; one is not enough, so the next oldest [5,6] follows, freeing 3 in all; - Asked to evict 1 more: only [7,8] is left to evict. The locked [1,2] and [3] stay.

End to end: four requests share one system prompt, the first prefills it in full and the other three compute only their own questions, which brings the prefill from 231 tokens down to 73.

Fitting into the request lifecycle

Looking back at cache_req from chapter 8, all four of its regions now really exist:

  1. At admission match_req finds the prefix that hit and locks it (matching up to the second-to-last token, so at least one token is computed);
  2. At the end of prefill, cache_req(finished=False) inserts the whole prompt (page-aligned) into the tree, unlocks the old handle and locks the new one; if another request inserted the same prefix in the meantime, this request's duplicate pages are freed right away;
  3. When the request finishes, cache_req(finished=True) inserts the generated text too (the next turn of the conversation will want it), unlocks, and frees the tail that does not fill a page.

A finished request's KV is not freed but stays in the tree as "evictable", and is only really evicted when the space is needed. So memory is always used to cache prefixes as much as it can be.

The official implementation

Difference from upstream: the integrity check

The official RadixPrefixCache.check_integrity is empty (pass). Ours walks the whole tree, recounts the evictable and protected lengths, reconciles them against the counters, and checks the invariant that a child is never locked more times than its parent. The scheduler calls it while idle, so any miscount shows up immediately.

minisgl/kvcache/radix_cache.py · RadixPrefixCache.check_integrity
def check_integrity(self) -> None:
    """Walk the whole tree again to reconcile the counters: the evictable size is the total length of the nodes with ref_count 0."""
    evictable = protected = 0
    stack = list(self.root_node.children.values())
    while stack:
        node = stack.pop()
        assert node.ref_count >= 0
        for child in node.children.values():
            assert child.ref_count <= node.ref_count, "child locked more than its parent"
            stack.append(child)
        if node.ref_count == 0:
            evictable += node.length
        else:
            protected += node.length
    assert (evictable, protected) == (self.evictable_size, self.protected_size), (
        f"size mismatch: tree=({evictable}, {protected}), "
        f"counter=({self.evictable_size}, {self.protected_size})"
    )

Tests

tests/test_ch09_radix.py · test_lock_protects_from_eviction_and_lru_order
def test_lock_protects_from_eviction_and_lru_order():
    c = _cache()
    c.insert_prefix(t(1, 2, 3), t(10, 11, 12))
    c.insert_prefix(t(7, 8), t(20, 21))
    c.insert_prefix(t(4, 5), t(30, 31))
    handle = c.match_prefix(t(7, 8)).cuda_handle  # [7, 8] was just visited
    c.lock_handle(handle)
    assert c.size_info == (5, 2)
    evicted = c.evict(3)  # the locked [7, 8] cannot be evicted, and of the other two [1, 2, 3] has gone unused longest
    assert evicted.tolist() == [10, 11, 12]
    c.lock_handle(handle, unlock=True)
    assert c.evict(2).tolist() == [30, 31]  # once unlocked, [4, 5] has gone unused longer than [7, 8]
    c.check_integrity()
tests/test_ch09_radix.py · test_shared_prefixes_are_reused_and_outputs_unchanged
def test_shared_prefixes_are_reused_and_outputs_unchanged():
    ENV.DISABLE_OVERLAP_SCHEDULING.value = True
    llm = build_llm(QWEN3, cache_type="radix")
    computed = []
    orig = llm._forward

    def spy(fi):
        if fi.batch.is_prefill:
            computed.append(sum(r.extend_len for r in fi.batch.reqs))
        return orig(fi)

    llm._forward = spy
    first = llm.generate(PROMPTS, greedy(6))
    cold = sum(computed)
    computed.clear()
    second = llm.generate(PROMPTS, greedy(6))  # the same prompts again: almost everything hits
    assert [o["token_ids"] for o in first] == [o["token_ids"] for o in second] == hf_greedy(QWEN3, PROMPTS, 6)
    assert sum(computed) == len(PROMPTS)  # each request only has to recompute its last token
    assert cold > 5 * len(PROMPTS)
    llm.cache_manager.check_integrity()
    llm.shutdown()
Running the same 5 prompts a second time, each request has to recompute only its last token, the output matches Hugging Face exactly, and the memory integrity check passes.

How to explain it

On prefix caching: a radix tree is a compressed trie whose node holds a stretch of tokens and their KV locations, with children indexed by the first token (or first page) so a match compares a whole stretch at once; a match can end mid-node, and the node is then split so the handle points exactly at the match, which is why matching also modifies the tree. Only whole pages are cached (with a page size of 16, at most 16 of a 30-token prompt). A running request locks the whole path by incrementing reference counts, and only unlocked leaves are evicted in LRU order; a finished request's KV stays in the tree and is evicted only when space is needed. Against vLLM's hashed-block prefix cache: a radix tree matches at any length and handles multi-turn chat naturally, while hashed blocks look up whole blocks in a table, which is simpler to implement and easier to make distributed or multi-tier.

Exercises

  1. vLLM uses "hashed blocks" for prefix caching: each KV block's hash comes from the previous block's hash plus this block's tokens, and lookup is a dictionary. What are the pros and cons against a radix tree? (See the prefix-caching chapter of the Inference Systems handbook.)
  2. Eviction currently walks the whole tree each time to collect evictable leaves. With a large tree of hundreds of thousands of nodes that is a bottleneck. How would you maintain an always-valid "set of evictable leaves"?
  3. Design a test that puts two requests with the same prefix in one prefill batch and verifies that cache_req frees the duplicate pages correctly.
Answers
  1. Hashed blocks: simple to implement, lookup is O(blocks) dictionary accesses, naturally block-aligned, and easy to share across processes (KV events); but the hashing starts from the beginning every time and matching is only block-granular. A radix tree: compares a whole stretch at once, stores a shared prefix only once, and knows the dependencies when evicting; but it is more complex, and splitting and merging nodes needs careful bookkeeping.
  2. Keep a set (or a heap ordered by timestamp with lazy deletion): add a node when it becomes a leaf with ref_count 0, and remove it when it is locked, gains a child or is evicted. SGLang 0.5.20's RadixCache does exactly this, maintaining an evictable_leaves set updated by _update_leaf_status after every lock, unlock, insertion and deletion.
  3. Use "The capital of France is" and "The capital of France is a city that" from PROMPTS: neither hits in the same prefill batch, the first to finish inserts the prefix, and the second sees cached_len > 0 on insertion and takes the "free the duplicate pages" branch. A passing check_integrity() afterwards shows there is no leak, and test_shared_prefixes_are_reused_and_outputs_unchanged already covers this.

Summary

  • A radix-tree node holds a stretch of tokens and their KV locations; children are indexed by the first token (or first page), and a match compares a whole stretch at once.
  • A match may split a node so the handle points exactly at where it ended; insertion caches only whole pages and hangs the new part on as a leaf.
  • Locking adds one to ref_count along the whole path; only unlocked leaves are evicted in LRU order, a node at a time.
  • A finished request's KV stays in the tree and is evicted only when space is needed; with a shared system prompt the prefill work drops sharply.

Exercises for this chapter (1)

Write the code and run the tests in your browser, or judge locally with python practice/judge.py. The exercises are in Chinese for now.

Found a mistake or a better way to say it? Edit this page on GitHubReport an issue