动态分配注意力缓存精度,4倍减少长文本推理内存占用。
ARKV: Adaptive and Resource-Efficient KV Cache Management under Limited Memory Budget for Long-Context Inference in LLMs
- 根据每层注意力特性与词元重要性,自适应分配缓存精度。
- 长文本任务保持97%基线准确率,内存降低4倍,吞吐损失极小。
- 无需重训练或改架构,适合实际部署的长序列推理场景。
大语言模型在需要超长上下文推理的场景(如智能体工作流、深度研究理解)中日益普及。然而,长上下文推理受限于KV缓存——这一随序列长度和批量大小线性增长的临时内存结构,迅速占据大量GPU内存。现有内存压缩技术(如淘汰与量化)多依赖静态启发式策略,在严苛内存预算下常导致性能下降。本文提出ARKV,一种轻量级自适应框架,基于每层注意力动态与词元级别重要性,动态分配缓存精度。在短预填充阶段,通过计算注意力熵、方差与峰度等统计得分,估计各层原始量化比(OQ)。解码阶段,词元被快速分配至原精度(Original)、低精度(Quantization)或淘汰(Eviction)三状态之一。在LLaMA3与Qwen3模型上,跨多样长/短上下文任务的实验表明,ARKV在长文本基准上保留约97%基线准确率,同时将KV内存使用降低4倍,吞吐损失极小;短文本任务表现媲美全精度基线;在GSM8K数学推理任务中显著优于均匀量化。结果证明ARKV在可扩展大模型部署中的实用性,实现无需重训练或架构修改的细粒度、数据驱动内存控制。源代码与资源见:https://github.com/Large-scale-Sustainable-Computing-LSC/ARKV
原文摘要 · Abstract (English)
Large Language Models (LLMs) are increasingly deployed in scenarios demanding ultra-long context reasoning, such as agentic workflows and deep research understanding. However, long-context inference is constrained by the KV cache, a transient memory structure that grows linearly with sequence length and batch size, quickly dominating GPU memory usage. Existing memory reduction techniques, including eviction and quantization, often rely on static heuristics and suffer from degraded quality under tight budgets. In this paper, we propose ARKV, a lightweight and adaptive framework that dynamically allocates precision levels to cached tokens based on per-layer attention dynamics and token-level importance. During a short prefill phase, ARKV estimates the original quantization (OQ) ratio of each layer by computing statistical scores such as attention entropy, variance and kurtosis. During decoding, tokens are assigned to one of three states, Original (full precision), Quantization (low precision), or Eviction, according to a fast heavy-hitter scoring strategy. Our experiments on LLaMA3 and Qwen3 models across diverse long- and short-context tasks demonstrate that ARKV preserves ~97% of baseline accuracy on long-context benchmarks while reducing KV memory usage by 4x, with minimal throughput loss. On short-context tasks, ARKV matches full-precision baselines; on GSM8K math reasoning, it significantly outperforms uniform quantization. These results highlight the practical viability of ARKV for scalable LLM deployment, offering fine-grained, data-driven memory control without retraining or architectural modifications. The source code and artifacts can be found in: https://github.com/Large-scale-Sustainable-Computing-LSC/ARKV
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。