最近在搞组里的推理服务时需要把 masked multi head attention 从 PyTorch 的封装里拆出来自己写一个 CUDA 实现。网上讲 Transformer 的文章很多但大部分都停在“调用nn.MultiheadAttention()就行”的层面真正到了要抠性能、抠显存、抠算子融合的时候文档和博客能直接参考的东西就很少。这篇算是我自己踩坑过程的记录从朴素 kernel 开始讲到可用的分块实现中间还包括 mask 处理、数值稳定、精度对齐和各种排查经验。如果你也在写或者准备写 attention 的 CUDA kernel那这篇应该能帮你省不少时间。1. 为什么非要从 CUDA 层面重新实现一遍1.1 现成 API 的坑不止是慢很多人会问PyTorch 不是已经封装得很好吗nn.MultiheadAttention、torch.nn.functional.scaled_dot_product_attention都有为什么还要自己写 CUDA我最初的想法也很简单项目的推理后端要在一个比较特殊的部署环境里跑batch size 可能很小seq len 却可能很长。PyTorch 的 SDPA 虽然已经集成了融合 attention 的优化但它的行为像个黑盒你很难精细控制 mask 的类型、数值精度、以及算子之间的融合边界。更麻烦的是一旦你需要在 attention 前后插入自定义的量化逻辑或者自定义 position encoding标准 API 就不够灵活了。另外一个更实际的问题是显存。对于[B, H, N, N]的 attention score 矩阵如果N4096H32B8光这一个中间结果就是 8×32×4096×4096×4 字节算下来 16GB 显存。这还只是 score 矩阵本身Q、K、V 和输出 O 还要各占一份。很多显卡根本吃不住这种显存开销。PyTorch 的融合 attention 已经尽力避免把 score 整块写回显存但一旦你为了支持某些特殊 mask 而回退到手动实现例如先算QK^T再masked_fill再 softmax再V显存占用立刻爆炸。我这边实测过一次同样的输入用这种“教科书式”写法比用融合实现多了 6 倍显存。所以自己写一个融合的 masked multi head attention CUDA kernel并不是为了炫技而是为了同时解决三个问题mask 的可控性、中间 score 矩阵的显存占用、以及算子融合带来的访存次数下降。1.2 从显存带宽看 attention 的真正开销在 GPU 上一个算子的耗时往往不取决于计算量而取决于显存读写量。attention 的计算复杂度是 O(N²D)看起来计算很重但实际上对于常见的head_dim64或者128只要把数据喂到位计算单元往往在“等数据”。举个直观例子如果不融合QK^T会把[B, H, N, N]的 score 写回显存然后 softmax 又读一遍、写一遍最后scoreV再读一遍。每多一次全量读写就是在烧显存带宽。假设N2048B8H16score 矩阵有 8×16×2048×2048×4 2GB。一个算子来回搬三次就是 6GB 的显存流量而一个旗舰级 GPU 的显存带宽也就 1-2TB/s光搬运就要几毫秒。如果把 score 留在寄存器或 shared memory 里计算完直接消费掉这部分开销就直接归零了。这就是为什么现代高性能 attention 实现都走“分块计算 在线 softmax”的路线。本质上就是把大矩阵切小块一小块一小块算算完的 score 不进显存直接从寄存器参与加权求和。我们的 CUDA 实现也要走这条路。1.3 mask 让问题进一步变复杂mask 是另一个容易低估的点。masked multi head attention 里的 mask可能是因果 maskcausal mask、padding mask、或者是两者叠加。因果 mask 要求位置的 query 只能看到它之前和它当前位置的 key表现在 score 矩阵上就是上三角区域被屏蔽。padding mask 则是因为 batch 里每条样本长度不同短样本的无效位置不能参与 attention。如果 mask 是在 PyTorch 层做的那很简单一个masked_fill就完事。但在 CUDA kernel 里你不能简单地把 mask 加载进来再填一个-inf就结束了因为 mask 的存储布局、加载方式、以及它是“运行时才知道”还是“静态可知”都会直接影响 kernel 的设计。这部分我放到后面专门讲。2. 计算结构拆解和 kernel 的初始布局2.1 先把维度写在纸上动手写代码之前我习惯先把所有张量的维度写在纸上。masked multi head attention 的输入通常是Q:[B, N, H, D]或者[B, H, N, D]取决于你用的框架。K、V 同理。mask可能是[B, N, N]的布尔矩阵也可能是每个样本的有效长度[B]。我这里统一按[B, H, N, D]来设计 kernel。B是 batchH是 head 数N是序列长度D是 head_dim。attention 公式是[ Attention(Q,K,V) softmax(\frac{QK^T}{\sqrt{D}} Mask) V ]QK^T的结果是[B, H, N, N]Mask会广播到这个矩阵上。被 mask 掉的位置在 softmax 之前让 score 变成一个非常小的数通常用负无穷这样 softmax 之后概率就是 0。这里有个关键点B和H其实是可以合并的。因为在[B, H, N, D]布局下每个 batch 里不同 head 的数据是独立的我们可以把BH当成一个统一维度来处理。kernel 的实际计算单元就是一个“batch-head 对”加上一段 query 位置这大大简化了索引计算。2.2 因果 mask、padding mask 和组合 mask我实现的时候分别处理了三种 mask 场景。第一种是因果 mask。它的特点是可以“隐式”处理不需要真的构造一个[N, N]的 mask 矩阵只需要在 score 计算时判断key_index query_index就屏蔽。这种判断在 kernel 里只是一个小分支代价很低。第二种是 padding mask。它有两种表示方式一种直接传入[B, N, N]的 mask 矩阵另一种只传每个样本的有效长度valid_len[B]。显然后者更省显存也更便于 kernel 加载。判断时只需要比较key_index valid_len[b]是否成立。第三种是组合 mask。通用做法是让 mask 变成一个float矩阵合法的位置为 0非法的位置是-inf或一个很大的负数直接在 score 上加 mask。这种做法的好处是逻辑简单坏处是 mask 本身要占显存而且访存开销大。如果N再大一点[B, N, N]的 float 矩阵也相当可观。我的建议是能隐式判断的尽量隐式判断实在不行再走显式 mask 矩阵。因果 mask 这种结构固定的场景完全没必要让 mask 占一块显存。2.3 线程/block 怎么映射到 attention 矩阵kernel 设计的第一个问题是block 和线程到底对应什么我采用的映射方式是一个 block 负责一个(batch_head_idx, query_chunk)。也就是说blockIdx.y表示处理第几个BHblockIdx.x表示处理第几个 query 分块。block 内部线程共同处理这个 query 分块与所有 key 之间的 attention 计算。这样做的理由是Q的一个分块可以被反复复用。假设每个 block 处理BLOCK_M个 query那这个 block 只需要加载一次对应的Qtile 到 shared memory然后循环遍历所有K、Vtile。随着序列变长Q的加载次数不会跟着涨整体访存效率会好很多。相反如果让一个 block 只处理一个 query虽然逻辑简单但每个 query 都要重新加载一遍Q对应的D维数据同时 softmax 的归约也变得更琐碎整体效率偏低。block 内部线程负责BLOCK_M行和D列。通常让一个 warp32 个线程负责一行 query 的若干列或者让一个 warp 负责多行的一列组。具体看D的大小。D64时一个 warp 刚好能覆盖两行或者让 32 个线程各负责一行 query 的 2 个列。3. 朴素 kernel 先跑通再聊优化3.1 朴素版本的计算流程写优化版本之前我强烈建议先写一个朴素版本。这个版本不追求性能只求逻辑正确、容易调试。我的朴素版本思路是这样的每个 block 负责一个(batch_head, query)位置。block 内线程数等于D也就是 64 个线程。每个线程负责D维中的一列先算Q和某个K的点积。整个计算分成三步每个线程遍历所有N个 key计算score(q, k)同时维护一个局部最大值max_score。基于最大值再次遍历所有 key计算exp(score - max_score)并求和得到sum。第三次遍历所有 key把exp(score - max_score) / sum作为权重累加V的对应行得到输出。三次遍历看起来蠢但它把 softmax 的数值稳定性问题处理得非常直观。先找最大值再算指数和exp(x - max)最大也不会超过 1不会出现exp(1000)这种溢出。这里的线程协作我用了 shared memory 存Q的一行然后用__syncthreads()同步。至于归约最大值和总和最简单的做法是让线程 0 去做一个串行归约或者用 warp shuffle。由于这是朴素版本串行归约完全可以接受跑通了再说。3.2 masked softmax 的数值稳定问题写这个朴素版的时候我踩了第一个坑mask 位置的数值到底该填多少。受 PyTorch 的masked_fill(-inf)影响我第一版直接给被 mask 的 score 赋了-INFINITY。结果在极端情况下——比如某个 query 的所有 key 都被 padding mask 屏蔽时——max_score也是-inf然后exp(-inf - (-inf))直接变成NaN。这个 NaN 会污染整个输出。正确做法是不要用真正的负无穷而是用一个很大的负数比如-1e30f。这样即使整行被 maskexp(-1e30 - (-1e30)) exp(0) 1也不会出现 NaN最多是分母为一堆 1最后输出变成无意义的数。但后续如果想彻底规避会在代码里加一个“如果整行无效直接输出 0”的分支。另一个容易踩的坑是 scale 的位置。公式里是QK^T / sqrt(D)如果你先算s QK^T再乘scale那么 mask 和scale的先后关系无所谓。但如果为了省事把scale乘到Q上例如先Q_scaled Q / sqrt(D)那softmax的结果不变这点是安全的。怕的是你用了 PyTorch 的nn.MultiheadAttention之后忘了它内部已经把scale写死了自己实现时又乘了一次结果精度怎么都对不上。我在对齐测试时就遇到过这个问题后面会详细说。3.3 一个容易被忽略的细节scale 的位置补充一下 scale 的实操经验。我在最初版本里把 scale 乘在了 K 上也就是K_scaled K / sqrt(D)这样点积的时候就不用每次都乘 scale。看起来没问题但乘以一个小于 1 的数可能会让 K 的数据精度略微下降。后来我改成了“在 shared memory 加载 Q 之后直接对 Q tile 做 scale”因为 Q tile 每个 block 只加载一次scale 的开销可以忽略同时点积结果天然就是缩放后的 score省掉了内层循环里的一次乘法。不过这个改动会带来对齐测试的一个坑如果你对比的 baseline 是先算完整 score 再 scale那么两者的浮点舍入顺序不同结果会出现细微差异。这个差异通常在1e-5量级但如果你的上游对精度要求很高可能需要刻意对齐舍入顺序。4. 分块 kernel 与在线 softmax 实现4.1 shared memory tile 的基本思路朴素版本能跑通之后我开始把它改成分块实现。目标是一个 block 处理一个(batch_head, query_chunk)把Q的BLOCK_M×Dtile 常驻在 shared memory然后循环加载K和V的BLOCK_N×Dtile逐步更新输出。为什么要用 shared memory因为矩阵乘法的数据复用率很高。Qtile 要与所有Ktile 相乘Ktile 要被多个 query 行复用。如果每次都从全局显存读带宽很快就会被吃满。shared memory 相当于一个软件管理的缓存把反复要用的数据放进去。我的第一版分块参数是BLOCK_M64、BLOCK_N32、D64。shared memory 使用量大致是Q tile64×64×4 16KBK tile32×64×4 8KBV tile32×64×4 8KB加起来 32KB对于大多数显卡的 48KB 或 96KB 共享内存限制来说余量很充足。这个参数下一个 block 需要 32 个 warp 来覆盖64×64的输出 tile也就是 1024 个线程。如果不想让 block 太大可以缩小到BLOCK_M32每个线程处理两个列或者用 256 线程处理32×64tile。4.2 核心循环代码和 mask 处理分块 kernel 的核心是内层循环。这里我给出一段核心逻辑的代码示意展示 online softmax 和 mask 的结合方式// 伪代码 / 核心循环结构 // 每个线程对应 output 中的一个 (query_idx, d) 位置 // q_cur 是当前 query 在 seq 维度上的全局下标 // d 是 head_dim 维度上的某一列 float m -1e30f; // 当前行看到的 score 最大值 float l 0.f; // 指数和 float acc 0.f; // 当前线程负责的那一列的输出累加 for (int k0 0; k0 N; k0 BLOCK_N) { // 1. 协作把 K[k0:k0BLOCK_N] 和 V[k0:k0BLOCK_N] 加载到 shared __syncthreads(); for (int j 0; j BLOCK_N; j) { int k k0 j; // 点积 float s 0.f; for (int d_i 0; d_i D; d_i) { s Qs[q_local][d_i] * Ks[j][d_i]; } s * scale; // mask 处理这里按因果 mask 举例 // 如果是 padding mask则改为 k valid_len[q_cur] if (k q_cur) { s -1e30f; } // online softmax 更新 float m_new fmaxf(m, s); float alpha expf(m - m_new); float beta expf(s - m_new); l l * alpha beta; acc acc * alpha beta * Vs[j][d]; m m_new; } __syncthreads(); } float out acc / l;这段代码的核心是 online softmax。它不要求一次性拿到整行的所有 score而是每看到一个 score 就更新当前的最大值m、分母l、分子acc。这样即使循环很多轮也只需要在最后做一次除法中间不需要把整行 score 存下来。这也是 flash attention 这类高性能实现的核心技巧。mask 在这里做得非常轻量只是给被屏蔽的 score 赋一个非常小的值。由于后续计算会走exp(s - m_new)这个极小值在指数上约等于 0所以被 mask 的位置对l和acc的贡献基本可以忽略。4.3 从逐元素更新到分块 rowmax 的演进上面代码里的 online softmax 是“逐元素”更新方式也就是每看到一个 score 就更新一次m和l。这种做法正确但每个 score 都要算两次指数一次是exp(m - m_new)一次是exp(s - m_new)。指数运算在 GPU 上不算便宜所以这个版本虽然能跑但性能一般。更好的做法是分块内部先做 rowmax。具体来说算出一个BLOCK_N长度的 score 小块之后先在小块内部找到局部最大值m_local然后用这个局部最大值统一计算指数和分母最后再把小块合并到全局的m和l上。这样每个 score 只需要一次exp指数运算量直接砍半。我当时先实现了逐元素版本目的是先把正确性跑通后来才优化成 rowmax 版本。实际测下来在N1024时逐元素版本比分块 rowmax 版本慢了接近 30%。如果你直接上正式代码建议跳过逐元素版本直接按“分块 rowmax 合并”来写。这里还有一个工程细节__expf和expf的差异。__expf是硬件快速版本速度更快但精度低一些误差大概在1e-5量级。做精度对齐测试时要用expf确认无误后再改成__expf提速。5. 编译、联动 PyTorch 与精度对齐5.1 环境准备与 CUDA 版本选择写 kernel 之前环境先得准备好。CUDA 版本选择主要看你的显卡架构和 PyTorch 编译时用的 CUDA 版本。我的经验是PyTorch 用的 CUDA 运行时版本和你自己nvcc编译用的版本不必完全一致但最好不要跨大版本。比如 PyTorch 是用 CUDA 12.1 发布的那你也用 12.x 的 toolkit 编译避免 ABI 兼容问题。如果你在 WSL2 环境下开发有个省心的地方Windows 侧的显卡驱动在 WSL2 里直接可用你只需要在 WSL2 里安装 CUDA toolkit不需要重复装驱动。装完记得检查一下nvcc --version和 PyTorch 的torch.version.cuda是否一致不一致时优先把LD_LIBRARY_PATH指到你 toolkit 的lib64目录。编译针对特定架构时用-arch参数指定计算能力。我手头测试卡是 Ampere 架构用-archsm_80。如果你的卡是更新一点的架构例如 Ada Lovelace 或 Hopper需要对应修改为sm_89或sm_90。记不清的话可以用torch.cuda.get_device_capability()查一下。5.2 编译和 launch 配置编译成 PyTorch 可以直接调用的扩展最省事的方式是用setuptools加CUDAExtension。我简化一下核心编译命令大概是nvcc -archsm_80 -O3 --shared -Xcompiler -fPIC attn_kernel.cu -o libattn.so然后在 Python 里用ctypes加载或者简单封装成 PyTorch 的 autograd Function。launch 配置是另一个容易错的地方。我的 kernel 启动参数大致是这样dim3 grid((N BLOCK_M - 1) / BLOCK_M, B * H); dim3 block(BLOCK_M * D / 32); // 或者根据线程模型调整 masked_attn_kernelgrid, block, shared_bytes(...);一个常见的错误是block 里线程数超过了 1024 的限制或者 shared memory 请求量超过了设备上限。遇到cudaErrorInvalidConfiguration时优先检查这两项。5.3 与 PyTorch baseline 的对比方法kernel 写完最重要的一步是跟 PyTorch 的 baseline 对齐。我习惯写一个独立的验证脚本随机生成几组 Q、K、V 和 mask然后对比输出。这里有几个坑按顺序排一下先对比 PyTorch 的scaled_dot_product_attention但要把allow_tf32关掉。默认情况下 Ampere 之后的卡在矩阵乘时会用 TF32数值和纯 FP32 不同。如果你的 CUDA kernel 是 FP32 累加那必须设置torch.backends.cuda.matmul.allow_tf32 False否则差异会让你误以为 kernel 写错了。对比时先测全 mask 的边界情况比如某个 query 整行都被 mask。这是最容易出 NaN 的场景先排除掉再测常规数据。数值误差阈值给多少合适一般head_dim64时allclose(atol1e-5, rtol1e-5)就算通过。但如果你用了__expf阈值可能要放宽到1e-4。我之前遇到过一个问题单独看每个输出元素误差都在1e-6以下但用 PyTorch 自带 test 的默认atol1e-8去比直接报错。后来查了下是 PyTorch 默认allow_tf32True导致的。先关掉再测误差就回到正常范围了。6. 性能实测和常见问题速查6.1 我实测的一组数据为了说明优化效果我在本地测了一组数据。配置是单卡 Ampere 架构N1024H8D64B1纯 FP32因果 mask。表格里的数据只是单次运行参考值不是严谨 benchmark但趋势很说明问题实现方式显存占用耗时相对PyTorch 手动实现先算 score再 mask再 softmax再乘 V高4.0x朴素 CUDA kernel三次遍历中3.2x分块 CUDA kernel逐元素 online softmax低1.4x分块 CUDA kernelrowmax 合并低1.0x显存这块手动实现的中间 score 矩阵要吃几 GB而分块 kernel 全程不落 score只占输入输出外加一小块 shared memory。个人体会是如果你的N只有一两百那手写 CUDA 的收益并不大直接调 PyTorch SDPA 更省事。但当N到 1024 以上或者你需要频繁改 mask 逻辑时融合 kernel 的优势就很明显了。6.2 数值错误、mask 错误和启动失败排查把自己踩过的坑整理成一个速查表方便后续排查现象可能原因处理方法输出出现 NaNmask 位置用了-INFINITY导致exp(-inf - (-inf))改用-1e30f或加整行 mask 保护分支精度和 PyTorch 对不上PyTorch 默认开了 TF32设置allow_tf32 False精度差但只有个别元素明显用了__expf换回expf验证确认后再说因果 mask 效果不对判断条件写反比如k q_cur检查key_index query_index是否被屏蔽launch 失败block 线程数超过 1024或 shared memory 超限用cudaOccupancyMaxActiveBlocksPerMultiprocessor检查性能比 PyTorch 还慢没有做分块显存反复读写优先做 shared memory tile 和 online softmax输出结果完全错QKV 索引算错head 维没有按连续布局处理在 kernel 里加几行 printf或者用cuda-gdb查还有一个很隐蔽的坑如果你把B*H合并成统一维度那么 blockIdx.y 和 blockIdx.z 的映射顺序要和 QKV 的内存布局严格一致。[B, H, N, D]布局下b和h应该用blockIdx.y / H和blockIdx.y % H拆开或者干脆让b由blockIdx.z承担。我当时就是这里写反了导致不同 batch 的数据串了输出看起来是“有点合理但整体不对”。6.3 经验总结和后续扩展整个写下来的过程我个人最深的体会是不要一上来就照着 flash attention 的论文抄实现。先写一个正确的朴素 kernel把掩码逻辑、数值稳定性、索引映射这些基础问题一次想清楚然后再逐步引入 shared memory tile、online softmax、向量化加载这些优化手段。每一步都做一次和 PyTorch baseline 的对齐出了问题很容易定位。后续如果想继续压性能可以考虑几条路一是用float4做向量化加载让显存读取更高效二是引入cp.async异步拷贝把 shared memory 加载和数据计算重叠起来三是在 Ampere 之后的卡上利用 tensor core 做矩阵乘。不过这些优化的投入产出比要看你实际的N和H规模。对我来说在N1024这个量级上分块加 rowmax 已经够用了。如果你的 seq len 更长或者有批量推理的需求那确实值得再往下探索。