LongCache Mechanism in TokenLake LLMs

• LongCache Mechanism is a segment-level prefix cache pooling system that enhances load balancing and deduplication in TokenLake.
• It divides long-context inputs into fixed-size segments (~500–600 tokens) to enable asynchronous, efficient cache reorganization and resource management.
• The design minimizes communication overhead and load imbalance, achieving throughput improvements up to 2.6× compared to traditional methods.

The LongCache Mechanism refers to the unified segment-level prefix cache pool introduced as part of TokenLake, a system designed for fine-grained, elastic serving of long-context LLMs on GPU clusters. TokenLake's architecture and mechanisms address pathological inefficiencies of traditional prefix caching schemes at cluster scale when serving multi-turn, long-context interactions. By segmenting prefix caches and pooling them asynchronously across peer GPUs, TokenLake’s LongCache Mechanism achieves improved load balance, deduplication, defragmentation, and communication minimization, while fully decoupling cache management from the request scheduler (Wu et al., 24 Aug 2025).

1. Motivation and System Architecture

Conventional prefix caching in LLM serving systems employs imperative APIs (put/get/transfer) that couple cache management tightly to request scheduling. This entanglement induces three principal pathologies as context and KV-cache sizes grow:

• Load Imbalance: Hot prefixes concentrate computation and KV-bandwidth on specific GPUs, overloading them.
• Redundancy: Shared prefixes (e.g., common system prompts) must be replicated across many GPUs to maintain compute balance, consuming excess memory.
• Fragmentation: With "strict locality" (one prefix per GPU), memory is divided into small unusable slots, which cannot be coalesced for large prefixes.

TokenLake restructures this model through a single, peer-to-peer, segment-level prefix cache pool. Prefixes are divided into fixed-size segments (typically 500–600 tokens) that may reside on any GPU. A declarative cache interface enables TokenLake to execute reorganization, balancing, deduplication, and defragmentation of cache segments without exposing these operations to the computation scheduler.

2. Declarative Cache Interface and Key Properties

TokenLake introduces two primary abstractions:

• Query tensors ($q$): Per-layer $Q$ projections driving prefix attention.
• Fixed-Size Prefix Cache Segments: Segments of $C$ tokens per segment, with $C \approx 500$–$600$.

The system presents distinct APIs for control and data planes:

• Control-plane: get_prefix_tree, get_cache_load, gen_plans.
• Data-plane: init_query, init_transfer, query, put.

All segments are sized $C$ and address-aligned. Theoretical constraints for segment size ensure efficiency, where

$$k_{\mathrm{comp}} = \max(4d/F, 4d/B_{\mathrm{mem}})$$

$T_{\mathrm{comp}}(C) = k_{\mathrm{comp}} \cdot C$

$$T_{\mathrm{comm}}(C) = 2\alpha_{\mathrm{net}} + \frac{4d}{B_{\mathrm{net}}}$$

with $d$ as hidden size; $Q$0 is chosen such that

$Q$1

3. Segment-Level Cache Operations

Each prefix in a multi-turn interaction is partitioned into $Q$2 segments, tracked globally via a directory mapping segment-IDs to the set of hosting GPU instances and last access times.

Key operations (expressed in pseudo-code in the canonical description (Wu et al., 24 Aug 2025)):

• Insertion: Upon generation, segments are registered in the global directory with their hosting instance(s) and access time.
• Lookup: For a given batch, required segments' owners are determined, and an optimal replica is selected for each segment (see Section 4).
• Eviction: When capacity is exceeded, global LRU is invoked—prioritizing eviction of redundant replicas before unique copies.

4. Heavy-Hitter-Aware Load Balancing

A breadth-first search on the global prefix trie identifies "heavy hitter" segments—the top $Q$3 segments by query access—but only these frequently accessed segments are selectively replicated. Normal segments are mapped via hash-based placement ($Q$4).

When selecting from available segment replicas during batch servicing, TokenLake employs the "power-of-two choices" strategy, choosing the less-loaded GPU from two random replicas. The objective is to minimize the variance $Q$5 of per-GPU loads, where $Q$6 is the current load on GPU $Q$7. Theoretical guarantees (Fan-Lim-Andersen) indicate that absorbing the top $Q$8 heavy hitters suffices for low load variance with high probability.

5. Deduplication and Defragmentation

The global directory enforces deduplication by identifying and coalescing identical segments via unique IDs. On eviction, TokenLake prefers discarding surplus replicas before removing the last copy. Segment sizing inherently caps internal fragmentation at $Q$9 tokens per segment.

Pooling segments cluster-wide means that any available memory slot, regardless of its GPU, can satisfy a segment allocation, eliminating fragmented "islands" of memory. Quantitatively:

• Before: $C$0
• After: $C$1

6. Communication Minimization Framework

All communication of query tensors and KV-writes occurs at the segment granularity ($C$2 per segment). Communication and local self-attention are overlapped, minimizing end-to-end data transfer.

Assignment of batch processing is formulated as a weighted bipartite matching problem: for each batch node $C$3 and GPU $C$4,

$C$5

where $C$6 is the set of GPUs holding required segments, $C$7 the destination GPUs for KV-writes, and $C$8 the number of KV-write operations. An $C$9 Hungarian algorithm computes the assignment that minimizes total communication.

7. Scheduler Decoupling and API Integration

The scheduler within TokenLake interacts only with stateless, high-level APIs (get_prefix_tree, get_cache_load, gen_plans), with all prefix cache placement, replica management, and memory operations abstracted away. For each scheduling iteration, the scheduler:

1. Queries the projected cache load ($C \approx 500$0) via get_cache_load.
2. Adjusts batch sizes and/or parallelism such that $C \approx 500$1.
3. Generates batch-specific query and transfer plans.
4. Invokes fully asynchronous data-plane calls: init_query, init_transfer, query, put.

This design enables elastic request scheduling, unburdened by cache coupling.

8. Performance Outcomes and Outlook

Empirical evaluation on real-world datasets (LooGLE, SCBench, ShareGPT) yields:

• Throughput Speedup: Up to $C \approx 500$2 (versus SGLang-Router) and $C \approx 500$3 (versus MoonCake PD disaggregation).
• Hit Rate: $C \approx 500$4–$C \approx 500$5 higher relative to both prior approaches.
• Load Imbalance: TokenLake reduces per-GPU coefficient of variation to approximately $C \approx 500$6, versus $C \approx 500$7 (SGLang-Router) and $C \approx 500$8–$C \approx 500$9 (MoonCake PD, depending on configuration).

The architectural innovations underpinning the LongCache Mechanism—segment-based global pooling, heavy-hitter-aware load balancing, deduplication, and communication minimization—support near-perfect resource utilization and operational elasticity. Potential extensions include online adaptive segment sizing, hierarchical (multi-tier) pooling (GPU:host), and integration with on-GPU KV-compression or selective overwrite to further accommodate extreme context lengths (Wu et al., 24 Aug 2025).