Radix Cache

Instructions

Implement the RadixCache class, which stores sequences of integers (think: token ID sequences) in a radix tree — a compressed prefix tree where each node holds a chunk of the sequence instead of a single element:

• RadixCache() — initializes an empty cache.
• insert(sequence) — adds sequence to the cache. A prefix shared with an already-stored sequence must be held once, in a shared node. When sequence partially overlaps an existing node's chunk, that node must split at the divergence point.
• exists(sequence) — returns True only if sequence was previously inserted exactly. A sequence that is merely a prefix of an inserted sequence, or an extension of one, returns False — mark terminal nodes so complete sequences can be told apart from internal prefixes.

Your implementation must actually compress: giving every integer its own node, or keeping a flat hash set of whole sequences, returns the right values but does not solve the problem.

Example:

rc = RadixCache()

rc.insert([1, 2, 3, 4, 5])
rc.insert([1, 2, 3, 9])       # shares prefix [1, 2, 3] — node splits after 3

rc.exists([1, 2, 3, 4, 5])    # True
rc.exists([1, 2, 3, 9])       # True
rc.exists([1, 2, 3])          # False — internal prefix, never inserted

rc.insert([1, 2, 3])
rc.exists([1, 2, 3])          # True — now a marked terminal
rc.exists([7])                # False

Constraints:

• 1 <= sequence.length <= 10^4
• Values fit in a 64-bit signed integer.
• Up to 10^5 total operations; many sequences share long prefixes.

Function Signature

from typing import List


class RadixCache:
    def __init__(self):
        pass

    def insert(self, sequence: List[int]) -> None:
        pass

    def exists(self, sequence: List[int]) -> bool:
        pass

Follow-ups

More Followup questions in our Private Crash Course. Check it out.

Background

What is the KV cache? What is a prefix cache?

The work being stored is the KV cache. When a model reads a prompt, every token is turned into a key vector and a value vector, kept in GPU memory so later tokens can attend to them.

The property everything else rests on: the vectors for token k depend only on tokens 1 through k, never on what comes after. So the same prefix always produces the same KV cache, no matter what follows it.

[You, are, a, helpful, assistant, .] + [What, is, Rust, ?]
[You, are, a, helpful, assistant, .] + [What, is, Go,   ?]
 └──── identical tokens ──────────┘     └── different, but it changes
       → identical KV cache                 nothing to the left

A prefix cache is the machinery that exploits this: keep the KV cache after a request finishes, and when the next request opens with the same tokens, hand back the stored vectors instead of recomputing them.

Why do we need a radix cache / prefix cache?

Because production traffic repeats prefixes constantly. Let's say a chatbot sends the same 800-token system prompt with every single request.

Without a prefix cache, the server recomputes that identical 800-token prefix on all 10,000 requests.

This turns serving into a lookup problem. The server is not asking "have I seen this exact prompt before?" It is asking "what is the longest prefix of these tokens that I already have cached?"

What is a radix tree?

A plain prefix tree (a trie) answers longest-prefix queries, but it spends one node per token. An 800-token system prompt becomes an 800-node chain, every node carrying pointer and allocation overhead, for a path that never branches.

A radix tree is the same idea with the boring parts squeezed out: any chain of single-child nodes is merged into one node holding a chunk of tokens. That 800-node chain becomes a single node.

That split is the core of the structure. When a new sequence agrees with a node's chunk for a while and then diverges, you cut the node in two: the agreeing part becomes a shared parent, the two tails become its children. The shared prefix ends up stored exactly once.

What is RadixAttention?

RadixAttention, is this tree wired into the inference server itself, mapping token sequences to the KV cache blocks that hold their vectors...