《Ascend C 进阶实战：高性能 Softmax 算子设计与数值稳定性优化》-编程实验室

《Ascend C 进阶实战：高性能 Softmax 算子设计与数值稳定性优化

1. 引言：Softmax 的挑战

Softmax 是分类任务中的核心算子，定义为：

Softmax(xi)=∑jexjexi

看似简单，但在 NPU 上高效实现却面临三大挑战：

数值溢出：当 xi 较大时，exi 会溢出为 inf。
归约操作（Reduce）：求和需跨整个向量，难以并行。
两次遍历：需先求 max，再求 exp 和 sum，最后归一化。

本文将基于 Ascend C，实现一个数值稳定、高吞吐的 Softmax 算子，并深入探讨其在昇腾 NPU 上的优化策略。

2. 数值稳定性：减去最大值

标准做法：令 m=max(x)，则

Softmax(xi)=∑jexj−mexi−m

这样可保证指数项 ≤ 0，避免溢出。

因此，Softmax 需分三步：

ReduceMax：求全局最大值 m
Exp & Sum：计算 exi−m 并累加
Divide：每个元素除以总和

3. Ascend C 实现策略

由于 ReduceMax 是全局操作，无法单个 Block 完成。我们采用两阶段归约：

Stage 1：每个 Block 计算局部 Max 和局部 Sum
Stage 2：Host 或额外 Kernel 合并局部结果（本文简化：假设单 Block 处理整个向量）

注：生产环境应使用多 Block + AllReduce，但为聚焦 Ascend C，本文假设输入长度 ≤ 2MB（可放入 UB）。

4. Kernel 代码实现

4.1 头文件与常量

cpp

编辑

#include "kernel_api.h" using namespace AscendC; constexpr int32_t BLOCK_SIZE = 1024; // 每次处理 1024 个元素

4.2 SoftmaxKernel 类

cpp

编辑

class SoftmaxKernel { public: __aicore__ inline void Init(GM_ADDR input, GM_ADDR output, uint32_t len) { this->input_gm = input; this->output_gm = output; this->len = len; // 分配 UB：输入、输出、临时 buffer DataShape shape{BLOCK_SIZE}; input_ub.Init(shape, FORMAT_ND, ACL_FLOAT, UB); output_ub.Init(shape, FORMAT_ND, ACL_FLOAT, UB); temp_ub.Init(shape, FORMAT_ND, ACL_FLOAT, UB); // 分配 SB：存放 max_val 和 sum_val max_val_sb.Init(DataShape{1}, FORMAT_ND, ACL_FLOAT, SB); sum_val_sb.Init(DataShape{1}, FORMAT_ND, ACL_FLOAT, SB); } __aicore__ inline void Process() { // Step 1: Find global max FindMax(); // Step 2: Compute exp(x - max) and sum ComputeExpAndSum(); // Step 3: Normalize Normalize(); } private: __aicore__ inline void FindMax() { float max_val = -FLT_MAX; int32_t loop = (len + BLOCK_SIZE - 1) / BLOCK_SIZE; for (int32_t i = 0; i < loop; ++i) { uint32_t offset = i * BLOCK_SIZE; uint32_t size = min(BLOCK_SIZE, len - offset); DataCopy(input_ub, input_gm[offset], size); // 在 UB 中找局部 max float local_max = -FLT_MAX; for (uint32_t j = 0; j < size; ++j) { local_max = fmax(local_max, TmpToFloat(input_ub[j])); } max_val = fmax(max_val, local_max); } // 将 max_val 存入 SB Cast(max_val_sb, max_val); } __aicore__ inline void ComputeExpAndSum() { float sum = 0.0f; float max_val = TmpToFloat(max_val_sb[0]); int32_t loop = (len + BLOCK_SIZE - 1) / BLOCK_SIZE; for (int32_t i = 0; i < loop; ++i) { uint32_t offset = i * BLOCK_SIZE; uint32_t size = min(BLOCK_SIZE, len - offset); DataCopy(input_ub, input_gm[offset], size); // 计算 exp(x - max) Sub(temp_ub, input_ub, max_val); // temp = x - max Exp(output_ub, temp_ub); // output = exp(temp) // 累加 sum for (uint32_t j = 0; j < size; ++j) { sum += TmpToFloat(output_ub[j]); } // 暂存 exp 结果到 GM（避免 UB 覆盖） DataCopy(output_gm[offset], output_ub, size); } Cast(sum_val_sb, sum); } __aicore__ inline void Normalize() { float sum_val = TmpToFloat(sum_val_sb[0]); int32_t loop = (len + BLOCK_SIZE - 1) / BLOCK_SIZE; for (int32_t i = 0; i < loop; ++i) { uint32_t offset = i * BLOCK_SIZE; uint32_t size = min(BLOCK_SIZE, len - offset); // 从 GM 读回 exp 结果 DataCopy(output_ub, output_gm[offset], size); // 除以 sum float inv_sum = 1.0f / sum_val; Muls(output_ub, output_ub, inv_sum); // 写回最终结果 DataCopy(output_gm[offset], output_ub, size); } } // 成员变量 GM_ADDR input_gm, output_gm; Tensor<UB> input_ub, output_ub, temp_ub; Tensor<SB> max_val_sb, sum_val_sb; uint32_t len; }; extern "C" __global__ void Softmax(GM_ADDR input, GM_ADDR output, uint32_t len) { SoftmaxKernel op; op.Init(input, output, len); op.Process(); }

关键点：
使用TmpToFloat()从 Tensor 读取标量
Exp,Sub,Muls为 Ascend C 内置向量化函数
中间结果暂存 GM，避免 UB 不足

5. 优化方向

5.1 避免 GM 中转（高级技巧）

若输入长度 ≤ UB 容量（如 512KB），可一次性载入，避免多次 GM 访问：

cpp

编辑

// 一次性拷贝全部输入到 UB（需确保 len * 4 <= UB_SIZE） DataCopy(full_input_ub, input_gm, len);

5.2 使用 Vector Unit 的 Reduce 指令

Ascend C 提供ReduceMax,ReduceSum等高效归约函数，比手动循环快 3~5 倍：

cpp

编辑

ReduceMax(max_ub, input_ub, REDUCE_LAST_AXIS);

5.3 多 Block 支持（略，需 Host 同步）

6. 测试与验证

python

编辑

import torch import numpy as np x = np.random.rand(1024).astype(np.float32) * 100 # 制造大值 y_ascend = run_softmax_on_ascend(x) y_torch = torch.softmax(torch.tensor(x), dim=-1).numpy() assert np.allclose(y_ascend, y_torch, rtol=1e-4) print("✅ Softmax numerical stable!")