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 模块清单

设计文档入口：

• 通用设计方法：development/epilogue-dev-guide.md
• MemBase 设计：development/epilogue-membase-design.md
• RegBase 设计：development/epilogue-regbase-design.md

RegBase 核心特性：

• 计算核心使用__VEC_SCOPE__块，内部声明Reg::RegTensor<T>和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等 API

MemBase 参考实现（MulEpilogue）：

• 使用标准 AscendC API（AscendC::Mul）
• 三路 UB buffer：cLocal_（matmul 结果）+dLocal_（第二路输入）+cLocalTmp_（计算结果）
• StageNum = 2（UB 分区数）

§2 RegBase vs MemBase 选择

决策建议：优先 RegBase。仅当融合操作为单一 vector op 且已有对应AscendC::API 时选 MemBase。

设计细节请查阅：

• development/epilogue-dev-guide.md
• development/epilogue-membase-design.md
• development/epilogue-regbase-design.md

§3 接口契约

普通 C+V 与 MX C+V 的 Epilogue 类必须提供以下四个接口，与MatmulKernelFused/MxMatmulKernelFused兼容：

Grouped C+V 使用 context-based Epilogue，operator()签名不同：

__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 C+V。

兼容性矩阵

Params 结构体约定：

• RegBase：GM_ADDR extraInputAddr（额外输入）+GM_ADDR outputGmAddr（输出）
• MemBase（MulEpilogue）：GM_ADDR multiplierGmAddr（乘数 D）+GM_ADDR outputGmAddr（输出）

BlockShape / ProblemShape：

• BlockShape = AscendC::Shape<int64_t, int64_t, int64_t, int64_t>—{curM, curN, 1, 1}
• ProblemShape = AscendC::Shape<int64_t, int64_t, int64_t, int64_t>—{m, n, k, 1}

辅助方法：

• InitParams(Params)—__host_aicore__静态方法，用于 host 端参数透传

§4 UB 占用情况

RegBase 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 循环按行切分（rowsThisStage），MemBase 按元素数切分（curStageSize）。

§5 命名空间说明

Epilogue 类定义在全局作用域，不归属Kernel::或Block::namespace。Kernel 层通过模板参数直接引用 Epilogue 类名（如MatmulKernelFused<..., EpilogueFusionRegBase>）。

【免费下载链接】cannbot-skillsCANNBot 是面向 CANN 开发的用于提升开发效率的系列智能体，本仓库为其提供可复用的 Skills 模块。项目地址: https://gitcode.com/cann/cannbot-skills