Attention 演进路线:5 大阶段,11 篇奠基论文

原文:​https://mp.weixin.qq.com/s/UuKiNeQkoNa7RWm2Llgllw

注意力机制走过了十余年的技术迭代。它最初只是循环神经网络里的一个对齐插件,后来成长为大模型时代的算力核心。梳理这条技术脉络,可以看到一条非常清晰的演进主线。早期研究聚焦于「怎么让模型学会抓重点 」。中期研究转向「怎么摆脱循环结构实现并行训练 」。近期研究则全面攻坚「怎么突破显存带宽与超长上下文的物理极限」。

全文目录

  • 阶段1:RNN 时代的软寻址起源(2014至2015)

  • 阶段2:抛弃循环与全自注意力革命(2017)

  • 阶段3:突破序列长度平方瓶颈(2020)

  • 阶段4:消解显存带宽瓶颈与KV缓存压缩(2019至2024)

  • 阶段5:硬件感知与系统级协同设计(2022至2025)

  • 演进总结与未来展望

为了直观呈现这一技术脉络,下表汇总了五个演进阶段的十一篇里程碑论文及其核心突破。

演进阶段 代表机制与年份 核心论文 解决的核心痛点 关键技术突破
软寻址起源 Bahdanau Attention(2014) arXiv 1409.0473 固定长度向量无法容纳长句信息 引入加性对齐分数与动态上下文向量
软寻址起源 Luong Attention(2015) arXiv 1508.04025 加性注意力计算偏慢且长句对齐分散 提出点积打分与局部滑动窗口对齐
全自注意力架构 Multi-Head Attention(2017) arXiv 1706.03762 循环网络无法并行训练且长程依赖易衰减 抛弃循环结构,用缩放点积与多头自注意力建模全局
复杂度降维探索 Longformer(2020) arXiv 2004.05150 全注意力计算量随序列长度呈平方级爆炸 结合局部滑动窗口与任务驱动的全局稀疏注意力
复杂度降维探索 Linear Attention(2020) arXiv 2006.16236 Softmax 锁死矩阵乘法顺序导致平方复杂度 用核函数近似改写乘法结合律,将复杂度降为线性
键值缓存压缩 Multi-Query Attention(2019) arXiv 1911.02150 自回归解码频繁搬运键值缓存导致显存带宽受限 所有查询头共享同一组键头与值头,大幅削减显存读取
键值缓存压缩 Grouped-Query Attention(2023) arXiv 2305.13245 多查询注意力在大规模模型上容易出现精度下滑 将查询头分组共享键值头,兼顾解码速度与模型质量
键值缓存压缩 Multi-head Latent Attention(2024) arXiv 2405.04434 分组共享仍需缓存较多键值状态,长文本显存压力大 利用低秩联合压缩将键值投影为紧凑潜向量并解耦位置编码
硬件与系统协同 FlashAttention(2022) arXiv 2205.14135 注意力中间矩阵频繁读写高带宽显存造成严重算力闲置 利用片上缓存分块计算与在线归一化,彻底省去中间矩阵落盘
硬件与系统协同 PagedAttention(2023) arXiv 2309.06180 连续显存预分配造成严重内存碎片与显存浪费 借鉴操作系统虚拟内存分页机制,按固定块动态管理键值缓存
硬件与系统协同 Native Sparse Attention(2025) arXiv 2502.11089 传统稀疏注意力难以端到端训练且硬件访存效率低下 采用块级压缩、块级精选与滑窗三路并行,实现硬件对齐原生训练

阶段1:RNN 时代的软寻址起源(2014 至 2015)

在注意力机制诞生之前,序列到序列模型全部依赖循环神经网络。编码器必须把整句话硬生生塞进一个固定长度的向量里。句子一旦变长,前面的信息就会被后面的词冲淡。这一阶段的两篇开创性论文打破了固定向量瓶颈。它们让解码器在生成每一个词时,都能回头翻看源句子的所有位置。

Bahdanau Attention(加性注意力与软对齐开创者)

早期的神经机器翻译采用经典的编码器与解码器框架。编码器像一个严苛的压缩机。无论输入的英文句子是五个词还是五十个词,它都必须压缩成一个固定维度的向量。解码器再仅凭这一个向量去还原整句法文。

你想想看,这就像让译员听完一整段长篇演讲后不许看笔记,只凭脑子里的笼统印象去逐字翻译。当句子长度超过二十个词时,翻译质量就会断崖式下跌。作者敏锐地指出,强行把变长序列压进定长向量,正是阻碍长句翻译的罪魁祸首。

这篇论文首次在神经机器翻译中引入软对齐机制。它让模型在生成每个目标词时动态检索源句中最相关的词,打破了定长编码向量的信息瓶颈。

作者提出了名为 RNNsearch 的全新架构。编码器不再只输出句末的那一个状态,而是采用双向循环神经网络。它把每个词的前后文信息都编码成各自的隐藏状态注解。整句话变成了一排随时可供翻阅的「词卡片」。

解码器在生成第 \(i\) 个目标词时,不再死盯着同一个向量。它会拿自己上一时刻的隐藏状态,去和源句里每一张「词卡片 」做匹配。匹配过程通过一个小型前馈网络实现。模型把两个状态向量相加后再经过非线性激活,计算出每个源词的相关性得分。这种打分方式后来被称为「加性注意力」。

接着,模型用归一化指数函数把这些得分转换成概率分布,作为每张卡片的权重。把所有卡片按权重加权求和,就得到了当前词专属的动态上下文向量。对齐模型与翻译模型完全融为一体,通过反向传播端到端联合训练。

这篇论文真正开启了深度学习的注意力时代。它证明神经网络不需要把所有历史信息死记硬背在单一状态里。相反,模型可以通过动态软寻址按需提取信息。这种「查询匹配并加权聚合」的思想,直接奠定了后续所有注意力变体的数学雏形。

Luong Attention(乘性注意力与局部窗口机制)

Bahdanau 的加性注意力虽然解决了长句瓶颈,但留下两个明显的工程痛点。① 加性打分需要额外的权重矩阵和非线性激活函数,计算开销偏大。② 每次生成一个词都要遍历源句的所有位置。当源文档非常长时,全局扫描不仅浪费算力,还会因为注意力过度分散而找不准重点。学术界迫切需要一种算得更快、聚焦更准的注意力机制。

这篇论文系统梳理了全局与局部两种注意力范式。它用极简的矩阵点积替代了繁琐的加性网络打分,为后续 Transformer 的缩放点积注意力铺平了道路。

作者从架构与打分函数两个维度对注意力机制做了大刀阔斧的简化。在架构上,模型不再借用解码器上一时刻的状态。它直接用解码器当前时刻的顶层隐藏状态去检索源端信息,缩短了计算路径。

在打分函数上,作者对比了三种简化方案。最核心的是「Dot点积打分 」,也就是直接把源端向量和目标端向量做内积。另一种是「通用双线性打分 」,在两个向量中间乘一个可学习的权重矩阵。相比于「加性注意力 」,点积打分直接转化为高度优化的矩阵乘法,也就是后来常说的「乘性注意力」。

为了解决长序列全局扫描的低效问题,作者进一步提出了「局部注意力 」(Local Attention)。模型先预测当前目标词对应源句中的一个锚点位置,然后只在这个锚点左右固定半径的窗口内计算注意力。此外,作者还设计了「输入馈送」(Input-feeding)机制。它把上一步算出的注意力向量拼接到下一步的输入中,让模型清楚记住哪些源词已经翻译过了。

Luong 的工作完成了注意力机制从「可行的对齐插件 」向「高效的矩阵算子」的关键蜕变。他证明了简单的向量点积足以捕捉复杂的语义对齐关系。两年后横空出世的 Transformer,正是直接沿用了这里的点积打分思想。而局部滑动窗口的设计,也成为了日后稀疏注意力的早期先驱。

阶段2:抛弃循环与全自注意力革命(2017)

尽管注意力机制在翻译任务上大放异彩,但它依然挂靠在循环神经网络身上。只要循环结构还在,模型就必须按时间步一个词接一个词地顺序计算。图形处理器的海量并行算力根本无法施展。2017 年,Google 团队提出了一个大胆的问题,如果彻底扔掉循环和卷积,只留注意力机制,模型还能运转吗?

Transformer 标准多头自注意力(MHA)

在 2017 年之前,长短期记忆网络和门控循环单元是序列建模的绝对主力。这类模型有着天然的物理缺陷,第 \(t\) 步的隐藏状态必须等第 \(t-1\) 步算完才能开始。序列越长,训练时的等待链条就越长,单样本内部根本无法并行化。

不仅如此,序列首尾两个词之间的信息传递需要跨越整条循环链路。距离一远,梯度就容易消失或爆炸。虽然 ByteNet 和 ConvS2S 尝试用多层卷积实现并行,但它们捕捉远距离依赖所需的层数仍然随距离增长。如何让任意两个位置的词都能在常数步内直接对话,同时让整句话一次性并行算完,成了当时的头号难题。

这篇里程碑论文彻底摒弃了循环与卷积结构。它首次完全基于缩放点积自注意力与多头注意力构建序列模型,开启了现代大语言模型的黄金时代。

作者提出了完全由注意力层和前馈网络堆叠而成的 Transformer 架构。其心脏是「缩放点积注意力」(Scaled Dot-Product Attention)。模型把输入序列中的每个词向量分别线性映射成三个角色,即查询向量(Query)、键向量(Key)和值向量(Value)。

所有位置的查询向量拼成矩阵 \(Q\),键向量拼成矩阵 \(K\),值向量拼成矩阵 \(V\)。模型直接计算 \(Q\) 与 \(K\) 转置的矩阵乘积,一次性算出任意两个词之间的相关性。这里有一个极其精妙的细节。当向量维度 \(d\_t\) 较大时,点积结果的方差会变大,很容易把 Softmax 函数推入梯度极小的饱和区。作者因此将点积结果除以缩放因子 \(\sqrt{d_k}\),稳住了训练梯度。

为了让模型同时关注不同维度的语法和语义信息,作者设计了「多头注意力 」(MHA)。它把高维的 \(Q\)、\(K\)、\(V\) 切分成 \(h\) 个低维子空间(原文取 8 个头)。每个头独立做缩放点积注意力,再把结果拼接起来做线性投影。这样一来,有的头可以专门盯着指代关系,有的头可以盯着动词搭配,而总计算量与单头全维注意力几乎持平。此外,由于自注意力本身对词序不敏感,作者在输入端加上了正弦余弦位置编码,为模型注入位置先验。

这篇论文彻底改写了人工智能的技术版图。它将任意两个词之间的交互路径长度缩短到了 \(O(1)\),并把训练过程完全转化为显卡最擅长的稠密矩阵乘法。此后无论是 BERT、GPT 系列,还是视觉领域的 ViT,全都建立在多头自注意力这块基石之上。但与此同时,矩阵 \(Q K^T\) 带来的 \(O(N^2)\) 时间与显存复杂度,也埋下了后续数年所有注意力改进工作争相攻克的新隐患。

阶段3:突破序列长度平方瓶颈(2020)

Transformer 一统江湖后,研究者很快撞上了一堵墙。标准自注意力需要计算每个词与其他所有词的配对得分。序列长度 \(N\) 每翻一倍,计算量和显存占用就会翻四倍。处理五百个词游刃有余,处理几万字的长篇小说或基因序列时,显卡显存瞬间就会爆掉。2020 年前后,学术界掀起了一场「高效注意力 」热潮,主要分化出「稀疏化剪枝 」与「核函数线性化」两条技术路线。

Longformer(局部滑窗与全局稀疏注意力)

以 BERT 和 RoBERTa 为代表的预训练模型,最大上下文长度通常被限制在 512 个词元。面对长篇新闻、法律合同或多文档问答时,工程师只能无奈地把长文切碎成小段分别处理。段与段之间的跨篇章逻辑因此完全断裂。

其实,自然语言具有很强的局部聚集性。绝大多数词只需要看懂它前后几十个词就能确定词义,只有极少数关键句或问题词才需要统揽全文。标准自注意力让所有词之间两两连线,绝大部分算力都浪费在了毫无关联的远距离词对上。

这篇论文将「局部滑动窗口 」与「少数任务驱动的全局节点」相结合。在不牺牲长文档理解能力的前提下,它将自注意力的计算复杂度从平方级成功降为线性级。

作者为长文档量身定制了一套混合稀疏注意力模式。第一层机制是「滑动窗口注意力」(Sliding Window Attention)。每个词只关注自己左右固定窗口大小 \(w\) 范围内的邻居。既然窗口大小 \(w\) 是固定的常数,整篇文档的计算复杂度就从 \(O(N^2)\) 直接降到了 \(O(N \times w)\),随序列长度呈线性增长。就像卷积神经网络一样,当多层滑动窗口堆叠起来后,顶层节点的感受野就会自然覆盖到数千词的范围。

为了进一步扩大感受野而不增加计算量,作者在高层引入了「空洞滑动窗口」(Dilated Sliding Window)。窗口内的注意力每隔几个词跳着采样一次,用同样的算力看清更远的距离。

仅有局部窗口还不够,像阅读理解任务中的「问题词 」或分类任务中的句首标记,必须能够随时读取全文信息。作者因此加入了不对称的「全局注意力」(Global Attention)。被预先指定为全局节点的少量词元,可以向全文所有词发起关注,全文所有词也会关注它们。由于全局节点数量极少,整体复杂度依然保持严格的线性规模。研究团队还用 TVM 编写了专门的分块对角矩阵 CUDA 算子,确保稀疏矩阵不会退化成缓慢的循环。

Longformer 确立了「局部滑窗加全局锚点」这一经典的稀疏注意力范式。它告诉业界,长文本建模根本不需要全连接图。只要保留局部连续性并留出少量全局信息枢纽,就能用线性成本撬动超长上下文。这一设计直接启发了后来的 BigBird、Mistral 滑动窗口机制以及现代原生稀疏架构。

Linear Attention(核函数近似与线性复杂度)

稀疏注意力虽然降低了计算量,但往往需要定制复杂的掩码规则和专用算子,而且视野受限于预设的稀疏图结构。另一批学者从纯代数角度思考,为什么标准注意力非得花费 \(O(N^2)\) 的代价不可?

症结其实出在 Softmax 函数上。在计算 \(\text{Softmax}(Q K^T) V\) 时,\(Q\) 和 \(K\) 的维度都是 \(N \times d\)。因为外面套了一层非线性的指数归一化,计算机必须先算 \(Q K^T\) 得到一个 \(N \times N\) 的大矩阵,再乘以 \(V\)。如果能拿掉外层的指数套壳,先算 K\^T V 得到一个小小的 \(d \times d\) 矩阵,再左乘 \(Q\),计算量不就瞬间从 \(O(N^2 d)\) 变成 \(O(N d^2)\) 了吗?

这篇论文利用核函数映射拆解了 Softmax 算子,通过改变矩阵乘法的先后顺序将自注意力复杂度降为线性。同时,它揭示了自回归线性注意力与循环神经网络的数学等价性。

作者从广义相似度的视角重构了注意力公式。标准注意力中的 \(\exp(q\_i^T k\_j)\) 其实可以看作一种衡量向量相似度的核函数。根据核方法理论,任何非负核函数都可以写成两个特征映射函数 \(\phi(q\_i)^T \phi(k\_j)\) 的内积。作者选用了一个极其轻巧的逐元素特征映射函数 \(\phi(x) = \text{elu}(x) + 1\),既保证了数值非负,又避免了高昂的映射开销。

一旦把相似度写成 \(\phi(Q) \\phi(K)^T\),矩阵乘法的结合律就被彻底释放了。模型不再计算巨大的 \(\phi(Q) \phi(K)^T\),而是先计算 \(\sum\_j \phi(k\_j) v\_j^T\)。这个聚合项的尺寸仅仅是特征维度 \(d \times d\),与序列长度 \(N\) 完全无关。每个查询向量 \(\phi(q\_i)\) 只需要和这个 \(d \\times d\) 的全局摘要矩阵做一次乘法,就能以 \(O(N)\) 的线性复杂度完成全文交互。

更妙的是在自回归生成场景下。当模型逐词生成时,每新增一个词,只需要把当前的 \(\phi(k_t) v_t^T\) 累加到上一时刻的状态矩阵里即可。作者由此证明,带有因果掩码的线性 Transformer,其实就是一个拥有矩阵隐藏状态的循环神经网络。它在训练时可以全序列并行,推理时则只需常数级显存和单步常数时间。

这篇论文打通了 Transformer 与 RNN 之间看似不可逾越的理论高墙。它不仅开创了线性注意力的数学流派,更直接催生了后来的 RetNet、RWKV、Mamba 状态空间模型以及 Gated DeltaNet 等混合线性架构。尽管纯核函数近似在超大规模语言模型的精确大海捞针任务上略逊于 Softmax,但它与标准注意力的混合堆叠已成为当下超长上下文模型的主流方案之一。

阶段4:消解显存带宽瓶颈与 KV 缓存压缩(2019 至 2024)

进入千亿参数大模型时代后,工程师们发现了一个反直觉的现象。大模型在线上陪用户聊天、逐字吐出回复时,瓶颈根本不是显卡算得不够快,而是显存带宽读得不够快。

在自回归解码阶段,为了避免重复计算历史词的键(Key)和值(Value),系统会把它们存进显存,这就是著名的「KV 缓存」(KV Cache)。每生成一个新词,显卡都必须把庞大的历史 KV 缓存从高带宽显存(HBM)搬运到计算核心里。算力花在等待数据搬运上的时间占了九成以上。这一阶段的核心演进路线,就是想方设法给 KV 缓存「瘦身」。

Multi-Query Attention(多查询注意力 MQA)

在标准的多头注意力(MHA)中,假如有 \(h\) 个注意力头,每个头都有自己专属的查询矩阵、键矩阵和值矩阵。在训练阶段,整句话几千个词一起并行算。读取一次权重和缓存可以服务大批量数据,计算密度很高,显存带宽完全够用。

但到了在线推理的解码阶段,模型每次只能蹦出一个词。为了算这一个词,显卡必须把过去所有词在全部 \(h\) 个头上的 KV 缓存完整读一遍。当批次增大和上下文拉长时,KV 缓存动辄占用几十 G 字节显存。每次解码都像为了喝一口水而把整个水库搬运一遍,内存访问与算术运算的比例严重失衡,解码延迟居高不下。

这篇论文提出让所有查询头共享唯一的一组键头和值头。它将解码阶段的 KV 缓存体积和显存带宽读取量直接缩减为原来的头数分之一,大幅提升了推理速度。

Transformer 原作者之一 Noam Shazeer 给出了一个极其干脆的手术方案,取名为「多查询注意力」(Multi-Query Attention,简称 MQA)。

在 MQA 中,查询(Query)依然保留原本的 \(h\) 个独立头,保证模型依然可以从多个不同的子空间视角去提问。但是,所有的查询头不再各自维护专属的键(Key)和值(Value),而是共同读取唯一的一组单头键和单头值。作者形象地把键和值称为「写入头」。他认为模型只需要一个写入头把历史信息写进缓存,再由多个查询头从不同角度去读取就足够了。

这一改动没有减少单步注意力的浮点乘法次数,却把需要保存和搬运的 KV 缓存张量直接缩小了 \(h\) 倍(例如 8 头模型直接缩小 8 倍,64 头模型直接缩小 64 倍)。显存带宽压力迎刃而解。

Shazeer 在 2019 年就预见了大模型推理的「内存墙」困境。MQA 首次把优化重心从「减少浮点计算量(FLOPs)」转向了「减少显存带宽字节搬运(Memory IO)」。PaLM、Falcon、StarCoder 等众多早期知名大模型纷纷采用 MQA 作为标配,它也直接引出了后续更加均衡的分组查询注意力。

Grouped-Query Attention(分组查询注意力 GQA)

MQA 虽然把推理速度拉满了,但在百亿乃至千亿参数的大模型实践中,工程师发现了新问题。直接把几十个键值头一刀砍到只剩一个头,压缩实在太过激进,往往会导致复杂推理任务上的模型质量退化和训练不稳定。

与此同时,业界已经花费巨资预训练好了大量基于标准多头注意力(MHA)的开源大模型。如果为了推理加速就得从零重新预训练一个 MQA 模型,算力成本根本无法承受。大家既想要 MHA 的高智商,又想要 MQA 的快速度,还希望能直接把现成的 MHA 模型低成本改造过来。

这篇论文在多头注意力与多查询注意力之间找到了黄金折中点。它通过将查询头分组并按组共享键值头,在几乎不损失精度的前提下获得了接近 MQA 的推理加速,还给出了已有模型的低成本转换方案。

Google 研究团队提出了「分组查询注意力 」(Grouped-Query Attention,简称 GQA),并配套了一套「增量升级训练」(Uptraining)配方。

在结构设计上,GQA 把全部 \(H\) 个查询头均匀划分为 \(G\) 个小组。每一个小组内部共享一组键(Key)和值(Value)头,整层一共保留 \(G\) 组键值头。你可以把 GQA 看作一个连续可调的旋钮。当组数 \(G=1\) 时,所有查询头共用一组键值,它就退化为最省显存的 MQA。当组数 \(G=H\) 时,每个查询头各自独立,它就还原为标准的 MHA。而当 \(G\) 取中间值(例如 64 个查询头分为 8 组)时,KV 缓存体积直接减少为原来的八分之一,同时保留了多组键值表征空间的多样性。

为了让已有的 MHA 模型无缝转换成 GQA,作者设计了两步升级法。①「均值池化投影 」,把同一组内原本多个独立的键头和值头权重矩阵直接按元素求平均,融合成该组共享的单个键值头。②「轻量继续预训练」,用原预训练仅约 5% 的算力比例对融合后的模型继续训练,让模型快速适应新的分组键值结构。

GQA 成为了当今开源与闭源大模型的绝对工业标准。从 Meta 的 Llama 2(70B)、Llama 3 全系列,到 Mistral、Qwen2 等主流模型,几乎清一色采用了 GQA 架构。它用极简的分组对称设计,彻底终结了推理速度与模型表征能力之间的二选一困境。

Multi-head Latent Attention(多头潜在注意力 MLA)

GQA 和 MQA 虽然通过「砍掉头数」压缩了 KV 缓存,但这种硬性共享不可避免地削弱了每个头的独立表征能力。当上下文长度扩展到 128K 甚至更长时,哪怕是 8 组的 GQA,单个请求的 KV 缓存依然非常庞大,严重限制了线上服务的并发批次大小。

DeepSeek 团队敏锐地意识到,与其粗暴地让多个头共用一模一样的键和值,不如利用线性代数中的「低秩压缩」原理。模型可以先将所有头的键和值联合压缩进一个低维度的潜在空间里存起来,用的时候再各自还原出每个头独有的键和值。

这篇论文提出了多头潜在注意力机制(MLA),通过低秩联合压缩将键和值投影为极小维度的潜在向量进行缓存。推理时它通过矩阵吸收省去解压开销,实现了比 MQA 更小的缓存体积和超越标准 MHA 的模型性能

作者设计了「多头潜在注意力 」(Multi-head Latent Attention,简称 MLA)。在生成键和值时,输入向量不再直接乘上庞大的投影矩阵生成全部头的 \(K\) 和 \(V\)。它先通过一个下投影矩阵,压缩成一个维度极小的「键值潜在向量 」\(c\_t^{KV}\)。在推理阶段,系统根本不需要缓存完整的多头 \(K\) 和 \(V\),只需要把这个紧凑的潜在向量 \(c\_t^{KV}\) 存进显存。

计算注意力时,模型再用上投影矩阵把潜在向量还原成各个头独立的键和值。更巧妙的是,由于上投影是纯线性变换,在推理时可以把键的上投影矩阵直接合并吸收到查询矩阵 \(W^Q\) 里,把值的上投影矩阵合并吸收到输出矩阵 \(W^O\) 里。这样一来,推理时甚至连「把潜在向量解压回完整键值」这一步计算都彻底省掉了。

不过这里藏着一个棘手的数学障碍,现代大模型普遍使用旋转位置编码(RoPE)。如果在键和查询中间夹了一个与位置强相关的旋转矩阵,上投影矩阵就没法跨越它与查询矩阵合并。为了攻克这个难题,作者提出了「解耦旋转位置编码」(Decoupled RoPE)。他们将查询和键拆成两部分。主体部分不加位置编码,专门走低秩潜在压缩和矩阵吸收通道。另外单独开辟极小维度的共享键头和查询分支来承载旋转位置编码,随后把两部分拼接起来算注意力,干净利落地打通了全流程。

MLA 是近年来注意力架构设计中最具原创性的突破之一。它跳出了「物理删减键值头 」的思维定式,改用「低秩潜空间压缩加矩阵吸收 」的代数技巧,第一次做到了「缓存比 MQA 还要小,效果比 MHA 还要好」。这一架构随后成为 DeepSeek-V3 与 DeepSeek-R1 的核心底座,彻底重塑了大模型高效推理的技术标杆。

阶段5:硬件感知与系统级协同设计(2022 至 2025)

前面几条路线都在改写注意力的数学公式或网络结构。但从 2022 年开始,以斯坦福大学、加州大学伯克利分校和 DeepSeek 为代表的研究团队意识到,纸面上的浮点运算量(FLOPs)往往具有欺骗性。算法跑得快不快,取决于它是否顺应现代 GPU 的硬件存储层次与操作系统调度逻辑。这一阶段开启了「算法与硬件深度协同」的新纪元。

FlashAttention 与 FlashAttention-2(IO 感知精确注意力)

为什么很多号称把复杂度降到 \(O(N)\) 的近似注意力算法,在实际显卡上跑起来反而比标准注意力还要慢?Tri Dao 团队一针见血地指出,以往算法只盯着算术运算次数,却忽视了 GPU 的存储金字塔。

现代 GPU(如 A100)有两个核心存储区。一个是高带宽显存(HBM),容量高达 40GB 到 80GB,但读写带宽只有每秒 1.5 到 2.0 TB。另一个是片上静态随机存储器(SRAM),读写带宽高达每秒 19 TB,比显存快十倍以上,但容量只有区区 20MB。标准注意力在计算时,必须先把 \(Q K^T\) 算出的 \(N \\times N\) 矩阵写回慢速的 HBM,再读出来算 Softmax 写回 HBM,再读出来乘 \(V\)。绝大部分时间全都耗在了慢速显存的来回倒腾上。

这篇论文在不做任何数学近似的前提下,通过片上缓存分块计算、在线 Softmax 归一化与反向重计算,彻底消除了中间矩阵对高带宽显存的读写。它将精确注意力的显存占用降为线性,并实现了数倍的真实硬件加速。

作者提出了具有硬件 IO 感知能力的「FlashAttention」,一分一毫都不改变标准注意力的输出结果,只改变数据在芯片内部的流动方式。核心武器是「分块计算」(Tiling)与「重计算」(Recomputation)。

算法将输入的 \(Q\)、\(K\)、\(V\) 切分成一个个小块,分批加载进极速的片上 SRAM 中。在 SRAM 内部一次性算完点积、Softmax 和乘 \(V\) 的算子融合全过程,只把最终的输出写回慢速显存 HBM。中间那个巨大的 \(N \\times N\) 注意力矩阵自始至终都不落地。

这里最大的拦路虎是 Softmax 需要对整行求和才能归一化,分块加载时看不全整行怎么办?作者引入了「在线归一化」(Online Softmax)技巧。每处理完一个小块,就动态更新当前记录的行最大值和指数累加和,并利用指数缩放因子实时修正之前算出的部分结果。扫描完所有小块后,结果与一次性算完整行完全一致。在反向传播时,模型也不保存 \(N \\times N\) 的中间矩阵,而是只存每一行的对数累加和统计量,反向时在 SRAM 里快速重算注意力矩阵。虽然多做了一点点浮点乘法,但由于省去了海量显存读取,反向速度反而大幅飙升。

一年后,作者又发表了 FlashAttention-2,进一步削减了非矩阵乘法(Non-matmul)算子的比例。它调换了内外层循环顺序,并在序列维度和线程束(Warp)之间重新划分了并行任务,减少了共享内存通信开销。

FlashAttention 是大模型基础设施史上当之无愧的分水岭。它证明了通过深入到底层硬件存储层次做算子融合,完全可以在零精度损失的前提下打破「显存平方墙」。如今,FlashAttention 已被直接内置进 PyTorch 官方底层、HuggingFace 与 Megatron-LM 等所有主流框架,成为全世界训练和部署大语言模型的默认底层引擎。

PagedAttention(虚拟内存分页注意力与 vLLM)

FlashAttention 解决了单次计算时的中间矩阵显存问题,但在大模型线上高并发服务(LLM Serving)中,另一个显存黑洞浮出水面。当成百上千个用户同时向大模型提问时,系统无法预知每个请求最终会生成五十个词还是两千个词。

传统的深度学习框架要求张量必须在物理显存里连续存放。因此,服务系统只能按每个请求可能达到的「最大长度」(比如 2048)提前划拨一大块连续显存。这带来了三种惨重的显存浪费。预留给未来生成词但还没用上的空间属于「预留浪费」。请求提前结束导致剩余空间永远用不上的属于「内部碎片」。不同长度的连续块之间夹杂着无法塞进新请求的空隙属于「外部碎片」。实测表明,现有推理系统中居然有 60% 到 80% 的 KV 缓存显存被白白浪费掉了,真正存有效数据的不到四成。

这篇论文借鉴操作系统虚拟内存与分页管理思想,允许将原本必须连续存储的 KV 缓存切分为固定大小的块并存储在非连续物理显存中。它将显存浪费率压低至 4% 以内,使大模型线上服务吞吐量提升二至四倍。

加州大学伯克利分校团队将操作系统的经典智慧搬进了注意力算子,提出了「PagedAttention」并开源了推理引擎 vLLM。

核心思路是把每个请求的 KV 缓存切分成固定大小的「KV 块」(Block)。每个块容纳固定数量(如 16 个)词元的键和值,就像操作系统里的「虚拟内存页 」。在逻辑视图上,一个请求的 KV 缓存依然是连续增长的块序列。但在 GPU 物理显存上,这些块可以散落在任意非连续的位置,完全不需要提前预订一大片连续空间。系统维护一张轻量级的「块表」(Block Table),负责记录逻辑块到物理块的映射关系。请求每生成几个新词,块管理器就按需动态分配一个新的物理块,只有排在末尾的那个未填满块会产生微不足道的尾部浪费。

在此基础上,定制的 PagedAttention CUDA 算子在计算注意力时,会根据块表动态去非连续的物理地址抓取对应的 KV 块做分块矩阵乘法。不仅如此,分页机制还顺手解锁了高效的「写时复制」(Copy-on-Write)内存共享。当多个请求共享同一段长系统提示词(Prompt)或使用并行采样、束搜索时,它们只需要把逻辑块表指向同一份物理块即可。只有需要修改时才复制新块,彻底省去了重复存储。

PagedAttention 从系统工程角度重新定义了注意力机制的存储形态。它没有改动模型哪怕一个权重参数,却通过解耦逻辑视图与物理显存,把大模型推理的硬件利用率推向了新高度。凭借这篇论文诞生的 vLLM 迅速成为全球最受欢迎的大模型开源推理引擎,分页 KV 缓存也成为了 TensorRT-LLM、SGLang 等所有现代推理框架的标配基石。

Native Sparse Attention(硬件对齐原生可训练稀疏注意力 NSA)

尽管早在 2020 年就有 Longformer 等稀疏注意力研究,但为什么后来的主流大模型在预训练时依然咬牙坚持用昂贵的全量注意力(Full Attention)?DeepSeek 团队剖析了传统稀疏注意力的两大死穴。

① 硬件不友好,理论快但实际慢。许多方法只盯着在推理阶段按词元(Token-level)离散挑选重要词。这种零散的随机内存访问严重违背了 GPU Tensor Core 连续分块读取的脾气。而且在结合 GQA 架构时,同一个查询组里的不同头往往要读取不同的离散键值,导致依然要读回整片显存,根本拿不到真实加速。

② 无法原生端到端训练。许多稀疏方法依赖不可导的离散剪枝或仅在推理阶段强行丢弃缓存。模型在预训练时从未适应过稀疏分布,一到长文推理就会出现明显的能力退化。

这篇论文提出了面向现代 GPU 硬件对齐且支持端到端原生预训练的动态分层稀疏注意力(NSA)。通过「粗粒度块压缩、细粒度块精选与局部滑动窗口」三路并行机制,它在 64K 长序列上实现了全生命周期数倍加速,同时在模型精度上追平甚至超越了全量注意力。

作者提出了「原生稀疏注意力」(Native Sparse Attention,简称 NSA),把算法设计与硬件算力密度深度绑定。NSA 为每个查询向量同时构建三条互补的上下文信息通路,并通过可学习的门控权重将三路注意力输出动态融合。

  • ① 粗粒度词元压缩(Token Compression)。模型把连续的键值序列按固定长度切分成块,通过含位置编码的多层感知机将每个块浓缩成一个代表性键值,提供低成本的全文宏观鸟瞰图。

  • ② 细粒度词元选择(Token Selection)。为了不漏掉关键细节,模型直接复用第一路算出的块级重要性评分,挑选出得分最高的若干个完整分块。注意,这里是「整块整块地挑」,并且强制让 GQA 同一个组内的所有查询头共享同一份精选块索引。这种「块级连续挑选加组内共享」的设计,完美契合了 FlashAttention 式的 SRAM 连续分块加载与 Tensor Core 矩阵运算,彻底消除了随机访存碎片。

  • ③ 局部滑动窗口(Sliding Window),专门负责盯紧当前词紧邻的局部上下文,防止局部高频信号塞满精选通道。由于三条通路的算子全部可导,模型可以直接从零开始进行端到端原生稀疏预训练。

NSA 标志着稀疏注意力终于走出了「近似折中方案」的阴影,正式迈入「原生训练、全周期加速且精度无损」的成熟阶段。它深刻揭示了下一代注意力演进的终极规律。只有同时尊重长文本的信息层级规律(全局概览加局部精选)与现代芯片的物理访存规律(连续分块与算术密度平衡),才能真正解锁百万级上下文与超长思维链推理的潜力。

演进总结与未来展望

回顾 Attention 机制这十余年的演进历程,我们可以清晰地看到技术焦点经历了四次深刻跃迁。从 Bahdanau 和 Luong 时代解决「信息装不下 」的软对齐寻址,到 Vaswani 时代解决「训练无法并行 」的全自注意力矩阵化,再到 Longformer 与 Linear Attention 探索「序列平方复杂度降维 」,直至近年来围绕大模型落地展开的「KV 缓存极限压缩(MQA、GQA、MLA) 」与「软硬件系统协同(FlashAttention、PagedAttention、NSA)」。

站在当前的时间节点往后看,注意力机制的演化正呈现出三大合流趋势。① 低秩潜在压缩与硬件原生稀疏的结合 ,以 MLA 和 NSA 为代表的技术路线正在进一步融合,让模型既能把单个词的缓存压到极小,又能在超长序列中只挑选关键块参与计算。② 混合架构的普及 ,将擅长精确检索的注意力层与擅长常数显存流式处理的线性注意力或状态空间模型按比例交替堆叠,兼顾大海捞针精度与极致吞吐。③ 算法与芯片微架构的原生绑定,未来的注意力机制将从诞生第一天起就围绕片上缓存带宽、张量核心指令集与分布式通信拓扑协同设计,持续推高智能计算的效率上限。

相关推荐
zx_741484815 小时前
BERT自然语言处理框架深度解析:从Transformer到双向预训练模型
自然语言处理·bert·transformer
布吉岛的石头1 天前
Java 程序员第 50 阶段13:对齐税 alignment tax 是什么
人工智能·深度学习·transformer
米饭是菜qy1 天前
深度学习Transformer详细介绍
人工智能·深度学习·transformer
Eloudy1 天前
《Attention Is All You Need》的 Encoder–Decoder Transformer 一次完整训练
llm·transformer
DevOps老兵2 天前
Transformer拆解01:Transformer为什么出现,RNN到底卡在哪里
rnn·深度学习·大模型·transformer·注意力机制
老唐7774 天前
大模型面试高频题:Transformer、MoE、RLHF、DPO 核心原理拆解
人工智能·深度学习·ai·大模型·transformer·智能体
JAI科研4 天前
把一个 Qwen3.8的聊天模型改造成 JEV 决策模型
人工智能·语言模型·自然语言处理·chatgpt·transformer·vllm·sglang
小易老师AI实战4 天前
Transformer掩码注意力机制原理深度解析
transformer·大模型基础·自回归生成·掩码注意力
wu_solo6 天前
6.1 Transformer架构概述
深度学习·架构·transformer