第 6 章 · FlashAttention / 6.1
FlashAttention：一次前向计算
读 FlashAttention 的 CuTe 代码时，很容易停在 partition、fragment、retile 上：每一行似乎都认识，连起来却不知道数据去了哪里，当前线程又算出了什么。
先看这个 kernel 最后完成的工作。假设 Q、K、V 各有 128 行，每行 64 个数，一个线程块只负责其中 64 行 query：
固定这 64 行 Q
→ 读取前 64 行 K、V，更新分数、权重和输出
→ 读取后 64 行 K、V，继续更新同一份输出
→ 统一归一化，写回这 64 行结果
FlashAttention 把一次完整 attention 拆成多轮计算，每轮结束只保留能继续累加的状态。 CuTe 再将每轮的加载、矩阵乘法和归约分给具体线程。
下面先用一个能手算的小矩阵走通两轮，再把真实的 64×64 分块拆到 block、warp、lane，最后回到源码。读线程布局时，我们会一直追问同一个问题：线程 0 手里有哪些值，它与哪些线程合作？
本文分析 tiny-flash-attention 面向 SM80（Ampere 架构）的 FP16 前向计算，也就是从 Q、K、V 得到 O 的过程；具体实现与版本见文末 [5]、[6]。前两节只需要矩阵乘法的基础，CuTe 的接口在用到时再解释。
先看 attention 最后要得到什么
Q、K、V 都是把向量按行排成的矩阵。Q 的每行称为 query，用来与 K 中各行的 key 比较；K 的每行又与 V 的同一行配对，提供对应的输出内容。对某一行 query，attention 做三件事：
- 与每行 K 做点积，得到一行分数。
- 将分数经过 softmax，变成一组总和为 1 的权重。
- 用这些权重对 V 的各行求加权和，得到一行输出。
记原始分数为 S、归一化权重为 P、输出为 O，公式是 [1]：
s 是缩放系数，标准 attention 使用 s=1/sqrt(d)，d 是每行 Q、K 的长度。下标 row 表示每行单独做 softmax：query 0 的权重与 query 1 的权重各自归一化。
本文的符号沿用 FlashAttention 原论文 [3]，跨轮累加和末尾归一化采用 FlashAttention-2 的 Algorithm 1 [4]：
| 符号 | 含义 |
|---|---|
| 序列长度、head dimension；完整自注意力的 | |
每块 query 行数、每轮 key/value 行数；对应源码 kBlockM、kBlockN | |
| query 分块编号、K/V 轮次，论文中从 1 开始； 是总轮数 | |
| 当前分数块 ，尚未乘缩放系数 | |
| 处理完第 轮后的逐行累计最大值、指数和 | |
| 当前轮未归一化的指数权重； 专指完整 attention 的归一化权重 | |
| 第 轮后的未归一化输出累加器，以及末尾归一化后的最终输出 | |
| 每行的 log-sum-exp，源码及工程讨论中也称 LSE |
论文算法为简洁省略了缩放系数。本文显式保留 ，并与源码一致，让 、 保存未缩放分数及其最大值，因此指数写为 。若将缩放吸收进论文分数，便有 、，两种写法等价（）。
原版 FlashAttention 每轮维护已归一化输出；本文对应的实现与 FlashAttention-2 一样，用 累加分子，最后才得到 。下文手算固定一个 query 块，省略下标 ；图注用 P̃⁽ʲ⁾、O⁽ʲ⁾、ℓ⁽ʲ⁾ 表示同一组符号，上标括号表示轮次，零次状态表示初始化。
如果先算出完整 S、P，就需要保存“每个 query 与每个 key”之间的结果。序列变长后，这两个矩阵会很大。FlashAttention 的做法是逐块读 K、V，在片上更新结果，避免把完整的 S、P 写入全局显存 [3]。
这里有一个必须解决的问题：第一块 key 算完时，后面的分数还没出现，softmax 的分母怎么确定？ 先用两轮具体数字来看。
用两轮小矩阵把结果算出来
输入：两行 query，四行 key
只展示四行自注意力输入中的前两行 Q，K/V 保留全部四行，每行两个数，每轮读取两个 key（）。输入见图1、图2、图3：
图中的 q、k、d 分别表示 query 编号、key 编号和行内维度，均从 0 开始。格内数字是实际参加运算的值。
为了方便手算，这一节单独取 s=ln(2)，于是 exp(s*x)=2^x，指数权重都能写成简单分数。这是教学用的缩放系数；后面真实 d=64 的标准 attention 应使用 s=1/8。本节先计算允许访问全部 key 的非 causal attention。
第一轮：先算前两个 key 的贡献
Q 与 K 做点积，得到图4。例如 S(0,2)=(1,0)·(2,1)=2。图中完整矩阵用于对照；第一轮实际只计算左侧两列。
query 0 在第一轮的分数为 (0,1)。直接取指数会得到 (1,2)；实现通常先减去这一行的最大值 m=1，再取指数，得到 (1/2,1)。
为什么可以减最大值？固定 query 行 q，设这一行的最大值为 。分子、分母同时乘上相同的系数，最终权重不变：
这样能避免指数值过大。这里先保留未归一化的指数权重，记为 ，与完整 attention 中总和为 1 的 P 区分开。第一轮两行的 max 都是 1，得到图5：
两行的指数和都是 ℓ⁽¹⁾=1+1/2=3/2。再用 P̃⁽¹⁾ 乘本轮 V 的前两行 V₁，计算输出分子 O⁽¹⁾=P̃⁽¹⁾V₁=P̃⁽¹⁾ @ V[0:2,:]：
例如 query 0 的输出分子为 (1/2)×(1,0)+1×(0,1)=(1/2,1)。
此时还不能把 O⁽¹⁾ 当成最终输出：后面两个 key 的贡献尚未加入。每行先保留表1中的三份状态即可。
| 状态 | 含义 | query 0 的当前值 |
|---|---|---|
| 已处理分数的最大值，尚未乘 s | 1 | |
| 相对于 的指数权重之和 | 3/2 | |
| 相对于 的加权和，即输出分子 | (1/2,1) |
这里 、 是每行各一个数， 是与输出维度等长的一行向量。后续每块 K、V 都更新同一份状态。
第二轮：新最大值出现后，先调整旧结果
第二轮读取 key 2、3。由分数矩阵可见，两行的最大值都从 1 变为 2。
第一轮的 P̃⁽¹⁾、ℓ⁽¹⁾、O⁽¹⁾ 都相对于旧最大值 1 计算。第二轮要把它们换到新最大值 2 下，统一乘：
α = exp(s*(旧 m - 新 m)) = 2^(1-2) = 1/2
例如第一轮 query 0 的原始权重 (1/2,1)，换到新基准后变成 (1/4,1/2)。所以旧分母从 3/2 变为 3/4，旧输出分子 O⁽¹⁾ 的每个分量也都乘 1/2。
再以新最大值 2 计算当前两个 key 的指数权重，得到图7：
当前两行的指数和分别是 5/4、3/2。与调整后的旧分母相加：
query 0：ℓ⁽²⁾ = 3/4 + 5/4 = 2
query 1：ℓ⁽²⁾ = 3/4 + 3/2 = 9/4
当前权重乘 V 的后两行，得到新增输出分子 P̃⁽²⁾V₂=P̃⁽²⁾ @ V[2:4,:]：
例如 query 0 是 1×(1,1)+(1/4)×(2,0)=(3/2,1)。加上已经调整的旧分子，O⁽²⁾=(1/2)*O⁽¹⁾+P̃⁽²⁾V₂：
旧分母和旧分子必须乘同一个 α。 如果只调整分母，或者直接把两轮未经调整的 O⁽ʲ⁾ 相加，结果就会使用不同的最大值基准。
所有 key 处理完，再统一归一化
现在四个 key 的贡献都已加入，逐行计算 O=O⁽²⁾/ℓ⁽²⁾。分别用 2、9/4 去除上一图的两行，得到最终输出：
例如 O(1,0)=3/(9/4)=4/3。两轮都没有单独归一化后再拼接；它们通过 m、ℓ、O⁽ʲ⁾ 合并，最后只除一次完整分母。
将刚才的过程写成通用更新
这种边读取、边更新 softmax 状态的方法叫 online softmax [2]。恢复论文的分块下标：固定 query 块 ，第 轮读取 ，维护行向量 和矩阵 。初始化为 、、。
先计算当前分数、新的累计最大值和调整系数：
再计算当前轮的未归一化权重，并更新分母与输出累加器：
这里 是逐元素乘法，减去行最大值时沿列广播； 使旧输出的每一行乘对应的调整系数。第一轮从当前分块建立状态，等价于取 ，无需缩放旧结果。
处理完全部 轮后，才得到最终输出和 log-sum-exp：
这正是前面手算的 。 始终是尚未乘 的原始分数最大值。
接下来保持这个计算过程不变，看看真实 kernel 怎样分配线程。
一个 64×64 分块，怎样分给四个 warp
从这里开始使用实际 kernel 的配置，只看一个 batch、一个 head：
Q、K、V：[128,64]，按行连续存储
序列长度 N=128，head dimension d=64
B_r=64（kBlockM）：一个 block 负责 64 行 query
B_c=64（kBlockN）：每轮读取 64 行 key、value
数学轮次 对应源码从 0 开始的 K/V 循环下标 。下面保留线程、lane、query/key 行号及数组切片的 0 起始编号；提到“第几轮”时仍按论文从 1 开始。源码标识符保持原名。
CUDA 的 block 是一个线程块。一个 warp 含 32 个线程，warp 内的线程编号叫 lane。本例每个 block 启动 4 个 warp，共 4×32=128 个线程，线程编号 threadIdx.x=0～127。
这个配置与前面的手算例子用途不同：小矩阵解释数值更新，本节的 64、128 用来说明真实线程分工。以下全部线程坐标都按真实配置计算。
先分给 block：各自固定 64 行 query
同一个 head 使用两个 block，分别负责 query 0～63、64～127。对第 0 个 block：
固定 Q[0:64,:]
第 1 轮：使用 K[0:64,:]、V[0:64,:]
第 2 轮：使用 K[64:128,:]、V[64:128,:]
最后写回 O[0:64,:]
图11中每格代表一个 64×64 分数块，高亮为第 0 个 block 的工作范围。Q₁、K₁ 表示前 64 行，Q₂、K₂ 表示后 64 行：
两个 block 各自遍历所需 K/V，只更新自己的 query 行。图示用于表示计算范围，完整分数矩阵不会在全局显存中分配。
再分给 warp：各自负责 16 行 query
一个 block 当前算的是 Q_tile[64,64] @ K_tile.T[64,64]。四个 warp 沿 query 行分工，每个得到一个 16×64 分数条带，并覆盖本轮全部 64 个 key。源码迭代下标为 t（数学上第 j=t+1 轮）时的范围见表2：
| warp | threadIdx.x | Q / S / O 的 query 行 | 当前 S 的 key 列 | 最终 O 的输出维度 |
|---|---|---|---|---|
| 0 | 0～31 | 0～15 | 64*t～64*t+63 | 0～63 |
| 1 | 32～63 | 16～31 | 64*t～64*t+63 | 0～63 |
| 2 | 64～95 | 32～47 | 64*t～64*t+63 | 0～63 |
| 3 | 96～127 | 48～63 | 64*t～64*t+63 | 0～63 |
第 1 个 block 的 query 行整体再加 64。每个 warp 会计算自己这 16 行的分数、max、分母和输出分子；不同 warp 负责不同 query，所以不需要再把它们的部分输出相加。
不过，“warp 负责哪些计算结果”还没有说明“一个线程持有哪些值”。要回答后一个问题，先看一条 MMA 的大小。
一个 warp 怎样算完 16×64 的分数
先算 16×8，再沿两个方向重复
Tensor Core 的 MMA 指令执行小块矩阵乘加。本例使用 m16n8k16，含义是：用长度为 16 的点积，更新一个 16 行、8 列的输出子块。
沿用 GEMM 的轴名，M 表示输出行、N 表示输出列、K 表示相乘后求和的归约维度。在 QK 中，三者分别是 query、key、head dimension。这里的“归约维度 K”与输入矩阵 K 要根据上下文区分。
warp 0 的一个 16×8 子块，例如 query 0～15、key 0～7，需要把 d=64 分成四段：
S_sub = Q[:, 0:16] @ K[:, 0:16].T
+ Q[:,16:32] @ K[:,16:32].T
+ Q[:,32:48] @ K[:,32:48].T
+ Q[:,48:64] @ K[:,48:64].T
这里 Q、K 都只取上述子块对应的行。四次 MMA 累加到同一个 16×8 结果中，才完成长度 64 的点积。沿 key 再处理 0～7、8～15、…、56～63 八组，就覆盖了整个 16×64 条带。
所以，一轮 QK 每个 warp 执行 8×4=32 条 MMA，四个 warp 合计 4×32=128 条。每条指令由一个 warp 的 32 个 lane 共同执行；计数不再乘 32，也不表示执行周期。
32 个 lane 共同提供一条 MMA 的输入
CuTe 的 GEMM 使用 A(M,K)、B(N,K)、C(M,N) 约定。因此 QK 可以直接把 K(key,d) 当作 B，表达数学上的 Q @ K.T，无需事先生成 K 的转置副本。
一条 m16n8k16 的 A 为 16×16、B 为 8×16、C 为 16×8。32 个 lane 各提供 8 个 FP16 A 元素、4 个 FP16 B 元素，并各持有 4 个 FP32 累加器。FP16 是 16 位浮点数，也叫 half；FP32 是 32 位浮点数。
先看 warp 0、key 0～7、d 0～15 的这一次 MMA，lane 0 持有：
Q 输入：query 0、8，每行取 d=0、1、8、9，共 8 个 half
K 输入：key 0，取 d=0、1、8、9，共 4 个 half
S 累加：S(0,0)、S(0,1)、S(8,0)、S(8,1)，共 4 个 float
lane 1 提供相同 Q 行在 d=2、3、10、11 的输入，提供 K 行 0 的同四个 d，并持有 query 0、8 的 key 2、3 分数。lane 4 则持有 query 1、9 的 key 0、1 分数。
一个 lane 提供的输入不足以独立完成整段点积；MMA 使用整个 warp 的操作数共同完成计算。遍历完全部八组 key 后，每个线程持有两行、每行 16 个分数，共 32 个 FP32。
将具体坐标推广到任意 lane
令 w=threadIdx.x/32、lane=threadIdx.x%32，再令 g=lane/4、a=lane%4。用 j 标记八组 key，用 h 标记四段 d，一次 MMA 的线程坐标为：
j=0,…,7；h=0,…,3；以下均为当前 tile 内坐标
A（Q）：query = 16*w + {g, g+8}
d = 16*h + {2*a, 2*a+1, 2*a+8, 2*a+9}
B（K）：key = 8*j + g
d = 16*h + {2*a, 2*a+1, 2*a+8, 2*a+9}
C（S）：query = 16*w + {g, g+8}
key = 8*j + {2*a, 2*a+1}
花括号表示取列出的所有坐标组合。例如 A 中两行各取四个 d，所以是八个值。源码迭代 t 的全局 key 行还要加 64*t，第 m_block 个 Q block 的全局 query 行还要加 64*m_block；文中的除法均为整数除法。
再看 TiledMMA 参数，就能对上这些范围
kernel_traits.h 的配置在本例中展开为：
using TiledMma = TiledMMA<
MMA_Atom<SM80_16x8x16_F32F16F16F32_TN>,
Layout<Shape<_4, _1, _1>>,
Tile<_64, _16, _16>>;
auto thr_mma = mma.get_slice(threadIdx.x);
第一项选择刚才的 MMA 指令；第二项把四个 warp 排在 M 方向；第三项指定 TiledMMA 的基本覆盖范围 64×16×16。其中 N=16 由两次 N=8 的 atom 覆盖，再沿 N 重复四次、沿 K 重复四次，完成本轮 QK。各层范围见表3：
| 层次 | 本例范围 | 参与线程与重复方式 |
|---|---|---|
| 一轮 QK | 64×64×64 | block 的 4 个 warp，共 128 线程 |
| TiledMMA 基本覆盖 | 64×16×16 | 4 个 warp 沿 M 排列；覆盖完整 tile 需沿 N 重复 4 次、沿 K 重复 4 次 |
| 一个 warp 的完整计算 | 16×64×64 | 32 线程；8 个 N 子块，每个累加 4 段 K |
一条 m16n8k16 | 16×8×16 | 一个 warp 的 32 个 lane 共同执行，累加到一个 16×8 输出子块 |
get_slice(threadIdx.x) 取得当前线程在这套分工中的位置。QK 完成后，线程持有刚才算出的 32 个分数；若开启 causal，先 mask，再求 max。
QK 之后，causal mask 怎样落到线程寄存器
前面的手算允许每个 query 访问全部 key。开启 causal 后，只允许 key<=query，也就是当前位置只能看自己和过去。实现分两层处理：先少遍历整块 K/V，再修改对角块内的分数。
先跳过整块位于未来的 K/V
对当前 64 行 query，源码计算：
int seqlen_start = m_block * 64;
int seqlen_end = (m_block + 1) * 64;
int n_block_max = Is_causal
? cute::ceil_div(seqlen_end, 64)
: cute::ceil_div(params.k_seqlen, 64);
nbi 从 0 遍历到 n_block_max-1。因此 query 0～63 的 block 只读取 key 0～63；query 64～127 的 block 读取 key 0～127。未来整块不进入循环，也就省掉它的 K/V 加载、QK 和 PV。
对于已经进入循环的块，本例 B_r=B_c=64（源码 kBlockM=kBlockN=64），nbi<m_block 时全部 key 都在当前 query 之前，可以直接使用。只有 nbi=m_block 的对角块既有有效位置，也有未来位置，需要逐元素判断。
再把线程自己的分数对应回全局坐标
QK 的结果已经在 tCrS 中。源码先为它建立行列视图：
Tensor tCrS_rowcol = make_tensor(
tCrS.data(), flash::convert_layout_acc_rowcol(tCrS.layout()));
这个视图仍只包含当前线程的 32 个 FP32。对本例，CuTe 打印的 layout 是：
shape = ((2,1),(2,8))
stride = ((2,0),(1,4))
行坐标 (mi,mj) 的范围是 (2,1)：mi 选当前线程负责的两行，mj 在本例只能为 0。列坐标 (b,nj) 的范围是 (2,8)：nj 选八个 N 子块，b 选其中相邻的两列。寄存器槽位为 2*mi+b+4*nj；这解释了为什么“看作两行、每行 16 个数”并不等于普通连续行存储。
沿用 w=tid/32、lane=tid%32，每个槽位对应的全局坐标是：
q = 64*m_block + 16*w + lane/4 + 8*mi
k = 64*nbi + 2*(lane%4) + b + 8*nj
mi=0,1；b=0,1；nj=0,…,7
这里 q、k 才是用于比较的 query、key 编号，不能拿本线程的行列下标直接比较。mask_within_nblock 的核心判断相当于：
if (k > q) {
tCrS_rowcol(make_coord(mi, 0), make_coord(b, nj)) = -INFINITY;
}
这段代码按当前 64 行配置简化；源码还保留了处理更多 M 重复块的 mj 循环。make_coord 保留行、列内部的层级，让坐标与 layout 的结构一致。
具体看 lane 0～3 各自保留哪些分数
取第 0 个 block 的第 1 轮、warp 0。lane 0～3 共同持有 query 0、8 的分数，mask 后的分工见表4：
| lane | 原来持有的 key 列 | query 0 保留 | query 8 保留 |
|---|---|---|---|
| 0 | 0,1,8,9,…,56,57 | 0 | 0、1、8 |
| 1 | 2,3,10,11,…,58,59 | 无 | 2、3 |
| 2 | 4,5,12,13,…,60,61 | 无 | 4、5 |
| 3 | 6,7,14,15,…,62,63 | 无 | 6、7 |
合起来，query 0 恰好只保留 key 0，query 8 恰好保留 key 0～8。每个线程只改写自己的分数，不需要 shuffle 或 shared memory；之后才通过四线程归约合并有效分数的 max。
为什么填负无穷？因为 exp(s*(-∞-m))=0，这些位置既不贡献分母，也不贡献输出。若填 0，它仍会产生正权重；有效分数全为负数时，还会错误地抬高 max。因此 mask 必须在 max 和指数变换之前执行。
给前面的手算矩阵开启同一规则，会得到图12。这张图单独展示 causal 分支，前面算出的非 causal 输出不再适用。
这里两个 query 都不能访问 key 2、3，因此第二个小 K/V tile 可以跳过；第一块内的 (query=0,key=1) 仍需单独 mask。本文实现采用左上对齐的 k<=q，query、key 长度不同时，需要留意它与右下对齐规则的区别。
对角块是在完整 QK 计算之后做 mask，块内被遮住的分数仍然执行过 MMA；省掉整块计算的是前面的循环范围控制。
同一行的 64 个分数，怎样在线程间求 max
手算时，我们直接对一行取 max。现在同一行的 64 个分数分散在四个线程中，必须先比较各线程自己的值，再把四份结果合并。这里把多个值合成一个 max 的过程叫归约。
先看完整过程，再展开线程编号：
每线程持有一行的 16 个分数
→ 在线程内取 max，得到一个局部结果
→ 与同组另一个 lane 交换，覆盖 32 列
→ 再交换一次，覆盖全部 64 列
→ 同组四个 lane 都保留完整的行 max
先找到负责同一行的四个线程
warp 0 的 lane 0～3 都持有 query 0、8 的分数，只是持有的 key 列不同。这四个线程构成一组；lane 4～7 是下一组，负责 query 1、9。八个组完整列在表5：
| 四线程组 g | lane / threadIdx.x | mi=0 对应 query | mi=1 对应 query | 每行的列覆盖 |
|---|---|---|---|---|
| 0 | 0～3 | 0 | 8 | 4 线程各 16 列，合计 64 列 |
| 1 | 4～7 | 1 | 9 | 同上 |
| 2 | 8～11 | 2 | 10 | 同上 |
| 3 | 12～15 | 3 | 11 | 同上 |
| 4 | 16～19 | 4 | 12 | 同上 |
| 5 | 20～23 | 5 | 13 | 同上 |
| 6 | 24～27 | 6 | 14 | 同上 |
| 7 | 28～31 | 7 | 15 | 同上 |
mi=0,1 是每个线程自己那两行的编号：在线程 0 中对应 query 0、8，在线程 4 中对应 query 1、9。两行分别求 max，始终互不比较。
其他 warp 沿用相同分工，只需给 query 行加上 warp 的行起点。例如线程 32～35 负责 query 16、24；最后一组线程 124～127 负责 query 55、63。整个 block 有 4×8=32 组，每组求两行的 max，共 64 个不同的行最大值。
每个 lane 先比较自己持有的 16 列
warp 0 第一组的列分工见表6。同一 lane 对 query 0、8 使用相同的列集合，但两行分别存储：
| lane | key 0～31 中持有的列 | key 32～63 中持有的列 |
|---|---|---|
| 0 | 0、1、8、9、16、17、24、25 | 32、33、40、41、48、49、56、57 |
| 1 | 2、3、10、11、18、19、26、27 | 34、35、42、43、50、51、58、59 |
| 2 | 4、5、12、13、20、21、28、29 | 36、37、44、45、52、53、60、61 |
| 3 | 6、7、14、15、22、23、30、31 | 38、39、46、47、54、55、62、63 |
每列恰好由一个 lane 持有，四个 lane 合起来覆盖一行的 64 列。每个 lane 的 16 列跨越八个 MMA 的 N 子块，每个子块贡献两列。
线程 0 因而先算：
query 0 的局部 max = max(S(0,0), S(0,1), S(0,8), S(0,9), …, S(0,56), S(0,57))
query 8 的局部 max = max(S(8,0), S(8,1), S(8,8), S(8,9), …, S(8,56), S(8,57))
对每行，以第一个分数初始化，再做 15 次比较。其他线程也比较各自那 16 列；这一阶段不读取其他线程的数据。
两次 XOR shuffle：从 16 列合并到 64 列
shuffle 是 warp 内读取另一 lane 寄存器值的指令。__shfl_xor_sync(mask,x,offset) 让当前 lane 读取编号为 lane XOR offset 的伙伴的 x。
令 x0、x1、x2、x3 是 lane 0～3 对 query 0 的局部 max，合并过程为：
开始：各覆盖 16 列
lane 0：x0 lane 1：x1
lane 2：x2 lane 3：x3
XOR 2：0↔2、1↔3，各覆盖 32 列
lane 0：max(x0,x2) lane 1：max(x1,x3)
lane 2：max(x2,x0) lane 3：max(x3,x1)
XOR 1：0↔1、2↔3，各覆盖全部 64 列
lane 0：max(max(x0,x2), max(x1,x3))
lane 1：max(max(x1,x3), max(x0,x2))
lane 2：max(max(x2,x0), max(x3,x1))
lane 3：max(max(x3,x1), max(x2,x0))
第二次读取的是伙伴在第一步已经合并的结果，因此每个 lane 最终都得到了四份列集合的 max。每次只传一个 FP32 标量，原来的分数仍保存在各线程中。
把数字代进去。单独构造一行 64 列的分数：前八列为 (-4,-3,-2,-1,0,1,2,3)，其余均为 -5。按前面的列分工，每个 lane 的结果见表7：
例如 lane 0 先得到 -3，再读取 lane 2 的 1，变成 1；第二步读取 lane 1 已经合并出的 3，最终得到 3。这一行专门演示真实 64 列布局下的归约通信，与前面用于手算 attention 的小矩阵分别使用。
query 8 在同样的四个线程中另做一次相同归约。与此同时，组 4～7 处理 query 1、9，其余组依次处理各自的行。
为什么不会串到另一组？写成 lane=4*g+a 后，XOR 2、1 只改变组内编号 a 的低两位，组号 g 保持不变。例如 lane 4 的伙伴是 6、5；lane 31 的伙伴是 29、30。shuffle 的通信范围本身又限定在当前 warp 中。
再看源码中的行视图与归约函数
源码用 convert_layout_acc_rowcol 给当前线程的 32 个分数建立两行、每行 16 列的视图 tCrS_rowcol，复用原来的寄存器：
以 lane 0 为例：
tCrS_rowcol(0,:) = S(query 0, [0,1,8,9,…,56,57])
tCrS_rowcol(1,:) = S(query 8, [0,1,8,9,…,56,57])
rRowMax(0) 保存 query 0 的 max
rRowMax(1) 保存 query 8 的 max
这里 ni=0,…,15 是当前线程的列下标，不能直接当成 key 编号。比如 lane 0 的 ni=2 对应 key 8，lane 1 的 ni=2 对应 key 10。
归约实现 正是先做线程内比较，再做四线程交换：
reduce_max<zero_init>(tCrS_rowcol, rRowMax)
└─ reduce_(..., MaxOp<float>)
├─ thread_reduce_<zero_init>：每行遍历本线程的 16 个分数
└─ quad_allreduce_：对两行分别调用 Allreduce<4>
└─ XOR 2，再 XOR 1
MaxOp<float> 比较两个 FP32 数并保留较大者。首轮在线程内用本行第一个分数初始化；四线程部分的执行效果可以写成：
// 所有 lane 都执行；每个线程的两行分别归约。
for (int mi = 0; mi < 2; ++mi) {
float x = rRowMax(mi);
x = max(x, __shfl_xor_sync(0xffffffffu, x, 2));
x = max(x, __shfl_xor_sync(0xffffffffu, x, 1));
rRowMax(mi) = x;
}
0xffffffffu 是整个 warp 的参与掩码，源码在所有 lane 都执行归约时使用它。所有 lane 参加这条指令，而 XOR 的伙伴仍只在各自四线程组中；参与掩码与数据合并范围各有作用。
结果为什么保留四份，下一轮又怎样继续
这是 all-reduce：同组四个 lane 都得到同样的行 max。线程 0～3 各自持有 query 0、8 的两个 max；全 block 用 128×2=256 个逻辑 FP32 槽位保存 64 个不同的行最大值。
这让每个线程都能直接用自己的 m，对自己那部分分数计算 exp(s*(S-m))，不必再等待某一个线程广播。
第一轮使用 zero_init=true，意思是从第一个分数开始，忽略旧 summary，并非把 max 初始化成 0，因此全负分数也能正确归约。后续轮使用 false，每个线程先算 max(旧 m, 自己当前的 16 个分数)，再做两次 shuffle。最终得到的是 max(旧 m, 当前全部 64 列)。
旧 m 已在四个 lane 中一致，重复参与不会改变 max。softmax_rescale_o 另存一份 rRowMax_prev，用旧、新 m 计算 α，缩放各线程持有的那部分旧输出分子。
当前布局下，一行的所有 key 已在同一个四线程组中，所以 max 不需要 shared memory、原子操作或跨 warp 汇总。reduce_sum 复用相同的线程内、四线程结构，把 max 换成加法，在指数变换后得到每行的 ℓ。其他 MMA 布局若改变了同一行的持有者，需要重新检查这套分组。
max 之后，online softmax 怎样更新寄存器状态
现在每个四线程组都知道自己两行的最大值。接下来要将前面手算的递推落实到寄存器：保留历史 m、ℓ、O，当前分数则原地变成 P̃。本节描述寄存器的就地更新，省略分块下标与轮次上标：循环中的 O 均指 对应的未归一化累加器，循环结束除以 ℓ 后才成为最终 。这里的 online 是“每收到一块分数，就更新已有状态”，不需要等完整一行分数到齐 [2]。
先分清哪些状态复制了，哪些元素分散存放
对 lane 0 来说，rRowMax、rRowSum 各有两个 FP32，分别对应 query 0、8。lane 1～3 也保留这两行的相同 m、ℓ；但四个线程各持有不同的 S/P̃ 列和 O 输出维度。每线程的存储为：
rRowMax：2 个 float，每行一个 m，同组四个 lane 各保留一份
rRowSum：2 个 float，每行一个 ℓ，同组四个 lane 各保留一份
tCrS ：32 个 float，两行各 16 个分数，随后原地改成 P̃
tCrO ：32 个 float，两行各 16 个输出分子元素 O
同一行的 O 是分散存放的：例如 lane 0 负责输出维度 0,1,8,9,…,56,57。因此“缩放旧 O”意味着四个线程各自缩放自己的 16 个元素，合起来覆盖整行。
第一轮：从当前分数建立状态
进入 K/V 循环前，tCrO 已被清零。第一轮没有旧贡献，前三步由 softmax_rescale_o<true> 完成，随后主循环执行 PV：
1. reduce_max<true>：16 个本地分数比较 → XOR 2 → XOR 1，得到 m
2. 每线程逐元素计算 P̃ = exp(s*S - s*m)，覆盖自己的 tCrS
3. reduce_sum：16 个本地 P̃ 求和 → XOR 2 → XOR 1，得到 ℓ
4. 返回主循环，将 P̃ 转为 FP16，执行 PV：O += P̃ @ V
第 2 步完全在线程内进行。同组 lane 已通过 max 的 all-reduce 得到相同的 m，所以它们计算的 P̃ 使用同一基准。mask 为负无穷的位置在这里变成 0。
源码函数虽名为 scale_apply_exp2，本路径实际调用 expf，传入的也是 params.softmax_scale。因此按 exp(s*(S-m)) 理解即可，不需要额外乘 log2(e)。
后续轮：先换基准，再加当前贡献
每行已有历史状态后，softmax_rescale_o 的后续轮分支（<false>） 按表8更新。表中的“每行”均指当前线程持有的两行分别执行。
| 顺序 | 更新内容 | 当前线程做什么 | 是否需要线程通信 |
|---|---|---|---|
| 1 | 保存旧 m | 把两个 rRowMax 复制到 rRowMax_prev | 否 |
| 2 | 得到新 m | 比较旧 m 与本线程新分数，再合并四个 lane | 两次 XOR max |
| 3 | 计算 α | 每行计算 exp(s*(旧 m-新 m)) | 否 |
| 4 | 调整旧 ℓ、O | 每行一个旧 ℓ 和自己的 16 个 O 都乘 α | 否 |
| 5 | 当前 S 改成 P̃ | 32 个分数逐个变为 exp(s*(S-新 m)) | 否 |
| 6 | 合并当前分母 | 本地求和，再归约得到当前块的 Δℓ | 两次 XOR sum |
| 7 | 更新 ℓ | 已调整的旧 ℓ 加 Δℓ | 否 |
| 8 | 更新 O | 函数外将 P̃ 转为 half，PV 累加 P̃ @ V | warp 共同执行 MMA |
α 为什么正好能调整旧结果？旧权重为 exp(s*(旧 S-旧 m))，乘上 exp(s*(旧 m-新 m)) 后，旧 m 抵消，得到 exp(s*(旧 S-新 m))。因此旧权重和 ℓ、旧加权和 O 都乘同一个 α，就统一到了新 m 的基准。新 m 没变时 α=1，旧贡献自然保持原值。
例如前面的第二轮，新 m 从 1 变为 2，s=ln(2)，两个行状态各算出 α=1/2。query 0 的 ℓ 先从 3/2 缩到 3/4，再加当前块的 5/4 得到 2；它的 O 则先缩成 (1/4,1/2)，等 PV 加入 (3/2,1) 后，才得到 (7/4,3/2)。
因此，softmax_rescale_o 返回时，m、ℓ 已更新，tCrS 已变成 P̃，tCrO 只完成了旧 O 的缩放；必须再经过 PV，当前轮的 O 才更新完整。
sum 同样归约四个 lane，为何历史 ℓ 要单独相加
sum 与 max 使用同一份行分工。每个线程先把本行的 16 个 P̃ 相加，再通过 XOR 2、XOR 1 做加法 all-reduce。举例，四个 lane 的局部和若为 (1/2,1,1/4,1/4)，两步结果是：
本地和： (1/2, 1, 1/4, 1/4)
XOR 2： (3/4, 5/4, 3/4, 5/4)
XOR 1： (2, 2, 2, 2 )
最后各线程都得到当前块的 Δℓ=2。源码将它存到新的 rRowSum_cur，然后才执行 rRowSum += rRowSum_cur。
这一步不能照搬 max 的历史合并方式。旧 m 在四个 lane 中重复出现不影响 max；旧 ℓ 若先放进每个 lane 的局部和，再参与 sum，就会被加四遍。正确顺序是：只归约当前块的 P̃，再让每个 lane 在自己的旧 ℓ 副本上加一次结果。
同理，最终统一归一化时，每线程直接用自己的行 ℓ 除自己那 16 个 O 元素；不需要重新汇总整行 O。
权重乘 V 时，线程分工怎样延续
求完 m、更新 ℓ 后，分数的寄存器存储已原地变成当前块的 P̃。现在执行手算中的 P̃ @ V，将新增贡献累加到 O。
源码把这一步称为 PV，输入变量也使用 P 的名字，但此时权重尚未除最终分母。四个 warp 仍各负责原来的 16 行 query，MMA 的三条轴则换成了表9中的含义：
| 项目 | QK | PV |
|---|---|---|
| A(M,K) | Q(query,d) | 未归一化权重 P̃(query,key) |
| B(N,K) | K(key,d) | sVt(d,key) |
| C(M,N) | S(query,key) | 输出分子 O(query,d) |
| 每个 warp 的输出 | 16×64 个分数 | 16×64 个输出分子元素 |
| 4 段 K 的含义 | 每段 16 个 head 维度 | 每段 16 个 key |
这里的 B 视图是 sVt(d,key)：同一份 V 在 shared memory 中按输出维度、key 访问。一个 warp 将 64 个输出维度分成八组，每组八列，再沿本轮 64 个 key 分四段累加。
所以 PV 也由每个 warp 执行 32 条 MMA，整个 block 128 条。一次 K/V 迭代中 QK、PV 各 128 条，共 256 条；非 causal 的两个迭代共 512 条。
为什么 QK 的权重可以在同一线程中继续使用
仍看 warp 0 的 lane 0。它持有 query 0、8 在 key 0、1、8、9、…、56、57 上的权重。PV 每段归约需要的 A 输入，恰好能从这些值中取出：
第 0 段 key： 0、 1、 8、 9，每行 4 个值，两行共 8 个 half
第 1 段 key：16、17、24、25
第 2 段 key：32、33、40、41
第 3 段 key：48、49、56、57
源码先将 FP32 权重转成 FP16，再把它们组织为 MMA 的 A 视图：
Tensor rP = flash::convert_type_f32_to_f16(tCrS);
Tensor tCrP = make_tensor(
rP.data(),
flash::convert_layout_rowcol_Aregs<TiledMma>(
tCrS_rowcol.layout()));
第一行发生数值转换；第二行复用 rP 的存储，在同一线程内重组访问方式。32 个 half 正好组成四段各八个值的 A 输入，无需把权重先搬到别的线程。
gemm_A_in_regs 加载 V 的 B fragment，并把乘积累加到 tCrO。PV 完成后每线程仍持有两行、每行 16 个输出维度，共 32 个 FP32；下一轮调整旧 O，最后除以 ℓ，都直接作用于这些寄存器。
到这里，我们已经知道一个 block 算哪些行、每条 MMA 算多大、每个线程持有哪些值。剩下的问题是：源码怎样把输入送到这些线程手中？
回到 CuTe：数据怎样送到计算线程手里
全局显存保存完整输入；shared memory 是同一个 block 的线程都能访问的共享内存；寄存器保存当前线程自己的 fragment，也就是它持有的那一小份数据。输入的数据流为：
全局显存 ──cp.async──→ shared memory ──ldmatrix──→ 线程寄存器 ──→ MMA
cp.async 异步搬运输入块，ldmatrix 再按矩阵指令的需要将 shared memory 数据加载并分发到 warp 内各 lane。先看各变量所在的位置，见表10：
| 变量 | 对应前文的什么 | 存储或视图范围 |
|---|---|---|
mQ | 当前 head 的完整 Q | 全局显存视图 |
gQ | 当前 block 使用的 Q 分块 | 全局显存中的 64×64 tile |
sQ | 协作加载后的 Q 分块 | block 共享的 shared memory |
tCrQ、tCrK | 当前线程提供的 MMA 输入 | 线程自己的寄存器 fragment |
tCrS | 的一部分，随后原地变成 | 每线程 32 个 FP32 |
tCrP | 转成 FP16 后的 MMA A 输入视图 | 与 rP 共享寄存器存储 |
rRowMax、rRowMax_prev | 当前轮 、上一轮 | 每线程各两行的 max，同组四个 lane 各一份 |
rRowSum | 累计分母 | 每线程两行，同组四个 lane 各一份 |
tCrO | 输出累加器 的一部分 | 每线程 32 个 FP32，最后逐行除以 ℓ 得到 |
rLSE | 最终 的一部分 | 每行一个 FP32，写回 gLSE |
local_tile：选出这个 block 的 Q
先用 make_tensor 把当前 head 的指针包装成 Tensor：
Tensor mQ = make_tensor(
make_gmem_ptr(q_ptr + bs_head_offset),
make_shape(params.q_seqlen, Int<kHeadDim>{}),
make_stride(Int<kHeadDim>{}, Int<1>{}));
这里省略指针类型转换。mQ 描述整个 [128,64] Q；K、V、O 也各自建立视图。再取得当前 Q 分块：
Tensor gQ = local_tile(
mQ,
make_tile(Int<kBlockM>{}, Int<kHeadDim>{}),
make_coord(m_block, _));
m_block=0,1 分别对应前、后 64 行 query。这两步创建地址和访问范围的视图，还没有读取数据。
sA 从哪里来：本例的 sQ 如何建立
通用 GEMM 示例常把 shared memory 中的 A 操作数写成 sA。这份 attention 源码按数据含义命名：QK 的 shared A 就是 sQ，shared B 是 sK；到了 PV，A 则是已经在寄存器中的 tCrP。先把名字对应起来，就能区分“存在哪”和“扮演哪个操作数”。
sQ 需要两样东西：一段属于当前 block 的 shared memory，以及把 (query,d) 坐标映射到这段内存的 layout。kernel 先取得启动时预留的动态共享内存：
extern __shared__ char smem_[];
using SharedStorage = SharedStorage<Element, SmemLayoutQ,
SmemLayoutK, SmemLayoutV>;
SharedStorage &shared_storage =
*reinterpret_cast<SharedStorage *>(smem_);
Tensor sQ = make_tensor(
make_smem_ptr(shared_storage.smem_q.data()), SmemLayoutQ{});
SharedStorage 中的 smem_q 是长度为 cosize_v<SmemLayoutQ> 的对齐数组。cosize 表示 layout 所需的地址跨度，保证分配的空间容得下所有映射地址。这里 Q 的空间是 4096 half=8192 字节，K、V 也各有同样大小的独立空间。
make_smem_ptr 标明这是 shared memory 指针；make_tensor 将指针与 layout 组合成视图。此时 sQ 还没有 Q 的数据，后面的 cp.async 才把 gQ 的值填进去。gQ 与 sQ 指向不同存储，不能把建立 sQ 理解成给 gQ 改名字。
SmemLayoutQ：先排成行，再做 swizzle
d=64 时，traits 中与 Q 相关的 layout 可以具体写成：
using SmemLayoutAtomQ = decltype(composition(
Swizzle<3,3,3>{},
Layout<Shape<_8,_64>, Stride<_64,_1>>{}));
using SmemLayoutQ = decltype(tile_to_shape(
SmemLayoutAtomQ{}, Shape<_64,_64>{}));
先看里面的普通 layout：Shape<8,64> 表示八行、每行 64 个元素，Stride<64,1> 让相邻行间隔 64 个 half、相邻列间隔一个 half。它把 (r,c) 变成线性元素偏移 p=64*r+c。
composition 再让 Swizzle<3,3,3> 变换这个偏移。第一个 3 表示处理三个位，第二个 3 表示目标位段从第 3 位开始，第三个 3 表示源位段比目标高三位。位编号从 0 开始，实际是将 p 的第 6～8 位异或到第 3～5 位。XOR 就是按位异或：两位不同得到 1，相同得到 0。公式为：
swizzle(p) = p XOR ((p & 0x1c0) >> 3)
最后 tile_to_shape 将这个八行原子铺到 64×64 的逻辑范围。对本例任意 0<=r,c<64，完整公式可写成：
元素偏移 = 64*r + (c XOR (8*(r%8)))
字节地址 = shared_base + 2*元素偏移
逻辑坐标仍然是第 r 行、第 c 列；变化的是它落在 shared memory 的哪个位置。表11列出几个地址，单位均为 half：
(r,c) | 普通行布局 64*r+c | swizzle 后偏移 |
|---|---|---|
(0,0) | 0 | 0 |
(0,8) | 8 | 8 |
(1,0) | 64 | 72 |
(1,8) | 72 | 64 |
(7,56) | 504 | 448 |
(8,0) | 512 | 512 |
例如逻辑行 1 的列 0～7 实际位于偏移 72～79，而列 8～15 位于 64～71。行内八个 half 为一组进行换位，最低三位不变，所以每组仍连续、仍适合 16 字节拷贝。swizzle 用这种排列配合矩阵加载的 shared memory bank 访问；是否发生 bank conflict 还取决于具体指令和访问模式。
因此，sQ 的完整 layout 是“行布局加 swizzle 的组合”，不能只写成普通的 (64,64):(64,1)。Q、K、V 的这一组加载与读取视图都要遵守同一地址映射，才能取到正确数据。
TiledCopy：先由整个 block 合作搬到 shared memory
这里可以直接沿用 make_tiled_copy 的读法：先看每线程的坐标，再看配置。
加载一个 64×64 FP16 tile 时，8 个线程合作搬一行，每线程一次搬连续八个 half，即 16 字节。128 个线程先覆盖 16 行，再沿行方向重复四次：
thread 0 ：行 0、16、32、48，列 0～7
thread 7 ：行 0、16、32、48，列 56～63
thread 8 ：行 1、17、33、49，列 0～7
thread 32：行 4、20、36、52，列 0～7
row = tid/8 + 16*r，r=0,1,2,3
col = 8*(tid%8) + v，v=0,…,7
traits 使用的线程布局为 Shape<16,8>、Stride<8,1>，值布局为 Shape<1,8>，恰好生成这份坐标分工。Q、当前 K、当前 V 都用它加载。
注意加载与计算的行范围：线程 0 搬来的 Q 行 16，随后由 warp 1 计算。输入先放入整个 block 共享的 shared memory，再按 MMA 的分工读取，所以搬某个值的线程不必就是最终使用它的线程。
代码为当前线程选出加载视图，并执行拷贝：
auto thr_copy_qkv = copy_qkv.get_thread_slice(threadIdx.x);
Tensor tQgQ = thr_copy_qkv.partition_S(gQ(_,_,0));
Tensor tQsQ = thr_copy_qkv.partition_D(sQ);
flash::copy(copy_qkv, tQgQ, tQsQ);
partition_S/D 分别建立源、目的视图，copy 才发起加载。项目的 flash::copy 包装了 CuTe 的 copy。比如线程 8 搬 Q 的逻辑行 1、列 0～7，源是 gQ 的这八个数，目的则是 sQ 的物理偏移 72～79；不是把它们写到普通行布局的 64～71。
partition_fragment 与 retile：按 MMA 分工读入寄存器
输入到 shared memory 后，还要送到 MMA 使用的寄存器。以 Q 为例：
Tensor tCrQ = thr_mma.partition_fragment_A(sQ);
Tensor tXsQ = s2r_thr_copy_q.partition_S(sQ);
Tensor tXrQ = s2r_thr_copy_q.retile_D(tCrQ);
第一行准备当前线程的 A fragment；第二行建立 shared memory 加载源；第三行给同一份 tCrQ 寄存器建立 Copy 目的视图。retile_D 本身不搬数据，随后的 copy 才装入值。K 也做相同准备。
Q 使用 make_tiled_copy_A，K 使用 make_tiled_copy_B，本例的 Copy Atom 都是 SM75_U32x4_LDSM_N，对应不转置的矩阵加载。Copy 分工将 shared memory 的数据分发到各 lane，retile_D 再让加载的目的槽位与 MMA fragment 对齐。
V 的加载还要区分两个视图：sVt(d,key) 与 sV(key,d) 指向同一物理地址，前者带着对应的 swizzle；sVtNoSwizzle 只用于 partition_fragment_B 构造 V 的寄存器 fragment。实际加载源 tXsVt 来自带 swizzle 的 sVt，并使用转置的 SM75_U16x8_LDSM_T。名字中的 NoSwizzle 不意味着实际 V 可以按普通列布局直接读取。
这些接口因此都能落到前面讲过的线程坐标上：local_tile 选整个分块，partition 选当前线程的范围，partition_fragment 准备线程自己的存储，retile 为已有存储提供另一种访问视图。
展开 GEMM：从分段加载到 mma.sync
GEMM 是矩阵乘加。前面已经知道一轮 QK 需要每 warp 执行 32 条 MMA；现在把源码中的一次函数调用拆开，看看这 32 条是怎样形成的。需要区分两层：项目的 flash::gemm_smem 负责 shared memory 加载和 K 分段，CuTe 的 cute::gemm 再遍历当前段的输出子块。
先读懂 fragment 的三个维度
partition_fragment 准备的是当前线程在整个 tile 中要用的存储，不是只留一条 MMA 的操作数。表12把嵌套 shape 各顶层 mode 的大小乘开，列出更容易读的三个维度：
| 视图 | 顶层各维大小 | 三个维度的含义 | 元素总数 |
|---|---|---|---|
tCrQ | (8,1,4) | 每 atom 的 A 值、M 重复、K 分段 | 32 half |
tCrK | (4,8,4) | 每 atom 的 B 值、N 重复、K 分段 | 128 half |
tCrS | (4,1,8) | 每 atom 的 C 值、M 重复、N 重复 | 32 float |
tXrQ | (8,1,4) | Copy 的值分组、M 重复、K 分段 | 与 tCrQ 共用存储 |
tXrK | (8,4,4) | Copy 的值分组、N 方向分组、K 分段 | 与 tCrK 共用存储 |
这里 M 重复为 1，因为每个 warp 只承担一组 16 行 query；N 重复为 8，因为 64 个 key 分为八组。K 的四段分别是 d=0～15、16～31、32～47、48～63。
tXrK 的前两维与 tCrK 不同，是因为一次 x4 矩阵加载按八个 half 分组，将两份 MMA B 的值放在一个 Copy 分组里。两者都覆盖 128 half，改变的是访问同一存储的方式，没有增加一份 K。
gemm_smem：每次取一个 K 段，更新八个输出子块
gemm_smem 先建立 tXrQ/tXrK，随后执行的加载与计算可以按本例写成：
// 第 0 段先装入寄存器。
cute::copy(s2r_copy_q, tXsQ(_,_,0), tXrQ(_,_,0));
cute::copy(s2r_copy_k, tXsK(_,_,0), tXrK(_,_,0));
for (int h = 0; h < 4; ++h) {
if (h + 1 < 4) {
cute::copy(s2r_copy_q, tXsQ(_,_,h+1), tXrQ(_,_,h+1));
cute::copy(s2r_copy_k, tXsK(_,_,h+1), tXrK(_,_,h+1));
}
cute::gemm(mma, tCrQ(_,_,h), tCrK(_,_,h), tCrS);
}
(_,_,h) 选出当前线程在这一段的全部输入：Q 为 8×1 个 half，K 为 4×8 个 half。它们与其他 lane 的操作数合起来，提供当前 warp 的 16×16 Q 和 64×16 K 数据，更新 16×64 分数条带。
加载顺序是“先加载段 0；加载段 1、算段 0；加载段 2、算段 1；加载段 3、算段 2；最后算段 3”。这些是 shared memory 到寄存器的矩阵加载；全局显存到 shared memory 的 cp.async 已在外层安排。各 h 有自己的寄存器槽位，不能据此把它理解为两份 shared memory 缓冲区。
tCrS 在每个 K/V tile 开始时只清零一次，四次 cute::gemm 累加到同一份分数。若在 h 循环中清零，最后就只剩 d=48～63 的贡献。
cute::gemm：当前段还要展开八个 N 子块
把 h 固定以后，cute::gemm(mma,A_h,B_h,C) 首先转成带独立输出参数的形式：
cute::gemm(mma, C, A_h, B_h, C); // D 与 C 使用同一存储
也就是 D = A_h @ B_h.T + C。CuTe 根据 Tensor 的 rank 选择重载；本例 A、B 各剩两个 mode，C 有三个 mode，进入 GEMM 的子块遍历。将本例的 M重复=1、N重复=8 代入，执行效果相当于：
for (int j = 0; j < 8; ++j) {
// 每个输入都取出当前 lane 对一条 MMA 提供的那份值。
cute::gemm(mma,
C(_,0,j), A_h(_,0), B_h(_,j), C(_,0,j));
}
这是专门代入当前尺寸后的示意；通用源码还处理更多 M 重复和不同遍历顺序。每个 j 现在只剩一份 atom 的寄存器向量：A 有 8 个 half，B 有 4 个 half，C/D 有 4 个 float。rank 为 1 的重载调用 mma.call，通过 mma_unpack 将这些值交给架构指令：
flash::gemm_smem：4 个 h
→ cute::gemm：每个 h 遍历 8 个 j
→ mma.call → mma_unpack
→ SM80_16x8x16_F32F16F16F32_TN::fma
→ mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32
进入底层 SM80 MMA 实现 前，mma_unpack 将每 lane 的 8 个 A half 按位打包成四个 32 位寄存器、4 个 B half 打包成两个 32 位寄存器，再与四个 FP32 累加器一同传给指令。打包是在满足指令操作数格式，不是把 half 数值转成整数。
每个 lane 都执行这条 warp 级 MMA 指令，整个 warp 共同完成 16×8×16 矩阵乘加，结果仍分散回各 lane 的四个 FP32。由此回到前面的计数：4段×8子块=32条/warp，整个 block 为 128 条。
PV 为什么只需要加载 B
gemm_A_in_regs 使用相同的 h 循环和 cute::gemm。区别是 tCrP 的四段 A 输入已经在寄存器中，函数只需将 sVt 的 B 输入逐段加载到 tCrVt，再执行：
cute::gemm(mma, tCrP(_,_,h), tCrVt(_,_,h), tCrO);
此时 h 沿 key 分段，j 沿输出维度分组。每 warp 同样执行 32 条 MMA，当前轮贡献加到 tCrO 的旧输出分子上。
两个 GEMM 的清零时机尤其要区分：tCrS 每轮 K/V 开始清零；tCrO 只在整个 K/V 循环之前清零。第二轮 PV 要保留已经乘过 α 的旧 O，才能完成 O_new=α*O_old+P̃@V。
将一轮的加载与计算放回时间顺序
前面解释了数据归属，现在看 kernel 怎样安排它们。主流程与“固定 Q、逐块更新 m、ℓ、O、最后归一化”的算法一致 [4]，同时提前加载后面将要使用的数据：
循环前：加载 Q、K[0]，提交异步拷贝组
源码迭代 t（数学上第 j=t+1 轮）：
等待 Q、K[t] 就绪，block 同步
发起 V[t] 的异步加载
计算 Q @ K[t].T
若为 causal 对角块，将 k>q 的分数设为负无穷
等待 V[t] 就绪，block 同步
若有下一轮，发起 K[t+1] 的异步加载
更新 m、ℓ，按 α 调整旧 O
计算 P̃⁽ʲ⁾ @ V[t]，累加当前贡献（V[t] 对应论文 Vⱼ）
循环后：逐行除以 ℓ，写回 O
V 的加载可以与 QK 重叠，下一块 K 的加载可以与当前 softmax、PV 重叠。Q 在多轮中保持不变，留在 shared memory 复用；本实现 Q、K、V 各只有一份 shared memory tile，K、V 在安全的时机覆盖更新。
cp_async_fence 提交拷贝组，cp_async_wait<0> 等待此前提交的组完成，__syncthreads() 让整个 block 安全使用协作加载的数据。同步也保证旧数据使用完后才复用对应空间。
最后的输出为什么也经过 shared memory
计算结束时，O 仍按 MMA 的分工分散在线程寄存器里。源码将 FP32 结果转为 FP16，经 shared memory 调整成适合连续写回的线程分工：
各线程的 FP32 输出
→ FP16 fragment
→ shared memory 的 sO
→ 写回分工下的临时寄存器
→ 全局显存中的 O
sO 复用已不再需要的 sQ 空间。与输入加载类似，计算分工和显存访问分工通过 shared memory 衔接。
kernel 还保存每行的 log-sum-exp，简称 LSE，论文记为 ，对应源码 rLSE：
前面两行手算结果对应 ln(8)、ln(9)，与直接对完整分数求 log(sum(exp(s*S))) 一致。
带着这些结果回到源码
现在可以沿 flash_attention_v2_cutlass_kernel 阅读：先找到当前 Q tile，再看 QK、online softmax、PV 的循环，最后看归一化与写回。遇到布局代码，就把它对应到前面具体的 block、warp 或 lane 坐标。
本文的线程分工针对 SM80、FP16、连续存储、N 为 64 的倍数、d 与模板尺寸一致的教学实现，不代表所有 FlashAttention kernel。接口默认 scale 为 1，复现标准 attention 时需显式传入 1/sqrt(d)；函数名虽然含 exp2，当前源码实际使用 expf。尾块保护、任意 stride、BF16 转换与 stream 支持仍有限制；本文各示例每行至少有一个有效 key，且只讨论 forward。
配套的 布局核对程序 检查加载坐标、shared memory 的 swizzle 地址、MMA 的 A/B/C 覆盖、同线程权重复用、causal mask 的全局坐标、四线程归约与跨轮 online softmax 状态。它在 host 上核对 CuTe 布局及归约模型，不执行 GPU 指令。文中的手算矩阵按精确分数计算，不模拟真实 FP16 舍入；两轮结果可与直接计算完整 attention 相互核对。
需要继续推导布局 API 时，可阅读 make_tiled_copy 和 retile_D 专题。数值图由本文输入推导，论文提供算法依据，具体调度以源码快照为准。
参考资料
[1] Vaswani, A., Shazeer, N., Parmar, N., et al. Attention Is All You Need. 2017. arXiv:1706.03762，§3.2.1：缩放点积 attention 与 1/sqrt(d) 缩放。
[2] Milakov, M., & Gimelshein, N. Online normalizer calculation for softmax. 2018. arXiv:1805.02867：在线维护最大值与指数和的递推。
[3] Dao, T., Fu, D. Y., Ermon, S., Rudra, A., & Ré, C. FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness. 2022. arXiv:2205.14135，Algorithm 1：分块 attention 与 online softmax。
[4] Dao, T. FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning. 2023. arXiv:2307.08691，Algorithm 1：保持未归一化输出分子，最后统一归一化；其线程划分讨论用于扩展阅读。
[5] NVIDIA. CUTLASS / CuTe Quickstart. 当前检出版本的 CuTe 文档。版本固定为 commit 0b55a2f691d69981583568fd9eb69687b1f0de8a，用于核对 Tensor、Layout 与分块操作的语义。
[6] 66RING. tiny-flash-attention，项目仓库。本文分析工作区快照的 flash_attention.cu、kernel_traits.h 及归约辅助文件 utils.h，网页链接提供对应的源码快照。工作区基于 commit dacda43160add8ed05c43b943b40787606a23f2f，其中 flash_attention.cu 含本地修改，不能仅凭该 commit 复现本文；前两文件的 SHA-256 分别为 06f07a11854fb69b7c0a62d4e893d0d4b851e0ae6847b722bdcac224376f0055、fbb88f1a90ee1c6a937e3b5ac5387c0c7be53c04da30ddf08a6cbf1c897fcf63；utils.h 的 SHA-256 为 d4f907265f4a109321b2d686f8accbd8839179b74e177096e36d2c03e7555501。