RegBase Epilogue 设计【免费下载链接】cannbot-skillsCANNBot 是面向 CANN 开发的用于提升开发效率的系列智能体本仓库为其提供可复用的 Skills 模块。项目地址: https://gitcode.com/cann/cannbot-skills适用场景GELU / SwiGLU / LayerNorm / 多输入 scale / 多中间值公式链等复杂后融合。1. 适用场景RegBase 适用于公式链存在多个中间值。需要减少 UB 中间值写回。需要在寄存器内完成多步运算并最终再写回输出。存在多路额外输入如 per-channel scale、per-token scale、residual 等。典型样例GELU SwiGLU LayerNorm matmul * scaleN * scaleM activation2. 从公式提取 RegBase API设计流程先把用户公式拆成有序的计算链。判断每个步骤需要的输入、输出和中间值类型。到ascendc-regbase-best-practice查找对应Reg::API。对不能直接确认的 API再去ascendc-api-best-practices和本地头文件核验。以 GELU_tanh 为例acc - cast float x acc * scaleN * scaleM u sqrt(2/pi) * x * (1 0.044715 * x * x) u_clip clamp(u, -10, 10) tanh (1 - exp(-2*u_clip)) / (1 exp(-2*u_clip)) y 0.5 * x * (1 tanh) bf16_out cast(y)常见 API 对应步骤RegBase APIloadReg::LoadAligncastReg::Castadd/subReg::Add/Reg::Sub/Reg::AddsmulReg::Mul/Reg::MulsdivReg::DivexpReg::ExpclampReg::Maxs/Reg::MinsstoreReg::StoreAlign3. UB 空间分配RegBase 的核心原则尽量让中间值停留在RegTensorUB 只保留输入 staging 和输出 staging。典型布局Offset 0 ┌──────────────────────────────────────────┐ │ cLocal_ MatMul 结果区 │ L0C2UB 写入不可释放 ├──────────────────────────────────────────┤ │ extraBufA_ 额外输入 A如 per-channel│ ├──────────────────────────────────────────┤ │ extraBufB_ 额外输入 B如 per-token │ ├──────────────────────────────────────────┤ │ outBuf_ 输出 staging │ └──────────────────────────────────────────┘预算nAlignL0C ceil(baseN / (32/sizeof(L0CDataType))) * (32/sizeof(L0CDataType)) splitMRows ceil(baseM / GetTaskRation()) matmulAreaBytes splitMRows * nAlignL0C * sizeof(L0CDataType) remainBytes TOTAL_UB_SIZE - matmulAreaBytes extraBytes sum(each extra input staging bytes) outBytes stageRows * nAlignOut * sizeof(OutputType) extraBytes outBytes remainBytes关键matmulAreaBytes的行步长是nAlignL0CUB 中的对齐行宽不是L0C cube 边长 16。L0C cube 边长是 L0C 硬件分型粒度与 UB buffer 行步长无关。常见错误是将nAlignL0C写成固定值 16导致matmulAreaBytes被低估、后续 buffer 重叠到 cLocal_ 数据区。若 buffer 不够优先减少stageRows。再考虑调整baseM/baseN。4. 多 extra input 设计复杂公式链经常有多路额外输入。每一路都要回答三个问题它依赖 N 维还是 M 维它是 tile 级加载还是 stage 级加载它是否可以跨 stage 复用典型模式输入类型例子加载方式列依赖输入per-channel scale / biastile 级加载可跨 stage 复用行依赖输入per-token scale / residual row scalarstage 级加载offset 依赖stageM0全局标量配置epsilon / alpha / beta可直接常量化或通过标量参数传入多 extra input 的 buffer 设计要避免互相覆盖且要和同步顺序匹配。tile 级与 stage 级 input 的同步分离tile 级 input加载一次、跨 stage 只读复用和 stage 级 input每 stage 覆盖必须使用不同 eventID的V_MTE2反向依赖避免 tile 级 Wait 阻塞 stage 级 Set 造成假依赖tile 级 extraBufA_: V_MTE2(eventID0) / MTE2_V(eventID0) stage 级 extraBufB_: V_MTE2(eventID1) / MTE2_V(eventID1) outBuf_: MTE3_V(eventID0) / V_MTE3(eventID0)eventID 分离原则详见references/fundamentals/blaze-sync-patterns.md§5.3。5. Vector 计算伪代码RegBase 伪代码模板for (stageRowOffset 0; stageRowOffset localRows; stageRowOffset stageRows) { curStageRows min(stageRows, localRows - stageRowOffset); // 加载额外输入 DataCopyPad(extraBufA_, extraInputA[row/col offset], ...); DataCopyPad(extraBufB_, extraInputB[row/col offset], ...); __VEC_SCOPE__ { for each row in curStageRows: for each vector lane group: LoadAlign(cLocal_ - vregL0C) Cast(vregL0C - vregCompute) LoadAlign(extraBufA_ - vregExtraA) LoadAlign(extraBufB_ - vregExtraB) // 公式链 // [USER] Reg::Mul/Add/Exp/Div/Maxs/Mins Cast(vregCompute - vregOut) StoreAlign(vregOut - outBuf_) } DataCopyPad(outBuf_ - output GM) }要求nAlign必须 per-call 从blockShapeN计算。row-dependent 输入 offset 必须按 SplitM 修正。Golden 与 device 最好同式避免公式差异干扰定位。6. MTE2/V/MTE3 同步伪代码以下伪代码展示 tile 级与 stage 级 input 的完整同步结构eventID 分离// ---- tile 级 extra input加载一次跨 stage 只读复用 ---- WaitFlagV_MTE2(TILE_EVENT_ID); DataCopyPad(extraBufA_, extraInputAGm[tileN0], ...); SetFlagMTE2_V(TILE_EVENT_ID); WaitFlagMTE2_V(TILE_EVENT_ID); // ---- stage 循环 ---- for (stageOffset 0; stageOffset localRows; stageOffset stageRows) { // stage 级 extra input每 stage 覆盖 WaitFlagV_MTE2(STAGE_EVENT_ID); DataCopyPad(extraBufB_, extraInputBGm[stageM0], ...); SetFlagMTE2_V(STAGE_EVENT_ID); WaitFlagMTE2_V(STAGE_EVENT_ID); // 等上一轮 MTE3 读完 outBuf_ WaitFlagMTE3_V(ZERO_FLAG); __VEC_SCOPE__ { // RegTensor load / cast / compute / store } // 通知 MTE2 可覆盖 extraBufB_ SetFlagV_MTE2(STAGE_EVENT_ID); SetFlagV_MTE3(ZERO_FLAG); WaitFlagV_MTE3(ZERO_FLAG); DataCopyPad(outputGm[outputOffset], outBuf_, ...); SetFlagMTE3_V(ZERO_FLAG); } // 通知 MTE2 可覆盖 extraBufA_tile 级只读复用完毕 SetFlagV_MTE2(TILE_EVENT_ID);DataCopyPad stride 单位DataCopyExtParams的srcStride/dstStride单位因搬运方向不同方向srcStridedstStrideGM → UBbytes32 字节单位UB → GM32 字节单位bytes当 UB 行按nAlign对齐排布时nAlign ceil(N / (32/sizeof(T))) * (32/sizeof(T))保证nAlign * sizeof(T)是 32 的倍数UB 侧 stride 恒等于 0直接传0即可。详见references/fundamentals/blaze-sync-patterns.md§9.5。同步规则以references/fundamentals/blaze-sync-patterns.md为准重点检查正向依赖MTE2_V、V_MTE3反向依赖V_MTE2、MTE3_V首轮等待是否需要在Init中预发射尾轮是否需要收尾等待7. SplitM offset若启用DUAL_DST_SPLIT_MRegBase 必须按 SubBlock 计算每一路 row-dependent 输入和输出 offset。GM offset需要 sub-block 偏移halfM ceilDiv(blockShapeM, GetTaskRation()) localRows (oddM) ? (halfM - GetSubBlockIdx()) : halfM tileM0 gmOffset / N tileN0 gmOffset % N subM0 tileM0 GetSubBlockIdx() * halfM stageM0 subM0 stageRowOffset rowDependentInputOffset stageM0 outputOffset stageM0 * N tileN0所有依赖 M 维的输入都必须使用stageM0计算偏移。UB 读取不需要 sub-block 偏移CopyL0C2UBSplitMDUAL_DST_SPLIT_M会硬件自动将 L0C 的 M 行对半切分每个 AIV 从自己的 UB offset 0 开始读取半份数据。因此srcAddr cLocal_.GetPhyAddr() // 从 offset 0 开始 rowSrc srcAddr row * nAlign // 不加 GetSubBlockIdx() * halfM * nAlign操作是否需要 SubBlock 偏移公式UB 读取 cLocal_否cLocal_.GetPhyAddr() row * nAlignGM 读取 row-dependent input是stageM0 tileM0 GetSubBlockIdx() * halfM stageRowOffsetGM 写回 output是gmRowOffset subM0 * N tileN0常见错误从 GM offset 公式推断UB 读取也需要加GetSubBlockIdx() * halfM * nAlign偏移导致 V1 跳过自己的数据读到后续 buffer 的垃圾值。localRows0 边界场景当curM为奇数且halfM1时如curM1V1 的localRows halfM - 1 0。此时 V1 应 early return不需要做任何 Vector 计算或 GM 写回。CV 同步由 kernel 层MatmulKernelFused统一处理不会因 V1 return 而挂死。8. Golden / Device 公式一致性提示复杂公式链容易出现 “device 公式”和 “golden 公式” 不一致。建议设计阶段就明确 golden 与 device 是否同式。若 device 使用clip exp-ratio tanhgolden 最好也使用同式。若不能同式必须在设计中写清误差来源和阈值依据。9. 常见错误错误后果修复把中间值写回 UBUB 压力大stageRows 降低中间值尽量保留在RegTensor多路 extra input 共用同一 buffer 无同步保护偶发错数或覆盖分离 buffer 或严格配对事件只修 output offset不修 row-dependent input offset下半块 scale/bias 错位所有依赖 M 的输入都按stageM0修正UB 读取 cLocal_ 时加 sub-block 偏移V1 读到垃圾值精度完全错误SplitM 已硬件分片从 offset 0 读取matmulAreaBytes行步长用 16 而非 nAlignL0CUB buffer 重叠bias/scale 覆写 matmul 结果行步长 nAlignL0C不是 L0C cube 边长 16UB→GM DataCopyPadsrcStride传 bytestail 场景行步长错误UB 侧 stride 是 32B 单位nAlign 对齐时传 0tile 级与 stage 级 input 共用同一 eventIDtile 级 Wait 阻塞 stage 级 Set假依赖使用不同 eventID如 0 和 1缺少 Init 预发射或析构排空首轮 hang 或尾轮 flag 泄漏Init 预发射所有反向 SetFlag析构排空所有 WaitFlag未检查 Golden 与 device 同式精度问题定位困难设计阶段明确公式一致性【免费下载链接】cannbot-skillsCANNBot 是面向 CANN 开发的用于提升开发效率的系列智能体本仓库为其提供可复用的 Skills 模块。项目地址: https://gitcode.com/cann/cannbot-skills创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考