企业官网建设流程全解析

从一行Softmax代码开始：前向算子里的exp(x)，在硬件上算出来的结果可能直接是Inf或NaN。不少刚接触CANN算子开发的读者，第一次在日志里看到这种溢出时，总是习惯性先去查硬件指令或算子框架的配置，其实问题通常出在Softmax本身的数值稳定性设计上。这篇文章专门拆解在CANN ops-math体系里，Softmax的数值稳定实现、溢出防护方案，以及log-sum-exp等数学技巧究竟是怎么落地到算子代码里的。适合正在做昇腾AI Core算子迁移、算子性能优化，或者准备把PyTorch里的自定义Softmax层改写到CANN算子接口的开发者参考，也适合想系统理解算子层数值稳定设计思路的读者。

1. Softmax数值不稳定的根源：从一段真实Debug经历说起

1.1 一个会算错的标准实现

先把最基础的Softmax公式摆出来。给定输入向量x = [x1, x2, ..., xn]，Softmax定义为：

softmax(x_i) = exp(x_i) / sum_j(exp(x_j))

看起来简单，但直接按这个公式写算子，在真实场景里会踩大坑。我之前调试某个视觉模型的自定义融合算子时，前向输出里出现了整片的NaN。第一反应是怀疑CANN算子的AI Core向量指令配置有问题，排查了一圈才发现纯粹的数值溢出问题。

问题出在exp(x)这一步。假设输入里有一个维度值是1000，那么exp(1000)的结果约等于1.97e434，而AI Core上的浮点数值范围上限（以常见的FP32为例）大约在3.4e38，这个结果直接溢出变成Inf。再继续计算分母sum(exp(x_j))，必然也是Inf。最后做除法Inf / Inf，得到的就是NaN。这个链条一旦启动，后面整个算子输出都受影响。

1.2 溢出的经典触发场景

实际业务里遇到这种问题，很少是输入恰好有1000这种夸张值，更多是以下几种情况叠加出来的。

一个是数值范围未被约束的中间层特征。尤其在Transformer类的模型里，经过多层残差连接和LayerNorm之前，特征向量的某些维度值可能涨到几百甚至上千。另一个是混合精度训练。模型在FP16下跑，FP16的最大表示范围大约65504，exp(11)就逼近溢出边界了，任何超过这个量级的输入都会直接产生Inf。当FP16的溢出结果被后续计算传播，转成FP32时依然保持Inf，最终导致Loss直接变成NaN。

还有一个容易被忽略的场景是输入数据本身存在异常值。比如在某些分布式训练集群上，某个节点的Loss不收敛，反向传播产生的梯度数值异常大，特征分布整体偏移，Softmax输入也跟着出现极端值。

1.3 为什么不直接给输入加个Clip

有人会问，既然怕溢出，能不能直接用clip(x, min, max)把输入限制在安全范围内？这种思路在工程上偶尔有人用，但本质上是有问题的。

Softmax的输出是概率分布，对输入做Clip相当于人为改了模型的行为。如果Clip的上限设得低，模型学到的特征分布会被截断，训练出的模型精度受损。如果设得高，溢出还是会发生，等于没解决。深度学习的训练中，我们很少能预先知道特征值的合理范围，尤其是迁移学习和fine-tune阶段，特征分布会动态变化。与其猜一个阈值，不如从算法层面彻底规避溢出风险。

正确的思路是通过数学等价变换，把计算过程中的数值范围拉回安全区间。这就是log-sum-exp技巧以及减最大值方案做的事情，稍后详细展开。

2. log-sum-exp：数学等价与工程化改造要点

2.1 先理解数学上为什么等价

log-sum-exp技巧的核心操作是：计算Softmax前，先从输入向量的所有分量中减去它们的最大值m。

m = max(x1, x2, ..., xn) softmax(x_i) = exp(x_i - m) / sum_j(exp(x_j - m))

数学上严格等价，因为分子分母同时乘了exp(-m)，这个缩放因子在除法中相互抵消。关键在于：x_i - m的结果永远小于等于0。只要输入的不是Inf，减去最大值后最多等于0，exp(0) = 1，其他情况都是exp(负数)，结果落在(0, 1]区间，彻底不再有溢出风险。

但有一个值得追问的细节：如果输入本身就是Inf，那减最大值也没用，因为Inf - Inf得到NaN。这种情况通常意味着之前的网络层已经发生异常，需要单独排查，不能指望Softmax这一层来兜底。

现实一点做个类比：这就像温度计的量程有限，如果室外温度是50摄氏度，直接测会爆表，于是先量一个基准温度，再测量与基准的差值。差值不会超出温度计量程，最后加上基准就得到真实温度。Softmax里的exp就是那支温度计，x_i - m就是差值。

2.2 log-sum-exp并不只是Softmax的专利

log-sum-exp这套数学方法不只是Softmax在用。在CANN算子的邻域里，很多地方都能看到它的影子。

比如在LogSoftmax实现中，直接算log(softmax(x_i))，如果把log和sum(exp(x_i))分开搞，中间结果的溢出问题照样存在。更稳妥的写法是：

log_softmax(x_i) = x_i - m - log(sum(exp(x_j - m)))

这个式子就是标准的log-sum-exp形式，先减最大值，再取对数，全程数值范围可控。

再比如在贝叶斯变分推断里的ELBO计算、强化学习里的soft Q-learning目标函数，甚至CTCLoss的前向计算里，都有类似的结构。算子开发人员不要只把log-sum-exp当成Softmax的附属品，它是一个底层数值工具，吃透之后能在处理大量涉及exp、log、sum组合的计算时游刃有余。

2.3 工程改造的注意点

数学上等价，工程上不能无脑替换。动手改代码前有几点得先想清楚。

第一点是最大值m的精度。如果输入向量各分量的数量级差距很大，比如一个元素是0.001，另一个是1000，那么m对所有元素减去后，小元素可能变成一个非常接近0的负数，这时候exp的结果依然是可以表示的，但浮点精度会受影响。实际中这种极端动态范围在训练初期不常见，但推理时量化模型里可能出现，值得留意。

第二点是第二遍遍历的计算效率。如果算法实现里每一层都重算m，遍历开销会翻倍。好在CANN算子场景里大都是向量化指令，ReduceMax一条指令就能完成找最大值的操作，OpenCL和CUDA里也都有现成的归约原语，并不会像CPU标量写法那么慢。

第三点是分母求和时的顺序。sum(exp(x_i - m))一般按照硬件向量的累加方式计算，如果向量长度不满足对齐要求，注意补零或者分块处理，避免溢出或精度损失。

3. CANN算子映射：从数学公式到AI Core指令

3.1 昇腾AI Core的向量化结构与指令选择

CANN算子开发面向的是昇腾AI Core，它是一个以向量计算为核心的加速单元，跟我们熟悉的CPU标量运算思路差别很大。在CPU上写Softmax，直观做法是循环遍历每个元素。但在AI Core上，我们需要把问题映射成向量或矩阵指令，用一次数据搬运加一条向量指令代替多层循环。

AI Core的指令集里，和Softmax直接相关的包括：

• ReduceMax：沿指定轴归约出最大值，用于计算log-sum-exp里的m。
• Add、Sub、Mul：向量间的逐元素运算，用于实现x_i - m。
• Exp：逐元素的指数运算，AI Core上有专门的向量指数指令。
• ReduceSum：沿指定轴归约求和，用于计算分母。
• Div：逐元素除法，用于最终归一化。

开发时第一步就是先把Softmax拆成这些原始指令的组合。不要硬套CPU上的循环逻辑，AI Core的向量化优势只有在整段向量操作里才能发挥出来。

3.2 两遍扫描策略：先求最大值，再归一化

两遍扫描是这个任务里最常见、也最稳妥的实现策略。

第一遍扫描：对输入向量执行ReduceMax，得到全局最大值m。

第二遍扫描：对每个元素执行x_i = x_i - m，紧跟着exp(x_i)，然后ReduceSum求和得到分母S，最后逐元素除以S。

两遍扫描的好处是实现简单、逻辑明确，每个中间结果的数值范围都是可控的，非常适合作为算子稳定版本的基线。缺点是数据要读两遍，对AI Core而言意味着多一次数据搬运开销。

如果对实时性要求极高，想极端压缩一遍扫描的版本，可以用EMA近似的在线式算法。但EMA引入的数值误差不易控制，还依赖于超参选择，在算子库这种追求确定性的场景下，我一般不建议。

3.3 融合场景的算子设计

实际工程里，Softmax很少单独存在。CANN算子开发时常见的做法是把它跟LayerNorm、残差连接、甚至Query乘Key的操作融合进一个算子，减少中间张量的片外内存搬运。融合之后，Softmax部分依然要用log-sum-exp技巧，但要注意融合算子里其他部分的数值范围处理。

举个例子，在融合的LayerNorm+Softmax算子中，LayerNorm输出通常是均值为0、方差为1的分布，此时Softmax输入基本不会出现极端值。但LayerNorm在FP16下计算时，方差做rsqrt之后可能产生较大的数值，这些细节要在算子日志里多做中间值打印，了解每一条指令出入数据的实际分布，而不是只对着公式做静态分析。

4. 实操过程：CANN算子中实现数值稳定Softmax

4.1 环境准备与接口基座

动手写代码之前，需要先确认CANN开发环境已经就绪。通常依赖于昇腾AI处理器驱动、CANN toolkit以及对应的算子开发框架。算子可以用Ascend C语言编写，也可以基于TBE或MindSpore自定义算子框架，甚至用PyTorch调用torch_npu的接口完成上层验证。下面这些步骤偏Ascend C风格，但思路通用。

4.2 核心实现：稳定版Softmax的算子逻辑

以简化示例说明。输入是一个一维向量x，长度N，输出y。

第一步，定义输入、输出的Tensor描述，并初始化算子上下文。包括分配Global Memory、确定数据格式和数据类型。

第二步，用ReduceMax归约得到全局最大值：

// 伪代码，示意AI Core向量指令调用 Tensor max_val = ReduceMax(x, axis = 0);

这里注意轴方向的正确性。如果是二维矩阵按行做Softmax，axis = 1；按列则axis = 0。CANN算子里的reduce接口通常支持指定axis。

第三步，逐元素减去最大值并求exp：

Tensor shifted = Sub(x, max_val); // x_i - m Tensor exp_shifted = Exp(shifted); // exp(x_i - m)

第四步，归约求和得到分母：

Tensor sum_exp = ReduceSum(exp_shifted, axis = 0);

第五步，逐元素除分母：

Tensor y = Div(exp_shifted, sum_exp);

这五步就是稳定版Softmax的全部核心逻辑。看起来极简，但要注意的是：

• ReduceMax和ReduceSum返回的标量Tensor形状必须能被广播，否则Sub和Div指令会因为形状不匹配报错。
• 如果输入是分块的，ReduceMax在分块间需要二次归约，不能直接拿单块的最大值当全局最大值。
• exp指令本身对输入范围有要求。多数硬件指令集对特别小的负数也容易产生接近0的超小值，但一般不会上溢，唯一需要防的还是输入正值过大。

4.3 数值对比实验：稳定版与Direct实现

为了验证稳定性，我做了一个简单的数值对比。构造两组极端的输入向量，一组是[0, 0, 1000]，另一组是更接近真实场景的[800, 809, 810]。

Direct实现（直接用exp(x_i)除以sum(exp(x_j))）在第一组输入里直接出现NaN，原因是exp(1000)溢出。第二组输入里，exp(810)虽然没溢出，但分母巨大，微小的相对精度损失在除法后会放大到每个输出分量上。

稳定版实现中，两组输入都正常输出概率分布。第一组结果是[0, 0, 1]，第二组结果在FP32精度下与理论值误差在1e-7量级以内。这说明减最大值策略没有引入额外精度损失，在正常输入范围内也能保持良好精度。

4.4 反向传播时的数值稳定考量

前向的Softmax用稳定版搞定，很多人就以为完事了，实际反向传播也要关心数值问题。

Softmax的反向传播公式是：dx_i = y_i * (dy_i - sum_j(dy_j * y_j))，其中y是前向输出，dy是上游梯度。这一步里没有exp操作，数值本身就比较稳定，但要注意的是y如果是直接从前向存储下来的概率值，它天然落在(0, 1)区间，不会导致反向的指数溢出。

所以核心的中间结果只需要存储前向的y，不必存x - m。

5. 溢出防护的进阶策略与细节打磨

5.1 下溢出风险同样不可忽视

极端值不仅会触发上溢，也会触发下溢。当x_i - m是一个非常小的负数，比如-800时，exp(-800)的结果在FP32下已经趋近于0，在FP16下直接就是0。这会导致Softmax输出的某些分量为0，后续的log计算（比如交叉熵Loss）产生NaN，因为log(0)无定义。

如果系统里存在这种情况，可以考虑为分母S加一个极小值epsilon，类似S + 1e-8，避免除零。或者对最终输出y做一次floor操作，保证最小值不低于某个阈值。

不过加epsilon也得看场景。Softmax的输出严格意义上要满足和为1、每个分量非负，加epsilon会轻微破坏这些特性。对模型训练来说通常可以接受，但对部署时的后处理（比如beam search打log）要小心。

5.2 量化场景的溢出防护

CANN算子经常要跑INT8量化或FP16推理。量化场景里，缩放因子（scale）和零点（zero point）的存在让输入范围被压缩到[-128, 127]（INT8）或[-65504, 65504]（FP16）。看起来范围可控，但实际上有两个隐患。

第一个隐患是错误的scale保证。当量化参数是从训练统计中推算出来的，如果推得的scale偏大，反量化后的浮点值可能超过模型设计时的预期范围，Softmax照样碰上极端值。建议在算子输入处增加动态范围校验，记录输入张量的min/max，跟训练时的统计量做比对。

第二个隐患是Softmax在量化模型里通常会把输入以浮点形式存储，或者用高精度累加器做中间计算。如果直接用低精度执行exp，误差会逐步累积。个人经验是计算x_i - m、exp、ReduceSum时统一用FP32，只在最终输出时再量化回INT8，精度收益非常明显。

5.3 单遍扫描的优化版本

两遍扫描是安全基线，但有一些推理场景对延迟极其敏感，希望尽量减小遍历开销。此时可以用one-pass版本，核心思路是维护一个动态最大值m和动态累加和S。

处理一个元素x_i时，更新逻辑是：

if x_i > m: S = S * exp(m - x_i) + 1 m = x_i else: S = S + exp(x_i - m)

这个版本理论上只需读一遍数据，但问题在于S * exp(m - x_i)中如果m和x_i差距很大，指数因子可能下溢；而且多次乘法累加会引入额外累积误差。实测下来，对于1000维以内的向量，one-pass在两遍扫描面前没有显著优势，反而代码更复杂。除非有特殊的内存带宽瓶颈，否则不值得。

5.4 多维度场景下的实现细节

现实中的Softmax输入通常是二维甚至三维。比如Transformer注意力矩阵是[batch, heads, seq_len, seq_len]，沿最后一个维度做Softmax。在CANN算子实现里，这意味着ReduceMax和ReduceSum的轴要精确指定为最后一个维度，并且每一批、每个头的归约是独立的。

相应地，内存布局也对性能有很大影响。如果头维度和序列维度在内存上不连续，ReduceMax指令的数据拾取效率会下降。开发时建议先用transpose或reshape把归约轴移到连续内存的低维方向，再执行归约操作。

6. 避坑指南与常见问题速查

6.1 常见问题速查表

6.2 调试技巧：拿住中间张量

CANN开发时很多人习惯直接看最终输出对不对，一不对就怀疑算子逻辑。但数值稳定类问题，中间张量才是关键。建议在关键节点显式输出x的min/max、exp(x - m)的min/max、S的数值，定位到具体是哪一步开始异常的。这比对着整段算子代码猜高效得多。

之前有个案例，某个融合算子跑FP16时输出偶尔抖动，查了好久，最后定位到LayerNorm里rsqrt的输出偶尔达到几千，喂给Softmax直接溢出。如果不打印中间值，这个问题极难复现和定位。

6.3 关于测试用例的边界设计

算子提交到算子库之前，测试用例有个容易被忽视的重点：一定要包含边界和极端数据，而不是只用随机正常数据。对于Softmax，至少要测以下几类输入：

• 所有元素相等的向量，比如[1, 1, 1]，输出应该均匀分布。
• 元素差距巨大的向量，比如[0, 1e-6, 1e6]。
• 元素为0和负数的向量，验证减最大值后依然正确。
• 单元素向量。此时Softmax恒为1，但ReduceMax和ReduceSum的shape处理容易出bug。
• 空维度或长度为1的维度。

我之前见过一个算子，所有正常测试都过了，但在batch为1、seq_len为1的推理配置上直接崩，就是单元素向量处理的边界问题。这类用例加上去，能省下很多线上事故排查时间。

6.4 log-sum-exp在算子库中的落地形态

最后提一句log-sum-exp在算子库里的落地形态。它不是单独提供给你调用的一个“log_sum_exp()”函数，而是内嵌在Softmax、LogSoftmax、CrossEntropyLoss等算子的实现里。作为开发者，需要掌握的是识别“哪些计算可以用这个技巧加固”，而不仅限于调用某个现成API。

比如当你写一个自定义的Top-K Softmax融合算子时，log-sum-exp照用不误。当你处理多分类任务的辅助损失函数时，也离不开这个框架。数值稳定的思维一旦内化成习惯，什么时候该减最大值、什么时候该加epsilon、什么时候该升级到FP32，都会变成直觉反应。

我个人在实际项目里的体会是：Softmax数值稳定性不是一次性能优化，而是一个贯穿前向、反向、推理、量化全链路的基本功。把它做对了，整个算子的质量下限就有了保障。遇到类似exp+sum为基本结构的数学表达式时，多想想能不能用减最大值、log-sum-exp这类等价变换保护数值精度，这比事后修NaN要划算得多。