Kernel Qbmm Mx【免费下载链接】ops-tensorops-tensor 是 CANN Compute Architecture for Neural Networks算子库中提供张量类计算的基础算子库采用模块化设计支持灵活的算子开发和管理。项目地址: https://gitcode.com/cann/ops-tensor代码位置功能说明MX 量化 Batch Matmul Kernel仅支持 AIC 计算支持 MxFP4/MxFP8 量化格式。集成 Scale 因子pertokenScale、scale处理、多 Batch 维度支持、L2 Cache 动态配置适用于量化推理场景。继承自Kernel Matmul 基础框架特殊约束量化格式支持支持以下量化数据类型MxFP4fp4x2_e2m1_t、fp4x2_e1m2_t4-bit 浮点MxFP8fp8_e5m2_t、fp8_e4m3fn_t8-bit 浮点Scale 因子要求必须提供两个 Scale 因子pertokenScaleGmAddrA 矩阵的 per-token scalefp8_e8m0_t类型scaleGmAddrB 矩阵的 per-group scalefp8_e8m0_t类型计算模式仅支持 AIC 模式不支持 AIV 计算AIV 核直接返回。BlockScheduler 限制仅支持BlockSchedulerQbmm调度器支持多 Batch 维度切分。L2 Cache 动态配置根据 tile 形状动态启用/禁用 L2 Cache大 tilecurBaseM problemShape.m禁用 L2 Cache小 tile启用 L2 CacheBatch 维度限制支持 4 维 BatchbatchA1/A2/A3/A4、batchB1/B2/B3/B4、batchC1/C2/C3/C4需满足广播规则。Atomic Add 模式可选 Atomic Add 模式isAtomicAdd true用于多核并行累加场景。特殊成员方法构造函数__aicore__ inline QuantBatchMmMx()功能构造 QuantBatchMmMx 对象。析构函数__aicore__ inline ~QuantBatchMmMx()功能析构 QuantBatchMmMx 对象。特殊模板参数template class ProblemShape, // 问题形状类型 class BlockMmad, // BlockMmadMX 组件 class BlockEpilogue, // 后处理组件通常为 void class BlockScheduler, // BlockSchedulerQbmm 调度器 bool isAtomicAdd // 是否启用 Atomic Add模板参数说明参数说明ProblemShape问题形状类型包含 m、n、k、bbatchBlockMmadBlockMmadMX 组件基于MatmulWithScaleMx调度策略BlockEpilogue后处理组件通常不使用传 voidBlockSchedulerBlockSchedulerQbmm 调度器isAtomicAdd是否启用 Atomic Add 模式多核并行累加特殊类型别名类型说明weightNzB 矩阵是否为 NZ 格式继承自 BlockMmadtransAA 矩阵是否转置继承自 BlockMmadtransBB 矩阵是否转置继承自 BlockMmadC0_SIZEC0 对齐大小FP4: 64FP8: 32SCALE_C0Scale C0 对齐大小固定为 2MakeLayoutScaleAScaleA Layout 构建器根据 transA 选择MakeLayoutScaleBScaleB Layout 构建器根据 transB 选择特殊数据结构Paramsstruct Params { ProblemShape problemShape; // 问题 shape (m, n, k, batch) BlockMmadParams mmadParams; // mmad 参数包含 GM 地址 L1Params l1Params; // L1 参数kL1, scaleKL1, l1BufNum BlockSchedulerParams schParams; // scheduler 参数 QBMMTiling qbmmParams; // QBMM 特有参数 };QBMMTilingstruct QBMMTiling { uint32_t batchA1, batchA2, batchA3, batchA4; // A 矩阵 Batch 维度 uint32_t batchB1, batchB2, batchB3, batchB4; // B 矩阵 Batch 维度 uint32_t batchC1, batchC2, batchC3, batchC4; // C 矩阵 Batch 维度 uint32_t biasThreeDim; // Bias 是否为 3 维 uint32_t baseM, baseN, baseK; // L0 tile 形状 uint32_t isBias; // 是否启用 bias uint32_t dbL0C; // L0C 双缓冲标志 };BlockMmadParamsMX 特有struct Params { GM_ADDR aGmAddr; // A 矩阵 GM 地址 GM_ADDR bGmAddr; // B 矩阵 GM 地址 GM_ADDR cGmAddr; // C 矩阵 GM 地址 GM_ADDR biasGmAddr; // Bias GM 地址可选 GM_ADDR pertokenScaleGmAddr; // A 矩阵 Scale GM 地址 GM_ADDR scaleGmAddr; // B 矩阵 Scale GM 地址 };特殊成员方法Init函数__aicore__ inline void Init(const Params params)功能初始化 Kernel提取问题规模、GM 地址、Batch 参数。 执行流程AIV 检查如果是 AIV 核则直接返回设置 Bias 标志根据qbmmParams.isBias判断设置 BiasThreeDim 标志根据qbmmParams.biasThreeDim判断调用ResetGmAddr设置 GM 地址Run函数__aicore__ inline void Run(const Params params)功能执行量化 Batch Matmul Kernel 计算。 执行流程Atomic Add 配置如果isAtomicAdd true调用SetAtomicAddfloat调用Init(params)设置参数创建 BlockScheduler 实例初始化 BlockMmadMX 组件判断 Batch 数量单 Batchb 1调用ProcessSingleBatch多 Batch调用ProcessWithBatch清理 Atomic Add如果启用调用SetAtomicNoneProcessSingleBatch函数__aicore__ inline void ProcessSingleBatch( const Params params, BlockScheduler bs, uint64_t batchCnt, bool isTailRound)功能处理单个 Batch 的矩阵乘计算。 执行流程构建 LayoutA、B、ScaleA、ScaleB、Bias、C创建 GM Tensor动态配置 L2 CacheTile 循环处理获取 tile 坐标和形状Slice GM Tensor 到当前 tile调用 BlockMmadMX 执行量化矩阵乘ProcessWithBatch函数__aicore__ inline void ProcessWithBatch(const Params params, BlockScheduler bs)功能处理多 Batch 的矩阵乘计算。 执行流程计算 Batch 广播倍数multiA1C1、multiB1C1 等4 维 Batch 循环batchC1/C2/C3/C4更新 Batch 偏移batchCOffset_、batchAOffset_、batchBOffset_调用AddBatchOffset更新 GM 地址调用ProcessSingleBatch处理当前 BatchSetL2Cache函数template typename TensorB, typename TensorScaleB, typename TensorC __aicore__ inline void SetL2Cache( const ProblemShape problemShape, uint64_t curBaseM, uint64_t baseN, TensorB gmB, TensorScaleB gmScaleB, TensorC gmC)功能动态配置 L2 Cache。 说明大 tile 场景禁用 B 和 ScaleB 的 L2 CacheAtomic Add 模式禁用 C 的 L2 Cache调用示例组件组装// 定义量化数据类型 using AType fp4x2_e2m1_t; // 或 fp8_e5m2_t using BType fp4x2_e2m1_t; // 或 fp8_e5m2_t using CType float; using BiasType float; using ScaleType fp8_e8m0_t; // 定义 Layout using LayoutA AscendC::Te::NDExtLayoutPtn; using LayoutB AscendC::Te::NZLayoutPtn; using LayoutC AscendC::Te::NDExtLayoutPtn; using LayoutBias AscendC::Te::NDExtLayoutPtn; // 定义调度策略 using DispatchPolicy Blaze::Gemm::MatmulWithScaleMxA_FULL_LOAD_MODE; // 定义 BlockMmadMX using BlockMmad Blaze::Gemm::Block::BlockMmad DispatchPolicy, AType, LayoutA, BType, LayoutB, CType, LayoutC, BiasType, LayoutBias; // 定义 BlockScheduler using BlockScheduler Blaze::Gemm::Block::BlockSchedulerQbmmProblemShape; // 定义 Kernel using QBMMKernel Blaze::Gemm::Kernel::QuantBatchMmMx ProblemShape, BlockMmad, void, BlockScheduler, false;参数准备using Params typename QBMMKernel::Params; Params params { {m, n, k, batch}, // problem shape {aGM, bGM, cGM, biasGM, scaleAGM, scaleBGM}, // mmad params {kL1, scaleKL1, l1BufNum}, // L1 params {baseM, baseN, baseK, mTailTile, nTailTile, ...}, // scheduler params {batchA1, batchA2, ... batchC4, biasThreeDim, baseM, baseN, baseK, isBias, dbL0C} // QBMM tiling };Kernel 执行QBMMKernel qbmm; qbmm(params); // 或 qbmm.Run(params);数据流存储层次GM (A/B/ScaleA/ScaleB/Bias) → L1 (量化数据 Scale) → L0A/L0B → L0C → GM (C)Batch 处理流程Batch 循环4 维 ↓ 更新 Batch 偏移batchA/B/COffset ↓ 重置 GM 地址AddBatchOffset ↓ 处理单个 BatchProcessSingleBatch ↓ Tile 循环 → BlockMmadMX 执行量化计算流程加载量化数据 (A/B) Scale (ScaleA/ScaleB) ↓ Dequantize自动在 Mmad 中完成 ↓ Mmad 计算MX 模式MmadTraitMX ↓ 输出 float 结果性能优化建议L1 缓冲配置l1BufNum建议 2 或 4平衡 L1 容量和流水线并行度kL1建议对齐到MXFP_DIVISOR_SIZE64Scale KL1 配置scaleKL1建议为kL1的整数倍如 2×kL1Scale 数据复用Scale 常驻 L1减少搬运开销全载模式选择非全载模式每次迭代重新加载 A/B 块A 全载模式A 矩阵常驻 L1适用于大 K、小 M 场景Batch 维度设计Batch 维度需满足广播规则batchA batchC × multiA、batchB batchC × multiB4 维 Batch 灵活支持多种广播场景L2 Cache 配置大 tile 场景自动禁用 L2 Cache避免缓存污染Atomic Add 模式自动禁用 C 的 L2 Cache适用场景量化推理MxFP4/MxFP8 量化矩阵乘Batch Matmul多 Batch 维度支持Scale 因子处理per-token 和 per-group scale广播 BatchA/B/C 不同 Batch 维度的广播计算【免费下载链接】ops-tensorops-tensor 是 CANN Compute Architecture for Neural Networks算子库中提供张量类计算的基础算子库采用模块化设计支持灵活的算子开发和管理。项目地址: https://gitcode.com/cann/ops-tensor创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考