
SGLang's RadixAttention, Explained
An examination of SGLang's RadixAttention mechanism, focusing on prefix caching, radix trees, and throughput optimization for ML infrastructure.
Large Language Model (LLM) serving is often bottlenecked not by the compute required to generate the next token, but by the memory bandwidth required to fetch the Key-Value (KV) cache for that token’s attention mechanism. When serving a batch of requests that share common prefixes—such as a system prompt, a code template, or a document context—standard inference engines recompute or re-fetch these identical states repeatedly. SGLang addresses this through a mechanism called RadixAttention, which leverages a radix tree (a compressed trie) to structure the KV cache, enabling zero-copy reuse of computation results across concurrent requests.
The Inefficiency of Linear Prefixes
To understand RadixAttention, one must first understand the cost of the Attention mechanism. In Transformer architectures, every generated token attends to all previous tokens. This requires storing the Key and Value vectors for each token in a KV cache. During decoding, the model computes the attention scores between the current token and all previous KV pairs.
Consider a scenario where you are serving 100 requests that all begin with the same system prompt: "You are a helpful assistant. Answer the following questions..." followed by unique user queries. In a naive implementation, the KV cache for the system prompt is computed once for the first request. For the subsequent 99 requests, the engine must either recompute these vectors (wasting GPU FLOPs) or fetch them from memory again (wasting HBM bandwidth), even though the values are identical.
Standard prefix caching solutions often use a hash table or an LRU cache to store these prefixes. However, they treat each prefix as an isolated block. If Request A has prefix ["A", "B", "C"] and Request B has prefix ["A", "B", "D"], a simple cache might store both separately if they diverge early, or require complex merging logic. SGLang’s approach is structural: it builds a tree of all executed prefixes.
The Radix Tree as a Physical Memory Map
SGLang organizes the KV cache into a radix tree. A radix tree is a space-optimized trie where nodes with only one child are merged with their parent. In the context of LLM serving, each node in this tree represents a sequence of tokens, and the edges represent the transitions between them.
The critical innovation is that memory is allocated at the node level, not the request level. When a request starts, SGLang traverses the existing radix tree. If the request’s prefix matches an existing path in the tree, the engine does not allocate new memory or compute new KV states. Instead, it attaches the request’s execution state to the existing nodes.
For example:
- Request 1 generates tokens
["Hello", "World"]. The engine creates a path:Root -> ["Hello"] -> ["Hello", "World"]. - Request 2 starts with
["Hello", "World", "How"]. The engine traversesRoot -> ["Hello"] -> ["Hello", "World"], finds the nodes already exist, and only computes the KV state for["How"], appending it to the shared path.
This structure allows for fine-grained sharing. Unlike a flat cache that might store ["Hello", "World"] as a single block, the radix tree preserves the hierarchy. This is crucial for fine-grained scheduling. SGLang allocates memory at the granularity of individual tokens rather than fixed-size blocks, and the radix tree maps these token sequences directly to their KV cache locations in GPU memory. When a request is preempted or scheduled, the engine can evict or reuse these cache entries without recomputing the attention keys and values.
How RadixAttention Operates
The term "RadixAttention" refers to the specific attention operator SGLang uses to handle this non-contiguous memory layout. Standard attention kernels like FlashAttention assume a contiguous sequence of tokens in memory. RadixAttention must gather KV pairs from potentially scattered locations in the radix tree.
The mechanism works in three steps:
- Path Traversal: For a given request, the scheduler identifies the sequence of token positions that constitute the current context in the radix tree.
- Gather Operation: Instead of copying data into a contiguous buffer, RadixAttention performs a gather operation. It fetches the KV tensors for the required tokens directly from global memory. This is optimized using CUDA kernels that handle strided and non-contiguous reads efficiently.
- Attention Computation: The attention mechanism computes scores using these gathered KV pairs. Because the KV states are immutable (once computed for a prefix, they don’t change), the engine can safely reuse them across multiple requests without synchronization overhead.
This approach eliminates the need to materialize the full context in a single contiguous array before attention, reducing memory bandwidth pressure. The attention computation itself becomes a function of the unique prefix length plus the new tokens, rather than the total sequence length repeated across all requests.
Throughput and Memory Trade-offs
The primary benefit of RadixAttention is a dramatic reduction in effective compute cost for shared prefixes. If 90% of your batch shares the same system prompt, you only compute attention for the last 10% of the context for 90% of the requests. This can lead to 2x–5x throughput improvements in workloads with high prefix overlap, such as API services with static prompts or code generation tasks with common boilerplate.
However, this optimization introduces complexity:
- Memory Fragmentation: The radix tree can lead to fragmented memory usage if blocks are not managed carefully. SGLang uses a block allocator to mitigate this, but peak memory usage can be higher than contiguous caching due to overhead.
- Scheduling Complexity: The scheduler must maintain the integrity of the radix tree. If a request is preempted, the engine must ensure that no other request is currently using the blocks it needs to evict. This requires reference counting or atomic operations on block states.
- Kernel Overhead: The gather operation in RadixAttention has higher latency than contiguous memory access. If the prefix overlap is low (e.g., <50%), the overhead of managing the tree and gathering non-contiguous data may outweigh the compute savings.
Conclusion
SGLang’s RadixAttention is a specialized optimization for high-concurrency LLM serving where prefix overlap is significant. By treating the KV cache as a shared radix tree rather than a collection of independent sequences, it transforms redundant computation into shared memory access. For infrastructure engineers, this means that workload design—specifically, maximizing prefix reuse through prompt templating—can directly impact hardware efficiency. The mechanism is not a universal fix for all inference patterns, but for the specific case of batched, prefix-heavy requests, it represents a fundamental shift from recomputing to reusing.