TokenLake: Unified Segment-Level Prefix Cache

• TokenLake is a unified segment-level prefix cache pool that manages GPU memory pooling for high-throughput, long-context LLM inference.
• It introduces a declarative interface that enables automatic load balancing, deduplication, and defragmentation across GPU clusters.
• Performance evaluations demonstrate significant throughput gains and hit rate improvements over conventional cache-aware routing systems.

TokenLake is a unified segment-level prefix cache pool designed to provide fine-grained, elastic, long-context LLM serving across GPU clusters. It introduces a declarative interface for managing prefix caches (layer-wise key/value tensors from prefill and decoding phases) at segment granularity, enabling pooled GPU memory, automatic load balancing, deduplication, defragmentation, and minimized inter-node communication. By fully decoupling cache management from compute scheduling, TokenLake addresses the fundamental inefficiencies of prior cache-aware routing and phase-disaggregation approaches in multi-turn, high-throughput LLM inference contexts (Wu et al., 24 Aug 2025).

1. Limitations of Prior Prefix Caching Systems

Large-context LLM deployments typically rely on prefix caching to avoid redundant computation in multi-turn interactions, where shared prefixes may extend over tens or hundreds of thousands of tokens. Traditional systems fall into two paradigms:

• Cache-Aware Routing: Each GPU maintains a local prefix-tree of KV cache slots, with a central router dispatching requests to optimize for hit-rate and load. This tight coupling induces:
  • Load imbalance (popular prefixes concentrated on few nodes)
  • Redundancy (prefixes replicated across multiple instances)
  • Memory fragmentation (underutilized slots cannot be utilized elsewhere)
• Phase-Disaggregation (PD): Prefill and decoding phases are segregated across pools, requiring transfer of full prefix states, which increases redundancy, fragmentation, and incurs high transfer overhead.

No existing approach enables true memory pooling since their imperative APIs force explicit, monolithic cache management, and always transfer increasingly long prefixes—precluding low-latency migration and fine-grained balancing (Wu et al., 24 Aug 2025).

2. Declarative Prefix Cache Interface

TokenLake introduces a declarative interface categorizing API calls into control-plane (stateless) and data-plane (declarative) operations.

• Control Plane:
  • get_prefix_tree(): returns the global prefix tree $T$
  • get_cache_load(reqs): estimates GPU resource fraction $L$ used by a request set $\mathrm{reqs}$
  • gen_plans(batches, \mathrm{DoP}): outputs per-batch segment access (query_plans) and storage (transfer_plans) directives for instance groups

Formally, for $\mathrm{reqs}$ with prefix lengths $p_r$, new input lengths $i_r$, and hidden dimension $d$:

$$M_r = \sum_{r \in \mathrm{reqs}} 4d(p_r + i_r) \ F_r = \sum_{r \in \mathrm{reqs}} 2d p_r i_r \ L = \max \left( \frac{M_r}{M_r + N B_\mathrm{mem} T_r},\ \frac{F_r}{F_r + N F T_r} \right)$$

with $N$ instances, $B_\mathrm{mem}$ GPU DRAM bandwidth, $L$0 GPU FLOPS, and $L$1 estimated per-request minimum execution time (from pre-profiled regression).

• Data Plane:
  • init_query(query_plan): initializes per-instance shared buffers for queries
  • init_transfer(transfer_plan): allocates buffers for new cache segments
  • query(q_\mathrm{buf}) and put(kv_\mathrm{buf}): coordinate scatter/gather and push operations across the cluster
  • All communication overlaps compute using a peer-to-peer, asynchronous design with CUDA MPS and NCCL-registered buffers.

This abstraction decouples the scheduler from explicit cache placement or movement, supporting stateless orchestration and enabling true GPU memory pooling (Wu et al., 24 Aug 2025).

3. Segment-Level Model and Execution

TokenLake fragments each prefix into contiguous segments of size $L$2 tokens, yielding $L$3 for a prefix of length $L$4:

$L$5

where $L$6, $L$7 is network latency, and $L$8 is network bandwidth. Empirically, for A100/InfiniBand ($L$9, $\mathrm{reqs}$0 TFLOPS, $\mathrm{reqs}$1 TB/s, $\mathrm{reqs}$2 GB/s, $\mathrm{reqs}$3 μs), $\mathrm{reqs}$4 tokens suffices.

With segment granularity, each segment is independently placed and managed, enabling:

• Fine-grained migration and partial cache sharing
• Defragmentation and deduplication
• Overlap of segment fetch with self-attention computation Remote fetching of a segment is never slower than local re-computation if $\mathrm{reqs}$5 satisfies the criterion above (Wu et al., 24 Aug 2025).

4. Heavy-Hitter-Aware Load Balancing and Eviction

Segment access skew is addressed via the Heavy-Hitter-Aware (HHA) algorithm:

• Segments are classified as:
  • Heavy hitters ($\mathrm{reqs}$6): $\mathrm{reqs}$7 most-accessed segments, tracked by breadth-first search over the prefix tree with access counts $\mathrm{reqs}$8
  • Normal segments ($\mathrm{reqs}$9): the remainder

Assignment proceeds as follows:

• Normal segments: Placed by uniform hashing ($\mathrm{reqs}$0) for intrinsic load balance
• Heavy hitters:
  • Monitor per-instance segment load $\mathrm{reqs}$1 by sliding window
  • Overloaded instances replicate hot segments to least-loaded peers until $\mathrm{reqs}$2
  • Queries use power-of-two-choices among replicas for minimal response time

Eviction is global least recently used (LRU): all replicas log last-access timestamps, and the least recently used is reclaimed on pool exhaustion. The procedure achieves a segment-access load coefficient of variation $\mathrm{reqs}$3 (compared to up to 122% in PD-disaggregation baselines) (Wu et al., 24 Aug 2025).

5. Communication Volume Minimization

Batch dispatch to instances affects the required communication for queries (scatter/gather) and new KV segment placement. TokenLake employs minimum-weight perfect matching via the Hungarian algorithm ($\mathrm{reqs}$4 time):

For $\mathrm{reqs}$5 batches $\mathrm{reqs}$6 and $\mathrm{reqs}$7 instances $\mathrm{reqs}$8:

• For each $\mathrm{reqs}$9, define $p_r$0 as the set of instances owning needed segments and $p_r$1 as those designated for new segments
• Edge weights $p_r$2 are negative of the sum of bytes that would be communicated if $p_r$3 is scheduled on $p_r$4
• The matching assigns each batch to an instance to minimize total inter-node communication

This step is entirely transparent to the LLM scheduler (Wu et al., 24 Aug 2025).

6. Integration with Stateless, Elastic Scheduling

Because TokenLake exposes cache management declaratively, conventional and modern schedulers (PD-disaggregation, chunked prefill, elastic sequence parallelism) require no modifications:

1. Batches are formed and degree-of-parallelism chosen as usual by the scheduler
2. TokenLake APIs provide cache layout and load estimation
3. Query and transfer plans are generated and initialized
4. Compute engines handle actual model execution, with all cache operations asynchronous and overlapped

Each GPU runs three concurrent processes: compute engine (core LLM computation), query engine (peer-to-peer query coordination), and transfer engine (segment migration), all using peer-to-peer zero-copy buffers. This architecture achieves isolation of TokenLake’s cache traffic from main compute, improving utilization (Wu et al., 24 Aug 2025).

7. Performance Evaluation and Contributions

Evaluation on Llama-2-7B (16K context), A100 GPUs, NVLink + 200G InfiniBand, and real-world workloads (LooGLE, SCBench, ShareGPT) demonstrates:

System	P90 Goodput (req/s)	Throughput (relative)	Hit Rate @512K tokens
SGLang-Router	9	×1	37%
SGLang-MoonCake-1P3D	12	×1.1	25%
SGLang-MoonCake-3P1D	11	×0.9	22%
TokenLake	24	×2.6	75% (+2.0×)

• Load balance CV: SGLang-Router 99%, MoonCake-1P3D 62%, MoonCake-3P1D 122%, TokenLake 11%
• Multi-node (16 × A800 GPUs): TokenLake achieves up to 28× higher throughput under comparable latency SLO
• Key contributions:
  • Declarative cache interface for decoupling and pooling
  • Segment-level cache management
  • HHA load balancing with power-of-two-choice routing
  • Communication-efficient batch dispatch
  • Demonstrated 2.6× throughput and 2.1× hit-rate improvement (Wu et al., 24 Aug 2025)

8. Limitations and Prospective Developments

• The segment size $p_r$5 is statically chosen based on hardware; further reductions in overhead for extreme prefix lengths may be possible with adaptive or hierarchical schemes
• Current pooling is GPU-local; extension to CPU/GPU hybrid memory is a viable direction
• The Hungarian algorithm is efficient for $p_r$6dozens of GPUs but may become intractable at cluster scale; approximate matchings or streaming approaches may be required
• Integration of advanced capabilities such as cross-request prefetching, cache compression, or retrieval-augmented methods into TokenLake’s declarative infrastructure is an open area

TokenLake demonstrates that declarative, segment-level, and globally managed prefix caches enable highly efficient, elastic, and low-latency LLM inference serving, with comprehensive improvements in throughput, hit-rate, and memory utilization over prior arts (Wu et al., 24 Aug 2025).