前言用 ops-math 的 MatMul7B 模型推理吞吐 72 tokens/s。自己写了一个 Ascend C MatMul吞吐掉到 38 tokens/s。差了快 2 倍。不是算子写得烂是 Cube/Vector 分配策略错了。昇腾的 AI Core 分 Cube Unit 和 Vector Unit。矩阵乘走 Cube逐元素运算走 Vector。MatMul 看起来只是矩阵乘其实 Cube/Vector 的分工、L1 缓存的预取、输出的对齐每个环节都能拉开 2 倍性能差距。Ascend C 编程模型Ascend C 是昇腾的算子编程语言核心概念AI Core一个计算单元 ├─ Cube Unit矩阵乘单元 │ └─ MAC 阵列 16×16一次算 16×16×16 的矩阵乘 ├─ Vector Unit逐元素运算单元 │ └─ 128-lane SIMD一次处理 128 个元素 └─ 内存层次 ├─ HBM全局内存1.2TB/s 带宽 ├─ L1 缓存1MB~10TB/s 带宽 ├─ L0A/L0BCube 输入缓冲各 64KB └─ L0CCube 输出缓冲128KBCube vs Vector 分工操作执行单元原因矩阵乘A×BCubeMAC 阵列专门算矩阵乘效率最高逐元素运算scale、add、reluVectorSIMD 并行处理 128 个元素标量运算循环、条件判断Scalar控制逻辑MatMul 只涉及矩阵乘应该全走 Cube。但实际实现中数据搬运、地址计算、边界处理都要 Scalar 和 Vector 参与。调度不好Cube 等 Vector 数据空转 40% 时间。MatMul 的 Tiling 策略大矩阵乘法例如 4096×4096不能一次算完必须拆成小 tile。Tiling 公式C[M][N] A[M][K] × B[K][N] 拆分 M M0 × tile_m K K0 × tile_k N N0 × tile_n 每次算 C_tile[tile_m][tile_n] A_tile[tile_m][tile_k] × B_tile[tile_k][tile_n]tile 大小的选择约束 1tile_m × tile_k × dtype L0A 容量64KB约束 2tile_k × tile_n × dtype L0B 容量64KB约束 3tile_m × tile_n × dtype L0C 容量128KB约束 4tile_m 和 tile_n 必须是 16 的倍数MAC 阵列 16×16FP16 下最优选择tile_m64, tile_k64, tile_n64为什么是 64L0A 容量 64KB64×64×2 8KB 64KB ✓L0B 容量 64KB64×64×2 8KB 64KB ✓L0C 容量 128KB64×64×2 8KB 128KB ✓64 是 16 的倍数 ✓工程经验tile_m16 时MAC 阵列只用了 1/416×16 阵列只填了 16 行。吞吐腰斩。tile_m64 时MAC 阵列利用率 89%。从 16 调到 64吞吐从 38 tokens/s 涨到 71 tokens/s。MAC 阵列的填充Cube Unit 的 MAC 阵列是 16×16 的乘法器阵列一次能算C[16][16] A[16][16] × B[16][16]关键如何填满 MAC 阵列。错误做法tile_m16tile_k16tile_n16每次只算 16×16 的 C_tileMAC 阵列利用率 100%但调度开销大。4096×4096 的矩阵要拆成 256×25665536 次 tile 计算调度开销吃掉 30% 时间。正确做法tile_m64tile_k64tile_n64每次算 64×64 的 C_tile拆成 4×416 次 MAC 阵列计算。调度开销降到 1/16。更激进的优化tile_m128tile_k64tile_n128每次算 128×128 的 C_tile拆成 8×864 次 MAC 阵列计算。但 L0A/L0B/L0C 容量不够要多次搬运数据性能反而掉。实测数据4096×4096 MatMul910B 单卡tile_mtile_ktile_n耗时 (ms)MAC 利用率1616168.2100%3232325.192%6464643.889%128641284.376%tile_m64 最优。tile_m16 MAC 利用率虽然 100%但调度开销大总时间反而长。L1 缓存预取HBM 带宽 1.2TB/s延迟 200ns。L1 带宽 ~10TB/s延迟 10ns。差距 20 倍。不预取的流程1. 从 HBM 读 A_tile 到 L0A200ns 2. 从 HBM 读 B_tile 到 L0B200ns 3. Cube 算 A_tile × B_tile50ns 4. 把 C_tile 写到 HBM200nsCube 算的时候数据搬运已经完成。但下一步要等数据搬运完成才能开始算空转 30% 时间。预取的流程1. 从 HBM 读 A_tile1 到 L0A200ns 2. 从 HBM 读 B_tile1 到 L0B200ns 3. Cube 算 A_tile0 × B_tile050ns// 用上一个 tile 的数据 4. 同时从 HBM 读 A_tile2 到 L1200ns// 预取下一个 tile 5. 把 C_tile 写到 HBM200nsCube 算上一个 tile 时DMA 在搬运下一个 tile 的数据。Cube 不等数据。如何实现预取Ascend C 用CopyAPI 搬运数据。关键参数cache_modeL1_CACHE表示同时存一份到 L1。// 从 HBM 读 A_tile 到 L0A同时缓存到 L1Copy(A_L0A,A_HBMoffset,tile_m*tile_k*sizeof(half),L1_CACHE);下次用到相同的 A_tile 时直接从 L1 读不需要再走 HBM。工程经验7B 模型推理时QKV 投影的权重矩阵被 3 次复用Q、K、V。预取到 L1 后第 2、3 次访问快 15 倍。吞吐从 61 tokens/s 涨到 71 tokens/s。输出对齐优化HBM 写入有对齐要求地址必须是 32 字节对齐。不对齐写入慢 15%。错误写法// 输出地址没对齐half*C_outputC_HBMoffset;// offset 可能不是 16 的倍数正确写法// 确保输出地址 32 字节对齐uint64_taligned_offset(offset15)/16*16;// 向上对齐到 16 个 half32 字节half*C_outputC_HBMaligned_offset;或者用 Ascend C 的AlignAPI// 自动对齐输出地址autoC_alignedAlign(C_HBM,32);// 32 字节对齐性能差异对齐方式写入带宽性能影响不对齐~1.0TB/s-15%32 字节对齐~1.2TB/s基准完整代码示例200 行 Ascend C MatMul精简版核心逻辑#includekernel_operator.hconstexprintTILE_M64;constexprintTILE_K64;constexprintTILE_N64;classMatMulKernel{public:__aicore__inlinevoidProcess(GM_ADDR a,GM_ADDR b,GM_ADDR c,intM,intK,intN){// 遍历所有 tilefor(intm0;mM;mTILE_M){for(intn0;nN;nTILE_N){// 初始化 C_tile 为 0InitC(cm*Nn,TILE_M,TILE_N);for(intk0;kK;kTILE_K){// 从 HBM 读 A_tile 到 L0A预取到 L1CopyA(am*Kk,TILE_M,TILE_K);// 从 HBM 读 B_tile 到 L0B预取到 L1CopyB(bk*Nn,TILE_K,TILE_N);// Cube 算 A_tile × B_tile累加到 C_tileMatMulTile(TILE_M,TILE_K,TILE_N);}// 把 C_tile 写回 HBMWriteC(cm*Nn,TILE_M,TILE_N);}}}private:TPipe pipe;TBufTPosition::A1A_L0A;// L0A bufferTBufTPosition::B1B_L0B;// L0B bufferTBufTPosition::C1C_L0C;// L0C buffer__aicore__inlinevoidCopyA(GM_ADDR a,intm,intk){// 从 HBM 读 A_tile 到 L0A同时缓存到 L1autolenm*k*sizeof(half);Copy(A_L0A,a,len,{.cache_modeL1_CACHE});}__aicore__inlinevoidCopyB(GM_ADDR b,intk,intn){// 从 HBM 读 B_tile 到 L0B同时缓存到 L1autolenk*n*sizeof(half);Copy(B_L0B,b,len,{.cache_modeL1_CACHE});}__aicore__inlinevoidMatMulTile(intm,intk,intn){// Cube 算矩阵乘结果累加到 L0CMatMul(C_L0C,A_L0A,B_L0B,m,k,n,{.accumulatetrue});}__aicore__inlinevoidWriteC(GM_ADDR c,intm,intn){// 从 L0C 写回 HBM确保 32 字节对齐autoaligned_cAlign(c,32);autolenm*n*sizeof(half);Copy(aligned_c,C_L0C,len);}};// 算子入口externC__global__ __aicore__voidmatmul_kernel(GM_ADDR a,GM_ADDR b,GM_ADDR c,intM,intK,intN){MatMulKernel op;op.Process(a,b,c,M,K,N);}编译和运行# 编译算子npu-smiset-tmm-s0-dmatmul_kernel.o matmul_kernel.cpp# 链接成动态库ld-sharedmatmul_kernel.o-olibmatmul.so# 在 ACL 中调用aclError retaclrtLaunchKernel(matmul_kernel, grid, block, args,0, stream);性能对比实现吞吐 (tokens/s)MAC 利用率L1 命中率初版tile_m163823%0%tile_m645256%0%L1 预取6772%45%输出对齐7189%45%ops-math官方7291%48%自己写的 MatMul 性能与 ops-math 持平。差距在 L1 命中率45% vs 48%ops-math 的预取策略更激进。踩坑实录坑 1tile_m16 吞吐腰斩MAC 阵列 16×16tile_m16 只填了一行利用率 23%。改成 tile_m64利用率拉到 89%。坑 2L1 没预取Cube 等 Vector 数据不预取时Cube 算的时候数据还在搬运空转 40% 时间。加预取后数据提前到 L1Cube 不等。坑 3输出没对齐写入慢 15%HBM 写入要 32 字节对齐。用AlignAPI 自动对齐性能提 15%。https://atomgit.com/cann/ops-mathhttps://atomgit.com/cann/opbasehttps://atomgit.com/cann/cann-samples
Ascend C 算子开发：10 分钟写一个高性能 MatMul
前言用 ops-math 的 MatMul7B 模型推理吞吐 72 tokens/s。自己写了一个 Ascend C MatMul吞吐掉到 38 tokens/s。差了快 2 倍。不是算子写得烂是 Cube/Vector 分配策略错了。昇腾的 AI Core 分 Cube Unit 和 Vector Unit。矩阵乘走 Cube逐元素运算走 Vector。MatMul 看起来只是矩阵乘其实 Cube/Vector 的分工、L1 缓存的预取、输出的对齐每个环节都能拉开 2 倍性能差距。Ascend C 编程模型Ascend C 是昇腾的算子编程语言核心概念AI Core一个计算单元 ├─ Cube Unit矩阵乘单元 │ └─ MAC 阵列 16×16一次算 16×16×16 的矩阵乘 ├─ Vector Unit逐元素运算单元 │ └─ 128-lane SIMD一次处理 128 个元素 └─ 内存层次 ├─ HBM全局内存1.2TB/s 带宽 ├─ L1 缓存1MB~10TB/s 带宽 ├─ L0A/L0BCube 输入缓冲各 64KB └─ L0CCube 输出缓冲128KBCube vs Vector 分工操作执行单元原因矩阵乘A×BCubeMAC 阵列专门算矩阵乘效率最高逐元素运算scale、add、reluVectorSIMD 并行处理 128 个元素标量运算循环、条件判断Scalar控制逻辑MatMul 只涉及矩阵乘应该全走 Cube。但实际实现中数据搬运、地址计算、边界处理都要 Scalar 和 Vector 参与。调度不好Cube 等 Vector 数据空转 40% 时间。MatMul 的 Tiling 策略大矩阵乘法例如 4096×4096不能一次算完必须拆成小 tile。Tiling 公式C[M][N] A[M][K] × B[K][N] 拆分 M M0 × tile_m K K0 × tile_k N N0 × tile_n 每次算 C_tile[tile_m][tile_n] A_tile[tile_m][tile_k] × B_tile[tile_k][tile_n]tile 大小的选择约束 1tile_m × tile_k × dtype L0A 容量64KB约束 2tile_k × tile_n × dtype L0B 容量64KB约束 3tile_m × tile_n × dtype L0C 容量128KB约束 4tile_m 和 tile_n 必须是 16 的倍数MAC 阵列 16×16FP16 下最优选择tile_m64, tile_k64, tile_n64为什么是 64L0A 容量 64KB64×64×2 8KB 64KB ✓L0B 容量 64KB64×64×2 8KB 64KB ✓L0C 容量 128KB64×64×2 8KB 128KB ✓64 是 16 的倍数 ✓工程经验tile_m16 时MAC 阵列只用了 1/416×16 阵列只填了 16 行。吞吐腰斩。tile_m64 时MAC 阵列利用率 89%。从 16 调到 64吞吐从 38 tokens/s 涨到 71 tokens/s。MAC 阵列的填充Cube Unit 的 MAC 阵列是 16×16 的乘法器阵列一次能算C[16][16] A[16][16] × B[16][16]关键如何填满 MAC 阵列。错误做法tile_m16tile_k16tile_n16每次只算 16×16 的 C_tileMAC 阵列利用率 100%但调度开销大。4096×4096 的矩阵要拆成 256×25665536 次 tile 计算调度开销吃掉 30% 时间。正确做法tile_m64tile_k64tile_n64每次算 64×64 的 C_tile拆成 4×416 次 MAC 阵列计算。调度开销降到 1/16。更激进的优化tile_m128tile_k64tile_n128每次算 128×128 的 C_tile拆成 8×864 次 MAC 阵列计算。但 L0A/L0B/L0C 容量不够要多次搬运数据性能反而掉。实测数据4096×4096 MatMul910B 单卡tile_mtile_ktile_n耗时 (ms)MAC 利用率1616168.2100%3232325.192%6464643.889%128641284.376%tile_m64 最优。tile_m16 MAC 利用率虽然 100%但调度开销大总时间反而长。L1 缓存预取HBM 带宽 1.2TB/s延迟 200ns。L1 带宽 ~10TB/s延迟 10ns。差距 20 倍。不预取的流程1. 从 HBM 读 A_tile 到 L0A200ns 2. 从 HBM 读 B_tile 到 L0B200ns 3. Cube 算 A_tile × B_tile50ns 4. 把 C_tile 写到 HBM200nsCube 算的时候数据搬运已经完成。但下一步要等数据搬运完成才能开始算空转 30% 时间。预取的流程1. 从 HBM 读 A_tile1 到 L0A200ns 2. 从 HBM 读 B_tile1 到 L0B200ns 3. Cube 算 A_tile0 × B_tile050ns// 用上一个 tile 的数据 4. 同时从 HBM 读 A_tile2 到 L1200ns// 预取下一个 tile 5. 把 C_tile 写到 HBM200nsCube 算上一个 tile 时DMA 在搬运下一个 tile 的数据。Cube 不等数据。如何实现预取Ascend C 用CopyAPI 搬运数据。关键参数cache_modeL1_CACHE表示同时存一份到 L1。// 从 HBM 读 A_tile 到 L0A同时缓存到 L1Copy(A_L0A,A_HBMoffset,tile_m*tile_k*sizeof(half),L1_CACHE);下次用到相同的 A_tile 时直接从 L1 读不需要再走 HBM。工程经验7B 模型推理时QKV 投影的权重矩阵被 3 次复用Q、K、V。预取到 L1 后第 2、3 次访问快 15 倍。吞吐从 61 tokens/s 涨到 71 tokens/s。输出对齐优化HBM 写入有对齐要求地址必须是 32 字节对齐。不对齐写入慢 15%。错误写法// 输出地址没对齐half*C_outputC_HBMoffset;// offset 可能不是 16 的倍数正确写法// 确保输出地址 32 字节对齐uint64_taligned_offset(offset15)/16*16;// 向上对齐到 16 个 half32 字节half*C_outputC_HBMaligned_offset;或者用 Ascend C 的AlignAPI// 自动对齐输出地址autoC_alignedAlign(C_HBM,32);// 32 字节对齐性能差异对齐方式写入带宽性能影响不对齐~1.0TB/s-15%32 字节对齐~1.2TB/s基准完整代码示例200 行 Ascend C MatMul精简版核心逻辑#includekernel_operator.hconstexprintTILE_M64;constexprintTILE_K64;constexprintTILE_N64;classMatMulKernel{public:__aicore__inlinevoidProcess(GM_ADDR a,GM_ADDR b,GM_ADDR c,intM,intK,intN){// 遍历所有 tilefor(intm0;mM;mTILE_M){for(intn0;nN;nTILE_N){// 初始化 C_tile 为 0InitC(cm*Nn,TILE_M,TILE_N);for(intk0;kK;kTILE_K){// 从 HBM 读 A_tile 到 L0A预取到 L1CopyA(am*Kk,TILE_M,TILE_K);// 从 HBM 读 B_tile 到 L0B预取到 L1CopyB(bk*Nn,TILE_K,TILE_N);// Cube 算 A_tile × B_tile累加到 C_tileMatMulTile(TILE_M,TILE_K,TILE_N);}// 把 C_tile 写回 HBMWriteC(cm*Nn,TILE_M,TILE_N);}}}private:TPipe pipe;TBufTPosition::A1A_L0A;// L0A bufferTBufTPosition::B1B_L0B;// L0B bufferTBufTPosition::C1C_L0C;// L0C buffer__aicore__inlinevoidCopyA(GM_ADDR a,intm,intk){// 从 HBM 读 A_tile 到 L0A同时缓存到 L1autolenm*k*sizeof(half);Copy(A_L0A,a,len,{.cache_modeL1_CACHE});}__aicore__inlinevoidCopyB(GM_ADDR b,intk,intn){// 从 HBM 读 B_tile 到 L0B同时缓存到 L1autolenk*n*sizeof(half);Copy(B_L0B,b,len,{.cache_modeL1_CACHE});}__aicore__inlinevoidMatMulTile(intm,intk,intn){// Cube 算矩阵乘结果累加到 L0CMatMul(C_L0C,A_L0A,B_L0B,m,k,n,{.accumulatetrue});}__aicore__inlinevoidWriteC(GM_ADDR c,intm,intn){// 从 L0C 写回 HBM确保 32 字节对齐autoaligned_cAlign(c,32);autolenm*n*sizeof(half);Copy(aligned_c,C_L0C,len);}};// 算子入口externC__global__ __aicore__voidmatmul_kernel(GM_ADDR a,GM_ADDR b,GM_ADDR c,intM,intK,intN){MatMulKernel op;op.Process(a,b,c,M,K,N);}编译和运行# 编译算子npu-smiset-tmm-s0-dmatmul_kernel.o matmul_kernel.cpp# 链接成动态库ld-sharedmatmul_kernel.o-olibmatmul.so# 在 ACL 中调用aclError retaclrtLaunchKernel(matmul_kernel, grid, block, args,0, stream);性能对比实现吞吐 (tokens/s)MAC 利用率L1 命中率初版tile_m163823%0%tile_m645256%0%L1 预取6772%45%输出对齐7189%45%ops-math官方7291%48%自己写的 MatMul 性能与 ops-math 持平。差距在 L1 命中率45% vs 48%ops-math 的预取策略更激进。踩坑实录坑 1tile_m16 吞吐腰斩MAC 阵列 16×16tile_m16 只填了一行利用率 23%。改成 tile_m64利用率拉到 89%。坑 2L1 没预取Cube 等 Vector 数据不预取时Cube 算的时候数据还在搬运空转 40% 时间。加预取后数据提前到 L1Cube 不等。坑 3输出没对齐写入慢 15%HBM 写入要 32 字节对齐。用AlignAPI 自动对齐性能提 15%。https://atomgit.com/cann/ops-mathhttps://atomgit.com/cann/opbasehttps://atomgit.com/cann/cann-samples
