arXiv:2606.11164cs.AI2026-06被引 1

针对推理模型的缓存瓶颈,提出分层动态分配缓存预算的新方法。

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

论文配图:ReasonAlloc: Hierarchical Decoding-Time KV Cache Budget Allocation for Reasoning Models
图 1 · 摘自论文原文
  • 分两层分配缓存:先按架构预设各层需求,再实时调整重要注意力头的资源
  • 在128-512缓存令牌下性能提升显著,数学推理任务准确率更高
  • 无需训练、可直接接入现有压缩策略,适合部署资源受限的推理系统

大语言模型在长链式思维(CoT)推理中因键值缓存快速膨胀导致严重推理瓶颈。现有解码时压缩方法通常假设各层和注意力头间均匀分配缓存预算,而现有非均匀分配方法多用于静态提示预填充阶段,无法捕捉自回归推理中逐步变化的上下文需求。为此,我们提出 ReasonAlloc,一种无需训练的框架,将解码时缓存压缩重构为分层预算分配问题。该方法包含两个互补层级:离线层级预分配策略捕捉由模型架构决定的“推理波”模式;在线头级策略根据实时信息价值重新分配资源至高价值注意力头。在 MATH-500、AIME~2024 数学推理基准上使用 DeepSeek-R1-Distill-Llama-8B、DeepSeek-R1-Distill-Qwen-14B、AceReason-14B 模型的评估表明,ReasonAlloc 在小缓存预算(128-512 令牌)下优于均匀预算的 R-KV、SnapKV 和 Pyramid-RKV(静态单调递减层预算基线),且与现有令牌淘汰策略无缝集成,推理开销可忽略不计。

原文摘要 · Abstract (English)

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.

缓存优化推理加速注意力机制分层分配

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。