已合并
feat: add matmul_streamk and quant_batch_matmul_mxa8w4 samples #82
feat: add matmul_streamk and quant_batch_matmul_mxa8w4 samples #82
已合并
weiyan_hw创建于 9月23日
weiyan_hw成员
9月23日

当前PR是否有AI参与

是。本 PR 的样例代码编写与仓内适配由 AI 辅助完成(平台:CANNBot / opencode,模型:GLM)。两个算子的内核设计、tiling 策略与性能调优为人工完成,AI 参与了代码落地、编码规范整改、测试与文档整理,全部改动已经过 133 个 NPU 精度用例回归验证。

PR功能描述 / 为什么需要这个合入

新增两个矩阵乘样例,补齐 samples/matmul/ 家族的调度与量化覆盖:

1. matmul_streamk(Stream-K / DPSK 矩阵乘,A16W16) — samples/matmul/matmul/matmul_streamk.py

  • matmul_streamk():C[M,N] = A[M,K] @ B[K,N],float16/bfloat16,transpose_a=False、transpose_b=False
  • DP + SK 混合调度:DP tiles 各 AIC 核独立计算完整 K 维直写输出 GM;SK tiles 按 K 维 split 写入 workspace,AIV 做 fp32 累加 reduce;纯 DP 路径启用 slide window + 行反转提升 L2 命中率
  • host 侧 tiling 按平台容量(get_mem_size)自适应推导 tile / 缓冲深度 / L2 cache 策略
  • 性能对比 CANN 包内置 torch_npu matmul:100 个纯 SK 用例(小 MN 大 K)geomean 0.968,小 M+N 场景最高 1.22x
  • 补齐 samples/matmul/matmul/README.md 已预埋的 Stream-K 文档章节

2. quant_batch_matmul_mxa8w4(混合量化矩阵乘,MXFP8 激活 × MXFP4 权重) — samples/matmul/quant_matmul/quant_batch_matmul_mxa8w4.py

  • matmul_mix_quant():C[M,N] = (A·sa)[M,K] @ (B·sb)[N,K]^T + bias,A 为 fp8 e4m3 + E8M0 paired-scale,B 为 uint8 打包 fp4 e2m1 + E8M0 paired-scale,输出 float16
  • AIV 权重 prologue 将 fp4→fp8 做纯位置换转换(2^-6 因子由 fixpipe deq_scale=64 与 host 侧 bias/64 补偿),经 CrossCore 通道交接到 AIC MX matmul
  • K 为 32 的倍数且无上限(片上 K 状态全部窗化),M/N/K 尾块由引擎在搬运途中补零
  • 性能对比 CANN 内置 aclnnQuantMatmulV5:100 形状 geomean 1.170,76/100 领先,小形状(K ≤ 1024)最高 2.28x

附属更新:两个样例的 pytest 精度测试(82 + 51 用例)、figures/ 性能对比图、根 README 算子表与目录说明、CHANGELOG 条目。

该PR关联的issue

fixes #(暂无关联 issue)

希望检视人员了解

  • 两个样例基于 cannbotdsl 0.6.0 显式 Channel API(produce/consume),为本仓首批 0.6.0 风格样例;仓内存量样例仍为隐式风格,待后续整体迁移
  • codecheck 遗留说明:MixQuantKernel.__call__ 与 _MixQuantLauncher.run 保留 6 参数 GM 操作数列表——已实测验证 cannbotdsl 内核绑定仅接受 Tensor/标量载体,张量结构体打包会被 ST101_UNSUPPORTED_DYNAMIC_ARGUMENT 拒绝,该签名无法收敛到 6 参数以下(约束已注释在 launcher 上)

改动类型 / Change Type

测试信息 / Testing

本地环境:Ascend 950PR(NPU ARCH 3510)+ cannbotdsl 0.6.0 + CANN 9.2.0,命令:

pytest test/matmul/matmul/test_matmul_streamk.py -v                    # 82 用例
pytest test/matmul/quant_matmul/test_quant_batch_matmul_mxa8w4.py -v   # 51 用例

结果:133/133 全部通过(0 失败),覆盖 DP/SK 混合调度、非对齐尺寸、极端形状、尾块/多窗/深环、bias 场景与输入校验;ruff 检查与 OAT 合规检查通过。

检查清单 / Checklist

likedislike
Pull Request已成功合入, 合并人@CANN-robot
(感谢 weiyan_hw 的贡献)
Wweiyan_hw成员
9月23日 创建了 pull request,commit ba8011de
CANN-robot
CANN-robot成员
9月23日 评论:

Hi @weiyan_hw, welcome to submitting your first PR to cannbot-dsl!

PR Merge Steps

1. CLA Signing

If the current PR label includes cann-cla/yes, it means you have signed the CLA and can proceed to the next step. If the label includes cann-cla/no, please sign the CLA first. If you have any questions, please refer to the FAQ.

2. CI Check

Please comment /compile to trigger the CI pipeline check. If the CI run is successful, the PR will be tagged with ci-pipeline-passed and you can proceed to the next step. If the CI run fails, the PR will be tagged with ci-pipeline-failed, please check the CI logs to fix the issues in the PR. If you have any questions, please refer to the FAQ.

3. Code Review

After CI passes, please refer to the PR Approval Progress and proactively @ the committers in the table to review the code. After approval, committers will comment /lgtm and /approve. Once the lgtm and approved labels are successfully added, the PR will be merged automatically.

likedislike
CANN-robotCANN-robot成员
9月23日 添加了label:cann-cla/yes
CANN-robot
CANN-robot成员
9月23日 评论:

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.
You can self-configure the PR merge rules for this repository. For more details, please refer to Here.
For more, you also can visit HICANN.


PR Approval Progress

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

Module Approval Details

module lgtm status approve status
repo-cann/cannbot-dsl ✅ liujun2025, yang-jiwei (2/2) ✅ liujun2025, yang-jiwei (2/1)

💡 Tip:

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

CLA Signature Pass

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

likedislike
weiyan_hw成员
9月23日 评论:

compile

likedislike
此处折叠了158条消息 查看更多
yang-jiwei成员
27 天前 评论:

/lgtm
/approve

likedislike
CANN-robotCANN-robot成员
27 天前 添加了label:approved
liujun2025成员
27 天前 评论:

/lgtm
/approve

likedislike
CANN-robotCANN-robot成员
27 天前 添加了label:lgtm
CANN-robotCANN-robot成员
27 天前 合入了pull request