通过复用相似查询的注意力计算,实现长上下文生成的高速高保真推理。
MAC-Attention: a Match-Amend-Complete Scheme for Fast and Accurate Attention Computation

- 基于匹配-修正-补全三阶段,复用近期相似查询的注意力结果
- 在128K上下文下降低99% KV访问,生成延迟减少60%以上
- 适用于多种模型架构,适合追求长文本生成性能的开发者
大语言模型在长上下文解码时受输入输出瓶颈制约:每个新标记需重读不断增长的键值缓存。现有加速方法通过压缩降低字节数,但损害精度;或通过选择/淘汰机制限制可访问内容,影响延迟召回与长文本生成质量。本文提出MAC-Attention,一种保持精度与访问完整性的替代方案,通过复用语义相近的近期查询所生成的注意力计算来加速解码。该方法包含三个阶段:首先在短局部窗口内进行预RoPE L2匹配;其次在匹配边界附近重新计算小范围内容以修正误差;最后通过数值稳定融合,将修正结果与新计算的KV尾部注意力合并。在匹配成功情况下,计算与带宽复杂度恒定,不随上下文长度增长。该方法与模型无关,可与IO感知内核、分页KV管理器及MQA/GQA兼容。在LongBench v2(120K)、RULER(120K)和LongGenBench(16K连续生成)上,相比最新FlashInfer库,MAC-Attention最多减少99%的KV访问,128K上下文下令牌生成延迟降低超过60%,注意力阶段速度提升超14.3倍,端到端提速达2.6倍,同时保持全注意力质量。代码已开源:https://github.com/YJHMITWEB/MAC-Attention.git
原文摘要 · Abstract (English)
Long-context decoding in LLMs is IO-bound: each token re-reads an ever-growing KV cache. Prior accelerations cut bytes via compression, which lowers fidelity, or selection/eviction, which restricts what remains accessible, and both can degrade delayed recall and long-form generation. We introduce MAC-Attention, a fidelity- and access-preserving alternative that accelerates decoding by reusing prior attention computations for semantically similar recent queries. It starts with a match stage that performs pre-RoPE L2 matching over a short local window; an amend stage rectifies the reused attention by recomputing a small band near the match boundary; and a complete stage fuses the rectified results with fresh attention computed on the KV tail through a numerically stable merge. On a match hit, the compute and bandwidth complexity is constant regardless of context length. The method is model-agnostic and composes with IO-aware kernels, paged-KV managers, and MQA/GQA. Across LongBench v2 (120K), RULER (120K), and LongGenBench (16K continuous generation), compared to the latest FlashInfer library, MAC-Attention reduces KV accesses by up to 99%, cuts token generation latency by over 60% at 128K, and achieves over 14.3x attention-phase speedups, up to 2.6x end-to-end, while maintaining full-attention quality. By reusing computation, MAC-Attention delivers long-context inference that is both fast and faithful. Code is available here: https://github.com/YJHMITWEB/MAC-Attention.git
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。