这个融合是我去年年底在优化一个线上推理链路时真正落地过的方案。当时模型里Attention占了接近四成耗时profile一打出来三个kernel来回倒腾中间矩阵显存带宽被白白烧掉一大截。这篇文章我把这段经验完整写出来先讲清楚三个算子为什么非融合不可然后从数学原理一步步推到Triton和CUDA的两种实现最后是实测数据和一些非常反直觉的调参坑。如果你也在做transformer推理加速这篇应该能帮你少走不少弯路。先交代一下背景。这里说的ops-transformer不是某个开源框架的名字而是我自己的一个算子优化实验项目目标很单纯把Attention里最常见的三阶段计算——QKV拼接计算、Softmax归一化、最后与V的矩阵乘——合并成一个融合Kernel消除中间结果的显存往返。为了表述方便下文统一叫这个项目名。1. Attention三算子的算路与访存瓶颈分离实现输在哪里1.1 从“三个算子”到完整计算链先把Attention的计算拆开。以单头为例输入X经过QKV投影后得到Q、K、V三个矩阵之后的核心计算是由Q和K计算注意力分数矩阵S Q · K^T再除以缩放因子。对S的每一行做Softmax归一化得到概率矩阵P。用P对V做加权求和O P · V。严格数一下这条链上有四个逻辑步骤投影、QK^T、Softmax、PV。但在实际工程里QKV投影往往通过拼接后的权重一次MatMul完成而且这一步天然可以和后面的QK^T连起来看。所以很多优化资料把它们归纳成“QKV拼接 Softmax MatMul”三大段这里的“QKV拼接”指的是Q、K、V数据在计算图中的汇合方式也可以是物理上把三个投影权重拼成一个矩阵来做合并运算。无论是哪种真正要解决的核心矛盾是同一个中间那个S矩阵太大而且它在分离实现中必须落回显存。我的优化起点就是盯着S矩阵的。当序列长度为512、头维度为64时单个头的S是一个512×512的矩阵float32下就是1MB。12个头就是12MB这还只是写一次下一轮QK^T的kernel读出来又要12MB一写一读就是24MB。这24MB里只有极小的比例是最终要保留的信息大部分带宽纯粹是传送带把数据运出去再运回来。1.2 访存量估算算力根本不是瓶颈用一张表格表达这个问题的直观程度可能比我上面说的一大段都清楚。以单头、序列长度512、头维度64、float32为例阶段中间张量大小单头写入显存读出显存QK^TS矩阵512×512×4B 1MB1MB1MBSoftmaxP矩阵512×512×4B 1MB1MB1MBPVO矩阵512×64×4B 128KB128KB忽略单头就产生了约3MB的显存往来。按12头计算就是36MB。我测试的这张推理卡有效带宽大约在1.5TB/s级别这些中间矩阵搬运就要消耗约24微秒。而真正这三个算子的数学运算量是多少呢QK^T是512×512×64约1600万次乘加PV也是约1600万次乘加两者合计约3200万次浮点操作在这块卡上纯算力执行不足3微秒。也就是说分离实现中超过85%的时间花在了把中间结果写出去再读回来。这个结论直接决定了融合的必要性只要能让S矩阵和P矩阵不出芯片内部的高速缓存性能收益就是结构性的。这还没有算kernel启动的开销。每个GPU kernel从启动到结束都有固定的调度开销几个微秒量级。三个kernel累计10到20微秒的额外延迟在长序列训练场景里也许还能被大计算量掩盖但在解码阶段的短序列场景中kernel启动开销甚至可能超过运算本身。这一点后面实测会再提到。1.3 解码场景下的放大效应训练阶段batch大、序列长计算密集访存问题相对没那么扎眼。真正被放大的是推理场景尤其是自回归解码。解码时每次只生成一个tokenQ只有一行但K和V会随着历史长度增长。分离实现会不断把越来越大的S矩阵写进显存序列一旦长到2048甚至4096单头S矩阵就是16MB到64MB多头累积起来的带宽消耗非常恐怖。我做优化时应用的线上场景正是这种变长解码推理序列长度经常在几百到两千之间波动所以对访存优化特别敏感。算力永远都在等待数据这不是某一个kernel写得不好而是算子边界划错了。2. 融合的数学底牌行级Softmax、Online Softmax与分块点积2.1 核心洞察Softmax只需要行内统计量任何一次算子融合本质都是在问一个问题这一步计算结果是否真的需要作为完整张量落到外部存储对于Attention来说答案显然是否定的。关键在于Softmax的行独立性。Softmax的操作对象是S矩阵的行。第i行的Softmax只依赖这一行的元素不涉及其他行的任何数值。这意味着S矩阵可以按行任意分块处理我先算第i行的一部分分数做局部Softmax再算另一部分最后把结果合并只要合并方式在数值上正确就不需要先物化完整行。这个特性给了融合一个非常重要的自由度QK^T这一步不需要把完整的S算完再交给Softmax而是可以算一小块分数、立刻做一次局部Softmax、把结果累加到最终输出中然后扔掉这一块分数。整个过程中S和P都只存在于片上缓存里。很多人第一次听说时会觉得这违背直觉因为数学公式里P Softmax(QK^T)明明需要先做完整矩阵乘再归一化。但Softmax的形式决定了它根本不需要“全局”信息行的最大值和归一化项才是关键。这也正是下面要讲的Online Softmax存在的理由。2.2 Online Softmax的迭代公式Online Softmax的妙处在于它允许我们用“流式”的方式逐块计算Softmax并且每一次合并都保持数值稳定。具体来说维护两个状态变量m当前见过的所有元素中的最大值初始为负无穷。l当前见过的元素的指数和初始为0。当新的一个数据块x_block到来时先求块内最大值m_block然后合并m_new max(m, m_block) l_new l * exp(m - m_new) sum(exp(x_block - m_new))注意新的归一化项由两部分组成旧部分的贡献需要乘以一个校正因子exp(m - m_new)因为归一化的基准从m变成了m_new新部分的每个元素都减去m_new再求指数避免直接计算exp造成溢出。这里的关键是只要维护好m和l合并顺序就无关紧要任意顺序都能得到正确的最终归一化项。这个迭代形式和经典的“双pass Softmax”相比只做了一次数据遍历但数值效果等价。我们用生活类比来记老数据的指数和中每一项都是以旧的最大值为基准缩放过的现在最大值变大了旧项就得统一乘上一个“缩水系数”把基准调整到新的最大值上。这跟调整天平滑板的过程极其相似——杆子抬高到新的位置所有旧刻度都要按比例重新标注。2.3 在Attention融合中怎么应用在融合Attention Kernel里Online Softmax的迭代过程是这样的取Q的第i行在内部循环里依次遍历K的行块。对每个K行块计算局部QK^T分数块得到s_block。更新m和l用新的m对历史累积输出做一次校正再把当前块的softmax概率乘以V的对应行块累加到输出。遍历完所有K行块后用最终的l对输出做最后一次除法得到Attention结果。每一步中s_block和p_block的生命周期都只存活在寄存器或线程块的共享内存里循环进入下一轮迭代后上一块的数据就可以被覆盖。这恰好和GPU的片上存储层级匹配shared memory容量有限但带宽远高于显存而且同一线程块内共享无需全局同步。分块大小也由此确定——要保证一个迭代步骤中的QK块、概率块和V块能同时放进shared memory。注意这里还有一层隐含的矩阵乘嵌套。最终输出O P · V而P本身来自QK^T和Softmax所以实际是一个两级矩阵乘的复合。如果不融合第一级结果S会被当作中转站写进显存融合后中转站变成片上缓冲。上一节访存量分析已经证明这个中转站恰恰是最大的带宽消耗源因此从显存中把它消灭是收益的核心来源。3. 第一版原型用Triton在半小时内跑通融合Kernel3.1 为什么先用Triton而不是直接写CUDA这个项目我一开始没直接写CUDA。虽然CUDA是最终性能的保证但调试周期长尤其是Online Softmax这类带跨迭代数据依赖的kernel稍不留神就是数值错误。Triton的优势在于你关注的是分块逻辑和计算顺序而不是线程怎么调度、shared memory怎么分配。Triton的编程模型以program为单位每个program处理一个计算块语言内置了块级别的matmul。更重要的是Triton的编译器能自动把这些块操作映射成底层的线程级操作shared memory的分配和复用大部分交给编译器处理。这对快速验证融合逻辑非常友好。我的建议是任何复杂的算子融合先花半小时用Triton做一个正确性可以验证的原型再根据Profile结果决定要不要手写CUDA。真实项目中我见过太多人一上来就写CUDA结果好几个星期耗在边界case上最后发现自己的算法路径一开始就错了。3.2 融合Kernel的核心代码下面这段是Triton版融合Attention的核心实现。为了阅读方便我做了简化假设序列长度S是BLOCK_S的整数倍不处理mask只做单头单query行的演示逻辑。实际多头的循环我在后面说明。import triton import triton.language as tl triton.jit def fused_attn_kernel( Q_ptr, K_ptr, V_ptr, O_ptr, S: tl.constexpr, D: tl.constexpr, BLOCK_S: tl.constexpr, BLOCK_D: tl.constexpr, scale: tl.constexpr, ): row tl.program_id(0) # 当前处理的query行 head tl.program_id(1) # 多头场景下的头编号 offs_d tl.arange(0, BLOCK_D) offs_s tl.arange(0, BLOCK_S) # 加载当前query行 q tl.load(Q_ptr head * S * D row * D offs_d) m_i -1e30 l_i 0.0 acc tl.zeros([1, BLOCK_D], dtypetl.float32) for start_s in range(0, S, BLOCK_S): offs_k start_s offs_s k tl.load(K_ptr head * S * D offs_k[:, None] * D offs_d[None, :]) v tl.load(V_ptr head * S * D offs_k[:, None] * D offs_d[None, :]) # QK^T 得到当前块分数 qk tl.dot(q[None, :], tl.trans(k)) # 形状 [1, BLOCK_S] qk tl.reshape(qk, (BLOCK_S,)) qk * scale # Online Softmax 更新 m_new tl.maximum(m_i, tl.max(qk, axis0)) p tl.exp(qk - m_new) l_new l_i * tl.exp(m_i - m_new) tl.sum(p, axis0) # 校正旧累积输出再加上当前块的PV correction tl.exp(m_i - m_new) * l_i / l_new acc acc * correction acc tl.dot(p[None, :] / l_new, v) m_i m_new l_i l_new out tl.reshape(acc, (BLOCK_D,)) tl.store(O_ptr head * S * D row * D offs_d, out)这段代码的逻辑和上一节的数学推导一一对应。有几个地方值得单独说。第一acc的形状是[1, BLOCK_D]它不是一个矩阵而是最终输出O的一行。tl.dot(p[None, :] / l_new, v)的形状恰好是[1, BLOCK_D]所以累加是逐块的。这里把除以l_new提前到矩阵乘里做而不是最后统一除原因是后面的校正因子需要基于当前l_new来修正旧输出提前归一化反而让数值更稳定。第二qk * scale使用的是缩放因子通常是1 / sqrt(D)。这个缩放必须在Softmax之前完成否则概率分布会被一个常数扭曲。第三也是最容易犯错的一点每次更新acc时校正因子是exp(m_i - m_new) * l_i / l_new注意这里同时包含了最大值修正和归一化项修正。如果漏掉exp那一项短序列上可能看不出问题但长序列上数值会逐渐漂移最终输出会明显偏离正确结果。3.3 原型验证与数值误差控制原型跑通后第一件事不是看性能而是验证数值正确性。我的做法是用PyTorch的标准Attention结果作为参照比较最大绝对误差和平均相对误差。用float32计算时融合版本和参考版本的误差通常能控制在1e-6以下如果误差超过1e-4几乎可以断定是Online Softmax的校正逻辑出了问题。这里有一个非常常见的问题为什么我把初始的m_i设成-1e30而不是负无穷用float表示真正的-inf当然也可以但在Triton的某些后端上exp处理-inf边界时存在优化路径可能导致第一个数据块的校正因子变成NaN。用-1e30可以安全规避这类问题代价仅仅是第一个块内exp(qk - m_i)可能会因为qk - m_i很大而上溢到inf——但实际上第一个块会把m_new更新为块内最大值随后一切都恢复正常。这个细节属于典型的“不做不知道做了才明白”的经验。数值验证通过后我才开始用triton.testing的基准工具测耗时。这时候发现原型已经比分离实现快了不少但还没有达到我的预期。这让我意识到Triton原型适合验证算法要在生产环境榨干性能还是得落到CUDA层面做更精细的调度。4. 从原型到CUDAshared memory、同步与mask的实战细节4.1 线程组织与shared memory的分配策略手写CUDA版本的融合Attention核心设计决策有三个线程怎么排列、shared memory放什么、以及何时需要同步。我采用的方案是一个线程块处理一个query行线程块内用一维128线程每个线程负责输出向量D维中的一部分。以D64为例64个线程分别持有最终输出的一个分量其余的线程用来加速QK^T计算中的并行规约。shared memory里每个迭代块需要放三样东西K的一个行块、V的对应行块、以及当前块计算出的概率向量p。K和V块的大小是BLOCK_S×D以BLOCK_S32、D64、float32计算两个矩阵加一起8KBqk向量和p向量不到1KB整个块需要的shared memory在10KB以内H100和主流推理卡上都不成问题。这也反向验证了一件事BLOCK_S不能无限增大因为shared memory是硬性上限。迭代流程是先把K块和V块从显存加载到shared memory所有线程协作计算当前行的qk部分和再通过块内归约得到局部分数、最大值和指数和。这里需要特别注意同步点——每次写shared memory后必须__syncthreads()否则线程A还在读旧数据、线程B已经在写新数据结果必错。4.2 Online Softmax在CUDA里的同步点处理Online Softmax放在CUDA里真正的难点是m_new和l_new是跨线程共享的标量更新它们需要在块内同步而同步会卡住所有线程的进度。我一开始的实现把同步放得很频繁结果发现每个迭代块都要等两次归约性能惨不忍睹。后来的优化策略是把一次迭代分成两个阶段。第一阶段所有线程并行计算各自负责的局部qk值并做线程桶内的部分归约得到这个线程的局部最大值和局部指数和。第二阶段通过共享内存做一次跨线程归约得到块级的m_new和l_new。归约完成后每个线程都拿到完整的m_new和l_new立即用它更新自己持有的acc分量。这样整个迭代只需要两次__syncthreads()一次是K、V块加载完毕后的数据就绪同步一次是归约结果写入共享内存后的可见性同步。两次同步两个阶段基本是这类融合kernel的标准范式。关键点在于acc这个输出累积向量按线程分布在寄存器里校正因子是标量所以每个线程更新自己的acc时不需要额外的跨线程通信。这是寄存器级数据复用带来的巨大优势——在分离实现中这一步对整个P矩阵做了全局遍历而在融合kernel里它只是寄存器上的一个标量乘加。// 伪代码省略了内存申请和指针计算细节 __global__ void fused_attn(float* Q, float* K, float* V, float* O, int S, int D, float scale) { int row blockIdx.y; int head blockIdx.x; int tid threadIdx.x; float acc[VDIM]; // 每个线程负责的VDIM维输出分量 float m_i -1e30f, l_i 0.f; float q[VDIM]; // 加载q... for (int start_s 0; start_s S; start_s BLOCK_S) { // 阶段1: 协作加载K、V块到shared memory // __syncthreads(); // 阶段2: 计算局部qk, 块内归约m_new, l_new // __syncthreads(); // 阶段3: 校正acc, 累加p * V的局部贡献 } // 写入输出: O acc / l_i }这个伪代码虽然省略了大量细节但结构是正确的骨架。我强烈建议自己动手从这版骨架开始填而不是直接抄完整的开源实现因为Attention的变体实在太多了——有mask、有GQA、有alibi位置编码每加一个变体kernel内部的数据流都要跟着改。骨架理解了改起来才不慌。4.3 Mask处理padding mask和因果mask塞进融合Kernel的两种姿势推理场景几乎都躲不开mask。padding mask是序列实际长度不满时需要在分数矩阵上把无效位置设成负无穷因果mask是解码时只允许关照历史位置。这两类mask在融合kernel里的处理方式略有不同。padding mask最简单。在计算qk分数时如果K的某一行对应的是padding位置直接把该位置的分数设为-INFINITY。这样做不需要额外存储mask矩阵只需要在加载K块时附带一个有效性标志可以用一个小的索引判断valid (start_s i) seq_len。在CUDA里这只是一个整数比较。因果mask更微妙。它要求第i行的query只能看到j i的key。在逐块遍历中对于当前query行row如果块内key索引block_idx全部大于row那么这个块可以完全跳过。如果块内部分有效、部分无效那就用j row ? 0 : -INFINITY的掩码逐元素乘到qk上。跳过整个块的优化对解码阶段尤其重要因为解码时每个query行只处理到目前为止的历史长度尾部大量块都是无效的跳过能省一大半计算。需要提醒的是任何引入-INFINITY的地方都要同时检查Online Softmax更新时会不会出现涉及无穷的数值异常。实际测试中只要序列不是全padding一般都能正常收敛但保险起见我会在exp之前把有效位置的值clip到一个安全范围比如[-85, 0]。4.4 更长序列两阶段切分策略当序列长度超过2048时即使BLOCK_S32单线程块要串行遍历的块数也太多了。更重要的是QK^T和PV在数学上对K和V的需求方向不同单纯增大BLOCK_S会让shared memory压力陡增。我的做法是引入两阶段切分外部按key方向分大块每个大块交给一个线程块线程块内部再对小块跑Online Softmax大块之间通过原子操作或者跨block归约合并部分结果。这种做法本质上就是FlashAttention的多核分区思路不过实现复杂度会高不少。对于绝大多数部署场景两阶段切分其实不是必需的。序列长度在1024以下时单block的串行版本已经能打到接近硬件极限。我之所以提它是因为很多人拿着长序列的benchmark来质疑短序列优化没有意义。我的经验是先做无分区的融合版本把基础收益拿下再根据最长序列需求决定要不要上分区。不要一开始就把架构搞得过于复杂。5. 实测收益与调参陷阱block大小、L2缓存与数值误差的取舍5.1 测试环境与对比基线原型和CUDA版本都跑通后我在自己的测试环境上做了完整的性能对比。测试卡的显存带宽属于主流推理卡水准batch size为1多头数12头维度64覆盖了128到2048五种序列长度。对比基线有两种一是PyTorch原生的Attention实现二是已经做了部分算子合并但仍有显存中转的中间版本。前者代表大多数人的起点后者模拟的是那些“只融合了一半”的优化方案。5.2 加速比数据序列长度分离实现耗时(us)部分融合耗时(us)完整融合耗时(us)完整融合加速比1284538182.5x2568267282.9x512152118453.4x1024310220784.0x20486804301384.9x这个趋势非常符合预期序列越长中间S矩阵越大融合省掉的访存量占比越高加速比越可观。短序列场景下完整融合的收益也没有消失因为分离实现多出来的kernel启动开销被完整吃掉了。这也回应了前面提到的解码场景——短序列虽然计算量小但三个kernel的启动延迟是固定的。我一度以为这个结果已经够了但深入profile之后发现还有潜力。5.3 第一个坑BLOCK_S不是越大越好最开始我把BLOCK_S设为64理由是每轮可以处理更多的K行减少循环次数。实测下来BLOCK_S64在序列长度512时性能反而退化了约8%在1024时退化更多。原因有两层。第一BLOCK_S变大K和V块在shared memory里的占用量翻倍这会挤压编译器做寄存器分配的余量导致寄存器溢出到local memory而local memory在物理上还是显存等于绕了一圈又回到了访存瓶颈。第二更大的块意味着更少的并行块数在batch size1的场景下GPU的线程块并发度本来就不高再减少块数会让SM的占用率雪上加霜。我在同样的硬件上试了BLOCK_S16、32、64三档。最优值稳定在3216时因为循环次数翻倍调度开销开始显现性能略低。最终的结论是BLOCK_S32对于头维度64、float32的典型配置是一个健康的默认值。5.4 第二个坑L2缓存命中率与“伪融合”有件很有意思的事我优化到一半时曾经怀疑融合到底有没有用因为把融合kernel单独跑耗时竟然比分离实现只快了一点点。后来抓了L2缓存命中率才发现分离实现在小序列下大量S矩阵幸运地命中了L2缓存显存访问被隐藏了。序列一旦拉长L2放不下了分离实现的短板才会彻底暴露。这个现象说明了为什么很多人做算子融合后感觉“收益不明显”——你可能是在小模型小序列上测的。融合的真正价值在大序列、大batch、长解码场景下才会充分显现。如果你只是做小规模的快速验证分离实现加上一个足够大的L2缓存确实可能掩盖问题。但注意这绝不意味着融合不重要。一方面缓存命中率是模型和数据依赖的你控制不了另一方面只要序列增加到超过L2容量分离实现的劣化是陡峭的而融合实现的曲线平缓得多。稳定性在生产系统里比峰值性能更重要。5.5 第三个坑数值误差和精度模式精度方面我对比了float32融合版本和float32参考实现的结果。最大绝对误差在1e-6量级对绝大多数推理场景足够。但如果你用的是float16混合精度事情就麻烦一些。float16下QK^T的求和会把累加误差放大Online Softmax的指数差也是如此。我测试时float16融合版本的误差比分离实现高一个数量级部分极端序列上最大误差到了5e-3。解决方案是qk的累加用float32中间精度只在输入输出和张量存储时使用float16。也就是说所有跨迭代的状态变量——m_i、l_i、acc——都必须保持float32。这点上Triton版实现天然占便宜因为Triton语言里标量累加默认就是高精度手写CUDA时就得自己非常小心一个不留神把acc声明成half后面的误差会让你排查到怀疑人生。6. 融合算子落地检查清单别让原型死在生产环境6.1 从原型到生产的四个额外工作原型验证通过只是第一步。我自己在把这套融合算子接进实际推理链路时发现至少还有四件事是原型的代码里完全没考虑的。第一动态shape支持。真实请求的序列长度每时每刻都在变Kernel不能只支持序列长度的整数倍。这需要在循环边界加mask判断对不足一个block的尾巴单独处理。我用的方法是循环次数取ceil(S / BLOCK_S)每次加载时判断索引是否越界越界填充零且对应qk分数设为负无穷。这个逻辑加进去后kernel的复杂度上升了一个台阶但这是生产环境没法回避的。第二多头数据的布局优化。最初原型里Q、K、V三个矩阵在内存里是分离的每访问一个头都要跳一次大指针。线上模型做完权重合并后QKV通常是连续存储的我调整了指针计算方式让一个线程块在处理多头时尽可能共享同一段缓存行。这个改动从性能上看贡献了约15%来自更友好的内存访问模式。第三算法选择的运行时分支。不是所有头都需要同样的处理比如kv cache里有些头是空的历史有些头已经累积了很长。我在kernel外做了一个轻量判断根据历史长度和序列长度动态选择BLOCK_S和是否启用因果mask单独分支。这个策略让短请求的延迟平均降低了10%。第四持续的回归验证。融合kernel的数值正确性必须纳入每次提交的自动测试里。前几次我还挺惊讶明明小序列测试全过怎么大序列就出了NaN后来定位到是一个线程块处理超长序列时block索引计算溢出。这种问题只有靠覆盖不同序列长度的回归测试才能拦住。6.2 检查自己的场景是否真的需要手动融合做完了这套优化有人问我既然现在FlashAttention和各种推理框架都有现成实现为什么还要自己写。我的看法是现成实现能覆盖80%的通用场景但当你模型结构有自定义mask逻辑、或者需要在attention输出之后马上接一个自定义算子、又或者目标硬件不在主流支持列表里时现成方案往往要么性能不理想要么根本没法改。这时候理解融合原理并有能力写自己的kernel就成了真正的生产力。但反过来我也要劝一句如果你的模型是标准结构、跑在标准框架里直接调优化过的FlashAttention即可没必要重复造轮子。手动融合适合那些已经确认瓶颈在Attention、且现成算子无法满足自身计算模式的场景。6.3 一点实操体会最后分享一个我自己重复踩过两次、最终总结出来的经验融合kernel调试时永远先在小规模数据上跑数值验证再上性能测试不要颠倒顺序。我最早一次做这个项目时直接拿线上序列长度跑测试发现耗时下降是下降了但输出完全不对最后定位到是校正因子里少乘了一项。这种问题在长序列上非常容易掩盖在小序列上因为小序列的误差会被舍入抵消掉一部分。还有一个小细节算是我自己的土办法在Triton原型里故意加入一个print语句把每次迭代的m_i和l_i打出来对照参考实现在相同位置的中间值。只要中间值的顺序一致最后输出就基本不会错。这个办法帮我定位过一次非常隐蔽的bug——那时候是因果mask在边界处把有效位置也一并mask掉了中间值全对但最终输出少了第一个token。如果少了这个步骤这个bug找起来恐怕要靠猜。这套融合算子后续我还在继续扩展比如尝试把LayerNorm和残差也融进同一个kernel里以及在量化场景下验证Online Softmax的中间精度对模型精度的影响。这些都是拿得到收益的方向有机会再单独写一篇。至少眼下这套QKV拼接加Softmax加MatMul的三算子融合方案已经稳定跑在我自己的推理链路里大半年性能曲线很干净没有出现意外劣化。 
长河君 · 2026/10/23 5:18:01 · 阅读 1.2k Attention三算子融合实战：从Triton到CUDA的推理加速优化
这个融合是我去年年底在优化一个线上推理链路时真正落地过的方案。当时模型里Attention占了接近四成耗时profile一打出来三个kernel来回倒腾中间矩阵显存带宽被白白烧掉一大截。这篇文章我把这段经验完整写出来先讲清楚三个算子为什么非融合不可然后从数学原理一步步推到Triton和CUDA的两种实现最后是实测数据和一些非常反直觉的调参坑。如果你也在做transformer推理加速这篇应该能帮你少走不少弯路。先交代一下背景。这里说的ops-transformer不是某个开源框架的名字而是我自己的一个算子优化实验项目目标很单纯把Attention里最常见的三阶段计算——QKV拼接计算、Softmax归一化、最后与V的矩阵乘——合并成一个融合Kernel消除中间结果的显存往返。为了表述方便下文统一叫这个项目名。1. Attention三算子的算路与访存瓶颈分离实现输在哪里1.1 从“三个算子”到完整计算链先把Attention的计算拆开。以单头为例输入X经过QKV投影后得到Q、K、V三个矩阵之后的核心计算是由Q和K计算注意力分数矩阵S Q · K^T再除以缩放因子。对S的每一行做Softmax归一化得到概率矩阵P。用P对V做加权求和O P · V。严格数一下这条链上有四个逻辑步骤投影、QK^T、Softmax、PV。但在实际工程里QKV投影往往通过拼接后的权重一次MatMul完成而且这一步天然可以和后面的QK^T连起来看。所以很多优化资料把它们归纳成“QKV拼接 Softmax MatMul”三大段这里的“QKV拼接”指的是Q、K、V数据在计算图中的汇合方式也可以是物理上把三个投影权重拼成一个矩阵来做合并运算。无论是哪种真正要解决的核心矛盾是同一个中间那个S矩阵太大而且它在分离实现中必须落回显存。我的优化起点就是盯着S矩阵的。当序列长度为512、头维度为64时单个头的S是一个512×512的矩阵float32下就是1MB。12个头就是12MB这还只是写一次下一轮QK^T的kernel读出来又要12MB一写一读就是24MB。这24MB里只有极小的比例是最终要保留的信息大部分带宽纯粹是传送带把数据运出去再运回来。1.2 访存量估算算力根本不是瓶颈用一张表格表达这个问题的直观程度可能比我上面说的一大段都清楚。以单头、序列长度512、头维度64、float32为例阶段中间张量大小单头写入显存读出显存QK^TS矩阵512×512×4B 1MB1MB1MBSoftmaxP矩阵512×512×4B 1MB1MB1MBPVO矩阵512×64×4B 128KB128KB忽略单头就产生了约3MB的显存往来。按12头计算就是36MB。我测试的这张推理卡有效带宽大约在1.5TB/s级别这些中间矩阵搬运就要消耗约24微秒。而真正这三个算子的数学运算量是多少呢QK^T是512×512×64约1600万次乘加PV也是约1600万次乘加两者合计约3200万次浮点操作在这块卡上纯算力执行不足3微秒。也就是说分离实现中超过85%的时间花在了把中间结果写出去再读回来。这个结论直接决定了融合的必要性只要能让S矩阵和P矩阵不出芯片内部的高速缓存性能收益就是结构性的。这还没有算kernel启动的开销。每个GPU kernel从启动到结束都有固定的调度开销几个微秒量级。三个kernel累计10到20微秒的额外延迟在长序列训练场景里也许还能被大计算量掩盖但在解码阶段的短序列场景中kernel启动开销甚至可能超过运算本身。这一点后面实测会再提到。1.3 解码场景下的放大效应训练阶段batch大、序列长计算密集访存问题相对没那么扎眼。真正被放大的是推理场景尤其是自回归解码。解码时每次只生成一个tokenQ只有一行但K和V会随着历史长度增长。分离实现会不断把越来越大的S矩阵写进显存序列一旦长到2048甚至4096单头S矩阵就是16MB到64MB多头累积起来的带宽消耗非常恐怖。我做优化时应用的线上场景正是这种变长解码推理序列长度经常在几百到两千之间波动所以对访存优化特别敏感。算力永远都在等待数据这不是某一个kernel写得不好而是算子边界划错了。2. 融合的数学底牌行级Softmax、Online Softmax与分块点积2.1 核心洞察Softmax只需要行内统计量任何一次算子融合本质都是在问一个问题这一步计算结果是否真的需要作为完整张量落到外部存储对于Attention来说答案显然是否定的。关键在于Softmax的行独立性。Softmax的操作对象是S矩阵的行。第i行的Softmax只依赖这一行的元素不涉及其他行的任何数值。这意味着S矩阵可以按行任意分块处理我先算第i行的一部分分数做局部Softmax再算另一部分最后把结果合并只要合并方式在数值上正确就不需要先物化完整行。这个特性给了融合一个非常重要的自由度QK^T这一步不需要把完整的S算完再交给Softmax而是可以算一小块分数、立刻做一次局部Softmax、把结果累加到最终输出中然后扔掉这一块分数。整个过程中S和P都只存在于片上缓存里。很多人第一次听说时会觉得这违背直觉因为数学公式里P Softmax(QK^T)明明需要先做完整矩阵乘再归一化。但Softmax的形式决定了它根本不需要“全局”信息行的最大值和归一化项才是关键。这也正是下面要讲的Online Softmax存在的理由。2.2 Online Softmax的迭代公式Online Softmax的妙处在于它允许我们用“流式”的方式逐块计算Softmax并且每一次合并都保持数值稳定。具体来说维护两个状态变量m当前见过的所有元素中的最大值初始为负无穷。l当前见过的元素的指数和初始为0。当新的一个数据块x_block到来时先求块内最大值m_block然后合并m_new max(m, m_block) l_new l * exp(m - m_new) sum(exp(x_block - m_new))注意新的归一化项由两部分组成旧部分的贡献需要乘以一个校正因子exp(m - m_new)因为归一化的基准从m变成了m_new新部分的每个元素都减去m_new再求指数避免直接计算exp造成溢出。这里的关键是只要维护好m和l合并顺序就无关紧要任意顺序都能得到正确的最终归一化项。这个迭代形式和经典的“双pass Softmax”相比只做了一次数据遍历但数值效果等价。我们用生活类比来记老数据的指数和中每一项都是以旧的最大值为基准缩放过的现在最大值变大了旧项就得统一乘上一个“缩水系数”把基准调整到新的最大值上。这跟调整天平滑板的过程极其相似——杆子抬高到新的位置所有旧刻度都要按比例重新标注。2.3 在Attention融合中怎么应用在融合Attention Kernel里Online Softmax的迭代过程是这样的取Q的第i行在内部循环里依次遍历K的行块。对每个K行块计算局部QK^T分数块得到s_block。更新m和l用新的m对历史累积输出做一次校正再把当前块的softmax概率乘以V的对应行块累加到输出。遍历完所有K行块后用最终的l对输出做最后一次除法得到Attention结果。每一步中s_block和p_block的生命周期都只存活在寄存器或线程块的共享内存里循环进入下一轮迭代后上一块的数据就可以被覆盖。这恰好和GPU的片上存储层级匹配shared memory容量有限但带宽远高于显存而且同一线程块内共享无需全局同步。分块大小也由此确定——要保证一个迭代步骤中的QK块、概率块和V块能同时放进shared memory。注意这里还有一层隐含的矩阵乘嵌套。最终输出O P · V而P本身来自QK^T和Softmax所以实际是一个两级矩阵乘的复合。如果不融合第一级结果S会被当作中转站写进显存融合后中转站变成片上缓冲。上一节访存量分析已经证明这个中转站恰恰是最大的带宽消耗源因此从显存中把它消灭是收益的核心来源。3. 第一版原型用Triton在半小时内跑通融合Kernel3.1 为什么先用Triton而不是直接写CUDA这个项目我一开始没直接写CUDA。虽然CUDA是最终性能的保证但调试周期长尤其是Online Softmax这类带跨迭代数据依赖的kernel稍不留神就是数值错误。Triton的优势在于你关注的是分块逻辑和计算顺序而不是线程怎么调度、shared memory怎么分配。Triton的编程模型以program为单位每个program处理一个计算块语言内置了块级别的matmul。更重要的是Triton的编译器能自动把这些块操作映射成底层的线程级操作shared memory的分配和复用大部分交给编译器处理。这对快速验证融合逻辑非常友好。我的建议是任何复杂的算子融合先花半小时用Triton做一个正确性可以验证的原型再根据Profile结果决定要不要手写CUDA。真实项目中我见过太多人一上来就写CUDA结果好几个星期耗在边界case上最后发现自己的算法路径一开始就错了。3.2 融合Kernel的核心代码下面这段是Triton版融合Attention的核心实现。为了阅读方便我做了简化假设序列长度S是BLOCK_S的整数倍不处理mask只做单头单query行的演示逻辑。实际多头的循环我在后面说明。import triton import triton.language as tl triton.jit def fused_attn_kernel( Q_ptr, K_ptr, V_ptr, O_ptr, S: tl.constexpr, D: tl.constexpr, BLOCK_S: tl.constexpr, BLOCK_D: tl.constexpr, scale: tl.constexpr, ): row tl.program_id(0) # 当前处理的query行 head tl.program_id(1) # 多头场景下的头编号 offs_d tl.arange(0, BLOCK_D) offs_s tl.arange(0, BLOCK_S) # 加载当前query行 q tl.load(Q_ptr head * S * D row * D offs_d) m_i -1e30 l_i 0.0 acc tl.zeros([1, BLOCK_D], dtypetl.float32) for start_s in range(0, S, BLOCK_S): offs_k start_s offs_s k tl.load(K_ptr head * S * D offs_k[:, None] * D offs_d[None, :]) v tl.load(V_ptr head * S * D offs_k[:, None] * D offs_d[None, :]) # QK^T 得到当前块分数 qk tl.dot(q[None, :], tl.trans(k)) # 形状 [1, BLOCK_S] qk tl.reshape(qk, (BLOCK_S,)) qk * scale # Online Softmax 更新 m_new tl.maximum(m_i, tl.max(qk, axis0)) p tl.exp(qk - m_new) l_new l_i * tl.exp(m_i - m_new) tl.sum(p, axis0) # 校正旧累积输出再加上当前块的PV correction tl.exp(m_i - m_new) * l_i / l_new acc acc * correction acc tl.dot(p[None, :] / l_new, v) m_i m_new l_i l_new out tl.reshape(acc, (BLOCK_D,)) tl.store(O_ptr head * S * D row * D offs_d, out)这段代码的逻辑和上一节的数学推导一一对应。有几个地方值得单独说。第一acc的形状是[1, BLOCK_D]它不是一个矩阵而是最终输出O的一行。tl.dot(p[None, :] / l_new, v)的形状恰好是[1, BLOCK_D]所以累加是逐块的。这里把除以l_new提前到矩阵乘里做而不是最后统一除原因是后面的校正因子需要基于当前l_new来修正旧输出提前归一化反而让数值更稳定。第二qk * scale使用的是缩放因子通常是1 / sqrt(D)。这个缩放必须在Softmax之前完成否则概率分布会被一个常数扭曲。第三也是最容易犯错的一点每次更新acc时校正因子是exp(m_i - m_new) * l_i / l_new注意这里同时包含了最大值修正和归一化项修正。如果漏掉exp那一项短序列上可能看不出问题但长序列上数值会逐渐漂移最终输出会明显偏离正确结果。3.3 原型验证与数值误差控制原型跑通后第一件事不是看性能而是验证数值正确性。我的做法是用PyTorch的标准Attention结果作为参照比较最大绝对误差和平均相对误差。用float32计算时融合版本和参考版本的误差通常能控制在1e-6以下如果误差超过1e-4几乎可以断定是Online Softmax的校正逻辑出了问题。这里有一个非常常见的问题为什么我把初始的m_i设成-1e30而不是负无穷用float表示真正的-inf当然也可以但在Triton的某些后端上exp处理-inf边界时存在优化路径可能导致第一个数据块的校正因子变成NaN。用-1e30可以安全规避这类问题代价仅仅是第一个块内exp(qk - m_i)可能会因为qk - m_i很大而上溢到inf——但实际上第一个块会把m_new更新为块内最大值随后一切都恢复正常。这个细节属于典型的“不做不知道做了才明白”的经验。数值验证通过后我才开始用triton.testing的基准工具测耗时。这时候发现原型已经比分离实现快了不少但还没有达到我的预期。这让我意识到Triton原型适合验证算法要在生产环境榨干性能还是得落到CUDA层面做更精细的调度。4. 从原型到CUDAshared memory、同步与mask的实战细节4.1 线程组织与shared memory的分配策略手写CUDA版本的融合Attention核心设计决策有三个线程怎么排列、shared memory放什么、以及何时需要同步。我采用的方案是一个线程块处理一个query行线程块内用一维128线程每个线程负责输出向量D维中的一部分。以D64为例64个线程分别持有最终输出的一个分量其余的线程用来加速QK^T计算中的并行规约。shared memory里每个迭代块需要放三样东西K的一个行块、V的对应行块、以及当前块计算出的概率向量p。K和V块的大小是BLOCK_S×D以BLOCK_S32、D64、float32计算两个矩阵加一起8KBqk向量和p向量不到1KB整个块需要的shared memory在10KB以内H100和主流推理卡上都不成问题。这也反向验证了一件事BLOCK_S不能无限增大因为shared memory是硬性上限。迭代流程是先把K块和V块从显存加载到shared memory所有线程协作计算当前行的qk部分和再通过块内归约得到局部分数、最大值和指数和。这里需要特别注意同步点——每次写shared memory后必须__syncthreads()否则线程A还在读旧数据、线程B已经在写新数据结果必错。4.2 Online Softmax在CUDA里的同步点处理Online Softmax放在CUDA里真正的难点是m_new和l_new是跨线程共享的标量更新它们需要在块内同步而同步会卡住所有线程的进度。我一开始的实现把同步放得很频繁结果发现每个迭代块都要等两次归约性能惨不忍睹。后来的优化策略是把一次迭代分成两个阶段。第一阶段所有线程并行计算各自负责的局部qk值并做线程桶内的部分归约得到这个线程的局部最大值和局部指数和。第二阶段通过共享内存做一次跨线程归约得到块级的m_new和l_new。归约完成后每个线程都拿到完整的m_new和l_new立即用它更新自己持有的acc分量。这样整个迭代只需要两次__syncthreads()一次是K、V块加载完毕后的数据就绪同步一次是归约结果写入共享内存后的可见性同步。两次同步两个阶段基本是这类融合kernel的标准范式。关键点在于acc这个输出累积向量按线程分布在寄存器里校正因子是标量所以每个线程更新自己的acc时不需要额外的跨线程通信。这是寄存器级数据复用带来的巨大优势——在分离实现中这一步对整个P矩阵做了全局遍历而在融合kernel里它只是寄存器上的一个标量乘加。// 伪代码省略了内存申请和指针计算细节 __global__ void fused_attn(float* Q, float* K, float* V, float* O, int S, int D, float scale) { int row blockIdx.y; int head blockIdx.x; int tid threadIdx.x; float acc[VDIM]; // 每个线程负责的VDIM维输出分量 float m_i -1e30f, l_i 0.f; float q[VDIM]; // 加载q... for (int start_s 0; start_s S; start_s BLOCK_S) { // 阶段1: 协作加载K、V块到shared memory // __syncthreads(); // 阶段2: 计算局部qk, 块内归约m_new, l_new // __syncthreads(); // 阶段3: 校正acc, 累加p * V的局部贡献 } // 写入输出: O acc / l_i }这个伪代码虽然省略了大量细节但结构是正确的骨架。我强烈建议自己动手从这版骨架开始填而不是直接抄完整的开源实现因为Attention的变体实在太多了——有mask、有GQA、有alibi位置编码每加一个变体kernel内部的数据流都要跟着改。骨架理解了改起来才不慌。4.3 Mask处理padding mask和因果mask塞进融合Kernel的两种姿势推理场景几乎都躲不开mask。padding mask是序列实际长度不满时需要在分数矩阵上把无效位置设成负无穷因果mask是解码时只允许关照历史位置。这两类mask在融合kernel里的处理方式略有不同。padding mask最简单。在计算qk分数时如果K的某一行对应的是padding位置直接把该位置的分数设为-INFINITY。这样做不需要额外存储mask矩阵只需要在加载K块时附带一个有效性标志可以用一个小的索引判断valid (start_s i) seq_len。在CUDA里这只是一个整数比较。因果mask更微妙。它要求第i行的query只能看到j i的key。在逐块遍历中对于当前query行row如果块内key索引block_idx全部大于row那么这个块可以完全跳过。如果块内部分有效、部分无效那就用j row ? 0 : -INFINITY的掩码逐元素乘到qk上。跳过整个块的优化对解码阶段尤其重要因为解码时每个query行只处理到目前为止的历史长度尾部大量块都是无效的跳过能省一大半计算。需要提醒的是任何引入-INFINITY的地方都要同时检查Online Softmax更新时会不会出现涉及无穷的数值异常。实际测试中只要序列不是全padding一般都能正常收敛但保险起见我会在exp之前把有效位置的值clip到一个安全范围比如[-85, 0]。4.4 更长序列两阶段切分策略当序列长度超过2048时即使BLOCK_S32单线程块要串行遍历的块数也太多了。更重要的是QK^T和PV在数学上对K和V的需求方向不同单纯增大BLOCK_S会让shared memory压力陡增。我的做法是引入两阶段切分外部按key方向分大块每个大块交给一个线程块线程块内部再对小块跑Online Softmax大块之间通过原子操作或者跨block归约合并部分结果。这种做法本质上就是FlashAttention的多核分区思路不过实现复杂度会高不少。对于绝大多数部署场景两阶段切分其实不是必需的。序列长度在1024以下时单block的串行版本已经能打到接近硬件极限。我之所以提它是因为很多人拿着长序列的benchmark来质疑短序列优化没有意义。我的经验是先做无分区的融合版本把基础收益拿下再根据最长序列需求决定要不要上分区。不要一开始就把架构搞得过于复杂。5. 实测收益与调参陷阱block大小、L2缓存与数值误差的取舍5.1 测试环境与对比基线原型和CUDA版本都跑通后我在自己的测试环境上做了完整的性能对比。测试卡的显存带宽属于主流推理卡水准batch size为1多头数12头维度64覆盖了128到2048五种序列长度。对比基线有两种一是PyTorch原生的Attention实现二是已经做了部分算子合并但仍有显存中转的中间版本。前者代表大多数人的起点后者模拟的是那些“只融合了一半”的优化方案。5.2 加速比数据序列长度分离实现耗时(us)部分融合耗时(us)完整融合耗时(us)完整融合加速比1284538182.5x2568267282.9x512152118453.4x1024310220784.0x20486804301384.9x这个趋势非常符合预期序列越长中间S矩阵越大融合省掉的访存量占比越高加速比越可观。短序列场景下完整融合的收益也没有消失因为分离实现多出来的kernel启动开销被完整吃掉了。这也回应了前面提到的解码场景——短序列虽然计算量小但三个kernel的启动延迟是固定的。我一度以为这个结果已经够了但深入profile之后发现还有潜力。5.3 第一个坑BLOCK_S不是越大越好最开始我把BLOCK_S设为64理由是每轮可以处理更多的K行减少循环次数。实测下来BLOCK_S64在序列长度512时性能反而退化了约8%在1024时退化更多。原因有两层。第一BLOCK_S变大K和V块在shared memory里的占用量翻倍这会挤压编译器做寄存器分配的余量导致寄存器溢出到local memory而local memory在物理上还是显存等于绕了一圈又回到了访存瓶颈。第二更大的块意味着更少的并行块数在batch size1的场景下GPU的线程块并发度本来就不高再减少块数会让SM的占用率雪上加霜。我在同样的硬件上试了BLOCK_S16、32、64三档。最优值稳定在3216时因为循环次数翻倍调度开销开始显现性能略低。最终的结论是BLOCK_S32对于头维度64、float32的典型配置是一个健康的默认值。5.4 第二个坑L2缓存命中率与“伪融合”有件很有意思的事我优化到一半时曾经怀疑融合到底有没有用因为把融合kernel单独跑耗时竟然比分离实现只快了一点点。后来抓了L2缓存命中率才发现分离实现在小序列下大量S矩阵幸运地命中了L2缓存显存访问被隐藏了。序列一旦拉长L2放不下了分离实现的短板才会彻底暴露。这个现象说明了为什么很多人做算子融合后感觉“收益不明显”——你可能是在小模型小序列上测的。融合的真正价值在大序列、大batch、长解码场景下才会充分显现。如果你只是做小规模的快速验证分离实现加上一个足够大的L2缓存确实可能掩盖问题。但注意这绝不意味着融合不重要。一方面缓存命中率是模型和数据依赖的你控制不了另一方面只要序列增加到超过L2容量分离实现的劣化是陡峭的而融合实现的曲线平缓得多。稳定性在生产系统里比峰值性能更重要。5.5 第三个坑数值误差和精度模式精度方面我对比了float32融合版本和float32参考实现的结果。最大绝对误差在1e-6量级对绝大多数推理场景足够。但如果你用的是float16混合精度事情就麻烦一些。float16下QK^T的求和会把累加误差放大Online Softmax的指数差也是如此。我测试时float16融合版本的误差比分离实现高一个数量级部分极端序列上最大误差到了5e-3。解决方案是qk的累加用float32中间精度只在输入输出和张量存储时使用float16。也就是说所有跨迭代的状态变量——m_i、l_i、acc——都必须保持float32。这点上Triton版实现天然占便宜因为Triton语言里标量累加默认就是高精度手写CUDA时就得自己非常小心一个不留神把acc声明成half后面的误差会让你排查到怀疑人生。6. 融合算子落地检查清单别让原型死在生产环境6.1 从原型到生产的四个额外工作原型验证通过只是第一步。我自己在把这套融合算子接进实际推理链路时发现至少还有四件事是原型的代码里完全没考虑的。第一动态shape支持。真实请求的序列长度每时每刻都在变Kernel不能只支持序列长度的整数倍。这需要在循环边界加mask判断对不足一个block的尾巴单独处理。我用的方法是循环次数取ceil(S / BLOCK_S)每次加载时判断索引是否越界越界填充零且对应qk分数设为负无穷。这个逻辑加进去后kernel的复杂度上升了一个台阶但这是生产环境没法回避的。第二多头数据的布局优化。最初原型里Q、K、V三个矩阵在内存里是分离的每访问一个头都要跳一次大指针。线上模型做完权重合并后QKV通常是连续存储的我调整了指针计算方式让一个线程块在处理多头时尽可能共享同一段缓存行。这个改动从性能上看贡献了约15%来自更友好的内存访问模式。第三算法选择的运行时分支。不是所有头都需要同样的处理比如kv cache里有些头是空的历史有些头已经累积了很长。我在kernel外做了一个轻量判断根据历史长度和序列长度动态选择BLOCK_S和是否启用因果mask单独分支。这个策略让短请求的延迟平均降低了10%。第四持续的回归验证。融合kernel的数值正确性必须纳入每次提交的自动测试里。前几次我还挺惊讶明明小序列测试全过怎么大序列就出了NaN后来定位到是一个线程块处理超长序列时block索引计算溢出。这种问题只有靠覆盖不同序列长度的回归测试才能拦住。6.2 检查自己的场景是否真的需要手动融合做完了这套优化有人问我既然现在FlashAttention和各种推理框架都有现成实现为什么还要自己写。我的看法是现成实现能覆盖80%的通用场景但当你模型结构有自定义mask逻辑、或者需要在attention输出之后马上接一个自定义算子、又或者目标硬件不在主流支持列表里时现成方案往往要么性能不理想要么根本没法改。这时候理解融合原理并有能力写自己的kernel就成了真正的生产力。但反过来我也要劝一句如果你的模型是标准结构、跑在标准框架里直接调优化过的FlashAttention即可没必要重复造轮子。手动融合适合那些已经确认瓶颈在Attention、且现成算子无法满足自身计算模式的场景。6.3 一点实操体会最后分享一个我自己重复踩过两次、最终总结出来的经验融合kernel调试时永远先在小规模数据上跑数值验证再上性能测试不要颠倒顺序。我最早一次做这个项目时直接拿线上序列长度跑测试发现耗时下降是下降了但输出完全不对最后定位到是校正因子里少乘了一项。这种问题在长序列上非常容易掩盖在小序列上因为小序列的误差会被舍入抵消掉一部分。还有一个小细节算是我自己的土办法在Triton原型里故意加入一个print语句把每次迭代的m_i和l_i打出来对照参考实现在相同位置的中间值。只要中间值的顺序一致最后输出就基本不会错。这个办法帮我定位过一次非常隐蔽的bug——那时候是因果mask在边界处把有效位置也一并mask掉了中间值全对但最终输出少了第一个token。如果少了这个步骤这个bug找起来恐怕要靠猜。这套融合算子后续我还在继续扩展比如尝试把LayerNorm和残差也融进同一个kernel里以及在量化场景下验证Online Softmax的中间精度对模型精度的影响。这些都是拿得到收益的方向有机会再单独写一篇。至少眼下这套QKV拼接加Softmax加MatMul的三算子融合方案已经稳定跑在我自己的推理链路里大半年性能曲线很干净没有出现意外劣化。