Kernel基础框架【免费下载链接】ops-tensorops-tensor 是 CANN Compute Architecture for Neural Networks算子库中提供张量类计算的基础算子库采用模块化设计支持灵活的算子开发和管理。项目地址: https://gitcode.com/cann/ops-tensor公共接口说明概述矩阵乘 Kernel 基础框架提供统一的模板参数、数据结构和核心流程。不同实现Basic、QBMM MX、StreamK在此基础上扩展特定功能其中 QBMM MX 支持 MxFP4/MxFP8 量化、Scale 因子处理、多 Batch 维度。详见README.md 查看 API 清单和实现对比。类模板概述模板参数参数说明ProblemShape_问题形状类型通常为AscendC::Te::Shapeint64_t, int64_t, int64_t, int64_t(m, n, k, batch)BlockMmad_BlockMmad 类矩阵乘计算组件BlockEpilogue_BlockEpilogue 类后处理组件BlockScheduler_BlockScheduler 类任务调度组件特殊模板参数量化 Kernel参数说明isAtomicAdd是否启用 Atomic Add 模式QBMM MX说明QBMM MX Kernel 支持isAtomicAdd参数用于多核并行累加场景。类模板概述模板参数参数说明ProblemShape_问题形状类型通常为AscendC::Te::Shapeint64_t, int64_t, int64_t, int64_t(m, n, k, batch)BlockMmad_BlockMmad 类矩阵乘计算组件BlockEpilogue_BlockEpilogue 类后处理组件BlockScheduler_BlockScheduler 类任务调度组件类型别名类型说明ProblemShape问题形状类型继承自模板参数BlockMmadBlockMmad 类型继承自模板参数BlockEpilogueBlockEpilogue 类型继承自模板参数BlockSchedulerBlockScheduler 类型继承自模板参数ATypeA 矩阵数据类型继承自 BlockMmadBTypeB 矩阵数据类型继承自 BlockMmadCTypeC 矩阵数据类型继承自 BlockMmadBiasTypeBias 数据类型继承自 BlockMmadLayoutAA 矩阵布局类型继承自 BlockMmadLayoutBB 矩阵布局类型继承自 BlockMmadLayoutCC 矩阵布局类型继承自 BlockMmadLayoutBiasBias 布局类型继承自 BlockMmadTupleShapeTile 形状类型AscendC::Te::Shapeint64_t, int64_t, int64_t, int64_tLayout 构建类型类型说明MakeLayoutAA 矩阵 Layout 构建器FrameLayoutFormatLayoutA, ...MakeLayoutBB 矩阵 Layout 构建器FrameLayoutFormatLayoutB, ...MakeLayoutCC 矩阵 Layout 构建器FrameLayoutFormatLayoutC, ...MakeLayoutBiasBias Layout 构建器FrameLayoutFormatLayoutBias, ...核心数据结构Paramsstruct Params { ProblemShape problemShape; // 问题 shape (m, n, k, batch) BlockMmadParams mmadParams; // mmad 参数包含 GM 地址 BlockEpilogueParams epilogueParams; // epilogue 参数 BlockSchedulerParams schedulerParams; // scheduler 参数 };Argumentsstruct Arguments { ProblemShape problemShape; // 问题 shape (m, n, k, batch) BlockMmadArguments mmadArgs; // mmad 参数 BlockEpilogueArguments epilogueArgs; // epilogue 参数 };核心成员变量变量类型说明problemShape_TupleShape问题规模 (m, n, k, batch)isBias_bool是否启用 bias 计算aGmAddr___gm__ AType*A 矩阵 GM 地址bGmAddr___gm__ BType*B 矩阵 GM 地址cGmAddr___gm__ CType*C 矩阵 GM 地址biasGmAddr___gm__ BiasType*Bias GM 地址可选核心成员方法构造函数__aicore__ inline KernelMatmul()功能构造 Kernel 对象。析构函数__aicore__ inline ~KernelMatmul()功能析构 Kernel 对象。Init函数__aicore__ inline void Init(Params const params)功能初始化 Kernel提取问题规模和 GM 地址。 执行流程设置问题规模problemShape_提取 BlockMmad 参数mmadParams设置 A、B、C 的 GM 地址判断 bias 地址是否为 nullptr设置isBias_和biasGmAddr_operator函数__aicore__ inline void operator()(Params const params)功能执行矩阵乘 Kernel 计算。公共执行流程初始化调用Init(params)设置参数BlockScheduler 初始化创建调度器获取 tile 信息Layout 构建构建 A、B、C、Bias 的 ND layoutGM Tensor 创建创建 A、B、C、Bias 的 GM TensorTile 循环处理遍历 tile 执行矩阵乘计算清理关闭 HF32/MM Layout Transform公共调用示例组件组装模板// 定义数据类型和布局 using AType half; using BType half; using CType float; using BiasType float; using LayoutA AscendC::Te::Layout::RowMajor; using LayoutB AscendC::Te::Layout::ColMajor; using LayoutC AscendC::Te::Layout::RowMajor; using LayoutBias LayoutC; // 定义问题 shape using ProblemShape AscendC::Te::Shapeint64_t, int64_t, int64_t, int64_t; // 定义 BlockScheduler using BlockScheduler Blaze::Gemm::Block::BlockSchedulerMatmulBasicProblemShape, FULL_LOAD_MODE; // 定义 Kernel根据需求选择 Basic 或 StreamK using MatmulKernel Blaze::Gemm::Kernel::KernelMatmulBasic...; // 或 using MatmulKernel Blaze::Gemm::Kernel::KernelMatmulStreamK...;参数准备模板using Params typename MatmulKernel::Params; Params params { {m, n, k, batch}, // problem shape {aGM, bGM, cGM, biasGM}, // mmad args {...}, // epilogue args {mL1, nL1, kL1, baseM, baseN, baseK, ...} // scheduler params };Kernel 执行模板MatmulKernel mm; mm(params);公共约束模板参数要求ProblemShape 必须为AscendC::Te::Shapeint64_t, int64_t, int64_t, int64_t类型BlockMmad 必须继承自相应的 BlockMmad 基类BlockEpilogue 必须与 Kernel 类型匹配BlockScheduler 必须提供 tile 切分和调度功能数据格式A、C、Bias支持 ND 格式B 矩阵支持 ND 和 NZ 格式Bias 支持可选 bias 输入通过biasGmAddr是否为 nullptr 判断Layout 构建使用FrameLayoutFormat根据数据类型自动适配 layout性能优化建议公共Tile 大小选择L1 tile充分利用 L1 容量通常 1MBL0 tile匹配 L0A/L0B 容量各 128KBBlock 数量配置根据问题规模合理设置 block 数量HF32 模式FP16 输入 FP32 输出场景建议启用数据布局权重矩阵B优先使用 NZ 格式【免费下载链接】ops-tensorops-tensor 是 CANN Compute Architecture for Neural Networks算子库中提供张量类计算的基础算子库采用模块化设计支持灵活的算子开发和管理。项目地址: https://gitcode.com/cann/ops-tensor创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考
