企业官网建设流程全解析

我拆过不少深度学习框架里的Softmax算子，CANN ops-nn这套算子库里的实现算是比较有代表性的一个。无论GPT还是BERT，注意力机制里QK^T之后的那一步，都要靠Softmax把注意力分数变成概率分布，这一步算得稳不稳、快不快，直接决定了整个模型训练能不能顺顺利利跑下去。说它是注意力机制的心脏，一点都不夸张。

这篇博文围绕CANN ops-nn中的Softmax算子，从数学原理、算子契约、内核实现到性能优化完整过一遍，包括很多资料里没有讲透的数值稳定性处理逻辑和轴（axis）语义，最后附上我实际调试过程中踩过的坑和排查思路。适合正在用CANN做算子开发、想搞懂Transformer推理训练性能调优的开发者，以及对NPU算子实现感兴趣的算法工程师。哪怕你完全不写算子，理解了Softmax数值稳定的处理方式，对你排查训练中反复出现的NaN发散问题也有直接帮助。

1. 注意力机制里的Softmax：先搞懂它到底在算什么

1.1 一个公式看懂Softmax在注意力中的位置

Transformer的自注意力路径可以缩写成这样：

Attention(Q, K, V) = softmax(Q * K^T / sqrt(d_k)) * V

Q和K做矩阵乘之后，得到的是一个形状为 [batch, head_num, seq_len, seq_len] 的注意力分数矩阵，每个元素代表某个query与某个key之间的匹配程度。这个分数是实数，可大可小，有正有负，范围没有约束。而后面V的加权求和需要一个"权重"——也就是归一化到0到1之间、且每行之和为1的分布。完成这个转换的，就是Softmax。

Softmax的数学定义非常简单：

softmax(x_i) = exp(x_i) / Σ_j exp(x_j)

对一行分数里的每个元素取指数，再除以整行指数之和。但小学算数级别的公式，放到工程实现里却藏着不少坑。

1.2 Softmax在Transformer里的特殊性

很多人觉得Softmax是个老掉牙的算子，没啥好讲的。但在Transformer场景里，它有几个非常独特的性质，直接决定了算子实现要往哪个方向优化。

第一，它是"按行"操作的。注意力分数矩阵最后一维是特征维度，Softmax严格沿着这个维度做归约。这和图像分类里常见的Softmax（沿着类别维度）在数据排布和并行策略上有明显差异。

第二，它处在矩阵乘和另一个矩阵乘的中间。前有QK^T，后有·V，性能上稍有拖沓，整个Attention的耗时就会放大。FlashAttention等一系列优化工作的核心之一，就是把Softmax和前面的矩阵乘融合起来减少访存。单独实现一个Softmax算子时，也必须考虑数据能不能尽量留在片上、少来回搬运。

第三，注意力分数在被Softmax处理之前，数值范围可能很夸张。尤其是训练初期或者长序列场景下，QK^T的结果方差很大，直接取exp很容易爆掉。这引出了Softmax实现里最经典的问题：数值稳定性。下面细说。

2. CANN ops-nn算子库：Softmax在体系里的位置

2.1 ops-nn是什么

CANN的全称是Compute Architecture for Neural Networks，它提供了从上层框架到NPU硬件的一套完整接口生态。ops-nn是这套体系里的神经网络算子库，里面封装了卷积、池化、归一化、Softmax、LayerNorm、矩阵乘等常用算子。它的定位有点像其他生态里的cudnn或oneDNN，但实现是针对NPU的异构架构专门做的。

我在实际开发中感受到，ops-nn里的算子通常分两类：一类是直接暴露给上层框架调用的标准算子，比如Softmax、LayerNorm，带完整的shape推导和类型推导；另一类是供其他融合算子内部使用的底层原语。Softmax属于前者，同时它又经常被融合进Attention的大算子中。

2.2 Softmax算子的输入输出契约

先规整一下算子的接口语义。源码里Softmax算子接受以下几个关键参数：

• input：输入张量，任意维度
• axis：沿哪个维度做Softmax，默认是-1，也就是最后一维
• output：输出张量，shape与输入完全一致
• 数据格式与精度：支持float16、float32，部分场景还涉及bfloat16

这里最容易被忽视的就是axis的语义。我们平时写PyTorch，torch.softmax(x, dim=-1)很自然。但在算子框架里，axis要翻译成对底层数据排布的具体操作。CANN内部有ND格式、NCHW格式，还有专门用于矩阵运算的NZ分形格式。不同的数据排布下，同一轴上的归约在内存里可能连续，也可能跨着很大的stride跳着取数，这直接影响实现策略。

我在用ops-nn的Softmax时发现，它内部其实把各种场景分成了几类：低维小数据量的直接暴力算、中等维度用向量指令做多趟归约、高维大shape则做多核切分。下面按这个分层思路展开。

3. Softmax实现原理拆解：从朴素公式到算子内核

3.1 朴素实现为什么不行

如果完全照搬数学公式，一个最简单的实现是：

for i in 每个元素: exp_x[i] = exp(x[i]) 输出[i] = exp_x[i] / Σ_j exp_x[j]

这个写法在数学上没问题，但在实际硬件上跑会有两个致命问题。

第一个是数值溢出。假设某一行注意力分数里有一个元素是100，exp(100)大约是2.7e43，这个数已经超出了float32的常见表示范围边缘，在float16下更是直接变成Inf。Inf再参与除法，结果就是NaN或者Inf，训练直接崩掉。

第二个是计算效率。朴素实现为了得到分母的sum，必须先完整遍历一遍整行数据，把exp结果存下来或者再读一遍，才能做除法。这意味着每行数据至少要进出内存两次以上。在NPU这种带宽受限、注重数据片上复用的架构上，多次访存的代价比多算几次加减法大得多。因此工程实现几乎不会做这种"教科书版"。

3.2 经典三步法：从减最大值开始

工业界标准的Softmax实现是"减最大值"的三步法：

第一步：x_max = max_j(x_j) // 找到该行的最大值 第二步：sum = Σ_j exp(x_j - x_max) // 对减掉最大值后的指数求和 第三步：out_i = exp(x_i - x_max) / sum // 归一化输出

为什么要减去最大值？因为exp是一个单调递增且增长极快的函数。我们把所有输入都平移同一个常数x_max，数学上Softmax结果完全不变，因为分子分母同时多乘了一个exp(-x_max)。但这个平移把数据的范围压到了(-∞, 0]，exp之后的结果最大也就1，彻底避免了溢出风险，也保证了数值精度。

CANN ops-nn的Softmax内核主体就是这套三步流程。但真正的工程细节在于：这三步在硬件上怎么排，怎么减少中间结果的落地。

3.3 内核里的数据流：一个tile一个tile地算

NPU上有个特点，片上内存（一般叫Unified Buffer或类似概念）容量有限，不可能把整个大张量一次性放进来。所以算子需要把数据切块（tile），一块一块搬到片上，算完再搬出去。对于Softmax这种沿最后维度归约的算子，切块方案直接影响性能。

对于 [batch, head, seq, seq] 这种四维注意力分数，最后两个维度是seq_len，而Softmax沿最后维度做。通常的切法是把前两个维度展开成一个大batch，每个"行"对应一次归约单位。假设seq_len=512，即每行512个元素，一个tile可以放多行。为了减少反复搬运，最优策略是让一个tile尽量覆盖完整的行——因为行内的数据必须全部算完才能归一化，如果一行被切成两半，两半之间就得跨tile同步max和sum，非常麻烦。

我在看内核实现时注意到一个细节：它先用向量指令找出当前tile内每行的局部最大值，再做exp和局部求和，最后才在寄存器或片上内存里做跨tile的归约。这种"先各自算局部、再合并"的思路，是为了避免把max和sum两个大数组长期驻留在片上，节省空间给后续的tile。

4. 性能优化：从两遍算法到一遍算法

4.1 两遍遍历带来的性能瓶颈

严格按三步法做，每一行数据需要被访问两次：第一次求max，第二次求sum。如果每个tile的数据可以完整缓存，那么两次访问都在片上完成，性能损失可控。但更大的模型、更长的序列会让tile无法完整覆盖整行，数据要从全局内存读两遍，这就很难受了。

针对这个问题，业界有一个经典优化，叫"在线Softmax"（Online Softmax），也叫流式Softmax。它的核心思想是：只遍历一遍数据，边遍历边维护运行中的最大值和运行中的指数和，发现更大的值时对旧和做修正。

初始化 running_max = -Inf, running_sum = 0 每处理一个新元素 x: 新max = max(running_max, x) 修正系数 = exp(running_max - 新max) running_sum = running_sum * 修正系数 + exp(x - 新max) running_max = 新max 结束 输出时：out_i = exp(x_i - running_max) / running_sum

这里的关键是那个exp(running_max - 新max)。你可以把它理解成：之前所有指数结果都是基于旧max算的，现在max变大了，需要乘一个小于1的系数，把它们整体"压低"到新max的尺度上，再加进新的exp贡献。我在实际项目中验证过，这种修正和一次遍历算出的结果是完全一致的（浮点误差范围内），但带宽消耗几乎减半。

4.2 多核切分：batch维度是最自然的并行面

NPU通常有多个AI Core，Softmax并行切分时最省事的维度是batch方向。比如 [8, 32, 512, 512] 的分数矩阵，完全可以按前两维（共256个独立的（batch, head）组合）分配给不同核，每个核独立处理自己负责的若干行，互不通信。这种切分的好处是天然无依赖，不需要跨核归约，性能和核数基本线性扩展。

但如果序列特别长、一行数据本身就大到单核放不下，那就得走分块-跨核归约的路线：多个核分别算同一行的部分max和sum，最后通过核间通信合并。CANN的算子实现里对这种场景也有考虑，不过复杂度明显上升——合并时的修正系数、同步点、负载均衡都要额外处理。我在实际调优中看到，对于seq_len在2048以下的场景，按行完整分配通常已经够用，没必要上跨核归约。

4.3 精度选择：看着简单的除法其实有讲究

再看精度问题。Transformer训练混合精度下，QK^T的输出可能是fp16的，注意力分数矩阵精度有限。Softmax内部的exp和求和如果用fp16算，误差会快速累积，最后归一化的分布可能出现明显的"尾部失真"。所以内核里普遍的做法是：输入输出保持fp16，但中间的max、sum和exp的计算结果用fp32累加器承接，最后再转回fp16输出。

这个选择背后的逻辑不复杂：exp本身是一个非线性放大，对输入误差很敏感；加法的累加如果也在低精度下做，长序列时误差随长度累积。而fp32累加器的成本在向量单元上并不高。我在对比实验里实测过，全fp16链路和fp32累加的版本，在对齐后的输出上能差出好几个ulp，尤其当一行里有动态范围很大的元素时，差异会被放大。

5. 实操中常见的Softmax问题与排查心得

5.1 输出出现NaN怎么办

如果网络前向里Softmax输出了NaN，首要怀疑对象不是Softmax自身，而是上游QK^T的结果里已经有Inf或者NaN。Softmax对输入做exp，输入是NaN则输出必是NaN。排查时我会先打印输入张量，统计是否有异常值，再看是特定step才出现还是从一开始就发散。

如果确认是Softmax内部溢出导致的，检查有没有做减最大值的处理，或者减最大值的实现是否正确。一个容易被忽略的细节是：有些优化实现为了省一次遍历，会把"减最大值"和"求sum"合并，但如果实现里对max的初始值设置成0而不是-Inf，遇到全负数的行就会算错。我调试过一个类似问题，一行全是负值的时候，初始max=0会让所有元素都减去0，exp输出依旧可能非常大，数值稳定性前提就失效了。

5.2 精度对不上参考实现

算子开发的同学经常会和PyTorch或参考实现做精度对齐。Softmax对不齐的场景，大概率出在两个方面：

一是axis理解不一致。输入是 [16, 32, 512, 512]，如果框架侧传的是dim=-1，算子侧解析成第3维，但如果中间有transpose或者reshape改变了数据排布，很容易把dim=-1悄悄映射错。做精度对齐前，先打印两边的输入shape和内存排布，永远比盯着diff快。

二是低精度累加。前面说过，fp16中间累加是误差大头。参考实现通常用fp32算得干干净净，算子侧如果没开fp32累加器，diff通常在1e-3量级，看起来不多，但在LayerNorm+Softmax层层叠加后，误差会被放大到下游矩阵乘里。建议统一用fp32累加器做底，再对比相对误差是否在可接受范围内。

5.3 性能比预期差

Softmax算子性能差，先别急着怪内核逻辑，先看数据排布。如果上游输出的数据格式是NCHW或者某种分形排布，而Softmax沿最后一维归约，恰好归约维度在内存上不连续，向量单元取数要带stride，效率直接掉一个档次。CANN里常见做法是让Softmax尽量吃它最熟悉的连续尾维排布，如果上游给了非连续排布，宁可先做一次转置，也不要硬着头皮带stride算。

还有一类性能问题来自"小算子"调用开销。如果Attention里每次只算一行几十个元素的小Softmax，却走一遍完整的算子下发流程，核间调度开销占比极高。这种情况下不要单独调Softmax，而是要用融合算子把QK^T、Softmax、乘V合成一个kernel。CANN社区里很多Attention融合示例就是这么做的，Softmax作为内核里的一个阶段直接在片上完成，不落全局内存，效果立竿见影。

6. 一层一层扒完之后，给你留个可复用的检查清单

文章写到这里，核心的Softmax实现原理和CANN ops-nn里的关键设计都过了一遍。最后给你整理一份我在实际算子开发和性能调优中反复对照的检查清单，直接拿去用：

• 实现有没有做减最大值？初始max是-Inf还是0？
• 中间累加用的是fp16还是fp32？长序列时是否验证过精度？
• axis语义和上层传参是否一致？有没有隐式reshape改变维度含义？
• 每一行数据是否完整落在单个tile内？跨tile的max/sum归约逻辑是否正确？
• 并行切分是否按batch/head进行？是否存在跨核通信热点？
• 是否为零散的"小Softmax"付出了过大的算子调度开销？

我个人最大的体会是：Softmax这种"看起来一秒就能写完"的算子，恰恰是最值得抠细节的地方。它的归约逻辑简单，但数值特性敏感、与上下游算子的衔接方式又多，任何一个环节放松都可能成为训练稳定性的隐患或性能瓶颈。你在自己项目里如果遇到和Softmax相关的奇怪问题，欢迎按上面几个维度去排查——大概率能在十分钟内定位到根因。