Optimizing AI Inference with RadixAttention and Prefix Trees

This article provides a quick overview of how RadixAttention optimizes large language model inference by reusing cached key-value prefix trees. We will explore the mechanics of the Key-Value (KV) cache, how organizing this cache into a radix tree structure eliminates redundant computations, and the specific ways this technology accelerates both shared system prompts and multi-turn conversations in artificial intelligence deployments.

Understanding the KV Cache Bottleneck

In autoregressive large language models (LLMs), generating text requires calculating attention scores for every previous token in a sequence. To avoid recomputing these values for every new token generated, AI systems use a Key-Value (KV) cache. This cache stores the intermediate attention states (keys and values) of processed tokens in GPU memory.

In traditional inference systems, the KV cache is allocated per request and immediately discarded once the generation is complete. This approach is inefficient because different requests often share identical starting sequences, forcing the GPU to repeatedly compute the exact same KV states.

How RadixAttention Works

RadixAttention solves this inefficiency by retaining the KV cache in GPU memory after a request finishes and organizing it into a radix tree—a space-optimized prefix tree.

Instead of storing individual tokens at each node, a radix tree stores sequences of tokens along its edges. In RadixAttention, each node and edge points to the corresponding KV cache tensors stored in paged GPU memory. When a new prompt arrives, the system searches the radix tree for the longest matching prefix of tokens. If a match is found, the model immediately reuses the existing KV cache for those tokens and only computes the attention states for the new, unseen tokens.

Reusing Shared System Prompts

In production AI deployments, applications routinely prepend a standard system prompt or few-shot examples to every user query. This shared text can easily span thousands of tokens.

With RadixAttention, the KV cache for the system prompt is computed only once during the very first request. It is stored as the root branch of the radix tree. When thousands of subsequent users send queries, the inference engine matches their requests to the root node. The system bypasses the prefill computation for the system prompt entirely, branching off the root node to process only the unique user input. This drastically reduces the time-to-first-token (TTFT) and saves massive amounts of compute power.

Accelerating Multi-Turn Conversations

Multi-turn chat interactions require the entire conversation history to be sent back to the LLM with every new user message. Without cross-request caching, a ten-turn conversation forces the model to re-process the first nine turns all over again.

RadixAttention treats a multi-turn conversation as a continuously growing path within the radix tree.

  1. Turn 1: The user asks a question, and the model replies. The KV cache for this entire exchange is saved as a branch in the tree.
  2. Turn 2: The user replies. The system receives the Turn 1 history plus the new Turn 2 prompt.
  3. Prefix Matching: RadixAttention identifies that the Turn 1 history already exists in the tree. It reuses that exact KV cache, appends a new child node for the Turn 2 input, and computes only the new tokens.

Furthermore, if a user edits a previous message or regenerates a response, the tree simply creates a new branch from the point of divergence, reusing the conversation history up to that exact token.

Memory Management via LRU Eviction

Because GPU memory is finite, the radix tree cannot grow indefinitely. RadixAttention pairs its prefix tree with a Least Recently Used (LRU) eviction policy.

Active requests lock the nodes they are currently using so they cannot be deleted. When the GPU runs out of memory for new requests, the system automatically evicts the leaf nodes (the ends of conversation branches) that have not been accessed for the longest time. This keeps frequently used prefixes—like global system prompts and active chats—in memory while seamlessly clearing out abandoned conversations.

Cache-Aware Scheduling

To maximize the hit rate of the radix tree, the inference engine uses cache-aware scheduling. Rather than processing incoming requests strictly in the order they arrive (First-In, First-Out), the scheduler analyzes the waiting queue and groups requests that share common prefixes in the radix tree. By executing these matching requests together, the system ensures that shared KV cache data is utilized to its fullest potential before any memory eviction occurs.