论文精选

ReasonAlloc:推理模型解码时KV缓存分层预算分配

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

精选理由

推理模型的长CoT导致KV缓存爆炸,做推理优化的开发者可以直接用ReasonAlloc替代均匀预算方案,在低预算下获得显著性能提升。

AI 摘要

ReasonAlloc 是一种无需训练的框架,针对大语言模型推理中长思维链(CoT)导致的KV缓存快速增长问题,提出分层预算分配方案。它通过离线层间预分配捕捉架构驱动的“推理波”模式,并结合在线头间实时重分配,将资源导向信息丰富的注意力头。在数学推理基准(MATH-500、AIME 2024)上,使用DeepSeek-R1-Distill-Llama-8B等模型测试,ReasonAlloc在低预算(128-512 tokens)下显著优于均匀预算方法(如R-KV、SnapKV)。该框架可即插即用于现有token驱逐策略,且推理开销极小。

原文 · arXiv: DeepSeek

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

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.

ReasonAlloc:推理模型解码时KV缓存分层预算分配 · AI 热点