FlashAttention 详解：从 GPU 硬件、普通 Attention 到分块计算与 Online Softmax

FlashAttention 的核心并不是改变 Attention 的数学公式，而是改变 Attention 的执行方式：

普通 Attention 会把巨大的中间矩阵 $S=QKTS=QK^T$ S=QKT 和 $P=softmax(S)P=\backslash operatorname\{softmax\}(S)$ P=softmax(S) 写入 HBM；FlashAttention 则把计算分块，在 GPU 片上高速存储中完成局部计算，并通过 online softmax 直接累积最终输出，从而避免保存完整的 $N\times NN\backslash times\; N$ N×N 中间矩阵。

要理解 FlashAttention，首先必须理解 GPU 的内存层次、Kernel、Block 和 SM。

---

GPU 硬件基础：HBM、SRAM、SM、Kernel、Block

1.1 HBM 是什么？

HBM 是 High Bandwidth Memory，也就是高带宽显存。

在 PyTorch 中，如果创建一个 CUDA tensor：

q = torch.randn(
    batch,
    heads,
    seq_len,
    head_dim,
    device="cuda",
    dtype=torch.float16,
)

这个 tensor 的主体数据通常位于 GPU 的 global memory 中。对于使用 HBM 的 NVIDIA GPU，可以粗略把 global memory 理解成 HBM。

HBM 的特点：

• 容量大，通常是几十 GB

• 所有 SM 都可以访问

• 用来保存模型参数、Q/K/V、输出和中间张量

• 带宽很高，但访问延迟仍然高于片上存储

• 适合保存大规模数据，不适合频繁保存和读取中间结果

例如，NVIDIA A100 的 HBM2 带宽可以达到约 2 TB/s，它属于 GPU 的全局主存储层。虽然 HBM 通过先进封装与 GPU 紧密连接，但仍不属于 SM 内部的片上存储。NVIDIA GPU Performance Background

---

1.2 SRAM 是什么？

SRAM 是 Static Random Access Memory，静态随机存储器。

GPU 内部常见的片上高速存储包括：

Register：寄存器
Shared Memory：共享内存
L1 Cache：一级缓存

这些存储位于 GPU 芯片内部，容量远小于 HBM，但离计算单元更近，访问速度更快。

可以把 GPU 存储层次粗略表示为：

容量越来越大
速度越来越慢

Register
   ↓
Shared Memory / L1 Cache
   ↓
L2 Cache
   ↓
HBM / Global Memory

NVIDIA 的典型层次关系是：

Register：
    每个线程私有，最快，容量最小

Shared Memory：
    同一个 block 内的线程共享

L2 Cache：
    多个 SM 可以访问

HBM / Global Memory：
    整个 GPU 上的线程都可以访问，容量最大

NVIDIA 对 GPU 内存层次的描述也明确指出：global memory 容量最大但速度较慢，shared memory 更小但更快，register 最小且由线程私有。NVIDIA CUDA Memory Hierarchy

---

1.3 为什么 SRAM 对 FlashAttention 很重要？

假设：

N = 4096
head_dim = 128

注意力分数矩阵：

$S=QKTS=QK^T$ S=QKT

形状是：

[4096, 4096]

如果使用 FP16：

$4096\times 4096\times 2\; bytes\approx 32\; MB4096\backslash times4096\backslash times2\backslash text\{\; bytes\}\; \backslash approx32\backslash text\{\; MB\}$ 4096×4096×2 bytes≈32 MB

softmax 后的概率矩阵 $PP$ P 也是同样大小。

但是 GPU 的 shared memory 并不是一个几十 MB 的统一内存池，而是分散在每个 SM 上、容量很小的片上存储。一个 block 通常只能使用几十 KB 到几百 KB 量级的 shared memory。

因此：

完整 S[4096,4096]：几十 MB
单个 block 的片上存储：几十 KB 级别

完整 $SS$ S 无法放入 SRAM。

FlashAttention 的策略不是把完整的 $SS$ S 从 HBM 搬到 SRAM，而是：

每次只加载 Q/K/V 的一个小 tile
在片上计算一个小的 score tile
立即处理它
处理完后丢弃

例如：

Q_tile: [64, 128]
K_tile: [64, 128]
V_tile: [64, 128]
S_tile: [64, 64]

这些 tile 可以被编译器映射到寄存器、shared memory 或其他片上资源中。实际布局取决于 kernel、GPU 架构、tile 大小和寄存器压力；如果寄存器压力过高，还可能发生 register spilling。

---

1.4 Kernel 是什么？

Kernel 是提交给 GPU 执行的函数。

例如 CUDA 代码：

flash_attention_kernel<<<grid, block>>>(Q, K, V, O);

这里：

flash_attention_kernel

就是一个 GPU kernel。

一次 kernel 启动后，GPU 会创建许多执行实例：

一个 Kernel
    ↓
很多个 Block
    ↓
每个 Block 中有很多 Thread

Kernel 可以理解成：

"请 GPU 用大量线程并行执行这个函数。"

---

1.5 Block 是什么？

Block，也叫 Thread Block，是一组协同工作的线程。

例如：

dim3 grid(64);
dim3 block(128);

kernel<<<grid, block>>>();

表示：

64 个 block
每个 block 有 128 个线程

同一个 block 内的线程可以：

• 共享 shared memory

• 进行同步

• 合作加载一个数据 tile

• 协作计算矩阵乘法

• 做 reduction，例如求最大值和求和

在 FlashAttention 中，一个 block 常常负责一块 query：

Block 0：负责 Q[0:64]
Block 1：负责 Q[64:128]
Block 2：负责 Q[128:192]

也就是说：

一个 block 通常负责输出矩阵 $OO$ O 的一部分行。

---

1.6 SM 是什么？

SM 是 Streaming Multiprocessor，可以理解为 GPU 内部的一个计算集群。

一个 GPU 通常由多个 SM 组成：

GPU
├── SM 0
├── SM 1
├── SM 2
├── ...
└── SM n

每个 SM 中包含：

• CUDA cores

• Tensor Cores

• registers

• shared memory

• warp scheduler

SM 是真正执行 block 的硬件单元。

例如：

Block 0 → SM 3
Block 1 → SM 7
Block 2 → SM 1

Block 本身不能指定运行在哪个 SM 上，GPU 调度器会自动安排。

同一个 block 的线程不会被拆到多个 SM 上。一个 block 会完整地驻留在某个 SM 中执行。

但是一个 SM 可以先后执行多个 block：

SM 3：
    先执行 Block 0
    Block 0 完成后
    再执行 Block 5

如果资源足够，一个 SM 也可以同时驻留多个 block。

---

1.7 Thread 和 Warp 是什么？

Thread 是最小执行单位。

例如一个 block 有 128 个线程：

Block 0
├── Thread 0
├── Thread 1
├── ...
└── Thread 127

GPU 通常以 32 个线程为一个 warp 执行：

128 threads = 4 warps

线程通常负责：

• 加载 Q/K/V 的一部分

• 计算矩阵元素的一部分

• 参与最大值和求和

• 写回输出的一部分

---

1.8 Block 和 SM 的关系

可以把它们理解为：

Block：软件层面的任务单元
SM：硬件层面的执行单元

一个 block 执行时占用某个 SM 的资源：

registers
shared memory
线程槽位

如果一个 block 使用的 shared memory 很多，那么一个 SM 同时能运行的 block 数量就会减少。

因此，FlashAttention 中的 tile 大小不仅影响计算，也影响：

• shared memory 占用

• register 占用

• occupancy

• 并行度

• 性能

---

普通 Attention 的数学公式

标准 Attention：

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

定义：

$S=\; QKTd\; S=\backslash frac\{QK^T\}\{\backslash sqrt\; d\}$ S=d QKT

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

$O=PVO=PV$ O=PV

其中：

Q: [N, D]
K: [N, D]
V: [N, D]

S: [N, N]
P: [N, N]
O: [N, D]

用 PyTorch 表示：

scores = q @ k.transpose(-1, -2)
scores = scores * scale
probs = torch.softmax(scores, dim=-1)
out = probs @ v

这三行代码在底层通常对应多个 GPU kernel。

---

普通 Attention 在 GPU 中的完整数据流

假设：

N = 4096
D = 128
BLOCK_M = 64
BLOCK_N = 64

将 Q、K、V 分块：

Q0 = Q[0:64]
Q1 = Q[64:128]
Q2 = Q[128:192]
...

K0 = K[0:64]
K1 = K[64:128]
K2 = K[128:192]
...

V0 = V[0:64]
V1 = V[64:128]
V2 = V[128:192]
...

每个 V tile 的形状是：

Vj: [64, 128]

也就是说，V 沿 sequence/token 维度分块，但保留完整的 head dimension。

---

3.1 普通 Attention 总体图

普通 Attention
                              │
       ┌──────────────────────┼──────────────────────┐
       │                      │                      │
       ▼                      ▼                      ▼
 Kernel 1                  Kernel 2               Kernel 3
 计算 S = QKᵀ             计算 P = softmax(S)     计算 O = P V
       │                      │                      │
       ▼                      ▼                      ▼
 S 写入 HBM                P 写入 HBM              O 写入 HBM

更完整地表示：

Q、K、V 初始保存在 HBM
        │
        ▼
┌───────────────────────────────────────────────┐
│ Kernel 1：计算完整 S = QKᵀ                    │
│                                               │
│ Block(0,0) 计算 S00                           │
│ Block(0,1) 计算 S01                           │
│ Block(0,2) 计算 S02                           │
│ Block(1,0) 计算 S10                           │
│ Block(1,1) 计算 S11                           │
│ ...                                           │
│                                               │
│ 每个 block 计算一个 S tile                    │
└───────────────────────────────────────────────┘
        │
        │ 每个 S tile 写回 HBM
        ▼
完整 S[4096,4096] 保存在 HBM
        │
        ▼
┌───────────────────────────────────────────────┐
│ Kernel 2：从 HBM 读取 S，计算 softmax          │
│                                               │
│ Block 0 处理 S 的一部分行                     │
│ Block 1 处理 S 的一部分行                     │
│ Block 2 处理 S 的一部分行                     │
│ ...                                           │
└───────────────────────────────────────────────┘
        │
        │ 每个 P tile 写回 HBM
        ▼
完整 P[4096,4096] 保存在 HBM
        │
        ▼
┌───────────────────────────────────────────────┐
│ Kernel 3：从 HBM 读取 P 和 V，计算 O = PV     │
│                                               │
│ Block 0 计算 O[0:64]                          │
│ Block 1 计算 O[64:128]                        │
│ Block 2 计算 O[128:192]                      │
│ ...                                           │
└───────────────────────────────────────────────┘
        │
        ▼
完整 O 写入 HBM

---

3.2 Kernel 1：计算 $S=QKTS=QK^T$ S=QKT

完整的 $SS$ S 可以按 block 表示：

S =
┌──────┬──────┬──────┬──────┐
│ S00  │ S01  │ S02  │ ...  │ ← Q0
├──────┼──────┼──────┼──────┤
│ S10  │ S11  │ S12  │ ...  │ ← Q1
├──────┼──────┼──────┼──────┤
│ S20  │ S21  │ S22  │ ...  │ ← Q2
├──────┼──────┼──────┼──────┤
│ ...  │ ...  │ ...  │ ...  │
└──────┴──────┴──────┴──────┘

每个 tile 的含义：

S00 = Q0 @ K0.T
S01 = Q0 @ K1.T
S02 = Q0 @ K2.T

S10 = Q1 @ K0.T
S11 = Q1 @ K1.T
S12 = Q1 @ K2.T

其中：

Q0: [64, 128]
K0: [64, 128]

Q0 @ K0.T:
[64,128] @ [128,64]
→ [64,64]

一个 block 的数据流：

从 HBM 读取 Qi
从 HBM 读取 Kj
        │
        ▼
加载到 SM 的寄存器/shared memory
        │
        ▼
计算 Sij = Qi @ Kj.T
        │
        ▼
得到 Sij[64,64]
        │
        ▼
把 Sij 写回 HBM

因此：

S00 → HBM 中 S[0:64,   0:64]
S01 → HBM 中 S[0:64,  64:128]
S10 → HBM 中 S[64:128, 0:64]
S11 → HBM 中 S[64:128,64:128]

最终，所有 block 把自己的 tile 写到 HBM 的不同位置，拼成完整的：

S[4096,4096]

---

3.3 Kernel 1 的 HBM 写入图

Kernel 1：QKᵀ
                          │
       ┌──────────────────┼──────────────────┐
       │                  │                  │
       ▼                  ▼                  ▼
    Block(0,0)         Block(0,1)         Block(1,0)
    计算 S00           计算 S01           计算 S10
       │                  │                  │
       ▼                  ▼                  ▼
写 HBM 的 S00 区域   写 HBM 的 S01 区域   写 HBM 的 S10 区域

最终 HBM 中保存：

S =
┌────────┬────────┬────────┐
│  S00   │  S01   │  S02   │
├────────┼────────┼────────┤
│  S10   │  S11   │  S12   │
├────────┼────────┼────────┤
│  S20   │  S21   │  S22   │
└────────┴────────┴────────┘

---

3.4 为什么 Kernel 1 必须把 S 写回 HBM？

因为 Kernel 1 和 Kernel 2 是两个独立的 kernel。

Kernel 1 中临时产生的：

register
shared memory

只能在 Kernel 1 执行期间使用。

Kernel 2 启动后，不能直接访问 Kernel 1 的：

register
shared memory

而且不同 block 可能运行在不同的 SM：

Block(0,0) → SM 0
Block(0,1) → SM 1
Block(1,0) → SM 2

不同 SM 之间不能直接把各自的 shared memory 当成一块全局共享内存使用。

因此，Kernel 1 必须把结果写入一个：

跨 SM
跨 block
跨 kernel

都可以访问的位置，这个位置就是 HBM/global memory。

---

3.5 Kernel 2：计算 $P=softmax(S)P=\backslash operatorname\{softmax\}(S)$ P=softmax(S)

Kernel 2 启动后，从 HBM 读取已经保存好的 $SS$ S。

例如：

Block 0：
    读取 S[0:64, :]
    对这些行做 softmax
    得到 P[0:64, :]
    写回 HBM

Block 1：
    读取 S[64:128, :]
    对这些行做 softmax
    得到 P[64:128, :]
    写回 HBM

数据流：

HBM 中的 S tile
       │
       ▼
加载到 SM 的寄存器/shared memory
       │
       ▼
计算一行或多行 softmax
       │
       ▼
得到 P tile
       │
       ▼
P tile 写回 HBM

结果：

P =
┌──────┬──────┬──────┬──────┐
│ P00  │ P01  │ P02  │ ...  │
├──────┼──────┼──────┼──────┤
│ P10  │ P11  │ P12  │ ...  │
├──────┼──────┼──────┼──────┤
│ P20  │ P21  │ P22  │ ...  │
├──────┼──────┼──────┼──────┤
│ ...  │ ...  │ ...  │ ...  │
└──────┴──────┴──────┴──────┘

又发生了一次：

HBM 读取 S
HBM 写回 P

---

3.6 Kernel 3：计算 $O=PVO=PV$ O=PV

Kernel 3 需要：

P：从 HBM 读取
V：从 HBM 读取

假设 Block 0 负责：

O0 = O[0:64, :]

它需要计算：

$O0=P0V\; O\_0=P\_0V$ O0=P0V

其中：

P0: [64, 4096]
V:  [4096, 128]
O0: [64, 128]

将中间的 4096 维切成 64 一块：

$O0=P00V0+P01V1+P02V2+\cdots +\; P0,63\; V63\; O\_0=\; P\_\{00\}V\_0+\; P\_\{01\}V\_1+\; P\_\{02\}V\_2+\backslash cdots+\; P\_\{0,63\}V\_\{63\}$ O0=P00V0+P01V1+P02V2+⋯+P0,63V63

这里必须区分：

P00 只对应 V0
P01 只对应 V1
P02 只对应 V2
...
P0,63 只对应 V63

完整流程：

acc0 = 0

读取 P00 和 V0
acc0 += P00 @ V0

读取 P01 和 V1
acc0 += P01 @ V1

读取 P02 和 V2
acc0 += P02 @ V2

...

读取 P0,63 和 V63
acc0 += P0,63 @ V63

O0 = acc0
写入 O[0:64, :]

每一项的形状：

P0j: [64, 64]
Vj:  [64, 128]

P0j @ Vj
    ↓
[64,128]

所以，普通 Attention 中，一个负责 $O0\; O\_0$ O0 的 block 需要遍历完整的 V：

V0, V1, V2, ..., V63

但它是分块读取的，不是一次把完整 V 放进 SRAM。

---

3.7 Kernel 3 的完整数据流

Block 0：负责 O[0:64]

初始化 acc0[64,128] = 0
        │
        ▼
读取 P00、V0
        │
        ▼
acc0 += P00 @ V0
        │
        ▼
读取 P01、V1
        │
        ▼
acc0 += P01 @ V1
        │
        ▼
读取 P02、V2
        │
        ▼
acc0 += P02 @ V2
        │
        ▼
...
        │
        ▼
读取 P0,63、V63
        │
        ▼
acc0 += P0,63 @ V63
        │
        ▼
得到完整 O0
        │
        ▼
写入 HBM：O[0:64]

Block 1 类似：

读取 P10、V0
acc1 += P10 @ V0

读取 P11、V1
acc1 += P11 @ V1

...

读取 P1,63、V63
acc1 += P1,63 @ V63

写入 HBM：O[64:128]

---

3.8 普通 Attention 的完整 HBM 读写路径

初始：

Q、K、V 保存在 HBM


第一阶段：Kernel 1

从 HBM 读取 Q、K
        │
        ▼
片上计算 S tile
        │
        ▼
S tile 写回 HBM
        │
        ▼
完整 S 保存在 HBM


第二阶段：Kernel 2

从 HBM 读取完整 S
        │
        ▼
片上计算 softmax
        │
        ▼
P tile 写回 HBM
        │
        ▼
完整 P 保存在 HBM


第三阶段：Kernel 3

从 HBM 读取 P
从 HBM 读取 V
        │
        ▼
片上计算 P @ V
        │
        ▼
O tile 写回 HBM

压缩成一张图：

Q ───────────────┐
                 ▼
K ───────> [Kernel 1: QKᵀ]
                 │
                 ├── S tile 写入 HBM
                 ▼
           完整 S 在 HBM
                 │
                 ├── S 从 HBM 读出
                 ▼
          [Kernel 2: softmax]
                 │
                 ├── P tile 写入 HBM
                 ▼
           完整 P 在 HBM
                 │
                 ├── P 从 HBM 读出
                 ▼
V ───────> [Kernel 3: P@V]
                 │
                 ├── O tile 写入 HBM
                 ▼
           完整 O 在 HBM

---

普通 Attention 的问题

4.1 中间矩阵是 $O(N2)O(N^2)$ O(N2)

S = [N, N]
P = [N, N]

序列长度加倍时，中间矩阵大小约增加四倍。

因此长序列下，显存占用快速增长。

---

4.2 S 和 P 都要写回 HBM

传统流程：

计算 S → 写 HBM
读取 S → 计算 P → 写 HBM
读取 P、V → 计算 O → 写 HBM

也就是：

HBM → 片上 → HBM
HBM → 片上 → HBM
HBM → 片上 → HBM

大量数据在 HBM 和片上存储之间来回搬运。

---

4.3 计算量不一定是主要瓶颈

标准 Attention 的主要矩阵乘法计算复杂度是：

$O(N2D)O(N^2D)$ O(N2D)

FlashAttention 并没有把这个数学计算量完全消除。它主要减少的是：

中间矩阵的显存读写

所以 FlashAttention 是一种 IO-aware，也就是面向内存访问优化的算法。

---

FlashAttention 的核心思想

FlashAttention 仍然计算：

$O=softmax(QKT)VO=\backslash operatorname\{softmax\}(QK^T)V$ O=softmax(QKT)V

但它不再按照：

完整 S
→ 完整 softmax
→ 完整 P
→ 完整 PV

的方式执行。

而是：

Q/K/V 分块读取
        ↓
片上计算一个 score tile
        ↓
立即完成局部 softmax 更新
        ↓
立即与 V tile 相乘
        ↓
更新输出累加器
        ↓
丢弃 score tile
        ↓
处理下一个 K/V tile

核心原则：

不生成需要长期保存的完整 $N\times NN\backslash times\; N$ N×N 矩阵。

---

FlashAttention 中的分块方式

假设：

Q, K, V: [N, D]
BLOCK_M = 64
BLOCK_N = 64

Q 分块：

Q0 = Q[0:64]
Q1 = Q[64:128]
Q2 = Q[128:192]
...

K/V 分块：

K0,V0 = [0:64]
K1,V1 = [64:128]
K2,V2 = [128:192]
...

Q tile 的形状：

Q_i: [BLOCK_M, D]

K/V tile 的形状：

K_j: [BLOCK_N, D]
V_j: [BLOCK_N, D]

score tile：

Sij = Qi @ Kj.T

形状：

[64,128] @ [128,64]
→ [64,64]

---

FlashAttention 的完整 GPU 数据流

7.1 总体图

FlashAttention Kernel
                                  │
        ┌─────────────────────────┼─────────────────────────┐
        │                         │                         │
      Block 0                   Block 1                   Block 2
      负责 Q0                   负责 Q1                   负责 Q2
        │                         │                         │
      SM 0                      SM 1                      SM 2
        │                         │                         │
        ▼                         ▼                         ▼
 ┌───────────────┐        ┌───────────────┐        ┌───────────────┐
 │ Q0 = Q[0:64]  │        │ Q1 = Q[64:128]│        │ Q2 = Q[128:192]│
 └───────┬───────┘        └───────┬───────┘        └───────┬───────┘
         │                        │                        │
         │ 遍历所有 K/V tile       │ 遍历所有 K/V tile       │ 遍历所有 K/V tile
         ▼                        ▼                        ▼
   得到 O0[0:64]            得到 O1[64:128]          得到 O2[128:192]
         │                        │                        │
         └──────────────┬─────────┴─────────┬──────────────┘
                        ▼                   ▼
                 写回 HBM 中不同位置
                        │
                        ▼
                 完整输出 O[0:N]

这里的"合并"不是求和，而是按输出行位置拼接：

Block 0 → O[0:64]
Block 1 → O[64:128]
Block 2 → O[128:192]

数学上：

$O=$ O0 O1 O2 ⋮ O= \begin{bmatrix} O_0\\ O_1\\ O_2\\ \vdots \end{bmatrix} O= O0O1O2⋮

而不是：

$O=O0+O1+O2\; O=O\_0+O\_1+O\_2$ O=O0+O1+O2

---

7.2 一个 Block 负责什么？

Block 0 负责：

Q0 = Q[0:64, :]

它最终只生成：

O0 = O[0:64, :]

Block 0 不负责：

O[64:128]
O[128:192]

它只负责自己的 query 行。

---

7.3 Block 0 的初始状态

Block 0 从 HBM 读取：

Q0 = Q[0:64, :]

并在片上初始化：

m = [-∞, -∞, ..., -∞]    # [64]
l = [0, 0, ..., 0]        # [64]
acc = zeros([64, 128])    # [64,128]

其中：

m：每一行目前见过的最大 score
l：修正后的 softmax 分母
acc：加权 V 的累加结果

---

FlashAttention 中的 Online Softmax

8.1 为什么不能对每个 tile 单独 softmax？

错误做法：

for k_block in blocks:
    scores = q @ k_block.T
    probs = softmax(scores, dim=-1)
    out += probs @ v_block

这会让每个 K block 单独归一化：

Block 0 的概率和 = 1
Block 1 的概率和 = 1
Block 2 的概率和 = 1

但完整 Attention 要求：

所有 K block 合起来后，每行概率和 = 1

因此不能直接对每个 tile 单独 softmax。

---

8.2 Online Softmax 维护三个量

对于某个 Q block，维护：

m：当前处理范围内的最大 score
l：重新缩放后的指数和
acc：重新缩放后的加权 V 累加值

初始：

$m=-\infty ,l=0,acc=0m=-\backslash infty,\backslash qquad\; l=0,\backslash qquad\; acc=0$ m=−∞,l=0,acc=0

处理新的 score tile $Sj\; S\_j$ Sj：

$mj=max(Sj)\; m\_j=\backslash max(S\_j)$ mj=max(Sj)

$mnew=max(m,mj)\; m\_\{\backslash text\{new\}\}=\backslash max(m,m\_j)$ mnew=max(m,mj)

旧结果的修正因子：

$\alpha =e\; m-mnew\; \backslash alpha=e^\{m-m\_\{\backslash text\{new\}\}\}$ α=em−mnew

当前 tile 的指数项：

$pj=e\; Sj-mnew\; p\_j=e^\{S\_j-m\_\{\backslash text\{new\}\}\}$ pj=eSj−mnew

更新 softmax 分母：

$lnew=\alpha l+\sum pj\; l\_\{\backslash text\{new\}\}\; =\; \backslash alpha\; l+\; \backslash sum\; p\_j$ lnew=αl+∑pj

更新累加结果：

$accnew=\alpha acc+pjVj\; acc\_\{\backslash text\{new\}\}\; =\; \backslash alpha\; acc+\; p\_jV\_j$ accnew=αacc+pjVj

最终：

$O=\; accl\; O=\backslash frac\{acc\}\{l\}$ O=lacc

这和完整 softmax 的数学结果一致，只是没有保存完整的 $SS$ S 和 $PP$ P。

---

FlashAttention 中 Block 0 的详细流程

9.1 处理第一个 K/V tile

从 HBM 读取：

K0 = K[0:64, :]
V0 = V[0:64, :]

片上计算：

$S00=Q0K0T\; S\_\{00\}=Q\_0K\_0^T$ S00=Q0K0T

形状：

Q0: [64,128]
K0: [64,128]
S00: [64,64]

然后：

计算 S00 的行最大值
更新 m
计算指数项
更新 l
计算当前 tile 对输出的贡献
更新 acc

此时：

acc 只包含 V0 的贡献

然后丢弃：

S00
当前局部概率 tile

它们不写回 HBM。

---

9.2 处理第二个 K/V tile

读取：

K1 = K[64:128, :]
V1 = V[64:128, :]

计算：

$S01=Q0K1T\; S\_\{01\}=Q\_0K\_1^T$ S01=Q0K1T

更新：

m
l
acc

此时：

acc 已经包含 V0 和 V1 的贡献

---

9.3 继续遍历完整的 K/V

完整流程：

Q0 固定在片上

K0,V0：
    S00 = Q0 @ K0.T
    online softmax
    更新 m,l,acc
    丢弃 S00

K1,V1：
    S01 = Q0 @ K1.T
    online softmax
    更新 m,l,acc
    丢弃 S01

K2,V2：
    S02 = Q0 @ K2.T
    online softmax
    更新 m,l,acc
    丢弃 S02

...

K63,V63：
    S0,63 = Q0 @ K63.T
    online softmax
    更新 m,l,acc
    丢弃 S0,63

处理完全部 K/V 后：

O0 = acc / l[:, None]

得到：

O0: [64,128]

然后写回：

O[0:64, :] = O0

---

V 是如何分块的？

V 沿 sequence/token 维度分块，保留完整的 head dimension。

假设：

V.shape = [N, D]
N = 4096
D = 128
BLOCK_N = 64

那么：

V0 = V[0:64, :]
V1 = V[64:128, :]
V2 = V[128:192, :]
...
V63 = V[4032:4096, :]

每个 V tile：

Vj.shape = [64,128]

完整 V：

V =
┌─────────────────────┐
│ V0: [0:64, :]       │
├─────────────────────┤
│ V1: [64:128, :]     │
├─────────────────────┤
│ V2: [128:192, :]    │
├─────────────────────┤
│ ...                 │
├─────────────────────┤
│ V63: [4032:4096, :] │
└─────────────────────┘

K 和 V 使用相同的 sequence 分块：

K0 对应 V0
K1 对应 V1
K2 对应 V2
...

原因是：

• K 用来计算 query 对 token 的注意力分数

• V 保存这些 token 的实际内容

---

11. $PP$ P 和 $VV$ V 的对应关系

完整的概率矩阵可以写成：

P =
┌──────┬──────┬──────┬──────┐
│ P00  │ P01  │ P02  │ ...  │ ← Q0
├──────┼──────┼──────┼──────┤
│ P10  │ P11  │ P12  │ ...  │ ← Q1
├──────┼──────┼──────┼──────┤
│ P20  │ P21  │ P22  │ ...  │ ← Q2
├──────┼──────┼──────┼──────┤
│ ...  │ ...  │ ...  │ ...  │
└──────┴──────┴──────┴──────┘

V：

V =
┌──────┐
│ V0   │
├──────┤
│ V1   │
├──────┤
│ V2   │
├──────┤
│ ...  │
└──────┘

计算 $O0=P0V\; O\_0=P\_0V$ O0=P0V：

P00 只和 V0 相乘
P01 只和 V1 相乘
P02 只和 V2 相乘
...
P0,63 只和 V63 相乘

因此：

$O0=P00V0+P01V1+\cdots +\; P0,63\; V63\; O\_0=P\_\{00\}V\_0+P\_\{01\}V\_1+\backslash cdots+P\_\{0,63\}V\_\{63\}$ O0=P00V0+P01V1+⋯+P0,63V63

需要特别注意：

P00 不会遍历 V0 到 V63

而是：

P00 → V0
P01 → V1
P02 → V2
...

是整个负责 $Q0\; Q\_0$ Q0 的 block 遍历：

P00、P01、P02、...、P0,63

---

FlashAttention 的单 Block 完整图

Block 0
              负责 Q[0:64, :]
                       │
                       ▼
             从 HBM 读取 Q0
                       │
                       ▼
          初始化 m、l、acc 到片上
                       │
       ┌───────────────┴────────────────┐
       │                                │
       ▼                                │
读取 K0,V0                              │
       │                                │
       ▼                                │
S00 = Q0 @ K0ᵀ                          │
       │                                │
       ▼                                │
online softmax                          │
       │                                │
       ▼                                │
更新 m、l、acc                          │
       │                                │
       ├── 丢弃 S00、局部概率 tile        │
       │                                │
       ▼                                │
读取 K1,V1                              │
       │                                │
       ▼                                │
S01 = Q0 @ K1ᵀ                          │
       │                                │
       ▼                                │
online softmax                          │
       │                                │
       ▼                                │
更新 m、l、acc                          │
       │                                │
       ├── 丢弃 S01、局部概率 tile        │
       │                                │
       ▼                                │
      ...                               │
       │                                │
       ▼                                │
读取 K63,V63                             │
       │                                │
       ▼                                │
S063 = Q0 @ K63ᵀ                         │
       │                                │
       ▼                                │
online softmax                          │
       │                                │
       ▼                                │
更新最终 m、l、acc                       │
       │                                │
       ▼                                │
O0 = acc / l                            │
       │                                │
       ▼                                │
写回 HBM：O[0:64, :]                     │

长期保留的是：

Q0
m
l
acc

短暂存在、用完即丢的是：

S00、S01、S02、...
局部概率 tile

---

FlashAttention 的总体数据流

FlashAttention Kernel
                                  │
        ┌─────────────────────────┼─────────────────────────┐
        │                         │                         │
      Block 0                   Block 1                   Block 2
      负责 Q0                   负责 Q1                   负责 Q2
        │                         │                         │
      SM 0                      SM 1                      SM 2
        │                         │                         │
        ▼                         ▼                         ▼
      Q0 固定                   Q1 固定                   Q2 固定
        │                         │                         │
        ├── K0,V0                 ├── K0,V0                 ├── K0,V0
        │   S00                   │   S10                   │   S20
        │   更新 acc0             │   更新 acc1             │   更新 acc2
        │                         │                         │
        ├── K1,V1                 ├── K1,V1                 ├── K1,V1
        │   S01                   │   S11                   │   S21
        │   更新 acc0             │   更新 acc1             │   更新 acc2
        │                         │                         │
        ├── ...                   ├── ...                   ├── ...
        │                         │                         │
        ├── K63,V63               ├── K63,V63               ├── K63,V63
        │   更新最终 acc0          │   更新最终 acc1          │   更新最终 acc2
        │                         │                         │
        ▼                         ▼                         ▼
      O0=acc0/l0                O1=acc1/l1                O2=acc2/l2
        │                         │                         │
        ▼                         ▼                         ▼
   写 O[0:64]              写 O[64:128]             写 O[128:192]

所有 block 最后按行组合：

O =
┌──────────────┐
│ O0           │ ← Block 0 写入 O[0:64]
├──────────────┤
│ O1           │ ← Block 1 写入 O[64:128]
├──────────────┤
│ O2           │ ← Block 2 写入 O[128:192]
├──────────────┤
│ ...          │
└──────────────┘

---

FlashAttention 如何解决普通 Attention 的问题？

普通 Attention

Kernel 1：
    Q/K → S_tile → 写 HBM

Kernel 2：
    HBM 中的 S → softmax → P_tile → 写 HBM

Kernel 3：
    HBM 中的 P/V → P@V → O_tile → 写 HBM

中间结果：

S：写入 HBM
P：写入 HBM

---

FlashAttention

一个 fused kernel：

Q_tile 从 HBM 读取
K_tile 从 HBM 读取
V_tile 从 HBM 读取
        │
        ▼
SM 片上计算 S_tile
        │
        ▼
立即做 online softmax
        │
        ▼
立即计算当前 tile 对输出的贡献
        │
        ▼
更新 acc
        │
        ▼
丢弃 S_tile 和局部概率
        │
        ▼
处理下一个 K/V tile
        │
        ▼
最终 O_tile 写回 HBM

核心差异：

普通 Attention：
    S_tile 写 HBM
    再从 HBM 读 S_tile
    P_tile 写 HBM
    再从 HBM 读 P_tile

FlashAttention：
    S_tile 和局部概率在片上用完即丢
    只把最终 O_tile 写回 HBM

---

为什么 FlashAttention 需要遍历完整的 V？

对于一个 query block $Qi\; Q\_i$ Qi：

$Oi=PiV\; O\_i=P\_iV$ Oi=PiV

如果：

P_i: [BLOCK_M, N]
V:   [N, D]

那么必须覆盖 V 的全部 N 行。

分块后：

$Oi=\; Pi0\; V0+\; Pi1\; V1+\cdots +\; Pi,m\; Vm\; O\_i\; =\; P\_\{i0\}V\_0+\; P\_\{i1\}V\_1+\backslash cdots+\; P\_\{i,m\}V\_m$ Oi=Pi0V0+Pi1V1+⋯+Pi,mVm

所以：

一个负责 Q_i 的 block
需要遍历完整的 V tile 序列

但是它不需要把完整 V 一次性加载到 SRAM：

读取 V0 → 累加
读取 V1 → 累加
读取 V2 → 累加
...
读取 Vm → 累加

FlashAttention 和普通 Attention 都需要覆盖所有 V tile，区别是：

普通 Attention：
    P 已经提前完整计算并保存

FlashAttention：
    当前 tile 的权重由 Qi、Kj 即时计算

---

Causal Attention 如何处理？

因果注意力要求：

$j>i\Rightarrow \; Sij\; =-\infty \; j>i\; \backslash Rightarrow\; S\_\{ij\}=-\backslash infty$ j>i⇒Sij=−∞

也就是当前位置不能看到未来 token。

对于 query 位置 q_pos 和 key 位置 k_pos：

if k_pos > q_pos:
    score = -inf

Triton 风格的写法：

causal_mask = offs_m[:, None] >= offs_n[None, :]

qk = tl.where(
    causal_mask,
    qk,
    -float("inf"),
)

实际高效实现会进一步优化：

完全位于对角线下方的 block：
    不需要 mask

完全位于对角线上方的 block：
    整个 block 跳过

位于对角线附近的 block：
    逐元素 mask

对于 causal Attention，一个 query 只需要访问当前位置及之前的 V，而不是完整的未来部分。

---

Triton 代码中对应什么？

Triton 官方教程提供了一个完整的 FlashAttention 2 教学实现：

• Triton 官方06-fused-attention.py

• 官方教程页面

重点函数通常包括：

_attn_fwd_inner
_attn_fwd
_attn_bwd
_attention

大致职责：

教学版 Triton forward 逻辑大致如下：

@triton.jit
def flash_attn_fwd_kernel(
    Q, K, V, O,
    N_CTX,
    scale,
    BLOCK_M: tl.constexpr,
    BLOCK_N: tl.constexpr,
    HEAD_DIM: tl.constexpr,
):
    pid_m = tl.program_id(0)

    offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
    offs_d = tl.arange(0, HEAD_DIM)

    # 加载一个 Q tile
    q = tl.load(...)

    # online softmax 状态
    m_i = tl.full(
        [BLOCK_M],
        -float("inf"),
        dtype=tl.float32,
    )

    l_i = tl.zeros(
        [BLOCK_M],
        dtype=tl.float32,
    )

    acc = tl.zeros(
        [BLOCK_M, HEAD_DIM],
        dtype=tl.float32,
    )

    # 遍历所有 K/V tile
    for start_n in range(0, N_CTX, BLOCK_N):
        offs_n = start_n + tl.arange(0, BLOCK_N)

        k = tl.load(...)
        v = tl.load(...)

        # 当前 score tile
        qk = tl.dot(q, tl.trans(k)) * scale

        # 当前 tile 最大值
        m_ij = tl.max(qk, axis=1)

        # 合并旧最大值和当前最大值
        m_ij_new = tl.maximum(m_i, m_ij)

        # 修正旧结果
        alpha = tl.exp(m_i - m_ij_new)

        # 当前 tile 的指数权重
        p = tl.exp(qk - m_ij_new[:, None])

        # 更新分母
        l_i = alpha * l_i + tl.sum(p, axis=1)

        # 更新累加器
        acc = alpha[:, None] * acc
        acc += tl.dot(p.to(v.dtype), v)

        m_i = m_ij_new

    # 最终归一化
    acc = acc / l_i[:, None]

    # 写回输出
    tl.store(O_ptrs, acc, mask=...)

这段代码中最重要的变量是：

m_i
l_i
acc

它们正是 online softmax 的三个核心状态。

---

FlashAttention 的反向传播

前向计算中没有保存完整的 $PP$ P，所以反向传播不能直接读取完整的 softmax 矩阵。

通常会保存：

• 输出 $OO$ O

• 每一行的 log-sum-exp，简称 LSE

LSE：

$LSEi\; =mi+log(li)\; \backslash operatorname\{LSE\}\_i=m\_i+\backslash log(l\_i)$ LSEi=mi+log(li)

反向时，kernel 会重新遍历 Q/K/V tile，重新计算局部概率。

对于：

$O=PVO=PV$ O=PV

有：

$dV=PTdOdV=P^TdO$ dV=PTdO

$dP=dOVTdP=dOV^T$ dP=dOVT

softmax 梯度：

$dS=P\odot (dP-rowsum(dP\odot P))dS\; =\; P\backslash odot\; \backslash left(\; dP-\backslash operatorname\{rowsum\}(dP\backslash odot\; P)\; \backslash right)$ dS=P⊙(dP−rowsum(dP⊙P))

最后：

$dQ=dSKdQ=dSK$ dQ=dSK

$dK=dSTQdK=dS^TQ$ dK=dSTQ

这体现了 FlashAttention 的一个重要取舍：

少保存中间结果，多做一些重计算。

---

FlashAttention 1、2、3 的大致区别

FlashAttention 1

核心贡献：

• 分块计算

• online softmax

• 不保存完整 attention matrix

• 降低 HBM 读写

FlashAttention 2

进一步优化：

• 更好的任务划分

• 减少线程块之间的同步

• 提高并行度

• 更充分利用 Tensor Core

FlashAttention 3

主要针对 Hopper 等新一代 GPU：

• warp specialization

• Tensor Core 和其他操作重叠

• 异步流水线

• 更好的低精度支持，例如 FP8

因此，FlashAttention 的性能不仅来自公式，还来自：

tile 大小
线程映射
shared memory 布局
register 使用
warp 分工
异步拷贝
Tensor Core

---

最后的完整心智模型

普通 Attention：

Q/K
  ↓
计算完整 S = QKᵀ
  ↓
S 写入 HBM
  ↓
读取 S
  ↓
计算完整 P = softmax(S)
  ↓
P 写入 HBM
  ↓
读取 P 和 V
  ↓
计算 O = PV
  ↓
O 写入 HBM

FlashAttention：

Q/K/V 保存在 HBM
        ↓
每个 block 负责一段 Q 行
        ↓
加载 Q tile
        ↓
遍历全部 K/V tile
        ↓
计算局部 QKᵀ
        ↓
online softmax
        ↓
立即累加当前 V tile 的贡献
        ↓
丢弃局部 score 和概率 tile
        ↓
得到当前 Q block 的输出
        ↓
写回对应的 O 行
        ↓
所有 block 按行拼成完整 O

最重要的三点是：

1. Block 通常负责一块 query 行和对应的输出行。

2. 一个 Q block 必须遍历完整的 K/V sequence tile，
   因为一个输出行要看到所有 key/value。

3. FlashAttention 不保存完整 S 和 P，
   而是在片上逐 tile 完成 score、softmax 和 PV 累加。

最终可以用一句话概括：

FlashAttention 让每个 block 固定处理一段 Q 行，从 HBM 分块读取全部 K/V，在 SM 的寄存器和 shared memory 中逐块计算 $QKTQK^T$ QKT、online softmax 和加权 V 累加，最后只把自己负责的输出行写回 HBM；所有 block 通过不同的输出行位置组成完整的 Attention 结果。

勘误与补充（只补充实现边界，不改变前文内容）

1. 数学阶段不等于 kernel 数量

前文的 Kernel 1、Kernel 2、Kernel 3 是朴素 Attention 的概念划分：

S = QKᵀ
P = softmax(S)
O = PV

实际框架可能将部分阶段融合，因此可能是三个或更多 kernel，也可能是融合 kernel：

朴素实现：QKᵀ → 写 S → softmax → 写 P → PV
部分融合：QKᵀ 与 softmax 融合，或 softmax 与 PV 融合
FlashAttention：QKᵀ、softmax、PV 在 tile 循环中融合

因此，"普通 Attention 会把 S/P 写入 HBM"应理解为"朴素、未融合的实现通常会这样做"，不是所有现代实现都严格如此。

2. 片上存储与 HBM 的严谨说法

本文用 SRAM/片上存储作为帮助理解的总称，但严格区分：

register file：线程私有
shared memory：block 内共享、软件可管理
L1 cache：硬件管理
L2 cache：多个 SM 共享
HBM/global memory：GPU 全局可访问

程序访问的是 global memory 地址；一次访问可能命中 L1/L2，也可能最终访问 HBM。从算法角度通常概括为 HBM/global-memory 访问。HBM 通过先进封装与 GPU 紧密连接，但不属于某个 SM 内部的寄存器或 shared memory。

Q/K/V tile、score tile 和 accumulator 可能被编译器映射到寄存器、shared memory 或其他片上资源；实际布局取决于 GPU 架构、tile 大小和寄存器压力。寄存器压力过高时还可能发生 register spilling。

3. Block 到 SM 的映射

文中的 Block 0 → SM 0 只是示意。真实映射由 GPU 调度器决定，程序通常不能指定 block 固定运行在哪个 SM。本文采用"一个 block/program 负责一个 Q tile"的教学映射；真实实现也可能使用 split-Q、split-K、persistent kernel 或多维 batch/head/sequence grid。

4. Online Softmax 的维度

对于：

Qi: [BLOCK_M, HEAD_DIM]
Kj: [BLOCK_N, HEAD_DIM]
Sj = Qi @ Kjᵀ: [BLOCK_M, BLOCK_N]

m_new 是每个 query 行的向量 [BLOCK_M]，减法沿 key 维广播：
$$Pj$$r,c=exp⁡(Sjr,c− mnew r) P_jr,c=\exp(S_jr,c-m_{new}r) Pjr,c=exp(Sjr,c−mnewr)

l 是数值稳定化后的指数和，而不是未平移的 sum(exp(score))：
$$li=\sum jexp(\; sij\; -mi)\; l\_i=\backslash sum\_j\backslash exp(s\_\{ij\}-m\_i)$$ li=j∑exp(sij−mi)

初始化为：

m_i=-∞
l_i=0
acc_i=0

非法的 padding/causal score 被设为 -inf，其指数为 0，不参与 l 和 acc。

5. Scale、bias 和 mask 的实际顺序

真实 score tile 通常经历：

Q_tile @ K_tileᵀ
        ↓
乘以 softmax scale，通常是 1/sqrt(head_dim)
        ↓
加入 attention bias（如果有）
        ↓
应用 causal/padding mask
        ↓
online softmax
        ↓
与 V_tile 相乘并累加

因此教学记号 Sij=Qi@Kjᵀ 更完整地表示为：

Sij = Qi @ Kjᵀ * scale + bias

然后再应用 mask 和 online softmax。

6. Causal mask 还需要 padding mask

教学代码常只写 causal mask：

causal = offs_m[:, None] >= offs_n[None, :]
qk = tl.where(causal, qk, -float("inf"))

实际 tile 可能越过序列尾部，还要处理边界：

valid_m = offs_m[:, None] < N_CTX
valid_n = offs_n[None, :] < N_CTX
causal = offs_m[:, None] >= offs_n[None, :]
mask = valid_m & valid_n & causal
qk = tl.where(mask, qk, -float("inf"))

"当前位置及之前的 V"默认是等长 self-attention；cross-attention 的 Q/K 长度和 mask 关系可能不同。

7. 复杂度与优化目标

项目	朴素 Attention	FlashAttention
数学计算复杂度	$O(N2D)O(N^2D)$ O(N2D)	$O(N2D)O(N^2D)$ O(N2D)
中间显存	$O(N2)O(N^2)$ O(N2)	近似 $O(ND)O(ND)$ O(ND)
完整 S 是否保存	通常保存	不保存
完整 P 是否保存	通常保存	不保存
backward 重计算	较少	较多
主要优化目标	通用计算	减少 HBM/global-memory IO

FlashAttention 主要减少中间结果的显存读写，并没有把矩阵乘法的理论复杂度从 $O(N2D)O(N^2D)$ O(N2D) 变成线性。

8. Arithmetic intensity、tile 和 Tensor Core

FlashAttention 使加载到片上的 tile 被反复复用：

Q tile 被多个 K/V tile 使用
K/V tile 被当前 Q tile 的多个线程复用

因此同一次 global-memory 读取支持更多计算，提高 arithmetic intensity。

tile 太小会造成 global-memory 访问次数多、复用不足；tile 太大会增加 shared memory 和 register 占用，降低 occupancy，甚至造成 register spilling。BLOCK_M、BLOCK_N 和 HEAD_DIM 需要结合 GPU 和输入形状调优。

QKᵀ 和 P@V 都是矩阵乘法，适合映射到 Tensor Core/MMA 指令。因此 tile 设计同时服务于减少 HBM IO 和适配 Tensor Core：

tile shape
warp 分工
shared-memory layout
register 使用
异步加载和流水线
Tensor Core 指令

9. Layout 和扩展功能

真实模型常用 [B,H,N,D]，也有实现使用 [B,N,H,D]。阅读 kernel 时必须确认 batch、head、sequence、head_dim 的位置、stride 和 contiguous 状态。

工程实现还可能支持：

dropout：训练时对注意力概率随机丢弃
attention bias：softmax 前加入 bias
ALiBi：位置相关 bias
GQA/MQA：Q head 数和 K/V head 数不同
variable-length：不同样本有效长度不同

这些功能不改变 FlashAttention 的基本 IO 思想，但会影响 mask、索引、layout、随机数和梯度实现。

10. Triton 教学代码的边界

Triton 官方 06-fused-attention.py 是学习入口，但主分支会随版本变化，可能使用 TensorDescriptor、TMA、warp specialization 或 FP8，部分路径依赖 Hopper 等新硬件。教学版代码用于说明 Q/K/V tile、QKᵀ、online softmax、P_tile@V_tile 和 O 写回，不保证直接复制即可运行。

真实 kernel 还要处理：

边界和 padding
causal mask
attention bias
dtype 转换
layout 和 stride
不同 head_dim
不同 GPU 架构

11. Backward 的实现边界

常见实现会保存 O、LSE 等较小状态，但具体保存哪些张量取决于实现。反向时重新遍历 tile，重新计算局部概率，得到 dQ、dK、dV。当多个 program 对同一 K/V tile 产生梯度贡献时，dK 和 dV 可能使用分块累加、split-K、workspace 或 atomic add。

12. FlashAttention 版本边界

FlashAttention 1、2、3 的数学核心都建立在 IO-aware tiling 和 online softmax 上，主要差异在并行划分、调度、流水线、Tensor Core 使用和硬件特化。FlashAttention 3 主要针对 Hopper 等架构进行专门优化，不能简单理解为所有 GPU 上都无条件适用的 FlashAttention 2 升级版。