1. 从一次线上推理延迟抖动说起去年帮一个团队排查线上推理服务的延迟抖动问题现象很典型P99延迟在流量高峰期会突然从80ms飙到400ms以上但GPU利用率却只有60%出头。一开始怀疑是批处理调度的问题折腾了两天没找到根因后来把profile打开一看问题出在注意力计算上——序列长度超过512之后注意力矩阵的显存占用和访存开销呈平方级增长直接把显存带宽打满了。这件事让我重新梳理了一遍Transformer架构里那些核心算法的演进脉络。很多人学Transformer停留在“知道有自注意力、有多头、有位置编码”的层面但真正到了工程落地你会发现每一个模块都有好几代技术方案在迭代选错了版本性能差距可能是数量级的。这篇内容我打算按“算法为什么这样演进”的思路把Transformer从原始版本到当前主流工程实践中的几个关键算法节点串一遍。涉及注意力机制的变体、KV Cache的引入逻辑、FlashAttention的工程价值、位置编码的几代方案对比以及这些演进背后的核心驱动力——计算效率和表达能力的博弈。适合已经了解Transformer基本结构、但在工程实践中遇到性能瓶颈或想深入理解设计动机的读者。2. 原始注意力机制的计算瓶颈到底卡在哪2.1 自注意力的数学本质与计算量拆解先把自注意力的计算过程摊开看。给定输入序列 $X \in \mathbb{R}^{n \times d}$其中 $n$ 是序列长度$d$ 是隐藏维度自注意力的计算分三步第一步线性投影得到Q、K、V$$Q XW_Q, \quad K XW_K, \quad V XW_V$$其中 $W_Q, W_K, W_V \in \mathbb{R}^{d \times d_k}$$d_k$ 是每个头的维度。第二步计算注意力分数矩阵$$S \frac{QK^T}{\sqrt{d_k}}$$这里 $S \in \mathbb{R}^{n \times n}$注意这个矩阵的尺寸只跟序列长度有关跟隐藏维度无关。第三步softmax归一化后加权求和$$\text{Attention}(Q,K,V) \text{softmax}(S) \cdot V$$最终输出 $\in \mathbb{R}^{n \times d_k}$。从计算量来看$QK^T$ 的矩阵乘法是 $O(n^2 \cdot d_k)$softmax是 $O(n^2)$加权求和也是 $O(n^2 \cdot d_k)$。当 $n$ 从512涨到2048时$n^2$ 从26万涨到420万翻了16倍。这就是为什么长序列场景下注意力计算会成为瓶颈。但计算量还不是最要命的真正卡脖子的是显存访问。$S$ 矩阵需要被完整地写入显存再读出来做softmax然后再写回去做加权求和。对于 $n4096$、$d_k64$ 的情况单个注意力头的 $S$ 矩阵就是 $4096 \times 4096 \times 4$ 字节FP32约64MB。如果是32个头那就是2GB——这还只是一个样本、一层的数据。2.2 为什么说注意力是“访存密集型”而非“计算密集型”这里有个反直觉的点GPU的浮点计算能力其实很强A100的FP16算力是312 TFLOPS但显存带宽只有约2TB/s。注意力计算中每做一次浮点运算需要搬运的数据量远高于矩阵乘法中的理想比例。我实测过一组数据在A100上跑 $n2048$、$d768$、12个头的注意力层纯计算时间只占总耗时的约35%剩下65%都花在显存读写和kernel启动开销上。这意味着即使你把计算单元升级到算力翻倍整体性能提升也不会超过35%。这就是为什么FlashAttention这类“IO感知”的算法能带来数倍加速——它优化的不是计算量而是显存访问模式。2.3 多头注意力表达能力与计算效率的第一次权衡原始Transformer用的是多头注意力Multi-Head Attention把 $d$ 维的Q、K、V拆成 $h$ 个头每个头独立做注意力最后拼接。这个设计的动机是让模型在不同子空间里学习不同的注意力模式。但多头带来一个工程上的问题每个头的 $S$ 矩阵都要单独计算和存储显存占用是单头的 $h$ 倍。而且头与头之间的计算是独立的kernel启动次数也翻倍。后来出现了一些优化方案比如MQAMulti-Query Attention和GQAGrouped-Query Attention。MQA让所有头共享同一组K、V只保留多组QGQA则是折中方案把头分成若干组每组共享K、V。这两个方案在推理阶段对KV Cache的压缩效果非常明显后面讲KV Cache时会展开。3. KV Cache自回归推理中绕不开的工程优化3.1 自回归解码的重复计算问题训练阶段整个序列是一次性喂进去的注意力可以并行计算。但推理阶段尤其是自回归生成比如文本生成是一个token一个token往外吐的。每生成一个新token都需要拿它跟前面所有token做注意力计算。如果每次都重新计算前面所有token的K、V那计算量会随序列长度线性增长而且大量计算是重复的。比如生成第100个token时前99个token的K、V跟生成第99个token时算出来的完全一样。KV Cache的思路很直接把已经算过的K、V缓存下来生成新token时只需要计算当前token的Q、K、V然后拿当前Q去跟缓存里的K做注意力再把当前V追加到缓存里。3.2 KV Cache的显存开销与优化方向KV Cache省了计算但吃显存。缓存大小是$$\text{Cache Size} 2 \times n \times d \times h \times \text{bytes per element}$$其中2是K和V两份$n$是当前序列长度$d$是每头维度$h$是头数。以LLaMA-7B为例32个头每头维度128FP16存储序列长度2048时单层KV Cache约32MB32层就是1GB。如果batch size是16那就是16GB——直接把一张24GB的卡吃满。这就是MQA和GQA在推理场景下价值凸显的地方。MQA把所有头的K、V合并成一份KV Cache直接缩小到原来的 $1/h$GQA按组共享压缩比例取决于分组数。LLaMA-2 70B用的就是GQA8组KV对应64个Q头KV Cache压缩了8倍。3.3 PagedAttention与显存碎片治理KV Cache还有一个容易被忽略的问题显存碎片。传统实现里每个请求的KV Cache需要预分配一块连续显存按最大可能长度来分配。但实际生成的序列长度是动态的预分配会导致大量浪费。vLLM提出的PagedAttention借鉴了操作系统虚拟内存分页的思路把KV Cache切成固定大小的block按需分配不要求连续。这个方案把显存利用率从原来的60%左右提升到了90%以上吞吐量提升非常明显。我在实际部署中对比过同样一张A100 80GB跑LLaMA-7B推理用传统KV Cache管理方式batch size最多到24换成PagedAttention之后batch size能到48吞吐量翻了一倍。4. FlashAttention不改变数学结果的前提下重写计算顺序4.1 标准注意力的显存访问模式问题前面提到标准注意力实现需要把 $S QK^T$ 完整写入显存再做softmax再读出来做加权求和。这个过程中$S$ 矩阵的读写是最大的开销来源。具体来说标准实现需要写 $S$ 到显存$n^2 \times d_k$ 次写读 $S$ 做softmax$n^2 \times d_k$ 次读写softmax结果$n^2 \times d_k$ 次写读softmax结果做加权求和$n^2 \times d_k$ 次读总共4次 $O(n^2)$ 级别的显存访问。而GPU的SRAM共享内存/寄存器容量有限A100的SRAM只有约20MB根本放不下完整的 $S$ 矩阵。4.2 Tiling与在线softmax的核心思路FlashAttention的核心洞察是既然SRAM放不下完整的 $S$那就把它切成小块tile每次只处理一块在SRAM里完成计算避免中间结果写回显存。但softmax有个麻烦它需要全局的归一化因子。标准softmax是$$\text{softmax}(s_i) \frac{e^{s_i}}{\sum_j e^{s_j}}$$分母需要所有 $s_j$ 的和。如果分块计算每块只能看到部分 $s_j$没法直接算。FlashAttention用了在线softmaxonline softmax的技巧维护一个运行中的最大值 $m$ 和运行中的指数和 $l$每处理一个新块时更新这两个值并相应地缩放之前累积的结果。这样就不需要一次性看到所有分数也能得到正确的softmax结果。数学上可以证明在线softmax的结果跟标准softmax完全一致不是近似。这也是FlashAttention跟一些近似注意力方法如Linformer、Performer的本质区别——它不改变数学结果只是重排了计算顺序。4.3 实测加速比与适用边界官方论文报告的数据在GPT-2上FlashAttention比标准注意力快2-4倍显存占用降低5-20倍。我自己在A100上实测的结果序列长度标准注意力耗时FlashAttention耗时加速比5121.2ms0.9ms1.3x10244.1ms2.2ms1.9x204815.3ms6.1ms2.5x409658.7ms18.4ms3.2x可以看到序列越长加速比越明显。这是因为短序列下 $S$ 矩阵本身就不大访存开销占比低FlashAttention的tiling反而引入了额外开销。注意FlashAttention对硬件有要求需要GPU支持特定的SRAM容量和指令集。在一些较老的卡上或者非NVIDIA平台上可能需要用适配版本性能收益也会打折扣。4.4 FlashAttention-2与v3的迭代方向FlashAttention-2主要优化了并行度和work partitioning。第一版在序列维度上的并行度不够第二版把序列也切分到不同线程块上提高了GPU利用率。官方数据是比第一版再快约2倍。FlashAttention-3则是针对Hopper架构H100做的优化利用了新的异步指令和FP8支持。不过这些优化跟硬件绑定较深通用性不如前两版。5. 位置编码从正弦函数到旋转位置嵌入5.1 为什么Transformer需要位置编码自注意力的计算本质是加权求和它对输入序列的顺序是不敏感的。你把序列打乱注意力输出只是跟着打乱不会感知到顺序变化。但语言是有顺序的“猫追老鼠”和“老鼠追猫”意思完全相反。所以需要在输入里注入位置信息。原始Transformer用的是正弦位置编码$$PE_{(pos, 2i)} \sin(pos / 10000^{2i/d})$$ $$PE_{(pos, 2i1)} \cos(pos / 10000^{2i/d})$$这个设计的好处是不需要学习参数而且理论上可以外推到训练时没见过的长度。但实际中外推效果并不好因为高频部分的周期太短稍微长一点就重复了。5.2 可学习位置编码与相对位置编码后来BERT用了可学习的位置编码就是给每个位置初始化一个向量跟着模型一起训练。这种方式在训练长度内效果不错但完全没法外推——训练时最大长度512推理时来了个600的序列第513到600的位置编码根本不存在。相对位置编码换了个思路不编码绝对位置而是编码token之间的相对距离。这样理论上可以处理任意长度因为相对距离的模式在训练和推理时是一致的。T5用的就是相对位置编码的简化版本。但相对位置编码有个工程问题它需要在注意力分数计算时额外加一个偏置项这个偏置项的计算和存储会带来额外开销而且实现起来比绝对位置编码复杂。5.3 RoPE的旋转矩阵构造与长度外推RoPERotary Position Embedding是目前主流大模型LLaMA、Qwen、ChatGLM等普遍采用的位置编码方案。它的核心思路是通过旋转矩阵把位置信息编码进Q、K向量里。具体来说对于位置 $m$ 的Q向量和位置 $n$ 的K向量RoPE让它们的点积只依赖于相对位置 $m-n$。这是通过把Q、K向量按维度分成若干对每对看作一个二维向量然后根据位置旋转一个角度来实现的。旋转角度 $\theta_i m \cdot \omega_i$其中 $\omega_i 10000^{-2i/d}$。这个设计跟正弦位置编码有相似之处但关键区别在于RoPE是作用在Q、K上的而不是加在输入上。这使得注意力分数天然携带了相对位置信息。RoPE的长度外推能力比正弦编码好但也不是无限的。当推理长度远超训练长度时高频维度的旋转角度会超出训练时见过的范围导致性能下降。后来出现了一些改进方案比如NTK-aware插值和YaRN通过调整频率基数来改善外推效果。5.4 位置编码方案选型的实际考量在实际项目中选位置编码我一般按这几个维度来评估方案训练长度内效果外推能力实现复杂度额外计算开销正弦编码中等差低无可学习编码好无低无相对位置编码好中等中有RoPE好较好中低如果是固定长度任务比如分类可学习编码就够了。如果需要处理变长序列或者有外推需求RoPE是目前综合最优的选择。相对位置编码在T5这类模型上表现不错但工程实现比RoPE麻烦。6. 注意力变体从Swin到通道注意力的场景分化6.1 视觉领域的窗口注意力与Swin TransformerTransformer从NLP扩展到视觉领域时遇到一个直接问题图像的token数量太大了。一张224x224的图如果按16x16的patch切有196个token如果按4x4切那就是3136个token。自注意力的 $O(n^2)$ 复杂度直接爆炸。Swin Transformer的解决方案是窗口注意力把图像划分成不重叠的窗口每个窗口内做自注意力窗口之间通过移位操作来交换信息。这样计算量从 $O(n^2)$ 降到 $O(n \cdot w)$其中 $w$ 是窗口大小。这个思路跟NLP里的滑动窗口注意力Longformer有相似之处但Swin的移位窗口设计更巧妙它保证了窗口之间的信息流通同时保持了计算效率。6.2 通道注意力与SE模块的融合另一条演进路线是通道注意力。SESqueeze-and-Excitation模块最早是给CNN设计的它通过全局池化得到每个通道的权重然后对通道做加权。这个思路被引入Transformer后出现了一些混合架构比如在注意力计算之前先用通道注意力做一次特征重标定。通道注意力的计算量很小因为它是 $O(d)$ 而不是 $O(n^2)$。在一些对延迟敏感的场景下用通道注意力替代部分自注意力层可以在精度损失很小的情况下显著降低计算量。6.3 时序注意力在序列建模中的特殊处理时序数据比如传感器读数、金融行情的注意力计算有一些特殊考量。时序数据往往有很强的局部相关性而且可能存在周期性模式。标准的自注意力对这种结构没有先验偏好需要大量数据才能学到。一些时序Transformer如Informer、Autoformer引入了稀疏注意力机制只计算每个位置跟部分关键位置的注意力而不是全连接。Informer用的ProbSparse注意力只保留Top-k个最重要的query把复杂度从 $O(n^2)$ 降到 $O(n \log n)$。不过这些稀疏注意力方案在实际落地中要谨慎评估。我见过一些团队直接套用Informer的稀疏注意力结果在自己的数据上精度掉得厉害。原因是稀疏化的假设注意力分布是长尾的在某些数据上不成立强行稀疏化会丢掉关键信息。7. 工程落地中的几个关键决策点7.1 训练与推理的算法选型差异训练阶段和推理阶段的最优算法选择往往不一样。训练时batch size大、序列长计算密集FlashAttention的收益最明显推理时batch size小、序列短KV Cache的管理策略比注意力计算本身更重要。我的一般建议是训练阶段默认开FlashAttention除非你的序列长度很短小于256推理阶段优先优化KV Cache用GQA或MQA压缩缓存再用PagedAttention管理显存。7.2 精度与速度的权衡策略FP16和BF16的选择也是个实际问题。FP16动态范围小容易溢出训练时通常需要loss scalingBF16动态范围跟FP32一样不需要loss scaling但精度略低。A100及以上建议用BF16V100及以下只能用FP16。FP8是更新的选项H100支持。但FP8对量化误差很敏感注意力分数在softmax之前如果精度不够会导致注意力分布偏移。目前FP8在注意力计算上的应用还在探索阶段生产环境慎用。7.3 长序列场景的显存预算方法做长序列推理时显存预算要提前算清楚。以LLaMA-7B为例序列长度8192batch size 1模型权重7B参数FP16存储约14GBKV Cache32层 × 2 × 8192 × 32头 × 128维 × 2字节 ≈ 4GB激活值取决于实现FlashAttention下约1-2GB其他开销约1GB总计约20-21GB一张24GB的卡勉强够用。如果序列长度翻倍到16384KV Cache翻倍到8GB总需求就超过24GB了。这时候要么用GQA压缩KV Cache要么上多卡。一个容易忽略的点KV Cache的显存占用是跟batch size线性相关的。batch size从1加到8KV Cache从4GB变成32GB直接爆显存。所以长序列场景下batch size往往只能设得很小。7.4 常见性能问题的排查清单最后整理一份我在实际工作中总结的排查清单遇到Transformer推理性能问题时可以按这个顺序过一遍先看序列长度分布。如果P99序列长度远大于平均长度考虑做长度分桶避免长序列拖累短序列。检查KV Cache管理方式。是否预分配了过大空间是否用了PagedAttention确认FlashAttention是否生效。有些框架需要显式开启或者对某些注意力变体不支持。看GPU利用率。如果计算利用率低但显存带宽利用率高说明是访存瓶颈FlashAttention会有帮助。检查是否有不必要的同步操作。比如每层都做一次all-reduce会严重拖慢速度。最后看batch size是否合理。太小浪费算力太大爆显存需要做trade-off。这套流程帮我定位过好几次线上问题大部分性能瓶颈都能在前三步找到线索。