第二十二课：Softmax 为什么难以高效，以及 FlashAttention 如何减少显存访问

导语：这一课把 Reduction 应用到 Softmax 和 Attention，并推导 FlashAttention 如何用分块与在线 Softmax 避免保存完整的 $N\times N$ 注意力矩阵。

Softmax 在分类、归一化和 Transformer Attention 中非常常见：

$$p_i=\frac{e^{x_i}}{\sum_{j=0}^{N-1}e^{x_j}}$$

它看起来只是：

指数
求和
除法

但在 GPU 上，Softmax 同时包含：

• 最大值归约；
• 求和归约；
• 指数函数；
• 多次读取同一行数据；
• 线程间同步；
• 数值稳定性处理。

而在 Attention 中，Softmax 的输入还是一个可能非常巨大的 $N\times N$ 矩阵。因此，Attention 的性能瓶颈往往不只是矩阵乘法，更是中间结果在显存中的反复读写。

---

一、为什么不能直接计算普通 Softmax

直接实现：

$$p_i=\frac{e^{x_i}}{\sum_j e^{x_j}}$$

存在数值溢出风险。

例如：

$$x_i=100$$

那么：

$$e^{100}$$

对于 FP32 来说可能已经非常大。

如果输入中出现：

1000
999
998

直接计算指数会溢出成：

inf
inf
inf

最后可能出现：

$$\frac{\infty}{\infty}$$

得到 NaN。

---

二、数值稳定 Softmax

Softmax 有一个重要性质：

$$\frac{e^{x_i-c}} {\sum_j e^{x_j-c}} = \frac{e^{x_i}} {\sum_j e^{x_j}}$$

因为分子分母同时乘了：

$$e^{-c}$$

通常选择：

$$c=\max_jx_j$$

因此数值稳定版本为：

$$m=\max_jx_j$$ $$s=\sum_j e^{x_j-m}$$ $$p_i=\frac{e^{x_i-m}}{s}$$

因为：

$$x_i-m\leq0$$

所以：

$$0<e^{x_i-m}\leq1$$

这样可以避免指数上溢。

---

三、稳定 Softmax 至少需要哪些步骤

对于一行输入 $x$，朴素实现需要：

第一次遍历

寻找最大值：

$$m=\max_i x_i$$

第二次遍历

计算指数和：

$$s=\sum_i e^{x_i-m}$$

第三次遍历

生成输出：

$$p_i=\frac{e^{x_i-m}}{s}$$

即：

读取输入 → Max Reduction
再次读取 → Exp + Sum Reduction
再次读取 → Exp + Normalize + 写输出

如果中间指数结果不保存，则同一个 $x_i$ 会被读取多次，并且指数可能计算两次。

如果保存中间指数：

exp(x_i - m)

又会增加一次中间数组写入和读取。

---

四、Softmax 为什么通常不是纯计算受限

假设处理一个长度为 $N$ 的 FP32 行。

至少需要读取：

$$4N\text{ Bytes}$$

并写出：

$$4N\text{ Bytes}$$

如果输入被读取三次，实际流量会更大。

每个元素主要执行：

• 比较；
• 指数；
• 加法；
• 除法或乘倒数。

对于较短的行，线程同步和 Kernel 启动成本明显。

对于较长的行，则可能受到：

• Global Memory 带宽；
• 指数函数吞吐率；
• Reduction 延迟；
• 同步次数；

共同限制。

---

五、Softmax 的并行结构

假设张量形状为：

$$[B,N]$$

通常每一行独立做 Softmax：

Row 0 → Softmax
Row 1 → Softmax
Row 2 → Softmax
...

因此可以映射为：

一个 Warp 或一个 Block
处理一行

具体选择取决于 $N$。

---

六、Warp Softmax

如果一行较短，例如：

$$N\leq32$$

可以让一个 Warp 处理一行。

Lane 0 读取 x₀
Lane 1 读取 x₁
...
Lane 31 读取 x₃₁

然后通过 Warp Shuffle 完成 Max Reduction：

float max_value = value;

for (int offset = 16; offset > 0; offset >>= 1) {
    max_value = max(
        max_value,
        __shfl_down_sync(
            0xffffffff,
            max_value,
            offset));
}

再将最终最大值广播给所有 Lane：

max_value =
    __shfl_sync(
        0xffffffff,
        max_value,
        0);

随后每个线程计算：

float exp_value = expf(value - max_value);

再进行 Sum Reduction。

---

七、Warp Softmax 的完整流程

每个 Lane 读取一个或多个元素
            ↓
寄存器中计算局部最大值
            ↓
Warp Shuffle Max Reduction
            ↓
计算 exp(x - max)
            ↓
Warp Shuffle Sum Reduction
            ↓
每个 Lane 除以总和
            ↓
写回输出

主要优点：

• 不需要 Shared Memory；
• 不需要 __syncthreads()；
• 数据主要保存在寄存器；
• Warp 内交换速度快。

---

八、当一行超过32个元素

如果：

$$N>32$$

一个线程可以处理多个元素。

例如 $N=128$，一个 Warp 有32个线程：

Lane 0：处理元素 0、32、64、96
Lane 1：处理元素 1、33、65、97
...

每个线程先在寄存器中求局部最大值：

$$m_t=\max(x_t,x_{t+32},x_{t+64},x_{t+96})$$

然后 Warp 对32个局部最大值归约。

求和阶段同理。

这种方式仍然可以由一个 Warp 处理一行，但会增加每线程寄存器需求。

---

九、Block Softmax

如果一行非常长，单个 Warp 处理会导致：

• 每线程处理元素过多；
• 寄存器使用过大；
• 单 Warp 并行度不足；
• 指数和内存访问延迟难以隐藏。

此时可以让一个 Block 处理一行。

例如：

Block：256 Threads
一行：4096个元素

每个线程循环处理：

$$\frac{4096}{256}=16$$

个元素。

---

十、Block Softmax 的两级归约

第一步，每个线程计算局部最大值：

$$m_t=\max(\text{elements handled by thread})$$

第二步，每个 Warp 使用 Shuffle 求出 Warp 最大值。

第三步，每个 Warp 的 Lane 0 把结果写入 Shared Memory。

第四步，第一个 Warp 对所有 Warp 的结果继续归约。

最终得到整行最大值。

求和阶段使用相同结构：

线程局部和
    ↓
Warp Reduction
    ↓
Shared Memory 保存 Warp 部分和
    ↓
第一个 Warp 得到整行和

---

十一、Softmax 为什么需要保存输入或中间结果

计算最大值时，只知道：

$$m$$

但之后还需要每个原始元素：

$$e^{x_i-m}$$

如果输入行足够短，可以让每个线程把自己负责的 $x_i$ 保存在寄存器中。

这样：

第一次读取输入
    ↓
输入保存在寄存器
    ↓
完成 Max 和 Sum Reduction
    ↓
直接写结果

不需要反复读取 Global Memory。

但如果一行很长，所有元素无法全部放入寄存器，就需要在以下方案中选择：

• 重新读取输入；
• 使用 Shared Memory；
• 分块处理；
• 使用 Online Softmax。

---

十二、Attention 的基本公式

Transformer 的单头 Attention 为：

$$S=\frac{QK^T}{\sqrt d}$$ $$P=\operatorname{softmax}(S)$$ $$O=PV$$

其中：

$$Q,K,V\in\mathbb{R}^{N\times d}$$ $$S,P\in\mathbb{R}^{N\times N}$$ $$O\in\mathbb{R}^{N\times d}$$

这里：

• $N$：序列长度；
• $d$：每个 Attention Head 的维度。

---

十三、Attention 在做什么

对于第 $i$ 个 Query：

$$q_i$$

它与所有 Key 计算相似度：

$$s_{ij}=\frac{q_i\cdot k_j}{\sqrt d}$$

然后对整行进行 Softmax：

$$p_{ij} = \frac{e^{s_{ij}}} {\sum_{t=0}^{N-1}e^{s_{it}}}$$

最后对 Value 加权：

$$o_i = \sum_{j=0}^{N-1} p_{ij}v_j$$

所以每个 Query 都需要观察全部 Key 和 Value。

---

十四、朴素 Attention 的执行方式

最直接的实现可以拆成三个 Kernel 或库操作：

Kernel 1：
S = QKᵀ

Kernel 2：
P = Softmax(S)

Kernel 3：
O = PV

数据流为：

Q、K
  ↓ GEMM
S 写入 HBM
  ↓
S 从 HBM 读取
  ↓ Softmax
P 写入 HBM
  ↓
P 从 HBM 读取
V 从 HBM 读取
  ↓ GEMM
O 写入 HBM

主要问题是：

巨大的 $N\times N$ 矩阵 $S$ 和 $P$ 被写入显存，再读取回来。

---

十五、$N\times N$ 中间矩阵有多大

假设：

$$N=4096$$

那么：

$$N^2=16\,777\,216$$

如果使用 FP16，每个元素2字节：

$$16\,777\,216\times2 \approx32\text{ MB}$$

这只是：

• 一个 Batch；
• 一个 Head；
• 一个中间矩阵。

假设有32个 Head：

$$32\times32\text{ MB}=1024\text{ MB}$$

也就是约1 GB。

实际训练还可能需要保存：

• Softmax 概率；
• Dropout Mask；
• 反向传播中间状态；
• 多个 Batch 的结果。

显存压力非常大。

---

十六、Attention 的计算量和内存量

矩阵乘法：

$$QK^T$$

计算复杂度约为：

$$O(N^2d)$$

后面的：

$$PV$$

也是：

$$O(N^2d)$$

FlashAttention 并没有消除这个二次计算量。

它优化的重点是：

不把完整的 $N\times N$ 分数矩阵和概率矩阵写入 HBM。

因此它是一种 IO-Aware 算法，而不是降低理论 Attention 计算复杂度的稀疏近似算法。

---

十七、为什么普通 Kernel Fusion 还不够

可以尝试把：

QKᵀ
Softmax
PV

融合到一个 Kernel。

但 Softmax 存在全行依赖：

必须先知道一整行的最大值和总和
才能得到该行的每个概率

假设只计算了部分 Key：

K Tile 0

此时不能确定这一部分的 Softmax 概率，因为后面的 Key Tile 可能出现更大的分数，从而改变：

• 行最大值；
• 分母；
• 前面所有概率的归一化结果。

要实现真正的分块融合，需要一种可以逐块更新 Softmax 的方法。

这就是 Online Softmax。

---

十八、Online Softmax 的基本目标

假设一行分数被拆成两个部分：

第一块：A
第二块：B

先处理 A 时得到：

$$m_A=\max(A)$$ $$l_A=\sum_{x\in A}e^{x-m_A}$$

处理 B 时得到：

$$m_B=\max(B)$$ $$l_B=\sum_{x\in B}e^{x-m_B}$$

现在希望不重新读取 A，就合并成整行结果。

---

十九、Online Softmax 的最大值更新

合并后的最大值为：

$$m=\max(m_A,m_B)$$

但 $l_A$ 是相对于 $m_A$ 计算的：

$$l_A=\sum_{x\in A}e^{x-m_A}$$

需要转换为相对于新最大值 $m$：

$$\sum_{x\in A}e^{x-m} = e^{m_A-m} \sum_{x\in A}e^{x-m_A}$$

因此：

$$l = e^{m_A-m}l_A + e^{m_B-m}l_B$$

这就是在线更新分母的核心公式。

---

二十、按元素流式更新 Online Softmax

也可以逐元素处理。

假设处理前 $k$ 个元素时保存：

$$m_k=\max(x_1,\ldots,x_k)$$ $$l_k=\sum_{i=1}^{k}e^{x_i-m_k}$$

加入新元素 $x_{k+1}$：

$$m_{k+1}=\max(m_k,x_{k+1})$$ $$l_{k+1} = e^{m_k-m_{k+1}}l_k + e^{x_{k+1}-m_{k+1}}$$

因此只需要维护：

当前最大值 m
当前归一化分母 l

就能逐步处理完整的一行。

---

二十一、仅有分母还不够

Attention 最终需要：

$$o_i = \sum_jp_{ij}v_j$$

展开：

$$o_i = \frac{ \sum_j e^{s_{ij}}v_j }{ \sum_j e^{s_{ij}} }$$

为了数值稳定，可以写成：

$$o_i = \frac{ \sum_j e^{s_{ij}-m_i}v_j }{ \sum_j e^{s_{ij}-m_i} }$$

因此除了维护分母 $l$，还可以维护未归一化输出累加器：

$$a=\sum_j e^{s_j-m}v_j$$

最终：

$$o=\frac{a}{l}$$

---

二十二、在线更新 Attention 输出

对于旧数据块 A，保存：

$$m_A$$ $$l_A=\sum_{j\in A}e^{s_j-m_A}$$ $$a_A=\sum_{j\in A}e^{s_j-m_A}v_j$$

对于新数据块 B：

$$m_B$$ $$l_B=\sum_{j\in B}e^{s_j-m_B}$$ $$a_B=\sum_{j\in B}e^{s_j-m_B}v_j$$

合并最大值：

$$m=\max(m_A,m_B)$$

合并分母：

$$l= e^{m_A-m}l_A + e^{m_B-m}l_B$$

合并输出累加器：

$$a= e^{m_A-m}a_A + e^{m_B-m}a_B$$

最终：

$$o=\frac{a}{l}$$

这使得之前处理过的数据块不需要重新读取。

---

二十三、FlashAttention 的核心思路

FlashAttention 将 $Q$、$K$、$V$ 分成较小 Tile。

例如：

Q Tile：Br × d
K Tile：Bc × d
V Tile：Bc × d

这些 Tile 被搬入：

• Shared Memory；
• 寄存器；

然后在片上完成：

Q Tile × K Tileᵀ
       ↓
局部 Score Tile
       ↓
局部 Softmax 统计
       ↓
与 V Tile 相乘
       ↓
更新输出累加器

处理完一个 $K,V$ Tile 后，不需要把完整 Score Tile 或 Probability Tile 写入 HBM。

---

二十四、FlashAttention 的高层数据流

对于一个 Q Tile：

加载 Q Tile 到片上
        ↓
遍历所有 K/V Tile
        ↓
加载 K Tile、V Tile
        ↓
计算局部 S = QKᵀ
        ↓
更新每行最大值 m
        ↓
更新每行分母 l
        ↓
更新输出累加器 O
        ↓
处理下一个 K/V Tile
        ↓
最终归一化
        ↓
只把最终 O 写入 HBM

完整的 $N\times N$ 注意力矩阵从未出现在 HBM 中。

它只以小 Tile 的形式短暂存在于片上存储中。

---

二十五、传统 Attention 与 FlashAttention 的数据流比较

传统实现

Q、K
 ↓
QKᵀ
 ↓
S 写入 HBM
 ↓
S 读回
 ↓
Softmax
 ↓
P 写入 HBM
 ↓
P、V 读回
 ↓
PV
 ↓
O 写入 HBM

FlashAttention

Q、K、V 分块读入
       ↓
片上计算局部分数
       ↓
片上在线 Softmax
       ↓
片上累加输出
       ↓
只写最终 O

减少的是：

S 的完整写入与读取
P 的完整写入与读取

---

二十六、为什么片上 Tile 能提高性能

GPU 存储层次大致为：

寄存器
速度最快，容量最小

Shared Memory / L1
速度较快，容量有限

L2

HBM / GDDR
容量大，访问成本最高

传统 Attention 不断在 HBM 和 SM 之间搬运 $N^2$ 中间数据。

FlashAttention 用更多片上计算和更复杂的调度，换取更少的 HBM 流量。

其基本思想是：

$$\text{trade compute for IO}$$

在现代 GPU 中，Tensor Core 算力增长通常快于 HBM 带宽增长，因此减少显存访问往往非常有价值。

---

二十七、FlashAttention 是否是近似算法

不是。

在忽略浮点运算顺序差异的情况下，FlashAttention 计算的是与标准 Attention 相同的数学结果：

$$\operatorname{softmax} \left( \frac{QK^T}{\sqrt d} \right)V$$

它没有：

• 删除部分 Key；
• 使用低秩近似；
• 把 Attention 改成线性形式；
• 减少理论上的全部 $QK^T$ 计算。

不同实现可能因为浮点舍入顺序不同，产生最后几位差异，但这不属于算法近似。

因此 FlashAttention 通常被称为：

Exact Attention 的 IO-Aware 实现。

---

二十八、FlashAttention 的计算复杂度

标准 Attention 的主要计算量为：

$$O(N^2d)$$

FlashAttention 仍然是：

$$O(N^2d)$$

它并没有把二次复杂度变为线性。

主要改善是：

• 不显式保存 $N\times N$ 矩阵；
• 降低 HBM 读写；
• 降低中间激活显存占用；
• 提高 Tensor Core 与数据搬运重叠能力。

因此序列长度继续增加时，计算量仍按 $N^2$ 增长。

---

二十九、显存复杂度的改善

朴素 Attention 前向过程中，可能显式保存：

$$S\in\mathbb{R}^{N\times N}$$

和：

$$P\in\mathbb{R}^{N\times N}$$

中间激活规模为：

$$O(N^2)$$

FlashAttention 不保存完整的 $S$ 和 $P$，只保存例如：

• 输入 $Q,K,V$；
• 输出 $O$；
• 每行最大值或 LogSumExp；
• 少量其他状态。

因此与注意力矩阵相关的额外存储可以从：

$$O(N^2)$$

降到接近：

$$O(N)$$

或总体张量存储中的 $O(Nd)$ 量级。

---

三十、为什么反向传播也不保存完整概率矩阵

训练时，普通实现通常会保存：

$$P=\operatorname{softmax}(S)$$

供反向传播使用。

但 $P$ 是一个巨大的 $N\times N$ 矩阵。

FlashAttention 的反向传播策略通常是：

保存少量每行统计量
        ↓
反向时重新计算局部 QKᵀ
        ↓
重建局部 Softmax 概率
        ↓
计算 dQ、dK、dV

这称为 Recompute 或 Rematerialization。

---

三十一、为什么重新计算反而可能更快

直觉上，重新计算会增加运算量。

但如果选择是：

方案 A：
从 HBM 读取巨大的概率矩阵

方案 B：
使用 Tensor Core 重新计算小块 QKᵀ

在现代 GPU 上，方案 B 可能更快。

因为：

• Tensor Core 算力非常高；
• HBM 带宽相对稀缺；
• 重新计算发生在片上 Tile 流水中；
• 避免保存和读取 $N^2$ 激活。

这再次体现：

算力可能便宜，数据搬运更昂贵。

---

三十二、FlashAttention 的伪代码

忽略 Batch、Head 和边界，可以写成：

for each Q tile Qi:

    load Qi

    m = -∞
    l = 0
    acc = 0

    for each K/V tile Kj, Vj:

        load Kj, Vj

        S = Qi × Kjᵀ / sqrt(d)

        block_max = row_max(S)
        new_m = max(m, block_max)

        P_local = exp(S - new_m)

        correction = exp(m - new_m)

        l = correction × l
            + row_sum(P_local)

        acc = correction × acc
              + P_local × Vj

        m = new_m

    Oi = acc / l

    store Oi

这里的：

correction = exp(m - new_m)

用于重新缩放前面已经累积的结果。

---

三十三、一个简单数值例子

假设第一块分数是：

$$[2,1]$$

第一块最大值：

$$m_1=2$$

第一块分母：

$$l_1=e^{2-2}+e^{1-2} =1+e^{-1}$$

现在第二块分数是：

$$[4,3]$$

第二块引入了更大的最大值：

$$m_2=4$$

此前的分母是相对于2计算的，需要重新缩放：

$$e^{2-4}l_1$$

第二块的贡献是：

$$e^{4-4}+e^{3-4}$$

所以新分母为：

$$l_2 = e^{-2}(1+e^{-1}) + 1+e^{-1}$$

它等价于直接计算：

$$e^{2-4}+e^{1-4}+e^{4-4}+e^{3-4}$$

但不需要重新访问第一块的两个分数。

---

三十四、Causal Attention 如何处理

自回归模型中，第 $i$ 个 Token 不能看到未来 Token：

$$j>i$$

因此需要因果掩码：

$$s_{ij}=-\infty,\quad j>i$$

Softmax 后：

$$e^{-\infty}=0$$

所以未来位置概率为0。

FlashAttention 在处理 Score Tile 时，可以根据 Query 和 Key 的位置关系：

• 完全跳过位于因果区域之外的 Tile；
• 对对角线附近的 Tile 使用局部 Mask。

例如：

Q Tile 0 只访问 K Tile 0
Q Tile 1 访问 K Tile 0、1
Q Tile 2 访问 K Tile 0、1、2

这样不仅保证正确性，还可以减少上三角区域的无效计算。

---

三十五、Padding Mask 如何处理

不同序列长度可能被 Padding 到统一长度。

无效 Key 位置可以设置为：

$$-\infty$$

使其 Softmax 权重为0。

实际 Kernel 中通常通过：

• 边界谓词；
• Mask；
• Tile 有效范围；

处理。

但过多不规则长度可能导致：

• 分支；
• 无效计算；
• Warp 利用率下降。

因此现代框架还会使用：

• Packed Sequence；
• Variable-Length Attention；
• 按长度分组；

减少 Padding 浪费。

---

三十六、Dropout 如何处理

训练时 Attention 概率可能经过 Dropout：

$$\tilde P_{ij} = \frac{M_{ij}P_{ij}}{1-p}$$

其中 $M_{ij}$ 是随机掩码。

高效实现不希望保存完整的 $N\times N$ Dropout Mask。

一种常见思路是：

• 使用可重现的计数器型随机数生成；
• 前向保存随机数种子和偏移；
• 反向时重新生成相同掩码。

这样用少量随机数状态代替完整 Mask 存储。

---

三十七、Tensor Core 在 FlashAttention 中做什么

FlashAttention 中主要有两类矩阵乘法：

$$QK^T$$

和：

$$PV$$

它们都可以映射到 Tensor Core。

执行流程可能类似：

Global Memory
      ↓
Shared Memory
      ↓
寄存器 Fragment
      ↓
Tensor Core MMA
      ↓
寄存器累加器

但 Softmax 相关的：

• Max；
• Exp；
• Sum；
• Rescale；

通常仍由普通标量/向量执行单元完成。

所以 FlashAttention 是多种硬件协作：

Tensor Core：矩阵乘法
CUDA Core/SFU：指数、归约、缩放
Shared Memory：Tile 暂存
寄存器：局部统计和输出累加
LSU/TMA：数据搬运

---

三十八、为什么 Tile 大小很重要

Tile 太小：

• 矩阵乘法效率低；
• Tensor Core 利用不足；
• 循环次数增加；
• 元数据和同步开销增大。

Tile 太大：

• Shared Memory 不够；
• 寄存器压力增加；
• Occupancy 降低；
• Spill 风险增加；
• 一次处理的 Score Tile 太大。

因此需要在以下因素间平衡：

Tensor Core 效率
Shared Memory 容量
寄存器容量
Occupancy
数据复用
同步开销

这也是为什么不同：

• Head Dimension；
• 序列长度；
• 数据类型；
• GPU 架构；

需要不同 Kernel 配置。

---

三十九、为什么 Head Dimension 会影响实现

Attention Head 维度常见为：

64
80
96
128
256

$d$ 决定：

• Q/K/V 每行长度；
• Tensor Core K 维循环次数；
• 每个线程保存的寄存器数据；
• Shared Memory Tile 大小；
• 输出累加器规模。

例如 $d=128$ 相比 $d=64$：

• 矩阵计算量增加；
• Q/K/V Tile 更大；
• 寄存器和 Shared Memory 压力增加；
• Block 配置可能需要变化。

---

四十、FlashAttention 为什么不一定在所有情况都更快

对于很短的序列，例如：

$$N=32$$

普通实现的中间矩阵很小。

此时 FlashAttention 的：

• 复杂 Tile 调度；
• 在线缩放；
• Kernel 启动；
• 边界处理；

可能占比较高。

其他可能限制包括：

• Head Dimension 不适合当前 Tile；
• GPU 较旧，片上资源较少；
• Batch 很小；
• 输入布局需要额外转换；
• 特殊 Mask 模式不适配；
• Kernel 实现未针对当前架构优化。

因此 FlashAttention 的优势通常在中长序列、大 Batch 或训练场景中更加明显。

---

四十一、FlashAttention-2 改进了什么

第一代 FlashAttention 已经大幅降低 HBM 访问，但仍存在一些非矩阵运算和并行映射开销。

FlashAttention-2 的主要方向可以概括为：

• 减少非矩阵乘法 FLOPs；
• 改进线程块和 Warp 的工作划分；
• 降低 Warp 间通信；
• 提高 Tensor Core 利用率；
• 更好地并行不同 Query 行。

核心算法仍然是：

Tiling
+
Online Softmax
+
不显式保存 N×N 矩阵

它主要优化的是 GPU 上的工作分配和微架构映射。

---

四十二、后续版本为什么强调异步数据流

在更新 GPU 上，矩阵计算速度进一步提高。

此时需要更好地把：

加载 K/V Tile
计算当前 Tile
处理 Softmax
写出结果

进行流水化。

可能使用：

• 异步 Global-to-Shared 搬运；
• 多缓冲；
• Warp Specialization；
• TMA；
• 更细粒度生产者—消费者同步；
• 新型 Tensor Core 指令。

例如：

Warp Group A：
搬运下一块数据

Warp Group B：
执行 QKᵀ

Warp Group C：
处理 Softmax 和 PV

不同角色形成流水线，尽量减少 Tensor Core 等待数据的时间。

---

四十三、Warp Specialization 是什么

普通设计中，每个 Warp 执行相似工作。

Warp Specialization 则让不同 Warp 承担不同角色：

Producer Warp：
负责数据搬运

Consumer Warp：
负责矩阵计算

其他 Warp：
负责 Softmax、归约或输出

优点：

• 任务分工清晰；
• 有利于异步流水；
• 减少每个 Warp 的控制复杂度；
• 可以让数据搬运与计算重叠。

缺点：

• Warp 间同步更复杂；
• 负载平衡困难；
• 某一角色速度不足会阻塞整个流水线；
• 需要较强硬件和编译器支持。

---

四十四、FlashAttention 与普通 GEMM 的区别

普通 GEMM：

$$C=AB$$

特点：

• 输出矩阵每个 Tile 可以独立计算；
• 不需要跨整行的 Softmax 归一化；
• 数据流相对规则。

FlashAttention：

QKᵀ
  ↓
每行 Max/Sum Reduction
  ↓
Softmax
  ↓
PV

其困难在于：

• QKᵀ 的输出不能直接写出；
• Softmax 需要整行统计量；
• 统计量会随 K Tile 更新；
• 输出累加器需要重新缩放；
• 同时混合 GEMM、Reduction 和特殊函数。

所以它不是简单地把两个 GEMM 合并。

---

四十五、FlashAttention 与线性 Attention 的区别

线性 Attention 通常尝试通过数学变换，把复杂度从：

$$O(N^2d)$$

降到接近：

$$O(Nd^2)$$

或其他较低形式。

代价通常是改变 Attention 形式或使用核函数近似。

FlashAttention：

• 不改变标准 Softmax Attention；
• 仍然计算所有 Query-Key 组合；
• 仍然具有二次计算复杂度；
• 主要减少 HBM IO 和中间存储。

因此：

FlashAttention：
精确，但计算仍为 O(N²)

线性/稀疏 Attention：
试图减少计算量，但可能改变模型行为

---

四十六、使用 Nsight Compute 分析 Attention

分析普通 Attention 时，可能看到：

QKᵀ GEMM

Tensor Core 吞吐率较高

Softmax

Memory Throughput 较高
Reduction 和 SFU 压力明显

PV GEMM

再次读取大规模 P
Tensor Core 工作

三个阶段之间有大量 HBM 中间流量。

---

FlashAttention 的预期表现

可能看到：

• 单个融合 Kernel 持续时间较长；
• Tensor Core 活跃；
• Shared Memory 使用较高；
• 寄存器使用较高；
• Occupancy 未必很高；
• HBM 中间流量显著下降；
• 算术强度提高；
• Barrier 和异步流水指标较重要。

因此不能仅根据 Occupancy 判断 FlashAttention Kernel 是否高效。

---

四十七、FlashAttention 的性能权衡

FlashAttention 通常采用较大的片上 Tile 和输出累加器。

这会导致：

寄存器增加
Shared Memory 增加
Occupancy 下降

但同时：

HBM 访问大幅减少
数据复用提高
Tensor Core 利用率提高
Kernel 间中间写回消失

因此它是典型的：

通过较低 Occupancy 换取更高单 Block 工作效率和更低 IO。

---

四十八、从 ASIC 角度看 FlashAttention

如果把类似 Attention 部署到专用硬件，核心问题仍然相同：

片上 SRAM 容量有限
外部 DRAM 访问昂贵
矩阵乘法阵列吞吐率很高
Softmax 需要全行归约

硬件数据流可以设计为：

Q Tile Buffer
K Tile Buffer
V Tile Buffer
      ↓
矩阵乘法阵列
      ↓
Score Tile
      ↓
Online Max / Sum
      ↓
概率与 V 的乘加
      ↓
Output Accumulator

核心目标同样是：

• 不把完整 Score Matrix 写回 DRAM；
• 在片上完成在线 Softmax；
• 复用 Q Tile；
• 流式读取 K/V；
• 使用双缓冲隐藏 DRAM 延迟。

---

四十九、在线 Softmax 需要保存哪些状态

对一个 Query 行，最基本状态为：

$$m_i$$

当前最大值。

$$l_i$$

当前缩放后的指数和。

$$a_i\in\mathbb{R}^{d}$$

当前未归一化输出向量。

所以每个 Query 行需要保存：

一个标量最大值
一个标量分母
一个长度 d 的输出累加器

而不是长度为 $N$ 的完整注意力概率行。

当：

$$d\ll N$$

时，存储减少非常明显。

---

五十、本课核心结论

第一，数值稳定 Softmax 需要先减去行最大值：

$$p_i= \frac{e^{x_i-m}} {\sum_je^{x_j-m}}$$

第二，普通稳定 Softmax 包含 Max Reduction 和 Sum Reduction，并可能多次读取同一行数据。

第三，短行适合 Warp Softmax，较长行适合 Block Softmax。

第四，Warp Shuffle 可以在寄存器之间完成 Max 和 Sum Reduction，减少 Shared Memory 与 Barrier。

第五，标准 Attention 为：

$$O= \operatorname{softmax} \left( \frac{QK^T}{\sqrt d} \right)V$$

第六，朴素 Attention 会把完整的 $N\times N$ Score 和 Probability 矩阵写入 HBM，再读取回来。

第七，序列长度增加时，$N\times N$ 中间矩阵会造成巨大的显存容量和带宽压力。

第八，Online Softmax 可以只维护当前最大值和归一化分母，逐块合并 Softmax 统计量。

第九，Attention 还可以维护输出累加器：

$$a=\sum_je^{s_j-m}v_j$$

最终输出：

$$o=\frac{a}{l}$$

第十，FlashAttention 将 Q、K、V 分块放入片上存储，在片上完成局部分数、在线 Softmax 和输出累加。

第十一，FlashAttention 不保存完整的 $N\times N$ 注意力矩阵，只写最终输出和少量统计量。

第十二，FlashAttention 没有降低标准 Attention 的：

$$O(N^2d)$$

计算复杂度，主要降低的是 HBM IO 和中间激活存储。

第十三，FlashAttention 是精确标准 Attention 的实现，不是稀疏或低秩近似算法。

第十四，反向传播可以重新计算局部分数和概率，用额外计算换取更少的激活存储和显存访问。

第十五，FlashAttention 依赖 Tensor Core、Shared Memory、寄存器、Reduction、异步数据搬运和多级流水线共同工作。

第十六，它可能具有较高寄存器和 Shared Memory 使用量，以及不高的 Occupancy，但整体性能仍然更好。

第十七，FlashAttention 的核心思想可以概括为：

$$\boxed{ \text{tiling} + \text{online Softmax} + \text{avoid }N^2\text{ intermediate HBM writes} }$$

下一课将进一步分析 Transformer 中另一个典型融合算子：

第二十三课：LayerNorm 和 RMSNorm 如何在 GPU 上实现

主要包括：

均值与方差的 Reduction
两遍 LayerNorm
Welford 在线算法
数值稳定方差计算
Warp/Block 归约
RMSNorm 为什么更简单
Scale 和 Bias 融合
残差连接融合
为什么归一化算子通常是内存受限
如何设计 ASIC 归一化数据通路