通过低秩分解减少大模型长文本推理的显存占用。
Efficient Low Rank Attention for Long-Context Inference in Large Language Models
- 预填充阶段将查询和键矩阵分解为低秩因子,解码时用低维投影计算注意力。
- 在RULER和LongBench上,显存消耗降低超60%,精度损失小于1%。
- 适合资源受限设备部署,尤其适用于长文本生成任务。
随着输入文本长度增加,大型语言模型中的键值(KV)缓存带来巨大的显存开销,限制了资源受限设备上的长上下文推理。现有方法如KV量化和剪枝虽可降低内存使用,但存在数值精度损失或关键键值对保留不足的问题。本文提出低秩查询与键注意力(LRQK),一种两阶段框架:在预填充阶段将全精度查询和键矩阵联合分解为紧凑的秩-")r" 因子,并在每个解码步骤中以 \\(\mathcal{O}(lr)\\) 时间利用这些低维投影计算代理注意力分数。通过仅保留前-")k" 个最相关标记和少量近期标记,采用混合GPU-CPU缓存与命中-未命中机制,仅传输缺失的全精度键值对,从而在保持精确注意力输出的同时减少跨设备数据传输。在LLaMA-3-8B和Qwen2.5-7B模型上,基于RULER与LongBench的大量实验表明,LRQK在长上下文设置下表现优于或媲美主流稀疏注意力方法,同时实现显著的内存节省且精度损失极小。代码已开源:https://github.com/tenghuilee/LRQK。
原文摘要 · Abstract (English)
As the length of input text increases, the key-value (KV) cache in LLMs imposes prohibitive GPU memory costs and limits long-context inference on resource constrained devices. Existing approaches, such as KV quantization and pruning, reduce memory usage but suffer from numerical precision loss or suboptimal retention of key-value pairs. In this work, Low Rank Query and Key attention (LRQK) is introduced, a two-stage framework that jointly decomposes full-precision query and key matrices into compact rank-\(r\) factors during the prefill stage, and then employs these low-dimensional projections to compute proxy attention scores in \(\mathcal{O}(lr)\) time at each decode step. By selecting only the top-\(k\) tokens and a small fixed set of recent tokens, LRQK employs a mixed GPU-CPU cache with a hit-and-miss mechanism where only missing full-precision KV pairs are transferred, thereby preserving exact attention outputs while reducing CPU-GPU data movement. Extensive experiments on the RULER and LongBench benchmarks with LLaMA-3-8B and Qwen2.5-7B demonstrate that LRQK matches or surpasses leading sparse-attention methods in long context settings, while delivering significant memory savings with minimal accuracy loss. Our code is available at https://github.com/tenghuilee/LRQK.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。