CANN/cannbot-skills：计算层（Blaze Matmul）参考

计算层Blaze Matmul参考【免费下载链接】cannbot-skillsCANNBot 是面向 CANN 开发的用于提升开发效率的系列智能体本仓库为其提供可复用的 Skills 模块。项目地址: https://gitcode.com/cann/cannbot-skills本文档承载 MC2 skill 的计算子能力。涵盖Blaze 模板选型、BlockMmad 接入、Tiling 数据流、与通信层的 buffer 协议。Blaze 是 Ascend C 的 CUTLASS 风格模板库由tensor_api/提供。本 skill 与ascendc-blaze-best-practice共享 Blaze 基底但聚焦 MC2 场景下的差异点。禁止使用 asc-devkit 的 matmul APIAscendC::Matmul等。1. MC2 场景下的 Blaze 子集参考工程用到的 Blaze 组件层级组件文件作用KernelBlaze::Gemm::Kernel::QuantMatmulMxKernelSwatinclude/kernel/qbmm_mx_kernel.hSWAT 量化 Matmul kernel 包装遍历所有 rank 累加 L0CBlockBlaze::Gemm::Block::BlockMmadblaze/gemm/block/block_mmad_qbmm_mx.htoolkit 内单 Block 的 MMAD 计算L0A×L0B→L0C→GMBlock SchedulerBlaze::Gemm::Block::Block_schedulerQuantBatchMatmulV3include/block/quant_matmul_mx_block_scheduler_swat.h多 Block 间任务切分Dispatch PolicyBlaze::Gemm::MatmulWithScaleMxblaze/gemm/policy/dispatch_policy.htoolkit 内流水策略含 scale 处理TileBlaze::Gemm::Tile::*include/tile/*.hL1→L0 搬运、Scale padLayout/TensorAscendC::Te::*tensor_api/Tensor / Layout 抽象Agent 开发原则include/block/、include/tile/、include/policy/下的文件[REUSE]常规 MC2 算子不需要改。需要改的是include/kernel/qbmm_mx_kernel.hScale 处理、A/B 来源切换include/kernel/all_to_all_matmul_impl.h通算流水编排include/tiling/quant_matmul_mx_tiling_swat.hTiling 字段。2. 从 Blaze 到 MC2 的桥接参考工程在qbmm_mx_kernel.h定义QuantMatmulMxKernelSwat把标准 BlazeBlockMmad接入 MC2 通算流水。AIC 在通信 buffer 上遍历所有 rank把各 rank 的部分和在 L0C 上累加最后一次mmadOp_触发 fixpipe 输出 GM。// qbmm_mx_kernel.h ProcessSingleBatch 核心逻辑简化伪码 for (uint64_t rank 0; rank rankSize; rank) { auto actualMPos rank * oriM mPos; // local A 按 oriM 分段 auto actualCommMPos rank * headMSize mPos; // 通信 buffer 按 headMSize 分段 // 默认从通信 buffer 读 Arank rankId 时改从本卡 GM 读 auto gmBlockA gmA.Slice(actualCommMPos, ...); auto gmBlockScaleA gmScaleA.Slice(actualMPos, ...); if (rank rankId) { gmBlockA gmALocal.Slice(actualMPos, ...); // 本卡 GM gmBlockScaleA gmScaleALocal.Slice(actualMPos, ...); } // B / ScaleB 始终从本卡 GM 读按 rank 切 K 轴段 auto gmBlockB gmB.Slice(rank * K nPos, ...); // L0C 上累加remoteRankCnt0 时 L0C reset最后一次触发 fixpipe mmadOp_(gmBlockA, gmBlockB, gmBlockScaleA, gmBlockScaleB, gmBlockBias, gmBlockC, singleShape, remoteRankCnt); remoteRankCnt; }2.1 为什么 rank rankId 时从本卡 GM 读AllToAllComm::PutToAllRanks中每个 Block 只 Put 给remoteRank ! rankId的对端——本 rank 不 Put 给自己。所以通信 buffer 中本 rank 的段是未填充的AIC 算到rank rankId时必须切换到本卡 GMlocalAGmAddr_否则会读到未初始化数据。2.2 remoteRankCnt 的语义mmadOp_的最后一个参数remoteRankCnt控制 L0C 累加位置 0L0C reset开始新的累加序列 0在已有结果上累加 splitKNum - 1这里 splitKNum rankSize触发 fixpipe把 L0C 累加结果写回 GM。参考工程在SetupParams中固定splitKNum rankSize保证for rank循环最后一次正好触发 fixpipe。3. Layout 与 Tensor 构造参考工程在qbmm_mx_kernel.h::ProcessSingleBatch构造 Blaze Tensor 句柄// Layout 构造 auto layoutA MakeLayoutA{}(rankSize * Te::GetMNK_M(problemShape), Te::GetMNK_K(problemShape)); auto layoutB MakeLayoutB{}(rankSize * Te::GetMNK_K(problemShape), Te::GetMNK_N(problemShape)); auto layoutC MakeLayoutC{}(Te::GetMNK_M(problemShape), Te::GetMNK_N(problemShape)); // Tensor 句柄 auto gmA Te::MakeTensor(Te::MakeMemPtrTe::Location::GM(aGmAddr_), layoutA); auto gmB Te::MakeTensor(Te::MakeMemPtrTe::Location::GM(bGmAddr_), layoutB); auto gmC Te::MakeTensor(Te::MakeMemPtrTe::Location::GM(cGmAddr_), layoutC);MC2 场景的 layout 维度注意 layout 维度包含rankSize *因为 A 和 B 矩阵在 MC2 中是按 rank 切分后逻辑拼起来的A 矩阵每卡持有自己的 M 段但 SHMEM buffer 收齐所有 rank 的 M 段后逻辑上是rankSize * headMSize行的矩阵B 矩阵每卡持有自己的 K 段kPerRank逻辑上是rankSize * kPerRank K行的矩阵C 矩阵本卡只输出自己 M 段的结果维度是headMSize * N不含 rankSize。Slice 时按 rank 切// 取远程 rank 的 A 段 auto gmBlockA gmA.Slice( Te::MakeCoord(rank * GetMNK_M(problemShape) mPos, kPos), Te::MakeShape(GetMNK_M(singleShape), GetMNK_K(problemShape)));4. Tiling 数据流参考工程的 Tiling 链路Host: QuantMatmulTilingSwat::GetTilingData() ↓ 填充 QuantMatmulTilingData (baseM/baseN/baseK/...) Device: AllToAllQuantMatmulImpl::SetupParams() ↓ 转成 Blaze Params (BlockMmadParams L1Params ...) Device: QuantMatmulMxKernelSwat::Run() ↓ 调用 BlockMmad Device: BlockMmad::operator() 做 MMADTiling 字段速查include/tiling/quant_matmul_tiling_data.h字段含义m, n, k问题 shape注意 k 是单卡的 kPerRankbaseM, baseN, baseK单 block 的 tile shapemTailTile, nTailTile尾块切分mBaseTailSplitCnt, nBaseTailSplitCnt尾块分裂策略mTailMain, nTailMain尾块主区大小usedCoreNum实际使用的 AIC 核数dbL0cL0C DoubleBuffer 深度1 或 2scaleKL1Scale 在 L1 的复用深度stepKK 轴 stepnBufferNumL1 buffer 数MC2 专属 Tiling 字段include/tiling/all_to_all_matmul_tiling_data.hstruct AllToAllCommTilingData { uint32_t tileCnt; // M 轴切分块数headMSize * tileCnt M uint32_t bufferSize; // 通信流水深度典型 4 };参考工程 host 侧的 Tiling 计算src/all_to_all_matmul.cppuint32_t headMSize 512; // 参考工程默认值非最优——详见 pipeline_tuning.md uint32_t tileCnt (m - tailMSize) / headMSize; tilingData.commTilingData.tileCnt tileCnt; tilingData.commTilingData.bufferSize 4; // 每块 matmul 的 Blaze tiling 由 GetTilingData 根据 headMSize 自动推导 tilingEngine.GetTilingData(headMSize, n, ka, false, true, tilingData.tileQbmmTilingData);headMSize512只是参考工程经验起点实际最优tileCnt即headMSize M/tileCnt以msprof op实测为准。Step 2-4 设计/审查阶段建议先用tileCnt1headMSizem做串行基线Step 6 再扫描tileCnt找最优——详见pipeline_tuning.md。Tiling 算法include/tiling/quant_matmul_mx_tiling_swat.h实现了 SWATSoft Wassenaar Allocation of Tiling软分配 tiling算法CalcBasicBlock()从 256 出发对齐到 CUBE_BLOCK16和 L1 对齐粒度OptimizeEdgeBasicBlock()合并 K 对齐时的尾块CalcTailBasicBlock()尾块按 M/N 双向切分CalcPathSpecificL1()搜索 L1 深度A/B 对称起步必要时打破。新算子一般不改这些算法只调 host 侧的headMSize等价于调tileCnt M/headMSize等参数。5. DispatchPolicy 选择参考工程用MatmulWithScaleMx带 MX 量化 scale 的 dispatch policy// all_to_all_matmul_impl.h using 声明段 using DispatchPolicy Blaze::Gemm::MatmulWithScaleMxNONE_FULL_LOAD_MODE, false;NONE_FULL_LOAD_MODEA/B 不全载 L1与ascendc-blaze-best-practice的模式选择一致第二个模板参数false是isKClonedToMTE1K 是否克隆到 MTE1 缓存。可选的 DispatchPolicy详见ascendc-blaze-best-practice的matmul_pattern.md§10MatmulMultiBlockPolicyNO_FULL_LOAD_MODE通用多 block SWATMatmulMultiBlockPolicyA_FULL_LOAD_MODEA 全载N≫M 时用MatmulWithScaleMx...MX 量化 matmul参考工程用。MC2 算子若非量化场景可改用MatmulMultiBlockPolicy。6. Scale量化系数处理参考工程是 MX FP8 量化 matmulScale 处理是关键。两条路径6.1 Scale 的 SHMEM PutScale 不参与 M 轴流水AIV 启动时一次性 Putall_to_all_matmul_impl.h::AllToAllProcess开头allToAllComm_.PutScaleToAllRanks(0, axisM_); // offset0, 全 M 行6.2 Scale 的 L0B 加载参考工程在include/tile/copy_scale_l1_to_l0b.h处理 Scale 从 L1 到 L0B 的搬运。qbmm_mx_kernel.h的SetScaleL2CacheProcessSingleBatch内控制 Scale 的 L2 cache 行为数据对齐时禁用 cache 走 stream不对齐时走 normal。6.3 非 MX 场景若新算子不做量化纯 BF16/FP16 matmul可以去掉allToAllComm_.PutScaleToAllRanks把BlockMmad模板从量化版block_mmad_qbmm_mx.h换成普通版block_mmad.hTiling 字段去掉 scaleKL1、stepK 等。7. 排错速查现象可能原因排查方向编译报BlockMmad模板参数错误DispatchPolicy 与 BlockScheduler 的*_LOAD_MODE不一致检查NONE_FULL_LOAD_MODE是否两侧都用精度错局部对整体差remoteRankCnt没从 0 起算 / rankrankId 时未切换到本卡 GM核对 ProcessSingleBatch 中的 if (rank rankId) 分支L0C 累加结果不对splitKNum与实际 rank 遍历数不一致检查 SetupParams 中splitKNum rankSize_是否正确fixpipe 时崩溃C 地址偏移错mOffset 算错打印cGm_ mOffset * axisN_ * sizeof(CType)与预期对比性能不达标cube ratio 低L1 深度过浅K 轴反复加载调大scaleKL1或换MatmulMultiBlockPolicyA_FULL_LOAD_MODEblaze 头文件找不到target_include_directories漏了_BLAZE_COMMON_DIR参考CMakeLists.txt的target_include_directories段8. 与 ascendc-blaze-best-practice 的关系ascendc-blaze-best-practiceskill 是 Blaze 单算子无跨卡通信的完整指南。本 skill 复用其 Blaze 基底但在以下方面不同维度ascendc-blaze-best-practice本 skillMC2通信无SHMEM/UDMA 跨卡数据来源全部本卡 GMrankrankId 走本卡 GM其他 rank 走 SHMEM bufferTiling单卡 L1/L0 容量 SHMEM 空间预算Scale可选量化场景必需参考工程是 MX FP8新算子设计时先读ascendc-blaze-best-practice选 Blaze 模板再用本 skill 把模板接入 MC2 通算流水。9. 后续阅读想了解读SHMEM/UDMA 通信层comm_shmem.mdMC2 整体架构mc2_architecture.mdBlaze 单算子细节模板选型、Tiling 算法ascendc-blaze-best-practice/references/matmul_pattern.md参考工程改造食谱codebase_map.md【免费下载链接】cannbot-skillsCANNBot 是面向 CANN 开发的用于提升开发效率的系列智能体本仓库为其提供可复用的 Skills 模块。项目地址: https://gitcode.com/cann/cannbot-skills创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考