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，然后再写回去做加权求和。对于 $n=4096$、$d_k=64$ 的情况，单个注意力头的 $S$ 矩阵就是 $4096 \times 4096 \times 4$ 字节（FP32），约64MB。如果是32个头，那就是2GB——这还只是一个样本、一层的数据。
2.2 为什么说注意力是“访存密集型”而非“计算密集型”
这里有个反直觉的点：GPU的浮点计算能力其实很强，A100的FP16算力是312 TFLOPS，但显存带宽只有约2TB/s。注意力计算中，每做一次浮点运算，需要搬运的数据量远高于矩阵乘法中的理想比例。
我实测过一组数据：在A100上跑 $n=2048$、$d=768$、12个头的注意力层，纯计算时间只占总耗时的约35%，剩下65%都花在显存读写和kernel启动开销上。这意味着即使你把计算单元升级到算力翻倍，整体性能提升也不会超过35%。
这就是为什么FlashAttention这类“IO感知”的算法能带来数倍加速——它优化的不是计算量，而是显存访问模式。
2.3 多头注意力：表达能力与计算效率的第一次权衡
原始Transformer用的是多头注意力（Multi-Head Attention），把 $d$ 维的Q、K、V拆成 $h$ 个头，每个头独立做注意力，最后拼接。这个设计的动机是让模型在不同子空间里学习不同的注意力模式。
但多头带来一个工程上的问题：每个头的 $S$ 矩阵都要单独计算和存储，显存占用是单头的 $h$ 倍。而且头与头之间的计算是独立的，kernel启动次数也翻倍。
后来出现了一些优化方案，比如MQA（Multi-Query Attention）和GQA（Grouped-Query Attention）。MQA让所有头共享同一组K、V，只保留多组Q；GQA则是折中方案，把头分成若干组，每组共享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个头，每头维度128，FP16存储，序列长度2048时，单层KV Cache约32MB，32层就是1GB。如果batch size是16，那就是16GB——直接把一张24GB的卡吃满。
这就是MQA和GQA在推理场景下价值凸显的地方。MQA把所有头的K、V合并成一份，KV Cache直接缩小到原来的 $1/h$；GQA按组共享，压缩比例取决于分组数。LLaMA-2 70B用的就是GQA，8组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用了在线softmax（online softmax）的技巧：维护一个运行中的最大值 $m$ 和运行中的指数和 $l$，每处理一个新块时更新这两个值，并相应地缩放之前累积的结果。这样就不需要一次性看到所有分数，也能得到正确的softmax结果。
数学上可以证明，在线softmax的结果跟标准softmax完全一致，不是近似。这也是FlashAttention跟一些近似注意力方法（如Linformer、Performer）的本质区别——它不改变数学结果，只是重排了计算顺序。
4.3 实测加速比与适用边界
官方论文报告的数据：在GPT-2上，FlashAttention比标准注意力快2-4倍，显存占用降低5-20倍。我自己在A100上实测的结果：
| 序列长度 | 标准注意力耗时 | FlashAttention耗时 | 加速比 |
|---|---|---|---|
| 512 | 1.2ms | 0.9ms | 1.3x |
| 1024 | 4.1ms | 2.2ms | 1.9x |
| 2048 | 15.3ms | 6.1ms | 2.5x |
| 4096 | 58.7ms | 18.4ms | 3.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, 2i+1)} = \cos(pos / 10000^{2i/d})$$
这个设计的好处是不需要学习参数，而且理论上可以外推到训练时没见过的长度。但实际中外推效果并不好，因为高频部分的周期太短，稍微长一点就重复了。
5.2 可学习位置编码与相对位置编码
后来BERT用了可学习的位置编码，就是给每个位置初始化一个向量，跟着模型一起训练。这种方式在训练长度内效果不错，但完全没法外推——训练时最大长度512，推理时来了个600的序列，第513到600的位置编码根本不存在。
相对位置编码换了个思路：不编码绝对位置，而是编码token之间的相对距离。这样理论上可以处理任意长度，因为相对距离的模式在训练和推理时是一致的。T5用的就是相对位置编码的简化版本。
但相对位置编码有个工程问题：它需要在注意力分数计算时额外加一个偏置项，这个偏置项的计算和存储会带来额外开销，而且实现起来比绝对位置编码复杂。
5.3 RoPE的旋转矩阵构造与长度外推
RoPE（Rotary 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 Transformer
Transformer从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模块的融合
另一条演进路线是通道注意力。SE（Squeeze-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 scaling；BF16动态范围跟FP32一样，不需要loss scaling，但精度略低。A100及以上建议用BF16，V100及以下只能用FP16。
FP8是更新的选项，H100支持。但FP8对量化误差很敏感，注意力分数在softmax之前如果精度不够，会导致注意力分布偏移。目前FP8在注意力计算上的应用还在探索阶段，生产环境慎用。
7.3 长序列场景的显存预算方法
做长序列推理时，显存预算要提前算清楚。以LLaMA-7B为例，序列长度8192，batch size 1：
- 模型权重：7B参数，FP16存储，约14GB
- KV Cache：32层 × 2 × 8192 × 32头 × 128维 × 2字节 ≈ 4GB
- 激活值：取决于实现，FlashAttention下约1-2GB
- 其他开销：约1GB
总计约20-21GB，一张24GB的卡勉强够用。如果序列长度翻倍到16384，KV Cache翻倍到8GB，总需求就超过24GB了。这时候要么用GQA压缩KV Cache，要么上多卡。
一个容易忽略的点：KV Cache的显存占用是跟batch size线性相关的。batch size从1加到8，KV Cache从4GB变成32GB，直接爆显存。所以长序列场景下，batch size往往只能设得很小。
7.4 常见性能问题的排查清单
最后整理一份我在实际工作中总结的排查清单，遇到Transformer推理性能问题时可以按这个顺序过一遍：
- 先看序列长度分布。如果P99序列长度远大于平均长度，考虑做长度分桶，避免长序列拖累短序列。
- 检查KV Cache管理方式。是否预分配了过大空间？是否用了PagedAttention？
- 确认FlashAttention是否生效。有些框架需要显式开启，或者对某些注意力变体不支持。
- 看GPU利用率。如果计算利用率低但显存带宽利用率高，说明是访存瓶颈，FlashAttention会有帮助。
- 检查是否有不必要的同步操作。比如每层都做一次all-reduce，会严重拖慢速度。
- 最后看batch size是否合理。太小浪费算力，太大爆显存，需要做trade-off。
这套流程帮我定位过好几次线上问题，大部分性能瓶颈都能在前三步找到线索。
