训练营简介 2025年昇腾CANN训练营第二季，基于CANN开源开放全场景，推出0基础入门系列、码力全开特辑、开发者案例等专题课程，助力不同阶段开发者快速提升算子开发技能。获得Ascend C算子中级认证，即可领取精美证书，完成社区任务更有机会赢取华为手机，平板、开发板等大奖。

报名链接：https://www.hiascend.com/developer/activities/cann20252#cann-camp-2502-intro

摘要：在深度学习面试中，“手写 Softmax”是道送分题；但在 NPU 算子开发中，它是一道送命题。FP16 的数值范围极其有限，稍有不慎就会导致 $e^x$ 溢出变成 INF，最终输出一片 NaN。本文将深入 Ascend C 的指令集细节，手把手带你实现一个**数值稳定（Safe）且性能极致（Vector-Only）**的 Softmax 算子。

前言：脆弱的 FP16 与狂暴的指数

在 Ascend 310/910 上，我们主力使用 half (FP16) 进行计算以获得最高算力。但 FP16 的最大表示范围只有 65504。

Softmax 的核心是指数运算 $e^x$。

• $e^{11} \approx 59874$ (安全)

• $e^{12} \approx 162754$ (溢出！)

这意味着，只要你的输入 Tensor 中有一个数值超过 12，你的算子就会立即“爆炸”，产生 INF。在动辄数百层的大模型中，中间激活值超过 12 简直是家常便饭。因此，教科书版的 Exp(x) / Sum(Exp(x)) 在工程上是绝对不可用的。

一、 核心原理：给数据“降温”

解决溢出的标准解法是利用 Softmax 的平移不变性：

$$\text{Softmax}(x_i) = \frac{e^{x_i}}{\sum e^{x_j}} = \frac{e^{x_i - M}}{\sum e^{x_j - M}}$$

通常取 $M = \max(x)$。 经过这一步变换，所有的输入 $x_i - M$ 都变成了 $\le 0$ 的数。 $e^{\le 0}$ 的范围恒定在 $(0, 1]$ 之间。 结论：无论输入多大，计算过程永远不会上溢！

二、 Ascend C 避坑指南：标量与向量的战争

知道了原理，很多新手（包括原来的我）会写出类似这样的代码：

// 典型的“低效”写法
ReduceMax(maxTensor, inputTensor, ...); 
half maxVal = maxTensor.GetValue(0); // 坑点：从 Vector 搬运数据到 Scalar
Duplicate(broadcastTensor, maxVal, ...); // 坑点：从 Scalar 搬运回 Vector
Sub(output, input, broadcastTensor, ...);

为什么这样慢？

达芬奇架构（Da Vinci Architecture）内部有 Vector Unit（负责并行计算）和 Scalar Unit（负责控制流）。

• GetValue(0) 会强制 CPU/Scalar 单元等待 Vector 单元计算完成，并将数据通过总线读回来。这不仅打断了流水线，还引入了巨大的通信开销。

 王者写法：全向量化 (Vector-Only)

我们要尽量让数据停留在 Vector 单元内部，不要“落地”。 Ascend C 提供了 Brcb (Broadcast from Block) 指令，它可以直接将 Vector 里的某一块数据广播到另一个 Vector，全程不经过 Scalar 单元。

三、 实战：打造工业级 Softmax Kernel

下面是经过优化的 Kernel 代码片段，包含了精度保护和性能优化技巧。

3.1 精度保护：FP16 -> FP32 -> FP16

除了溢出，FP16 还有精度问题。在做 ReduceSum 累加时，如果数据量很大，FP16 的精度（有效位太少）会导致严重的“大数吃小数”现象。 工程铁律：凡是涉及 ReduceSum 和 Exp 的操作，尽量强转为 float (FP32) 进行。

3.2 完整流水线代码

template <typename T>
__aicore__ inline void ComputeSoftmax(LocalTensor<T>& input, LocalTensor<T>& output, int32_t len) {
    // 假设 input 已经在 UB 中
    
    // Step 1: 找最大值 (ReduceMax)
    // 结果 maxLocal 虽然是 Tensor，但只有第0个元素有效
    ReduceMax(maxLocal, input, workBuf, len);

    // Step 2: 广播最大值 (Brcb - 性能关键点！)
    // 直接在 Vector 内部，将 maxLocal[0] 广播填满 maxBroadcasted
    Brcb(maxBroadcasted, maxLocal, 1, len);

    // Step 3: 减去最大值 (Sub)
    // x_safe = x - max
    Sub(tmpLocal, input, maxBroadcasted, len);

    // Step 4: 求指数 (Exp)
    // 建议：此处若追求精度，可先 Cast 到 FP32
    Exp(tmpLocal, tmpLocal, len);

    // Step 5: 求和 (ReduceSum)
    // sumVal 同样只有第0个元素有效
    ReduceSum(sumLocal, tmpLocal, workBuf, len);

    // Step 6: 广播和 (Brcb)
    // 准备做除法，先将分母 sum 广播
    Brcb(sumBroadcasted, sumLocal, 1, len);
    
    // Step 7: 归一化 (Div)
    // out = exp(x-max) / sum
    Div(output, tmpLocal, sumBroadcasted, len);
}

四、 进阶思考：FlashAttention 里的 Online Softmax

上述算法有一个前提：一行数据能完整塞进 UB (Unified Buffer)。 如果模型上下文长度（Context Length）达到 32K 甚至 128K，一行数据比 UB 还大，没法一次性做完 ReduceMax，怎么办？

这就涉及到了 FlashAttention 的核心——Online Softmax 技巧。我们需要把长数据切块（Tiling），在处理每个小块时，动态更新全局的 Global Max 和 Global Sum。

$$m_{new} = \max(m_{old}, m_{block})$$$$d_{new} = d_{old} \times e^{m_{old} - m_{new}} + d_{block} \times e^{m_{block} - m_{new}}$$

这是大模型推理算子开发中最高阶的技法之一，Ascend C 提供了极其灵活的 Tensor 操作能力，让我们能够手动控制这种细粒度的数值更新。

五、 总结

写好一个 Softmax，不仅考验你对数学公式的理解，更考验你对硬件特性的掌握。

1. 安全第一：时刻警惕 FP16 的溢出，必须使用 $x - \max(x)$ 策略。

2. 拒绝通信：能用 Brcb 解决的广播，绝不用 GetValue。

3. 精度为王：关键累加路径上，毫不吝啬地使用 FP32。

当你开始关注 0.001 的精度误差和 1us 的指令流水开销时，恭喜你，你已经跨入了算子优化的深水区。

本文基于昇腾 CANN 8.0 编写，文中涉及的指令周期可能随硬件版本（910A/910B）有所不同。