通过聚类压缩固定上下文,加速长文本推理。
Squeezed Attention: Accelerating Long Context Length LLM Inference
- 用聚类将固定上下文键值分组,仅计算相关部分注意力。
- 在LongBench上减少3.1倍KV缓存占用,速度提升4倍以上。
- 适合文档分析、代码生成等长上下文任务的部署场景。
新兴的大语言模型应用需要长输入上下文以完成文档分析、代码生成等复杂任务。然而,输入序列越长,推理开销呈线性增长。针对许多场景中大部分上下文内容固定的特点,我们提出Squeezed Attention,通过离线K-means聚类将固定上下文的键按语义相似性分组,并用中心点代表每类。推理时,用户输入查询与中心点比对,预测相关键,仅对关键键计算精确注意力,显著降低带宽和计算成本。还提出层级版本,使注意力复杂度从线性降至对数级。在LongBench等长上下文基准上,对LLaMA-2-7B-32K、LWM-Text-Chat-1M、Longchat-7B-v1.5-32K模型实现3.1×的KV预算减少,无明显精度损失;最多8×减少下仅0.5点精度下降。我们实现了中心点比对与重要键稀疏FlashAttention内核,在预填充和生成阶段均实现超4×加速。代码已开源。
原文摘要 · Abstract (English)
Emerging Large Language Model (LLM) applications require long input context in order to perform complex tasks like document analysis and code generation. For these long context length applications, the length of the input prompt poses a significant challenge in terms of inference efficiency since the inference costs increase linearly with sequence length. However, for many of these applications, much of the context in the prompt is fixed across different user inputs, thereby providing the opportunity to perform offline optimizations in order to process user inputs quickly, as they are received. We propose Squeezed Attention to accelerate LLM applications where a large portion of the input context is fixed. We first leverage K-means clustering offline to group the keys for the fixed context based on semantic similarity and represent each cluster with a single centroid value. During inference, we compare query tokens from the user input with the centroids to predict which keys from the fixed context are semantically relevant, and then compute exact attention using only the important keys, thereby reducing bandwidth and computational costs. We also present a hierarchical version of our algorithm which can reduce the complexity of attention from linear to logarithmic with respect to the fixed context length. We evaluate our method on long-context benchmarks including LongBench, where it achieves a 3.1$\times$ reduction in KV budget with no noticeable accuracy loss and up to an 8$\times$ reduction with only a 0.5 point accuracy gap for the LLaMA-2-7B-32K, LWM-Text-Chat-1M, and Longchat-7B-v1.5-32K models. Futhermore, we implement kernels for centroid comparison and sparse FlashAttention with important keys, achieving more than 4$\times$ speedups during both the prefill and generation phases for long-context inference. Our code is available at https://github.com/SqueezeAILab/SqueezedAttention.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。