CUTLASS 从零重写 FlashAttention：Split KV 并行优化实战

1. 为什么要折腾 Split KV 这件事

搞过 FlashAttention 推理优化的朋友大概都有体会，当 batch size 小、序列长度又特别长的时候，GPU 的 SM 利用率会掉得很难看。我最初跑一个 seq_len 到 8K 以上的场景，batch 只有 1 到 4，结果发现整卡占用率连 30% 都上不去，大部分 SM 都在空转。这个现象背后的原因其实不复杂：标准 FlashAttention 的并行维度是 batch × head，当这两个维度乘起来不够大时，能分出去的 thread block 数量就撑不满硬件。

Split KV 就是冲着这个痛点来的。它的核心思路是把 KV 序列维度再切一刀，让每个 attention head 的计算被拆到多个 thread block 上并行执行，最后再做一个跨 block 的 reduction 把结果合并回来。这样一来，并行度从 batch × head 变成了 batch × head × kv_split，SM 占用率能拉回来一大截。

我这次用 CUTLASS 从零重写 FlashAttention，第一站就选了 Split KV 作为切入点。原因有两个：一是它足够独立，可以单独验证正确性和性能；二是它涉及到的 reduction 逻辑、workspace 管理、kernel 拆分策略，都是后面做更复杂优化（比如 PagedAttention、连续 batching）的基础设施。把这部分吃透，后面再往上叠功能会顺很多。

这篇文章适合谁看？如果你已经写过或者读过基础版的 FlashAttention kernel，对 online softmax 那套流程不陌生，想进一步了解怎么在 CUTLASS 框架下做多 block 拆分和跨 block 归约，那这篇内容应该能帮到你。如果你还没接触过 FlashAttention，建议先把单 block 版本啃一遍再来看，不然有些细节会接不上。

2. Split KV 的整体设计思路拆解

2.1 标准 FlashAttention 的并行瓶颈在哪

先把问题说清楚。标准 FlashAttention 前向的并行粒度是每个 (batch, head) 对一个 thread block。假设你有 B 个 batch、H 个 head，那总共能开出的 block 数就是 B × H。在训练场景下，B 和 H 通常都不小，比如 B=32、H=32，那就是 1024 个 block，足够填满任何一张卡的 SM。但推理场景完全是另一回事——在线服务里 batch 经常是 1 到 8，head 数取决于模型结构，7B 模型常见的是 32 个 head。乘起来也就 32 到 256 个 block，而一张中高端 GPU 的 SM 数量轻松上百。这就导致大量 SM 闲置，延迟下不去。

更麻烦的是，当 seq_len 很长时，单个 block 要处理的 KV 长度也跟着涨。比如 seq_len=16K，一个 block 要遍历 16K 个 KV 位置做 online softmax 更新，计算量集中在一个 block 上，其他 block 早就跑完了在等它。这种负载不均衡进一步拉低了整体效率。

2.2 Split KV 的核心思想：把序列维度切开

Split KV 的做法很直接：既然 (batch, head) 维度不够分，那就把 KV 序列维度也拿来分。具体来说，对于每个 (batch, head) 对，把它的 KV 序列切成 S 段，每段由一个独立的 thread block 负责计算局部的 attention 输出和局部 softmax 统计量（running max 和 running sum）。等所有 S 个 block 都算完之后，再启动一个 reduction kernel，把这 S 份局部结果按照 online softmax 的合并公式归并成最终输出。

这里的关键在于，FlashAttention 的 online softmax 本身就是可结合的。也就是说，你先算前一半 KV 的 softmax 结果，再算后一半，然后把两份结果合并，和一次性算完整序列的结果在数学上是等价的。这个性质是 Split KV 能成立的理论基础。合并公式大概长这样：

• 设两份局部结果分别为 (O1, m1, l1) 和 (O2, m2, l2)，其中 m 是 running max，l 是 running sum of exp，O 是未归一化的加权输出。
• 合并后的 max：m = max(m1, m2)
• 合并后的 sum：l = l1 * exp(m1 - m) + l2 * exp(m2 - m)
• 合并后的输出：O = (O1 * l1 * exp(m1 - m) + O2 * l2 * exp(m2 - m)) / l

这个公式在 reduction kernel 里会被反复用到，后面讲实现的时候会展开。

2.3 为什么选 CUTLASS 而不是手写 CUDA

有人可能会问，Split KV 的逻辑用纯 CUDA 也能写，为什么要套 CUTLASS？我的考虑是这样的：CUTLASS 提供了一套比较完整的 tile 抽象和 pipeline 原语，尤其是它的cute库对 layout 和 copy 的抽象，能让你在写 kernel 的时候少操心底层的 swizzle、bank conflict 这些琐事。另外 CUTLASS 的 epilogue 和 mainloop 分离结构，天然适合把 attention 的 QK^T 和 PV 两个 GEMM 拆开写，代码组织上更清晰。

当然代价是学习曲线陡。cute 的 layout 代数不是一两天能玩明白的，我一开始也被Layout、Tensor、MMA_Traits这些概念绕得够呛。但一旦过了这个坎，后面改 kernel 的效率会高很多。而且 CUTLASS 的社区和文档相对完善，遇到问题查起来方便。

2.4 拆分策略的选择：按 KV 长度均分还是按 block 数均分

拆分的时候有个细节要决定：是把 KV 序列按长度均分，还是按 block 数量均分？这两种做法在负载均衡上表现不一样。

按长度均分就是每段 KV 长度一样，比如 seq_len=8192、split=4，那每段 2048。这种做法实现简单，但有个问题：FlashAttention 的计算量随 KV 长度线性增长，如果每段长度一样，那每个 block 的计算量理论上也一样，看起来是均衡的。但实际上因为 causal mask 的存在，靠前的 KV 段参与的计算更少（对于 causal attention），所以按长度均分会导致前面的 block 先跑完、后面的 block 拖后腿。

按 block 数均分则是让每个 split 负责相同数量的 KV block（比如每 128 个 KV 位置为一个 block），这样在 causal 场景下，靠前的 split 负责的 block 数虽然一样，但实际有效计算量更少。不过这种做法需要额外处理边界，实现起来稍微麻烦一点。

我这次先用了按长度均分的方案，因为第一版目标是跑通正确性，负载均衡的优化留到后面。实测下来在非 causal 场景下均分效果就挺好，causal 场景下确实有负载不均，但配合后面的 reduction 重叠还能接受。

3. 核心细节解析与实操要点

3.1 Workspace 的分配与管理

Split KV 相比标准 FlashAttention 多了一个绕不开的东西：workspace。因为每个 split block 算出来的局部结果不能直接丢掉，得先写到全局内存里，等 reduction kernel 来读。所以你需要一块显存来存这些中间结果。

具体要存什么？每个 (batch, head, split) 三元组需要存三样东西：局部的未归一化输出 O（形状是 [B, H, S, D]，D 是 head dim）、局部的 running max m（形状 [B, H, S]）、局部的 running sum l（形状 [B, H, S]）。注意 O 是未归一化的，也就是还没除以 l 的版本，这样 reduction 的时候才能正确合并。

workspace 的大小计算要小心。以 B=8、H=32、S=4、D=128、fp16 为例，O 的大小是 8 × 32 × 4 × 128 × 2 字节 = 256KB，m 和 l 各是 8 × 32 × 4 × 4 字节 = 4KB。加起来不到 300KB，对显存来说微不足道。但如果 B 和 H 很大、S 也开得大，比如 B=64、H=64、S=16、D=128，那 O 就是 64 × 64 × 16 × 128 × 2 = 16MB，也不算大。真正需要注意的是 workspace 的对齐和访问模式，后面会讲。

提示：workspace 建议用at::empty在 PyTorch 侧预分配好传进来，不要在 kernel 里动态申请。一是避免频繁 cudaMalloc 的开销，二是方便做内存复用。

3.2 局部 softmax 统计量的计算

每个 split block 内部的计算流程和标准 FlashAttention 基本一致，只是 KV 的遍历范围变成了自己负责的那一段。核心还是 online softmax 那套：维护 running max 和 running sum，每处理一个 KV block 就更新一次。

这里有个容易踩的坑：初始值的设定。标准 FlashAttention 里 running max 初始化为负无穷，running sum 初始化为 0。但在 Split KV 里，如果某个 split 负责的 KV 段全是 masked 掉的（比如 causal 场景下最靠后的 split 在某些 query 位置上没有任何有效 KV），那这个 split 的 running max 会一直是负无穷，running sum 一直是 0。这种 block 在 reduction 的时候要特殊处理，否则 exp(-inf - m) 会算出 NaN。

我的处理方式是在 reduction kernel 里加一个判断：如果某个 split 的 l 是 0，就跳过它，不参与合并。这样即使有全 masked 的 split 也不会污染结果。

3.3 Reduction kernel 的设计

Reduction kernel 的任务是把 S 个 split 的局部结果合并成最终输出。它的并行粒度是每个 (batch, head) 对一个 block，block 内部再按 head dim 分线程。每个线程负责输出 O 的一个或多个元素，需要遍历所有 S 个 split，按照前面说的合并公式累加。

这里有个性能上的考量：reduction kernel 要读 S 份 O、m、l，如果 S 比较大，读的量不小。为了减少全局内存访问，可以把 m 和 l 先读到 shared memory 或者寄存器里，因为它们的量很小（每个 split 就一个标量）。O 的读取则没法避免，但可以用向量化 load 来提升带宽利用率。

另一个细节是数值稳定性。合并公式里有 exp(m_i - m) 这一项，当 m_i 远小于 m 时，这一项会趋近于 0，这是正常的。但如果 m 本身是负无穷（所有 split 都 masked），那就要直接输出 0。我在 kernel 里加了一个全局的 max 检查，如果合并后的 m 还是负无穷，就直接把输出置零。

3.4 CUTLASS 中的 layout 选择

用 CUTLASS 写 kernel，layout 的选择直接影响性能和代码复杂度。对于 attention 的 Q、K、V，我用的 layout 是 row-major，也就是 [seq_len, head_dim] 这种排布。这样每个 thread 处理一行 query 的时候，访问 K 和 V 的列方向是连续的，方便做向量化。

但 CUTLASS 的 MMA 指令对 layout 有要求，比如 sm80 的mma.sync要求操作数是特定的 fragment layout。所以在把数据喂给 MMA 之前，需要做一次 layout 转换。CUTLASS 的cute库提供了make_tiled_mma和make_fragment_like这些工具来帮你处理，但具体怎么用还是得看例子。

我踩过的一个坑是：一开始没注意 K 的 layout，直接按 row-major 喂进去，结果 MMA 算出来的结果全错。后来查了半天才发现 K 需要转置，因为 QK^T 里的 K 是以列向量形式参与计算的。这个转置可以在 load 的时候用cute的make_transposed_layout来做，也可以在 shared memory 里做 swizzle。我选了前者，代码更干净。

4. 实操过程与核心环节实现

4.1 环境准备与依赖

我用的环境是 CUDA 12.1 + CUTLASS 3.x + PyTorch 2.1。CUTLASS 直接从 GitHub 拉最新版就行，它主要是 header-only 的，不需要编译安装，把 include 路径加到编译选项里即可。PyTorch 这边主要是用来做正确性对比和性能 benchmark 的，kernel 本身是纯 CUDA/CUTLASS 写的。

编译的时候要注意架构参数。我用的卡是 sm80 架构，所以编译选项里加了-gencode arch=compute_80,code=sm_80。如果你用的是更新的架构，记得改成对应的 compute capability。另外 CUTLASS 有些特性在不同架构上支持程度不一样，比如 TMA 在 sm90 上才有，写的时候要注意条件编译。

4.2 Kernel 拆分：split kernel 和 reduce kernel

整个实现拆成两个 kernel：一个是split_kv_attention_kernel，负责每个 split 的局部计算；另一个是split_kv_reduce_kernel，负责合并。两个 kernel 串行执行，中间通过 workspace 传递数据。

split kernel 的 grid 是 (B × H × S)，每个 block 处理一个 (batch, head, split) 三元组。block 内部的线程组织参考了标准 FlashAttention 的做法：按 query 的行来分，每个 warp 负责若干行，warp 内部再按 head dim 分 lane。

reduce kernel 的 grid 是 (B × H)，每个 block 处理一个 (batch, head) 对，遍历所有 S 个 split 做合并。block 内部的线程按 head dim 分，每个线程负责输出的一部分元素。

4.3 Split kernel 的关键代码结构

split kernel 的 mainloop 大概是这样的：

// 伪代码，展示结构 for (int kv_block = kv_start; kv_block < kv_end; kv_block += BLOCK_KV) { // load K block load_k_block(K, kv_block, k_shared); // compute QK^T gemm_qk(Q_frag, k_shared, s_frag); // apply mask if causal apply_mask(s_frag, kv_block, q_pos); // online softmax update update_softmax(s_frag, running_max, running_sum, scale); // load V block load_v_block(V, kv_block, v_shared); // compute PV gemm_pv(p_frag, v_shared, o_frag); } // write O, m, l to workspace write_workspace(o_frag, running_max, running_sum, workspace, bid);

这里kv_start和kv_end是根据 split id 算出来的。假设 seq_len 是 N，split 数是 S，那第 i 个 split 负责的 KV 范围是[i * N / S, (i + 1) * N / S)。注意边界处理，最后一个 split 要保证覆盖到 N。

update_softmax是核心，它做的事情是：先算当前 block 的局部 max，和 running max 比较得到新的 max，然后用 exp 缩放之前的 running sum 和 O，再加上当前 block 的贡献。这部分和标准 FlashAttention 完全一样，我就不展开了。

4.4 Reduce kernel 的合并逻辑

reduce kernel 的核心是遍历 S 个 split，按合并公式累加。代码结构大概是这样：

// 伪代码 float m_final = -inf; float l_final = 0; float o_acc[D] = {0}; for (int s = 0; s < num_splits; s++) { float m_s = load_m(workspace, bid, s); float l_s = load_l(workspace, bid, s); if (l_s == 0) continue; // skip masked split float m_new = max(m_final, m_s); float scale_old = exp(m_final - m_new); float scale_new = exp(m_s - m_new); l_final = l_final * scale_old + l_s * scale_new; for (int d = 0; d < D; d++) { o_acc[d] = o_acc[d] * scale_old + load_o(workspace, bid, s, d) * scale_new; } m_final = m_new; } // normalize and write output for (int d = 0; d < D; d++) { output[bid][d] = o_acc[d] / l_final; }

这里有个优化点：scale_old和scale_new的计算涉及 exp，而 exp 在 GPU 上不算便宜。如果 S 不大（比如 4 或 8），这个开销可以接受。但如果 S 很大，可以考虑用查表或者近似。不过实测下来 S 一般不会超过 16，所以直接算 exp 就行。

4.5 参数选择与调优

Split 数 S 的选择是个 trade-off。S 太小，并行度不够，SM 填不满；S 太大，reduction 的开销和 workspace 的读写量都会增加。我的经验是 S 取到让 B × H × S 刚好能填满 SM 数量再稍微多一点比较合适。比如卡有 108 个 SM，B × H = 64，那 S 取 2 到 4 就够；如果 B × H = 16，那 S 可能要取 8 甚至 16。

BLOCK_KV 的大小也影响性能。太小的话，每次 load 的数据量不够，带宽利用率低；太大的话，shared memory 放不下，或者 register 压力大。我一般取 64 或 128，具体看 head dim 和 shared memory 容量。

还有一个容易忽略的参数是 reduction kernel 的 block size。如果 head dim 是 128，那 reduce kernel 的 block 至少要 128 个线程才能覆盖所有维度。我一般取 128 或 256，让每个线程处理 1 到 2 个元素。

5. 常见问题与排查技巧实录

5.1 正确性对不上：从哪查起

第一次跑通的时候，输出和 PyTorch 的参考实现差得离谱，数值完全不对。排查这类问题我一般按这个顺序来：

先查 mask 逻辑。causal mask 的边界最容易出错，尤其是 split 之后每个 block 的 query 位置和 KV 范围的对应关系。我当时的 bug 就是 split 之后 KV 的起始位置算错了，导致 mask 应用到了错误的位置。

再查 softmax 的数值稳定性。如果 running max 的更新逻辑有问题，exp 可能会溢出或者下溢。可以在 kernel 里加一些 printf 调试，把中间结果打出来和 CPU 参考对比。

最后查 reduction 的合并公式。合并公式里的 exp 缩放很容易写反，比如把exp(m_old - m_new)写成了exp(m_new - m_old)。这种错误不会导致 NaN，但结果会偏。

5.2 性能不升反降：可能的原因

Split KV 理论上应该提升性能，但如果你实测发现比标准版还慢，可能是这几个原因：

一是 S 选得太大，reduction 的开销超过了并行度带来的收益。可以试着把 S 调小，看性能曲线怎么变。

二是 workspace 的访问模式不好，导致全局内存带宽成为瓶颈。检查一下 O 的写入和读取是不是连续的，有没有做向量化。

三是 split kernel 内部的负载不均。如果是 causal attention，靠前的 split 计算量小，靠后的 split 计算量大，导致整体被最慢的 split 拖累。这种情况可以考虑按 block 数均分而不是按长度均分。

5.3 常见问题速查表

5.4 几个我踩过的坑

第一个坑是 workspace 的对齐。我一开始用at::empty分配 workspace，没注意对齐，结果在某些尺寸下访问会触发 misaligned access，性能掉得厉害。后来改成用at::empty加上手动对齐到 128 字节，问题就解决了。

第二个坑是 reduction kernel 的线程组织。我一开始让每个线程负责一个 head dim 元素，结果 head dim=128 的时候只有 128 个线程，block 利用率不高。后来改成每个线程负责多个元素，用循环展开，性能提升明显。

第三个坑是 split kernel 里 shared memory 的复用。K 和 V 的 shared memory 如果分开申请，可能会超出容量限制。我后来让它们复用同一块 shared memory，因为 K 用完就可以被 V 覆盖，这样能省一半 shared memory。

提示：调试 Split KV 的时候，建议先用小尺寸（比如 seq_len=128、split=2）跑通正确性，再逐步放大。大尺寸下调试信息很难定位。

6. 后续可以继续挖的方向

Split KV 跑通之后，我下一步打算做几件事。一是把负载均衡优化一下，试试按 block 数均分的方案，看 causal 场景下能提升多少。二是把 reduction kernel 和 split kernel 做 overlap，用 stream 或者 cooperative groups 让 reduction 在部分 split 算完的时候就开始，减少串行等待。三是把 Split KV 和 PagedAttention 结合起来，因为实际推理场景里 KV cache 往往是分页存储的，这两者结合才能落地到生产环境。

另外我还在考虑把 split 的粒度做得更细，比如支持动态 split——根据当前 batch 和 seq_len 自动决定 S 的大小，而不是写死。这样在不同负载下都能拿到比较好的性能。

这些方向每一个都够写一篇单独的总结，等我把它们都跑通了再回来补。如果你也在做类似的优化，欢迎交流踩坑经验。