已合并
增加grouped_mx_A8w4_matmul模版 #843
UniquePG创建于 7月9日
增加grouped_mx_A8w4_matmul模版 #843
已合并
UniquePG创建于 7月9日
UniquePG
UniquePG
7月9日

描述

此 PR 新增了 Ascend950 平台上的 A8W4 grouped MX matrix multiplication 内核模板,支持分组矩阵乘法(多专家场景),其中 A 矩阵为 FP8(E4M3)量化、B 矩阵为 FP4(E2M1)量化带 MX scale。核心改动包括:新增 A8W4GroupedMxMatmul 内核类实现 AIC/AIV 双核协同的分组计算逻辑、新增 BlockPrologue 的 MxGroupedA8W4Prologue 特化负责 AIV 侧 B 矩阵反量化、为 MmadA8W4Mx 引入 L1_SCALE_FACTOR_K 参数支持跨 K 迭代的 scale 复用优化;同时修改example 59 A8W4MxMatmul相关kernel与host侧代码,使其共用block层代码。

性能测试数据如链接所示:https://dcnucy7udkr0.feishu.cn/wiki/Mpqrw4t0Mi31O4kAAl9ctM1Lnxf?sheet=s4xNTJ
本例较标杆性能提升10%

关联的Issue

文档更新

重构 block_mmad_mx_a8w4.hpp; block_mx_a8w4_prologue.hpp 优化AIV双缓冲逻辑;增加scale合并加载优化;重构原prologue的反量化逻辑,修改传入nZ格式B矩阵时的反量化逻辑;修改原nd格式下双缓冲等bug问题
重构 a8w4_mx.cpp; a8w4_mx_matmul.hpp 修改原example 59相关的kernel与host侧代码,共用block层代码

类型标签

likedislike
Pull Request已成功合入, 合并人@CANN-robot
(感谢 UniquePG 的贡献)
atomgit-bot
atomgit-bot
7月9日 评论:

变更摘要

此 PR 新增了 Ascend950 平台上的 A8W4 grouped MX matrix multiplication 内核模板,支持分组矩阵乘法(多专家场景),其中 A 矩阵为 FP8(E4M3)量化、B 矩阵为 FP4(E2M1)量化带 MX scale。核心改动包括:新增 A8W4GroupedMxMatmul 内核类实现 AIC/AIV 双核协同的分组计算逻辑、新增 BlockPrologueMxGroupedA8W4Prologue 特化负责 AIV 侧 B 矩阵反量化、为 MmadA8W4Mx 引入 L1_SCALE_FACTOR_K 参数支持跨 K 迭代的 scale 复用优化,同时提供完整的数据生成、精度测试和性能测试配套工具链。

主要改动

  • 新增 A8W4GroupedMxMatmul 内核模板:位于 include/catlass/gemm/kernel/weight_quant_a8w4_grouped_mx_matmul.hpp,实现 AIC 核上按 group_list 遍历各组执行 block MMAD 计算,AIV 核上同步执行 B 矩阵 prologue 反量化,通过 CrossCoreSetFlag/CrossCoreWaitFlag 实现双核流水同步。

  • 新增 BlockPrologue<MxGroupedA8W4Prologue> 特化:位于 include/catlass/gemm/block/block_mx_grouped_a8w4_prologue.hpp,负责将 B 矩阵从 packed FP4 格式(float4_e2m1x2_t)通过 AntiQuantCompute 解包并反量化为 FP8,支持 nZ/ND 两种输入 layout,含 UB 缓冲区的 cast 和 L1 搬移逻辑。

  • 重构 BlockMmadA8W4Mx 引入 scale 复用:在 include/catlass/gemm/block/block_mmad_mx_a8w4.hpp 中增加 L1_SCALE_FACTOR_K 模板参数,scale 数据改为每 L1_SCALE_FACTOR_K 次 K 迭代才加载一次;新增 l1MxScaleAListId/l1MxScaleBListId 独立管理 scale buffer 的流水 ID;修改 AIC/AIV 同步 flag 编号为 AIV_SYNC_AIC_FLAG + l1BListId 以支持多阶段流水。

  • 扩展 dispatch policy 定义:在 include/catlass/gemm/dispatch_policy.hpp 中新增 MxGroupedA8W4Prologue 策略结构体(含 L1B_STAGES),MmadA8W4Mx 模板参数顺序调整为插入 L1_SCALE_FACTOR_K_(默认16)。

  • 修复 PackedMxA8W4TileCopyTlaLayoutB 类型别名:在 include/catlass/gemm/tile/tile_copy.hpp 中将 LayoutBElementPrologueB_ 修正为 ElementB_,确保 layout 推导与反量化后的实际元素类型一致。

likedislike
atomgit-bot
atomgit-bot
7月9日 评论:

代码审查

Now let me write the closing summary.


审查关闭摘要

已审查全部 20 个变更文件。发现 12 个问题,按优先级分布:

优先级 数量 关键问题
P0 1 跨核同步标志位方向错误,AIC 不再等待 AIV 完成 B 数据反量化处理
P1 1 主机端 lenMxScaleA/lenMxScaleB 在 uint32_t 域溢出后赋值 size_t
P2 6 kernel 假保护类型转换、gen_data 维度错乱/范围退化为单点值、gen_data_nd 越界访问、run.sh 缺少 set -e、设备内存泄漏
P3 4 多余分号、README 编号错误、gen_data_nd 同源范围 bug、accur_test 去重逻辑未 strip 值

各文件审查结果:

文件 结果
CMakeLists.txt (example) 无问题(缺少末尾换行属格式噪音)
README.md P3: 引用错误示例编号
accur_test.py P3: 已完成用例去重值未 strip
csv/benchmark.csv 无问题
csv/generate_benchmark.py 无问题
csv/result.csv 无问题(数据文件,前导空格由 accur_test.py 的 strip 问题覆盖)
gen_data.py P2: build_random 缩放系数错误 + reshape 维度错乱
gen_data_nd.py P2: build_random 缩放系数错误 + 对角线赋值越界
perf_test.py 无问题
run.sh P2: 缺少 set -euo pipefail
weight_quant_a8w4_grouped_mx_matmul.cpp P1: 32位溢出 + P2: deviceGroupList 泄漏
examples/CMakeLists.txt 无问题
copy_gm_to_ub_tla.hpp 无问题(删除空行)
copy_ub_to_l1_tla.hpp 无问题(删除空行,末尾换行属格式噪音)
block_mmad.hpp 无问题(新增 include)
block_mmad_mx_a8w4.hpp P0: 跨核同步标志错误(最严重)
block_mx_grouped_a8w4_prologue.hpp P3: 多余分号
dispatch_policy.hpp 无问题
weight_quant_a8w4_grouped_mx_matmul.hpp P2: 假保护类型转换
tile_copy.hpp 无问题(LayoutB 类型修正为反量化后的 ElementB_)

总体风险判断:高风险

P0 的跨核同步缺陷可能导致非确定性的计算精度问题(数据竞争),应在合入前修复。P1 的整数溢出在当前典型参数下未触发,但不符合算子规范要求,需修正。其余 P2/P3 问题建议在后续迭代中修复。

类型 数量
🔴 阻塞 3
🟡 建议 6

⛔ 需要修改

likedislike
CANN-robotCANN-robot成员
7月9日 添加了label:stat/needs-squash
CANN-robotCANN-robot成员
7月9日 添加了label:cann-cla/no
CANN-robot
CANN-robot成员
7月9日 评论:

Thanks for your pull-request.
The full list of commands accepted by me can be found at here
You can get sig-info at here


PR Approval Progress

Congratulations! All modules have met the lgtm and approve requirements.

Module Approval Details

module lgtm status approve status
repo-cann/catlass weixin_42818618, 黄鑫 (2/2) weixin_42818618 (1/1)

💡 Tip:

  • Committer can comment /approve or /lgtm
  • Commenting /approve implies both code review (lgtm) and intent to merge (approve)

CLA Signature Pass

UniquePG, thanks for your pull request. All authors of the commits have signed the CLA. 👍

likedislike
此处折叠了193条消息 查看更多
CANN-robotCANN-robot成员
7月23日 添加了label:approved
huangxin361成员
7月24日 评论:

/lgtm

likedislike
CANN-robotCANN-robot成员
7月24日 添加了label:lgtm
CANN-robotCANN-robot成员
7月24日 合入了pull request
CANN-robot
CANN-robot成员
7月24日 评论:

Pull Request 已合并或已关闭。

If you want to solve this problem, you can click here to do it in the FAQs.

likedislike