1. 为什么一个Softmax还要专门做数值稳定
第一次在某个异构计算架构的算子库（就是这里说的CANN ops-math）里看到专门为Softmax写数值稳定方案时，我心里其实有点不以为然——Softmax不就是exp(x) / sum(exp(x))吗，教材上写了几十年的公式，照着实现还能出什么问题？
结果第一次在极端数据下跑测试，我的看法就彻底变了：输入里有几个元素超过了fp16能表示的数值上限（大约 65504），整个输出直接变成NaN，那一整批数据全都废了。后来排查下来才发现，问题根本不是Softmax公式本身，而是在计算机的有限精度表示下，朴素的实现方式会把中间结果撑爆。
这个问题的本质用一句大白话讲就是：exp(100)在数学上是有限值，但具体到某个浮点格式，它能表示的最大数就那么一点，超过了就是Inf，再往下做除法自然就是NaN。Softmax的理论公式适用于连续数学世界，但计算机里的浮点数是个有界离散集合，你的实现必须迁就这个现实。
这篇博文就围绕这件事展开。我会从数学推导、实际算子代码、测试验证三个层面，把Softmax数值稳定技术完整拆开讲透。适合的人群是正在做算子开发、AI框架算子适配的工程同学，以及想搞懂为什么框架里Softmax实现看起来“多此一举”的研究者。如果你是刚接触异构计算的新手，也能看懂，我尽量把每一步“为什么这么做”都交代清楚。
2. 溢出场景的完整解剖：数值不稳定的根源
2.1 一个案例看懂朴素实现的隐患
先做一个快速的心理实验：假设有一个向量[1.0, 2.0, 3.0, 100.0]，用支持fp32的编程环境去算exp(100)，结果是2.688117e43。这个数字本身没问题，因为fp32的最大值大约是3.4e38，能装得下。
但如果你用的是fp16，那exp(100)是一个什么结果？fp16最大只能表示到65504，而exp(100)约等于2.7e43，远远超出上限。硬件会把它饱和处理为Inf，也就是正无穷。然后exp(100)/Inf这种运算必然制造NaN，因为规范里就规定Inf/Inf是 NaN。一个元素出问题，整组归一化后的结果全部失效。
我实际碰到的一个场景是在做某个跨平台模型迁移时，模型某层的输出经过多层累加后，出现了几个接近80以上的离群值。但我们的推理推理阶段使用的是fp16计算，Feed到Softmax的输入本来就超过了安全水位线。这种情况下，百分之百复现NaN结果，模型输出完全不可用。
这种情况不是偶发。很多模型训练或推理的前几轮看起来正常，但一旦batch里出现极端数据（特征分布重尾、叠加了数值误差的梯度等），Softmax就成了最脆弱的环节。
2.2 fp16下数值范围被压缩的问题
fp16的数字分布和fp32不同。它的数值范围窄、精度也低。把fp16的正数范围列出来看：
- 最大值：65504
- 最小正规格化数：约
6.1e-5 - 最小正非规格化数：约
6e-8
如果输入维度是[batch, seq_len, vocab]这种结构，seq_len和vocab动辄几千，Softmax 内部的中间值只被极端输入拖着走，很快就超出fp16的表达能力。这不光是训练场景的问题，推理阶段如果用了 fp16 或混合精度，同样要面对。
即使不用fp16，用bf16的时候情况更麻烦。bf16的取值范围和fp32一样大，但精度极低。一旦输入动态范围大，exp后的值和分母一加，误差会瞬间膨胀。这也是为什么数值稳定技巧不是“fp16专属”，而是所有低位宽浮点都需要的基础防护。
3. log-sum-exp：化解溢出的数学武器
3.1 恒等式的推导过程
真正稳定Softmax的数学基础，是一个看似简单但威力很大的恒等变换。核心思想是：既然exp(x_i)在x_i比较大的时候会溢出，那我就先让所有输入整体“下移”，把它控制在安全区间内。
数学上使用的技巧是Log-Sum-Exp（LSE）：
softmax(x_i) = exp(x_i) / sum_j exp(x_j) = exp(x_i - M) * exp(M) / (sum_j exp(x_j - M) * exp(M)) = exp(x_i - M) / sum_j exp(x_j - M)这里M是向量中的最大值。因为exp(M)在分子分母中都出现了，可以约掉。而x_i - M <= 0恒成立，所以exp(x_i - M) <= 1，永远不会超过 1，也就从数学上避免了溢出。
用生活里的例子类比：想象你要比较30个人中谁的生日最接近今天，你不会直接看每个人出生那一刻的全球日期（相当于x_i本身），而是先问“今天几号”，然后拿每个人的生日跟今天做差（相当于x_i - M）。差值的绝对值最大也就是“今天到最远生日”的距离，永远不会炸掉。这就是同一种思想：先找一个参考点，把数字整体搬到安全的区间，再做后续计算。
这个恒等式在数学世界里是“恒等”的，不会改变结果；但在浮点数世界里，它改变了中间计算的数值范围，直接影响了结果的正确性。
3.2 减法之后为什么就安全了
你可能会想：那如果x_i极小（负几百），x_i - M也会是一个绝对值很大的负数，exp的结果会下溢为0。下溢确实存在，但它不是灾难性的——一个元素下溢为0，最多让对应位置的Softmax输出变成0，其他位置的输出依然正确。因为分母会的正数项主导，0 不会污染全局。
这一点和溢出完全不同。溢出产生的是Inf和NaN，它们会传染：Inf - Inf = NaN，NaN参与一切运算都会产生NaN。而0参与加法运算只会保持不变，不会传染。
所以在数值稳定设计中，我们的目标是“避免溢出”，下溢可以允许，必要时再做精度补偿。这就像一个精算原则：优先级最高的是不让系统整体崩溃，局部损失（某些位置输出为0）是完全可以接受的。
3.3 极端情况下下溢是否会带来大误差
为了避免让你觉得我在回避问题，也把下溢可能带来的精度影响说清楚。假设向量里有100万个数，大部分是极端负数，只有一个是0。做x_i - M后，那些极端负数会变成巨大的负值，exp后为0。分母sum仍然是1（只有那个0位置贡献exp(0)=1），其他位置的Softmax输出为0。这个结果从数学上讲也是对的。
真正需要警惕的精度问题是：如果输入的动态范围非常大，而且我们使用低精度计算，那么把一个很大正数和几个很小正数累加时，小数的贡献可能被“吃掉”。这是累加精度的问题，不是log-sum-exp变换本身的问题。解决累加精度的方法通常是在做分母求和时用更高精度的中间累加器（比如fp32），或者在求和时做分段求和。这个后面工程部分会展开。
4. 算子在CANN平台上的落地实现
4.1 CANN算子开发的基础思路
CANN是面向特定AI处理器的一个异构计算架构。写算子不是像写普通CPU代码那样一个for循环搞定，而是要考虑到：数据从主存搬运到片上、在算子上逐级计算、结果再写回主存。整个过程由多个“循环”和“流水”构成。
在CANN框架下实现Softmax算子，通常有两种思路：
- 整体实现：在一个算子内部完成
reduce_max、sub、exp、reduce_sum、div全套操作。 - 组合实现：把Softmax拆成多个基础算子，靠算子间调度完成。
组合实现的好处是复用已有算子，开发量小，但缺点是要多次读写数据，访存开销大。整体实现则相反，一次数据搬运后所有计算都在片上完成，性能更好，工程复杂度也高。我最终采用的方案是整体实现，核心理由是：性能要求和算子融合趋势都在朝这个方向走，一次搬运永远好过中间结果来回倒腾。
4.2 从朴素的过桥算法到稳定enabled的实现
我把整个开发过程梳理成两个阶段：
第一阶段：过桥算法（朴素版）。先求最大值M，然后每个元素减去M，再做exp，求和，最后逐元素除法。这个版本在fp32下大部分场景能工作，但与硬件流水线配合得一般，且没有做任何数据布局优化。
第二阶段：稳定版。重点改造三件事。第一，把最大值计算和元素迁移放到一起，利用“两趟循环”的模式——先扫一遍求max，第二遍计算时直接减去max并做exp累加。第二，累加分母时使用更高位宽累加器，避免大量小数值被截断。第三，对特殊shape做Tiling分块，避免一次处理过大张量导致显存或片上内存爆炸。
4.3 Tiling策略：超越shape上限的关键
通常一个算子不能假设输入数据完全能塞进片上内存。假设你的硬件片上内存只有256KB，但输入是一个非常大的[1, 50000, 128]张量，一次装载必然失败。因此Tiling策略是算子实现里的必修课。
Softmax的Tiling有个特殊约束：分母的计算需要“全局视角”，也就是整个向量上的exp和。如果分块了，常见的做法是分两次遍历：
- 第一遍：按块扫描，每块只求局部最大值。
- 第二遍：按块计算
exp(x_i - M_global)，并且把每一块的局部和记录到临时空间。 - 第三遍：汇总所有局部和，得到全局分母，再做逐块归一化输出。
这里有优化空间：大多数算子实现里没有第三遍真正“等待”全局分母完成再输出，而是把第二步的中间结果先暂存，等第二步完成后原地做除法，省掉了一次额外数据搬移。如果你是新手，建议按三步走先把功能跑通，再考虑融合优化。
关于Tiling的块大小选择，我的经验是“不要贪大”。块越大虽然循环次数少，但一旦出现尾块（数据量不能被整块整除）或者对齐问题，处理起来麻烦且容易踩性能坑。实测下来，选择片上内存能装下的最大且对齐友好的块大小，性能最稳。
4.4 尾块处理和边界情况
尾块是Tiling绕不开的问题。比如总长度为10000，块大小定为4096，前两块各处理4096个元素，剩下1808个元素就是尾块。尾块处理方式有两种：一种是单独写一个非对齐分支处理，另一种是把尾块用0值填充成完整块，在计算时跳过这些位置。前者实现更直观，后者在某些硬件上可以避免分支跳转带来的性能损失。
我个人的实践是：如果硬件支持自然的掩码操作，优先用掩码；否则用显式分支处理尾块。因为掩码方式让代码更统一，分支分歧少，整体吞吐也更好统计。
需要注意：Softmax计算中尾块填充不能用任意值。如果填充了某个很大的正数，它会被当做一个“假元素”参与max和exp计算，污染结果。所以填充必须用-Inf或者一个足够小的负数（比如-1e38），这样它对max的贡献可忽略，exp后也会趋于0。这里一定要在实现时保护到位。
5. 工程优化与踩坑实录
5.1 中间结果的数据类型选择
这是所有实现里最容易“看起来没事、实际错得离谱”的点。我踩过一个大坑：在某个版本里，把max和sub的中间结果直接存成输入同类型（当时是fp16），然后在累加分母时也用fp16累加，结果小概率出现比较大的误差。
原因很简单：fp16只有大约3位十进制精度，分母累加时如果个别项特别大，后面所有小项的贡献都会被截断。举个具体例子：分母累加过程中已经累计到100.0，再往里加0.1，fp16下100.0 + 0.1会被舍入成100.0，小数贡献直接丢失。当数量级差异超过100倍时，这种丢失是必然的。
我的建议：内部计算统一用较高精度。具体来说，exp的输入使用fp32，同时分母累加器使用fp32。如果硬件对fp32的支持弱而fp16计算单元快，那至少也应当使用fp32累加器，因为一个累加器的开销是无足轻重的，但精度收益是很大的。
实际工程里还有一个更低成本的补偿方案：对累加结果做Kahan求和补偿，或者并行树状求和。但如果你刚起步，先把高精度累加器做了，能解决90%以上的精度问题。
5.2 适配Kernel的细节校正
CANN算子开发通常使用类C++语言，配合向量化指令。一个容易犯的错是：写算子时假设内存连续性完美、元素一一对应，结果在真实硬件上因为数据对齐问题出现性能断崖。
以我们常用的[batch, row, col]布局为例，Softmax通常是作用在最后一个维度上的。如果col长度不是硬件向量长度的整数倍，向量化加载就变得很别扭。常见解决方案：
- 如果
col是向量宽度的倍数：直接顺序访问，最简单。 - 如果不匹配：优先对
col做尾部填充对齐，而不是写复杂的内部循环。 - 如果
col太长需要分块Tiling：保证子块也满足对齐要求。
另一个校正点是：指数计算通常有硬件指令支持，但这类指令的输入范围有限制。以某些低精度硬件为例，即使你计算了x_i - M，当x_i - M还是正数或0时，安全性没问题；但需要确认硬件指令对负指数范围的处理是否精确，有些硬件会对极小输入直接返回0，误差更大但毕竟不会产生NaN。理解底层指令的边界行为，是实现质量的关键。
5.3 性能剖析：一次搬运与多次搬运的差距
理论上Softmax计算应该在“一趟流水”内完成，但实际会因为前面说的三步流程带来访存差异。我拿同一个[2048, 64]的输入做过对比：
- 未做融合的组合算子：3次完整读写（max一次、sub/exp一次、sum/div一次），耗时大约是融合方案的1.8倍。
- 融合但未做空间优化：2次读写，性能居中。
- 完整融合加原地更新：1次读写 + 少量片上临时缓存，性能最好。
在CN/推理场景，这种差异会直接影响端到端延迟。如果你只是想快速验证功能，用组合算子完全可行；但如果要做推理部署，整体融合算子几乎是必选项。
6. 数值稳定性测试与排查实录
6.1 测试矩阵的构建方法
不要只测常规数据。一个数值稳定方案如果只在“正常范围”内验证通过，根本证明不了什么。我给Softmax算子建的测试矩阵包含以下几类：
- 常规输入：均匀分布，值域在
[-1, 1]。 - 大正数输入：值在
[80, 100]范围内的少量离群点。 - 大负数输入：值在
[-100, -80]范围内的输入。 - 混合极端：同时包含
100和-100的向量。 - 动态范围极大：从
[0, 0.001]跨越到[100, 1000]。 - 零向量和全等向量：检验分母是否为0或均匀分布。
测试标准我定得很死：输出必须全部在[0, 1]区间内，整行求和必须等于1.0（允许fp16下千分之一的误差），不得出现NaN或Inf。
6.2 一例典型的“NaN出厂”问题排查
有一次在跑混合极端测试时出现NaN，排查过程值得记录。第一步，我怀疑是exp硬件指令的输入溢出，于是打印了所有进入exp的数值，结果发现最大不过0.0，理论上不会溢出。这个排查方向上存在盲区：问题不在exp本身，而在分母的计算。
后来我定位到：在某个分支里，如果整块数据全为-Inf填充尾块，局部max计算返回-Inf，后续x_i - M出现Inf - Inf，直接生成NaN。
这个陷阱非常隐蔽。修法是在max计算处加一个判断：如果局部最大值仍然是-Inf，则把该块的max标记为一个安全的负数（如-1e38）。这样即使整块数据都是极端负值，x_i - M也只是一个普通的正数，不会产生Inf。
这里你也能看出这类数值问题的特点：几乎总是藏在边界条件里，而常规测试根本覆盖不到。没有极端数据的测试矩阵，等于没测。
6.3 快速查错的工具技巧
如果你面对一个已经产出NaN的算子，不要直接从汇编或者底层指令开始查。那是最慢的方式。我习惯的做法是：
- 第一步：用模拟器或CPU参考实现跑同一组输入。对比结果，确定
NaN是硬件特有行为还是算法本身的问题。 - 第二步：在关键计算节点（max、减、exp、累加、除法）后分别打印统计量：最大值、最小值、NaN个数。能把问题定位到某一层。
- 第三步：如果是fp16的问题，把中间结果用fp32重算一遍，如果不出现NaN，就基本确认是精度导致的，不需要去查指令集。
这套流程90%的情况下都能快速定位问题。剩下10%才是指令异常、内存越界这类刁钻问题，那才需要把目光下沉到硬件调试工具和汇编层。
7. 几个值得后续扩展的方向
如果你把基础版Softmax稳定算子已经写通，我建议接下来可以从这几个方向继续深入。
第一个方向是Online Softmax（在线Softmax）。这主要是解决input非常长、无法一次全量装载的问题。它的核心是分块滚动更新：每处理一块，就用新的局部max更新全局max，同时对之前分子分母值做缩放校正。这样就不需要“先求全局max再第二次遍历”的多趟结构，访存效率有质的提升。
第二个方向是混合精度训练/推理下的自动缩放。有些模型会在Softmax前插入一个scale层，目的就是提前把大数值压下来。你可以把Softmax算子和前面的缩放逻辑融合，做成一个算子，减少一次数据搬运，同时让数值更稳定。
第三个方向是通用化：不只是Softmax，很多需要exp、log、求和归一化的算子都有同样的数值问题。比如log_softmax、sigmoid交叉熵、注意力分数归一化等。这些算子的稳定方案几乎同源，一旦你吃透了log-sum-exp的思路，迁移起来非常快。
8. 模块化复用：让数值稳定方案沉淀成库
如果你在公司或团队里负责算子基础库，我强烈建议不要把数值稳定逻辑写到每个算子内部，而是抽成公共模块。CANN场景里就可以封装一个“SafeExp”函数，它接收输入向量、最大值、精度模式，返回已经经过防溢出处理的exp结果。这样任何需要exp的算子都能一键接入，不用重复造轮子。
这个模块的接口设计我建议包含三类参数：
- 输入数据指针。
- 最大值处理方法：自动计算、外部传入、保守预设。
- 精度模式：fp16快速、fp32精确、混合累加。
有了这个封装，新算子开发时数值稳定部分直接复用，剩下的精力可以专注在业务逻辑和性能优化上。实测新的算子开发时间可以从两三天缩短到小半天，而且正确性更有保障。
从更宏观的角度看，数值稳定不应该是一个“出了bug再修”的应急事项，而是算子开发初期的必备设计维度。建立一套统一的防护层，才能保证边缘case不出问题，这一点在高性能和规模化场景下的价值会越来越高。
