​作者​：昇腾实战派
​知识地图​：https://blog.csdn.net/Lumos_Lovegood/article/details/161601003

背景概述

在深度学习模型训练中，MSELoss 作为常用的损失函数，其计算性能直接影响整体训练效率。本文针对大规模向量（如 33.5M 元素）下的 MSELoss 计算，从基线实现中存在的性能瓶颈出发，逐步优化，最终在昇腾硬件上实现接近甚至超越 PyTorch 原生实现的吞吐性能。文章记录了优化思路、关键改进及性能数据，旨在为类似场景下的算子优化提供参考。

基线问题

原始实现（mse_loss-base.py）采用变长 Grid 策略：

grid = lambda meta: (triton.cdiv(n_elements, meta['BLOCK_SIZE']), )

每个 program 仅处理一个 block，并直接通过 tl.atomic_add 将结果写入输出。当 N=33.5M 时，共启动 32768 个 program，每个 program 均竞争同一个 atomic_add 地址，导致严重的写冲突和性能下降。

优化历程

1. Grid-Strided Loop（核心优化）

维度 基线 优化后
Grid 策略 变长 ceil(N/BLOCK_SIZE) 固定 num_vector_cores
Atomic 竞争 32768 次竞争 56 次（每个 program 循环累加后一次写入）
耗时占比 向量等待 81.5% 向量等待 38~66%

本质优化：将数千次 kernel 启动和数万次 atomic_add 合并为 1 次 kernel 启动 + 56 次写入，大幅降低竞争开销。

2. BLOCK_SIZE 自适应

N 基线 优化后
≤16K 256 1024
≤128K 512 2048
≤4M 512 4096
>4M 1024 8192

依据 UB 192KB 容量上限，最大化每次 DMA 搬运粒度，摊薄固定开销，提升整体吞吐。

3. tl.store 替代 atomic_add

将 tl.atomic_add 改为 tl.store(output + pid, ...) + host 侧 sum()。L2 write hit 从 0% 提升到 100%，但由于输出数据仅占总数据量的 0.002%，对总吞吐影响可忽略。

Profiler 指标解读

指标 含义 理想值
mte2_ratio L1 ↔ UB 搬运占时间比 > 80%
mte3_ratio GM ↔ UB 搬运占时间比 < 1%（数据在 L1 中复用）
vec_ratio 向量计算占时间比 < 30%（MSELoss 是 memory-bound）
vec_wait 向量单元等待比例 < 40%（越高越空闲）
L2 write hit L2 写命中率 越高越好

优化瓶颈

当前 Atlas 800I A2 上最大瓶颈是 MTE2（DMA 管道）占空比：

可尝试方向：

  1. 增大 BLOCK_SIZE 至 16384/32768（受 UB 限制，需确认）
  2. 双缓冲/软件流水线（Triton DSL 中需编译器支持）
  3. 使用向量 load 指令替代 DMA（需 Triton Ascend 后端支持）
# 性能优化
Logo
人工智能6S服务平台

作为“人工智能6S店”的官方数字引擎，为AI开发者与企业提供一个覆盖软硬件全栈、一站式门户。

更多推荐

基于 Triton-Ascend 的 MSELoss 算子设计

在深度学习模型训练中，均方误差损失（MSELoss，又称 L2 Loss）是回归任务中最常用的损失函数之一。随着模型规模的不断扩大，对算子的计算效率和精度要求也越来越高。Triton 作为一种高效的 GPU 编程语言，能够帮助开发者编写高性能的自定义算子。本文基于 Triton-Ascend 框架，设计并实现了一个支持多种归约模式、具备动态优化策略的 MSELoss 算子，旨在解决现有实现中精度不

avatar 人工智能6S服务平台
cover

【上新预告】CANN社区任务8月3日上新！体验昇腾950系列产品算力，冲刺万元大奖！

avatar 人工智能6S服务平台

昇腾与英伟达硬件架构对比

上图展示有点问题，实际一个Node的Level1里是一个交换芯片的7个虚拟平面，不是7个交换芯片。UB交换芯片Level1使用量0.392*2/2*8*2=6.272 Tb/s, Level2 交换芯片使用量0.056/2*8*48=10.752Tb/s，芯片性能对标博通TH3.B300出1.8TB双向带宽，Nvswitch单芯片容量28.8Tb，整机57.6Tb（性能对标博通TH5)，交换芯片容

avatar 人工智能6S服务平台