ReasonAlloc: Hierarchical Decoding-Time KV Cache Budget Allocation for Reasoning Models

• URL: http://arxiv.org/abs/2606.11164v1
• Date: Tue, 09 Jun 2026 17:44:23 GMT
• Title: ReasonAlloc: Hierarchical Decoding-Time KV Cache Budget Allocation for Reasoning Models
• Abstract summary: ReasonAlloc recasts decoding-time KV compression as a hierarchical budget allocation problem.<n>An offline layer-wise preallocation strategy captures an architecture-driven demand pattern.<n>An online head-wise strategy reallocates resources during decoding to information-rich heads based on real-time utility.
• Score: 51.947287846779545
• License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/
• Abstract: Long chain-of-thought (CoT) trajectories in large language model (LLM) reasoning cause severe inference bottlenecks due to rapid key-value (KV) cache growth. Current decoding-time compression methods mitigate this issue via token eviction, but typically assume a uniform budget distribution across all layers and heads. In contrast, existing non-uniform budget allocation methods are predominantly designed for the static prompt prefill phase, and they do not capture the stepwise context demands of autoregressive reasoning. To bridge this gap, we propose ReasonAlloc, a training-free framework that recasts decoding-time KV compression as a hierarchical budget allocation problem. ReasonAlloc operates at two complementary levels: an offline layer-wise preallocation strategy captures an architecture-driven demand pattern which we call ``\textit{Reasoning Wave}'', while an online head-wise strategy reallocates resources during decoding to information-rich heads based on real-time utility. Evaluations on mathematical reasoning benchmarks (MATH-500, AIME~2024) using DeepSeek-R1-Distill-Llama-8B, DeepSeek-R1-Distill-Qwen-14B, and AceReason-14B show that ReasonAlloc outperforms uniform-budget R-KV, SnapKV, and Pyramid-RKV (a baseline enforcing a static, monotonically decreasing layer budget), with the largest gains at small budgets (128-512 tokens). ReasonAlloc is plug-and-play with existing token-eviction policies and introduces negligible inference-time overhead.

Related papers

• Thinking Economically: A Hierarchical Framework for Adaptive-Complexity Reasoning in LLMs [36.73612668202756]
  HAB learns step-specific token budgeting signals from PPL-derived step comparisons.<n>Experiments on GSM8K and MATH500 show that HAB not only surpasses standard CoT in accuracy but also reduces token usage.
  arXiv  Detail & Related papers  (2026-05-31T11:20:00Z)
• Make Each Token Count: Towards Improving Long-Context Performance with KV Cache Eviction [65.710271475739]
  We introduce a global retention-based KV eviction method that learns each token's future utility under a unified memory budget.<n>We show that our method substantially reduces KV memory while matching or surpassing full-cache inference.<n>These results suggest that learned, globally calibrated KV eviction is not only a compression technique, but also a mechanism for improving long-context reasoning.
  arXiv  Detail & Related papers  (2026-05-10T16:47:50Z)
• DepthKV: Layer-Dependent KV Cache Pruning for Long-Context LLM Inference [11.462434448123169]
  We propose a layer-dependent pruning framework that allocates a fixed global KV budget across layers based on sensitivity.<n>Across multiple models and tasks, DepthKV consistently outperforms uniform pruning at the same global pruning ratio.
  arXiv  Detail & Related papers  (2026-04-27T16:15:37Z)
• Kwai Summary Attention Technical Report [69.40814939510126]
  Long-context ability, has become one of the most important direction of next-generation Large Language Models.<n>Standard softmax attention exhibits quadratic time complexity with respect to sequence length.<n>We propose Kwai Summary Attention (KSA), a novel attention mechanism that reduces sequence modeling cost by compressing historical contexts.
  arXiv  Detail & Related papers  (2026-04-27T12:59:53Z)
• CASK: Core-Aware Selective KV Compression for Reasoning Traces [0.0]
  CASK partitions the decode-time reasoning trace into a protected core that anchors answer formation and intermediate state, and mergeable scratch with high redundancy.<n>On the H100 reasoning gate, CASK shows higher full- KV continuation fidelity than TriAttention at matched budgets on both AIME24 and AIME25, with recurring cask@384 > triattention@512 crossings.
  arXiv  Detail & Related papers  (2026-04-13T02:03:16Z)
• HybridKV: Hybrid KV Cache Compression for Efficient Multimodal Large Language Model Inference [23.016627145448467]
  Multimodal Large Language Models (MLLMs) have advanced unified reasoning over text, images, and videos.<n>Each visual input expands into thousands of tokens, causing caches to scale linearly with context length.<n>A common solution is to compress caches under a fixed allocated budget at different granularities.<n>We propose HybridKV, a hybrid KV cache compression framework that integrates complementary strategies in three stages.
  arXiv  Detail & Related papers  (2026-04-07T13:51:07Z)
• Towards Efficient Large Language Reasoning Models via Extreme-Ratio Chain-of-Thought Compression [55.63153956934198]
  Chain-of-Thought (CoT) reasoning successfully enhances the reasoning capabilities of Large Language Models (LLMs)<n>Existing CoT compression methods often suffer from a critical loss of logical fidelity at high compression ratios.<n>We propose a novel EXTreme-RAtio Chain-of-Thought Compression framework, termed Extra-CoT, which aggressively reduces the token budget while preserving answer accuracy.
  arXiv  Detail & Related papers  (2026-02-09T06:57:15Z)
• OjaKV: Context-Aware Online Low-Rank KV Cache Compression with Oja's Rule [54.37983890753086]
  We introduce OjaKV, a framework that integrates a strategic hybrid storage policy with online subspace adaptation.<n>OjaKV preserves crucial first and most recent tokens in full-rank, maintaining high-fidelity anchors for attention.<n>It applies low-rank compression by incrementally adapting the projection basis using Oja's algorithm for online principal component analysis.
  arXiv  Detail & Related papers  (2025-09-25T21:42:27Z)
• Lookahead Q-Cache: Achieving More Consistent KV Cache Eviction via Pseudo Query [48.52389201779425]
  KV cache memory usage grows substantially with longer text sequences.<n>Existing KV cache eviction methods prune tokens using prefilling-stage attention scores.<n>Lookahead Q-Cache generates low-cost pseudo lookahead queries to better approximate the true decoding-stage queries.
  arXiv  Detail & Related papers  (2025-05-24T10:34:38Z)
• Ada-KV: Optimizing KV Cache Eviction by Adaptive Budget Allocation for Efficient LLM Inference [37.94892570127548]
  Large Language Models have excelled in various domains but face efficiency challenges due to the growing Key-Value (KV) cache.<n>Recent efforts aim to reduce KV cache size by evicting vast non-critical cache elements during runtime.<n>We propose Ada-KV, the first head-wise adaptive budget allocation strategy.
  arXiv  Detail & Related papers  (2024-07-16T09:53:32Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.