音乐
暂未播放
MAC-Attention:Match-Amend-Complete 注意力复用,长上下文解码 KV 读取最高减少 99%(MLSys 2026)

背景:每生成一个 token,都要把 16 GiB 的 KV 重新读一遍#
长上下文解码到底慢在哪,值得先用一组数字钉死。以 LLaMA-3.1-8B 为例:32 层、8 个 KV 头、每个头 128 维,FP16 存储。每新增一个 token,它的 K/V 要写进缓存,大小是 32×8×128×2×2 B=128 KiB(K、V 各一份)。当上下文长到 128K token 时,每生成一个 token 就要把这 128 KiB×131072≈16 GiB 的 KV 缓存从 HBM 全部读出来做一次注意力。H100 SXM5 的 HBM3 带宽是 3.35 TB/s,光是把这些数据流式读一遍就要约 5 ms;而 8B 稠密模型每 token 的计算量只有约 2×8×109=1.6×1010 FLOP(16 GFLOP),在 H100 约 990 TFLOPS 的 BF16 稠密算力下大约 16 μs 就算完了。两者差出两个数量级——这就是教科书式的 memory-bound:解码的瓶颈不是算力,是内存带宽,而且这个带宽开销随上下文长度线性增长。
FlashAttention 由 Tri Dao 等人于 2022 年提出(NeurIPS 2022),它用分块和在线 softmax 避免了中间注意力矩阵的显式物化,把 IO 模式改造成”每个块只读一次”,大幅降低了常数因子;本博客前文《FlashAttention-4:面向 Blackwell 的算法-流水线协同设计》讲过这条路线能榨出什么。但 FlashAttention 系列不改变一个事实:每个解码步仍然要把整个前缀的 KV 从 HBM 读一遍——它优化的是”怎么读”,而不是”读多少”。PagedAttention(UC Berkeley,SOSP 2023)解决了显存碎片化,让长上下文的 KV 装得下、管得好,但同样没有减少读取总量。
于是过去两年,削减 KV 读取量的工作分成两个家族,各自的代价也很清晰:
- 压缩(compression):用低秩投影(Palu,2024)或量化(RocketKV,佐治亚理工团队 2025 年提出)缩小 KV 本身的体积。代价是保真度——量化误差、低秩截断误差直接进入注意力计算。
- 选择/驱逐(selection/eviction):StreamingLLM(MIT 韩松课题组,2023,ICLR 2024)只保留”注意力汇点”和最近的 token;H2O(2023,NeurIPS 2023)逐 token 保留 heavy hitter;SnapKV(2024,NeurIPS 2024)用观察窗口选出重要 KV;Quest(MIT 韩松课题组,2024)按 query 自适应地选 KV 块。代价是可达性——被驱逐的 token 永远失去了被注意的机会,而”延迟回忆”(delayed recall)、跨文档关联、长文生成恰恰依赖这种可达性。RULER 基准(NVIDIA 团队,2024)就是专门为暴露这类退化设计的。
论文 Figure 1 给出了这两条路线的全景:横轴是 KV 预算(每个解码步实际使用的 KV 比例),纵轴是精度,全注意力是那条灰色虚线。可以看到,多数方法在把 KV 预算压到 10% 以下时,精度曲线明显下坠。
MAC-Attention 由俄亥俄州立大学的 Jinghan Yao、Dhabaleswar K. (DK) Panda 与微软的 Sam Adé Jacobs、Walid Krichene 以及 Anyscale 的 Masahiro Tanaka 合作完成,发表于 MLSys 2026(第九届 MLSys,美国 Bellevue),arXiv 预印本 2604.00235 于 2026 年 3 月 30 日公开,代码开源在 GitHub(YJHMITWEB/MAC-Attention)。它走的是第三条路:既不压缩也不驱逐,而是复用——把”相似的 query 之前算过的注意力”缓存起来,新 query 直接继承,只补算少量新鲜部分。论文报告的成果是:KV 访问量最高减少 99%,128K 上下文下每 token 延迟下降 60% 以上,注意力阶段加速至少 14.3 倍(256K 下最高约 46 倍),端到端生成最高加速 2.6 倍,同时保持全注意力级别的质量。
这条”复用”路线与本博客前几篇讲的方案恰好互补:BLASST 是”跳过低贡献块”(稀疏),QEvict 是”驱逐但可恢复”(量化+晋升),MAC-Attention 是”缓存结果、复用计算”——三个答案针对同一个 IO 瓶颈,机制完全不同。
核心思想:把”算过一遍的注意力”记在便签上#
注意力的可分解表示#
先建立形式化。记 Qt,Kt,Vt 为位置 t 的 query、key、value 编码(省略层与头的下标),Q~t 为施加位置编码之前的 query(pre-RoPE query),当前解码位置为 m。注意力对每个位置 t 先算 logits:
ℓt(m)=d1Qm⊤Kt注意力的输出可以写成两个量的比值:加权和 S(m)=∑t≤meℓt(m)Vt 与归一化常数 Z(m)=∑t≤meℓt(m),即 om=S(m)/Z(m)。这两个标量/向量对就是 softmax 的全部信息量,论文称之为注意力摘要(attention summary):
ASI(m)=(SI(m),ZI(m))=(t∈I∑eℓt(m)Vt,t∈I∑eℓt(m))这里 S 是 d 维向量,Z 是标量。摘要的妙处在于可结合性:两个不相交 token 集合的摘要可以直接相加合并,因此任意位置 p 处,注意力都可以拆成前缀与后缀:
om=Z1:p(m)+Zp+1:m(m)S1:p(m)+Sp+1:m(m)这个分解在数值上是稳定的,p=m−1 的特殊情形就是 FlashAttention 所用的在线归一化恒等式(Milakov 与 Gimelshein 2018 年提出的 online normalizer calculation 的经典结果)。换句话说,只要保存了某个集合的摘要,之后任何时刻想把这个集合的注意力”接入”结果,都是常数时间的加法合并。
复用:相似的 query 会看到相似的注意力#
现在提出 MAC 的核心观察:生成式解码中,一个请求内部短时间窗口内的 query 高度自相似。多轮对话、代码补全、长推理这些真实负载里,模型经常在”反复查”同一段上下文——内部状态不是从均匀分布里随机抽的,而是带有重复的检索模式。两个语义相近的 query Qm≈Qp 在共享前缀上诱导出的注意力分布高度相关:logits 只差一个小的扰动 δt=ℓt(m)−ℓt(p),因此位置 p 当时算好的前缀摘要 S1:p(p),Z1:p(p) 可以当作位置 m 所需前缀摘要 S1:p(m),Z1:p(m) 的近似。
于是解码步 m 的计算变成:
om=Z1:p−r(m)+Zp−r+1:m(m)S1:p−r(m)+Sp−r+1:m(m)≈Z1:p−r(p)+Zp−r+1:m(m)S1:p−r(p)+Sp−r+1:m(m)右侧第一项直接取缓存(位置 p 时算好的、截断到 p−r 的摘要),第二项对当前 query 现算 [p−r+1,m] 这一段。命中时前缀的计算和访存复杂度与上下文长度无关——KV 的 [1,p−r] 区间根本不需要读。这正是标题里的 99% KV 读取削减的来源。
一个直观类比:长对话里模型反复”查”同一段前文,等价于一个图书管理员每天被问相似的问题。常规做法是每次把整柜书重新翻一遍(全注意力);MAC 的做法是把查书结果记在便签上,相似的问题直接翻便签,只对”便签边界附近最容易变化的那几页”和”新到的书页”重新查——便签就是注意力摘要环,重新查的那几页就是修正带。

上图(论文 Figure 2)把流程画全了,三个阶段各有分工:
- Match(匹配):解码步 m,把当前 pre-RoPE query Q~m 与近期 query 环里的一小批候选逐一比较;最近邻 p 通过 L2 阈值检验就算命中,取出它的注意力摘要 AS1:p−r;否则回退到全注意力(输出与基线完全一致)。
- Complete(补全,关键路径):用施加位置编码后的真实 query qm 只对 [p−r+1,m](修正带+新尾部)计算注意力,然后与缓存摘要做 log 域合并(图中 ⊕)。KV 的 [1,p−r] 区间完全不读。
- Amend(修正,异步):为”未来被复用”做准备——构造刚产出 token 的修正后摘要,连同它的 query 一起写入两个环(图 (c))。它跑在辅助流上,不在关键路径。
注意”复用”和”修正”的顺序:缓存的摘要在创建时就已经把匹配点附近的短带”剪掉”了(截断到 p−r),所以复用方拿到的永远是”修正后摘要”,重算带是新 query 的义务。
与既有复用范式的边界#
MAC 不是第一个”复用注意力”的想法,但此前的复用都发生在别的维度上:
- 跨请求字面复用:Prompt Cache(2023 年提出,MLSys 2024)和 vLLM 的 prefix caching 复用”文本前缀完全相同”的跨请求计算;树形解码(如 DeFT,ICLR 2025)复用树形候选之间的共享前缀。它们依赖字面前缀重叠,MAC 的复用发生在单个请求流内部、基于语义相似性,与字面重叠无关——两者正交,可以叠加。
- 步间调度复用:Recycled Attention(Xu 等人 2024 年提出)让全注意力步与部分步交替执行,部分步只对上一次注意力模式的 top-k token 计算。MAC 每步都是完整的 Match-Amend-Complete,没有交替调度,且保留了全前缀的可达性。
- token 级投机复用:本博客前文《DSpark:半自回归生成与置信度调度的投机解码》讲的投机解码是在 token 层面”猜并验证”,MAC 是在注意力层面”复用并修正”——二者作用于推理的不同层级,同样可以组合。
原理详解:Match、Amend、Complete 各回答一个”为什么”#
三个阶段背后各有几个非显然的设计决策,逐个拆开。
Match(一):为什么在 pre-RoPE 空间匹配,而不是 post-RoPE#
从操作顺序看(QKV 投影 → RoPE → 注意力),在 post-RoPE 空间做匹配似乎更自然——毕竟注意力消费的就是位置编码之后的向量,而且有一个漂亮的数学保证:对任意 key Kt,由 Cauchy–Schwarz 不等式,
ℓt(m)−ℓt(p)=d(Qm−Qp)⊤Kt≤d∥Qm−Qp∥⋅∥Kt∥即 post-RoPE 距离小,就能对所有 key 统一地限制 logits 误差上界。这正是想要的性质:匹配度量直接约束最终误差。
但 RoPE 会毁掉这个度量。RoPE 由苏剑林等人于 2021 年在 RoFormer 中提出,它的本质是对 query/key 施加一个随位置旋转的正交变换 Qm=R(m)Q~m。把两个位置的 query 距离展开:
∥Qm−Qp∥2=∥Q~m−R(m−p)Q~p∥2即使 Q~m≈Q~p,距离里也被注入了相对旋转 R(m−p)。RoPE 的分块结构是 2×2 旋转的直积:
R(Δ)=diag(Rω1(Δ),…,Rωd/2(Δ)),x⊤R(Δ)x=j=1∑d/2cos(ωjΔ)∥x(j)∥2最后这个等式是说:向量 x 与其旋转 Δ 个位置后的自身做内积,等于按各频率分量能量加权的余弦平均。定义
avgcos(Δ)=∥x∥2∑jcos(ωjΔ)∥x(j)∥2则 ∥x−R(Δ)x∥2=2∥x∥2(1−avgcos(Δ))。RoPE 的频率 ωj 通常按几何级数铺开,间隔 Δ 稍大时,不同频率的余弦以不同速度去相位,加权平均很快趋近于零——相位破坏性混合。于是距离 ∥x−R(Δ)x∥2→2∥x∥2:两个语义完全相同的 query,只要位置隔得稍远,在 post-RoPE 空间就”长得完全不像”,匹配率大幅下降。这正是论文在实验里发现的:post-RoPE 匹配会显著减少近邻命中。
所以 MAC 的解法是把两件事拆开:匹配在 pre-RoPE 空间做——Q~m 不含位置相位,距离纯粹度量语义相似性,命中率稳健;保真度交给 Amend 阶段——在复用边界附近重算一小段高注意力质量的 token,把”post-RoPE 空间才能看到的误差”就地修掉。论文的表述很准确:这个组合”在维持稳健匹配频率的同时,恢复了 post-RoPE 匹配应有的实际保真度”。若没有命中,回退全注意力,输出与基线逐位一致——不存在”错误地复用”。
Match(二):为什么用 L2 而不是余弦,阈值从哪来#
相似度度量的候选是余弦和 L2。余弦只捕捉方向:两个共线但模长不同的 query 有相同的余弦值,却会诱导出不同尺度的 logits,进而产生不同的 softmax 锐度——注意力分布的形状变了,复用误差随之而来。L2 同时约束角度和模长,且直接对接上面 Cauchy–Schwarz 的误差上界:∥Qm−Qp∥ 小,最坏情况的 logits 误差就小。所以选 L2。
下一个问题是”多近算近”。论文用一个高斯零模型把阈值解析出来:假设两个无关的 pre-RoPE query 之差是各向同性高斯 Q~m−Q~p∼N(0,2Id),则其平方范数服从 2χd2 分布,在 d 较大时高度集中在均值 2d 附近,即随机巧合的典型距离是 2d。于是接受匹配的条件设为显著低于这个基线:
∥Q~m−Q~p∥<2d(1−τ),τ∈[0,1)参数 τ 的含义是”相对随机巧合的余量”。拿 LLaMA 的 d=128 算:2d=16,τ=0.45 时阈值为 8.8——只有距离不到随机典型值六成的候选才配被复用。论文在附录里进一步给出了 τ 与虚警率 α 的对应关系(基于 χ2 分布函数的解析式):
α≈Fχdq2(dq(1−τ)2),τ≈1−Fχdq2−1(α)/dq实际取值上,论文对上下文密集型任务(LongBench v2、RULER)用较宽松的 τ=0.45(更多复用),对生成密集型任务(LongGenBench)用更严格的 τ=0.75——注意 τ 越大阈值越小、命中越难、回退越多,这是”生成任务对注意力扰动更敏感”这一直觉的参数化。附录建议的实用区间是 τ∈[0.4,0.8],阈值在所有头上取同一个固定值,避免逐头调参的复杂性。
Match(三):为什么只在最近 K 个 token 的窗口里找#
候选集被限制在每个(分组)头最近 K 个 token 的窗口内,理由有三层:
- 复用的价值正比于匹配前缀的长度。命中的 p 离当前 m 越近,被跳过的前缀 [1,p−r] 越长,省下的 KV 读取越多;匹配一个远古位置的 query 收益微乎其微。
- 成本必须恒定。窗口固定,环的容量、匹配的访存量就都是常数 O(K),与序列长度无关——这是”命中时 O(1)“能成立的系统前提。
- 时间局部性。相似的 query 大概率出现在邻近位置,扩大窗口对命中率的边际贡献很小(论文的窗口扫描显示 K 超过约 4096 后收益递减,默认 K=1024)。
匹配 kernel 的实现同样围绕”轻”字:对每个候选只需累积两个内积,按平方距离展开式计算:
∥Q~m−Q~p∥2=∥Q~m∥2+∥Q~p∥2−2Q~m⊤Q~pbf16 读入、fp32 累加,算术强度极低,本质是访存受限的小 kernel。实现上按环分块扫描、用仅 warp 级的 shuffle 归约选出每个(请求,头)的最近邻(无 barrier 停顿)。实测它在 32K–120K 所有上下文长度下都是恒定的 9.1 μs(论文 Table 5)——这是整个方案”固定前端 + 可变尾端”成本结构的基石。需要注意 GQA 下的一个工程细节:匹配是按 query 头逐个做的,它要流式读全部 32 个 query 头的数据,而注意力 kernel 只按 8 个 KV 头流式读。所以当 GQA 压缩比很大(如 40:10)时,匹配 kernel 的相对访存量反而更高(论文 Figure 8a),这也是匹配 kernel 必须极致轻量化的原因。
Amend:为什么复用必须”修”,误差的一阶分析#
最朴素的做法是把缓存的前缀摘要整体照搬(等价于令所有 δt=0)。论文用 softmax 的一阶展开精确指出了它的病根。对 logits 的微小扰动,注意力权重 αt=eℓt/Z 的变化近似为
Δαt≈αt(p)(δt−Eu∼α(p)[δu])括号里减去扰动均值的一项来自归一化常数 Z 的同步变化。关键在于系数 αt(p):位置 t 的注意力权重越大,该位置的近似误差越大。而 softmax 的质量在解码游标附近高度集中——原因有二:位置编码(RoPE、ALiBi)带来的近因偏差,以及相邻 token 的语义相似性。于是误差几乎全部集中在匹配位置 p 附近的短带上。
修正方案呼之欲出:记 R={p−r+1,…,p} 为匹配点附近宽度 r 的短带,ρ=∑t∈Rαt(p) 为带内注意力质量占比。缓存的摘要在创建时就把这条带剪掉(只存 [1,p−r]),复用方则对 [p−r+1,m] 用当前 query 现算:
Sp−r+1:m(m)=t=p−r+1∑meℓt(m)Vt,Zp−r+1:m(m)=t=p−r+1∑meℓt(m)再与缓存摘要合并得到 om。论文用指数函数的 Lipschitz 界配合上式证明了修正后的前缀误差满足
o1:preuse−o1:p(m)≲(eΔ−1)(1−ρ)⋅Et∼αRˉ(p)[∥Vt∥]逐一解释这个界里的因子:Δ 是 logits 漂移的上界(∣δt∣≤Δ),它度量 query 到底有多”相似”;(1−ρ) 是带外残余注意力质量——只要带宽 r 选得够大、把大部分质量装进带内,误差就随残余质量线性衰减;E[∥Vt∥] 是 value 的尺度项。这个界把”修哪里、修多宽”从玄学变成了参数选择:让 ρ→1,误差就去零。

论文 Figure 4 的逐层热力图(横轴复用间隔 Δ=m−p,纵轴带宽 r,颜色越浅误差越小)给出了经验印证:带宽加到 r≈8 就消除大部分误差,r≥256 时与全注意力几乎不可区分;间隔越大越敏感,深层需要的带宽略宽于浅层。实用的带宽区间是 r∈[128,512],论文默认 r=256。附带一个诚实的小观察:缓存摘要在命中步上由”旧 query 的缓存部分 + 新 query 的新算部分”合并而成,所以环里的摘要是近似摘要,连续命中链上的误差如何逐级传播,论文没有给出严格分析(见”局限”一节)——目前的质量结论由经验实验支撑。
Complete:log 域合并与成本账#
Complete 阶段就是前面的合并恒等式。数值稳定性上,摘要在创建与合并时都减去当前最大值(running max),缓存里存的 Z 以对数形式(ln-LSE 标量)维护,因此多次链式合并不会产生指数溢出——这沿用了 FlashAttention 的在线 softmax 手法(Milakov & Gimelshein 2018)。合并是结合的,先复用再复用、或者带+尾各自计算再相加,结果一致。
最后是经济账。设 BKV 为每个缓存 token 的读取字节数,Bq 为匹配时每个候选 query 的读取字节数,匹配位置为 p、带宽 r、窗口 K,则省下的前缀读取必须覆盖匹配与重算的开销:
前缀省下的读取pBKV≳匹配开销KBq+带+尾重算rBKV论文把 K≤1024、r≤512 作为工程约束,此时盈亏平衡点只需要 p 在几十个 token 量级——对 128K 上下文而言几乎必然满足。反过来,当上下文短、批量小、复用不充分时,这个不等式会倒挂,论文在 32K 上下文、batch 1、低复用率(MAC-60%)的设定下实测到了小幅回退(Figure 6),并明确指出这是一个”该关掉 MAC、用全注意力”的回归点。
系统实现:三步微流水线,演进为融合持久化 kernel#
论文版本:K1 / K2 / K3 三步微流水线#
每个请求维护两个容量为 K 的环缓冲区:pre-RoPE query 环(最近 K 个 token 的 Q~)与修正摘要环(对应的 (Srect,Zrect) 与 ln-LSE)。环按全局长度模 K 寻址,插入是 O(1),容量与上下文长度无关。辅助显存的开销是 O(K) 而非 O(L):论文实测 128K 上下文下,K=1024 时约 4.7%(120K 下约 5%),K=2048 时 9.4%。可以粗算验证:query 环每层 K×32 头×128 维×2 B≈8 MiB,摘要环每层量级相同,32 层合计约 500 MiB,对比 120K 上下文约 15 GiB 的 KV 缓存(128K 时为 16 GiB)——“约 5%“与论文实测在同一量级。
每个解码步按最小同步原则执行三个 kernel:
- K1 Match:分块的 L2 最近邻扫描(上文已述),恒定 9.1 μs。
- K2 Amend & Complete:这里有个容易忽略的负载均衡问题——每个头的命中位置 p、带宽 r 各不相同,带+尾的重算跨度随头而异。实现把 (请求, 头) 轴拍平,按各头的工作跨度按比例分配 CTA,跨度大的头多分线程块。实测该调度策略相比朴素均分削减 29%–38% 的注意力阶段延迟,能追回与”完美均衡神谕”差距的 75%–80%;剩余 16%–21% 的失衡来自 CTA 粒度——一个 CTA 一旦被部分填满,在 GQA 的工作映射下就无法再细分。
- K3 Rectify-append:在辅助 CUDA 流上异步执行,与主流程只有一个事件依赖——它的结果要等下一解码步回到同一层才被消费。K3 做两件事:把当前步合并得到的完整摘要用 log 域减法(⊖)去掉最近 r 个 token 的贡献,得到位置 m 的修正后摘要 AS1:m−r;然后连同 q~m 一起写入两个环。全程原地操作,延迟可以忽略,稳态下可捕获进 CUDA Graph。
整个微流水线用伪代码概括(每请求、每层):
1# 每个解码步 m:2K1 Match: 对 query 环中最近 K 个候选 q_tilde[p] 累加3 d2 = ||q_tilde[m]||^2 + ||q_tilde[p]||^2 - 2 * q_tilde[m]·q_tilde[p]4 tile 级取最小 + warp shuffle 归约选出最近邻 p5 若 d2 >= 2d(1-tau)^2:未命中 → 全注意力(输出与基线逐位一致)6K2 Amend & Complete:7 取出摘要环中修正后摘要 (S_1:p-r, Z_1:p-r)8 对 [p-r+1, m] 用当前 query 现算 S、Z9 (CTA 按各头 band+tail 跨度比例分配,跨度大者多分线程块)10 o_m = (S_1:p-r + S_p-r+1:m) / (Z_1:p-r + Z_p-r+1:m) # log 域合并11K3 Rectify-append(辅助流,异步,原地操作):12 合并摘要 ⊖ 最近 r 个 token 的贡献 → 修正后摘要 AS_1:m-r13 将 (q_tilde[m], AS_1:m-r, ln-LSE) 写入环(mod K 寻址,O(1))
上图的三个面板对应:每请求环上的 L2 匹配(SplitQ 设计)、绿色块所示的按跨度 CTA 分配(很多头只需要一次廉价合并)、以及完整的数据流总览。
GQA 分组内还有一个被论文诚实报告的坑:匹配按 query 头做,同一 KV 组内的 4 个 query 头可能命中不同的 p;而实现把线程块映射到 KV 组、整组共享一次 KV 加载,只要组内有一个头复用率低,整组的工作量就被抬高。论文在 LongBench v2 上实测了组内匹配位置偏移(Table 6):第 0 层出奇地一致(所有头命中同一个历史 query,标准差 0),第 24 层的标准差达到 121——逐层行为差异巨大,这也是”逐层自适应调参”被列为未来工作的原因。
生产版本:一个融合的持久化 BF16 kernel#
论文描述的三 kernel 结构在开源仓库里已经演进为更激进的形式。2026 年 7 月的 kernel 更新(见仓库 Release Notes)把整个解码路径融合成一个持久的 BF16 CUDA kernel mac_persistent_decode_bf16,kernel 内部完成:环内匹配、逐头/逐组命中分类、命中与未命中混合的负载调度、带+尾的局部注意力、log-sum-exp 合并、输出写回、以及为下一个 token 准备的 query/摘要/LSE 缓存更新。旧版 0.1.0 参考实现(独立环匹配扩展 + 包 FlashInfer 式分页 KV 注意力的解码包装 + 分离的缓存更新路径)与之并存,融合版的动机是消除 token 解码热路径上的 host 侧编排开销。
工程细节里有几处值得学习:
- GQA 专用分块路径:GQA 组大小不超过 4 时用 block-128,GQA-8(如 Qwen3-30B 的 32 查询头/4 KV 头)走专门的 block-256 路径。
- 匹配扫描重写:向量化的成对槽位探测(paired-slot probes)配合精确下界过滤,替代朴素扫描;统一的动态共享内存池提高 SM 占用率;更深的
cp.async暂存环;波次量化(wave-quantized)的兜底/尾部调度。 - 注意力汇点隔离:环境变量
MAC_FRONT_SINK_TOKENS=F把上下文最前面 F 个 token(StreamingLLM 发现的 attention sink)排除在复用摘要之外,每次命中都精确重算它们——防止陈旧的汇点贡献被摘要带着跨步传播。默认 F=0(纯带宽修正)。 - CUDA Graph 安全的就绪门控:每个请求在设备侧维护 ready/epoch 状态,摘要缓存未就绪或已失效的请求即使在捕获好的 Graph 里也会被降级为精确全注意力;模型上下文超过
MAC_PERSISTENT_MAX_CONTEXT(默认 131072)时拒绝进入 Graph。这是把”回退逻辑”塞进静态 Graph 的典型做法。 - 接入方式:以 SGLang 可移植插件的形式提供运行时钩子,不修改 SGLang 源码,验证于 SGLang 0.5.13.dev84(CUDA 13.0,H100)。官方 SGLang 仍负责模型执行、请求调度、分页 KV 分配与 FlashInfer 后端集成,MAC 钩子只负责保留 query 状态、维护环缓存、拦截支持的 BF16 分页 KV 解码调用并发射融合 kernel。
- 可审计性:
MAC_WORKFLOW_AUDIT=1时每个 MAC 解码步之后重跑一遍精确 FlashInfer 并记录误差统计(诊断专用,禁止在性能测量时开启)。融合 kernel 在一个 853 例矩阵(上下文至 127K、GQA 4/8、前汇点 0/64/128、命中率 0–1、CUDA Graph 重放)上与参考实现逐例对拍:命中决策完全一致,输出与 FP32 oracle 的相对 L2 误差在 1e-2 容差内。
生产默认配置与论文略有出入:MAC_THRESHOLD=0.45、MAC_LOOKBACK_TOKENS_LEFT=512(窗口,论文默认 1024)、MAC_SEMANTIC_POS_AHEAD=256(修正带)、MAC_GEN_MIN_LIMIT=2048(上下文不足 2048 时不启用 MAC)——窗口减半是工程化权衡,论文的窗口扫描显示 512 与 1024 的收益差有限。仓库给出的推荐环境变量如下:
1export MAC_THRESHOLD=0.45 # L2 匹配阈值(相对随机巧合的余量)2export MAC_LOOKBACK_TOKENS_LEFT=512 # 环搜索窗口 K3export MAC_SEMANTIC_POS_AHEAD=256 # 修正带宽 r4export MAC_GEN_MIN_LIMIT=2048 # 上下文不足 2048 不启用 MAC5export MAC_FRONT_SINK_TOKENS=0 # 前部注意力汇点隔离数(0=默认带宽修正)6export MAC_DISABLE_CUDA_GRAPH=0 # 保持 CUDA Graph 捕获7export MAC_PERSISTENT_PARTIAL_FP32=1 # 局部注意力路径使用 FP32 工作区8export MAC_USE_FUSED_KV_ROPE=1 # 融合 KV RoPE实验:注意力阶段快 14 倍以上,端到端快 2.6 倍,精度不降反升#
实验设置#
论文的评测在 SGLang 0.4.9 + FlashInfer 0.2.7 上完成,硬件为 NVIDIA H100 SXM5(CUDA 12.8.1),对比基线统一为最新 FlashInfer 全注意力。模型覆盖 LLaMA-3.1-8B / 70B、Phi-4-Mini、GLM-4-32B,外加 MoE 模型 Qwen3-30B-A3B-Instruct;基准为 LongBench v2(至 120K 上下文,含 QA/摘要/检索,按 Short/Medium/Long 分层)、RULER(固定 120K,专攻延迟回忆与长度外推的合成探针)与 LongGenBench(连续生成至 16K token,LLM-as-judge 评分)。默认参数 K=1024、r=256,上下文密集型任务 τ=0.45、生成密集型任务 τ=0.75。所有实验固定随机种子、temperature 为 0,唯一变量是注意力路径。
质量:在 1% 的 KV 预算上跑出更高精度#
LLaMA-3.1-8B 在 LongBench v2 上:全注意力 29.0 分,MAC-Attention 30.2 分,KV 预算只有 1%;作为对照,Quest 在 2% 预算(Quest-1K)下掉到 26.2 分。注意方向是反的:省掉 99% 的 KV 读取,分数反而涨了 1.2 分。70B 模型同样如此(31.4 → 32.2,KV↓99%)。LongGenBench 上 8B 的完成率从 93.1% 微降到 90.0%,但按完成样本加权的精度从 34.7 升到 38.2(KV↓80%)——生成质量没有随预算收紧而塌方。RULER 上 8B 为 79.8 → 78.8(KV↓95%),70B 为 80.1 → 78.0(KV↓99%),论文把这两处小幅回退如实列出。跨模型外推同样成立:Phi-4-Mini 在 LongBench v2 上 28.0 → 29.8(KV↓91%),GLM-4-32B 在 LongGenBench 上完成率 69.8 → 66.1、完成加权精度 29.7 → 30.1(KV↓70%)。
分数偶尔高于基线这件事,论文自己的解释是长上下文全注意力计算本身带有噪声,并明确标注为”推测”。更值得记取的是 LongGenBench 的意义:它是长文本连续生成,注意力上的微小扰动会经过数千个解码步级联放大——此前几乎所有 KV 高效方法都不报告这个基准,MAC 是少数正面迎战的,其完成加权精度不降反升,说明”复用+局部修正”的误差结构对级联放大有抵抗力。
延迟:别家方法比全注意力还慢,MAC 是唯一全线更快的#
论文 Table 3 在 120K 上下文、LongBench v2 上对比了多个高效注意力方法的端到端注意力延迟(μs):
| KV 预算 | 全注意力 | Quest | RocketKV | Multipole | MAC-Attention |
|---|---|---|---|---|---|
| 1% | 234.2 | 581.2 | 822.8 | 192.4 | 62.9 |
| 5% | 234.2 | 594.7 | 844.7 | 210.8 | 64.0 |
| 10% | 234.2 | 608.5 | 1042.5 | 265.4 | 78.1 |
| 20% | 234.2 | 640.5 | 1855.6 | 324.6 | 103.8 |
这个表的信息量很大:Quest、RocketKV 在每种预算下都比什么都不做的全注意力 FlashInfer 还慢(最夸张的是 RocketKV 在 20% 预算下慢 7.9 倍),Multipole 只在 1% 预算下勉强略快。原因正是它们的选择/压缩前置 pass 本身随上下文扩展,额外开销吃掉了理论收益。MAC 的固定前端(9.1 μs 匹配 + 约 30 μs 负载均衡)与上下文无关,这是它在所有预算下唯一全线低于基线的原因。
细粒度分解(论文 Table 5,K=1024)把这一点钉得更死:匹配 kernel 在 32K/60K/120K 下都是 9.1 μs;负载均衡器在 28.6–30.6 μs 之间几乎不动;真正随上下文增长的是带+尾注意力 kernel——但它在 1%–5% 预算下即使 120K 也只有约 25 μs。整体注意力路径延迟在 1% 预算下从 32K 的 63.9 μs 走到 120K 的 63.6 μs,几乎不随上下文变化;而全注意力同期从 71.6 μs 涨到 234.2 μs。
加速比:批量越大、上下文越长,收益越陡#
论文 Figure 6 展示了加速比随批量与上下文长度的系统性增长:32K 上下文、99% 复用率下,注意力阶段加速从 batch 1 的约 1.1 倍爬升到 batch 32 的约 13.5 倍;固定 batch 32 时,99% 复用率的加速比从 32K 的 13.5 倍增长到 64K 的 21 倍、128K 的 34 倍、256K 的 46 倍。方向完全符合 IO-bound 的预期:批量越大,基线越接近 HBM 带宽饱和,省下的每一字节 KV 读取都直接转化为墙钟时间;上下文越长,注意力占比越高,同样的削减比例撬动更大的绝对收益。反方向同样成立:64K–128K 下高复用率的收益远超 32K,因为固定前端在更长的上下文上被摊得更薄——这是 Amdahl 定律的另一面。
端到端层面(Figure 7),MAC 只改注意力路径,QKV+RoPE、FFN、运行时开销原封不动,这部分构成了端到端加速的上限。实测结果:128K 上下文下每 token 生成延迟下降 60% 以上,LLaMA 系列端到端最高加速 2.6 倍。收益随上下文长度单调增长——上下文越长,注意力在解码画像中的份额越大,MAC 能撬动的部分越多。
跨架构与逐层行为#
MoE 的动态专家路由会打乱 token 间的语义连续性,MAC 的匹配是否还成立?论文在 Qwen3-30B-A3B-Instruct 上验证:命中率 99.5%–99.7%,前缀跳过率 98.6%–99.0%,精度 36.6–37.6 对全注意力的 37.0,在 K∈{512,1024,2048,4096} 上全部稳定——专家路由没有破坏语义匹配结构。
逐层剖面(Figure 5)显示复用率强烈依赖层:浅层自相似性高、几乎步步命中;部分中层与高层波动大,需要更大窗口。阈值扫描(τ∈[0,0.9],K=2048)与窗口扫描(K∈{32,…,2048},τ=0.75)都指向同一结论:轻量级的逐层调参会进一步抬高命中率与跳过率,这也是论文列出的未来工作。
局限与未解决的问题#
- 命中率是全部前提,未命中即回退。MAC 没有最坏情况保证:未命中时执行完整全注意力,与基线同速但不会更快。对注意力分布平坦、缺少短期语义重复的负载(论文 B.5 明确列举了”全局平坦注意力或特异的头行为”),收益自然衰减。命中率 >99% 是长上下文场景下的经验观察,不是保证。
- 存在明确的”该关掉”区间。小批量 + 短上下文 + 低复用率时,约 40 μs 的固定前端(匹配+负载均衡)压过省下的 KV 读取,实测出现小幅回退(32K、batch 1、MAC-60%)。生产系统需要把这个经济学判断做成自动开关,论文与仓库目前都没有给出运行时自适应方案。
- GQA 组内失衡的残差。CTA 粒度决定了 16%–21% 的负载失衡无法再被修正;组内单个低复用头会拖累整个 KV 组的共享加载。
- prefill 阶段尚不可用。论文附录 A.3 的探索性实验:把 MAC 用于 prefill(全局窗口匹配 + 短带修正),即使命中率最高 99%、前缀跳过率最高 86%,总体精度仍低于全注意力(最好也只有 27.4 对 29.0)——原因是 prefill 块内的复用间隔经常横跨整个块、post-RoPE 漂移随间隔累积,短带修不过来。论文给出的方向是:复用距离上限(1–2K)、随间隔自适应的带宽、pre-RoPE L2 与 post-RoPE 漂移代理的双门限接受、以及更小的 prefill 分块。默认配置因此只作用于 decode。
- 连续命中链上的误差传播没有严格分析。如 Amend 一节所注,命中步写入环的摘要是”旧 query 缓存部分 + 新 query 新算部分”的混合体,长期链式复用下的误差行为论文未给出理论刻画,目前的质量结论依赖经验评测与仓库的审计工具(
MAC_WORKFLOW_AUDIT)。这是该方向值得做理论工作的空白点。 - 工程成熟度仍有缺口。仓库 Roadmap 明确列出:端到端 SGLang 服务的吞吐/延迟对比尚未发布;模型覆盖目前以 LLaMA-3.1 家族为主(另有 Qwen2/Qwen3 MoE 钩子);验证目标是 H100 级硬件;接入面目前只有 SGLang,尚无 vLLM 移植。
- 假设的适用范围。方案依赖近因偏向的注意力分布(RoPE/ALiBi 模型满足);对注意力汇点的处理上,生产版默认仍把汇点含在摘要里(
MAC_FRONT_SINK_TOKENS=0),只有显式开启隔离才精确重算。
小结#
MAC-Attention 的可贵之处不在某单个技巧,而在问题定义方式的转换:长上下文解码的 IO 瓶颈面前,压缩改动了”信息的表示”,驱逐改动了”信息的可达性”,二者都动了注意力计算的语义;MAC 选择复用”计算的结果”,把对完整序列的访问权原封不动地保留下来,误差只从一条可控的通道进入——匹配点附近的高质量带——再用固定宽度的重算把这条通道封死。Match 回答”复用谁的”,Amend 回答”哪里会出错”,Complete 回答”怎么把新旧结果拼起来”,三个阶段的每一个设计决策都有对应的解析式或实验依据:pre-RoPE 匹配有 RoPE 旋转结构的推导,L2 阈值有高斯零模型,带宽有 softmax 一阶展开和误差界,成本有盈亏不等式。
系统层面的启示同样清晰:把随上下文增长的成本压缩成”固定前端 + 可变尾端”,是长上下文推理优化的通用形态——BLASST 的标量阈值判定、QEvict 的晋升机制、MAC 的 9.1 μs 匹配 kernel,殊途同归。MAC 的开源实现进一步展示了这类方法的工程化路径:从论文的三 kernel 微流水线,演进为融合持久化 kernel、CUDA Graph 就绪门控、可审计的精确回退,以及不侵入框架源码的可移植插件。对想要动手的读者,论文附录 B 的调参指南(K∈{1024,2048} 起步、τ∈[0.4,0.8]、r 取”游标附近累计质量超过 1−ε“的最小带宽)和仓库的复现脚本是现成的入口。
参考资料#
- MAC-Attention: a Match-Amend-Complete Scheme for Fast and Accurate Attention Computation(arXiv:2604.00235,MLSys 2026)
- MAC-Attention 官方代码仓库(YJHMITWEB/MAC-Attention)
- MAC-Attention OpenReview 页面
- FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness
- FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning
- Online normalizer calculation for softmax(Milakov & Gimelshein 2018)
- Efficient Memory Management for Large Language Model Serving with PagedAttention
- RoFormer: Enhanced Transformer with Rotary Position Embedding
- Train Short, Test Long: Attention with Linear Biases Enables Input Length Extrapolation(ALiBi)
- RULER: What’s the Real Context Size of Your Long-Context Language Models?
- LongBench v2: Towards Deeper Understanding and Reasoning on Realistic Long-Context Multitasks
- Benchmarking Long-form Generation in Long-context LLMs(LongGenBench)
- Quest: Query-Aware Sparsity for Efficient Long-Context LLM Inference
- H2O: Heavy-Hitter Oracle for Efficient Generative Inference of Large Language Models
- SnapKV: LLM Knows What You Are Looking for Before Generation
- Efficient Streaming Language Models with Attention Sinks(StreamingLLM)
- RocketKV: Accelerating Long-Context LLM Inference via Two-Stage KV Cache Compression
- Multipole Attention for Efficient Long Context Reasoning
- Recycled Attention: Efficient Inference for Long-Context Language Models
- Prompt Cache: Modular Attention Reuse for Low-Latency Inference
- GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints
- Fast Transformer Decoding: One Write-Head is All You Need(MQA)
- FlashInfer: Efficient and Customizable Attention Engine for LLM Inference Serving
- SGLang: Efficient Execution of Structured Language Model Programs
- The Llama 3 Herd of Models
文章分享
如果这篇文章对你有帮助,欢迎分享给更多人!
部分内容可能已过时
评论区
分享你的想法,与大家交流讨论
音乐
暂未播放



