1. 从 PyTorch 的 SDPA 到裸 CUDA：为什么要自己写 MMHA

很多人第一次接触 masked multi head attention（下面统一简称 MMHA）都是在 PyTorch 里一行F.scaled_dot_product_attention搞定，或者干脆用nn.MultiheadAttention。但只要你的场景稍微偏一点——比如自回归推理的 decode 阶段、变长序列训练、或者要在自定义算子图里做融合——你就会发现框架封装的那层抽象开始漏风：要么是 causal mask 的广播方式和你预期不一致，要么是 batch 内 padding 位置参与了 softmax 归一化，要么是 decode 阶段 Q 长度只有 1 却还在按训练时的矩阵乘规模算，浪费大量算力。

MMHA 的本质其实不复杂，拆开就是四步：Q 乘 K 转置得到注意力分数、按缩放因子缩放、加 mask、softmax 归一化后乘 V。真正难的是把这四步在 GPU 上写得又快又对。我这次动手写 CUDA 版本，核心动机有三个：第一，想把 causal mask 和 padding mask 的处理逻辑完全握在自己手里，避免框架默认行为带来的隐性 bug；第二，想验证 split-K 和 GEMV 这两条路径在不同序列长度下的实际收益；第三，纯粹是想把 attention 的访存模式吃透，因为这块是后续做算子融合的地基。

这篇文章适合两类人看：一类是已经会写基础 CUDA kernel、想进阶到 attention 这种带 reduction 和 mask 的复杂算子的开发者；另一类是做推理引擎、需要自己实现 decode 阶段 attention 的工程师。我会把设计取舍、访存优化、mask 处理、split-K 的边界条件都讲清楚，代码片段可以直接参考。需要说明的是，下面涉及的硬件参数和性能数字都是我在自己机器上实测的量级，不同卡上会有差异，但优化思路是通用的。

2. MMHA 的计算拆解与 CUDA 线程映射思路

2.1 把 attention 公式翻译成可并行的维度

先把公式摆出来，标准的多头注意力是：

Attention(Q, K, V) = softmax(Q @ K^T / sqrt(d_head) + mask) @ V

其中 Q、K、V 的形状是[batch, num_heads, seq_len, head_dim]。masked 版本的关键在于那个mask项：causal mask 让位置 i 只能看到位置 j <= i，padding mask 则把无效 token 位置的分数置为负无穷。

并行化的第一刀切在哪里，决定了整个 kernel 的形态。最自然的切法是：每个 (batch, head) 组合是一个独立的注意力矩阵，形状[seq_len, seq_len]。一个 block 负责一个 (batch, head)，block 内的线程协作完成这一整块矩阵的计算。这种映射的优点是 Q、K、V 的加载局部性好，一个 head 的 head_dim 通常只有 64 或 128，能塞进 shared memory。

但这里有个坑：当 seq_len 很大（比如 4096），一个[4096, 4096]的分数矩阵就是 64MB（fp32），根本放不下。所以实际实现里要么分块（tiling），要么把行维度再切给不同 block。我这次采用的是行分块策略：每个 block 处理若干行 query，K 和 V 按需从 global memory 流式加载。

2.2 线程映射：一个 warp 负责一行还是多行

线程映射有两种常见方案。方案 A：一个 warp 负责一行 query，warp 内的 32 个 lane 分摊 head_dim 和 seq_len 方向。方案 B：一个 block 负责多行，每行由若干线程协作。

我实测下来，方案 A 在 head_dim 是 32 的倍数时非常顺，因为一个 lane 正好处理一个 head_dim 元素，做 Q·K 点积时用 warp shuffle 做 reduction，不需要 shared memory 中转。但如果 head_dim 是 64 或 128，一个 lane 要处理 2 到 4 个元素，shuffle reduction 的轮数就上去了。

方案 B 更适合 head_dim 较大的情况。我最终选的是折中：block 大小设为 128 线程，处理 4 行 query，每行 32 个线程（正好一个 warp）。这样 head_dim 为 64 时，每个 lane 处理 2 个元素，reduction 用两次 shuffle 完成。下面是一个简化的线程映射示意：

// blockDim.x = 32 (一个 warp 处理一行 query) // blockDim.y = 4 (一个 block 处理 4 行) int row = blockIdx.x * blockDim.y + threadIdx.y; // query 行索引 int lane = threadIdx.x; // warp 内 lane

这个映射的好处是，Q 的一行加载进来后，K 的每一行都可以被这个 warp 复用，访存比接近最优。

2.3 缩放因子与数值稳定性

1/sqrt(d_head)这个缩放因子看起来不起眼，但它是数值稳定性的第一道防线。如果不缩放，当 head_dim 是 128 时，Q·K 的点积结果量级会到几十甚至上百，softmax 的指数运算直接溢出成 inf。缩放之后，点积结果落在合理区间，exp 不会爆。

我在实现里把缩放因子提前乘到 Q 上，而不是在点积之后乘。这样做的好处是点积循环里少一次乘法，而且 Q 只加载一次，缩放成本可以忽略。但要注意：如果你后续还要用原始 Q 做别的计算，就不能这么干，得保留一份未缩放的副本。

3. Causal Mask 与 Padding Mask 的融合处理

3.1 两种 mask 的语义差异

causal mask 是位置相关的：query 位置 i 只能 attend 到 key 位置 j <= i。padding mask 是内容相关的：某些 key 位置是 padding，所有 query 都不应该 attend 到它们。这两者的处理方式完全不同，但经常被混在一起写，导致 bug。

causal mask 可以在计算分数时直接跳过：当 j > i 时，根本不计算点积，直接把分数设为负无穷。padding mask 则必须在算出分数之后再加一个负无穷偏置，因为它是内容决定的，无法从索引推断。

我踩过的一个坑是：在 decode 阶段（query 长度只有 1），causal mask 其实退化成"全部可见"，因为唯一的 query 位置是序列末尾，它能看到所有历史 key。如果这时还按训练时的 causal 逻辑去算，会多一次无用的比较。所以我在 kernel 里加了一个is_decode的编译期分支，decode 时直接跳过 causal 判断。

3.2 负无穷的表示与 softmax 的配合

把 mask 位置设为负无穷，标准做法是用-INFINITY或一个足够大的负数（比如-1e9）。这里有个细节：如果你用-INFINITY，softmax 时exp(-inf - max)会得到 0，没问题；但如果整行都是 mask（理论上不该发生，但 padding 处理不当会出现），max 就是负无穷，-inf - (-inf)是 NaN。所以我在 softmax 之前加了一个保护：如果一行的 max 是负无穷，直接把整行输出置 0。

另一个细节是 fp16 场景。fp16 能表示的最小负数是 -65504，用-INFINITY会被截断成 -65504，虽然效果上够用，但不够干净。我建议在 fp16 kernel 里用一个显式的-1e4作为 mask 值，既安全又不会和真实分数混淆。

3.3 mask 在 shared memory 中的布局

如果 mask 是 causal 的，其实不需要显式存储，用索引比较即可。但如果是 padding mask，就需要把 mask 向量加载到 shared memory。我的做法是：每个 block 在加载 K、V 的同时，把对应的 padding mask 也加载进来，存成一个[seq_len]的 bool 数组。计算分数时，先查 causal，再查 padding，两者都通过才保留。

这里有个访存优化点：padding mask 的加载可以和 K 的加载合并成一次向量化读取，减少 global memory 事务数。我实测下来，把 mask 和 K 打包成float2或int4一起加载，能省下大约 5% 到 8% 的加载时间，序列越长收益越明显。

4. Split-K 与 GEMV：两条路径的取舍

4.1 什么时候该用 split-K

split-K 的核心思想是：当 seq_len 很大时，一个 block 算一整行 query 的注意力分数太慢，因为要遍历所有 key。把 key 维度切成若干段，每个 block 算一段的部分和，最后再做一个 reduction 合并。这样能提高并行度，尤其当 batch * num_heads 不够多、GPU 占用率上不去的时候。

但 split-K 有代价：部分和需要额外的 global memory 往返，而且 softmax 的 max 和 sum 必须在所有段都算完之后才能确定。所以 split-K 通常配合 online softmax（也叫 flash attention 的思路）来做，边算边更新 running max 和 running sum，避免存下整个分数矩阵。

我实测的结论是：当 seq_len 超过 2048 且 batch * num_heads 小于 SM 数量的 4 倍时，split-K 的收益明显，能到 1.5 倍以上。但如果 batch * num_heads 已经足够大，split-K 反而因为额外的 reduction 开销变慢。

4.2 GEMV 路径：decode 阶段的专用优化

decode 阶段 query 长度是 1，整个计算退化成矩阵向量乘（GEMV）：一个[1, head_dim]的 Q 乘以[seq_len, head_dim]的 K 转置，得到[1, seq_len]的分数。这时候用通用的 tiling kernel 是浪费的，因为 query 方向没有并行度。

GEMV 路径的优化重点是访存：K 和 V 是流式读取的，每个元素只用一次，所以要尽量用向量化加载（float4）和足够的 in-flight 请求来打满带宽。我的实现里，一个 block 处理一个 (batch, head)，block 内 256 个线程，每个线程负责一段 key 的部分点积，最后用 shared memory 做 block 内 reduction。

这里有个容易忽略的点：decode 阶段通常配合 KV cache，K 和 V 不是重新计算的，而是从 cache 里读的。cache 的布局如果是[batch, num_heads, seq_len, head_dim]，那么读取时 head 维度的 stride 很大，访存效率低。我建议把 cache 布局改成[batch, seq_len, num_heads, head_dim]，让同一个 head 的连续 key 在内存里挨着，这样向量化加载才顺。

4.3 两条路径的切换阈值

我在代码里设了一个阈值：当 query 长度 <= 4 时走 GEMV 路径，否则走通用 tiling 路径。这个 4 是实测出来的，因为 query 长度到 4 时，tiling 路径的并行度已经够用，而 GEMV 路径的 reduction 开销开始显现。当然这个阈值和 head_dim、SM 数量都有关，建议在自己的硬件上重新测一遍。

5. 访存优化与 shared memory 的实战细节

5.1 K 和 V 的加载策略

attention 的访存瓶颈在 K 和 V：每个 query 行都要遍历所有 key，如果每次都从 global memory 读，带宽直接爆炸。所以必须把 K 和 V 分块加载到 shared memory，让 block 内的多个 query 行复用。

我的做法是：block 处理 4 行 query，K 和 V 按 key 维度分块，每块大小设为 64 或 128。加载时用float4向量化，配合cp.async（如果架构支持）做异步拷贝，让计算和加载重叠。实测下来，双缓冲（double buffering）能把 K/V 加载的延迟隐藏掉大半，整体吞吐提升 20% 到 30%。

这里有个 shared memory 容量的问题：head_dim 为 128、块大小为 128 时，一块 K 就是128 * 128 * 4 = 64KB，双缓冲就是 128KB，已经接近很多卡的 shared memory 上限。所以块大小要根据 head_dim 动态调整，head_dim 大时块就小一点。

5.2 bank conflict 的规避

shared memory 的 bank conflict 是 attention kernel 里最常见的性能杀手。K 的布局是[block_size, head_dim]，如果 head_dim 是 32 的倍数，那么同一列的元素会落在同一个 bank 上，warp 内多个 lane 同时读同一列就会冲突。

解决办法是给 K 的每一行加 padding，让行 stride 变成head_dim + 1。这样同一列的元素就分散到不同 bank 了。代价是多占一点 shared memory，但换来的性能提升很值。我实测 padding 之后，K 的读取冲突从 8 路降到 1 路，kernel 整体快了 15% 左右。

5.3 输出写回的合并

softmax 之后乘 V 得到输出，输出形状是[seq_len, head_dim]。写回时要注意合并写：每个线程算出的输出元素应该连续，用float4写回。如果输出布局和线程映射不匹配，会出现大量非合并写，带宽利用率掉到 30% 以下。

我的做法是让 warp 内的 lane 负责 head_dim 方向，这样写回时 32 个 lane 正好写连续的 32 个元素（head_dim 为 32 时）或分两次写（head_dim 为 64 时）。这个映射和前面 Q 的加载映射是对称的，所以整个 kernel 的访存模式很规整。

6. 实测中暴露的坑与排查过程

6.1 结果对不上：从 mask 位置开始查

第一次跑通 kernel 后，我发现输出和 PyTorch 参考实现对不上，误差在 1e-2 量级，不算大但也不小。排查思路是：先关掉 mask，看纯 attention 是否一致。结果关掉 mask 后误差降到 1e-6，说明问题出在 mask 处理上。

进一步定位：把 causal mask 换成全零 mask，误差又回来了。这说明 causal 的判断逻辑有问题。最后发现是行索引和列索引搞反了：我写的是j > i时 mask，但实际应该是j > i时保留（因为 j 是 key 位置，i 是 query 位置，key 位置大于 query 位置才要 mask）。这个 bug 很隐蔽，因为矩阵是对称的，反了之后结果看起来"差不多"，但语义完全错了。

6.2 性能不达标：先看 occupancy 再看访存

kernel 正确之后，性能只有理论峰值的 30%。我先用 profiler 看了 occupancy，发现只有 25%，原因是 shared memory 用量太大，每个 SM 只能驻留一个 block。把块大小从 128 降到 64 后，occupancy 上去了，但性能反而降了，因为 K/V 的复用率下降。

这说明 occupancy 不是越高越好，关键看访存和计算的平衡。我最后把块大小定在 96，配合双缓冲，occupancy 到 50%，性能提升到峰值的 55%。这个调参过程没有捷径，就是反复试。

6.3 长序列下的数值溢出

序列长度到 4096 时，fp16 路径开始出现 inf。原因是 softmax 的 max 减法在 fp16 下精度不够，running max 更新时出现舍入误差累积。解决办法是在 softmax 内部用 fp32 累加，只在最后输出时转回 fp16。这个改动让长序列的数值稳定性大幅提升，代价是少量性能损失，但完全值得。

提示：fp16 做 attention 时，softmax 的 max 和 sum 一定要用 fp32 累加，这是血泪教训。很多框架默认用 fp16 累加，短序列看不出问题，长序列必炸。

7. 编译、调试与性能验证的完整流程

7.1 编译环境的准备

CUDA kernel 的编译我推荐用nvcc直接编，不要一上来就套 PyTorch 的扩展机制，因为后者会引入一堆头文件依赖，编译慢且报错难读。先用一个独立的.cu文件把 kernel 跑通，再考虑集成。

编译命令的关键参数：-O3开优化，-arch=sm_XX指定架构（XX 根据你的卡填，比如 80 对应 A100，86 对应 3090，89 对应 4090），--use_fast_math可以开但要注意它会降低精度，attention 里慎用。我一般不开--use_fast_math，因为 softmax 的精度太关键。

调试阶段加-G开 device 端调试信息，配合compute-sanitizer查越界和竞态。compute-sanitizer是必用的，attention kernel 里 shared memory 的索引很容易越界，sanitizer 能直接定位到行号。

7.2 正确性验证的方法

验证分三层：第一层，和 CPU 上的朴素实现对比，用随机输入，误差阈值设 1e-5；第二层，和 PyTorch 的scaled_dot_product_attention对比，注意 PyTorch 默认的 mask 语义；第三层，用真实模型的一层 attention 权重跑，看端到端输出是否一致。

第一层最容易过，第二层经常因为 mask 语义差异挂掉，第三层能暴露数值精度问题。我建议三层都跑，尤其是第三层，很多在随机输入下看不出的问题，在真实权重下会暴露。

7.3 性能测量的注意事项

测性能一定要用 CUDA event 计时，不要用 CPU 的clock()，因为 kernel 是异步的，CPU 计时会把启动开销算进去。测量前先 warmup 至少 10 次，让 GPU 频率稳定。每次测量跑 100 次取平均，去掉前几次的冷启动。

还要注意测的是 kernel 时间还是端到端时间。如果测端到端，要把 H2D 和 D2H 的拷贝时间单独算出来，否则会误以为 kernel 慢。我一般先测纯 kernel 时间，确认 kernel 本身没问题，再看端到端。

8. 一些踩坑之后的个人体会

写这个 MMHA kernel 花了我大概两周的业余时间，其中真正写代码的时间不到一半，剩下的都在调 mask 逻辑和调性能。最大的体会是：attention 的难点不在数学，而在访存和边界条件。数学公式谁都看得懂，但把 causal mask、padding mask、decode 阶段的退化情况、fp16 的数值稳定性这些边界都处理对，才是真正花时间的地方。

另一个体会是关于 split-K 的：不要一上来就上 split-K，先把基础版本写对写快，测出瓶颈在哪里，再决定要不要 split。我一开始就想着上 split-K，结果基础版本都没跑通，白白浪费了几天。正确的顺序是：先写一个能跑对的朴素版本，再优化访存，最后才考虑 split-K 这种改变计算结构的优化。

最后分享一个小技巧：调试 attention kernel 时，把 seq_len 设成 4 或 8 这种小值，然后手动把 Q、K、V 设成单位矩阵或简单模式，这样输出有明确的预期，一眼就能看出哪里错了。比用随机输入瞎猜高效得多。这个方法帮我定位了好几个索引 bug，强烈推荐。