Crystal-KV: Efficient KV Cache Management for Chain-of-Thought LLMs via Answer-First Principle

TL;DR

The paper tackles the memory burden of think-stage KV caches in Chain-of-Thought reasoning by introducing Crystal-KV and the answer-first principle, which separates think-stage KV entries into CrystalKV (critical for the final answer) and SlipKV (aiding reasoning). It presents an attention-based LRFU eviction policy that identifies and preserves CrystalKV while discarding SlipKV, and an adaptive budget allocation mechanism that reallocates cache space to the most impactful layers and heads. Empirically, Crystal-KV delivers substantial memory savings (around 90.9%), significant throughput increases (about 7.6x), and latency reductions (up to 1.24x), while maintaining or improving final-answer accuracy on complex CodeForces and math tasks; it even achieves up to 105% of FullKV accuracy with only 10% of the KV budget on long sequences. The work demonstrates a practical, scalable approach to efficient CoT reasoning, with strong implications for deploying large-scale LLMs in latency- and memory-constrained environments.

Abstract

Chain-of-Thought (CoT) reasoning in large language models (LLMs) significantly improves accuracy on complex tasks, yet incurs excessive memory overhead due to the long think-stage sequences stored in the Key-Value (KV) cache. Unlike traditional generation tasks where all tokens are uniformly important, CoT emphasizes the final answer, rendering conventional KV compression strategies ineffective. In this paper, we present Crystal-KV, an efficient KV cache management framework tailored for CoT reasoning. Our key insight is the answer-first principle. By mapping answer preferences into think-stage attention map, we distinguish between SlipKV, which mainly maintains the reasoning flow but may occasionally introduce misleading context, and CrystalKV, which truly contributes to the correctness of the final answer. Next, we propose an attention-based Least Recently Frequently Used algorithm. It precisely identifies when a SlipKV entry's utility expires and evicts it, retaining CrystalKV without disrupting reasoning flow. Finally, we introduce an adaptive cache budget allocation algorithm. Based on the dynamic proportion of CrystalKV, it estimates the importance of each layer/head and adjusts the KV cache budget during inference, amplifying critical components to improve budget utilization. Results show that Crystal-KV achieves state-of-the-art KV cache compression, significantly improves throughput, and enables faster response time, while maintaining, or even improving, answer accuracy for CoT reasoning.

Figures (5)

• Figure 1: The Workflow of Chain-of-Thought Reasoning
• Figure 2: Distinguishing CrystalKV and SlipKV by Projecting Answer Preferences onto Think-Stage Attention Map
• Figure 3: Overview of Crystal-KV
• Figure 4: Accuracy Comparison on Math and Code Tasks across Three Reasoning LLMs
• Figure 5: Effect of Decay Rate and Sampling Threshold