CANN/ops-tensor量化矩阵乘法调度器

CANN/ops-tensor量化矩阵乘法调度器 Block Scheduler Quant Batch Matmul【免费下载链接】ops-tensorops-tensor 是 CANN Compute Architecture for Neural Networks算子库中提供张量类计算的基础算子库采用模块化设计支持灵活的算子开发和管理。项目地址: https://gitcode.com/cann/ops-tensor代码位置功能说明量化 Batch Matmul 调度器支持多 Batch 维度切分、尾块切分、负载均衡、Z 型扫描。适用于 Quant Batch Matmul MX Kernel 场景支持 MxFP4/MxFP8 量化数据类型。继承自Block Scheduler 公共框架模板参数template class ProblemShape_, // 问题规模类型 uint64_t FullLoadMode_, // 全载模式0非全载1A全载 class LayoutA_, // A 矩阵布局类型 class LayoutB_, // B 矩阵布局类型 class AType_ // A 矩阵数据类型用于判断 C0_SIZE class BlockSchedulerQuantBatchMatmulV3;全载模式支持两种全载模式FullLoadMode_ 0非全载模式默认FullLoadMode_ A_FULL_LOAD_MODE1A 矩阵全载模式C0_SIZE 计算根据数据类型自动计算 C0 对齐大小FP4fp4x2_e2m1_t, fp4x2_e1m2_tC0_SIZE 64FP8fp8_e5m2_t, fp8_e4m3fn_tC0_SIZE 32转置判断根据布局类型判断转置transAIsTransLayoutA_::valuetransBIsTransLayoutB_::value尾块切分支持尾块 tile 切分mTailTileM 轴尾块切分数量nTailTileN 轴尾块切分数量totalTailTile总尾块切分数量mTailTile × nTailTileZ 型扫描使用 Z 型扫描策略WINDOW_LEN 4扫描窗口大小正向扫描偶数行rowIdx % 2 0反向扫描奇数行rowIdx 1负载均衡支持负载均衡配置mBaseNormCntM 轴正常 tile 数量mBaseTailMainM 轴尾块主尺寸nBaseNormCntN 轴正常 tile 数量nBaseTailMainN 轴尾块主尺寸特殊静态常量常量说明C0_SIZEC0 对齐大小FP4: 64FP8: 32WINDOW_LENZ 型扫描窗口长度4transAA 矩阵是否转置transBB 矩阵是否转置特殊类型别名类型说明BlockShapeBlock 形状Shapeint64_t, int64_t, int64_t, int64_tBlockCoordBlock 坐标Coordint64_t, int64_t, int64_t, int64_tProblemShape问题规模类型模板参数ATypeA 矩阵数据类型模板参数特殊数据结构Paramsstruct Params { int64_t baseM; // L0 M 维度 base 大小 int64_t baseN; // L0 N 维度 base 大小 int64_t mTailTile; // M 轴尾块切分数量 int64_t nTailTile; // N 轴尾块切分数量 int64_t mBaseTailSplitCnt; // M 轴尾块 L1 切分数量 int64_t nBaseTailSplitCnt; // N 轴尾块 L1 切分数量 int64_t mTailMain; // M 轴尾块主尺寸 int64_t nTailMain; // N 轴尾块主尺寸 };特殊成员变量变量说明m_, n_, k_问题规模baseM_, baseN_L0 base 形状mCnt_, nCnt_, totalCnt_tile 数量mBaseNormCnt_, nBaseNormCnt_正常 tile 数量mBaseTailMain_, nBaseTailMain_尾块主尺寸mBaseTailLast_, nBaseTailLast_尾块最后尺寸mCoreNum_, mTailCoreNum_M 轴核心数量、尾核心数量blockIdx_, blockNum_当前 Block 索引、总 Block 数量startBlockIdx_, endBlockIdx_起始/结束 Block 索引roundIdx_, round_当前轮次、总轮次mTailTile_, nTailTile_, totalTailTile_尾块切分数量mSplitAddrOffset_, nSplitAddrOffset_尾块切分偏移mainRow_主行数特殊成员方法构造函数__aicore__ inline BlockSchedulerQuantBatchMatmulV3(const ProblemShape shape, const Params params)功能初始化 BlockSchedulerQuantBatchMatmulV3计算 tile 切分、尾块参数、轮次等。 参数说明 | 参数 | 类型 | 说明 | |------|------|------| | shape | ProblemShape | 问题规模(m, n, k)| | params | Params | 调度参数 |执行流程设置问题规模m_,n_,k_,baseM_,baseN_计算 tile 数量mCnt_,nCnt_,totalCnt_计算扫描窗口mCoreNum_,mainRow_,mTailCoreNum_计算轮次endBlockIdx_,round_计算尾块参数根据transA和transB计算尾块尺寸UpdateTailTile__aicore__ inline void UpdateTailTile(uint32_t mTailTile, uint32_t nTailTile)功能更新尾块切分数量重新计算结束 Block 索引和轮次。 参数说明 | 参数 | 类型 | 说明 | |------|------|------| | mTailTile | uint32_t | M 轴尾块切分数量 | | nTailTile | uint32_t | N 轴尾块切分数量 |GetTotalCnt__aicore__ inline int64_t GetTotalCnt()功能返回总 tile 数量totalCnt_。GetEndBlockIdx__aicore__ inline int64_t GetEndBlockIdx()功能返回结束 Block 索引endBlockIdx_。CalSingleCoreShapeByCoord__aicore__ inline void CalSingleCoreShapeByCoord(int64_t singleCoreM, int64_t singleCoreN, const BlockCoord blockCoord)功能根据 Block 坐标计算单核形状处理尾块。 参数说明 | 参数 | 类型 | 说明 | |------|------|------| | singleCoreM | int64_t | 单核 M 维度原地修改 | | singleCoreN | int64_t | 单核 N 维度原地修改 | | blockCoord | BlockCoord | Block 坐标 |GetBlockShapetemplate QuantBatchMatmul::QuantMode aQuantMode, QuantBatchMatmul::QuantMode bQuantMode, bool weightNz false __aicore__ inline BlockShape GetBlockShape(BlockCoord blockCoord)功能返回当前 Block 的形状支持量化模式和 NZ 格式。 参数说明 | 参数 | 类型 | 说明 | |------|------|------| | aQuantMode | QuantMode | A 矩阵量化模式PERGROUP/PERBLOCK | | bQuantMode | QuantMode | B 矩阵量化模式PERGROUP/PERBLOCK | | weightNz | bool | B 矩阵是否为 NZ 格式 | | blockCoord | BlockCoord | Block 坐标 |返回值BlockShape {singleCoreM, singleCoreN, mSplitAddrOffset_, nSplitAddrOffset_}特殊逻辑尾块切分判断totalTailTile_ 1 roundIdx_ round_FP4 对齐FP4 transA 时 M 对齐到 2FP4 !transB 时 N 对齐到 2PERBLOCK 模式对齐到 PER_BLOCK_SIZE 或 2 的幂次方NZ 格式对齐根据 transB 对齐到 C0_SIZE 或 BLOCK_CUBEGetLoadBalanceInfo__aicore__ inline AscendC::Std::tupleuint32_t, uint32_t, uint32_t, uint32_t GetLoadBalanceInfo()功能返回负载均衡信息{mBaseNormCnt_, mBaseTailMain_, nBaseNormCnt_, nBaseTailMain_}。UpdateNextBatchBlockRoundParams__aicore__ inline void UpdateNextBatchBlockRoundParams()功能更新下一 Batch 的 Block 轮次参数。 执行流程更新startBlockIdx_和endBlockIdx_重置roundIdx_ 0重新计算round_GetTileIdx__aicore__ inline bool GetTileIdx(BlockCoord blockCoord)功能获取当前轮次的 tile 索引更新 Block 坐标。 参数说明 | 参数 | 类型 | 说明 | |------|------|------| | blockCoord | BlockCoord | Block 坐标原地修改 |返回值true当前轮次有效返回 tile 坐标false当前轮次结束执行流程判断轮次结束roundIdx_ round_计算 tile 索引根据全载模式计算Z 型扫描计算blockCoordM,blockCoordN反向扫描奇数行反向blockCoordN nCnt_ - 1 - blockCoordN更新roundIdx_GetTileCoord__aicore__ inline void GetTileCoord(const BlockCoord blockCoord, int64_t mPos, int64_t nPos)功能根据 Block 坐标计算 GM 地址偏移。 参数说明 | 参数 | 类型 | 说明 | |------|------|------| | blockCoord | BlockCoord | Block 坐标 | | mPos | int64_t | M 轴 GM 偏移原地修改 | | nPos | int64_t | N 轴 GM 偏移原地修改 |调用示例组件组装using ProblemShape AscendC::Te::Shapeint64_t, int64_t, int64_t; using LayoutA AscendC::Te::NZLayoutPtn; using LayoutB AscendC::Te::NZLayoutPtn; using AType fp4x2_e2m1_t; constexpr uint64_t FULL_LOAD_MODE 0; using BlockScheduler Blaze::Gemm::Block::BlockSchedulerQuantBatchMatmulV3 ProblemShape, FULL_LOAD_MODE, LayoutA, LayoutB, AType;参数准备BlockScheduler::Params params { baseM, // L0 M 维度 base如 128 baseN, // L0 N 维度 base如 128 mTailTile, // M 轴尾块切分数量如 1 nTailTile, // N 轴尾块切分数量如 1 mBaseTailSplitCnt, // M 轴尾块 L1 切分数量如 1 nBaseTailSplitCnt, // N 轴尾块 L1 切分数量如 1 mTailMain, // M 轴尾块主尺寸如 1 nTailMain // N 轴尾块主尺寸如 1 };组件初始化ProblemShape shape{m, n, k}; BlockScheduler scheduler(shape, params);更新尾块切分scheduler.UpdateTailTile(mTailTile, nTailTile);获取 tile 数量int64_t totalCnt scheduler.GetTotalCnt();Tile 循环处理BlockCoord blockCoord{0, 0, 0, 0}; while (scheduler.GetTileIdx(blockCoord)) { // 获取 Block 形状 constexpr auto aQuantMode QuantBatchMatmul::QuantMode::PERGROUP_MODE; constexpr auto bQuantMode QuantBatchMatmul::QuantMode::PERGROUP_MODE; constexpr bool weightNz true; auto blockShape scheduler.GetBlockShapeaQuantMode, bQuantMode, weightNz(blockCoord); int64_t singleCoreM Get0(blockShape); int64_t singleCoreN Get1(blockShape); int64_t mSplitOffset Get2(blockShape); int64_t nSplitOffset Get3(blockShape); // 获取 GM 地址偏移 int64_t mPos, nPos; scheduler.GetTileCoord(blockCoord, mPos, nPos); // 执行 BlockMmadMX 计算 // ... }更新下一 Batchscheduler.UpdateNextBatchBlockRoundParams();获取负载均衡信息auto loadBalanceInfo scheduler.GetLoadBalanceInfo(); uint32_t mBaseNormCnt std::get0(loadBalanceInfo); uint32_t mBaseTailMain std::get1(loadBalanceInfo); uint32_t nBaseNormCnt std::get2(loadBalanceInfo); uint32_t nBaseTailMain std::get3(loadBalanceInfo);数据流Tile 切分流程问题规模 (m, n, k) ↓ L0 tile 切分 (baseM, baseN) ↓ tile 数量计算 (mCnt × nCnt) ↓ 尾块参数计算 (mBaseNormCnt, mBaseTailMain, nBaseNormCnt, nBaseTailMain) ↓ 轮次计算 (round, startBlockIdx, endBlockIdx) ↓ Z 型扫描 (mCoreNum, mainRow, mTailCoreNum) ↓ Block 形状/坐标 (singleCoreM, singleCoreN, mPos, nPos)量化模式对齐流程GetBlockShapeQuantMode, QuantMode, weightNz ↓ 尾块切分判断 (totalTailTile 1 roundIdx round) ↓ FP4 对齐transA → M 对齐到 2!transB → N 对齐到 2 ↓ PERBLOCK 模式对齐到 PER_BLOCK_SIZE 或 2 的幂次方 ↓ NZ 格式对齐!transB → C0_SIZEtransB → BLOCK_CUBE ↓ 返回 BlockShape {singleCoreM, singleCoreN, mSplitAddrOffset, nSplitAddrOffset}Z 型扫描流程tileIdx 计算 ↓ rowIdx tileIdx / nCnt / mCoreNum ↓ rowIdx mainRowblockCoordM rowIdx × mCoreNum tileIdx % mCoreNum ↓ rowIdx mainRow尾窗口计算 ↓ rowIdx 1反向扫描blockCoordN nCnt - 1 - blockCoordN性能优化建议baseM/baseN 配置建议值根据量化数据类型选择如 128C0 对齐确保 baseM/baseN 对齐到 C0_SIZEFP4: 64FP8: 32尾块切分配置mTailTile/nTailTile建议尾块切分数量不超过 4mBaseTailSplitCnt/nBaseTailSplitCnt建议 L1 尾块切分数量为 1不切分全载模式选择非全载模式FullLoadMode 0适用于一般场景A 全载模式FullLoadMode A_FULL_LOAD_MODE适用于大 K、小 M 场景量化模式选择PERGROUP_MODEper-group 量化适用于小规模量化PERBLOCK_MODEper-block 量化对齐要求更高NZ 格式优化weightNz trueB 矩阵 NZ 格式提升 L1/L0 搬运效率对齐要求根据 transB 选择 C0_SIZE 或 BLOCK_CUBE适用场景Quant Batch Matmul MX Kernel量化 Batch MatmulMxFP4/MxFP8 量化支持 FP4 E2M1/E1M2 和 FP8 E5M2/E4M3FN多 Batch 维度支持 4 维 BatchbatchA/A2/A3/A4负载均衡动态调整 tile 分配【免费下载链接】ops-tensorops-tensor 是 CANN Compute Architecture for Neural Networks算子库中提供张量类计算的基础算子库采用模块化设计支持灵活的算子开发和管理。项目地址: https://gitcode.com/cann/ops-tensor创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考