Pull Request已成功合入, 合并人@CANN-robot
(感谢 UniquePG 的贡献)变更摘要
此 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 复用优化,同时提供完整的数据生成、精度测试和性能测试配套工具链。
主要改动
-
新增
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)。 -
修复
PackedMxA8W4TileCopyTla的LayoutB类型别名:在include/catlass/gemm/tile/tile_copy.hpp中将LayoutB从ElementPrologueB_修正为ElementB_,确保 layout 推导与反量化后的实际元素类型一致。


代码审查
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 |
⛔ 需要修改


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
/approveor/lgtm- Commenting
/approveimplies 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. 👍


/lgtm


Pull Request 已合并或已关闭。
If you want to solve this problem, you can click here to do it in the FAQs.


描述
此 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层代码类型标签