Closed
[Requirement|需求建议]: Blaze 新增 SiLU + MXFP8 量化 epilogue #287
ddsszAdd on  Oct 8Closed at  16 days ago
ddssz
ddsszCollaborators
Oct 8 Create

一、背景信息(必填)

quant_matmul_activation_quant 需要新增 SiLU 融合后处理:量化矩阵乘及可选 bias 得到 FP32 矩阵后,在完整 M×N 上执行 SiLU(x) = x × sigmoid(x),随后执行 MXFP8 量化。需要在 Blaze 中提供可复用的 epilogue,供 ops-nn 的 ND 和 WeightNZ 路径接入。

二、价值/作用(必填)

复用现有 QBMM 的 Cube 主循环和 MX 量化组件,将 SiLU 与输出量化融合到 AIV 后处理,减少独立激活和量化算子的中间数据搬运;与现有 SwiGLU 的 dtype、转置及量化模式支持保持一致,SiLU 输出保留完整 N。

三、设计方案(必填)

3.1 使能方式

由 ops-nn 的 quant_matmul_activation_quant 通过 ACLNN ND/WeightNZ 接口及 PyTorch 包装入口选择 activationType/activation_type="silu"。

3.2 总体设计
3.2.1 算子支持的数据类型

Epilogue 输入为 FP32 UB 矩阵,SiLU 结果以 RINT 舍入为 BF16,再输出 FP8 E4M3FN 或 E5M2;scale 使用 E8M0。矩阵乘的 A/B dtype、ND/NZ 布局与转置组合由上层校验。

3.2.2 Host 侧设计

沿用 QBMM MX tiling,由 ops-nn 传入输出地址、baseM/baseN 和 scaleAlg。SiLU 保留完整 N,并适配单 batch、多 batch、A 全载和非全载分发。

3.2.3 Kernel 侧设计

新增 BlockEpilogueSiluMxQuant,使用 FP32 Exp/Div/Mul 计算 SiLU,随后 BF16 RINT,并复用 Tile::MxQuant 完成 OCP/BLAS MXFP8 量化。处理不足 32 元素的 N 尾块、M=1 的无效 AIV 分工,以及配对 E8M0 scale 的补齐字节,避免越界写入及无效值参与归约。

3.3 支持硬件

Ascend 950PR/950DT 系列,arch35。

3.4 算子约束限制

  • 仅支持 MXFP8 和 RINT,scaleAlg=0(OCP)或 1(BLAS)。
  • 输出保持 [..., M, N],scale 为 [..., M, ceil(N/64), 2]。
  • 集成到 quant_matmul_activation_quant 时,ND 支持四种 FP8 A/B 组合及四种转置;WeightNZ 的 B 为 E4M3FN、A 不转置,输出 dtype 与 A 一致。
  • 本需求不扩展 FP4 或其他舍入方式。

相关实现提交:smdbha/ops-tensor@a22eb2ccb1c243c0fb3b3806c1aaeda9af055b31。ops-nn 同步增加接口、tiling、kernel 分发、CSV 测试和文档。

likedislike
ddsszddsszCollaborators
Oct 8 Added label: requirement
ddsszddsszCollaborators
Oct 8 Associated pull request
ddsszddsszCollaborators
Oct 8 Associated pull request
ddsszddsszCollaborators
Oct 8 Assign smdbha as responsible person
CANN-robotCANN-robotCollaborators
16 days ago Close issue
CANN-robotCANN-robotCollaborators
16 days ago Added label: resolved