2. 自注意力机制算子：QKV投影、注意力分数计算、Softmax、加权求和，以及FlashAttention优化

自注意力机制，说白了就是Transformer的心脏。我当年第一次接触这个结构时，心里就一个想法：这玩意儿怎么算得动？后来真正上手做算子优化，才发现这里面的门道比想象中深得多。

今天咱们就把自注意力机制拆开揉碎了讲。从最基础的QKV投影开始，一步步走到FlashAttention这个革命性的优化方案。嗯，每个环节我都会结合自己的踩坑经历，希望能帮你少走弯路。

2.1 QKV投影：从输入到查询、键、值

先看最基础的一步。输入序列 X 的形状是 [B, S, D]，其中 B 是batch size，S 是序列长度，D 是隐藏维度。我们要把它映射到三个不同的空间：查询（Query）、键（Key）、值（Value）。

数学上很简单：

Q = X @ W_Q      // [B, S, D] @ [D, D_k] → [B, S, D_k]
K = X @ W_K      // [B, S, D] @ [D, D_k] → [B, S, D_k]
V = X @ W_V      // [B, S, D] @ [D, D_v] → [B, S, D_v]

这里 D_k 和 D_v 通常等于 D / h，h 是注意力头的数量。我习惯把这三个投影矩阵合在一起做一次大矩阵乘法，然后split。为什么？因为访存效率更高。

实战技巧： 在GPU上，把三个投影矩阵拼成 [D, 3 * D_k] 的大矩阵，一次matmul搞定。我在优化BERT推理时，这个改动让整体吞吐提升了约15%。原因很简单——减少了kernel launch的开销。

2.2 注意力分数计算：Q和K的点积

得到Q、K、V之后，下一步是计算注意力分数：

S = Q @ K^T      // [B, h, S, D_k] @ [B, h, D_k, S] → [B, h, S, S]

这里有个细节：实际计算时我们会除以 sqrt(D_k)，也就是 S = (Q @ K^T) / sqrt(D_k)。为什么要除？我刚开始做这个时也不理解，直到有一次训练一个深层模型，发现梯度爆炸了。查了半天，原来是点积的方差随着维度增长而变大，softmax之后几乎变成了one-hot，梯度根本传不过去。

核心要点： 缩放因子 1/sqrt(D_k) 是为了保持方差稳定。假设Q和K的每个元素都是独立同分布的标准正态分布，那么点积的方差就是 D_k。除以 sqrt(D_k) 后，方差变回1。

这个算子的计算量有多大？O(S^2 * D_k)。当S=4096时，单头就要算1600万个元素。这还只是一个头，一个batch。你想想看，这得多大的显存和算力。

2.3 Softmax：沿着序列维度归一化

注意力分数算出来后，要对每个query对应的所有key做softmax：

A = softmax(S, dim=-1)   // [B, h, S, S]

Softmax的经典实现分三步：

  1. 找最大值：m = max(s_i)
  2. 计算指数和：d = sum(exp(s_i - m))
  3. 归一化：a_i = exp(s_i - m) / d

为什么要减最大值？为了防止数值溢出。exp(1000) 直接爆float16，但 exp(1000 - 1000) = exp(0) = 1，稳稳的。我曾经在训练一个16位精度的模型时，忘记做这个数值稳定处理，结果loss直接变成NaN。排查了整整两天，最后发现是softmax的锅。

注意： 在混合精度训练中，softmax的中间结果很容易溢出。我建议在计算softmax时临时提升精度，或者使用专门为低精度优化的softmax实现。NVIDIA的cuDNN里就有针对性的优化。

2.4 加权求和：用注意力权重聚合Value

最后一步，用softmax输出的注意力权重对V做加权求和：

O = A @ V          // [B, h, S, S] @ [B, h, S, D_v] → [B, h, S, D_v]

这一步就是矩阵乘法，没什么特别的。但要注意，A 矩阵是 [S, S] 的，当S很大时，这个矩阵乘法非常昂贵。而且 A 本身就需要 O(S^2) 的显存来存储。

整个自注意力机制的计算流程，我画了张图帮你理解：

自注意力机制计算流程 输入 X [B,S,D] Q [B,S,D_k] K [B,S,D_k] V [B,S,D_v] S = Q@K^T/√D_k Softmax O = A@V MatMul 逐行 MatMul

2.5 标准实现的痛点：显存和带宽瓶颈

上面这套流程，每个步骤都产生中间结果。我算一笔账给你看：

中间张量 形状 显存占用 (B=1, h=16, S=4096, D=1024)
Q [B, h, S, D/h] 1 × 16 × 4096 × 64 × 2 = 8 MB
K [B, h, S, D/h] 8 MB
V [B, h, S, D/h] 8 MB
S [B, h, S, S] 1 × 16 × 4096 × 4096 × 2 = 512 MB
A [B, h, S, S] 512 MB
O [B, h, S, D/h] 8 MB

看到了吗？光 S 和 A 两个矩阵就占了1GB显存。这还只是一个transformer层。我做过一个12层的模型，光注意力部分就吃掉12GB。更别提还有FFN、layer norm那些。

而且，这些中间结果要从HBM读到SRAM，算完再写回去。访存带宽成了瓶颈。我记得有一次在A100上跑长序列，计算单元利用率不到20%，大部分时间都在等数据。

2.6 FlashAttention：革命性的优化方案

FlashAttention的出现，说白了就是解决了上面两个问题：显存爆炸和访存瓶颈。它的核心思想是——不要一次性算出整个 S 和 A，而是分块计算。

具体怎么做？我简单描述一下：

  1. 把Q、K、V都切成小块（tile），比如每块128行。
  2. 每次只加载一个Q块和多个K块到SRAM中，计算局部注意力分数。
  3. 在SRAM里直接做softmax，然后加权求和。
  4. 把结果写回HBM，释放SRAM空间。

这里有个技术难点：softmax是全局操作，需要知道所有key的分数才能归一化。分块计算时，你怎么知道当前块的归一化系数？

FlashAttention的解法很巧妙——在线softmax（online softmax）。它维护两个统计量：当前最大值 m 和当前指数和 d。每处理一个新块，就更新这两个值，同时修正之前的结果。

在线softmax的核心公式：
假设已处理块的最大值为 m_old，指数和为 d_old。
新块的最大值为 m_new，指数和为 d_new。
更新后的全局最大值 m = max(m_old, m_new)
更新后的全局指数和 d = d_old * exp(m_old - m) + d_new * exp(m_new - m)
之前的结果需要乘以 exp(m_old - m) 来修正。

这个算法让我印象深刻。我第一次看到时，心里想的是：还能这么玩？后来在项目中实现了一遍，发现效果确实好。在序列长度8192时，FlashAttention比标准实现快了3倍以上，显存占用降低了80%。

我的建议： 如果你在做长序列模型（比如文档级别的NLP任务、视频理解等），直接上FlashAttention。别犹豫。PyTorch 2.0已经内置了 torch.nn.functional.scaled_dot_product_attention，底层自动调用FlashAttention。用起来很方便。

2.7 算子层面的优化要点

最后，我总结一下自注意力机制中各个算子的优化要点：

算子 计算类型 瓶颈 优化方向
QKV投影 矩阵乘法 访存（权重加载） 权重融合、kernel合并
Q@K^T 矩阵乘法 计算（O(S^2)） 分块计算、利用tensor core
Softmax 逐行规约 访存（读写S矩阵） 在线softmax、fused kernel
A@V 矩阵乘法 访存（读A矩阵） 与softmax融合、分块计算

嗯，这里要特别强调一点：算子融合是自注意力优化的核心思路。把多个小算子合并成一个大的fused kernel，减少中间结果的读写。FlashAttention就是把这个思路发挥到了极致——它把Q@K^T、softmax、A@V三个算子完全融合在一起，中间结果根本不落HBM。

我曾经在一个项目中，手动实现了QKV投影和注意力计算的融合。虽然代码写起来很痛苦（要处理各种边界条件和分块逻辑），但效果立竿见影——端到端推理速度提升了40%。从那以后，我养成了一个习惯：看到连续的几个小算子，第一反应就是能不能融合。

好了，自注意力机制算子这块就讲到这里。内容不少，但核心就两句话：标准实现简单直接但显存爆炸，FlashAttention分块计算又省又快。你在实际开发中，根据序列长度和硬件条件选择合适的方案就行。

← 上一章 📖 返回目录 下一章 →