IceCache: A New Efficient KV Cache Management Scheme for Long-Sequence Large Language Models

IceCache is a new scheme proposed to address the KV cache memory bottleneck in long-sequence Large Language Model (LLM) inference. By integrating semantic clustering and paged attention mechanisms, it achieves near-original model inference accuracy with only 25% of the cache budget, providing a practical memory optimization solution for long-sequence LLM inference.

In LLM inference, KV cache stores intermediate attention states to accelerate inference, but memory usage grows linearly with sequence length, easily leading to memory bottlenecks when processing long texts. As LLM applications expand, the demand for long sequences (such as long documents, multi-turn dialogues, and chain-of-thought reasoning) increases. Traditional KV cache strategies face issues like high hardware upgrade costs or trade-offs between performance and memory efficiency.

Existing KV cache optimization schemes (e.g., partial offloading to CPU) have shortcomings: 1. Token selection is based on heuristics/simple statistics, lacking semantic understanding, which tends to retain non-critical information; 2. Performance degradation is obvious in long-sequence chain-of-thought scenarios; 3. CPU-GPU data transfer bandwidth easily becomes a bottleneck, and frequent transfers slow down inference speed.

The core innovations of IceCache include: 1. Semantic-aware token clustering: Organize tokens based on semantic similarity and select cache based on semantic importance; 2. Hierarchical dynamic data structure: Dynamically adjust cache content to ensure relevant semantic information stays in GPU memory; 3. Deep integration with paged attention: Optimize memory page allocation and CPU-GPU transfer modes.

In LongBench benchmark tests: 1. Maintains 99% of the original model's accuracy with a 256-token cache budget; 2. With only 25% of the KV cache budget, latency and accuracy are better than other offloading methods; 3. Strong adaptability to long-sequence scenarios, stably retaining key information of long-distance dependencies.

IceCache implementation details: 1. Semantic encoding and similarity calculation: Encode tokens through embedding models or attention weights, then calculate semantic similarity; 2. Dynamic clustering reorganization: Dynamically adjust clusters during inference, adding new tokens to existing groups or forming new groups; 3. Intelligent prefetching and eviction: Predict semantic clusters that need to be loaded, and prioritize evicting low-correlation clusters.

Application value of IceCache: 1. Edge device deployment: Enables consumer GPUs/edge devices to run large models; 2. Long document processing: Expands the ability to process long documents in fields like law and medicine; 3. Multi-turn dialogue and reasoning: Better retains key context, improving interaction quality and reasoning accuracy.

The IceCache code has been open-sourced (Project website: https://yuzhenmao.github.io/IceCache/). Future directions: Explore more fine-grained semantic representations, expand to multi-modal scenarios, combine model quantization to optimize memory, and develop adaptive budget allocation mechanisms.