已开启
feat(scatter_add): add NPU-native ScatterAdd kernel for Ascend 910B3 #21
feat(scatter_add): add NPU-native ScatterAdd kernel for Ascend 910B3 #21
已开启
lhp_lhp创建于 2 天前
lhp_lhp成员
2 天前

Description

为 Ascend 910B3 (dav-2201) 实现原生 ScatterAdd kernel,替换原有的 CPU 回退路径(D2H → CPU ScatterAdd → H2D)。

参照 CUDA 实现 src/array/cuda/segment_reduce.cuh:54ScatterAddKernel,采用 AscendC MemBase 范式实现。

修改功能

ScatterAddDispatchsrc/array/kernel.cc)中 Ascend 设备的调度路径从 CPU 回退改为原生 NPU kernel 调用ATEN_XPU_SWITCH_CUDA_ASCEND),与仓内 SegmentReduce/SpMM/SDDMM 等 Ascend 算子的 dispatch 范式一致。

适配算子

ScatterAdd(scatter add on first dimension)

  • 数学语义:out[idx[i], *] += feat[i, *]
  • 累加方式:add(仅支持加法累加,不支持 max/min 等其他 reduce op)
  • kernel 不清零 out,由调用方预清零(F.zeros),与 CUDA/CPU 实现一致

支持的数据类型

IdType Dtype 状态
int32 float32 ✅ 原生支持
int64 float32 ✅ 支持(int64 idx 内部转为 int32,与 segment_reduce.cc 同范式)
int32 half (float16) ❌ LOG(FATAL) 占位
int64 half (float16) ❌ LOG(FATAL) 占位
int32 bfloat16 ❌ LOG(FATAL) 占位
int64 bfloat16 ❌ LOG(FATAL) 占位
int32 double (float64) ❌ LOG(FATAL) 占位
int64 double (float64) ❌ LOG(FATAL) 占位

实现方案

  • 架构:SIMD / MemBase,参照 segment_reduce_sum_kernel.cpp
  • 累加策略SetAtomicAdd<float>() + DataCopyPad UB→GM 原子累加(与 CUDA atomicAdd 语义一致)
  • 多核切分:按 N(feat 行数)切分,blockDim = min(N, 40)
  • 同步机制:idx 预搬 UB + SetFlag/WaitFlag<MTE2_MTE3> 事件同步 + PipeBarrier 多队列汇聚 + SetAtomicNone 复位

验收结果

指标 结果
功能用例 46/46 通过(42 精度 + 4 异常 dtype),×3 reps 全确定性
精度 float32 混合容差(rtol=2⁻¹⁰, atol=2⁻¹⁶),最差 max_abs_error=1.03e-4(阈值 1e-2)
性能 原生 vs CPU 回退加速比 17x–44x
代码检视 通过

Checklist

Changes

likedislike
合并受阻
Llhp_lhp成员
2 天前 修改了pull request 的描述
Llhp_lhp成员
2 天前 将 knightgd 设为合并人
Llhp_lhp成员
2 天前 强制推送  1 个提交:c483a000-feat(scatter_add): add NPU-native ScatterAdd kernel for Ascend 910B3
Llhp_lhp成员
2 天前 强制推送  1 个提交:2d43651a-feat(scatter_add): add NPU-native ScatterAdd kernel for Ascend 910B3
Llhp_lhp成员
1 天前 修改了pull request 的描述
Llhp_lhp成员
1 天前 修改了pull request 的描述
Llhp_lhp成员
1 天前 强制推送  1 个提交:69af1c02-feat(scatter_add): add NPU-native ScatterAdd kernel for Ascend 910B3
Llhp_lhp成员
1 天前 强制推送  1 个提交:137bf8e0-feat(scatter_add): add NPU-native ScatterAdd kernel for Ascend 910B3
Llhp_lhp成员
1 天前 强制推送  1 个提交:bcd484a3-feat(scatter_add): add NPU-native ScatterAdd kernel for Ascend 910B3