wenyi-li Triton Ascend 性能优化指南
优化策略 Checklist
- Grid 1D 化:
grid=(CORE_NUM,)+ 核内交错循环for block_id in range(pid, total, CORE_NUM) - Grid 维度选择:
- 对于计算密集型算子（矩阵乘、卷积、大块 reduce等），考虑 2D/3D grid，利用硬件调度优势
- 对于大量小规模计算（element-wise、pointwise等），考虑 1D grid + 核内循环，减少启动开销
- 核内循环: 无需 for 的场景添加额外循环，编译器自动多级流水
- 尝试不同 BLOCK_SIZE: 从较大 tile 开始，
ub overflow则缩小；在核内循环中平衡并行度和资源占用：- 尝试小切分策略，使读写能够并行进行
- 尝试大切分策略，提升 UB 使用率
- 列多组参数配置，添加 @triton.autotune
- 算子拆分: 复杂融合算子可拆为多 kernel 顺序执行，有时性能更优
- Autotune: 列多组 tile 参数配置（不含 num_warps/num_stages）
- Reduction 用标量累加: 每个核心标量累加 + 单次 atomic 写入
- 内存对齐: matmul 的 K 维度按 512B 对齐提升带宽
- 避免 host 端 permute: 非最后维 reduce 在 kernel 内用多维索引处理
- 隐式广播: 用
a[:, None] * b替代tl.broadcast_to，减少临时 tensor - load 时直接 mask:
tl.load(ptr, mask=m, other=0.0)优于先加载再tl.where - 减少冗余精度转换: 避免反复在 fp16/fp32 转换 即
.to(float16)和.to(float32)，一次转换多次复用 - 核心数配置: grid 数设为核心数（VEC/CUBE），过大时启动开销反增
- 256B 对齐: 数据搬运以 256B 为单位，对齐可提升带宽
Reduction 优化
每个核心先局部标量累加，最后一次原子写入：
core_sum = 0.0
for block_start in range(pid, total_blocks, CORE_NUM):
data = tl.load(...)
core_sum += tl.sum(data, axis=0)
tl.atomic_add(output_ptr, core_sum)
数值稳定性
防溢出
max_val = tl.max(scores, axis=0)
scores = scores - max_val
p = tl.math.exp2(scores)
防负值开方
- 任何 sqrt 前确保非负:
max(input, 0.)或max(input, eps)
精度提升
- matmul 使用 fp32 累加器：
acc = tl.zeros([M, N], dtype=tl.float32) - 最后再转回目标精度：
result = acc.to(tl.float16)