三亩地 三亩地SAN MU DI · CODE DIARY
ARTICLE DETAIL

日记详情

真实记录编程学习的某一天,欢迎挑你感兴趣的翻一翻。

CANN Catlass后处理组件

CANN Catlass后处理组件

Catlass 后处理(Epilogue)组件清单

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

导航:新架构文档已建立 → 设计时优先阅读 architecture/02-block-layer.md §2.3 BlockEpilogue 部分和 kernels/epilogue-patterns.md。本文档作为 catlass 内置组件索引补充。

若下表无所需运算:先走 custom-epilogue.md(设计:确认 catlass 是否已有 Tile → 无则先设计自定义 Tile),再在 codegen 中落盘op_kernel/custom_epilogue/

Epilogue 组件分为Tile 级(单 tile 内操作)和Block 级(协调多个 tile)两层,按功能分为以下类别。

Tile 级组件

类别可选组件说明
激活函数TileElemWiseGelu,TileElemWiseSiluGELU、SILU 激活;现成TileElemWiseTanh时在自定义 Tile 中用Exp/Muls/Adds/Div实现 (\tanh)(与 Silu 同双缓冲范式),或评估AscendC::Tanh与固定COMPUTE_LENGTH、小 (M\times N) 的运行期 UB 风险
逐元素运算TileElemWiseAdd,TileElemWiseMul,TileElemWiseMuls加法、乘法、标量乘
广播运算TileBroadcastAdd,TileBroadcastMul,TileBroadcastInplaceByColumn,TileBroadcastInplaceByRow广播加/乘、按列/行原地广播
数据搬运TileCopy(含CopyGMToUBCopyUBToGM变体)GM ↔ UB 搬运
类型转换TileCast精度转换
数据重排TileSwizzle内存布局重排
量化反量化TileRowBroadcastMul,TileBroadcastOneBlk,TileOneBlkColumnBroadcastMul,TilePerTokenDequant行广播乘、单块广播、列广播乘、per-token 反量化

Block 级组件

类别可选组件说明
通用后处理BlockEpilogue组合 Tile 组件的标准流水
单源逐元素BlockEpilogueElemwiseOneSource带一个额外输入(如加偏置)
无源逐元素BlockEpilogueElemwiseNoSource无额外输入的后处理
SoftmaxBlockEpilogueOnlineSoftmaxNoMask在线 Softmax
量化反量化BlockEpiloguePerTokenDequant,BlockEpilogueW4A4PerTokenPerChannelDequant,BlockEpilogueDequant反量化
量化反量化+激活BlockEpilogue(组合 dequant tiles +TileElemWiseGelu量化乘 + 反量化 + 激活融合
GEMM/GEMVBlockEpilogueGEMM,BlockEpilogueGEMVGEMM/GEMV 专用
Flash AttentionBlockEpilogueFARescaleO,BlockEpilogueFASoftmaxFA 后处理
MLABlockEpilogueMLARescaleO,BlockEpilogueMLASoftmaxMLA 后处理

组装模式

BlockEpilogue将多个 Tile 组件组合为后处理流水线,典型模式:

模式流水线典型用例
纯矩阵乘无 Epilogue(BlockEpilogue = void基础 Matmul
激活函数Copy入 → GELU/SILU → Copy出Matmul + GELU
加偏置Copy入 → Add(额外源) → Copy出Matmul + Bias
加偏置+激活Copy入 → Add → GELU → Copy出Matmul + Bias + GELU
广播+运算Copy入 → BroadcastAdd → Mul → Copy出融合后处理
量化反量化TileRowBroadcastMul → TileBroadcastOneBlk → TileOneBlkColumnBroadcastMul → TileCopyW8A8 量化乘 + per-token 反量化
量化反量化+激活TileRowBroadcastMul → TileBroadcastOneBlk → TileOneBlkColumnBroadcastMul → TileElemWiseGelu → TileCopyW8A8 量化乘 + 反量化 + GELU

选型原则

  1. 按需组合:只选需要的环节,不要全部引入
  2. Tile 组件按序排列:Copy入 → 计算 → Copy出
  3. Block 组件接收 Tile 列表BlockEpilogue<DispatchPolicy, InType, OutType, Tile...>
  4. 无后处理时设BlockEpilogue = void

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

创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考

← 返回列表