通过复用相似注意力图，加速大模型预填充阶段推理。
AttnCache: Accelerating Self-Attention Inference for LLM Prefill via Attention Cache
- 构建注意力图缓存库，按相似性检索复用
- CPU上平均提速1.2倍，GPU上达3倍注意力计算加速
- 适合仅需预填充的场景，如分类、问答、推荐
大型语言模型广泛应用于对话、代码生成和推理等生成任务。然而，许多实际应用如分类、问答、推荐和文本嵌入仅依赖推理的预填充阶段，即模型编码输入序列而不进行自回归解码。在此类仅预填充场景中，由于自注意力计算随序列长度呈二次复杂度，成为主要性能瓶颈。本文观察到语义不同的句子在多层多头中常产生相似的注意力图。基于此，我们提出AttnCache框架，通过检索并重用相似的注意力图来加速大模型预填充阶段的推理。该框架基于注意力图记忆数据库，采用高效的缓存与相似性搜索技术，在推理时识别并复用预缓存的注意力图，从而降低自注意力的计算开销。实验表明，AttnCache在CPU上实现平均1.2倍端到端加速和2倍注意力加速，在GPU上实现1.6倍端到端加速和3倍注意力加速，且精度损失可忽略。
原文摘要 · Abstract (English)
Large Language Models (LLMs) are widely used in generative applications such as chatting, code generation, and reasoning. However, many realworld workloads such as classification, question answering, recommendation, and text embedding rely solely on the prefill stage of inference, where the model encodes input sequences without performing autoregressive decoding. In these prefill only scenarios, the self-attention computation becomes the primary performance bottleneck due to its quadratic complexity with respect to sequence length. In this paper, we observe that semantically different sentences often produce similar attention maps across layers and heads. Building on this insight, we propose AttnCache, a framework that accelerates the prefill stage of LLM inference by retrieving and reusing similar attention maps. Based on an attention map memorization database, AttnCache employs efficient caching and similarity search techniques to identify and reuse pre-cached attention maps during inference, thereby reducing the computational overhead of self-attention. Experimental results show that AttnCache achieves an average of 1.2x end-to-end and 2x attention speedup on CPU, and 1.6x end-to-end and 3x attention speedup on GPU, with negligible accuracy degradation.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。