Block Mmad Qgmm Mx【免费下载链接】ops-tensorops-tensor 是 CANN Compute Architecture for Neural Networks算子库中提供张量类计算的基础算子库采用模块化设计支持灵活的算子开发和管理。项目地址: https://gitcode.com/cann/ops-tensor代码位置功能说明MX 量化 Grouped Matmul 的 Block 组件基于 Tensor API 实现仅支持 AIC 计算。 组件负责单个 tile 的 A/B 搬运、ScaleA/ScaleB 搬运、MX Mmad 计算以及 bias 处理适用于 grouped matmul 的 mx tensor_api kernel。继承思路参考Block Mmad 公共框架特殊约束调度策略限制仅支持GroupedMatmulWithScaleMx调度策略不用于 Basic、QBMM 或 StreamK 路径。量化数据类型支持支持以下典型 MX 量化输入类型fp4x2_e2m1_tfp4x2_e1m2_tfp8_e5m2_tfp8_e4m3fn_tScale 类型ScaleA 和 ScaleB 固定使用fp8_e8m0_t。计算模式仅支持 AIC 模式不支持 AIV 计算。输出目标结果直接输出到 GM不依赖 workspace。特殊静态常量常量说明transAA 是否转置transBB 是否转置C0_SIZE数据 C0 对齐大小FP4 为 64FP8 为 32SCALE_C0Scale 布局的 C0 对齐大小固定为 2SCALE_BUFFER_NUMScale 双缓冲数量HALF_L0_SIZEL0A/L0B 半缓冲大小HALF_L0C_SIZEL0C 半缓冲大小特殊类型别名别名含义ProblemShape当前 group 的问题规模BlockShape单 tile 形状MxL0ATypeA 在 L0 中的数据类型MxL0BTypeB 在 L0 中的数据类型特殊数据结构Paramsstruct Params { GM_ADDR aGmAddr; GM_ADDR bGmAddr; GM_ADDR cGmAddr; GM_ADDR biasGmAddr; GM_ADDR scaleAGmAddr; GM_ADDR scaleBGmAddr; };参数说明参数说明aGmAddrA 的 GM 地址bGmAddrB 的 GM 地址cGmAddrC 的 GM 地址biasGmAddrbias 的 GM 地址可为空scaleAGmAddrA 对应的 per-token scale 地址scaleBGmAddrB 对应的 per-group scale 地址L1Paramsstruct L1Params { uint64_t kAL1; uint64_t kBL1; uint64_t scaleKL1; };参数说明参数说明kAL1A 的 L1 K 轴切分kBL1B 的 L1 K 轴切分scaleKL1ScaleA/ScaleB 共享的 L1 K 轴切分参数约束kAL1、kBL1、scaleKL1均需大于 0。kAL1、kBL1建议按MXFP_DIVISOR_SIZE64对齐末尾 K tail 由实现内部 padding 处理。当kAL1 kBL1时外层按kAL1复用 A要求kAL1是kBL1的整数倍。当kBL1 kAL1时外层按kBL1复用 B要求kBL1是kAL1的整数倍。scaleKL1必须不小于外层 L1 K 窗口max(kAL1, kBL1)且应为该外层窗口的整数倍Scale 只在scaleKL1边界搬运一次并在窗口内复用。MmadParamsstruct MmadParams { BlockShape l0TileShape; L1Params l1Params; bool isBias; bool enableL0cPingPong; };特殊成员方法Init 函数__aicore__ inline void Init( const ProblemShape problemShape, const MmadParams params)功能初始化当前 group 的m/n/k设置 L0 tile 和 L1 切分参数重置组内双缓冲状态UpdateParamsForNextProblem 函数__aicore__ inline void UpdateParamsForNextProblem(const ProblemShape problemShape)功能在 grouped matmul 切换到下一个 group 时刷新m/n/k保持同一个 block 组件在不同 group 间复用operator() 函数template typename TensorA, typename TensorB, typename TensorScaleA, typename TensorScaleB, typename TensorBias, typename TensorC __aicore__ inline void operator()( TensorA gmA, TensorB gmB, TensorScaleA gmScaleA, TensorScaleB gmScaleB, TensorBias gmBias, TensorC gmC, const BlockShape singleShape)功能处理一个 tile 的 MX grouped matmul 计算输入 tensor 由 kernel 层完成 slice 后传入调用示例组件组装using AType fp8_e4m3fn_t; using BType fp8_e4m3fn_t; using CType float; using BiasType float; using LayoutA AscendC::Te::NDExtLayoutPtn; using LayoutB AscendC::Te::NZLayoutPtn; using LayoutC AscendC::Te::NDExtLayoutPtn; using LayoutBias AscendC::Te::NDExtLayoutPtn; using DispatchPolicy Blaze::Gemm::GroupedMatmulWithScaleMx0; using BlockMmad Blaze::Gemm::Block::BlockMmad DispatchPolicy, AType, LayoutA, BType, LayoutB, CType, LayoutC, BiasType, LayoutBias;组件初始化BlockMmad blockMmad; BlockMmad::ProblemShape problemShape{m, n, k, 0}; BlockMmad::BlockShape l0TileShape{baseM, baseN, baseK, 0}; BlockMmad::L1Params l1Params{kAL1, kBL1, scaleKL1}; BlockMmad::MmadParams params{l0TileShape, l1Params, isBias, enableL0cPingPong}; blockMmad.Init(problemShape, params);组件执行auto gmBlockA gmA.Slice(...); auto gmBlockB gmB.Slice(...); auto gmBlockScaleA gmScaleA.Slice(...); auto gmBlockScaleB gmScaleB.Slice(...); auto gmBlockBias gmBias.Slice(...); auto gmBlockC gmC.Slice(...); BlockMmad::BlockShape singleShape{tileM, tileN, tileK, 0}; blockMmad(gmBlockA, gmBlockB, gmBlockScaleA, gmBlockScaleB, gmBlockBias, gmBlockC, singleShape);数据流GM(A/B) GM(ScaleA/ScaleB) GM(Bias) - L1 - L0A/L0B Scale Buffer - MmadTraitMX - L0C - GM(C)适用场景MX 量化 grouped matmul 的 tensor_api kernelgroup 间m/n/k动态变化的场景同时处理输入 scale 与权重 scale 的 grouped matmul 场景【免费下载链接】ops-tensorops-tensor 是 CANN Compute Architecture for Neural Networks算子库中提供张量类计算的基础算子库采用模块化设计支持灵活的算子开发和管理。项目地址: https://gitcode.com/cann/ops-tensor创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考