★ FEATURED ARTICLE 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 ipadding 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]的分数矩阵就是 64MBfp32根本放不下。所以实际实现里要么分块tiling要么把行维度再切给不同 block。我这次采用的是行分块策略每个 block 处理若干行 queryK 和 V 按需从 global memory 流式加载。2.2 线程映射一个 warp 负责一行还是多行线程映射有两种常见方案。方案 A一个 warp 负责一行 querywarp 内的 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 长度只有 1causal mask 其实退化成全部可见因为唯一的 query 位置是序列末尾它能看到所有历史 key。如果这时还按训练时的 causal 逻辑去算会多一次无用的比较。所以我在 kernel 里加了一个is_decode的编译期分支decode 时直接跳过 causal 判断。3.2 负无穷的表示与 softmax 的配合把 mask 位置设为负无穷标准做法是用-INFINITY或一个足够大的负数比如-1e9。这里有个细节如果你用-INFINITYsoftmax 时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-Ksplit-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 cacheK 和 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 数量都有关建议在自己的硬件上重新测一遍。场景query 长度推荐路径关键优化训练512-4096通用 tilingshared memory 复用 K/V推理 prefill128-2048通用 tiling向量化加载 双缓冲推理 decode1-4GEMVKV cache 布局优化超长序列 2048split-K online softmaxrunning max/sum5. 访存优化与 shared memory 的实战细节5.1 K 和 V 的加载策略attention 的访存瓶颈在 K 和 V每个 query 行都要遍历所有 key如果每次都从 global memory 读带宽直接爆炸。所以必须把 K 和 V 分块加载到 shared memory让 block 内的多个 query 行复用。我的做法是block 处理 4 行 queryK 和 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开优化-archsm_XX指定架构XX 根据你的卡填比如 80 对应 A10086 对应 309089 对应 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 本身没问题再看端到端。验证层次对比对象误差阈值主要暴露的问题第一层CPU 朴素实现1e-5索引错误、mask 逻辑第二层框架 SDPA1e-4mask 语义差异第三层真实模型权重1e-3数值精度、溢出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强烈推荐。
阅读完成 · 觉得有帮助？
