Optimizing LLM Context Caching via Dynamic Radix Trees in Custom Rust Inference Servers

In production LLM applications, multi-turn conversations and agentic workflows are plagued by a silent latency and cost killer: the redundant processing of prefix tokens. When a user sends a new message in a chat session containing 8,000 tokens of system prompt, retrieved context documents, and prior turn history, standard inference engines re-evaluate the entire prompt sequence from scratch. This prefill phase is computationally expensive ($O(N^2)$ attention complexity) and leads to massive Time-to-First-Token (TTFT) spikes—often exceeding 1.5 seconds on modern GPU clusters. By implementing a dynamic radix tree cache directly within a custom Rust inference server, we can map logical prefix token segments to physical GPU Key-Value (KV) cache blocks. This allows incoming requests with overlapping prefixes to bypass the prefill phase entirely, dropping TTFT to under 50 milliseconds and boosting system throughput by over 2x.

The Anatomy of LLM Context Latency and KV Caching

To understand why prefix caching is critical, we must look at how transformer-based models process tokens. LLM inference is split into two phases:

1. Prefill Phase: The engine processes all input prompt tokens in parallel, computing their Key-Value projections and storing them in the KV cache. This phase is highly compute-bound due to dense matrix operations.
2. Decode Phase: The engine generates tokens one by one. Each new token requires the KV cache of all previous tokens to compute the self-attention layer. This phase is highly memory-bandwidth bound.

For a model like Llama-3-70B running at FP8 precision, the KV cache footprint is substantial. Each token requires:

\[\text{KV Cache Size per Token} = 2 \times (\text{layers}) \times (\text{kv\_heads}) \times (\text{head\_dim}) \times (\text{bytes per element})\]

For Llama-3-70B (80 layers, 8 KV heads, 128 head dimension, 1 byte for FP8):

\[\text{Size} = 2 \times 80 \times 8 \times 128 \times 1 = 163,840 \text{ bytes (160 KB) per token}\]

For an 8,000-token context, a single user session consumes 1.28 GB of VRAM. In a multi-turn chat, if the system prompt and conversation history are re-processed at every turn, the GPU spends precious FLOPS re-computing the exact same KV values.

Static caching mechanisms fall short because conversations are dynamic. Users branch off, agent loops inject different tool call outputs, and context windows slide. We need a dynamic cache that matches arbitrary prefixes on the fly and integrates directly with a paged memory manager like PagedAttention.

Designing the Logical Radix Tree for Tokens

A radix tree (or compact trie) is the natural data structure for prefix matching. Unlike standard tries where each edge represents a single character, radix tree edges represent sequences of characters—or in our case, sequences of token IDs.

We index our radix tree by token IDs (u32) rather than raw UTF-8 strings. String-based caching is fragile: tokenizers can group whitespace and characters differently depending on the preceding text, causing logical cache misses on identical text. By caching token IDs directly, we align the cache state with the exact input tensor fed to the GPU.

Each node in our dynamic radix tree represents a segment of tokens that has been processed and resides in physical GPU memory. The logical tree acts as a mapping layer over the physical memory block pool.

Prefix Traversal and Matching Algorithm

When a request arrives, the server tokenizes the input prompt into a flat array of TokenIds. We traverse the radix tree starting from the root to find the longest matching path.

If the incoming request partially matches a node’s prefix, we must perform a node split. The node split divides the parent node at the mismatch index, creating a child node to hold the suffix of the original prefix and its associated physical GPU block references. This is a critical operation: physical blocks are mapped at a specific granularity (typically 16 tokens per block). When we split a node, we must also split its physical block list along block boundaries.

Physical Paged Cache Management & Mapping

To avoid memory fragmentation, we adopt a virtual memory scheme similar to operating system paging. VRAM is divided into a pool of physical blocks of a fixed size (e.g., 16 tokens).

The BlockAllocator manages these physical blocks. It is decoupled from the logical tree structure, maintaining a free list of indices. When the radix tree needs to cache new tokens, it requests blocks from this allocator. Conversely, when nodes are evicted, their blocks are returned to the free pool.

Thread Safety and Lock Contention in High-Throughput Servers

In a production server handling hundreds of concurrent tokens/sec, the cache lookup and insertion paths are hot spots. A naive design using a single global Mutex over the entire radix tree introduces severe lock contention, turning your multi-core GPU server into a single-threaded bottleneck.

To solve this, we use fine-grained locking. Each node in the tree is wrapped in a parking_lot::RwLock. Readers can traverse down different branches of the tree concurrently without blocking each other. We only acquire a write lock when a node split is required or when inserting a new child branch.

The lookup operation traverses the tree, dropping read locks as it progresses down the path to avoid holding locks longer than necessary.

Eviction Policies: LRU under GPU Memory Pressure

VRAM is a finite, expensive resource. When the physical block pool runs dry, we must evict cached prefixes to reclaim physical block space.

We implement a Least Recently Used (LRU) eviction strategy. However, we cannot evict nodes that are currently active. Active requests pin their matched nodes by incrementing the node’s ref_count. If a node’s ref_count > 0, it and its ancestors are protected from eviction.

Our eviction pipeline executes the following steps:

1. Search the tree recursively for candidate leaf nodes where ref_count == 0.
2. Sort candidates by their last_accessed timestamp in ascending order.
3. Select the oldest node, return its physical blocks back to the BlockAllocator, and prune the node from the parent’s children map.
4. If freeing a node turns its parent into an unused leaf node, the parent becomes eligible for eviction in the next cycle.

Production Realities, Edge Cases, and Performance Metrics

Deploying context caching in production exposes several critical edge cases that can compromise cache efficiency or cause runtime errors.

1. Tokenizer Non-Determinism & Formatting Sensitivity

A major failure mode occurs when application code dynamically formats prompt templates. A single extra whitespace character, a change in system template formatting (e.g., swapping \n for \r\n), or subtle variations in ChatML markers can alter the tokenization sequence of the entire prompt. Because a radix tree relies on an exact sequence prefix, even a single token mismatch at position 5 invalidates the cache for the remaining 8,000 tokens.

Production Fix: Enforce strict validation on prompt schemas. Sanitize all incoming text fields to remove trailing/leading whitespace and standardize delimiters at the API gateway layer before tokenization.

2. CUDA Memory Pinning & Page Swapping

When the GPU memory is full, the engine cannot allocate block arrays for new requests. If you block the inference thread to wait for LRU eviction, you degrade concurrency. In highly loaded servers, we pre-allocate the entire KV block array in CUDA virtual memory at startup using cudaMalloc or Triton’s custom allocator, bypassing standard OS allocations during runtime.

3. Comparison with vLLM / SGLang

Modern runtimes like vLLM and SGLang offer built-in prefix caching. However, implementing a custom dynamic radix tree cache in Rust provides two main advantages:

• Application Routing Integration: By maintaining the radix cache state on a lightweight Rust gateway, we can route incoming requests to specific GPU worker nodes that already have the matched system prompts loaded in VRAM, eliminating the inter-GPU KV cache transfer overhead.
• Custom Eviction Logic: We can write domain-specific eviction policies (e.g., prioritizing cache retention for high-paying user tiers or long-running agent loops over one-off API calls).

4. Benchmark Performance Metrics

In our production testing using Llama-3-70B (FP8) on a node with 8x H100 GPUs (80GB VRAM), we measured performance improvements under a simulated multi-turn workload (average prompt size: 6,144 tokens, average response size: 256 tokens, block size: 16 tokens):

Metric	Without Dynamic Caching	With Dynamic Radix Tree Caching (Hit)	Improvement
Time-to-First-Token (TTFT)	1,240 ms	42 ms	29.5x Reduction
P99 Decode Latency	22 ms/tok	22 ms/tok	Unchanged
Maximum Request Throughput	18 req/sec	41 req/sec	2.27x Speedup
Average GPU Compute Cost	$0.045 / 1k tok	$0.024 / 1k tok	46.6% Cost Saving

By skipping the expensive attention calculation on reused system contexts, the engine scales throughput linearly with GPU compute constraints, shifting the bottleneck exclusively to memory bandwidth during the decode generation phase. Writing the routing and cache manager in Rust guarantees that the latency overhead of checking the cache tree remains under 150 microseconds, ensuring that cache operations never impact the inference path.