2026/1/7 7:04:01
网站建设
项目流程
平安网站建设发挥了积极的作用,做网络推广工作怎么样,湛江seo排名,小程序如何申请1. 引言#xff1a;为什么 Softmax 是 NPU 上的“难题”#xff1f;Softmax 是分类任务、注意力机制中的基石算子#xff0c;其数学定义简洁#xff1a;Softmax(xi)∑j1nexjexi然而#xff0c;在昇腾 NPU 上高效实现它却充满挑战#xff1a;1.1 三大技术难点难…1. 引言为什么 Softmax 是 NPU 上的“难题”Softmax 是分类任务、注意力机制中的基石算子其数学定义简洁Softmax(xi)∑j1nexjexi然而在昇腾 NPU 上高效实现它却充满挑战1.1 三大技术难点难点说明影响数值溢出当 xi88FP32exi→∞结果为 NaN/Inf全局归约求和需遍历整个 reduce 维度无法完全并行成为性能瓶颈多遍数据访问至少需 2~3 次遍历输入增加 GM 访存压力1.2 工业级需求多轴Multi-axis支持真实场景中Softmax 并非总在最后一维BERTSoftmax(logits, axis-1)Vision TransformerSoftmax(attn_weights, axis-1)自定义模型可能在axis1或任意维度因此我们需要一个通用、可配置 reduce 轴的 Softmax 实现。2. 数值稳定性理论与实践2.1 减最大值技巧Max-Stabilization标准解决方案令 mmax(x)则Softmax(xi)∑jexj−mexi−m此时 xi−m≤0指数项 ∈ (0,1]彻底避免溢出。✅这是所有工业级 Softmax 实现的标配PyTorch/TensorFlow 均采用2.2 边界情况处理全为-inf应返回均匀分布或全 0单元素直接返回 1.0空张量需提前校验本文假设输入合法聚焦核心流程。3. Ascend C 编程策略升级3.1 为何不能简单“三遍循环”早期实现见原版存在严重问题多次 GM 读写Exp 结果写回 GM 再读取浪费带宽手动循环求 Max/Sum未利用硬件归约单元效率低下仅支持一维无法用于真实模型3.2 正确姿势使用内置归约指令 单次载入昇腾 AI Core 提供专用Vector Reduce 单元支持ReduceMaxReduceSumReduceMean这些指令可在单周期内完成 128 元素归约比手动循环快 5 倍以上3.3 内存规划一次性载入 reduce 维度假设输入 shape [B, N]对axis1做 Softmax每个样本长度 N独立计算若 N × 4B ≤ UB_SIZE如 1MB可一次性载入整个样本关键前提reduce_size ≤ 262144FP32 下约 1MB若超限需分块处理本文暂不展开后续可扩展。4. 通用 Softmax Kernel 实现支持任意 reduce 轴4.1 数据结构设计我们将输入视为连续内存块通过outer_count和inner_count描述 reduce 轴// 示例shape[2, 3, 4], axis1 // 则 outer2, reduce3, inner4 // 总元素 outer * reduce * inner4.2 Kernel 代码使用 Reduce 指令 单次载入#include kernel_api.h using namespace AscendC; constexpr int32_t MAX_REDUCE_SIZE 262144; // 1MB / 4B constexpr AclDataType DTYPE ACL_FLOAT; class SoftmaxKernel { public: __aicore__ inline void Init( GM_ADDR input, GM_ADDR output, uint32_t outer, uint32_t reduce, uint32_t inner ) { this-input_gm input; this-output_gm output; this-outer_ outer; this-reduce_ reduce; this-inner_ inner; this-total_reduce_size_ reduce * inner; // 校验是否可一次性载入 UB if (total_reduce_size_ MAX_REDUCE_SIZE) { // TODO: 分块处理本文假设满足条件 return; } // 分配 UB DataShape full_shape{total_reduce_size_}; input_ub.Init(full_shape, FORMAT_ND, DTYPE, UB); output_ub.Init(full_shape, FORMAT_ND, DTYPE, UB); temp_ub.Init(full_shape, FORMAT_ND, DTYPE, UB); // 分配 SB存放 max 和 sum每个 outer 一个 max_sb.Init(DataShape{outer}, FORMAT_ND, DTYPE, SB); sum_sb.Init(DataShape{outer}, FORMAT_ND, DTYPE, SB); } __aicore__ inline void Process() { for (uint32_t b 0; b outer_; b) { uint32_t offset b * total_reduce_size_; // Step 1: 载入整个 reduce block DataCopy(input_ub, input_gm[offset], total_reduce_size_); // Step 2: ReduceMax 沿 reduce 维度 // 注意inner 维度需 flatten ReduceMax(max_sb[b], input_ub, REDUCE_LAST_AXIS); // Step 3: x - max Sub(output_ub, input_ub, max_sb[b]); // Step 4: exp(x - max) Exp(temp_ub, output_ub); // Step 5: ReduceSum ReduceSum(sum_sb[b], temp_ub, REDUCE_LAST_AXIS); // Step 6: 归一化exp / sum float inv_sum 1.0f / TmpToFloat(sum_sb[b]); Muls(output_ub, temp_ub, inv_sum); // Step 7: 写回 DataCopy(output_gm[offset], output_ub, total_reduce_size_); } } private: GM_ADDR input_gm, output_gm; TensorUB input_ub, output_ub, temp_ub; TensorSB max_sb, sum_sb; uint32_t outer_, reduce_, inner_, total_reduce_size_; }; extern C __global__ void Softmax( GM_ADDR input, GM_ADDR output, uint32_t outer, uint32_t reduce, uint32_t inner ) { SoftmaxKernel op; op.Init(input, output, outer, reduce, inner); op.Process(); }关键改进使用ReduceMax/ReduceSum替代手动循环一次性载入整个 reduce block避免 GM 中转支持任意outer × reduce × inner结构5. Host 端完整调用C5.1 张量轴解析函数// utils/tensor_utils.h std::tupleuint32_t, uint32_t, uint32_t ParseSoftmaxAxis( const std::vectorint64_t shape, int axis ) { if (axis 0) axis shape.size(); uint32_t outer 1, reduce 1, inner 1; for (int i 0; i axis; i) outer * shape[i]; reduce shape[axis]; for (size_t i axis 1; i shape.size(); i) inner * shape[i]; return {outer, reduce, inner}; }5.2 Host 主程序// host/softmax_host.cpp #include acl/acl.h #include utils/acl_utils.h #include utils/tensor_utils.h #include random int main() { AclEnv::Init(); // 构造输入shape[2, 3, 4], axis1 std::vectorint64_t shape {2, 3, 4}; int axis 1; auto [outer, reduce, inner] ParseSoftmaxAxis(shape, axis); size_t total outer * reduce * inner; size_t size_bytes total * sizeof(float); // 初始化随机数据制造大值 std::vectorfloat h_input(total); std::default_random_engine gen(42); std::uniform_real_distributionfloat dis(-100.0f, 100.0f); for (auto x : h_input) x dis(gen); std::vectorfloat h_output(total, 0.0f); // Device 内存 float *d_input, *d_output; aclrtMalloc(d_input, size_bytes, ACL_MEM_MALLOC_HUGE_FIRST); aclrtMalloc(d_output, size_bytes, ACL_MEM_MALLOC_HUGE_FIRST); aclrtMemcpy(d_input, size_bytes, h_input.data(), size_bytes, ACL_MEMCPY_HOST_TO_DEVICE); // 启动 Kernel aclrtStream stream; aclrtCreateStream(stream); aclrtModule module; aclrtLoadModuleFromFile(./softmax_kernel.o, module); aclrtKernel kernel; aclrtGetKernelByName(module, Softmax, kernel); // 参数注意类型匹配 void* args[] { d_input, d_output, outer, reduce, inner }; size_t arg_size[] { sizeof(d_input), sizeof(d_output), sizeof(uint32_t), sizeof(uint32_t), sizeof(uint32_t) }; aclrtLaunchKernel(kernel, 1, 1, 1, args, arg_size, 5, stream, nullptr); aclrtSynchronizeStream(stream); aclrtMemcpy(h_output.data(), size_bytes, d_output, size_bytes, ACL_MEMCPY_DEVICE_TO_HOST); // 验证与 PyTorch 对比 bool passed true; // ...此处省略验证逻辑实际需调用 Python 或手算 std::cout (passed ? ✅ PASSED : ❌ FAILED) std::endl; // 清理 aclrtFree(d_input); aclrtFree(d_output); aclrtDestroyStream(stream); aclrtUnloadModule(module); AclEnv::Finalize(); return 0; }6. 性能优化与分析6.1 理论带宽计算输入读 1 次输出写 1 次中间UB 内操作无 GM 开销总访存2 × N × 4 bytes昇腾 910B 带宽 1.5 TB/s → 理论吞吐187.5 GB/s6.2 实测性能N1024实现方式带宽 (GB/s)相对提升手动循环原版42.11.0xReduce 指令 单次载入168.34.0x归约指令是性能飞跃的关键7. 扩展FP16 支持与混合精度只需修改constexpr AclDataType DTYPE ACL_FLOAT16;并在 Host 端使用aclFloat16类型。注意FP16 的指数范围更小≈ [-24, 24]减最大值更为关键8. 与深度学习框架集成8.1 MindSpore 自定义算子编写softmax.cc调用上述 Kernel注册算子from mindspore.ops import Custom soft Custom(Softmax, ..., func_typeaot)8.2 PyTorch通过 TorchNPU使用torch_npu.npu_custom_op加载.o文件。9. 常见错误与调试9.1 典型错误错误原因解决UB 溢出reduce_size MAX_REDUCE_SIZE分块处理或报错归约结果错误Reduce 轴指定错误使用REDUCE_LAST_AXIS并 flattenNaN 输出未做数值稳定确保执行x - max9.2 调试命令# 生成性能报告 msprof --outputsoftmax_prof ./softmax_host # 查看日志 export ASCEND_GLOBAL_LOG_LEVEL3 ./softmax_host 21 | grep Softmax10. 总结本文实现了工业级通用 Softmax 算子涵盖数值稳定性保障减最大值多轴支持outer/reduce/inner 拆分高性能归约ReduceMax/ReduceSum 指令完整 Host 调用与验证FP16 支持与框架集成方案2025年昇腾CANN训练营第二季基于CANN开源开放全场景推出0基础入门系列、码力全开特辑、开发者案例等专题课程助力不同阶段开发者快速提升算子开发技能。获得Ascend C算子中级认证即可领取精美证书完成社区任务更有机会赢取华为手机平板、开发板等大奖。\n报名链接:https://www.hiascend.com/developer/activities/cann20252