ReasonAlloc: 推理模型的分层解码时KV缓存预算分配
ReasonAlloc: Hierarchical Decoding-Time KV Cache Budget Allocation for Reasoning Models
AI总结：
针对长链式推理中KV缓存快速增长导致的推理瓶颈，提出ReasonAlloc框架，通过离线层预分配和在线头重分配的分层预算分配策略，在不增加训练开销下显著提升小预算下的推理性能。
摘要：
大语言模型推理中的长链式思维轨迹由于键值缓存的快速增长导致严重的推理瓶颈。当前的解码时压缩方法通过令牌驱逐缓解此问题，但通常假设所有层和头之间均匀分配预算。相比之下，现有的非均匀预算分配方法主要针对静态提示预填充阶段设计，未能捕捉自回归推理的逐步上下文需求。为弥补这一差距，我们提出ReasonAlloc，一个无需训练的框架，将解码时KV压缩重新表述为分层预算分配问题。ReasonAlloc在两个互补层面运作：离线层预分配策略捕捉一种架构驱动的需求模式，我们称之为“推理波”；在线头策略根据实时效用将资源重新分配给信息丰富的头。在数学推理基准（MATH-500、AIME 2024）上使用DeepSeek-R1-Distill-Llama-8B、DeepSeek-R1-Distill-Qwen-14B和AceReason-14B的评估表明，ReasonAlloc优于均匀预算的R-KV、SnapKV和Pyramid-RKV（一种强制静态单调递减层预算的基线），在小预算（128-512令牌）下增益最大。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.
机构、总结、中文摘要均由AI分析而来，可能存在错误，请以论文正文为准。
