尧图网站设计 尧图网站设计YAOTU DESIGN
ARTICLE DETAIL

资讯详情

深耕网站设计与一线实操的经验洞察。

Cube算子终极实现:ops-nn mat_mul_v3与quant_batch_matmul源码深度剖析

Cube算子终极实现:ops-nn mat_mul_v3与quant_batch_matmul源码深度剖析 Cube算子终极实现ops-nn mat_mul_v3与quant_batch_matmul源码深度剖析【免费下载链接】ops-nn本项目是CANN提供的神经网络类计算算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-nn在昇腾NPU上跑大模型矩阵乘法matmul占用了绝大部分算力。CANN 算子库ops-nn提供了神经网络类高阶算子让网络在 NPU 上加速计算其中mat_mul_v3负责通用矩阵乘quant_batch_matmul_v3负责 INT8/INT4 低比特量化矩阵乘。本文将从源码结构、Tiling 策略到量化实现带你完整看懂这对Cube 算子双子星的设计思路。为什么 Cube 算子是 NPU 加速的核心NPU 的硬件核心是 Cube 单元矩阵乘法硬件它专门负责矩阵乘加运算。Transformer 中的线性层、注意力打分、FFN 全部是矩阵乘谁能把矩阵乘喂给 Cube 单元且喂得不饿谁的性能就高。ops-nn 中matmul/目录收录了 30 多个矩阵乘算子而最基础的通用入口就是 mat_mul_v3低比特推理场景则交给 quant_batch_matmul_v3。mat_mul_v3 源码剖析通用矩阵乘的实现骨架功能与公式一个算子承载多种调用方式mat_mul_v3实现通用矩阵乘计算公式为$$C op(A) op(B) bias$$其中 A 为 (M, K)B 为 (K, N)bias 为 (N,)。它支持 FLOAT16/BFLOAT16/FLOAT32覆盖 Atlas A2/A3、Atlas 推理系列、Ascend 950PR/950DT、Kirin X90/9030 等多款产品。对外它统一了aclnnMatmul、aclnnMm、aclnnAddmm等多种 API 入口源码见 op_api 目录 与 aclnnMatmul 文档。三层目录结构从 API 到 Cube kernel每个算子都遵循统一的工程结构mat_mul_v3也不例外op_host主机侧逻辑包括形状推断 mat_mul_v3_infershape.cpp 和 Tiling 计算op_tiling 目录。Tiling 决定矩阵如何切块、如何分配到各 AI Core是性能调优的关键战场。op_kernelNPU 侧 kernel 实现op_kernel 目录 下能看到 full_load、splitk、streamk、pingpong 等多种数据流模板。op_graph图模式下的融合 Passfusion_pass 目录 中实现了 GEMM 与 MatMul 互转、matmulrelu融合等规则框架图会自动改写成最优形态。Tiling 策略split-K 与 Stream-K 如何榨干 Cube大模型中经常出现 K 维极大、M/N 维较小的瘦长矩阵乘单核串行计算会让 Cube 利用率暴跌。mat_mul_v3在 Tiling 阶段提供了多种策略split-K把 K 维拆到多个 AI Core 上并行累加见 mat_mul_sc_splitk_tiling.cpp 同目录下的 splitk 相关文件Stream-K按块流水式分发计算任务避免 split-K 末尾的空泡核心 kernel 为 mat_mul_stream_k_kernel.h配套 Tiling 为 qbmm_streamk_tiling 风格的通用策略注册。此外mat_mul_v3还支持 ND 与 FRACTAL_NZ 两种数据格式互转mat_mul_nd2nz.h以及 runtime 配置知识库——不同芯片形态各有一份预调优参数表开箱即得最优 Tiling。quant_batch_matmul 源码剖析低比特量化矩阵乘为什么要做量化 matmulINT8 权重的矩阵乘可以把访存带宽需求降到 FP16 的 1/4计算吞吐翻倍。quant_batch_matmul_v3正是为此设计支持 INT8/INT4 输入、2~6 维 batch 矩阵乘并内置反量化dequantepilogue一次 kernel 内完成量化乘 还原省去单独 dequant 算子的搬运开销。其核心公式无 pertoken、无 bias 场景为$$out (x1 x2) \times scale offset$$完整参数表与多产品支持矩阵见 READMEAPI 文档见 aclnnQuantMatmulV3。kernel 组织按量化粒度分模板op_kernel 目录 的命名直接体现了量化粒度的演进quant_batch_matmul_v3_bf16_basic.h基础版 INT8×INT8→BF16pertensor 量化quant_batch_matmul_v3_pertoken.hper-token 动态量化每行一个 scale是 LLM 动态量化的主力模板quant_batch_matmul_v3_basic_epilogue.h统一的反量化/加 bias 尾处理把 scale、offset、bias 组合逻辑收敛到一处arch35 目录下还有 INT4 预处理qbmm_int4_to_int8_preprocess.h、MX 系列 fp8/fp4 等新数据类型模板支撑 A8W4 等混合精度组合。Tiling 侧同样按策略分层quant_batch_matmul_v3_basic_tiling.cpp 负责通用切分arch35 子目录提供自适应滑动窗口、iterbatch、streamk 等高性能 Tiling思路与mat_mul_v3一脉相承。快速上手从示例代码开始两个算子都提供了完整的 aclnn 调用示例照着改参数即可跑通通用矩阵乘test_aclnn_matmul.cpp、test_aclnn_addmm_aclnninplace_addmm.cpp量化矩阵乘test_aclnn_quant_matmul_v3_at.cpp、INT4 场景 test_aclnn_quant_matmul_weight_nz_a8w4.cpp项目根目录的 docs/QUICKSTART.md 提供了零基础编译与调用指南docs/zh/develop/aicore_develop_guide.md 则适合想深入开发自定义 kernel 的进阶读者。总结维度mat_mul_v3quant_batch_matmul_v3定位通用矩阵乘 bias低比特量化矩阵乘 反量化数据类型FP16/BF16/FP32INT8/INT4/HIFLOAT8/FP8 等亮点Stream-K/split-K Tiling、图融合per-token 动态量化、A8W4 混合精度适用场景训练与高精度推理大模型低比特推理加速ops-nn 的这两个算子展示了 Cube 算子开发的完整范式统一 API 入口、按芯片分层 Tiling、按数据流模板实现 kernel、用融合 Pass 自动优化。理解它们就掌握了在 NPU 上开发高性能矩阵乘算子的钥匙。【免费下载链接】ops-nn本项目是CANN提供的神经网络类计算算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-nn创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考
返回列表