feat: add matmul_streamk and quant_batch_matmul_mxa8w4 samples-cannbot-dsl-AtomGit

已合并

feat: add matmul_streamk and quant_batch_matmul_mxa8w4 samples #82

weiyan_hw创建于 9月23日从weiyan_hw/cannbot-dsl： feat/matmul-streamk-mxa8w4 合入到 cann/cannbot-dsl： master

feat: add matmul_streamk and quant_batch_matmul_mxa8w4 samples #82

已合并

weiyan_hw创建于 9月23日

从weiyan_hw/cannbot-dsl feat/matmul-streamk-mxa8w4合入到cann/cannbot-dsl master

共 11 个文件变更+2355-2

MCHANGELOG.md+4-0

	@@ -6,6 +6,10 @@
6	#### 新特性 New Features	6	#### 新特性 New Features
7	- 【batch_matmul】新增非量化批量矩阵乘算子 `batch_matmul()`：$C[c\_batch, M, N] = A[a\_batch, M, K] @ B[b\_batch, N, K]^T + bias$，batch 维按 numpy/torch.matmul 语义右对齐广播（size-1 维重复、多维同时广播、混合秩），支持 rank 2~6（双 2-D 输入自动升维并压回 2-D）；fp16/bf16/fp32（核内 fp32 累加），fp32 可选 HF32 快速模式（TF32 档）。转置由 stride 自动推导：canonical `transpose(-1,-2)` 视图（K-major 存储）零拷贝翻转到对应 kernel 路径，其余非连续布局拒绝；bias 经 BT 折入 init MMAD，支持 `[N]` 共享与 `[*c_batch, N]` 逐 batch。面向 NPU ARCH 3510（Ascend 950PR / Ascend 950DT），host tiling 对齐主线 `batch_mat_mul_v3`（ResetBaseDav3510 + CalL1TilingDefault + GetBaseK），性能对比 `torch.bmm/baddbmm`（msprof Task Duration）。测试覆盖 fp16/bf16/fp32/HF32、bias（shared/per-batch）、K-major 转置视图、多维 batch 广播（rank 2~6）、非对齐尾块、2-D 输入与不可广播拒绝共 18 个 NPU 用例。	7	- 【batch_matmul】新增非量化批量矩阵乘算子 `batch_matmul()`：$C[c\_batch, M, N] = A[a\_batch, M, K] @ B[b\_batch, N, K]^T + bias$，batch 维按 numpy/torch.matmul 语义右对齐广播（size-1 维重复、多维同时广播、混合秩），支持 rank 2~6（双 2-D 输入自动升维并压回 2-D）；fp16/bf16/fp32（核内 fp32 累加），fp32 可选 HF32 快速模式（TF32 档）。转置由 stride 自动推导：canonical `transpose(-1,-2)` 视图（K-major 存储）零拷贝翻转到对应 kernel 路径，其余非连续布局拒绝；bias 经 BT 折入 init MMAD，支持 `[N]` 共享与 `[*c_batch, N]` 逐 batch。面向 NPU ARCH 3510（Ascend 950PR / Ascend 950DT），host tiling 对齐主线 `batch_mat_mul_v3`（ResetBaseDav3510 + CalL1TilingDefault + GetBaseK），性能对比 `torch.bmm/baddbmm`（msprof Task Duration）。测试覆盖 fp16/bf16/fp32/HF32、bias（shared/per-batch）、K-major 转置视图、多维 batch 广播（rank 2~6）、非对齐尾块、2-D 输入与不可广播拒绝共 18 个 NPU 用例。
8		8
		9	+- 【matmul_streamk】新增 Stream-K（DPSK）矩阵乘算子 `matmul_streamk()`：$C[M,N] = A[M,K] \times B[K,N]$，采用 DP + SK 混合调度——DP tiles 各核独立计算完整 K 维并直写输出 GM，SK tiles 按 K 维 split 写入 workspace，AIV 从 workspace 做 fp32 累加 reduce；纯 DP 路径启用 slide window + 行反转提升 L2 命中率。A/B/C 支持 float16、bfloat16（核内仅支持 transpose_a=False、transpose_b=False），host 侧 tiling 推导 tile/缓冲深度/L2 cache 参数。面向 NPU ARCH 3510（Ascend 950PR / 950DT），5 组代表性 shape × 2 dtype 共 10 个精度用例，性能对比 CANN 包内置 torch_npu matmul（100 个纯 SK 用例）。
		10	+
		11	+- 【quant_batch_matmul_mxa8w4】新增混合精度量化矩阵乘算子 `matmul_mix_quant()`（MXA8W4）：$C[M,N] = (A \cdot sa)[M,K] \times (B \cdot sb)[N,K]^T + bias$，激活为 MXFP8（float8_e4m3fn + E8M0 paired-scale），权重为 MXFP4（uint8 打包 float4_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 尾块由引擎在搬运途中补零。面向 NPU ARCH 3510（Ascend 950PR / 950DT），11 个精度用例（5 组代表性形状 × 含 bias + 输入校验），性能对比 CANN 内置 `aclnnQuantMatmulV5`。
		12	+
9	### 【2026-09-21】	13	### 【2026-09-21】
10	#### 文档 Documentation	14	#### 文档 Documentation
11	- 【API 文档】扩充设备侧接口文档并重组 Kernel API 分类：Host API 新增平台信息（`get_platform_info`、`get_mem_size`）；Kernel API 划分为数据搬运、系统变量访问、同步与缓存控制三类，新增系统变量访问 9 个与同步与缓存控制 4 个接口页，原 `api/operations/` 下的接口页迁至 `api/kernel/`，侧边栏同步更新。	15	- 【API 文档】扩充设备侧接口文档并重组 Kernel API 分类：Host API 新增平台信息（`get_platform_info`、`get_mem_size`）；Kernel API 划分为数据搬运、系统变量访问、同步与缓存控制三类，新增系统变量访问 9 个与同步与缓存控制 4 个接口页，原 `api/operations/` 下的接口页迁至 `api/kernel/`，侧边栏同步更新。

MREADME.md+4-2

Afigures/matmul_streamk.png+0-0

C chen-shuai9月24日

建议左边柱子 dsl 右边柱子 built-in 顺序换一下

Afigures/quant_batch_matmul_mxa8w4.png+0-0

C chen-shuai9月24日

同理

Msamples/matmul/matmul/README.md+73-0

	@@ -65,3 +65,76 @@ pytest test/matmul/matmul/test_matmul.py -v
65	基于 CANNBot-DSL 实现的 matmul 算子与 CANN 包内置 matmul 算子基础模板在部分用例上的性能对比结果如下：	65	基于 CANNBot-DSL 实现的 matmul 算子与 CANN 包内置 matmul 算子基础模板在部分用例上的性能对比结果如下：
66		66
67	![matmul_perf_compare](../../../figures/matmul.png)	67	![matmul_perf_compare](../../../figures/matmul.png)
		68	+
		69	+## Stream-K 实现
		70	+
		71	+Stream-K 实现采用 DP + SK 混合调度：
		72	+
		73	+- DP tiles（能被核数整除的 tile）：各 AIC 核独立计算完整 K 维度，结果直接写回 output GM。纯 DP 路径启用 slide window 调度（WINDOW_LEN=4）和行反转，使相邻核共享 B 矩阵同一行，提升 L2 命中率。
		74	+- SK tiles（尾部不能整除的 tile）：按 K 维 split，多个 AIC 核分别计算部分 K 段，结果写入 workspace GM。
		75	+- AIV reduce：AIV 核从 workspace 读取各 K 段部分和，fp32 累加后 cast 回输入 dtype 写回 output GM。
		76	+
		77	+数据流为 GM → L1（MTE2, nd2nz, L2 cache 控制）→ L0A/L0B（MTE1）→ MMAD（M）→ L0C → GM（FIXPIPE）。AIV reduce 数据流为 workspace GM → UB（MTE3）→ add（V）→ cast → output GM（MTE3）。实现详见 `matmul_streamk.py`。
		78	+
		79	+\| 特性与约束 \| 说明 \|
		80	+\| :--------- \| :--- \|
		81	+\| 数据类型 \| float16、bfloat16（A/B/C 一致） \|
		82	+\| transpose \| 当前核内仅支持 transpose_a=False、transpose_b=False \|
		83	+\| 输入 \| 二维 contiguous 矩阵 \|
		84	+\| 多核并行 \| DP（Data Parallel）+ SK（Split-K）混合调度 \|
		85	+\| 调度策略 \| AIC 负责 DP/SK tile 计算，AIV 负责 workspace reduce \|
		86	+\| L2 cache \| 按矩阵复用情况自适应开关 \|
		87	+\| 滑动窗口 \| 纯 DP 路径启用 slide window + 行反转 \|
		88	+\| 支持架构 \| NPU ARCH 3510（Ascend 950PR / Ascend 950DT） \|
		89	+
		90	+### 快速开始
		91	+
		92	+```python
		93	+import torch
		94	+import torch_npu
		95	+from matmul_streamk import matmul_streamk
		96	+
		97	+M, K, N = 1024, 8192, 1024
		98	+dtype = torch.float16
		99	+
		100	+a = torch.randn(M, K, dtype=dtype).npu() # (M, K)
		101	+b = torch.randn(K, N, dtype=dtype).npu() # (K, N)
		102	+
		103	+c = matmul_streamk(a, b, transpose_a=False, transpose_b=False)
		104	+```
		105	+
		106	+`matmul_streamk()` 参数说明：
		107	+
		108	+\| 参数 \| shape \| dtype \| 说明 \|
		109	+\| :--- \| :---: \| :---: \| :--- \|
		110	+\| `a` \| (M, K) \| float16/bfloat16 \| 左矩阵 \|
		111	+\| `b` \| (K, N) \| float16/bfloat16 \| 右矩阵 \|
		112	+\| `transpose_a` \| - \| bool \| 左矩阵是否转置（当前核内仅支持 False） \|
		113	+\| `transpose_b` \| - \| bool \| 右矩阵是否转置（当前核内仅支持 False） \|
		114	+\| 返回值 `c` \| (M, N) \| float16/bfloat16 \| 矩阵乘结果，与输入同 dtype \|
		115	+
		116	+### 精度测试
		117	+
		118	+测试代码位于 `test/matmul/matmul/test_matmul_streamk.py`，使用 pytest 驱动，运行命令如下：
		119	+
		120	+```bash
		121	+pytest test/matmul/matmul/test_matmul_streamk.py -v
		122	+```
		123	+
		124	+包含 5 组代表性 shape × 2 dtype = 10 个用例，比对 CPU golden（fp32 累加），覆盖 DP、SK、DP+SK 混合调度、大 K 与极端形状。精度容差为相对误差 `< 1e-2 * max(\|ref_max\|, 1.0)`。
		125	+
		126	+### 性能数据
		127	+
		128	+基于 CANNBotDSL 实现的 `matmul_streamk` 算子与 CANN 包内置 torch_npu matmul 算子在 100 个纯 SK 用例上的性能对比结果如下：
		129	+
		130	+![matmul_streamk_perf_compare](../../../figures/matmul_streamk.png)
		131	+
		132	+测试条件：100 个纯 Stream-K shape（小 MN 大 K，MN tile 数 ≤ 16、K ≥ 8192），msprof 采集 Task Duration（每 shape 重复 10 次取 min），fp16 输入输出。
		133	+
		134	+\| 加速比区间 \| 用例数 \|
		135	+\| :--------- \| :----: \|
		136	+\| < 0.8 \| 4 \|
		137	+\| 0.8 ~ 0.9 \| 11 \|
		138	+\| ≥ 0.9 \| 85 \|
		139	+
		140	+DSL 在 42/100 个用例上快于 torch_npu，geomean 加速比 0.968。DSL 在小 M+N（如 M=37,N=32）场景优势显著（最高 1.22x），在 M 或 N=1 的退化形状上存在劣势。

Asamples/matmul/matmul/matmul_streamk.py+703-0

Msamples/matmul/quant_matmul/README.md+91-0

	@@ -205,3 +205,94 @@ pytest test/matmul/quant_matmul/test_quant_batch_matmul_hif8_tt.py -v
205	15 次取中位数；图中不包含编译、Launch、同步和输入构造时间。	205	15 次取中位数；图中不包含编译、Launch、同步和输入构造时间。
206		206
207	![quant_batch_matmul_hif8_tt_perf_compare](../../../figures/quant_batch_matmul_hif8_tt.png)	207	![quant_batch_matmul_hif8_tt_perf_compare](../../../figures/quant_batch_matmul_hif8_tt.png)
		208	+
		209	+# QuantBatchMatmul MXA8W4
		210	+
		211	+基于 CANNBot-DSL 实现的混合精度量化矩阵乘算子（MXA8W4）：激活为 MXFP8
		212	+（float8_e4m3fn + E8M0 Scale），权重为 MXFP4（float4_e2m1 + E8M0 Scale），
		213	+输出 float16，面向 Ascend NPU。
		214	+
		215	+## 算子介绍
		216	+
		217	+计算公式：
		218	+
		219	+$$
		220	+C[M,N] = (A[M,K] \cdot sa) @ (B[N,K] \cdot sb)^T + bias
		221	+$$
		222	+
		223	+其中 A/B 沿 K 轴每 32 个数据使用一个 E8M0 Scale（MX group size = 32）。
		224	+公开 Scale 接口将相邻两个 Scale 组成一个 pair，因此 Scale 的分组轴长度为
		225	+`ceil(K / 64)`，最后一维为 2。B 的 fp4 数据以 uint8 打包（每字节 2 个
		226	+e2m1 元素，偶数 K 在低 nibble）。
		227	+
		228	+\| 特性与约束 \| 说明 \|
		229	+\| :--- \| :--- \|
		230	+\| 数据类型 \| A 为 float8_e4m3fn，B 为打包 float4_e2m1（uint8 载体）；输出 float16 \|
		231	+\| Scale \| ScaleA/ScaleB 为 float8_e8m0fnu 三维 paired-scale Tensor \|
		232	+\| Shape \| M/N 任意；K 为 32 的倍数且无上限（片上 K 状态全部窗化） \|
		233	+\| 尾块 \| M/N/K 尾块由引擎在搬运途中补零；K 尾子块运行期缩窄（`k_extent`），读界不超过 `align64(K)` \|
		234	+\| 多核并行 \| 平衡连续 tile 分配，自适应核数（grid 上限 32） \|
		235	+\| 架构 \| 双核混合：AIV 权重转换 prologue + AIC MX matmul，CrossCore 通道交接 \|
		236	+
		237	+算子由 Host 侧 cost-model tiling 推导 baseM/baseN/baseK、L1 K 窗
		238	+（`k_l1 = step_k × base_k`，与 L0/mmad 块解耦）、环深与核数。device 侧
		239	+AIV 将 B 的 fp4 → fp8 做纯位置换转换（`[s e e m] → [s 0 0 e e m 0 0]`，
		240	+数值恒为 e2m1 的 2^-6，含次正规精确），经 UB 环形缓冲写入 CrossCore L1
		241	+通道；AIC 完成 A/B 的 MX 矩阵乘与 fixpipe 输出，转换的 2^-6 因子由
		242	+fixpipe `deq_scale=64` 与 host 侧 `bias/64` 补偿。数据流为
		243	+GM → L1（MTE2，nd2nz）→ L0A/L0B（MTE1，MX scale 联动加载）→ MMAD（MX）
		244	+→ L0C → GM（FIXPIPE）。实现详见 `quant_batch_matmul_mxa8w4.py`。
		245	+
		246	+## 快速开始
		247	+
		248	+```python
		249	+import torch
		250	+import torch_npu
		251	+
		252	+from quant_batch_matmul_mxa8w4 import matmul_mix_quant
		253	+
		254	+M, K, N = 256, 352, 256
		255	+
		256	+a = torch.randn(M, K).to(torch.float8_e4m3fn).npu() # (M, K)
		257	+codes = torch.randint(0, 16, (N, K), dtype=torch.uint8).npu()
		258	+b = (codes[:, 0::2] & 0xF) \| ((codes[:, 1::2] & 0xF) << 4) # (N, K // 2)
		259	+
		260	+# [M, ceil(K / 64), 2] 为 ScaleAND，[N, ceil(K / 64), 2] 为 ScaleBDN。
		261	+scale_a = torch.full((M, (K + 63) // 64, 2), 127, dtype=torch.uint8).npu()
		262	+scale_b = torch.full((N, (K + 63) // 64, 2), 127, dtype=torch.uint8).npu()
		263	+scale_a = scale_a.view(torch.float8_e8m0fnu)
		264	+scale_b = scale_b.view(torch.float8_e8m0fnu)
		265	+
		266	+c = matmul_mix_quant(a, b, scale_a, scale_b) # (M, N) float16
		267	+```
		268	+
		269	+`matmul_mix_quant()` 输入/输出说明：
		270	+
		271	+\| 参数 \| shape \| dtype \| 说明 \|
		272	+\| :--- \| :---: \| :---: \| :--- \|
		273	+\| `a` \| (M, K) \| float8_e4m3fn \| A，左矩阵 \|
		274	+\| `b` \| (N, K // 2) \| uint8 \| B，右矩阵（fp4 e2m1 打包，偶数 K 低 nibble） \|
		275	+\| `scale_a` \| (M, ceil(K / 64), 2) \| float8_e8m0fnu \| ScaleA（ScaleAND） \|
		276	+\| `scale_b` \| (N, ceil(K / 64), 2) \| float8_e8m0fnu \| ScaleB（ScaleBDN） \|
		277	+\| `bias` \| (N,) \| torch.float32 \| 可选，bias（内部按 2^-6 补偿） \|
		278	+\| 返回值 `c` \| (M, N) \| float16 \| 矩阵乘结果 \|
		279	+
		280	+## 精度测试
		281	+
		282	+测试代码位于 `test/matmul/quant_matmul/test_quant_batch_matmul_mxa8w4.py`，
		283	+11 个用例（5 组代表性形状 × 含 bias + 输入校验，覆盖尾块/部分尾窗/深环/退化形状），pytest 驱动：
		284	+
		285	+```bash
		286	+pytest test/matmul/quant_matmul/test_quant_batch_matmul_mxa8w4.py -v
		287	+```
		288	+
		289	+## 性能数据
		290	+
		291	+与 CANN 内置 `aclnnQuantMatmulV5`（QuantBatchMatmulV4 官方 kernel +
		292	+SWAT tiling solver）在 100 形状（方形/常规/瘦长/尾块 × M 扫描/宽 N/
		293	+奇 M,N/K 尾残差/大 K ≤ 16384/种子随机）上的 msprof task duration
		294	+对比（比值 = aclnn / 本实现，大于 1 表示本实现更快）：几何平均
		295	+1.170，76/100 形状领先，无落后超过 9% 的形状；小形状（K ≤ 1024）
		296	+最高 2.28x，尾块形状（K 非 256 倍数）1.07-1.38x。
		297	+
		298	+![quant_batch_matmul_mxa8w4_perf_compare](../../../figures/quant_batch_matmul_mxa8w4.png)

Asamples/matmul/quant_matmul/quant_batch_matmul_mxa8w4.py+1145-0

Atest/matmul/matmul/test_matmul_streamk.py+73-0

@@ -0,0 +1,73 @@
	1	+# Copyright (c) 2026 Huawei Technologies Co., Ltd.
	2	+# This program is free software, you can redistribute it and/or modify it under the terms and conditions of
	3	+# CANN Open Software License Agreement Version 2.0 (the "License").
	4	+# Please refer to the License for details. You may not use this file except in compliance with the License.
	5	+# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
	6	+# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
	7	+# See LICENSE in the root of the software repository for the full text of the License.
	8	+
	9	+"""Tests for the matmul_streamk kernel (DPSK, transpose_a=False, transpose_b=False).
	10	+
	11	+Formula: C[M,N] = A[M,K] @ B[K,N] (fp16/bf16 inputs, fp32 accumulator)
	12	+
	13	+Constraints honoured by the shape list:
	14	+ * Inputs are 2-D contiguous matrices.
	15	+ * Current kernel path only implements transpose_a=False, transpose_b=False.
	16	+"""
	17	+
	18	+from __future__ import annotations
	19	+
	20	+import os
	21	+
	22	+import pytest
	23	+
	24	+import torch
	25	+
	26	+from _samples_path import load_sample
	27	+
	28	+_sample = load_sample("matmul/matmul/matmul_streamk.py")
	29	+MatmulStreamK = _sample.MatmulStreamK
	30	+matmul_streamk = _sample.matmul_streamk
	31	+
	32	+
	33	+_TRANSPOSE_A = False
	34	+_TRANSPOSE_B = False
	35	+
	36	+
	37	+# 5 representative shapes (full 41-shape regression list kept locally):
	38	+# DP-dominant large, SK-dominant small MN, degenerate 1x1, mixed
	39	+# DP+SK with big K, and a tall-aspect grid.
	40	+_TEST_SHAPES = [
	41	+ pytest.param(1024, 1024, 8192, id="M1024-N1024-K8192"),
	42	+ pytest.param(37, 37, 8192, id="M37-N37-K8192"),
	43	+ pytest.param(1, 1, 8192, id="M1-N1-K8192"),
	44	+ pytest.param(300, 300, 16384, id="M300-N300-K16384"),
	45	+ pytest.param(2048, 300, 8192, id="M2048-N300-K8192"),
	46	+]
	47	+
	48	+
	49	+@pytest.mark.npu
	50	+@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16], ids=["fp16", "bf16"])
	51	+@pytest.mark.parametrize("m,n,k", _TEST_SHAPES)
	52	+def test_matmul_streamk_precision(m: int, n: int, k: int, dtype: torch.dtype) -> None:
	53	+ pytest.importorskip("torch_npu")
	54	+ if not os.environ.get("ASCEND_HOME_PATH"):
	55	+ pytest.skip("ASCEND_HOME_PATH not set")
	56	+ torch.npu.empty_cache()
	57	+ torch.npu.synchronize()
	58	+
	59	+ op = MatmulStreamK(m, n, k, dtype="fp16" if dtype == torch.float16 else "bf16")
	60	+ op.print_config()
	61	+ a = torch.randn(m, k, dtype=dtype).npu()
	62	+ b = torch.randn(k, n, dtype=dtype).npu()
	63	+ out = torch.zeros(m, n, dtype=dtype).npu()
	64	+ ws = torch.zeros(op.ws_rows, op.ws_cols, dtype=torch.float32).npu()
	65	+ op.run(out, ws, a, b)
	66	+ torch.npu.synchronize()
	67	+
	68	+ ref = a.cpu().float() @ b.cpu().float()
	69	+ err = (out.cpu().float() - ref).abs().max().item()
	70	+ ref_max = ref.abs().max().item()
	71	+ assert err < 1e-2 * max(abs(ref_max), 1.0), (
	72	+ f"DPSK {m}x{n}x{k}: err={err:.4e}, ref_max={ref_max:.4e}"
	73	+ )

Atest/matmul/quant_matmul/test_quant_batch_matmul_mxa8w4.py+243-0

@@ -0,0 +1,243 @@
	1	+# Copyright (c) 2026 Huawei Technologies Co., Ltd.
	2	+# This program is free software, you can redistribute it and/or modify it under the terms and conditions of
	3	+# CANN Open Software License Agreement Version 2.0 (the "License").
	4	+# Please refer to the License for details. You may not use this file except in compliance with the License.
	5	+# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
	6	+# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
	7	+# See LICENSE in the root of the software repository for the full text of the License.
	8	+
	9	+"""Precision tests for the MXA8W4 mix-quantized matmul kernel.
	10	+
	11	+Formula:
	12	+ C[M,N] = (A[M,K] * sa) @ (B[N,K] * sb)^T + bias
	13	+
	14	+ A: fp8 e4m3 with per-32-group E8M0 scales (MXFP8)
	15	+ B: packed fp4 e2m1 with per-32-group E8M0 scales (MXFP4)
	16	+ C: fp16
	17	+
	18	+Golden path mirrors the kernel: bit-exact fp4/fp8 decode, fp32 dequant
	19	+and accumulation, fp16 output.
	20	+
	21	+A/B/C are never padded on GM (engines pad tails in flight); only the
	22	+E8M0 scale and bias descriptors are aligned to the tile grid. Both AIVs
	23	+convert B in parallel (row halves of the CrossCore L1 ring slot). K
	24	+domain: multiple of 32, otherwise unbounded — aligned and mixed-tail K
	25	+alike run the window ring; K tails need no padding.
	26	+"""
	27	+
	28	+from __future__ import annotations
	29	+
	30	+import dataclasses
	31	+import logging
	32	+import os
	33	+import sys
	34	+
	35	+import numpy as np
	36	+import pytest
	37	+import torch
	38	+from ml_dtypes import float4_e2m1fn, float8_e4m3fn
	39	+
	40	+sys.path.insert(
	41	+ 0,
	42	+ os.path.join(
	43	+ os.path.dirname(__file__),
	44	+ "..",
	45	+ "..",
	46	+ "..",
	47	+ "samples",
	48	+ "matmul",
	49	+ "quant_matmul",
	50	+ ),
	51	+)
	52	+
	53	+from quant_batch_matmul_mxa8w4 import matmul_mix_quant
	54	+
	55	+
	56	+# ---------------------------------------------------------------------------
	57	+# Test cases
	58	+# ---------------------------------------------------------------------------
	59	+
	60	+# 5 representative shapes (full 25-shape precision matrix kept locally):
	61	+# degenerate minimal, aligned multi-tile, mixed M/N/K tails, a partial
	62	+# last L1 window with bias, and a deep-window big grid.
	63	+_TEST_SHAPES = [
	64	+ pytest.param(1, 32, 1, id="M1-K32-N1-minimal"),
	65	+ pytest.param(256, 256, 256, id="M256-K256-N256-multiTile"),
	66	+ pytest.param(256, 736, 96, id="M256-K736-N96-mixedTail3"),
	67	+ pytest.param(2048, 2176, 2048, id="M2048-K2176-N2048-b1tail"),
	68	+ pytest.param(2048, 4608, 2048, id="M2048-K4608-N2048-kw128s36"),
	69	+]
	70	+
	71	+
	72	+# ---------------------------------------------------------------------------
	73	+# Golden kit
	74	+# ---------------------------------------------------------------------------
	75	+
	76	+
	77	+def _align(value: int, alignment: int) -> int:
	78	+ return (value + alignment - 1) // alignment * alignment
	79	+
	80	+
	81	+@dataclasses.dataclass
	82	+class Mxa8w4Inputs:
	83	+ """Deterministic MXA8W4 golden inputs (e4m3 A, packed fp4 B, E8M0 scales, optional bias)."""
	84	+
	85	+ a: np.ndarray
	86	+ codes: np.ndarray
	87	+ packed: np.ndarray
	88	+ sa: np.ndarray
	89	+ sb: np.ndarray
	90	+ bias: np.ndarray \| None
	91	+
	92	+
	93	+def make_inputs(m: int, k: int, n: int, *, bias: bool, seed: int = 42) -> Mxa8w4Inputs:
	94	+ rng = np.random.default_rng(seed)
	95	+ # activations: e4m3-exact values
	96	+ a_vals = rng.choice(
	97	+ np.array([0.5, 1.0, -1.0, 2.0, -0.5, 1.5, -2.0], dtype=np.float32),
	98	+ size=(m, k),
	99	+ )
	100	+ a = a_vals.astype(float8_e4m3fn)
	101	+ # weights: random fp4 codes (all 16 codes exercised)
	102	+ codes = rng.integers(0, 16, size=(n, k), dtype=np.uint8)
	103	+ packed = (codes[:, 0::2] & 0x0F) \| ((codes[:, 1::2] & 0x0F) << 4)
	104	+ # scales: E8M0 exponents in [-3, 3]
	105	+ scale_cols = _align(k, 64) // 32
	106	+ sa = (rng.integers(-3, 4, size=(m, scale_cols)) + 127).astype(np.uint8)
	107	+ sb = (rng.integers(-3, 4, size=(n, scale_cols)) + 127).astype(np.uint8)
	108	+ bias_np = (
	109	+ rng.choice(np.array([0.5, -1.0, 2.0, -0.25], dtype=np.float32), size=(n,))
	110	+ if bias
	111	+ else None
	112	+ )
	113	+ return Mxa8w4Inputs(a=a, codes=codes, packed=packed, sa=sa, sb=sb, bias=bias_np)
	114	+
	115	+
	116	+def golden(inputs: Mxa8w4Inputs) -> np.ndarray:
	117	+ a = inputs.a
	118	+ codes = inputs.codes
	119	+ _, k = a.shape
	120	+ a_deq = a.astype(np.float32)
	121	+ sa_exp = inputs.sa.astype(np.int32) - 127
	122	+ sa_deq = np.repeat(np.power(2.0, sa_exp.astype(np.float32)), 32, axis=1)[:, :k]
	123	+ b_deq = codes.view(float4_e2m1fn).astype(np.float32)
	124	+ sb_exp = inputs.sb.astype(np.int32) - 127
	125	+ sb_deq = np.repeat(np.power(2.0, sb_exp.astype(np.float32)), 32, axis=1)[:, :k]
	126	+ c = (a_deq * sa_deq) @ (b_deq * sb_deq).T
	127	+ if inputs.bias is not None:
	128	+ c = c + inputs.bias[None, :]
	129	+ return c.astype(np.float16)
	130	+
	131	+
	132	+def assert_isclose(actual: np.ndarray, golden_np: np.ndarray) -> None:
	133	+ actual = actual.astype(np.float32)
	134	+ gold = golden_np.astype(np.float32)
	135	+ # NaN guard FIRST: rel-error comparisons are NaN-blind (NaN > tol
	136	+ # is False), so a NaN-producing kernel can pass the tolerance check
	137	+ err = np.abs(actual - gold)
	138	+ nan_mask = np.isnan(actual) \| np.isinf(actual)
	139	+ if nan_mask.any():
	140	+ idx = np.argwhere(nan_mask)[:5]
	141	+ details = "\n".join(
	142	+ f" [{i},{j}] actual={actual[i, j]} golden={gold[i, j]:.4f}" for i, j in idx
	143	+ )
	144	+ raise AssertionError(
	145	+ f"matmul_mix_quant produced {nan_mask.sum()} NaN/Inf elements "
	146	+ f"({nan_mask.mean() * 100:.3f}%)\n{details}"
	147	+ )
	148	+ rel = err / np.maximum(np.abs(gold), 1.0)
	149	+ bad = rel > 2e-2
	150	+ if bad.mean() > 1e-3:
	151	+ idx = np.argwhere(bad)[:5]
	152	+ details = "\n".join(
	153	+ f" [{i},{j}] actual={actual[i, j]:.4f} golden={gold[i, j]:.4f}"
	154	+ for i, j in idx
	155	+ )
	156	+ raise AssertionError(
	157	+ f"matmul_mix_quant mismatch: bad_ratio={bad.mean():.6f}, "
	158	+ f"max_abs_err={err.max():.4e}\n{details}"
	159	+ )
	160	+
	161	+
	162	+# ---------------------------------------------------------------------------
	163	+# NPU accuracy tests
	164	+# ---------------------------------------------------------------------------
	165	+
	166	+
	167	+@pytest.mark.npu
	168	+@pytest.mark.parametrize("m,k,n", _TEST_SHAPES)
	169	+@pytest.mark.parametrize("bias", [False, True], ids=["nobias", "bias"])
	170	+def test_quant_batch_matmul_mxa8w4(m: int, k: int, n: int, bias: bool) -> None:
	171	+ pytest.importorskip("torch_npu")
	172	+ if not os.environ.get("ASCEND_HOME_PATH"):
	173	+ pytest.skip("ASCEND_HOME_PATH not set")
	174	+ if torch.npu.device_count() == 0:
	175	+ pytest.skip("no NPU device is available")
	176	+ if "Ascend950" not in torch.npu.get_device_name(0):
	177	+ pytest.skip("MXA8W4 MX matmul requires Ascend950")
	178	+
	179	+ inputs = make_inputs(m, k, n, bias=bias)
	180	+ gold = golden(inputs)
	181	+
	182	+ gm_a = torch.from_numpy(inputs.a.view(np.uint8)).npu().view(torch.float8_e4m3fn)
	183	+ gm_b = torch.from_numpy(np.ascontiguousarray(inputs.packed)).npu()
	184	+ gm_sa = torch.from_numpy(inputs.sa.copy()).npu()
	185	+ gm_sb = torch.from_numpy(inputs.sb.copy()).npu()
	186	+ gm_bias = (
	187	+ torch.from_numpy(inputs.bias).npu()
	188	+ if inputs.bias is not None
	189	+ else torch.zeros(n, dtype=torch.float32, device="npu")
	190	+ )
	191	+
	192	+ c = matmul_mix_quant(gm_a, gm_b, gm_sa, gm_sb, gm_bias)
	193	+ torch.npu.synchronize()
	194	+
	195	+ actual = c.cpu().numpy()
	196	+ assert_isclose(actual, gold)
	197	+ max_err = float(np.abs(actual.astype(np.float32) - gold.astype(np.float32)).max())
	198	+ logging.info(
	199	+ "matmul_mix_quant (M=%s,K=%s,N=%s,bias=%s) max\|err\|=%.4e finished!",
	200	+ m,
	201	+ k,
	202	+ n,
	203	+ bias,
	204	+ max_err,
	205	+ )
	206	+
	207	+
	208	+@pytest.mark.npu
	209	+def test_input_validation() -> None:
	210	+ pytest.importorskip("torch_npu")
	211	+ if not os.environ.get("ASCEND_HOME_PATH"):
	212	+ pytest.skip("ASCEND_HOME_PATH not set")
	213	+ if torch.npu.device_count() == 0:
	214	+ pytest.skip("no NPU device is available")
	215	+ if "Ascend950" not in torch.npu.get_device_name(0):
	216	+ pytest.skip("MXA8W4 MX matmul requires Ascend950")
	217	+
	218	+ m, k, n = 64, 128, 64
	219	+ inputs = make_inputs(m, k, n, bias=False)
	220	+ gm_a = torch.from_numpy(inputs.a.view(np.uint8)).npu().view(torch.float8_e4m3fn)
	221	+ gm_b = torch.from_numpy(np.ascontiguousarray(inputs.packed)).npu()
	222	+ gm_sa = torch.from_numpy(inputs.sa.copy()).npu()
	223	+ gm_sb = torch.from_numpy(inputs.sb.copy()).npu()
	224	+ gm_bias = torch.zeros(n, dtype=torch.float32, device="npu")
	225	+
	226	+ # K not a multiple of 32 (MX group size)
	227	+ bad_b = torch.zeros(n, 33, dtype=torch.uint8, device="npu")
	228	+ with pytest.raises(ValueError):
	229	+ matmul_mix_quant(gm_a, bad_b, gm_sa, gm_sb, gm_bias)
	230	+
	231	+ # K is unbounded (windowed on-chip state; tails host-padded) —
	232	+ # mixed-tail and aligned large K alike run the window ring; no
	233	+ # NotImplementedError caps remain
	234	+
	235	+ # wrong scale width
	236	+ bad_sa = torch.zeros(m, 3, dtype=torch.uint8, device="npu")
	237	+ with pytest.raises(ValueError):
	238	+ matmul_mix_quant(gm_a, gm_b, bad_sa, gm_sb, gm_bias)
	239	+
	240	+ # wrong bias length
	241	+ bad_bias = torch.zeros(n + 1, dtype=torch.float32, device="npu")
	242	+ with pytest.raises(ValueError):
	243	+ matmul_mix_quant(gm_a, gm_b, gm_sa, gm_sb, bad_bias)

Mtest/test_config.yaml+19-0

	@@ -52,6 +52,25 @@ flash_kda:
52	options:	52	options:
53	- flash_kda	53	- flash_kda
54		54
		55	+# 文件级路由：matmul/matmul 与 matmul/quant_matmul 目录内混有尚未适配
		56	+# cannbotdsl 0.6.0 接口的存量样例，新增样例按精确文件路径触发各自的
		57	+# 测试文件，避免连带触发同目录存量样例的用例。
		58	+matmul_streamk:
		59	+ module: True
		60	+ src:
		61	+ - samples/matmul/matmul/matmul_streamk.py
		62	+ - test/matmul/matmul/test_matmul_streamk.py
		63	+ options:
		64	+ - matmul/matmul/test_matmul_streamk.py
		65	+
		66	+quant_batch_matmul_mxa8w4:
		67	+ module: True
		68	+ src:
		69	+ - samples/matmul/quant_matmul/quant_batch_matmul_mxa8w4.py
		70	+ - test/matmul/quant_matmul/test_quant_batch_matmul_mxa8w4.py
		71	+ options:
		72	+ - matmul/quant_matmul/test_quant_batch_matmul_mxa8w4.py
		73	+
55	matmul:	74	matmul:
56	module: True	75	module: True
57	src:	76	src:

	@@ -28,8 +28,10 @@ CANNBot 是 [CANN](https://hiascend.com/software/cann) 社区的 Infra 智能体
28	\| flash_kda_metadata \| FlashKDA 调度 metadata 生成（AICPU） \| [samples/flash_kda_metadata](samples/flash_kda_metadata) \|	28	\| flash_kda_metadata \| FlashKDA 调度 metadata 生成（AICPU） \| [samples/flash_kda_metadata](samples/flash_kda_metadata) \|
29	\| fused_recurrent_kda_snapshot \| Kimi Delta Attention decode（1～8 token，状态快照） \| [samples/fused_recurrent_kda_snapshot](samples/fused_recurrent_kda_snapshot) \|	29	\| fused_recurrent_kda_snapshot \| Kimi Delta Attention decode（1～8 token，状态快照） \| [samples/fused_recurrent_kda_snapshot](samples/fused_recurrent_kda_snapshot) \|
30	\| matmul \| $C[M,N] = A[M,K] @ B[N,K]^T$ \| [samples/matmul/matmul](samples/matmul/matmul) \|	30	\| matmul \| $C[M,N] = A[M,K] @ B[N,K]^T$ \| [samples/matmul/matmul](samples/matmul/matmul) \|
		31	+\| matmul_streamk \| $C[M,N] = A[M,K] @ B[K,N]$，DP + SK 混合调度 \| [samples/matmul/matmul](samples/matmul/matmul) \|
31	\| batch_matmul \| $C[c\_batch, M, N] = A[a\_batch, M, K] @ B[b\_batch, N, K]^T$ \| [samples/matmul/batch_matmul](samples/matmul/batch_matmul) \|	32	\| batch_matmul \| $C[c\_batch, M, N] = A[a\_batch, M, K] @ B[b\_batch, N, K]^T$ \| [samples/matmul/batch_matmul](samples/matmul/batch_matmul) \|
32	\| quant_matmul \| $C[M,N] = Dequant(A)[M,K] @ Dequant(B)[N,K]^T$ \| [samples/matmul/quant_matmul](samples/matmul/quant_matmul) \|	33	\| quant_matmul \| $C[M,N] = Dequant(A)[M,K] @ Dequant(B)[N,K]^T$ \| [samples/matmul/quant_matmul](samples/matmul/quant_matmul) \|
		34	+\| quant_batch_matmul_mxa8w4 \| $C[M,N] = (A \cdot sa)[M,K] @ (B \cdot sb)[N,K]^T + bias$，MXFP8 激活 × MXFP4 权重 \| [samples/matmul/quant_matmul](samples/matmul/quant_matmul) \|
33	\| grouped_matmul \| $y_i[m_i,n_i] = x_i[m_i,k_i] \times weight_i[k_i,n_i]$ \| [samples/grouped_matmul](samples/grouped_matmul) \|	35	\| grouped_matmul \| $y_i[m_i,n_i] = x_i[m_i,k_i] \times weight_i[k_i,n_i]$ \| [samples/grouped_matmul](samples/grouped_matmul) \|
34	\| pointnet_sa \| $\text{feat}[K, D_{out}] = \max_{j} \text{MLP}(\text{points}[K, j, D_{in}])$ \| [samples/pointnet_sa](samples/pointnet_sa) \|	36	\| pointnet_sa \| $\text{feat}[K, D_{out}] = \max_{j} \text{MLP}(\text{points}[K, j, D_{in}])$ \| [samples/pointnet_sa](samples/pointnet_sa) \|
35	\| rms_norm \| $y = x \cdot rstd \cdot \gamma$ \| [samples/rms_norm](samples/rms_norm) \|	37	\| rms_norm \| $y = x \cdot rstd \cdot \gamma$ \| [samples/rms_norm](samples/rms_norm) \|
	@@ -45,9 +47,9 @@ CANNBot 是 [CANN](https://hiascend.com/software/cann) 社区的 Infra 智能体
45	│ ├── flash_kda_metadata/ # FlashKDA 调度 metadata（AICPU）	47	│ ├── flash_kda_metadata/ # FlashKDA 调度 metadata（AICPU）
46	│ ├── fused_recurrent_kda_snapshot/ # KDA decode 状态快照	48	│ ├── fused_recurrent_kda_snapshot/ # KDA decode 状态快照
47	│ ├── matmul/ # 矩阵乘	49	│ ├── matmul/ # 矩阵乘
48	-│ │ ├── matmul/ # 非量化矩阵乘	50	+│ │ ├── matmul/ # 非量化矩阵乘（基础调度 + Stream-K）
49	│ │ ├── batch_matmul/ # 非量化批量矩阵乘（batch 维广播）	51	│ │ ├── batch_matmul/ # 非量化批量矩阵乘（batch 维广播）
50	-│ │ └── quant_matmul/ # MXFP8/MXFP4 全量化矩阵乘	52	+│ │ └── quant_matmul/ # 量化矩阵乘（MXFP8/MXFP4、MXA8W4、HiFloat8 TT）
51	│ ├── grouped_matmul/ # 非量化分组矩阵乘	53	│ ├── grouped_matmul/ # 非量化分组矩阵乘
52	│ ├── pointnet_sa/ # PointNet Set Abstraction	54	│ ├── pointnet_sa/ # PointNet Set Abstraction
53	│ ├── rms_norm/ # RmsNorm 归一化	55	│ ├── rms_norm/ # RmsNorm 归一化