ReasonAlloc: Hierarchical Decoding-Time KV Cache Budget Allocation for Reasoning Models
作者: Wenhao Liu, Hao Shi, Yunhe Li, Weizhi Fei, Xiangyuan Wang, Mengzhe Ruan, Hanxu Hou, Peisong Wang, Linqi Song, Shuang Qiu
分类: cs.AI
发布日期: 2026-06-09
💡 一句话要点
提出ReasonAlloc以解决推理模型中的KV缓存预算分配问题
🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)
关键词: 长链推理 KV缓存 预算分配 推理优化 大型语言模型
📋 核心要点
- 现有的解码时间压缩方法在KV缓存预算分配上存在均匀假设，无法适应推理过程中的动态需求。
- 本文提出ReasonAlloc，通过分层预算分配策略，优化解码时间的KV缓存使用，提升推理效率。
- 实验表明，ReasonAlloc在小预算条件下（128-512个令牌）显著优于均匀预算的基线方法，提升效果明显。
📝 摘要（中文）
长链推理在大型语言模型中导致了严重的推理瓶颈，主要由于快速增长的键值（KV）缓存。现有的解码时间压缩方法通过令牌驱逐来缓解这一问题，但通常假设所有层和头的预算分配是均匀的。与此不同，现有的非均匀预算分配方法主要针对静态提示预填充阶段，未能捕捉自回归推理的逐步上下文需求。为了解决这一问题，本文提出了ReasonAlloc，一个无训练的框架，将解码时间的KV压缩重新定义为一个分层预算分配问题。ReasonAlloc在两个互补层面上运作：离线的层级预分配策略捕捉了架构驱动的需求模式，而在线的头级策略则根据实时效用重新分配资源。实验结果表明，ReasonAlloc在多个数学推理基准上优于现有的均匀预算方法。
🔬 方法详解
问题定义：本文旨在解决大型语言模型推理过程中KV缓存的快速增长问题，现有方法在动态上下文需求下的预算分配不够灵活，导致推理效率低下。
核心思路：ReasonAlloc通过将解码时间的KV压缩视为分层预算分配问题，提出了离线层级预分配和在线头级资源重新分配的双重策略，以适应推理过程中的动态需求。
技术框架：ReasonAlloc的整体架构包括两个主要模块：离线层级预分配策略和在线头级资源分配策略。前者基于架构驱动的需求模式进行预分配，后者在解码过程中根据实时效用动态调整资源分配。
关键创新：ReasonAlloc的核心创新在于其分层预算分配方法，能够有效捕捉推理过程中的上下文需求变化，与现有的静态预算分配方法形成鲜明对比。
关键设计：在设计中，ReasonAlloc采用了基于推理波动的需求模式，结合了离线和在线策略，确保了在不同推理阶段的资源分配灵活性，且引入的推理时间开销极小。
🖼️ 关键图片
📊 实验亮点
在多个数学推理基准（如MATH-500和AIME 2024）上的实验结果显示，ReasonAlloc在小预算条件下（128-512个令牌）显著优于均匀预算的R-KV、SnapKV和Pyramid-RKV等基线方法，提升幅度明显，验证了其有效性。
🎯 应用场景
ReasonAlloc的研究成果可广泛应用于大型语言模型的推理优化，尤其是在需要高效处理长链推理任务的场景中，如数学推理、自然语言理解等领域。其灵活的预算分配策略将为未来的推理模型设计提供重要参考，提升模型的实际应用价值。
📄 摘要（原文）
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.