用 CuTe 读懂 FlashAttention：一次 forward 的完整流程

读 CuTe 写的 FlashAttention，容易被 partition、fragment、retile 这些名字打断。看完几段布局代码，反而忘了当前正在算什么。

先抓住整个 kernel 要做的事：一个线程块固定一块 Q，依次读取 K、V 的各个分块，每轮计算分数、更新 softmax、累加输出，最后归一化并写回。 CuTe 的这些接口，是在为加载和计算准备合适的数据视图。

下面沿着一个 Q 分块走完一次 forward。遇到 CuTe API 时，只解释它在当前步骤中的作用，布局与指令的内部推导留给专题文章。

本文以 tiny-flash-attention 当前工作区的 flash_attention.cu 和 kernel_traits.h 为主线，讨论 SM80、FP16 路径。CuTe 接口对照当前 CUTLASS 检出版本 0b55a2f6。实现与版本范围见文末参考资料 [5]、[6]。

先固定一个例子

只看一个 batch、一个 head，取连续存储的输入：

Q、K、V：[128,64]
序列长度 L = 128
head dimension D = 64

BlockM = 64：一个 block 处理 64 行 query
BlockN = 64：每轮处理 64 行 key、value
一个 block 有 4 个 warp，共 128 个线程

缩放点积 attention 的定义见 Vaswani 等 [1]。本文采用以下记号：

$$S=QK^T,\qquad P=\operatorname{softmax}_{\mathrm{row}}(sS),\qquad O=PV, \qquad s=1/\sqrt D.$$

完整计算的矩阵形状为：

Q[128,64] @ K.T[64,128] → S[128,128]
P[128,128] @ V[128,64]  → O[128,64]

FlashAttention 将 K、V 分块读取，在片上完成分数、softmax 和输出累加，避免把完整的 S、P 写入 global memory [3]。下文“固定 Q、遍历 K/V、最后统一归一化”的阅读顺序也可对照 FlashAttention-2 的算法描述 [4]；具体线程分工与调度以本文采用的教学源码为准。

对负责前 64 行 query 的 block，工作顺序是：

固定 Q[0:64,:]

第 0 轮：读取 K[0:64,:]、V[0:64,:]
         计算当前分数，更新 softmax 状态，累加输出

第 1 轮：读取 K[64:128,:]、V[64:128,:]
         计算当前分数，更新 softmax 状态，继续累加输出

最后：归一化，写出 O[0:64,:]

另一个 block 独立处理 query 64～127。每个 block 都遍历所需的 K、V，但只写自己负责的输出行。

图1把整个分数矩阵粗分成四块，展示线程块负责的范围。这里 Q0、K0 表示前 64 行，Q1、K1 表示后 64 行，每一格代表一个 64×64 的分数块：

Q/K	K0	K1
Q0	S00	S01
Q1	S10	S11

图1. 128×128 分数矩阵的分块示意；每格代表 64×64 子块，高亮为负责 Q0 的线程块的计算范围

图1只表示工作分块，不表示 kernel 会在 global memory 中分配完整的 S。

贯穿后文的小矩阵

为了把每轮的数值算出来，另外取两行 query、四行 key/value，head dimension 为 2，每轮处理两个 key。图2、图3、图4以及后文的数值矩阵都属于这个缩小的教学例子，不改变前面实际 kernel 的 BlockM=BlockN=64。

这个数值例子采用非 causal 模式，并单独取 s=ln(2)，使 exp(s*x)=2^x，便于手算。它不是标准缩放系数；前面实际 D=64 的标准 attention 应使用 s=1/8。除 causal 小节专门展示 mask 外，后文的输出均按这个非 causal 例子计算。

q/d	0	1
0	1	0
1	0	1

图2. 教学输入 Q（2×2）；行轴 q 为 query 编号，列轴 d 为 head 维度

k/d	0	1
0	0	1
1	1	0
2	2	1
3	0	2

图3. 教学输入 K（4×2）；行轴 k 为 key 编号，高亮为第 0 轮读取的 key 0、1

k/d	0	1
0	1	0
1	0	1
2	1	1
3	2	0

图4. 教学输入 V（4×2）；按 key 行与 K 同步分块，高亮为第 0 轮读取的两行

图2～图4的格内数字是参与运算的数据值。所有轴编号均从 0 开始；q、k 分别标识 query 与 key 行，d 标识 head/输出维度。后文 E0、E1 与 U0、U1 的后缀表示第 0、1 轮，并非矩阵元素下标。本文按精确实数运算展示分数，不包含 FP16 转换舍入。第 0 轮读取 key 0、1，第 1 轮读取 key 2、3。

第一步：取得当前 block 的数据分块

进入 kernel，先用 make_tensor 将当前 head 的裸指针包装成 CuTe 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 描述完整 Q 的地址与布局，本例为 [128,64]。K、V、O 同样建立 mK、mV、mO。

接着取得当前 block 负责的 Q tile：

Tensor gQ = local_tile(
    mQ,
    make_tile(Int<kBlockM>{}, Int<kHeadDim>{}),
    make_coord(m_block, _));

m_block=0 时对应前 64 行，m_block=1 时对应后 64 行。后续每轮也通过 local_tile 取得当前 K、V tile。

这两步都只创建视图，还没有搬运数据。此时确定的是“当前 block 要处理哪一块”。

变量的视图范围与存储位置汇总于表1：

表1. CuTe 变量的视图范围与存储位置

名字	表示什么
mQ	当前 head 的完整 Q
gQ	当前 block 使用的 global memory Q tile
sQ	shared memory 中的 Q tile
tCrQ	当前线程供 MMA 使用的 Q 寄存器 fragment
tCrS、tCrO	当前线程的分数与输出累加器

第二步：把 Q 和第一块 K 加载到 shared memory

一个 Q tile 由整个 block 协作加载。CuTe 用 TiledCopy 描述线程分工，再为当前线程选出源、目的视图：

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 回答当前线程从哪里读，partition_D 回答往哪里写；最后的 copy 才执行加载。flash::copy 是项目对 cute::copy 的包装。

这条 SM80 路径使用 cp.async，异步将数据从 global memory 搬到 shared memory。循环前先加载 Q 和第 0 块 K，再提交拷贝组。

Q 在后续各轮中保持不变，因此留在 shared memory 复用。K、V 则随着循环更新。shared memory 的具体布局由 traits 指定，加载与后续读取都遵循同一布局。

第三步：从 shared memory 取值，计算 QK

每轮开始，kernel 等待 Q、当前 K 的异步加载完成，再做 block 同步。随后便可以计算：

Q_tile[64,64] @ K_tile.T[64,64] → S_tile[64,64]

矩阵乘法由 Tensor Core 的 MMA 指令完成，输入需要按指令要求放在各线程的寄存器中。CuTe 先准备 fragment，再建立加载视图：

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);

这三行可以按用途读：

• tCrQ 准备当前线程参与 MMA 所需的寄存器。
• tXsQ 描述 shared memory 加载的源。
• tXrQ 为已有寄存器建立 Copy 目的视图。

retile_D 只改变访问视图，随后调用 copy 才装入数值。K 也采用相同的准备过程。

项目将这些加载和矩阵乘法封装在 gemm_smem 中：

flash::gemm_smem(tCrS, tCrQ, tCrK, tXsQ, tXsK, /* ... */);

它沿 head dimension 分段加载、累加。本例 D=64，每条 MMA 的归约长度为 16，因此沿 D 处理四段后，得到当前 64×64 的原始分数。

CuTe 的 GEMM 使用 A(M,K)、B(N,K)、C(M,N) 约定。这里可以直接把 K(key,dim) 作为 B，表达数学上的 Q @ K.T，无需提前生成 K 的转置副本。

用图2的 Q 与图3的 K 做点积，可得图5中的原始分数 S。这里尚未乘缩放系数 s；每行对应一个 query，每列对应一个 key：

q/k	0	1	2	3
0	0	1	2	0
1	1	0	1	2

图5. 原始分数 S=QKᵀ（2×4）；高亮为第 0 轮的两列，右侧为第 1 轮的两列

例如 S(0,2)=(1,0)·(2,1)=2。图5完整的 2×4 矩阵只是为了对照；实际执行时，每轮只保留当前两列的分数。

第四步：更新 online softmax

当前 S tile 只覆盖 64 个 key，后面还有下一块，尚不能得到最终 softmax 分母。因此 kernel 为每行 query 保存表2中的三份可继续更新的状态。在线更新归一化因子的思路见 Milakov 与 Gimelshein [2]，其与分块 attention 的组合见 [3]、[4]。

表2. 每行 query 的 online softmax 状态及其源码变量

状态	含义	源码变量
m	已处理原始分数的最大值，尚未乘 s	rRowMax
ℓ	相对于 m 计算的指数和	rRowSum
U	相对于 m 计算的输出分子	tCrO

对某一固定 query 行，令 $J$ 为已处理且未被 mask 的 key 集合，$x_j$ 为该行对 key $j$ 的原始点积，$v_j$ 为 V 的第 $j$ 行。以下省略 query 下标，所有归约均独立逐行执行。对于正缩放系数 $s$，$m=\max_{j\in J}x_j$，状态满足：

$$\ell=\sum_{j\in J} e^{s(x_j-m)},\qquad U=\sum_{j\in J} e^{s(x_j-m)}v_j.$$

新 tile 的有效 key 集合记为 $B$，更新后的最大值为 $m'=\max(m,\max_{j\in B}x_j)$。若最大值从 m 变成 m′，先计算：

$$\alpha=e^{s(m-m')}.$$

旧分母和旧输出分子都乘 α，就切换到了新的最大值基准。再加上当前块的贡献：

$$\ell'=\alpha\ell+\sum_{j\in B}e^{s(x_j-m')}, \qquad U'=\alpha U+\sum_{j\in B}e^{s(x_j-m')}v_j.$$

第一块存在有效 key 时，直接以该块最大值建立 m，并由该块权重建立 ℓ、U；此时没有需要重缩放的旧贡献。后续块再应用上述递推。本文各示例每行至少有一个有效 key，不讨论整行被 mask 时的额外处理。

这样，每一轮只需保存 m、ℓ、U，前一轮的完整分数可以丢弃。

源码中的 softmax_rescale_o 完成最大值与分母更新、缩放旧 U，并将当前分数原地变成指数权重。新的权重乘 V，则交给紧接着的第二个 GEMM。

做行归约前，源码为分数 fragment 建立行列视图。同一行的值分散在多个 lane 中，归约会先合并线程内的值，再通过 shuffle 合并线程间结果。这里先记住它最终得到的是每行的 m 和 ℓ 即可。

用两轮数值看清重缩放

从图5左侧两列出发，第 0 轮的两行原始分数分别为 (0,1) 和 (1,0)，行最大值都是 m=1。按 s=ln(2) 计算 exp(s*(x-m))，得到图6中的未归一化指数权重：

q/k	0	1
0	1/2	1
1	1	1/2

图6. 第 0 轮指数权重 E0（2×2）；s=ln(2)，两行最大值均为 1，尚未归一化

两行的分母都是 ℓ0=3/2。完成本轮 PV 后，还会得到当前输出分子 U0，具体乘法在下一步展开。

第 1 轮读取 key 2、3，新的最大值都变为 m′=2。因此两行都需要 α=2^(1-2)=1/2：旧分母从 3/2 缩放成 3/4，旧输出分子 U0 的每个分量也都乘 1/2。

本轮相对新最大值的指数权重见图7：

q/k	2	3
0	1	1/4
1	1/2	1

图7. 第 1 轮指数权重 E1（2×2）；两行累计最大值均为 2，列轴为全局 key 编号

当前两行的指数和分别为 5/4 和 3/2。加上已缩放的旧分母，最终得到 ℓ1=(2,9/4)。这里始终没有把某个 tile 单独归一化后再拼接：两轮通过同一份行状态合并。

第五步：让当前权重乘 V，累加输出分子

softmax 更新后，当前分数存储里已经是指数权重。源码先将它们转换为 FP16，再建立第二个 GEMM 的 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()));

第一行有实际数值转换，第二行复用这些 half 的存储。当前 MMA 分工允许第一个 GEMM 的结果在同一线程内组成第二个 GEMM 所需的 A 输入。

这里的 P 尚未除最终分母，它表示当前块的指数权重。计算的是：

P_tile[query,key] @ V_tile[key,dim] → 当前输出分子的增量

第二个 GEMM 仍采用 CuTe 的 B(N,K) 约定，此时 N 是输出维度、K 是 key。因此源码为同一份 V shared memory 建立 sVt(dim,key) 视图，并配合转置形式的加载指令准备 B fragment。

gemm_A_in_regs 使用已经在寄存器中的 tCrP，加载 V 并累加：

flash::gemm_A_in_regs(tCrO, tCrP, tCrVt, tXsVt, /* ... */);

这一轮结束后，U 已经包含当前块的贡献，m、ℓ 也已经更新，可以继续处理下一块 K、V。

将权重乘到具体的 V 上

第 0 轮使用图4中 V 的前两行，与图6中的 E0 相乘，计算 U0=E0 @ V[0:2,:]，得到图8：

q/d	0	1
0	1/2	1
1	1	1/2

图8. 第 0 轮输出分子 U0=E0·V[0:2,:]（2×2）；列轴 d 为输出维度

第 1 轮将图7的 E1 与图4中 V 的后两行相乘，得到新贡献 ΔU=E1 @ V[2:4,:]，见图9：

q/d	0	1
0	3/2	1
1	5/2	1/2

图9. 第 1 轮输出分子的新增量 ΔU=E1·V[2:4,:]（2×2）；未包含旧分子的贡献

例如第一行是 1×(1,1)+(1/4)×(2,0)=(3/2,1)。然后将新贡献加到已经缩放的旧结果上，U1=(1/2)*U0+ΔU，结果见图10：

q/d	0	1
0	7/4	3/2
1	3	3/4

图10. 两轮合并后的输出分子 U1=(1/2)·U0+ΔU（2×2）；尚未除以行分母

可见，第二轮不是简单地将两次 PV 结果相加；旧分子的缩放与旧分母的缩放必须使用相同的 α。

主循环怎样把加载与计算重叠

前面按计算依赖解释了每一步。实际调度还会提前加载后面要用的数据：

循环前：加载 Q、K[0]，提交异步拷贝组

第 i 轮：
  等待 Q、K[i] 就绪，block 同步
  发起 V[i] 的异步加载
  计算 Q @ K[i].T

  等待 V[i] 就绪，block 同步
  若有下一轮，发起 K[i+1] 的异步加载
  更新 online softmax，缩放旧输出分子
  计算 P[i] @ V[i]，累加新输出分子

因此，V 的加载可以与 QK 计算重叠，下一块 K 的加载可以与当前 softmax、PV 计算重叠。

cp_async_fence 提交拷贝组，cp_async_wait<0> 等待此前提交的组完成，__syncthreads() 让整个 block 安全使用协作加载的数据。同步也保护 shared memory 的复用，确保旧数据使用完之后才被覆盖。

本实现的 Q、K、V 各只有一份 shared memory tile。Q 留在其中，K、V 在合适的时机更新。

causal 模式在哪一步生效

开启 causal 时，第 query 行只允许访问 key <= query 的位置。实现分两步：

1. 跳过整块位于未来的 K、V tile，减少主循环次数。
2. 对与对角线相交的分数 tile，在 QK 之后、softmax 之前，将未来位置设为 -INFINITY。

这些位置经过指数函数后权重为 0，后续 online softmax 与 PV 流程保持一致。

本例中，query 0～63 的 block 只需遍历第一个 K、V tile；query 64～127 的 block 需要遍历两个。对角线所在 tile 内仍需逐元素 mask。

若给前面的缩小数值例子开启 causal，并令 query、key 都从 0 编号，原始分数变为图11的形式。这里采用左上对齐的规则 k≤q；当 query 与 key 长度不同时，它与右下对齐的 causal 规则不同。高亮处是未来位置：

q/k	0	1	2	3
0	0	−∞	−∞	−∞
1	1	0	−∞	−∞

图11. 左上对齐的 causal 原始分数（2×4）；高亮的 k>q 位置置为 −∞

这里两个 query 的 key 2、3 都在未来，因此第二个 K、V tile 可以整块跳过；第一个 tile 中的 (query=0,key=1) 仍需单独 mask。这个 causal 分支的输出会与前面非 causal 的 U1 不同，不能混用两者的归一化结果。

最后：归一化并写回 O

遍历完所有 K、V tile 后，寄存器中保存的是完整输出分子 U。逐行除以 ℓ，得到最终结果：

$$O=U/\ell.$$

同时计算并保存：

$$\mathrm{LSE}=s\,m+\log\ell.$$

回到非 causal 数值例子，分别用两行的 ℓ1=2、9/4 去除图10中 U1 的各列，得到图12：

q/d	0	1
0	7/8	3/4
1	4/3	1/3

图12. 非 causal 最终输出 O（2×2）；将图10逐行除以 ℓ1=(2,9/4)

例如 O(1,0)=3/(9/4)=4/3。两行的 LSE 分别为 ln(8) 和 ln(9)，与直接对完整原始分数计算 log(sum(exp(s*S))) 一致。

输出此时仍按 MMA 的线程分工分散在寄存器中。源码先转换成 FP16，写到 shared memory，再按适合 global memory 写回的线程分工读出并存储：

FP32 输出累加器
  → FP16 fragment
  → shared memory 的 sO
  → 写回分工下的临时寄存器
  → global memory 的 O

sO 复用已经不再需要的 sQ 空间。这个过程让计算分工与写回分工衔接起来。

回到 kernel，可以按这个顺序阅读

建立完整矩阵视图，取得当前 Q tile
  ↓
加载 Q、第一块 K 到 shared memory
  ↓
循环遍历 K、V tile
  ├─ QK：从 shared memory 加载到寄存器，计算分数
  ├─ mask：若开启 causal，屏蔽未来位置
  ├─ online softmax：更新 m、ℓ，缩放旧 U，生成当前权重
  └─ PV：权重乘当前 V，累加到 U
  ↓
逐行计算 O=U/ℓ，保存 LSE
  ↓
整理输出的线程分工，写回 global memory

读到 CuTe API 时，只需先确认它服务于哪一步：local_tile 选 block 的数据块，partition 选线程的访问范围，partition_fragment 准备计算所需的存储，retile 为已有存储建立另一种视图，copy 和 gemm 执行实际的数据操作与计算。

本文使用 FP16、连续存储、L 为 64 的倍数、D 与模板分派尺寸一致的输入。当前教学实现的尾块保护、任意 stride、BF16 转换与 stream 支持仍有限制；返回 LSE 也不表示已经实现 backward。当前指数计算实际使用 expf，接口的默认 scale 为 1，复现标准 attention 时需要显式传入 1/sqrt(D)。

主流程可直接对照 flash_attention_v2_cutlass_kernel。需要继续追布局细节时，可阅读 retile_D 专题，或运行配套的 布局核对程序。配套程序的 host 布局检查已重新运行；新增小矩阵的点积、两轮 online softmax、PV、causal 结果与一次性 softmax 已在 CPU 上核对。这里的数值表采用精确分数，不模拟真实 FP16 转换舍入，也没有执行 GPU forward 数值验证。

图1为分块示意；图2～图12的数值由本文选定输入推导，未转载论文插图。论文提供算法依据，具体 API、精度转换与流水调度以参考资料 [6] 的源码快照为准。

参考资料

[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，网页链接提供对应的源码快照。工作区基于 commit dacda43160add8ed05c43b943b40787606a23f2f，其中 flash_attention.cu 含本地修改，不能仅凭该 commit 复现本文；两文件的 SHA-256 分别为 06f07a11854fb69b7c0a62d4e893d0d4b851e0ae6847b722bdcac224376f0055、fbb88f1a90ee1c6a937e3b5ac5387c0c7be53c04da30ddf08a6cbf1c897fcf63。