cannbot-skills Blaze后处理模块
Blaze Custom Epilogue 层模块手册【免费下载链接】cannbot-skillsCANNBot 是面向 CANN 开发的用于提升开发效率的系列智能体本仓库为其提供可复用的 Skills 模块。项目地址: https://gitcode.com/cann/cannbot-skills定位blaze_custom 路径下 Epilogue 层各模块的使用手册。Epilogue 层运行在 AIV 核心上接收 AIC 通过 CopyL0C2UB 写入 UB 的 matmul 结果执行 vector 后处理并写回 GM。§1 模块清单模块头文件类型说明EpilogueFusionRegBaseepilogue/epilogue_fusion_regbase.h通用 RegBase 框架使用__VEC_SCOPE__Reg::API中间值不写回 UBMulEpilogueepilogue/epilogue_fusion_membase.h参考 MemBase 实现Output matmul(A,B) × D使用AscendC::MulCvSync常量epilogue/cv_sync_constants.h同步常量集MODE4, AIC_TO_AIV_FLAG8, AIV_TO_AIC_FLAG5, COUNT_ID_MAX15, COUNT_FLAG3设计文档入口通用设计方法development/epilogue-dev-guide.mdMemBase 设计development/epilogue-membase-design.mdRegBase 设计development/epilogue-regbase-design.mdRegBase 核心特性计算核心使用__VEC_SCOPE__块内部声明Reg::RegTensorT和Reg::MaskReg中间值全程驻留向量寄存器不落 UB减少 UB 中间 buffer 数量通过Reg::LoadAlign加载、Reg::StoreAlign写回支持Reg::Add/Sub/Mul/Div、Reg::Adds/Muls/Divs、Reg::Axpy、Reg::Exp/Log/Abs/Sqrt、Reg::Compare/Select等 APIMemBase 参考实现MulEpilogue使用标准 AscendC APIAscendC::Mul三路 UB buffercLocal_matmul 结果dLocal_第二路输入cLocalTmp_计算结果StageNum 2UB 分区数§2 RegBase vs MemBase 选择判断维度选 RegBase选 MemBase公式复杂度复杂公式GELU / SwiGLU / LayerNorm3 中间值简单 vector 操作Mul / Add / Div可用 API需要Reg::系列 APIExp/Log/Axpy 等已有现成AscendC::API如AscendC::MulUB 占用较少2-3 blocks较多3-4 blocksAIV 效率更高寄存器级操作无 UB 读写开销一般需 UB 中间 buffer 读写开发难度较高需理解 RegTensor/MaskReg/VF 循环较低类标准 AscendC 开发决策建议优先 RegBase。仅当融合操作为单一 vector op 且已有对应AscendC::API 时选 MemBase。设计细节请查阅development/epilogue-dev-guide.mddevelopment/epilogue-membase-design.mddevelopment/epilogue-regbase-design.md§3 接口契约普通 CV 与 MX CV 的 Epilogue 类必须提供以下四个接口与MatmulKernelFused/MxMatmulKernelFused兼容接口签名说明Paramsstruct Params { ... }参数结构体含 GM_ADDR 字段Initvoid Init(Params, baseM, baseN, ProblemShape)初始化 UB 布局、绑定 GM 地址operator()void operator()(BlockShape, gmOffset, flagId)逐 tile 执行后处理GetTensorauto GetTensor()返回cLocal_UB Tensor供 Kernel 层 CopyL0C2UB 使用Grouped CV 使用 context-based Epilogueoperator()签名不同__aicore__ inline void operator()(BlockShape blockShape, TileContext tileContext, int64_t flagId)TileContext由GroupMatmulKernel传入包含groupIdx、prefixM、mOffset、nOffset、writeM、totalM、totalN、totalK等字段。普通 linear offset Epilogue 不直接兼容 Grouped CV。兼容性矩阵Epilogue 类型普通 CVMX CVGrouped CVMulEpilogue/ 普通 linear offset Epilogue可用可用但必须按GetTaskRation()/GetSubBlockIdx()消费 SplitM 后的 UB不直接可用MX SplitM-aware Epilogue不需要可用不直接可用context-based grouped Epilogue不需要不需要必需Params 结构体约定RegBaseGM_ADDR extraInputAddr额外输入GM_ADDR outputGmAddr输出MemBaseMulEpilogueGM_ADDR multiplierGmAddr乘数 DGM_ADDR outputGmAddr输出BlockShape / ProblemShapeBlockShape AscendC::Shapeint64_t, int64_t, int64_t, int64_t—{curM, curN, 1, 1}ProblemShape AscendC::Shapeint64_t, int64_t, int64_t, int64_t—{m, n, k, 1}辅助方法InitParams(Params)—__host_aicore__静态方法用于 host 端参数透传§4 UB 占用情况路径UB 分区说明RegBase2-3 blockscLocal_matmul 结果固定extraBuf_额外输入1 行bf16Out_输出stageRows 行MemBase3-4 blockscLocal_matmul 结果固定dLocal_第二路输入stageSizecLocalTmp_计算结果stageSizeRegBase UB 布局计算matmulAreaBytes ceilDiv(baseM, taskRation) × nAlign × sizeof(L0CDataType) extraBufBytes nAlign × sizeof(ComputeType) remainBytes TOTAL_UB_SIZE - matmulAreaBytes - extraBufBytes stageRows_ remainBytes / (nAlign × sizeof(OutputType))MemBase UB 布局计算matmulArea ceilDiv(baseM, taskRation) × nAlign // 元素数 lastUBBytes TOTAL_UB_SIZE - matmulArea × sizeof(DataType) stageSize_ min(lastUBBytes / StageNum / sizeof(DataType), matmulArea)RegBase 通过ReinterpretCast划分 UB 区域MemBase 通过LocalTensor偏移划分。RegBase 的 stage 循环按行切分rowsThisStageMemBase 按元素数切分curStageSize。§5 命名空间说明模块命名空间说明EpilogueFusionRegBase全局作用域不在任何 namespace 内MulEpilogue全局作用域不在任何 namespace 内CvSyncCvSync::namespace同步常量集合Epilogue 类定义在全局作用域不归属Kernel::或Block::namespace。Kernel 层通过模板参数直接引用 Epilogue 类名如MatmulKernelFused..., EpilogueFusionRegBase。【免费下载链接】cannbot-skillsCANNBot 是面向 CANN 开发的用于提升开发效率的系列智能体本仓库为其提供可复用的 Skills 模块。项目地址: https://gitcode.com/cann/cannbot-skills创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考