有一次我在 A100 上分析长文本推理的 kernel 耗时，看到 attention 内核只吃掉了 GPU 15% 的算力。第一反应是显存带宽不够，但算了算 HBM 流量之后发现不是。真正的原因是 并行度不够：GPU 有一百多个 SM，而当时的 batch×head 只有十几个，大部分 SM 在排队等活。
这个现象在推理场景特别常见：batch 小、序列长、用了 GQA 之后 head 数也少。FlashAttention 默认把 batch×head 分给 thread block，已经用不满 GPU。于是就有了两个进阶策略：split-K 和 split-V。这两个词只差一个字母，但拆的东西、付的代价完全不同。我最早也以为它们是同一种策略的两个叫法，直到把 FlashAttention-2 论文 的图 3 和 FlashDecoding 的实现 对照起来看，才算想通。
attention 到底有哪些地方可以拆
注意力计算就是三个矩阵：Q [M, D]，K [N, D]，V [N, D]，输出 O [M, D] = softmax(QK^T) V。最自然的拆法是拆 M（query 序列），也就是 split-Q：每个 thread block 负责一块 query，遍历全部 K/V，最后写回。batch 和 head 本身就是并行维度，不用你额外操作。
问题在于，decode 阶段一个 token 只产生一个 query，M=1，split-Q 就没得拆了。剩下的潜力维度只有两个：K/V 共享的序列长度 N，和 V 的列维度 D。前者是 split-K，后者是 split-V。
split-K：拆 K/V 序列，最后必须全局归约
先把 N 维度切开。假设切成 C 段，每个 CTA 拿一块 query 和一段 KV，各自做一轮 FlashAttention。问题立刻出现：softmax 不是逐元素操作，它要先在整个序列上算归一化分母。每个 CTA 只看到一段 keys，它求出的 softmax 是局部的，不能直接写回输出。
这件事的解法是 online softmax。每个 CTA 在遍历自己那一段 KV 时，维护 running max m 和 running sum l，得到一个“未彻底归一化”的局部输出。
接下来每个 split 手里有 (m_i, l_i, O_i)，要做一次跨 CTA 的归约：
- 找出全局最大值 m_total = max(m_i)；
- 把每段的 exp(m_i – m_total) 作为权重，累加所有 l_i 和 O_i；
- 最后用总权重做一次除法，得到最终输出。
这一步在实现上可以用原子加直接累加到全局 buffer，也可以另起一个小 kernel 做 reduce。FlashDecoding 正是用 split-K 把很长的 KV cache 切成多段，让每个 SM 负责一段，从而填满 GPU。
代价有三个。第一，多一次全局归约，不管用原子操作还是 reduce kernel，都有额外开销和等待。第二，同一个 Q 块被多个 CTA 重复读取，如果 Q 很大，HBM 读取量会上去。第三，用原子操作做浮点累加时，累加顺序不确定，同一个输入跑两次可能得到逐 bit 不同的结果。
split-V：拆 V 的列，让每个 warp 各算各的
Split-K 是 CTA 之间的并行。再往下一层，一个 CTA 内部的 warps 怎么分活？FlashAttention-1 的分法是按 Q 的行：每个 warp 负责几行 query，大家共享 KV，每算完一块 query-key 分数，就要在 warp 之间同步一次 running max/sum。同步多了，性能就下来了。
FlashAttention-2 换了一刀：按 V 的列分。head_dim 通常是 64 或 128，比如 128 列分给 4 个 warp，每个 warp 负责 32 个输出列。每个 warp 自己保存一份完整的 Q 块，独立遍历全部 KV 块，独立维护 running max/sum，最后直接把自己那 32 列写回。因为每个 warp 看到的是完整 keys，它的统计量天然是全局的，完全不需要跨 warp 归约。
代价也很直白：每个 warp 都要完整算一遍 QK^T。4 个 warp 就意味着 QK^T 在同一个 CTA 里重复算了 4 遍。但 FlashAttention-2 认为这笔买卖划算：矩阵乘在 GPU 上是计算密集型的，多算几遍小矩阵，比一遍遍同步、在共享内存里搬来搬去要快得多。这也是 FlashAttention-2 比 v1 快的一个重要原因。
这等于用“多算几遍矩阵乘”去换“少同步几次”。对 GPU 来说，矩阵乘是它最擅长的业务，而跨线程同步是所有人都想躲的事。
一张表说清两个策略的取舍
| 对比点 | Split-K | Split-V |
|---|---|---|
| 拆哪个维度 | K/V 共用的序列长度 N | V 的列维度（head_dim） |
| 并行层级 | CTA / SM 之间 | CTA 内部 warp 之间 |
| 主要额外开销 | 全局归约 + 重复读取 Q | 每个 warp 重复计算 QK^T |
| 需要归约吗 | 必须 | 不需要 |
| 典型适用 | batch×head 太少、KV 超长（长上下文解码） | head_dim 适中、要减少 warp 同步时 |
| 代表实现 | FlashDecoding、FA2/FA3 的 split_k 模式 | FlashAttention-2 默认 warp 划分 |
实际工程里怎么选
这两个策略不是对立的。CTA 之间可以 split-K，CTA 内部可以 split-V，两者叠加，就是三层并行：batch/head → 序列长度 → V 的列。FlashAttention-2 的默认实现里，CTA 之间按 Q 拆分，CTA 内部按 V 拆列；当 head 数不够时，再打开 split-K 模式。
一个粗略的决策路径：写 kernel 时先算一下 batch × head 和你的 GPU SM 数，如果前者已经填满，用默认 split-Q 就好；如果差一个数量级，优先尝试 split-K；如果你自己写 CUDA/Triton 内核，发现 warp 同步占了大量时间，就想想 FlashAttention-2 的 split-V 思路。
边界也很清楚：Split-K 的原子归约会让结果失去确定性，调试、或者要做逐 bit 对比的场合不友好；Split-V 在 head_dim 很小时不划算，因为重复计算的分数太多，而且拆成列也不够分给 warp。还有一个更隐蔽的坑：这些策略看起来是在“省时间”，实际是在“拿另一种资源换并行度”，如果换来的并行度并没有让你跑满 SM，反而多花了通信和计算，那就是负优化。
我后来再看那个 15% 利用率的 kernel，用 split-K 把 KV cache 切成 8 段，GPU 占有率提到了 60% 多。不是因为这个技巧高级，而是终于把问题定位到了 SM 数量上。理解 split-K 和 split-V，本质上就是理解 GPU 上“并行度从哪里来、通信往哪里去”这件事。
原创文章，作者：guanweilu，如若转载，请注明出处：https://guanweilu.cn/article/499.html