# 用 CuTe 读懂 FlashAttention：从一个 64 行的 Query 分块开始

第一次读 CuTe 写的 attention，往往还没看到 softmax，就被这几行拦住了：

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

同一个 Q，为什么要准备这么多 tensor？`partition` 之后为什么还有 `retile`？它们什么时候真的搬了数据？

先看这几行最后要完成的事：一个线程块负责 64 个 query，把它们依次与 K、V 的各个分块计算，最后写出 64 行结果。上面的三个变量分别描述“计算需要哪些寄存器”“从 shared memory 的哪里读”“加载指令往哪些寄存器位置写”。真正加载时才调用 `copy`。

下面就跟着这 64 行 query 走一遍：先确定整个线程块的任务，再看数据怎样进入 shared memory 和寄存器，最后接上两个矩阵乘法之间的 softmax。遇到布局变换时，把当前线程实际持有的元素列出来，再读代码。

本文以 tiny-flash-attention 当前工作区的 [`flash_attention_cutlass/csrc/flash_attention.cu`](flash_attention_cutlass/csrc/flash_attention.cu) 为主线，对照 [`kernel_traits.h`](flash_attention_cutlass/csrc/kernel_traits.h) 和 [`utils.h`](flash_attention_cutlass/csrc/utils.h)。实现以 `dacda43` 为基础，当前工作区已统一变量命名；CuTe 接口另用本机 CUTLASS 4.8 检出版本 `0b55a2f6` 核对。仓库记录的 CUTLASS 子模块提交仍是 `751eb9a8`，因此下文会明确区分当前接口与旧教程里的写法。

阅读前只需要认识 CUDA 的 thread、warp、block，以及矩阵乘法和 softmax。这里选定 FP16、SM80 的指令路径，不展开 Hopper、Blackwell 专用的 attention 实现。后文的 C++ 片段用于对照源码，省略了重复的模板参数；能单独运行的例子会另行标明。

本文与项目中两份 CuTe 实现的局部变量均参照官方 [`sgemm_sm80.cu`](https://github.com/NVIDIA/cutlass/blob/0b55a2f6/examples/cute/tutorial/sgemm_sm80.cu) 命名：完整矩阵用 `mQ`，global memory 分块用 `gQ`，MMA 寄存器视图用 `tCrQ`，Copy 寄存器视图用 `tXrQ`。Q、K、V 等后缀保留 attention 的数据含义。旧版名称列在[文末对照表](#旧版命名与当前源码对照)，方便阅读早期版本。

## 先固定一个能算清楚的例子

设 Q、K、V 都是连续存储的 `[B,H,L,D]` tensor：B 是 batch 数，H 是 head 数，L 是序列长度，D 是每个 head 的维度。先只看其中一个 batch、一个 head：

```text
Q、K、V：128×64

L = 128
D = 64
BlockM = 64
BlockN = 64
每个线程块：4 个 warp，共 128 个线程
```

`BlockM` 表示一次处理多少行 query，`BlockN` 表示一次读入多少行 key、value。这正是当前 Python 绑定路径使用的 `64×64、4 warp` 分块配置，head dimension 则由模板分派选择。

普通 attention 写成：

$$
S = sQK^T,\qquad P = \operatorname{softmax}_{\mathrm{row}}(S),\qquad O = PV,
$$

其中通常取：

$$
s = \frac{1}{\sqrt{D}}
$$

本例的矩阵形状为：

```text
Q[128,64] @ K.T[64,128]  -> S[128,128]
P[128,128] @ V[128,64]   -> O[128,64]
```

FlashAttention 将 K、V 分块读入，并在片上完成 softmax 和后面的矩阵乘法，使完整的 `128×128` 中间矩阵不必写入 global memory。序列更长时，省下的这份中间存储和读写才更可观；矩阵乘法的工作量仍随 L² 增长。

对于前 64 行 Q，线程块的任务是：

```text
固定 Q[0:64,:]

第 0 轮：Q[0:64,:] @ K[0:64,:].T
         更新 softmax 的状态
         当前权重 @ V[0:64,:]，更新输出分子

第 1 轮：Q[0:64,:] @ K[64:128,:].T
         更新 softmax 的状态，并修正旧输出分子
         当前权重 @ V[64:128,:]，继续累加

最后：用分母归一化，写回 O[0:64,:]
```

这里有一个暂时留下的问题：第一轮还没见过后面的 key，softmax 的分母怎么确定？先记住源码保存了行最大值、分母和输出分子三份状态，等数据流走到 softmax 时再把它们算清楚。

## 这 64 行归哪个线程块，哪个线程

Q 的不同输出行可以独立计算，因此源码按 Q 的行分块并行启动线程块：

```cpp
const int num_m_block =
    (params.q_seqlen + Kernel_traits::kBlockM - 1)
    / Kernel_traits::kBlockM;

dim3 grid(num_m_block, params.bs * params.head, 1);
dim3 block(Kernel_traits::kNThreads);
```

本例 `B=H=1`，所以 grid 中有两个 block：

| 线程块 | 固定的 Q | 遍历的 K、V | 写回的 O |
|---|---|---|---|
| `blockIdx.x=0` | 第 0～63 行 | 全部 128 行 | 第 0～63 行 |
| `blockIdx.x=1` | 第 64～127 行 | 全部 128 行 | 第 64～127 行 |

**64 行 Q 是整个 block 的任务。** block 内的 128 个线程还要协作加载数据、执行矩阵乘法、做行归约。一个线程只保存其中一部分元素。

当前 MMA 配置把 4 个 warp 排在 M 方向，所以 warp 0 负责这个分块中的第 0～15 行，warp 1 负责第 16～31 行，依此类推。warp 内的线程编号称为 lane，取值为 0～31；这些线程再按 MMA 指令规定的方式分担各行。

同一批 64 行 query 会反复使用，每一轮只更换 K、V 分块。这样组织之后，Q 可以留在 shared memory，分数和输出累加器可以留在寄存器里。

## Tensor 先描述地址，local_tile 再选出分块

进入 kernel 后，先找到当前 batch、head 的起点，再把裸指针包装成 CuTe Tensor：

```cpp
const int m_block = blockIdx.x;
const int base_id = blockIdx.y;
const int bs_head_offset = base_id * params.head_stride;

Tensor mQ = make_tensor(
    make_gmem_ptr(reinterpret_cast<Element *>(params.q_ptr)
                  + bs_head_offset),
    make_shape(params.q_seqlen, Int<kHeadDim>{}),
    make_stride(Int<kHeadDim>{}, Int<1>{}));
```

本例 `mQ` 的 shape 是 `(128,64)`，stride 是 `(64,1)`。因此访问 `mQ(m,d)`，对应的元素偏移为：

```text
offset = m * 64 + d
```

Shape 描述每个维度有多大，Stride 描述该维度前进一步时偏移多少个元素，Layout 则把坐标映射到偏移。Tensor 在这套映射之外，再带上存储位置。

`mQ` 对应当前 batch、head 的完整 Q，K、V 的完整视图同样命名为 `mK`、`mV`。这里传入的是已有的 global memory 指针，`make_tensor` 创建访问视图。它没有另行分配一份 Q，也没有执行加载。后面遇到 `make_tensor<Element>(shape)` 时要留意：那种带元素类型、没有传入数据指针的写法，会创建拥有自己存储的 fragment。

代码里的 `_64`、`Int<64>` 表示编译期常量 64，`_` 表示保留这个维度。CuTe 可以利用这些静态信息，在编译时推导分块和布局。

接着，用 `local_tile` 取当前 block 的 Q 分块：

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

对 `mQ[128,64]` 使用 `tile[64,64]`，第一个分块坐标固定为 `m_block`，第二个分块坐标保留。所以本例的结果是：

```text
gQ 的 shape：(64,64,1)

m_block=0：gQ(m,d,0) 对应 mQ(m,d)
m_block=1：gQ(m,d,0) 对应 mQ(64+m,d)
```

最后的 1 表示 D 方向只剩一个 tile。它不是 K、V 的迭代次数。源码每次更换 K 分块时，会重新调用：

```cpp
gK = local_tile(
    mK,
    make_tile(Int<kBlockN>{}, Int<kHeadDim>{}),
    make_coord(nbi + 1, _));
```

`local_tile` 也只产生视图。到现在，我们确定了要读哪一块数据，还没有把它搬进 shared memory。

## 128 个线程怎样把 Q 搬进 shared memory

CuTe 用 TiledCopy 描述一组线程如何协作拷贝。对于本例，`kernel_traits.h` 里的配置可以展开成：

```cpp
using GmemTiledCopyQKV = decltype(make_tiled_copy(
    Copy_Atom<SM80_CP_ASYNC_CACHEGLOBAL<cute::uint128_t>,
              cutlass::half_t>{},
    Layout<Shape<_16,_8>, Stride<_8,_1>>{},
    Layout<Shape<_1,_8>>{}));
```

这里有三个部分：

| 配置 | 本例中规定的事情 |
|---|---|
| Copy Atom | 使用 128 位 `cp.async`，从 global memory 异步拷到 shared memory |
| Thread layout `(16,8):(8,1)` | 将 128 个线程组织为 16 行、每行 8 个线程 |
| Value layout `(1,8)` | 每个线程一次负责同行的 8 个 half，即 16 字节 |

于是一次覆盖 `16×64` 个元素：

```text
逻辑第 0 行：thread 0 读列 0～7，thread 1 读列 8～15，……，thread 7 读列 56～63
逻辑第 1 行：thread 8 读列 0～7，……，thread 15 读列 56～63
……
逻辑第 15 行：thread 120～127
```

处理完整的 `64×64` Q tile，需要沿行方向重复四次。相邻线程读取相邻的 16 字节区间，形成连续的访问。

这些数字来自当前 `D=64` 的配置。源码实际上通过 `kBlockKSmem`、`kGmemElemsPerLoad` 等常量计算 Thread layout，换 D 后要重新推导。

### partition_S 和 partition_D 分别选出哪一端

kernel 中先取得当前线程的拷贝分工，再将它分别应用到源、目的 tensor：

```cpp
typename Kernel_traits::GmemTiledCopyQKV copy_qkv;
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` 回答“当前线程往哪里写”。它们返回各自的访问视图，最后一行才执行拷贝。`flash::copy` 是项目自己的小包装，会遍历分块中的剩余维度，再调用 `cute::copy`。

### 按官方例子的方式拆开变量名

官方 [GEMM 教程的命名说明](https://github.com/NVIDIA/cutlass/blob/0b55a2f6/media/docs/cpp/cute/0x_gemm_tutorial.md#copy-partitioning) 将 `tAsA` 读成“把分工 `tA` 应用到 tensor `sA`”。因此可以把一个名字拆成两段：

```text
tQgQ = tQ + gQ：Q 拷贝分工下的 global memory 视图
tQsQ = tQ + sQ：同一拷贝分工下的 shared memory 视图

tCrQ = tC + rQ：MMA 分工下的 Q 寄存器视图
tXsQ = tX + sQ：shared memory 加载 Copy 的 Q 源视图
tXrQ = tX + rQ：同一个 Copy 的 Q 目的寄存器视图
```

前面的 `tQ`、`tC`、`tX` 区分线程分工，后面的 `g`、`s`、`r` 区分 global memory、shared memory、register。`t` 不需要解释成英文的 “to”。

官方 SM80 例子用 `tCrA/B/C` 表示 MMA 相关的 fragment，用 `tXsA/B`、`tXrA/B` 表示加载 Copy 的源、目的视图。本文沿用这套前缀，再将数据后缀扩展成 Q、K、V、P、S、O：

| 本文变量 | 表达的内容 |
|---|---|
| `mQ`、`gQ`、`sQ` | 当前 head 的完整 Q、当前 block 的 global memory Q 分块、shared memory Q 分块 |
| `tQgQ`、`tQsQ` | 同一 global→shared Copy 下，当前线程的源、目的视图 |
| `tCrQ`、`tCrK` | 当前线程参与 MMA 所需的 Q、K fragment |
| `tXsQ`、`tXrQ` | shared→register Copy 下，当前线程的源、目的视图 |
| `tCrS`、`tCrO` | 分数 S 和输出 O 的 MMA 累加器 |

注意 `tCrQ` 和 `tXrQ` 都是寄存器中的 Q，区别在于访问它们的分工；名字中的 C 也不表示这个 tensor 一定是 GEMM 的 C 操作数。具体如何分配元素，仍要看 `partition`、`retile` 和实际布局。

Copy 对象也沿用官方例子的写法：`copy_qkv`、`thr_copy_qkv` 表示 global→shared 的整组拷贝及其线程切片；`s2r_copy_q`、`s2r_thr_copy_q` 表示 shared→register 的相应对象，其中 `s2r` 是 shared to register。输出时用 `r2s` 表示相反方向。它们是对象名，不参与上面的 tensor 名字拆分。

## shared memory 为什么还要 Swizzle

global memory 中，Q 是普通行主序。shared memory 中，源码使用的是：

```cpp
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>{}));
```

这是将本例的常量代入后的写法。`tile_to_shape` 将一个 `8×64` 的基本布局铺到整个 `64×64` tile。

shared memory 分为多个 bank。若一组线程需要访问不同地址，却集中落在相同 bank 上，访问可能被拆开处理，产生 bank conflict。适合 global memory 连续加载的排布，未必适合后面的 `ldmatrix`。

Swizzle 为逻辑坐标选择另一套物理偏移。以本例基本布局的一行起点为例，原偏移为 `x=row*64`，变换是：

```text
swizzle(x) = x XOR ((x AND 0x1c0) >> 3)

row        0    1    2    3    4    5    6    7
原偏移      0   64  128  192  256  320  384  448
新偏移      0   72  144  216  288  360  432  504
```

这里的单位是 half 元素。低 3 位保持不变，因此 8 个 half 组成的 16 字节小组仍保持连续；不同小组的起点被重新安排，以配合 warp 的加载方式。

`partition_D(sQ)` 写入时使用这套布局，后面的 `partition_S(sQ)` 读取时也使用它。同一个 `sQ(m,d)` 始终代表同一个逻辑元素。

所以，Swizzle 的作用要结合具体的拷贝和加载指令理解。换 tile 大小或线程分工后，需要重新检查 bank 分布；不能仅凭加上 `Swizzle` 就断定所有访问都没有冲突。

## 先看一条 MMA，寄存器 fragment 才有具体含义

Q 进入 shared memory 后，接下来要计算 `Q_tile @ K_tile.T`。源码选择的 SM80 MMA Atom 是：

```cpp
MMA_Atom<SM80_16x8x16_F32F16F16F32_TN>
```

MMA 是矩阵乘加。这个类型封装的 PTX 指令为：

```ptx
mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32
    {d0,d1,d2,d3},
    {a0,a1,a2,a3},
    {b0,b1},
    {c0,c1,c2,c3};
```

一整个 warp 协作计算：

```text
D = A @ B + C

A：16×16
B：16×8       按通常数学写法
C、D：16×8
```

`m16n8k16` 表示 M、N、K 的尺寸，类型顺序是 D、A、B、C。每个 lane 提供的寄存器为：

| 操作数 | 每个 lane 持有的逻辑值 | PTX 中的操作数 |
|---|---|---|
| A | 8 个 half | 4 个 32 位打包值 |
| B | 4 个 half | 2 个 32 位打包值 |
| C、D | 各 4 个 float | 各 4 个 FP32 值 |

两个 half 装在一个 32 位操作数中。因此后文说“8 个 half 槽位”，并不表示 8 个独立的硬件寄存器。

每个 lane 提供矩阵中的哪些位置，由指令约定。CuTe 的 MMA traits 将这份约定写成 `(thread,value) -> 矩阵位置` 的布局，常见名称是 `LayoutA_TV`、`LayoutB_TV`、`LayoutC_TV`。`TV` 就是 Thread、Value。

还有一个坐标顺序需要提前约定：CuTe 的 MMA 接口用 `A(M,K)`、`B(N,K)`、`C(M,N)`。于是第一个 GEMM 可以直接把 `K[n,d]` 当作 B tensor，表示：

```text
C(m,n) += sum_d Q(m,d) * K(n,d)
```

这正好计算数学上的 `Q @ K.T`，不需要先在内存中复制出一份 K 的转置。类名中的 `TN` 和指令的 `.row.col` 对应操作数约定，不能把它们直接读成“先执行一次内存转置”。

### 从一条指令扩大到整个 block

当前 traits 使用：

```cpp
using TiledMma = TiledMMA<
    MMA_Atom<SM80_16x8x16_F32F16F16F32_TN>,
    Layout<Shape<_4,_1,_1>>,
    Tile<_64,_16,_16>>;
```

第二个参数将 4 个 warp 沿 M 方向组织起来。第三个参数在当前 CuTe 接口中名为 `PermutationMNK`，可以描述各方向的 tiling、排列；本例是整数 tiler，使 `tile_size` 得到 `(64,16,16)`。

这个 `64×16×16` 仍然由多条 `16×8×16` 指令完成。一个 `64×64 @ 64×64` 的 GEMM，还要沿 N、K 方向继续重复。每个 warp 在 N 方向执行多个 MMA，并复用已经加载的 Q 值。

旧文章里有时用 `ValLayoutMNK = Layout<Shape<_1,_2,_1>>` 讲“每线程重复”。当前项目虽保留了这个类型名，实际传给 `TiledMMA` 的第三个参数已经是上面的 `Tile`。读新接口时，应当沿当前调用路径查看 `PermutationMNK` 和 `tile_size`。

## partition_fragment 创建了什么

取得当前线程的 MMA 分工之后，源码创建 Q、K fragment：

```cpp
auto thr_mma = mma.get_slice(threadIdx.x);

Tensor tCrQ = thr_mma.partition_fragment_A(sQ);
Tensor tCrK = thr_mma.partition_fragment_B(sK);

auto tCrS = partition_fragment_C(
    mma, Shape<_64,_64>{});
clear(tCrS);
```

这里的 fragment 可以理解为“当前线程要持有的矩阵元素”。对这条 SM80 MMA，`partition_fragment_A/B` 根据线程分工创建寄存器存储，刚创建时还没有 Q、K 的值。`clear` 则确实将分数累加器清零。

打印本例的寄存器布局，shape 如下。为便于阅读，省略打印结果中的静态常量前缀 `_`：

```text
tCrQ：((2,2,2),1,4)
tCrK：((2,2),8,4)
tCrS：((2,2),1,8)
```

三个顶层维度分别表示：

```text
A fragment：(单条 MMA 的每线程 A 值，M 方向重复，K 方向重复)
B fragment：(单条 MMA 的每线程 B 值，N 方向重复，K 方向重复)
C fragment：(单条 MMA 的每线程 C 值，M 方向重复，N 方向重复)
```

CuTe 把一个维度叫作 mode，mode 自己还可以包含子维度。例如 `tCrQ` 最外层有三个 mode，所以 rank 是 3；它的第一维 shape 是 `(2,2,2)`，大小为 8。

本例每个线程保存 32 个 Q half、128 个 K half、32 个分数 float。K fragment 会在不同 warp 间重复，因为 4 个 warp 处理不同的 Q 行，却需要相同的 K。分数输出则恰好由 `128×32=4096` 个 float 覆盖整个 `64×64` tile。

这些数量只描述 fragment 的逻辑值。编译器最后用了多少寄存器，还要看生命周期、优化和是否发生溢出。

## ldmatrix 和 retile_D 怎样把寄存器填好

源码从 shared memory 读取 Q、K 时使用：

```cpp
using SmemCopyAtom =
    Copy_Atom<SM75_U32x4_LDSM_N, cutlass::half_t>;
```

虽然名字带 `SM75`，这条加载指令也用于这里的 SM80 路径。它封装的是：

```ptx
ldmatrix.sync.aligned.x4.m8n8.shared.b16
    {r0,r1,r2,r3}, [addr];
```

warp 一次协作加载四个 `8×8`、16 位元素的矩阵。每 8 个 lane 为其中一个矩阵提供八行的地址，32 个 lane 都接收结果，每个 lane 最终得到 4 个 32 位值，也就是 8 个 half。

先只看其中一个非转置的 `8×8` 加载。lane 0 提供第 0 行的地址，但这一行的 8 个 half 会分给 lane 0～3：

```text
X(0,0) X(0,1) | X(0,2) X(0,3) | X(0,4) X(0,5) | X(0,6) X(0,7)
     lane 0   |      lane 1   |      lane 2   |      lane 3
```

因此，“我提供哪一行的地址”与“我的寄存器收到哪一对值”有不同的规则。CuTe 为 Copy 分别定义源端、目的端布局，正是为了表达这件事。`ldmatrix` 的行地址还要求 16 字节对齐，这也是前面保留 8 个 half 连续分组的原因之一。

### Q：加载和计算访问同一份寄存器

源码先让 Copy 采用 MMA 的 A 参考分工，再为当前线程创建源视图和目的视图：

```cpp
auto s2r_copy_q = make_tiled_copy_A(
    typename Kernel_traits::SmemCopyAtom{}, mma);
auto s2r_thr_copy_q = s2r_copy_q.get_thread_slice(threadIdx.x);

Tensor tXsQ = s2r_thr_copy_q.partition_S(sQ);
Tensor tXrQ = s2r_thr_copy_q.retile_D(tCrQ);
```

以 warp 0 的 lane 0 为例，一条 MMA 的 8 个 Q 值按顺序是：

```text
槽位 0、1：Q(0,0)、Q(0,1)
槽位 2、3：Q(8,0)、Q(8,1)
槽位 4、5：Q(0,8)、Q(0,9)
槽位 6、7：Q(8,8)、Q(8,9)
```

本例的 `x4` 一次正好写入这 8 个 half。原来的 MMA 视图用 `(2,2,2)` 表示它们，Copy 视图将它们组织为一组 8 个值：

```text
tCrQ 的 shape：((2,2,2),1,4)
tXrQ 的 shape：((8,1),1,4)
```

两个 tensor 使用同一个 `data()`。通过 Copy 视图写入后，MMA 视图马上就能访问这些值。`retile_D` 本身只计算访问布局，没有执行加载，也没有在线程间搬运寄存器数据。这里的 D 表示 Copy 的 destination。

### K：一次 x4 要从两个 N fragment 取值

一条 MMA 每个 lane 只需要 4 个 B half，而 `x4` 每个 lane 要写入 8 个。当前配置的 Copy 会将两个 N 方向的 B fragment 组织到一起。

对 warp 0 的 lane 0，第一组加载涉及：

| 来自哪个 MMA fragment | K 的逻辑元素 | 原 B fragment 中的 half 槽位 |
|---|---|---|
| 第 0 个 N fragment | `K(0,0),K(0,1),K(0,8),K(0,9)` | `0,1,2,3` |
| 第 1 个 N fragment | `K(8,0),K(8,1),K(8,8),K(8,9)` | `16,17,18,19` |

这里的行号 0、8 是 key 的位置。后四个值落在槽位 16～19，是因为原 `tCrK` 的 K 重复维度 stride 为 4，N 重复维度 stride 为 16。

`retile_D` 保留这个对应关系，让一次 Copy 可以写入这八个位置。随后 MMA 仍通过原来的 B 视图取值。

这也是不能随手将 `retile` 换成“把连续 8 个寄存器分为一组”的原因：对本例的 K，这一组在原 fragment 存储中并不连续。参考文章里 `x1` 的简单分组是理解入口，换成这里的 `x4` 后，需要结合实际的 MMA 布局重新看。

## 第一个 GEMM：沿 head dimension 累加分数

Q、K 的 Copy 视图准备好后，`gemm_smem` 负责 shared memory 到寄存器的加载，以及 QK 的矩阵乘加。下面沿用官方 GEMM 例子中的通用 A、B、C 命名：本次 A 对应 Q、B 对应 K、C 对应分数 S。`tXsA/B` 是 Copy 源视图，`tCrC` 是累加器。核心代码是：

```cpp
Tensor tXrA = s2r_thr_copy_a.retile_D(tCrA);
Tensor tXrB = s2r_thr_copy_b.retile_D(tCrB);

cute::copy(s2r_copy_a,
           tXsA(_, _, _0{}), tXrA(_, _, _0{}));
cute::copy(s2r_copy_b,
           tXsB(_, _, _0{}), tXrB(_, _, _0{}));

#pragma unroll
for (int k_block = 0; k_block < size<2>(tCrA); ++k_block) {
    if (k_block < size<2>(tCrA) - 1) {
        cute::copy(s2r_copy_a,
                   tXsA(_, _, k_block + 1), tXrA(_, _, k_block + 1));
        cute::copy(s2r_copy_b,
                   tXsB(_, _, k_block + 1), tXrB(_, _, k_block + 1));
    }
    cute::gemm(mma, tCrA(_, _, k_block), tCrB(_, _, k_block), tCrC);
}
```

在 QK 这次乘法中，M 是 query 行，N 是 key 行，K 是 head dimension。因此本例 `size<2>(tCrA)=4`，循环依次处理 D 的 `0～15、16～31、32～47、48～63`。

循环前先准备第 0 段；循环中先加载下一段，再计算当前段。这给编译器留下安排加载和计算的空间。这里的 shared memory 加载使用 `ldmatrix`，与 global memory 到 shared memory 的 `cp.async` 是两种指令，不能将所有 `copy` 都理解为异步拷贝。

四段累加结束后，寄存器中得到当前 `64×64` 分数块。Q 仍留在 shared memory，本项目会在后续 K、V 迭代中再次从 shared memory 加载它，并没有让整块 Q 跨所有迭代一直留在寄存器里。

## 分数分散在各个线程里，怎样逐行做 softmax

普通二维矩阵可以直接按行归约。MMA 的结果却先以 fragment 的方式排列。对 warp 0、lane 0，前几个分数槽位是：

```text
r0=S(0,0)  r1=S(0,1)  r2=S(8,0)  r3=S(8,1)
r4=S(0,8)  r5=S(0,9)  r6=S(8,8)  r7=S(8,9)
……
```

源码先为同一份累加器创建按行、列访问的视图：

```cpp
Tensor tCrS_rowcol = make_tensor(
    tCrS.data(),
    flash::convert_layout_acc_rowcol(tCrS.layout()));
```

本例 shape 的变化为：

```text
MMA 视图：((2,2),1,8)
行列视图：((2,1),(2,8))
           行       列
```

对于当前线程，这表示 2 行、每行 16 个值。lane 0 的两行是第 0、8 行；它持有的列为 `0,1,8,9,…,56,57`。lane 1 持有相同行的 `2,3,10,11,…,58,59`，lane 2、3 再补齐其余列。

因此，一整行的 64 个分数分布在同一 warp 内相邻的 4 个 lane 中。`utils.h` 的归约分两步进行：

```text
thread_reduce_：当前线程先归约自己持有的 16 个列值
Allreduce<4>：再用 shuffle 合并四个 lane 的局部结果
```

这样得到完整行的最大值、分母和。`convert_layout_acc_rowcol` 只是换视图，跨 lane 的信息合并由后面的 shuffle 完成。

## 没看完所有 K，softmax 怎样继续算

现在回到开头留下的问题。对于一行 query，设已经处理过的 key 的原始分数是 $x_j=q\cdot k_j$，缩放系数为正数 $s$。保存三份状态：

$$
m=\max_j x_j,\qquad
\ell=\sum_j e^{s(x_j-m)},\qquad
U=\sum_j e^{s(x_j-m)}v_j.
$$

$m$ 是原始 QK 的最大值，$\ell$ 是当前分母，$U$ 是长度 D 的输出分子。最终结果是 $O=U/\ell$。

新来一个 K、V tile 后，按下面的顺序更新：

$$
\begin{aligned}
m' &= \max\left(m,\max_{j\in\mathrm{new}}x_j\right),\\
\alpha &= e^{s(m-m')},\\
p_j &= e^{s(x_j-m')},\quad j\in\mathrm{new},\\
\ell' &= \alpha\ell+\sum_{j\in\mathrm{new}}p_j,\\
U' &= \alpha U+\sum_{j\in\mathrm{new}}p_jv_j.
\end{aligned}
$$

关键是第二行。旧分子、旧分母都以旧最大值 m 为基准；最大值变成 m′ 后，给它们乘上同一个 $\alpha$，就能和新一块相加。旧分数不需要重新读出来。

### 用四个分数算一次 rescale

把 head dimension 暂时缩成 1，取 `s=1`，只是为了手算。第一块的分数为 `[1,2]`，V 为 `[10,20]`：

```text
m = 2
l = exp(-1) + 1
U = 10*exp(-1) + 20
```

第二块的分数为 `[3,4]`，V 为 `[30,40]`，新最大值变成 4：

```text
alpha = exp(2-4) = exp(-2)

l' = exp(-2)*(exp(-1)+1) + exp(-1)+1
   = exp(-3)+exp(-2)+exp(-1)+1

U' = exp(-2)*(10*exp(-1)+20) + 30*exp(-1)+40
   = 10*exp(-3)+20*exp(-2)+30*exp(-1)+40
```

`U'/l' ≈ 34.93`，正好等于对 `[1,2,3,4]` 一次做完整 softmax 后的加权和。第一块只留下了 `m、l、U`，就足够继续计算。

### 和 softmax_rescale_o 对上号

采用本文命名后，对应关系为：

| 数学量 | 本文变量 |
|---|---|
| 原始 QK 分数，随后原地变成当前块的指数权重 | `tCrS_rowcol`，与 `tCrS` 共享存储 |
| m | `rRowMax` |
| l | `rRowSum` |
| U | `tCrO` |
| s | `params.softmax_scale` |

第一块调用 `softmax_rescale_o<true>`，直接初始化最大值、分母；输出累加器此前已清零。后续块调用 `<false>`，更新最大值，用 `alpha` 缩放旧分母和旧输出，再加上当前块的分母。

当前块的 `P @ V` 则由紧接着的第二个 GEMM 完成。所以函数返回时，`tCrO` 的旧值已经 rescale，但本轮新的输出贡献还没加进去。

还有一个容易被函数名误导的地方：当前 `scale_apply_exp2` 的函数体实际调用 `expf`，kernel 也传入普通的 `params.softmax_scale`。虽然参数结构还保留 `softmax_scale_log2`，这条路径没有使用它。阅读时应按自然指数的公式理解，不能额外乘一次 `log2(e)`。

## 第二个 GEMM：分数怎样直接变成输入

第一个 GEMM 的累加器是 FP32，第二个 GEMM 的 A 输入需要 FP16。源码先转换类型，再变换寄存器视图：

```cpp
Tensor rP = flash::convert_type_f32_to_f16(tCrS);

Tensor tCrP = make_tensor(
    rP.data(),
    flash::convert_layout_rowcol_Aregs<TiledMma>(tCrS_rowcol.layout()));
```

此时 `tCrS` 已经被 softmax 函数原地改成当前块的指数权重。这里变量名叫 P，但还没有除最终的分母；归一化被推迟到所有 K、V 分块计算结束。

两行分别做了两件事：第一行把 float 数值转换成 half，创建新 fragment；第二行沿用 `rP.data()`，为这些 half 建立第二次 MMA 的 A 视图。

对于本例、warp 0 的 lane 0，前八个 half 按存储顺序是：

```text
P(0,0), P(0,1), P(8,0), P(8,1),
P(0,8), P(0,9), P(8,8), P(8,9)
```

与前面的 Q fragment 对照，就能看到它们正好是 A fragment 需要的那八个位置。布局变化为：

```text
第一个 GEMM 的 C：((2,2),1,8)   共 32 个值
按行列访问：     ((2,1),(2,8))
第二个 GEMM 的 A：((2,2,2),1,4) 共 32 个值
```

这里将两个相邻 N fragment 的值，组织成第二个 GEMM 中一段 K=16 的 A 输入。值都已经在当前 lane 里，因此这一步可以通过改视图完成，无需先写到 shared memory 再由其他线程重新加载。

这解释了融合两个 GEMM 为什么要考虑 MMA 的线程分工：需要检查第一个 GEMM 生成的值，能否在相同 lane 中组成第二个 GEMM 的 A 输入。仅仅看到两个 tensor 的元素总数相等，还不足以保证可以这样衔接。

### V 为什么需要转置视图和转置加载

第二个 GEMM 的数学形状是：

```text
权重[64,64] @ V[64,64] -> 输出分子[64,64]
       query,key       key,dim           query,dim
```

虽然本例的三个尺寸碰巧都是 64，它们的含义已经改变：第二个 GEMM 的 M 是 query，N 是输出的 dim，K 是 key。

CuTe 的 B tensor 要按 `(N,K)` 索引，所以这里需要按 `(dim,key)` 访问 V。源码在同一份 shared memory 上建立两种视图：

```cpp
Tensor sV = make_tensor(
    make_smem_ptr(shared_storage.smem_v.data()), SmemLayoutV{});

Tensor sVt = make_tensor(
    make_smem_ptr(shared_storage.smem_v.data()), SmemLayoutVt{});
```

本例中，它们满足：

```text
sV(key,dim) 与 sVt(dim,key) 指向同一个 half
```

没有在 shared memory 中再生成一份转置矩阵。写 V 时按 `sV` 的布局，给第二个 GEMM 加载时则按 `sVt` 的布局选择地址，并使用转置形式的加载 Atom：

```cpp
using SmemCopyAtomTransposed =
    Copy_Atom<SM75_U16x8_LDSM_T, cutlass::half_t>;
```

它对应 `ldmatrix.x4.trans`，将 shared memory 中的 16 位元素按转置规则分发到寄存器，配合 MMA 所需的 B 分工。逻辑转置视图和硬件的转置加载共同完成这段衔接。

源码中的 `sVtNoSwizzle` 用来创建 B 寄存器 fragment，提供所需的逻辑形状；真正加载 shared memory 的源视图来自带 Swizzle 的 `sVt`。两者在这里承担不同的工作。

最终，`gemm_A_in_regs` 使用已经在寄存器里的 `tCrP`，加载 V 的 B fragment，执行本轮 `tCrO += P @ V`。

FP16 输入和 FP32 累加能减少存储并使用 Tensor Core，但权重在进入第二个 GEMM 前会舍入为 FP16。因此上述 online softmax 公式在实数运算中等价，实现结果与完整 FP32 attention 通常有小幅数值差异。

## 把异步加载放回整条时间线

现在各个计算步骤都有了含义，可以读 kernel 中的调度顺序了。先看非 causal 的主循环：

```text
循环前：异步加载 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[i]` 时计算 `QK[i]`，加载 `K[i+1]` 时计算 softmax 和 `P[i]V[i]`。

源码用下面几种操作分别控制它们：

| 操作 | 作用 |
|---|---|
| `flash::copy`，使用 `cp.async` Atom | 发起 global memory 到 shared memory 的异步拷贝 |
| `cute::cp_async_fence()` | 对应 `cp.async.commit_group`，将此前未提交的异步拷贝组成一组 |
| `flash::cp_async_wait<0>()` | 等待当前线程此前提交的拷贝组全部完成 |
| `__syncthreads()` | 让 block 内线程会合，安全使用协作加载的数据或复用缓冲区 |

`fence` 不表示“现在才开始拷贝”，也不表示拷贝已经完成。`wait` 等待当前线程的异步操作，后面还需要 block 同步，才能消费其他线程加载的部分。

本实现的 shared storage 分别只有一份 Q、一份 K、一份 V tile：K 的旧数据在本轮 QK 完成后被下一轮 K 覆盖，V 的旧数据在下一轮加载 V 时被覆盖。这里通过两个 GEMM 的先后关系安排预取，源码没有为 K、V 各自开两份 ping-pong 缓冲区。

这些同步点除了等待数据，也保护缓冲区的复用。例如，等待 V 后的 block 同步保证所有线程已经完成本轮 QK，随后才可以覆盖 `sK`。修改调度顺序时，需要同时跟踪“数据是否已到达”和“旧数据是否已被所有使用者读完”。

## causal mask 怎样找到当前元素的位置

开启 causal 后，第 q 行只能访问 `key <= q` 的元素。当前配置下，源码分两层处理。

首先少遍历整块位于未来的 K、V。对当前 Q block：

```cpp
int seqlen_end = (m_block + 1) * kBlockM;
int n_block_max = cute::ceil_div(seqlen_end, kBlockN);
```

本例的第 0 个 Q block 只访问 `K[0:64]`，第 1 个 Q block 访问两个 K tile。

其次，在与对角线相交的分数块内，将 `key > query` 的位置设为 `-INFINITY`。这些位置经过指数函数后权重为 0。

难点在于，当前线程的 fragment 坐标不是整块矩阵的行列坐标。`mask_within_nblock` 根据当前 MMA 分工恢复位置。对本例，可以写成：

```text
warp = threadIdx.x / 32
lane = threadIdx.x % 32

query = m_block*64 + warp*16 + lane/4 + 8*mi
key   = nbi*64 + (lane%4)*2 + j + 8*nj
```

其中 `mi=0、1` 表示当前 lane 的两行，`j=0、1` 表示一对相邻列，`nj=0～7` 表示 N 方向重复。比较 `key > query`，就能决定是否屏蔽该值。

例如 warp 0、lane 1 持有第 0、8 行的第 2、3、10、11 等列。对 query 0，这些 key 全在未来；对 query 8，key 2、3 可以保留，key 10、11 必须屏蔽。

这份位置公式依赖当前的 TiledMMA。换 warp 排布或 MMA 指令后，需要一起修改。它实现的是本例等长 self-attention 的 causal 语义，也没有代替序列末尾的边界检查。

## 最后一轮之后，归一化并写回结果

循环结束时，`tCrO` 保存的是 U，`rRowSum` 保存的是 l。源码先创建输出的行列视图，再给每行乘上 `1/l`，得到最终 O。

同时计算：

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

也就是这一行缩放后分数的 LogSumExp。这里的 m 是没有乘 s 的 QK 最大值，所以公式里还要补乘 s。

LSE 比完整的注意力矩阵小得多，可用于 backward 时重建 softmax 所需的信息；当前仓库的这条路径只实现 forward。由于一行的归约结果在四个 lane 中重复保存，写回 LSE 时还要选出代表 lane，避免四个线程同时写同一行。源码用 identity tensor 的矩阵坐标判断谁负责写。

### 输出为什么要经过 shared memory

MMA 输出的值仍分散在各个线程中。直接沿这个分工写 global memory，不一定适合连续、向量化的存储。源码先将结果转换成 FP16，再通过 shared memory 调整写回分工：

```text
FP32 输出累加器
  -> FP16 的 tCrO_fp16
  -> shared memory 的 sO
  -> 按 global memory 写回分工读取到各线程的临时寄存器
  -> global memory 的 O
```

其中 `sO` 复用已经不再使用的 `sQ` 空间。写入 shared memory 的那一步使用：

```cpp
Tensor tCrO_fp16 = flash::convert_type_f32_to_f16(tCrO);

auto r2s_copy_o = make_tiled_copy_C(
    typename Kernel_traits::SmemCopyAtomO{}, mma);
auto r2s_thr_copy_o = r2s_copy_o.get_thread_slice(threadIdx.x);

Tensor tXrO = r2s_thr_copy_o.retile_S(tCrO_fp16);
Tensor tXsO = r2s_thr_copy_o.partition_D(sO);
cute::copy(r2s_copy_o, tXrO, tXsO);
```

这里改成 `retile_S`，因为已有的寄存器 fragment 是 Copy 的 source。它与前面的 `retile_D` 一样，只创建共享存储的访问视图。

block 同步之后，再按 `GmemTiledCopyO` 的分工从 `sO` 读入临时寄存器，写到 global memory。shared memory 在这里使计算分工和写回分工得以衔接。

## 自己核对一次：先打印布局，再比较数值

读完这些代码，最值得先试的是固定一组小尺寸，观察 layout 和结果。直接扩大维度、切换类型，会同时引入编译、边界和数值问题。

### 一个可以独立运行的布局打印程序

下面的程序在 CPU 上打印布局，不执行 GPU 加载。它显式指定 SM80 MMA，避免项目 traits 根据 `__CUDA_ARCH__` 在 host 编译时选到另一条分支。

项目根目录的 [`layout_probe.cu`](layout_probe.cu) 已包含下面的完整程序：

```cpp
#include <cute/tensor.hpp>
#include <cstdio>
using namespace cute;

int main() {
    using TiledMma = TiledMMA<
        MMA_Atom<SM80_16x8x16_F32F16F16F32_TN>,
        Layout<Shape<_4,_1,_1>>,
        Tile<_64,_16,_16>>;
    TiledMma mma;
    auto thr_mma = mma.get_slice(0);

    // 只借用地址构造布局，不访问这些数据。
    half_t storage[64 * 64]{};
    auto sA = make_tensor(
        make_smem_ptr(storage),
        make_layout(Shape<_64,_64>{}, Stride<_64,_1>{}));
    auto sB = sA;  // 本例只打印布局，可以共用这份占位存储。

    auto tCrA = thr_mma.partition_fragment_A(sA);
    auto tCrB = thr_mma.partition_fragment_B(sB);
    auto tCrC = partition_fragment_C(mma, Shape<_64,_64>{});

    auto s2r_copy_a = make_tiled_copy_A(
        Copy_Atom<SM75_U32x4_LDSM_N, half_t>{}, mma);
    auto s2r_copy_b = make_tiled_copy_B(
        Copy_Atom<SM75_U32x4_LDSM_N, half_t>{}, mma);
    auto s2r_thr_copy_a = s2r_copy_a.get_slice(0);
    auto s2r_thr_copy_b = s2r_copy_b.get_slice(0);
    auto tXrA = s2r_thr_copy_a.retile_D(tCrA);
    auto tXrB = s2r_thr_copy_b.retile_D(tCrB);

    print("tCrA: "); print(tCrA.layout()); print("\n");
    print("tCrB: "); print(tCrB.layout()); print("\n");
    print("tCrC: "); print(tCrC.layout()); print("\n");
    print("tXrA: "); print(tXrA.layout()); print("\n");
    print("tXrB: "); print(tXrB.layout()); print("\n");
}
```

在项目根目录中编译、运行。下面使用本文核对布局时的本机 CUTLASS 目录：

```bash
nvcc -std=c++17 -I/home/x/workspace/cutlass/include layout_probe.cu -o /tmp/layout_probe
/tmp/layout_probe
```

换机器时，将 include 路径改成自己的 CUTLASS 目录。省略静态常量前缀后，关键输出为：

```text
tCrA: ((2,2,2),1,4):((1,2,4),0,8)
tCrB: ((2,2),8,4):((1,2),16,4)
tCrC: ((2,2),1,8):((1,2),0,4)
tXrA: ((8,1),1,4):((1,0),0,8)
tXrB: (((4,2),1),4,4):(((1,16),0),32,4)
```

冒号前是 shape，冒号后是 stride。最后一行的 `(4,2):(1,16)` 正好说明：一组八个 B half，来自四个连续槽位，以及相隔 16 的另一组四个槽位。

想进一步核对“当前 lane 持有矩阵的哪个位置”，可以创建 `make_identity_tensor(Shape<_64,_64>{})`，然后调用 `thr_mma.partition_A/B/C`。identity tensor 的值就是坐标，得到的分区会直接显示 `(行,列)`，比从指针偏移猜测更方便。

### 在已有 CUDA 扩展环境中验证 forward

仓库的 `make build` 会构建 Python 扩展，依赖 CUDA、PyTorch 和完整的 CUTLASS 头文件。首次使用需要先初始化仓库记录的子模块：

```bash
git submodule update --init --recursive
cd flash_attention_cutlass
make build
```

这会使用项目锁定的子模块版本；若有意改用另一版 CUTLASS，应同时检查构建配置和 API。前面的独立布局程序可以直接指向另外一份 CUTLASS 检出目录。

扩展安装后，在项目根目录运行 [`flash_attention_example.py`](flash_attention_example.py)，用 FP16、连续存储、L 为 64 倍数的等长输入先做小规模比较：

```bash
python flash_attention_example.py
```

脚本的完整内容如下：

```python
import math
import torch
from attention_cutlass import flash_attention_v2_cutlass

torch.manual_seed(0)
B, H, L, D = 1, 1, 128, 64
q, k, v = [
    torch.randn(B, H, L, D, device="cuda", dtype=torch.float16)
    for _ in range(3)
]
scale = 1.0 / math.sqrt(D)

for causal in (False, True):
    out, lse = flash_attention_v2_cutlass(q, k, v, causal, scale)

    score = (q.float() @ k.float().transpose(-1, -2)) * scale
    if causal:
        mask = torch.ones(L, L, device="cuda", dtype=torch.bool).triu(1)
        score.masked_fill_(mask, -torch.inf)
    ref_out = torch.softmax(score, dim=-1) @ v.float()
    ref_lse = torch.logsumexp(score, dim=-1)

    torch.testing.assert_close(out.float(), ref_out, atol=3e-3, rtol=3e-3)
    torch.testing.assert_close(lse, ref_lse, atol=1e-4, rtol=1e-4)
    print("passed, causal =", causal)
```

这里显式传入 `1/sqrt(D)`。当前接口默认 scale 为 1；仓库 `test.py` 的示例则用了 `1/sqrt(SEQLEN)`，都不能直接当作常规 attention 的默认缩放方式。

本文的布局输出已用 CUTLASS `0b55a2f6` 编译核对，也逐项检查了全部 128 个线程的 C→A 元素坐标，以及 V 转置视图的地址对应关系。另从当前 `csrc` 提取原样的 forward 设备代码，绕过 PyTorch 绑定，在本机 GPU 上检查了 L=128、D=64/96/128，以及 L=192、D=32 的 causal、非 causal 两种模式。对这些 FP16 小规模随机输入，8 组检查的输出最大绝对误差低于 `3.22e-4`，LSE 最大绝对误差低于 `7.55e-7`。这些检查用于核对本文讨论的布局和计算路径；Python 扩展的完整构建、安装需在对应环境中另行验证。

## 当前教学实现有哪些边界

读源码时需要把“参数里有这个字段”和“kernel 已经实现这种能力”分开。当前版本保留了一些完整 FlashAttention 的接口和类型分派，但本文能沿代码确认的路径有明确限制：

| 项目 | 当前实现中需要留意的地方 |
|---|---|
| 序列尾块 | `flash::copy` 没有启用边界 predicate，输出和 LSE 写回也没有完整尾块保护。本文使用 64 的整数倍长度 |
| head dimension | 分派会向上选择 32 的倍数，但 kernel 用编译期 `kHeadDim` 作为行跨度，没有补齐任意 D 的 padding。本文选用与分派尺寸相等的 D |
| 布局、不同长度 | Q、K、V 使用同一个 head 起点偏移，行跨度也由模板固定。不能据此宣称已支持任意 stride、不同 Q/K 长度或 GQA/MQA |
| BF16 | 虽然有 BF16 MMA 分派，权重和输出转换仍写死为 `convert_type_f32_to_f16`。本文只讨论 FP16 路径 |
| CUDA stream | `run_flash_fwd` 接收 stream，但实际 launch 没有把它传入；Python 包装还执行设备同步 |
| backward | 返回 LSE，并不表示已经实现 backward |

这些限制也解释了为什么 `ceil_div`、模板分派、参数结构本身不足以证明一个输入已经受到完整支持。要扩展输入范围，需要沿加载、mask、计算、写回整条路径检查。

`standalone_src` 提供了无需 PyTorch 绑定的完整示例，但它和 `csrc` 是两份实现，参数结构、入口和默认配置并不完全同步。阅读本文时以 `csrc` 为准，切换到 standalone 后应重新核对相应配置。

## 回到开头的三行代码

现在再读它们：

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

第一行准备当前线程参与 MMA 所需的寄存器；第二行按加载指令的地址分工选择 shared memory 源；第三行给已有寄存器建立加载时使用的目的视图。随后调用 `copy`，数据才进入寄存器，再由 `gemm` 消费。

沿着同样的方法看整段 kernel，就能把那些长变量名落到具体工作上：固定一个 Q tile，流式读入 K、V；第一个 GEMM 产生分数，online softmax 保存可续算的状态，第二个 GEMM 把权重变成输出分子；最后归一化，将各线程的结果整理成适合写回的排列。

## 旧版命名与当前源码对照

本文、`csrc` 和 `standalone_src` 的 CuTe 实现已使用相同的命名。下表保留旧版名称，方便对照早期文章和代码；Q、K、V 等数学符号与 `partition_A/B/C` 这些 API 名称继续沿用。

| 旧版源码名称 | 当前源码与本文名称 | 对应的对象 |
|---|---|---|
| `Q`、`K`、`V`、`O` | `mQ`、`mK`、`mV`、`mO` | 当前 head 的完整矩阵视图 |
| `tiled_mma` | `mma` | TiledMMA 对象 |
| `gmem_tiled_copy_QKV` | `copy_qkv` | global→shared Copy |
| `gmem_thr_copy_QKV` | `thr_copy_qkv` | 上述 Copy 的线程切片 |
| `smem_tiled_copy_Q/K/V`、`smem_thr_copy_Q/K/V` | `s2r_copy_q/k/v`、`s2r_thr_copy_q/k/v` | Q、K、V 的 shared→register Copy 及线程切片 |
| `tSrQ`、`tSrK` | `tCrQ`、`tCrK` | MMA 的 Q、K 寄存器 fragment |
| `tSsQ` | `tXsQ` | Q 加载 Copy 的 shared memory 源视图 |
| `tSsK`、`tOrVt`、`tOsVt` | `tXsK`、`tCrVt`、`tXsVt` | K 的 Copy 源视图，以及 V 的 MMA、Copy 视图 |
| helper 内的 `tCrA_copy_view`、`tCrB_copy_view` | `tXrA`、`tXrB`；Q 的例子写作 `tXrQ` | 已有 fragment 的 Copy 视图 |
| helper 内的 `smem_tiled_copy_A/B`、`smem_thr_copy_A/B` | `s2r_copy_a/b`、`s2r_thr_copy_a/b` | 通用 A、B 加载 Copy 及线程切片 |
| helper 内的 `tCsA`、`tCsB` | `tXsA`、`tXsB` | 上述 Copy 的 shared memory 源视图 |
| helper 内的 `acc` | `tCrC` | 通用 GEMM 的累加器 |
| `rAccScore`、`rAccOut` | `tCrS`、`tCrO` | 分数和输出的累加器 |
| `scores` | `tCrS_rowcol` | 分数的行列视图 |
| `scores_max`、`scores_sum` | `rRowMax`、`rRowSum` | 每行最大值、分母的寄存器状态 |
| `tOrP` | `tCrP` | 第二个 GEMM 的 A fragment 视图 |
| `rO` | `tCrO_fp16` | 转换成 FP16、保留 MMA 排列的输出 fragment |
| `smem_tiled_copy_O`、`smem_thr_copy_O` | `r2s_copy_o`、`r2s_thr_copy_o` | 输出的 register→shared Copy 及线程切片 |
| `taccOrO`、`taccOsO` | `tXrO`、`tXsO` | 输出 Copy 的寄存器源、shared memory 目的视图 |
| `gmem_tiled_copy_O`、`gmem_thr_copy_O` | `copy_o`、`thr_copy_o` | 输出的 global memory 写回 Copy 及线程切片 |
| `LSE`、`lse` | `mLSE`、`rLSE` | LSE 的完整矩阵视图和寄存器状态 |
| `caccO`、`taccOcO`、`taccOcO_row` | `cO`、`tCcO`、`tCcO_row` | 输出元素坐标及其 MMA 分区、行坐标视图 |

`gQ/sQ`、`tQgQ/tQsQ`、`thr_mma`、`rP` 等本来就符合本文采用的命名方式，继续沿用。`rRowMax`、`rRowSum` 和 `_rowcol`、`_fp16` 后缀是本文用于说明状态、视图和类型的扩展，并非 CuTe API 的强制规则。

继续阅读源码时，可以按下面这几个入口定位：

| 想核对的内容 | 源码入口 |
|---|---|
| 当前分块和指令选择 | [`kernel_traits.h`](flash_attention_cutlass/csrc/kernel_traits.h)：`Flash_fwd_kernel_traits` |
| block 分配、两次 GEMM、调度和写回 | [`flash_attention.cu`](flash_attention_cutlass/csrc/flash_attention.cu)：`flash_attention_v2_cutlass_kernel` |
| QK、PV 的寄存器加载 | 同文件：`gemm_smem`、`gemm_A_in_regs` |
| online softmax 与布局衔接 | 同文件：`softmax_rescale_o`、`convert_layout_acc_rowcol`、`convert_layout_rowcol_Aregs` |
| 四个 lane 的行归约 | [`utils.h`](flash_attention_cutlass/csrc/utils.h)：`thread_reduce_`、`Allreduce` |
| CuTe Copy 分区和 retile | [`copy_atom.hpp`](https://github.com/NVIDIA/cutlass/blob/0b55a2f6/include/cute/atom/copy_atom.hpp) |
| TiledMMA 的分工和 tiling | [`mma_atom.hpp`](https://github.com/NVIDIA/cutlass/blob/0b55a2f6/include/cute/atom/mma_atom.hpp) |
| 实际 MMA、ldmatrix 指令 | [`mma_sm80.hpp`](https://github.com/NVIDIA/cutlass/blob/0b55a2f6/include/cute/arch/mma_sm80.hpp)、[`copy_sm75.hpp`](https://github.com/NVIDIA/cutlass/blob/0b55a2f6/include/cute/arch/copy_sm75.hpp) |
| 指令规定的线程和值分布 | [`mma_traits_sm80.hpp`](https://github.com/NVIDIA/cutlass/blob/0b55a2f6/include/cute/atom/mma_traits_sm80.hpp)、[`copy_traits_sm75.hpp`](https://github.com/NVIDIA/cutlass/blob/0b55a2f6/include/cute/atom/copy_traits_sm75.hpp) |

本文实现参考 [FlashAttention](https://github.com/Dao-AILab/flash-attention)。算法背景可继续读 [FlashAttention](https://arxiv.org/abs/2205.14135) 和 [FlashAttention-2](https://arxiv.org/abs/2307.08691) 论文；想系统补 CuTe 的 Layout 运算，可对照 [CuTe 官方教程](https://docs.nvidia.com/cutlass/media/docs/cpp/cute/index.html)。