Star497
281
代码
代码
Issues73
Pull Requests50
Actions
项目讨论
Wiki
分析
Star497
281
已合并
增加grouped_mx_A8w4_matmul模版 #843
UniquePG创建于 7月9日
增加grouped_mx_A8w4_matmul模版 #843
已合并
UniquePG创建于 7月9日
从UniquePG/catlassnd_change合入到cann/catlassmaster
共 30 个文件变更+2783-361
Mexamples/59_ascend950_a8w4_mx_matmul/a8w4_mx.cpp+11-5
@@ -161,7 +161,13 @@ static void Run(const Options& options)
161 161 
162 // archtag uniflag162 // archtag uniflag
163 using ArchTag = Arch::Ascend950;163 using ArchTag = Arch::Ascend950;
164- constexpr bool enableUnitFlag = false;164+ constexpr bool enableUnitFlag = true;
165+ static constexpr uint32_t L1_SCALE_FACTOR_K = 16;
166+ static constexpr uint32_t L1A_STAGES = 2;
167+ static constexpr uint32_t L1B_STAGES = 2;
168+ static constexpr uint32_t L0A_STAGES = 2;
169+ static constexpr uint32_t L0B_STAGES = 2;
170+ static constexpr uint32_t L0C_STAGES = 1;
165 171 
166 // shape & type172 // shape & type
167 using L1TileShape = Shape<Int<128>, Int<128>, Int<128>>;173 using L1TileShape = Shape<Int<128>, Int<128>, Int<128>>;
@@ -170,8 +176,9 @@ static void Run(const Options& options)
170 using PrologueDstType = Gemm::GemmType<ElementB, LayoutB>;176 using PrologueDstType = Gemm::GemmType<ElementB, LayoutB>;
171 177 
172 // DispatchPolicy178 // DispatchPolicy
173- using DispatchPolicyMmad = Gemm::MmadA8W4Mx<ArchTag, enableUnitFlag>;179+ using DispatchPolicyMmad = Gemm::MmadA8W4Mx<
174- using DispatchPolicyPrologue = Gemm::MxA8W4Prologue<ArchTag>;180+ ArchTag, enableUnitFlag, false, L1_SCALE_FACTOR_K, L0C_STAGES, L1A_STAGES, L1B_STAGES, L0A_STAGES, L0B_STAGES>;
181+ using DispatchPolicyPrologue = Gemm::MxA8W4Prologue<ArchTag, L1B_STAGES>;
175 182 
176 // layout (tla)183 // layout (tla)
177 auto layoutA = tla::MakeLayout<ElementA, LayoutA>(m, k);184 auto layoutA = tla::MakeLayout<ElementA, LayoutA>(m, k);
@@ -188,8 +195,7 @@ static void Run(const Options& options)
188 195 
189 // BlockMmad196 // BlockMmad
190 using BlockMmad = Gemm::Block::BlockMmadA8W4Mx<197 using BlockMmad = Gemm::Block::BlockMmadA8W4Mx<
191- DispatchPolicyMmad, L1TileShape, L0TileShape, ElementA, ElementB, ElementC, ElementPrologueB, ElementBias,198+ DispatchPolicyMmad, L1TileShape, L0TileShape, ElementA, ElementB, ElementC, ElementBias, TileCopy>;
192- TileCopy>;
193 199 
194 using BlockPrologue =200 using BlockPrologue =
195 Gemm::Block::BlockPrologue<DispatchPolicyPrologue, PrologueSrcType, PrologueDstType, L1TileShape, TileCopy>;201 Gemm::Block::BlockPrologue<DispatchPolicyPrologue, PrologueSrcType, PrologueDstType, L1TileShape, TileCopy>;
Aexamples/74_ascend950_weight_quant_a8w4_grouped_mx_matmul/CMakeLists.txt+12-0
@@ -0,0 +1,12 @@
1+# ----------------------------------------------------------------------------
2+# This program is free software, you can redistribute it and/or modify.
3+# Copyright (c) 2026 Huawei Technologies Co., Ltd.
4+# This file is a part of the CANN Open Software.
5+# Licensed under CANN Open Software License Agreement Version 2.0 (the "License").
6+# Please refer to the License for details. You may not use this file except in compliance with the License.
7+# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+# See LICENSE in the root of the software repository for the full text of the License.
9+# ----------------------------------------------------------------------------
10+ 
11+set_source_files_properties(weight_quant_a8w4_grouped_mx_matmul.cpp PROPERTIES LANGUAGE ASC)
12+catlass_example_add_executable(74_ascend950_weight_quant_a8w4_grouped_mx_matmul mix weight_quant_a8w4_grouped_mx_matmul.cpp)
Aexamples/74_ascend950_weight_quant_a8w4_grouped_mx_matmul/README.md+92-0
@@ -0,0 +1,92 @@
1+# WeightQuantA8W4GroupedMxMatmul Example Readme
2+> **注意**：社区包暂不支持 950 能力，后续支持的版本敬请期待。
3+## 功能介绍
4+- 演示 Ascend 950 上的伪量化场景下的 Group Mx Matmul 矩阵乘法：左矩阵 A 与伪量化后的右矩阵 B 经 MX 缩放（`float8_e8m0`）后在 Cube 上完成乘加，输出为 FP16。
5+- 本示例中 A 元素类型为 `float8_e4m3_t`，B 元素类型为 `float4_e2m1x2_t`；缩放因子为 `float8_e8m0_t`。未启用 Bias（`ElementBias` 为 `void`）。
6+- 默认布局为 A `RowMajor`、B `Weight4BitnZ`、C `RowMajor`。
7+- B矩阵布局仅支持`Weight4BitnZ`, `ColumnMajor`。通常的fp4内部分型为16x64，但这里的Weight4BitnZ内部分型为16x32，和fp8分型保持一致，目的是为了方便vector将fp4 cast到fp8，省去的分型转换，提高算子性能，Weight4BitnZ分型结构如下图：
8+ 
9+![image-20260721093458714](https://raw.gitcode.com/weixin_42818618/picture0/raw/main/image-20260721093458714.png)
10+ 
11+## 代码组织
12+```text
13+├── 74_ascend950_weight_quant_a8w4_grouped_mx_matmul
14+│ ├── CMakeLists.txt # CMake 编译文件
15+│ ├── README.md
16+│ ├── gen_data.py
17+│ └── weight_quant_a8w4_grouped_mx_matmul.cpp # 主文件
18+```
19+## 使用示例
悠悠的云朵
悠悠的云朵悠悠的云朵7月12日

文档中补充nZ格式的使用说明。

likedislike
20+- 获取代码之后编译相应的算子可执行文件，可参考 [quickstart](../../docs/zh/1_Practice/01_quick_start.md#编译执行)。本用例为 Ascend 950（3510）算子，编译时需加 `-DCATLASS_ARCH=3510`。
21+- 执行算子。
22+**编译指定用例：**
23+```bash
24+bash scripts/build.sh 74_ascend950_weight_quant_a8w4_grouped_mx_matmul -DCATLASS_ARCH=3510
25+```
26+**生成测试样例**（在 `examples/74_ascend950_weight_quant_a8w4_grouped_mx_matmul/data` 下生成 `input/` 与 `golden/`）：
27+输入支持两种模式，同时支持生成两种布局的 B 矩阵。
28+**模式一：group_list mode** —— 通过命令行参数显式输入 `group_m_list`。
29+```bash
30+python3 examples/74_ascend950_weight_quant_a8w4_grouped_mx_matmul/gen_data.py group_list 128,128 256 256 256 1
31+```
32+**模式二：expect_m_per_group mode** —— 按专家数和期望平均值随机生成 `group_m_list`。
33+```bash
34+python3 examples/74_ascend950_weight_quant_a8w4_grouped_mx_matmul/gen_data.py expect_m_per_group 2 128 256 256 256 1
35+```
36+**gen_data.py脚本参数说明：**
37+| 参数 | 含义 |
38+| :--- | :--- |
39+| `group_m_list` | 每个专家对应的分组大小 |
40+| `expect_m_per_group` | 随机分组模式，按每组期望分组大小随机生成分组 |
41+| `group_num` | 专家数 / 分组数 |
42+| `expect_m_per_group` | 每组期望平均分组大小 |
43+| `m` | 总的 M 上限，要求 `m >= sum(group_m_list)` |
44+| `k` | 矩阵乘的 k 维 |
45+| `n` | 矩阵乘的 n 维 |
46+| `isNz` | 0 为 nd 布局，1 为 nZ 布局 |
47+ 
48+需要注意的是，脚本可以生成任意大小n、k数据。
49+对于本例matmul，nd支持任意shape，Weight4BitnZ需要按分型对齐
50+ 
51+**执行测试样例：**
52+```bash
53+./output/bin/74_ascend950_weight_quant_a8w4_grouped_mx_matmul 2 256 256 256 0
54+# 可执行文件名 | 矩阵 m 轴（总m） | n 轴 | k 轴 | Device ID
55+# Device ID 可选，默认为 0
56+```
57+执行结果如下，说明精度比对成功：
58+```text
59+Compare success.
60+```
61+## 使用说明
62+**B 矩阵布局切换**
63+本案例支持 B 矩阵切换`ColumnMajor`和`Weight4BitnZ`格式的输入，默认为 `nZ` 格式。如需切换，请在 `weight_quant_a8w4_grouped_mx_matmul.cpp` 文件中将：
64+```cpp
65+using LayoutPrologueB = layout::Weight4BitnZ
66+```
67+修改为：
68+```cpp
69+using LayoutPrologueB = layout::ColumnMajor
70+```
71+生成 nd 布局的 B 矩阵时，请修改 `gen_data.py` 传入参数：
72+```bash
73+python3 examples/74_ascend950_weight_quant_a8w4_grouped_mx_matmul/gen_data.py expect_m_per_group 2 128 256 256 256 0
74+```
75+**算子计算逻辑**
76+本 example 完成 Grouped mx 量化矩阵乘：
77+```
78+C = (MxScaleA x A) * (MxScaleB x B) + Bias
79+```
80+- A、B 支持数据类型为 `float8_e4m3` 和 `float4_e2m1`，B 矩阵Cast为 `float8_e4m3` 后参与 Cube 计算。
81+- `MxScaleA`、`MxScaleB` 支持数据类型为 `float8_e8m0`。
82+## 性能测试
83+本例性能对比如下表格所示
84+| group_num | m_per_expert | n | k | catlass(us) | 标杆(us) | 标杆/catlass |
85+| :--- | :--- | :--- | :--- | :--- | :--- | :--- |
86+| 48 | 64 | 4096 | 1024 | 89.573 | 164.111 | 1.83214808 |
87+| 48 | 80 | 4096 | 1024 | 94.001 | 167.312 | 1.779895959 |
88+| 48 | 96 | 4096 | 1024 | 96.111 | 178.953 | 1.861940881 |
89+| 48 | 112 | 4096 | 1024 | 101.063 | 181.987 | 1.800728259 |
90+| 48 | 128 | 4096 | 1024 | 106.075 | 162.171 | 1.528833373 |
91+ 
92+本例采用了动态 tiling、L2 Cache Hit、Double Buffer 等优化，实现性能提升 10%+。
Aexamples/74_ascend950_weight_quant_a8w4_grouped_mx_matmul/gen_data.py+621-0
@@ -0,0 +1,621 @@
1+#!/usr/bin/env python3
2+# -*- coding: utf-8 -*-
3+# ----------------------------------------------------------------------------
4+# This program is free software, you can redistribute it and/or modify.
5+# Copyright (c) 2026 Huawei Technologies Co., Ltd.
6+# This file is a part of the CANN Open Software.
7+# Licensed under CANN Open Software License Agreement Version 2.0 (the "License").
8+# Please refer to the License for details. You may not use this file except in compliance with the License.
9+# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
10+# See LICENSE in the root of the software repository for the full text of the License.
11+# ----------------------------------------------------------------------------
12+ 
13+import argparse
14+import math
15+import os
16+import sys
17+from typing import Dict, List, Tuple
18+ 
19+import torch
20+import numpy as np
21+ 
22+WORKSPACE = os.path.dirname(os.path.abspath(__file__))
23+ 
24+_BLOCK_SIZE = 32
25+_EPSILON = 1e-12
26+_MIN_SCALE_EXP = -128
27+_MAX_SCALE_EXP = 127
28+ 
29+_FP8_FORMATS = {
30+ "E4M3": {
31+ "torch_dtype": torch.float8_e4m3fn,
32+ "exp_bits": 4,
33+ "mantissa_bits": 3,
34+ "bias": 7,
35+ "emax": 8,
36+ "max_value": 448.0,
37+ "min_value": -448.0,
38+ },
39+ "E5M2": {
40+ "torch_dtype": torch.float8_e5m2,
41+ "exp_bits": 5,
42+ "mantissa_bits": 2,
43+ "bias": 15,
44+ "emax": 15,
45+ "max_value": 57344.0,
46+ "min_value": -57344.0,
47+ },
48+}
49+ 
50+_FP4_FORMATS: Dict[str, Dict[str, float]] = {
51+ "E2M1": {
52+ "exp_bits": 2,
53+ "mantissa_bits": 1,
54+ "bias": 1,
55+ "emax": 2,
56+ "max_value": 6.0,
57+ "min_value": -6.0,
58+ },
59+ "E1M2": {
60+ "exp_bits": 1,
61+ "mantissa_bits": 2,
62+ "bias": 1,
63+ "emax": 0,
64+ "max_value": 1.75,
65+ "min_value": -1.75,
66+ },
67+}
68+ 
69+ 
70+GROUP_LIST_MODE = "group_list"
71+EXPECT_M_PER_GROUP_MODE = "expect_m_per_group"
72+ 
73+def parse_group_m_list(arg: str) -> List[int]:
74+ """解析逗号分隔的 group M 列表"""
75+ values: List[int] = []
76+ for item in arg.split(","):
77+ item = item.strip()
78+ if not item:
79+ raise ValueError("group_m_list contains an empty item")
80+ value = int(item)
81+ if value < 0:
82+ raise ValueError("Each group M value must be greater than or equal to 0")
83+ values.append(value)
84+ if not values:
85+ raise ValueError("group_m_list must not be empty")
86+ return values
87+ 
88+ 
89+def build_random_group_m_list(group_num: int, expect_m_per_group: int, m: int) -> List[int]:
90+ """在 [0.7*expect, 1.3*expect] 范围内随机生成 group M 列表，保证总和不超过 m"""
91+ if group_num <= 0:
92+ raise ValueError("group_num must be greater than 0")
93+ if expect_m_per_group < 0:
94+ raise ValueError("expect_m_per_group must be greater than or equal to 0")
95+ if m < 0:
96+ raise ValueError("m must be greater than or equal to 0")
97+ 
98+ low = int(math.floor(expect_m_per_group * 0.7))
99+ high = int(math.ceil(expect_m_per_group * 1.3))
100+ low = max(0, low)
101+ high = max(low, high)
102+ 
103+ min_total_m = group_num * low
104+ if m < min_total_m:
105+ raise ValueError(
106+ f"m must be greater than or equal to group_num * floor(0.7 * expect_m_per_group)={min_total_m}"
107+ )
108+ 
109+ if high == 0:
110+ return [0] * group_num
111+ 
112+ rng = np.random.default_rng()
113+ for _ in range(200):
114+ group_m_arr = rng.integers(low, high + 1, size=group_num)
115+ if int(group_m_arr.sum()) <= m:
116+ return group_m_arr.astype(int).tolist()
117+ 
118+ # Fallback: 从 low 开始分配剩余
119+ group_m_list = [low] * group_num
120+ remaining = m - sum(group_m_list)
121+ if remaining <= 0:
122+ return group_m_list
123+ 
124+ capacities = [high - low for _ in range(group_num)]
125+ order = rng.permutation(group_num).tolist()
126+ while remaining > 0:
127+ progressed = False
128+ for idx in order:
129+ if capacities[idx] <= 0:
130+ continue
131+ group_m_list[idx] += 1
132+ capacities[idx] -= 1
133+ remaining -= 1
134+ progressed = True
135+ if remaining == 0:
136+ break
137+ if not progressed:
138+ break
139+ return group_m_list
140+ 
141+ 
142+def parse_cli_args(argv: List[str]) -> Tuple[List[int], int, int, int, int]:
143+ """
144+ 解析命令行参数，返回 (group_list, m, n, k, isNz)
145+ 支持两种模式：
146+ 1. group_list mode: gen_data.py group_list group_m_list m n k isNz
147+ 2. expect_m_per_group mode: gen_data.py expect_m_per_group group_num expect_m_per_group m n k isNz
148+ """
149+ if len(argv) >= 6 and argv[1] == GROUP_LIST_MODE:
150+ group_m_list = parse_group_m_list(argv[2])
151+ m = int(argv[3])
152+ n = int(argv[4])
153+ k = int(argv[5])
154+ isNz = bool(int(argv[6]))
155+ if m < sum(group_m_list):
156+ raise ValueError(f"m must be greater than or equal to sum(group_m_list)={sum(group_m_list)}")
157+ 
158+ return group_m_list, m, n, k, isNz
159+ 
160+ if len(argv) >= 7 and argv[1] == EXPECT_M_PER_GROUP_MODE:
161+ group_num = int(argv[2])
162+ expect_m_per_group = int(argv[3])
163+ m = int(argv[4])
164+ n = int(argv[5])
165+ k = int(argv[6])
166+ isNz = bool(int(argv[7]))
167+ group_m_list = build_random_group_m_list(group_num, expect_m_per_group, m)
168+ 
169+ return group_m_list, m, n, k, isNz
170+ 
171+ raise ValueError(
172+ "Usage:\n"
173+ " python gen_data.py group_list group_m_list m n k\n"
174+ " python gen_data.py expect_m_per_group group_num expect_m_per_group m n k\n"
175+ "Example:\n"
176+ " python gen_data.py group_list 64,128,96 400 256 128\n"
177+ " python gen_data.py expect_m_per_group 4 100 400 256 128"
178+ )
179+ 
180+ 
181+def _build_e4m3_lut() -> torch.Tensor:
182+ bias = _FP8_FORMATS["E4M3"]["bias"]
183+ fp8_max = _FP8_FORMATS["E4M3"]["max_value"]
184+ values = []
185+ for i in range(256):
186+ if i < 128:
187+ sign, val = 1, i
188+ else:
189+ sign, val = -1, i - 128
190+ if val == 0:
191+ v = 0.0
192+ elif val == 127:
193+ v = sign * fp8_max
194+ else:
195+ exp = (val >> 3) & 0x0F
196+ mantissa = val & 0x07
197+ if exp == 0:
198+ v = (mantissa / 8.0) * (2.0 ** (1 - bias))
199+ else:
200+ v = (1.0 + mantissa / 8.0) * (2.0 ** (exp - bias))
201+ v *= sign
202+ v = max(min(v, fp8_max), -fp8_max)
203+ values.append(v)
204+ return torch.tensor(values, dtype=torch.float32)
205+ 
206+ 
207+def _build_e5m2_lut() -> torch.Tensor:
208+ bias = _FP8_FORMATS["E5M2"]["bias"]
209+ fp8_max = _FP8_FORMATS["E5M2"]["max_value"]
210+ values = []
211+ for i in range(256):
212+ if i < 128:
213+ sign, val = 1, i
214+ else:
215+ sign, val = -1, i - 128
216+ if val == 0:
217+ v = 0.0
218+ elif 124 <= val <= 127:
219+ v = sign * fp8_max
220+ else:
221+ exp = (val >> 2) & 0x1F
222+ mantissa = val & 0x03
223+ if exp == 0:
224+ v = (mantissa / 4.0) * (2.0 ** (1 - bias))
225+ else:
226+ v = (1.0 + mantissa / 4.0) * (2.0 ** (exp - bias))
227+ v *= sign
228+ v = max(min(v, fp8_max), -fp8_max)
229+ values.append(v)
230+ return torch.tensor(values, dtype=torch.float32)
231+ 
232+ 
233+_FP8_LUT_BUILDERS = {"E4M3": _build_e4m3_lut, "E5M2": _build_e5m2_lut}
234+_FP8_LUT_CACHE = {}
235+_FP8_LUT_POS_CACHE = {}
236+ 
237+ 
238+def _get_fp8_lut(format_name: str) -> torch.Tensor:
239+ if format_name not in _FP8_LUT_CACHE:
240+ _FP8_LUT_CACHE[format_name] = _FP8_LUT_BUILDERS[format_name]()
241+ return _FP8_LUT_CACHE[format_name]
242+ 
243+ 
244+def _get_fp8_lut_pos(format_name: str) -> torch.Tensor:
245+ if format_name not in _FP8_LUT_POS_CACHE:
246+ full = _get_fp8_lut(format_name)
247+ pos = full[:128].contiguous()
248+ diffs = pos[1:] - pos[:-1]
249+ if (diffs < 0).any():
250+ raise AssertionError(
251+ f"{format_name} positive LUT half is not non-decreasing")
252+ _FP8_LUT_POS_CACHE[format_name] = pos
253+ return _FP8_LUT_POS_CACHE[format_name]
254+ 
255+ 
256+def _build_fp4_lut(format_name: str) -> torch.Tensor:
257+ config = _FP4_FORMATS[format_name]
258+ exp_bits = int(config["exp_bits"])
259+ mantissa_bits = int(config["mantissa_bits"])
260+ bias = float(config["bias"])
261+ 
262+ values = []
263+ for i in range(16):
264+ sign = (i >> 3) & 0x01
265+ exp = (i >> mantissa_bits) & ((1 << exp_bits) - 1)
266+ mantissa = i & ((1 << mantissa_bits) - 1)
267+ 
268+ if exp == 0:
269+ if mantissa == 0:
270+ value = 0.0
271+ else:
272+ value = (mantissa / float(1 << mantissa_bits)) * (2.0 ** (1.0 - bias))
273+ else:
274+ value = (1.0 + mantissa / float(1 << mantissa_bits)) * (2.0 ** (float(exp) - bias))
275+ 
276+ if sign == 1:
277+ value = -value
278+ values.append(value)
279+ 
280+ return torch.tensor(values, dtype=torch.float32)
281+ 
282+ 
283+_FP4_LUT = {
284+ "E2M1": _build_fp4_lut("E2M1"),
285+ "E1M2": _build_fp4_lut("E1M2"),
286+}
287+ 
288+ 
289+def _e8m0_exp(max_abs: torch.Tensor, emax: int,
290+ epsilon: float = _EPSILON) -> torch.Tensor:
291+ assert max_abs.dtype == torch.float32, max_abs.dtype
292+ zero_mask = max_abs < epsilon
293+ safe = torch.where(zero_mask, torch.ones_like(max_abs), max_abs)
294+ bits = safe.contiguous().view(torch.int32)
295+ exp_bits = (bits >> 23) & 0xFF
296+ exp = exp_bits - 127 - emax
297+ exp = exp.clamp(_MIN_SCALE_EXP, _MAX_SCALE_EXP)
298+ return torch.where(zero_mask, torch.zeros_like(exp), exp)
299+ 
300+ 
301+def _vectorized_lut_quantize_fp8(scaled: torch.Tensor, format_name: str,
302+ fp8_dtype: torch.dtype) -> torch.Tensor:
303+ lut_pos = _get_fp8_lut_pos(format_name)
304+ last_idx = lut_pos.numel() - 1
305+ 
306+ sign = torch.sign(scaled)
307+ mag = scaled.abs()
308+ 
309+ upper_idx = torch.searchsorted(lut_pos, mag).clamp(max=last_idx)
310+ lower_idx = (upper_idx - 1).clamp(min=0)
311+ 
312+ upper_val = lut_pos[upper_idx]
313+ lower_val = lut_pos[lower_idx]
314+ 
315+ pick_lower = (mag - lower_val) <= (upper_val - mag)
316+ chosen_mag = torch.where(pick_lower, lower_val, upper_val)
317+ 
318+ snapped_fp32 = sign * chosen_mag
319+ 
320+ zero_mask = chosen_mag == 0
321+ snapped_fp32 = torch.where(
322+ zero_mask, torch.zeros_like(snapped_fp32), snapped_fp32)
323+ 
324+ return snapped_fp32.to(fp8_dtype)
325+ 
326+ 
327+def _quantize_to_fp4_lut(values: torch.Tensor, format_name: str) -> Tuple[torch.Tensor, torch.Tensor]:
328+ lut = _FP4_LUT[format_name].to(values.device)
329+ min_value = _FP4_FORMATS[format_name]["min_value"]
330+ max_value = _FP4_FORMATS[format_name]["max_value"]
331+ 
332+ clamped = values.clamp(min_value, max_value)
333+ 
334+ distances = (clamped.unsqueeze(-1) - lut).abs()
335+ indices = torch.argmin(distances, dim=-1)
336+ quantized = lut[indices]
337+ 
338+ return quantized, indices.to(torch.uint8)
339+ 
340+ 
341+def _pack_fp4_nibbles(index_matrix: torch.Tensor) -> torch.Tensor:
342+ rows, cols = index_matrix.shape
343+ if cols % 2 != 0:
344+ index_matrix = torch.cat(
345+ [index_matrix, torch.zeros((rows, 1), dtype=torch.uint8, device=index_matrix.device)],
346+ dim=1,
347+ )
348+ 
349+ low = index_matrix[:, 0::2]
350+ high = index_matrix[:, 1::2] << 4
351+ packed = low | high
352+ return packed.to(torch.uint8)
353+ 
354+ 
355+def _quantize_fp8_axis_last(matrix: torch.Tensor, format_name: str,
356+ block_size: int = _BLOCK_SIZE
357+ ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
358+ M, N = matrix.shape
359+ fmt = _FP8_FORMATS[format_name]
360+ fp8_dtype = fmt["torch_dtype"]
361+ fp8_emax = fmt["emax"]
362+ fp8_max = fmt["max_value"]
363+ 
364+ num_blocks = (N + block_size - 1) // block_size
365+ padded_n = num_blocks * block_size
366+ if padded_n != N:
367+ padded = torch.zeros(M, padded_n, dtype=matrix.dtype)
368+ padded[:, :N] = matrix
369+ else:
370+ padded = matrix
371+ 
372+ blocks = padded.view(M, num_blocks, block_size)
373+ 
374+ max_abs = blocks.abs().amax(dim=-1)
375+ exp = _e8m0_exp(max_abs, fp8_emax)
376+ scale = torch.exp2(exp.to(torch.float32))
377+ 
378+ scaled = blocks / scale.unsqueeze(-1)
379+ scaled_clamped = scaled.clamp(-fp8_max, fp8_max)
380+ 
381+ quant_fp8 = _vectorized_lut_quantize_fp8(scaled_clamped, format_name, fp8_dtype)
382+ 
383+ dequant = quant_fp8.to(torch.float32) * scale.unsqueeze(-1)
384+ 
385+ if padded_n != N:
386+ quant_fp8 = quant_fp8.reshape(M, padded_n)[:, :N].contiguous()
387+ dequant = dequant.reshape(M, padded_n)[:, :N].contiguous()
388+ else:
389+ quant_fp8 = quant_fp8.reshape(M, N)
390+ dequant = dequant.reshape(M, N)
391+ 
392+ padded_blocks = ((num_blocks + 1) // 2) * 2
393+ if padded_blocks != num_blocks:
394+ scale_padded = torch.ones((M, padded_blocks), dtype=torch.float32)
395+ scale_padded[:, :num_blocks] = scale
396+ scale = scale_padded
397+ 
398+ return quant_fp8, scale, dequant
399+ 
400+ 
401+def _quantize_fp8_axis_first(matrix: torch.Tensor, format_name: str,
402+ block_size: int = _BLOCK_SIZE
403+ ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
404+ qt, st, dt = _quantize_fp8_axis_last(
405+ matrix.t().contiguous(), format_name, block_size)
406+ return (qt.t().contiguous(),
407+ st.t().contiguous(),
408+ dt.t().contiguous())
409+ 
410+ 
411+def _quantize_fp8(matrix: torch.Tensor, format_name: str, axis: int,
412+ block_size: int = _BLOCK_SIZE
413+ ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
414+ if axis == 0:
415+ return _quantize_fp8_axis_first(matrix, format_name, block_size)
416+ if axis == 1:
417+ return _quantize_fp8_axis_last(matrix, format_name, block_size)
418+ raise ValueError(f"axis must be 0 or 1, got {axis}")
419+ 
420+ 
421+def _quantize_fp4_axis_last(matrix: torch.Tensor, format_name: str,
422+ block_size: int = _BLOCK_SIZE
423+ ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
424+ m, n = matrix.shape
425+ padded_n = ((n + block_size - 1) // block_size) * block_size
426+ num_blocks = padded_n // block_size
427+ 
428+ if padded_n != n:
429+ padded = torch.zeros((m, padded_n), dtype=matrix.dtype, device=matrix.device)
430+ padded[:, :n] = matrix
431+ else:
432+ padded = matrix
433+ 
434+ blocks = padded.view(m, num_blocks, block_size)
435+ max_abs = blocks.abs().amax(dim=-1)
436+ 
437+ exp = torch.floor(torch.log2(torch.clamp(max_abs, min=_EPSILON))) - _FP4_FORMATS[format_name]["emax"]
438+ exp = torch.where(max_abs < _EPSILON, torch.zeros_like(exp), exp)
439+ exp = exp.clamp(_MIN_SCALE_EXP, _MAX_SCALE_EXP)
440+ scale = torch.pow(torch.tensor(2.0, dtype=torch.float32, device=matrix.device), exp)
441+ 
442+ scaled = blocks / scale.unsqueeze(-1)
443+ quantized_blocks, _ = _quantize_to_fp4_lut(scaled, format_name)
444+ dequant_blocks = quantized_blocks * scale.unsqueeze(-1)
445+ 
446+ quantized = quantized_blocks.reshape(m, padded_n)
447+ dequantized = dequant_blocks.reshape(m, padded_n)
448+ if padded_n != n:
449+ quantized = quantized[:, :n].contiguous()
450+ dequantized = dequantized[:, :n].contiguous()
451+ 
452+ padded_blocks = ((num_blocks + 1) // 2) * 2
453+ if padded_blocks != num_blocks:
454+ scale_padded = torch.ones((m, padded_blocks), dtype=torch.float32, device=matrix.device)
455+ scale_padded[:, :num_blocks] = scale
456+ scale = scale_padded
457+ 
458+ return quantized, scale, dequantized
459+ 
460+ 
461+def _quantize_fp4_axis_first(matrix: torch.Tensor, format_name: str,
462+ block_size: int = _BLOCK_SIZE
463+ ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
464+ quantized_t, scale_t, dequantized_t = _quantize_fp4_axis_last(
465+ matrix.t().contiguous(), format_name, block_size)
466+ return quantized_t.t().contiguous(), scale_t.t().contiguous(), dequantized_t.t().contiguous()
467+ 
468+ 
469+def _quantize_fp4(matrix: torch.Tensor, format_name: str, axis: int,
470+ block_size: int = _BLOCK_SIZE
471+ ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
472+ if axis == 0:
473+ return _quantize_fp4_axis_first(matrix, format_name, block_size)
474+ if axis == 1:
475+ return _quantize_fp4_axis_last(matrix, format_name, block_size)
476+ raise ValueError(f"axis must be 0 or 1, got {axis}")
477+ 
478+ 
479+def gen_data_fp8_e4m3(row, col, axis):
480+ 
481+ matrix = torch.randn((row, col), dtype=torch.float32) * 10
482+ quant_fp8, scale_fp32, dequant_fp32 = _quantize_fp8(matrix, "E4M3", axis)
483+ return (quant_fp8.to(torch.float8_e4m3fn),
484+ scale_fp32.to(torch.float8_e8m0fnu),
485+ dequant_fp32)
486+ 
487+ 
488+def gen_data_fp8_e5m2(row, col, axis):
489+ matrix = torch.randn((row, col), dtype=torch.float32)
490+ quant_fp8, scale_fp32, dequant_fp32 = _quantize_fp8(matrix, "E5M2", axis)
491+ return (quant_fp8.to(torch.float8_e5m2),
492+ scale_fp32.to(torch.float8_e8m0fnu),
493+ dequant_fp32)
494+ 
495+ 
496+def gen_data_fp4_e2m1(row, col, axis, trans):
497+ matrix = torch.randn((row, col), dtype=torch.float32)
498+ quantized_matrix, scale_matrix, dequantized_matrix = _quantize_fp4(matrix, "E2M1", axis)
499+ 
500+ if trans == 1:
501+ quantized_matrix = quantized_matrix.t().contiguous()
502+ 
503+ _, fp4_indices = _quantize_to_fp4_lut(quantized_matrix, "E2M1")
504+ quantized_matrix_uint8 = _pack_fp4_nibbles(fp4_indices)
505+ 
506+ return quantized_matrix_uint8, scale_matrix.to(torch.float8_e8m0fnu), dequantized_matrix
507+ 
508+def trans_nd2nz(input_data):
509+ g, n_pad, k_pad = input_data.shape
510+ return input_data.reshape(g, n_pad // 16, 16, k_pad // 32, 32).permute(0, 3, 1, 2, 4)
511+ 
512+def gen_data(group_m_list, m, n, k, isNz, trans_a = 0, trans_b = 1) -> None:
513+ data_dir = os.path.join(WORKSPACE, "data")
514+ input_dir = os.path.join(data_dir, "input")
515+ golden_dir = os.path.join(data_dir, "golden")
516+ os.makedirs(input_dir, exist_ok=True)
517+ os.makedirs(golden_dir, exist_ok=True)
518+ 
519+ a_fp8, a_scale, a_fp32 = gen_data_fp8_e4m3(m, k, 1)
520+ 
521+ a_scale = a_scale.reshape(a_scale.shape[0], a_scale.shape[1] // 2, 2)
522+ 
523+ if trans_a == 1:
524+ a_scale = a_scale.permute(1, 0, 2)
525+ 
526+ g = len(group_m_list)
527+ b_indices_list = [] # 存储未打包的 FP4 索引矩阵 (k, n)，元素为 0~15
528+ b_scale_list = []
529+ b_fp32_list = []
530+ b_packed = None
531+ if isNz:
532+ for _ in range(g):
533+ # 生成随机 FP32 矩阵
534+ matrix = torch.randn((k, n), dtype=torch.float32)
535+ # 量化：axis=0（按行量化），返回量化后的浮点值、scale、反量化结果
536+ quantized_vals, scale, dequantized = _quantize_fp4(matrix, "E2M1", axis=0)
537+ 
538+ if trans_b == 1:
539+ quantized_vals = quantized_vals.t().contiguous()
540+ 
541+ # 从浮点量化值获取 FP4 索引（整数）
542+ _, fp4_indices = _quantize_to_fp4_lut(quantized_vals, "E2M1")
543+ b_indices_list.append(fp4_indices) # 形状 (k, n)，整数类型
544+ b_scale_list.append(scale.to(torch.float8_e8m0fnu))
545+ b_fp32_list.append(dequantized)
546+ 
547+ b_indices_stacked = torch.stack(b_indices_list, dim=0) # (g, k, n)，整数
548+ 
549+ pad_k = (32 - k % 32) % 32
550+ pad_n = (16 - n % 16) % 16
551+ if pad_k > 0 or pad_n > 0:
552+ b_indices_stacked = torch.nn.functional.pad(b_indices_stacked, (0, pad_k, 0, pad_n), "constant", 0)
553+ 
554+ b_quant_nd = b_indices_stacked.reshape(g, n + pad_n, k + pad_k) # (g, n_pad, k_pad)
555+ 
556+ b_quant_nz = trans_nd2nz(b_quant_nd) # (g, k_blk, n_blk, 16, 32)
557+ b_quant_flat = b_quant_nz.reshape(-1, 32) # (total_rows, 32) 整数
558+ b_packed = _pack_fp4_nibbles(b_quant_flat) # 打包为 nibble
559+ else:
560+ for _ in range(g):
561+ b_fp4, b_scale, b_fp32 = gen_data_fp4_e2m1(k, n, 0, trans_b)
562+ b_indices_list.append(b_fp4)
563+ b_scale_list.append(b_scale)
564+ b_fp32_list.append(b_fp32)
565+ 
566+ b_packed = torch.stack(b_indices_list, dim=0) # (g, k, n) 或 (g, n, k)
567+ 
568+ b_scale_stacked = torch.stack(b_scale_list, dim=0)
569+ b_fp32_stacked = torch.stack(b_fp32_list, dim=0)
570+ 
571+ b_scale_processed_list = []
572+ for i in range(g):
573+ scale = b_scale_stacked[i]
574+ scale = scale.reshape(scale.shape[0] // 2, 2, scale.shape[1])
575+ if trans_b == 1:
576+ scale = scale.permute(2, 0, 1)
577+ else:
578+ scale = scale.permute(0, 2, 1)
579+ b_scale_processed_list.append(scale)
580+ b_scale_processed = torch.stack(b_scale_processed_list, dim=0)
581+ 
582+ a_np = torch.tensor(a_fp8.flatten().untyped_storage(), dtype=torch.int8).numpy()
583+ b_np = torch.tensor(b_packed.flatten().untyped_storage(), dtype=torch.int8).numpy()
584+ a_np.tofile(os.path.join(input_dir, "a_8.bin"))
585+ b_np.tofile(os.path.join(input_dir, "b_4.bin"))
586+ 
587+ a_scale_np = torch.tensor(a_scale.flatten().untyped_storage(), dtype=torch.int8).numpy()
588+ b_scale_np = torch.tensor(b_scale_processed.flatten().untyped_storage(), dtype=torch.int8).numpy()
589+ a_scale_np.tofile(os.path.join(input_dir, "a_scale.bin"))
590+ b_scale_np.tofile(os.path.join(input_dir, "b_scale.bin"))
591+ 
592+ group_list = np.array(group_m_list, dtype=np.int64)
593+ group_list.tofile(os.path.join(input_dir, "group_list.bin"))
594+ 
595+ c_fp32_list = []
596+ m_offset = 0
597+ for i, group_m in enumerate(group_list):
598+ if group_m == 0:
599+ continue
600+ end_m = m_offset + group_m
601+ a_group = a_fp32[m_offset:end_m] # (group_m, k)
602+ b_group = b_fp32_stacked[i] # (k, n) 或 (n, k)
603+ c_group = a_group @ b_group # (group_m, n)
604+ c_fp32_list.append(c_group)
605+ m_offset = end_m
606+ 
607+ c_fp32 = torch.cat(c_fp32_list, dim=0) # (m, n)
608+ 
609+ c_np = c_fp32.numpy()
610+ 
611+ c_np.tofile(os.path.join(golden_dir, "expected_data.bin"))
612+ 
613+ 
614+if __name__ == "__main__":
615+ try:
616+ group_list, m, n, k, isNz = parse_cli_args(sys.argv)
617+ except ValueError as error:
618+ print(error)
619+ sys.exit(1)
620+ 
621+ gen_data(group_list, m, n, k, isNz)
Aexamples/74_ascend950_weight_quant_a8w4_grouped_mx_matmul/weight_quant_a8w4_grouped_mx_matmul.cpp+329-0
@@ -0,0 +1,329 @@
1+/**
2+ * This program is free software, you can redistribute it and/or modify.
3+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
4+ * This file is a part of the CANN Open Software.
5+ * Licensed under CANN Open Software License Agreement Version 2.0 (the "License").
6+ * Please refer to the License for details. You may not use this file except in compliance with the License.
7+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING
8+ * BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. See LICENSE in the root of
9+ * the software repository for the full text of the License.
10+ */
11+ 
12+#ifndef K_MAX_SHAPE_DIM
13+#define K_MAX_SHAPE_DIM 0
14+#endif
15+ 
16+#include "catlass/gemm/kernel/weight_quant_a8w4_grouped_mx_matmul.hpp"
17+ 
18+#include "catlass/arch/arch.hpp"
19+#include "catlass/catlass.hpp"
20+#include "catlass/gemm/block/block_mmad.hpp"
21+#include "catlass/gemm/block/block_swizzle.hpp"
22+#include "catlass/gemm/device/device_gemm.hpp"
23+#include "catlass/gemm/dispatch_policy.hpp"
24+#include "catlass/gemm/gemm_type.hpp"
25+#include "catlass/layout/layout.hpp"
26+#include "catlass/status.hpp"
27+#include "tla/layout.hpp"
28+ 
29+#include "golden.hpp"
30+#include "helper.hpp"
31+ 
32+using namespace Catlass;
33+using namespace tla;
34+ 
35+using Options = GroupedGemmOptions;
36+ 
37+// Default data root when running from build output (e.g. output/bin), aligned with gen_data.py (WORKSPACE/data).
38+static const std::string kDataRoot = "./examples/74_ascend950_weight_quant_a8w4_grouped_mx_matmul/data";
39+ 
40+static void Run(const Options& options)
41+{
42+ aclrtStream stream{nullptr};
43+ 
44+ ACL_CHECK(aclInit(nullptr));
45+ ACL_CHECK(aclrtSetDevice(options.deviceId));
46+ ACL_CHECK(aclrtCreateStream(&stream));
47+ 
48+ // group_num、m、n、k
49+ uint32_t problemCount = options.problemCount;
50+ uint32_t m = options.problemShape.m();
51+ uint32_t n = options.problemShape.n();
52+ uint32_t k = options.problemShape.k();
53+ uint32_t mxScaleK = CeilDiv<MX_SCALE_GROUP_NUM>(k);
54+ 
55+ using ElementA = float8_e4m3_t;
56+ using ElementB = float8_e4m3_t;
57+ using ElementPrologueB = float4_e2m1x2_t;
58+ using ElementMxScale = float8_e8m0_t;
59+ using ElementC = half;
60+ using ElementBias = void;
61+ 
62+ using ElementGroupList = int64_t;
63+ 
64+ using ElementBiasType = std::conditional_t<std::is_void_v<ElementBias>, uint8_t, ElementBias>;
65+ 
66+ // basic layout
67+ using LayoutA = layout::RowMajor;
68+ using LayoutB = layout::nZ;
69+ using LayoutPrologueB = layout::Weight4BitnZ;
70+ using LayoutMxScaleB = layout::ColumnMajor;
71+ using LayoutC = layout::RowMajor;
72+ 
73+ // makeLayout
74+ LayoutA tagA = LayoutA::MakeLayout<ElementA>(m, k);
75+ LayoutPrologueB tagPrologueB = LayoutPrologueB::MakeLayout<ElementPrologueB>(k, n);
76+ LayoutC tagC = LayoutC::MakeLayout<ElementC>(m, n);
77+ 
78+ static constexpr uint32_t MX_k_ALIGN = 2;
79+ static constexpr uint32_t SIZE_MAGNIFICATION = 2;
80+ 
81+ // data length
82+ size_t lenA = tagA.Capacity();
83+ size_t lenPrologueB = tagPrologueB.Capacity() * problemCount;
84+ uint32_t mxScaleAlignedK = RoundUp<MX_k_ALIGN>(mxScaleK);
85+ size_t lenMxScaleA = static_cast<size_t>(m) * mxScaleAlignedK;
86+ size_t lenMxScaleB = static_cast<size_t>(mxScaleAlignedK) * n * problemCount;
87+ size_t lenC = tagC.Capacity();
88+ size_t lenBias = static_cast<size_t>(n);
89+ 
90+ // data size(len * sizeof(element))
91+ size_t sizeA = lenA * sizeof(ElementA);
92+ size_t sizeB = lenPrologueB / SIZE_MAGNIFICATION;
93+ size_t sizeMxScaleA = lenMxScaleA * sizeof(ElementMxScale);
94+ size_t sizeMxScaleB = lenMxScaleB * sizeof(ElementMxScale);
95+ size_t sizeC = lenC * sizeof(ElementC);
96+ size_t sizeBias = lenBias * sizeof(ElementBiasType);
97+ size_t sizeGroupList = problemCount * sizeof(ElementGroupList);
98+ size_t sizeWorkspace;
99+ 
100+ // host
101+ std::vector<int8_t> hostA(sizeA);
102+ std::vector<int8_t> hostB(sizeB);
103+ std::vector<int8_t> hostMxScaleA(lenMxScaleA);
104+ std::vector<int8_t> hostMxScaleB(lenMxScaleB);
105+ std::vector<ElementBiasType> hostBias(lenBias);
106+ std::vector<ElementGroupList> hostGroupList(sizeGroupList);
107+ 
108+ const auto releaseAclEarly = [&]() {
109+ ACL_CHECK(aclrtDestroyStream(stream));
110+ ACL_CHECK(aclrtResetDevice(options.deviceId));
111+ ACL_CHECK(aclFinalize());
112+ };
113+ // file read
114+ if (!ReadFile(kDataRoot + "/input/a_8.bin", hostA.data(), sizeA)) {
115+ releaseAclEarly();
116+ return;
117+ }
118+ if (!ReadFile(kDataRoot + "/input/b_4.bin", hostB.data(), sizeB)) {
119+ releaseAclEarly();
120+ return;
121+ }
122+ if (!ReadFile(kDataRoot + "/input/a_scale.bin", hostMxScaleA.data(), sizeMxScaleA)) {
123+ releaseAclEarly();
124+ return;
125+ }
126+ if (!ReadFile(kDataRoot + "/input/b_scale.bin", hostMxScaleB.data(), sizeMxScaleB)) {
127+ releaseAclEarly();
128+ return;
129+ }
130+ if (!ReadFile(kDataRoot + "/input/group_list.bin", hostGroupList.data(), sizeGroupList)) {
131+ releaseAclEarly();
132+ return;
133+ }
134+ if constexpr (!std::is_void_v<ElementBias>) {
135+ if (!ReadFile(kDataRoot + "/input/bias.bin", hostBias.data(), sizeBias)) {
136+ releaseAclEarly();
137+ return;
138+ }
139+ }
140+ 
141+ // device
142+ uint8_t* deviceGroupList{nullptr};
143+ ACL_CHECK(aclrtMalloc(reinterpret_cast<void**>(&deviceGroupList), sizeGroupList, ACL_MEM_MALLOC_HUGE_FIRST));
144+ ACL_CHECK(
145+ aclrtMemcpy(deviceGroupList, sizeGroupList, hostGroupList.data(), sizeGroupList, ACL_MEMCPY_HOST_TO_DEVICE));
146+ 
147+ uint8_t* deviceA{nullptr};
148+ ACL_CHECK(aclrtMalloc(reinterpret_cast<void**>(&deviceA), sizeA, ACL_MEM_MALLOC_HUGE_FIRST));
149+ ACL_CHECK(aclrtMemcpy(deviceA, sizeA, hostA.data(), sizeA, ACL_MEMCPY_HOST_TO_DEVICE));
150+ 
151+ uint8_t* deviceB{nullptr};
152+ ACL_CHECK(aclrtMalloc(reinterpret_cast<void**>(&deviceB), sizeB, ACL_MEM_MALLOC_HUGE_FIRST));
153+ ACL_CHECK(aclrtMemcpy(deviceB, sizeB, hostB.data(), sizeB, ACL_MEMCPY_HOST_TO_DEVICE));
154+ 
155+ uint8_t* deviceMxScaleA{nullptr};
156+ ACL_CHECK(aclrtMalloc(reinterpret_cast<void**>(&deviceMxScaleA), sizeMxScaleA, ACL_MEM_MALLOC_HUGE_FIRST));
157+ ACL_CHECK(aclrtMemcpy(deviceMxScaleA, sizeMxScaleA, hostMxScaleA.data(), sizeMxScaleA, ACL_MEMCPY_HOST_TO_DEVICE));
158+ 
159+ uint8_t* deviceMxScaleB{nullptr};
160+ ACL_CHECK(aclrtMalloc(reinterpret_cast<void**>(&deviceMxScaleB), sizeMxScaleB, ACL_MEM_MALLOC_HUGE_FIRST));
161+ ACL_CHECK(aclrtMemcpy(deviceMxScaleB, sizeMxScaleB, hostMxScaleB.data(), sizeMxScaleB, ACL_MEMCPY_HOST_TO_DEVICE));
162+ 
163+ uint8_t* deviceC{nullptr};
164+ ACL_CHECK(aclrtMalloc(reinterpret_cast<void**>(&deviceC), sizeC, ACL_MEM_MALLOC_HUGE_FIRST));
165+ 
166+ uint8_t* deviceBias{nullptr};
167+ if constexpr (!std::is_void_v<ElementBias>) {
168+ ACL_CHECK(aclrtMalloc(reinterpret_cast<void**>(&deviceBias), sizeBias, ACL_MEM_MALLOC_HUGE_FIRST));
169+ ACL_CHECK(aclrtMemcpy(deviceBias, sizeBias, hostBias.data(), sizeBias, ACL_MEMCPY_HOST_TO_DEVICE));
170+ }
171+ 
172+ uint8_t* deviceWorkspace{nullptr};
173+ 
174+ // Get the number of cube cores of the current hardware
175+ auto aicCoreNum = platform_ascendc::PlatformAscendCManager::GetInstance()->GetCoreNumAic();
176+ 
177+ // archtag uniflag
178+ using ArchTag = Arch::Ascend950;
179+ constexpr bool enableUnitFlag = true;
180+ static constexpr uint32_t L1_SCALE_FACTOR_K = 16;
181+ static constexpr uint32_t L1A_STAGES = 2;
182+ static constexpr uint32_t L1B_STAGES = 2;
183+ static constexpr uint32_t L0A_STAGES = 2;
184+ static constexpr uint32_t L0B_STAGES = 2;
185+ static constexpr uint32_t L0C_STAGES = 1;
186+ 
187+ // shape & type
188+ using L1TileShape = Shape<Int<256>, Int<256>, Int<256>>;
189+ using L0TileShape = Shape<Int<256>, Int<256>, Int<128>>;
190+ using PrologueSrcType = Gemm::GemmType<ElementPrologueB, LayoutPrologueB>;
191+ using PrologueDstType = Gemm::GemmType<ElementB, LayoutB>;
192+ 
193+ // DispatchPolicy
194+ using DispatchPolicyMmad = Gemm::MmadA8W4Mx<
195+ ArchTag, enableUnitFlag, false, L1_SCALE_FACTOR_K, L0C_STAGES, L1A_STAGES, L1B_STAGES, L0A_STAGES, L0B_STAGES>;
196+ using DispatchPolicyPrologue = Gemm::MxA8W4Prologue<ArchTag, L1B_STAGES>;
197+ 
198+ // layout (tla)
199+ auto layoutA = tla::MakeLayout<ElementA, LayoutA>(m, k);
200+ auto layoutprologueB = tla::MakeLayout<ElementPrologueB, LayoutPrologueB>(k, n);
201+ auto layoutMxScaleA = tla::MakeMxScaleLayout<ElementMxScale, LayoutA, false>(m, mxScaleK);
202+ auto layoutMxScaleB = tla::MakeMxScaleLayout<ElementMxScale, LayoutMxScaleB, true>(mxScaleK, n);
203+ auto layoutC = tla::MakeLayout<ElementC, LayoutC>(m, n);
204+ 
205+ // tile
206+ using TileCopy = Gemm::Tile::PackedMxA8W4TileCopyTla<
207+ ArchTag, ElementA, LayoutA, ElementPrologueB, LayoutPrologueB, ElementB, LayoutB, ElementMxScale,
208+ decltype(layoutMxScaleA), ElementMxScale, decltype(layoutMxScaleB), ElementC, LayoutC, ElementBias, false,
209+ Gemm::Tile::ScaleGranularity::PER_TENSOR>;
210+ 
211+ // BlockMmad
212+ using BlockMmad = Gemm::Block::BlockMmadA8W4Mx<
213+ DispatchPolicyMmad, L1TileShape, L0TileShape, ElementA, ElementB, ElementC, ElementBias, TileCopy>;
214+ 
215+ using BlockPrologue =
216+ Gemm::Block::BlockPrologue<DispatchPolicyPrologue, PrologueSrcType, PrologueDstType, L1TileShape, TileCopy>;
217+ 
218+ // Epilogue
219+ using BlockEpilogue = void;
220+ 
221+ // Swizzle offset is 3 and direction is 0.
222+ using BlockScheduler = typename Gemm::Block::GemmIdentityBlockSwizzle<3, 0>;
223+ 
224+ // kernel level
225+ using MatmulKernel =
226+ Gemm::Kernel::A8W4GroupedMxMatmul<BlockMmad, BlockPrologue, BlockEpilogue, BlockScheduler, ElementGroupList>;
227+ 
228+ using MatmulAdapter = Gemm::Device::DeviceGemm<MatmulKernel>;
229+ 
230+ MatmulKernel::Arguments arguments{
231+ options.problemShape, options.problemCount, deviceGroupList, deviceA, layoutA, deviceB, layoutprologueB,
232+ deviceMxScaleA, layoutMxScaleA, deviceMxScaleB, layoutMxScaleB, deviceC, layoutC, deviceBias};
233+ 
234+ uint32_t taskNum = CeilDiv(options.problemShape.m(), tla::get<0>(L1TileShape{})) *
235+ CeilDiv(options.problemShape.n(), tla::get<1>(L1TileShape{}));
236+ uint32_t aicCoreUsed = min(aicCoreNum, taskNum);
237+ 
238+ MatmulAdapter matmulOp;
239+ matmulOp.CanImplement(arguments);
240+ sizeWorkspace = matmulOp.GetWorkspaceSize(arguments);
241+ if (sizeWorkspace > 0) {
242+ ACL_CHECK(aclrtMalloc(reinterpret_cast<void**>(&deviceWorkspace), sizeWorkspace, ACL_MEM_MALLOC_HUGE_FIRST));
243+ }
244+ matmulOp.Initialize(arguments, deviceWorkspace);
245+ matmulOp(stream, aicCoreUsed);
246+ ACL_CHECK(aclrtSynchronizeStream(stream));
247+ 
248+ std::vector<fp16_t> hostC(lenC);
249+ ACL_CHECK(aclrtMemcpy(hostC.data(), sizeC, deviceC, sizeC, ACL_MEMCPY_DEVICE_TO_HOST));
250+ 
251+ std::vector<float> hostGolden(lenC);
252+ std::string outputFileName = kDataRoot + "/golden/expected_data.bin";
253+ if (!ReadFile(outputFileName, hostGolden.data(), sizeof(float) * hostGolden.size())) {
254+ ACL_CHECK(aclrtFree(deviceA));
255+ ACL_CHECK(aclrtFree(deviceB));
256+ ACL_CHECK(aclrtFree(deviceMxScaleA));
257+ ACL_CHECK(aclrtFree(deviceMxScaleB));
258+ ACL_CHECK(aclrtFree(deviceC));
259+ ACL_CHECK(aclrtFree(deviceGroupList));
260+ if constexpr (!std::is_void_v<ElementBias>) {
261+ ACL_CHECK(aclrtFree(deviceBias));
262+ }
263+ if (sizeWorkspace > 0) {
264+ ACL_CHECK(aclrtFree(deviceWorkspace));
265+ }
266+ ACL_CHECK(aclrtDestroyStream(stream));
267+ ACL_CHECK(aclrtResetDevice(options.deviceId));
268+ ACL_CHECK(aclFinalize());
269+ return;
270+ }
271+ 
272+ std::vector<uint64_t> errorIndices = golden::CompareData(hostC, hostGolden, k);
273+ if (errorIndices.empty()) {
274+ std::cout << "Compare success." << std::endl;
275+ } else {
276+ std::cerr << "Compare failed. Error count: " << errorIndices.size() << std::endl;
277+ 
278+ for (uint32_t i = 0; i < 10; ++i) {
279+ std::cout << "Index: " << errorIndices[i] << " npu:" << (float)hostC[errorIndices[i]]
280+ << " cpu:" << hostGolden[errorIndices[i]] << std::endl;
281+ }
282+ 
283+ std::cout << std::endl;
284+ 
285+ for (uint32_t i = 0; i < 20; ++i) {
286+ std::cout << "Index: " << i << " npu:" << (float)hostC[i] << " cpu:" << hostGolden[i] << std::endl;
287+ }
288+ 
289+ uint32_t* hostAInt = reinterpret_cast<uint32_t*>(hostA.data());
290+ uint32_t* hostBInt = reinterpret_cast<uint32_t*>(hostB.data());
291+ uint32_t* hostMxScaleAInt = reinterpret_cast<uint32_t*>(hostMxScaleA.data());
292+ uint32_t* hostMxScaleBInt = reinterpret_cast<uint32_t*>(hostMxScaleB.data());
293+ for (uint32_t i = 0; i < 10; ++i) {
294+ std::cout << "index: " << i << " hostA: " << hostAInt[i] << " hostB: " << hostBInt[i]
295+ << " hostMxScaleA: " << hostMxScaleAInt[i] << " hostMxScaleB: " << hostMxScaleBInt[i]
296+ << std::endl;
297+ }
298+ }
299+ // for (uint32_t i = 0; i < 20; ++i){
300+ // std::cout<< "Index: " << i << " npu:" << (float) hostC[i]<< " cpu:" << hostGolden[i]<<std::endl;
301+ // }
302+ 
303+ ACL_CHECK(aclrtFree(deviceA));
atomgit-bot
atomgit-botatomgit-bot7月9日

🟡 Medium Priority

weight_quant_a8w4_grouped_mx_matmul.cpp 第 143 行通过 aclrtMalloc 分配了 deviceGroupList，但在正常的清理路径（第 299-309 行）和 golden 文件读取失败的错误路径（第 253-267 行）中均未调用 aclrtFree(deviceGroupList)。其他设备指针（deviceA、deviceB、deviceMxScaleA、deviceMxScaleB、deviceC、deviceBias、deviceWorkspace）在两个路径中都被正确释放，唯独遗漏了 deviceGroupList。每次 Run() 调用都会泄漏 sizeGroupList 字节的 NPU 设备内存。

建议：在两个清理路径（正常路径第 299 行之前和错误路径第 253 行之前）均添加 ACL_CHECK(aclrtFree(deviceGroupList));。

likedislike
不准确？
304+ ACL_CHECK(aclrtFree(deviceB));
305+ ACL_CHECK(aclrtFree(deviceMxScaleA));
306+ ACL_CHECK(aclrtFree(deviceMxScaleB));
307+ ACL_CHECK(aclrtFree(deviceC));
308+ ACL_CHECK(aclrtFree(deviceGroupList));
309+ if constexpr (!std::is_void_v<ElementBias>) {
310+ ACL_CHECK(aclrtFree(deviceBias));
311+ }
312+ if (sizeWorkspace > 0) {
313+ ACL_CHECK(aclrtFree(deviceWorkspace));
314+ }
315+ 
316+ ACL_CHECK(aclrtDestroyStream(stream));
317+ ACL_CHECK(aclrtResetDevice(options.deviceId));
318+ ACL_CHECK(aclFinalize());
319+}
320+ 
321+int main(int argc, const char** argv)
322+{
323+ Options options;
324+ if (options.Parse(argc, argv) != 0) {
325+ return -1;
326+ }
327+ Run(options);
328+ return 0;
329+}
Mexamples/CMakeLists.txt+1-0
@@ -169,6 +169,7 @@ set(EXAMPLE_ASCEND950
169 69_ascend950_tail_multi_core_splitk_matmul169 69_ascend950_tail_multi_core_splitk_matmul
170 70_ascend950_flash_attention_chunk_prefill170 70_ascend950_flash_attention_chunk_prefill
171 73_ascend950_matmul_full_loadA171 73_ascend950_matmul_full_loadA
172+ 74_ascend950_weight_quant_a8w4_grouped_mx_matmul
172 80_grouped_matmul_slice_m_gelu173 80_grouped_matmul_slice_m_gelu
173)174)
174 175 
Minclude/catlass/detail/tag_to_layout.hpp+10-0
@@ -108,6 +108,16 @@ struct TagToLayout<float8_e8m0_t, layout::nN> {
108 tla::Stride<108 tla::Stride<
109 tla::Stride<tla::Int<1>, tla::Int<ELE_NUM_PER_FRACTAL>>, tla::Stride<tla::Int<ELE_NUM_PER_C0>, int64_t>>>;109 tla::Stride<tla::Int<1>, tla::Int<ELE_NUM_PER_FRACTAL>>, tla::Stride<tla::Int<ELE_NUM_PER_C0>, int64_t>>>;
110};110};
111+ 
112+template <>
113+struct TagToLayout<float4_e2m1x2_t, layout::Weight4BitnZ> {
114+ static constexpr uint32_t ELE_NUM_PER_C0 = 32;
115+ static constexpr uint32_t ELE_NUM_PER_FRACTAL = 512;
116+ using type = tla::Layout<
117+ tla::Shape<tla::Shape<tla::Int<ELE_NUM_PER_C0>, uint32_t>, tla::Shape<tla::Int<C0_NUM_PER_FRACTAL>, uint32_t>>,
118+ tla::Stride<
119+ tla::Stride<tla::Int<1>, int64_t>, tla::Stride<tla::Int<ELE_NUM_PER_C0>, tla::Int<ELE_NUM_PER_FRACTAL>>>>;
120+};
111#endif121#endif
112 122 
113// Convenience aliases123// Convenience aliases
Minclude/catlass/gemm/block/block_mmad.hpp+1-1
@@ -71,7 +71,7 @@ struct BlockMmadTla {
71 71 
72template <72template <
73 class DispatchPolicy, class L1TileShape, class L0TileShape, class ElementA, class ElementB, class ElementC,73 class DispatchPolicy, class L1TileShape, class L0TileShape, class ElementA, class ElementB, class ElementC,
74- class ElementPrologueB, class ElementBias = void,74+ class ElementBias = void,
75 class TileCopy = Gemm::Tile::PackedTileCopyTla<75 class TileCopy = Gemm::Tile::PackedTileCopyTla<
76 typename DispatchPolicy::ArchTag, ElementA, layout::zN, ElementB, layout::zN, ElementC, layout::zN,76 typename DispatchPolicy::ArchTag, ElementA, layout::zN, ElementB, layout::zN, ElementC, layout::zN,
77 ElementBias>,77 ElementBias>,
Minclude/catlass/gemm/block/block_mmad_mx_a8w4.hpp+110-60
dxdx7月17日

🟡 代码质量：ENABLE_UNIT_FLAG 分支硬编码访问 l0CTensorList [0] 存在隐患 此处开启 ENABLE_UNIT_FLAG 时直接用魔法数字 0 下标访问缓存数组，仅靠现有断言弱约束参数关联，模板传参错误会引发数组越界，且和其余分支索引写法不统一； 关注点： 补充 static_assert 强绑定 ENABLE_UNIT_FLAG 与 L0C_STAGES==1，编译期拦截非法参数； 使用 l0CListId 替换魔法数字 0，统一全文件数组访问方式。

likedislike
dxdx7月17日

🟡 代码质量：裸指针 lastAddrA/lastAddrB 强制转换无空值防护，存在未定义行为 UB 隐患 此处将 GM 张量物理地址强转为全局裸指针存入数组，未对 GetPhyAddr () 返回地址做空指针校验，若 GM 张量为空会存储 nullptr，后续 L1 驻留模式指针判读逻辑异常引发访存崩溃； 关注点： 赋值前增加物理地址空判断，仅非空地址才存入 lastAddr 数组； 统一在 RestoreStatus 构造初始化阶段批量清零指针，避免野指针残留。

likedislike
@@ -27,20 +27,19 @@ namespace Catlass::Gemm::Block {
27#if (defined(CATLASS_ARCH) && CATLASS_ARCH == 3510)27#if (defined(CATLASS_ARCH) && CATLASS_ARCH == 3510)
28 28 
29template <29template <
30- class ArchTag_, bool ENABLE_UNIT_FLAG_, uint32_t L0C_STAGES_, bool ENABLE_L1_RESIDENT_, uint32_t L1A_STAGES_,30+ class ArchTag_, bool ENABLE_UNIT_FLAG_, bool ENABLE_L1_RESIDENT_, uint32_t L1_SCALE_FACTOR_K_, uint32_t L0C_STAGES_,
31- uint32_t L1B_STAGES_, uint32_t L0A_STAGES_, uint32_t L0B_STAGES_, class L1TileShape_, class L0TileShape_,31+ uint32_t L1A_STAGES_, uint32_t L1B_STAGES_, uint32_t L0A_STAGES_, uint32_t L0B_STAGES_, class L1TileShape_,
32- class ElementA_, class ElementB_, class ElementC_, class ElementPrologueB_, class ElementBias_, class TileCopy_,32+ class L0TileShape_, class ElementA_, class ElementB_, class ElementC_, class ElementBias_, class TileCopy_,
33 class TileMmad_>33 class TileMmad_>
34struct BlockMmadA8W4Mx<34struct BlockMmadA8W4Mx<
35 MmadA8W4Mx<35 MmadA8W4Mx<
36- ArchTag_, ENABLE_UNIT_FLAG_, L0C_STAGES_, ENABLE_L1_RESIDENT_, L1A_STAGES_, L1B_STAGES_, L0A_STAGES_,36+ ArchTag_, ENABLE_UNIT_FLAG_, ENABLE_L1_RESIDENT_, L1_SCALE_FACTOR_K_, L0C_STAGES_, L1A_STAGES_, L1B_STAGES_,
37- L0B_STAGES_>,37+ L0A_STAGES_, L0B_STAGES_>,
38- L1TileShape_, L0TileShape_, ElementA_, ElementB_, ElementC_, ElementPrologueB_, ElementBias_, TileCopy_,38+ L1TileShape_, L0TileShape_, ElementA_, ElementB_, ElementC_, ElementBias_, TileCopy_, TileMmad_> {
39- TileMmad_> {
40public:39public:
41 using DispatchPolicy = MmadA8W4Mx<40 using DispatchPolicy = MmadA8W4Mx<
42- ArchTag_, ENABLE_UNIT_FLAG_, L0C_STAGES_, ENABLE_L1_RESIDENT_, L1A_STAGES_, L1B_STAGES_, L0A_STAGES_,41+ ArchTag_, ENABLE_UNIT_FLAG_, ENABLE_L1_RESIDENT_, L1_SCALE_FACTOR_K_, L0C_STAGES_, L1A_STAGES_, L1B_STAGES_,
43- L0B_STAGES_>;42+ L0A_STAGES_, L0B_STAGES_>;
44 using ArchTag = typename DispatchPolicy::ArchTag;43 using ArchTag = typename DispatchPolicy::ArchTag;
45 using TileCopy = TileCopy_;44 using TileCopy = TileCopy_;
46 using L1TileShape = L1TileShape_;45 using L1TileShape = L1TileShape_;
@@ -54,7 +53,6 @@ public:
54 using ElementC = ElementC_;53 using ElementC = ElementC_;
55 using LayoutC = typename TileCopy::LayoutC;54 using LayoutC = typename TileCopy::LayoutC;
56 using ElementBias = ElementBias_;55 using ElementBias = ElementBias_;
57- using ElementPrologueB = ElementPrologueB_;
58 using ElementL0A = typename helper::GetL0Element<ElementA, true>::Element;56 using ElementL0A = typename helper::GetL0Element<ElementA, true>::Element;
59 using ElementL0B = typename helper::GetL0Element<ElementB, true>::Element;57 using ElementL0B = typename helper::GetL0Element<ElementB, true>::Element;
60 58 
@@ -90,6 +88,9 @@ public:
90 static constexpr uint32_t L0A_STAGES = DispatchPolicy::L0A_STAGES;88 static constexpr uint32_t L0A_STAGES = DispatchPolicy::L0A_STAGES;
91 static constexpr uint32_t L0B_STAGES = DispatchPolicy::L0B_STAGES;89 static constexpr uint32_t L0B_STAGES = DispatchPolicy::L0B_STAGES;
92 static constexpr uint32_t L0C_STAGES = DispatchPolicy::L0C_STAGES;90 static constexpr uint32_t L0C_STAGES = DispatchPolicy::L0C_STAGES;
91+ 
92+ static constexpr uint32_t L1_SCALE_FACTOR_K = DispatchPolicy::L1_SCALE_FACTOR_K;
93+ 
93 static constexpr uint32_t L1_TILE_M = tla::get<0>(L1TileShape{});94 static constexpr uint32_t L1_TILE_M = tla::get<0>(L1TileShape{});
94 static constexpr uint32_t L1_TILE_N = tla::get<1>(L1TileShape{});95 static constexpr uint32_t L1_TILE_N = tla::get<1>(L1TileShape{});
95 static constexpr uint32_t L1_TILE_K = tla::get<2>(L1TileShape{});96 static constexpr uint32_t L1_TILE_K = tla::get<2>(L1TileShape{});
@@ -100,8 +101,10 @@ public:
100 // L1 tile size101 // L1 tile size
101 static constexpr uint32_t L1A_TILE_SIZE = L1_TILE_M * L1_TILE_K * SizeOfBits<ElementA>::value / 8;102 static constexpr uint32_t L1A_TILE_SIZE = L1_TILE_M * L1_TILE_K * SizeOfBits<ElementA>::value / 8;
102 static constexpr uint32_t L1B_TILE_SIZE = L1_TILE_N * L1_TILE_K * SizeOfBits<ElementB>::value / 8;103 static constexpr uint32_t L1B_TILE_SIZE = L1_TILE_N * L1_TILE_K * SizeOfBits<ElementB>::value / 8;
103- static constexpr uint32_t L1SCALEA_TILE_SIZE = L1_TILE_M * L1_TILE_K / MX_SCALE_GROUP_NUM * sizeof(ElementMxScaleA);104+ static constexpr uint32_t L1SCALEA_TILE_SIZE =
104- static constexpr uint32_t L1SCALEB_TILE_SIZE = L1_TILE_N * L1_TILE_K / MX_SCALE_GROUP_NUM * sizeof(ElementMxScaleB);105+ L1_TILE_M * L1_TILE_K / MX_SCALE_GROUP_NUM * sizeof(ElementMxScaleA) * L1_SCALE_FACTOR_K;
106+ static constexpr uint32_t L1SCALEB_TILE_SIZE =
107+ L1_TILE_N * L1_TILE_K / MX_SCALE_GROUP_NUM * sizeof(ElementMxScaleB) * L1_SCALE_FACTOR_K;
105 static constexpr uint32_t L1_USED_SIZE = L1A_TILE_SIZE * L1A_STAGES + L1B_TILE_SIZE * L1B_STAGES +108 static constexpr uint32_t L1_USED_SIZE = L1A_TILE_SIZE * L1A_STAGES + L1B_TILE_SIZE * L1B_STAGES +
106 L1SCALEA_TILE_SIZE * L1A_STAGES + L1SCALEB_TILE_SIZE * L1B_STAGES;109 L1SCALEA_TILE_SIZE * L1A_STAGES + L1SCALEB_TILE_SIZE * L1B_STAGES;
107 // L0 tile size110 // L0 tile size
@@ -151,9 +154,9 @@ public:
151 static constexpr auto L1B_LAYOUT =154 static constexpr auto L1B_LAYOUT =
152 tla::MakeLayout<ElementB, LayoutTagL1B>(tla::Int<L1_TILE_K>{}, tla::Int<L1_TILE_N>{});155 tla::MakeLayout<ElementB, LayoutTagL1B>(tla::Int<L1_TILE_K>{}, tla::Int<L1_TILE_N>{});
153 static constexpr auto L1SCALEA_LAYOUT = tla::MakeMxScaleLayout<ElementMxScaleA, LayoutTagL1MxScaleA, false>(156 static constexpr auto L1SCALEA_LAYOUT = tla::MakeMxScaleLayout<ElementMxScaleA, LayoutTagL1MxScaleA, false>(
154- tla::Int<L1_TILE_M>{}, tla::Int<L1_TILE_K / MX_SCALE_GROUP_NUM>{});157+ tla::Int<L1_TILE_M>{}, tla::Int<L1_TILE_K / MX_SCALE_GROUP_NUM * L1_SCALE_FACTOR_K>{});
155 static constexpr auto L1SCALEB_LAYOUT = tla::MakeMxScaleLayout<ElementMxScaleB, LayoutTagL1MxScaleB, true>(158 static constexpr auto L1SCALEB_LAYOUT = tla::MakeMxScaleLayout<ElementMxScaleB, LayoutTagL1MxScaleB, true>(
156- tla::Int<L1_TILE_K / MX_SCALE_GROUP_NUM>{}, tla::Int<L1_TILE_N>{});159+ tla::Int<L1_TILE_K / MX_SCALE_GROUP_NUM * L1_SCALE_FACTOR_K>{}, tla::Int<L1_TILE_N>{});
157 static constexpr auto L1BIAS_LAYOUT = tla::MakeLayout(tla::Int<L1_TILE_N>{});160 static constexpr auto L1BIAS_LAYOUT = tla::MakeLayout(tla::Int<L1_TILE_N>{});
158 static constexpr auto L0BIAS_LAYOUT = tla::MakeLayout(tla::Int<L0_TILE_N>{});161 static constexpr auto L0BIAS_LAYOUT = tla::MakeLayout(tla::Int<L0_TILE_N>{});
159 162 
@@ -171,17 +174,29 @@ public:
171 lastAddrB[i] = nullptr;174 lastAddrB[i] = nullptr;
172 lastCoordB[i] = MatrixCoord{0U, 0U};175 lastCoordB[i] = MatrixCoord{0U, 0U};
173 }176 }
177+ l1MxScaleAListId = 0;
178+ l1MxScaleBListId = 0;
174 }179 }
175 180 
176 /// Construct181 /// Construct
177 CATLASS_DEVICE182 CATLASS_DEVICE
178- BlockMmadA8W4Mx(Arch::Resource<ArchTag>& resource, uint32_t l1BufAddrStart)183+ BlockMmadA8W4Mx(Arch::Resource<ArchTag>& resource, uint32_t l1BufAddrStart = 0)
179 {184 {
180 if constexpr (ENABLE_UNIT_FLAG && tla::detail::isRowMajor<LayoutC>::value) {185 if constexpr (ENABLE_UNIT_FLAG && tla::detail::isRowMajor<LayoutC>::value) {
181 AscendC::SetMMLayoutTransform(true);186 AscendC::SetMMLayoutTransform(true);
182 }187 }
183- uint32_t l1Offset = l1BufAddrStart;188+ 
184 // Init buffers189 // Init buffers
190+ uint32_t l1Offset = l1BufAddrStart;
191+ for (uint32_t i = 0; i < L1B_STAGES; i++) {
192+ // Assign L1/L0A/L0B space for each stages
193+ l1BTensorList[i] = resource.l1Buf.template GetBufferByByte<ElementB>(l1Offset);
194+ l1Offset += L1B_TILE_SIZE;
195+ // Assign event ID for each stages
196+ l1BEventList[i] = i + L1A_STAGES;
197+ // The event id that needs to be set before the loop
198+ AscendC::SetFlag<AscendC::HardEvent::MTE1_MTE2>(l1BEventList[i]);
199+ }
185 for (uint32_t i = 0; i < L1A_STAGES; i++) {200 for (uint32_t i = 0; i < L1A_STAGES; i++) {
186 // Assign L1/L0A/L0B space for each stages201 // Assign L1/L0A/L0B space for each stages
187 l1ATensorList[i] = resource.l1Buf.template GetBufferByByte<ElementA>(l1Offset);202 l1ATensorList[i] = resource.l1Buf.template GetBufferByByte<ElementA>(l1Offset);
@@ -191,12 +206,6 @@ public:
191 // The event id that needs to be set before the loop206 // The event id that needs to be set before the loop
192 AscendC::SetFlag<AscendC::HardEvent::MTE1_MTE2>(l1AEventList[i]);207 AscendC::SetFlag<AscendC::HardEvent::MTE1_MTE2>(l1AEventList[i]);
193 }208 }
194- for (uint32_t i = 0; i < L1B_STAGES; i++) {
195- // Assign event ID for each stages
196- l1BEventList[i] = i + L1A_STAGES;
197- // The event id that need to be set before the loop
198- AscendC::SetFlag<AscendC::HardEvent::MTE1_MTE2>(l1BEventList[i]);
199- }
200 for (uint32_t i = 0; i < L0A_STAGES; i++) {209 for (uint32_t i = 0; i < L0A_STAGES; i++) {
201 // Assign L1/L0A/L0B space for each stages210 // Assign L1/L0A/L0B space for each stages
202 l0ATensorList[i] = resource.l0ABuf.template GetBufferByByte<ElementL0A>(L0A_TILE_SIZE * i);211 l0ATensorList[i] = resource.l0ABuf.template GetBufferByByte<ElementL0A>(L0A_TILE_SIZE * i);
@@ -220,7 +229,8 @@ public:
220 AscendC::SetFlag<AscendC::HardEvent::FIX_M>(l0CEventList[i]);229 AscendC::SetFlag<AscendC::HardEvent::FIX_M>(l0CEventList[i]);
221 }230 }
222 } else {231 } else {
223- l0CTensorList[0] = resource.l0CBuf.template GetBufferByByte<ElementAccumulator>(0);232+ l0CTensorList[l0CListId] =
233+ resource.l0CBuf.template GetBufferByByte<ElementAccumulator>(L0C_TILE_SIZE * l0CListId);
224 }234 }
225 if constexpr (HAS_BIAS) {235 if constexpr (HAS_BIAS) {
226 l1BiasTensor = resource.l1Buf.template GetBufferByByte<uint8_t>(l1Offset);236 l1BiasTensor = resource.l1Buf.template GetBufferByByte<uint8_t>(l1Offset);
@@ -275,12 +285,11 @@ public:
275 285 
276 /// Perform a block-scoped matrix multiply-accumulate286 /// Perform a block-scoped matrix multiply-accumulate
277 template <287 template <
278- class TensorA, class TensorL1B, class TensorC, class TensorMxScaleA = EmptyClass,288+ class TensorA, class TensorC, class TensorMxScaleA = EmptyClass, class TensorMxScaleB = EmptyClass,
279- class TensorMxScaleB = EmptyClass, class TensorBias = EmptyClass>289+ class TensorBias = EmptyClass>
280 CATLASS_DEVICE void operator()(290 CATLASS_DEVICE void operator()(
281- TensorA& tensorA, TensorC& tensorC, GemmCoord const& actualShape, TensorL1B& tensorL1B,291+ TensorA& tensorA, TensorC& tensorC, GemmCoord const& actualShape, TensorMxScaleA const& tensorMxScaleA = {},
282- TensorMxScaleA const& tensorMxScaleA = {}, TensorMxScaleB const& tensorMxScaleB = {},292+ TensorMxScaleB const& tensorMxScaleB = {}, TensorBias const& tensorBias = {})
283- TensorBias const& tensorBias = {})
284 {293 {
285 // Check L1TileShape294 // Check L1TileShape
286 if constexpr (HAS_BIAS) {295 if constexpr (HAS_BIAS) {
@@ -301,6 +310,8 @@ public:
301 CopyGmToL1MxScaleB copyGmToL1MxScaleB;310 CopyGmToL1MxScaleB copyGmToL1MxScaleB;
302 CopyL0CToDst copyL0CToDst;311 CopyL0CToDst copyL0CToDst;
303 312 
313+ uint32_t l1ScaleTileK = L1_TILE_K * L1_SCALE_FACTOR_K;
314+ 
304 uint32_t mBlockActual = actualShape.m();315 uint32_t mBlockActual = actualShape.m();
305 uint32_t kBlockActual = actualShape.k();316 uint32_t kBlockActual = actualShape.k();
306 uint32_t nBlockActual = actualShape.n();317 uint32_t nBlockActual = actualShape.n();
@@ -324,12 +335,13 @@ public:
324 335 
325 // main loop336 // main loop
326 uint32_t kL1Loop = CeilDiv<L1_TILE_K>(kBlockActual);337 uint32_t kL1Loop = CeilDiv<L1_TILE_K>(kBlockActual);
338+ uint32_t kL1ScaleLoop = CeilDiv<L1_SCALE_FACTOR_K>(kL1Loop);
339+ 
327 for (uint32_t kL1Idx = 0; kL1Idx < kL1Loop; kL1Idx++) {340 for (uint32_t kL1Idx = 0; kL1Idx < kL1Loop; kL1Idx++) {
328- l1AListId = (l1AListId + 1 < L1A_STAGES) ? (l1AListId + 1) : 0;
329- l1BListId = (l1BListId + 1 < L1B_STAGES) ? (l1BListId + 1) : 0;
330- 
331 auto kL1Actual = (kL1Idx < kL1Loop - 1) ? L1_TILE_K : (kBlockActual - kL1Idx * L1_TILE_K);341 auto kL1Actual = (kL1Idx < kL1Loop - 1) ? L1_TILE_K : (kBlockActual - kL1Idx * L1_TILE_K);
332- 342+ uint32_t kL1ScaleIdx = kL1Idx / L1_SCALE_FACTOR_K;
343+ uint32_t kL1MxScaleActual =
344+ (kL1ScaleIdx < kL1ScaleLoop - 1) ? l1ScaleTileK : (kBlockActual - kL1ScaleIdx * l1ScaleTileK);
333 auto l1ATensor = l1ATensorList[l1AListId];345 auto l1ATensor = l1ATensorList[l1AListId];
334 auto tensorL1A = tla::MakeTensor(l1ATensor, L1A_LAYOUT, Arch::PositionL1{});346 auto tensorL1A = tla::MakeTensor(l1ATensor, L1A_LAYOUT, Arch::PositionL1{});
335 347 
@@ -340,47 +352,69 @@ public:
340 // load matrix A tile from GM to L1352 // load matrix A tile from GM to L1
341 AscendC::WaitFlag<AscendC::HardEvent::MTE1_MTE2>(l1AEventList[l1AListId]);353 AscendC::WaitFlag<AscendC::HardEvent::MTE1_MTE2>(l1AEventList[l1AListId]);
342 if constexpr (ENABLE_L1_RESIDENT) {354 if constexpr (ENABLE_L1_RESIDENT) {
343- if (lastAddrA[l1AListId] != tensorTileA.data().GetPhyAddr() ||355+ auto currentAddrA = tensorTileA.data().GetPhyAddr();
344- tla::get<0>(tensorTileA.coord()) != lastCoordA[l1AListId].row() ||356+ if (currentAddrA != nullptr && (lastAddrA[l1AListId] != currentAddrA ||
345- tla::get<1>(tensorTileA.coord()) != lastCoordA[l1AListId].column()) {357+ tla::get<0>(tensorTileA.coord()) != lastCoordA[l1AListId].row() ||
358+ tla::get<1>(tensorTileA.coord()) != lastCoordA[l1AListId].column())) {
346 copyGmToL1A(tensorL1A, tensorTileA);359 copyGmToL1A(tensorL1A, tensorTileA);
347 lastCoordA[l1AListId] =360 lastCoordA[l1AListId] =
348 MatrixCoord{tla::get<0>(tensorTileA.coord()), tla::get<1>(tensorTileA.coord())};361 MatrixCoord{tla::get<0>(tensorTileA.coord()), tla::get<1>(tensorTileA.coord())};
349- lastAddrA[l1AListId] = const_cast<__gm__ typename AscendC::GlobalTensor<ElementA>::PrimType*>(362+ lastAddrA[l1AListId] =
350- tensorTileA.data().GetPhyAddr());363+ const_cast<__gm__ typename AscendC::GlobalTensor<ElementA>::PrimType*>(currentAddrA);
351 }364 }
352 } else {365 } else {
353 copyGmToL1A(tensorL1A, tensorTileA);366 copyGmToL1A(tensorL1A, tensorTileA);
354 }367 }
355- 
356 // Init Zero for k axis368 // Init Zero for k axis
357 InitZeroInL1A(tensorL1A, tla::MakeShape(mL1Actual, kL1Actual));369 InitZeroInL1A(tensorL1A, tla::MakeShape(mL1Actual, kL1Actual));
358 370 
359 // load mxScaleA tile from GM to L1371 // load mxScaleA tile from GM to L1
360- auto tensorL1MxScaleA =372+ if (kL1Idx % L1_SCALE_FACTOR_K == 0) {
361- tla::MakeTensor(l1MxScaleATensorList[l1AListId], L1SCALEA_LAYOUT, Arch::PositionL1{});373+ // load mxScaleA tile from GM to L1
362- auto tensorTileMxScaleA = GetTile(374+ auto tensorL1MxScaleA = tla::MakeTensor(
363- tensorMxScaleA, tla::MakeCoord(0, kL1Idx * L1_TILE_K / MX_SCALE_GROUP_NUM),375+ l1MxScaleATensorList[l1MxScaleAListId], L1SCALEA_LAYOUT, Catlass::Arch::PositionL1{});
364- tla::MakeShape(mBlockActual, CeilDiv<MX_SCALE_GROUP_NUM>(kL1Actual)));376+ 
365- copyGmToL1MxScaleA(tensorL1MxScaleA, tensorTileMxScaleA);377+ auto tensorTileMxScaleA = GetTile(
378+ tensorMxScaleA,
379+ tla::MakeCoord(0, kL1ScaleIdx * L1_TILE_K * L1_SCALE_FACTOR_K / Catlass::MX_SCALE_GROUP_NUM),
380+ tla::MakeShape(mBlockActual, CeilDiv<Catlass::MX_SCALE_GROUP_NUM>(kL1MxScaleActual)));
381+ copyGmToL1MxScaleA(tensorL1MxScaleA, tensorTileMxScaleA);
382+ }
383+ 
366 AscendC::SetFlag<AscendC::HardEvent::MTE2_MTE1>(l1AEventList[l1AListId]);384 AscendC::SetFlag<AscendC::HardEvent::MTE2_MTE1>(l1AEventList[l1AListId]);
367 385 
368 AscendC::WaitFlag<AscendC::HardEvent::MTE1_MTE2>(l1BEventList[l1BListId]);386 AscendC::WaitFlag<AscendC::HardEvent::MTE1_MTE2>(l1BEventList[l1BListId]);
387+ 
369 // load mxScaleB tile from GM to L1388 // load mxScaleB tile from GM to L1
370- auto tensorL1MxScaleB =389+ if (kL1Idx % L1_SCALE_FACTOR_K == 0) {
371- tla::MakeTensor(l1MxScaleBTensorList[l1BListId], L1SCALEB_LAYOUT, Arch::PositionL1{});390+ // load mxScaleB tile from GM to L1
372- auto tensorTileMxScaleB = GetTile(391+ auto tensorL1MxScaleB = tla::MakeTensor(
373- tensorMxScaleB, tla::MakeCoord(kL1Idx * L1_TILE_K / MX_SCALE_GROUP_NUM, 0),392+ l1MxScaleBTensorList[l1MxScaleBListId], L1SCALEB_LAYOUT, Catlass::Arch::PositionL1{});
374- tla::MakeShape(CeilDiv<MX_SCALE_GROUP_NUM>(kL1Actual), nBlockActual));393+ auto tensorTileMxScaleB = GetTile(
375- copyGmToL1MxScaleB(tensorL1MxScaleB, tensorTileMxScaleB);394+ tensorMxScaleB,
376- AscendC::SetFlag<AscendC::HardEvent::MTE2_MTE1>(l1BEventList[l1BListId]);395+ tla::MakeCoord(kL1ScaleIdx * L1_TILE_K * L1_SCALE_FACTOR_K / Catlass::MX_SCALE_GROUP_NUM, 0),
396+ tla::MakeShape(CeilDiv<Catlass::MX_SCALE_GROUP_NUM>(kL1MxScaleActual), nBlockActual));
397+ copyGmToL1MxScaleB(tensorL1MxScaleB, tensorTileMxScaleB);
398+ }
377 399 
378 // Get the loop nums on L0400 // Get the loop nums on L0
379- AscendC::CrossCoreWaitFlag<AIC_SYNC_AIV_MODE, PIPE_MTE1>(AIV_SYNC_AIC_FLAG);401+ AscendC::CrossCoreWaitFlag<AIC_SYNC_AIV_MODE, PIPE_MTE1>(AIV_SYNC_AIC_FLAG + l1BListId);
380- AscendC::CrossCoreWaitFlag<AIC_SYNC_AIV_MODE, PIPE_MTE1>(AIC_SYNC_AIV_FLAG);402+ AscendC::CrossCoreWaitFlag<AIC_SYNC_AIV_MODE, PIPE_MTE1>(AIV_SYNC_AIC_FLAG + FLAG_ID_MAX + l1BListId);
atomgit-bot
atomgit-botatomgit-bot7月9日

🔴 Critical

在 block_mmad_mx_a8w4.hpp 的 main loop 中，AIC 侧通过 CrossCoreWaitFlag 等待跨核同步。旧代码使用 AIV_SYNC_AIC_FLAG（AIV→AIC 方向，旧值 0），正确等待 AIV prologue 完成 B 数据的反量化处理后通知 AIC。新代码改为使用 AIC_SYNC_AIV_FLAG + l1BListId（AIC→AIV 方向，值 8/9），同时 AIC 自身在 loop 底部又以 CrossCoreSetFlag 设置同一组标志（8/9）。而 AIV prologue 在 ProcessL1Tile 中设置的是 AIV_SYNC_AIC_FLAG + l1BListId（值 6/7），无人等待。

证据链：

失败模式：AIC 不再阻塞等待 AIV 的反量化处理完成，当 AIC 执行速度超过 AIV 时，AIC 可能读到 AIV 尚未写完的 L1B 数据，导致计算结果错误（非确定性精度问题）。触发条件：AIC 与 AIV 核之间的实际执行时序差。

建议：将 line 462-463 的 CrossCoreWaitFlag 等待标志从 AIC_SYNC_AIV_FLAG 改为 AIV_SYNC_AIC_FLAG，恢复 AIC 对 AIV 处理完成的等待语义。同时检查第二组等待标志（原 FLAG_ID_MAX 组）是否仍需要，若不需要则应移除。

改动建议
402
+ AscendC::CrossCoreWaitFlag<AIC_SYNC_AIV_MODE, PIPE_MTE1>(AIV_SYNC_AIC_FLAG + l1BListId);
402
403
  AscendC::CrossCoreWaitFlag<AIC_SYNC_AIV_MODE, PIPE_MTE1>(AIV_SYNC_AIC_FLAG + FLAG_ID_MAX + l1BListId);
应用建议
likedislike
不准确？
UniquePG
UniquePG
7月14日 评论：
381- uint32_t kL0Loop = CeilDiv<L0_TILE_K>(kL1Actual);403+ 
404+ AscendC::SetFlag<AscendC::HardEvent::MTE2_MTE1>(l1BEventList[l1BListId]);
405+ 
382 AscendC::WaitFlag<AscendC::HardEvent::MTE2_MTE1>(l1AEventList[l1AListId]);406 AscendC::WaitFlag<AscendC::HardEvent::MTE2_MTE1>(l1AEventList[l1AListId]);
383 AscendC::WaitFlag<AscendC::HardEvent::MTE2_MTE1>(l1BEventList[l1BListId]);407 AscendC::WaitFlag<AscendC::HardEvent::MTE2_MTE1>(l1BEventList[l1BListId]);
408+ 
409+ uint32_t kL0Loop = CeilDiv<L0_TILE_K>(kL1Actual);
410+ auto tensorL1B = tla::MakeTensor(l1BTensorList[l1BListId], L1B_LAYOUT, Arch::PositionL1{});
411+ auto tensorL1MxScaleA =
412+ tla::MakeTensor(l1MxScaleATensorList[l1MxScaleAListId], L1SCALEA_LAYOUT, Arch::PositionL1{});
413+ auto tensorL1MxScaleB =
414+ tla::MakeTensor(l1MxScaleBTensorList[l1MxScaleBListId], L1SCALEB_LAYOUT, Arch::PositionL1{});
415+ 
416+ uint32_t const l0kOffset = L1_TILE_K * (kL1Idx % L1_SCALE_FACTOR_K);
417+ 
384 for (uint32_t kL0Idx = 0; kL0Idx < kL0Loop; kL0Idx++) {418 for (uint32_t kL0Idx = 0; kL0Idx < kL0Loop; kL0Idx++) {
385 uint32_t kL0Actual = (kL0Idx < kL0Loop - 1) ? L0_TILE_K : (kL1Actual - kL0Idx * L0_TILE_K);419 uint32_t kL0Actual = (kL0Idx < kL0Loop - 1) ? L0_TILE_K : (kL1Actual - kL0Idx * L0_TILE_K);
386 420 
@@ -397,7 +431,7 @@ public:
397 431 
398 // Locate the current tile of matrix mxScaleA on L1432 // Locate the current tile of matrix mxScaleA on L1
399 auto tensorTileL1MxScaleA = GetTile(433 auto tensorTileL1MxScaleA = GetTile(
400- tensorL1MxScaleA, tla::MakeCoord(0, kL0Idx * L0_TILE_K / MX_SCALE_GROUP_NUM),434+ tensorL1MxScaleA, tla::MakeCoord(0, (l0kOffset + kL0Idx * L0_TILE_K) / MX_SCALE_GROUP_NUM),
401 tla::MakeShape(mL1Actual, CeilDiv<MX_SCALE_GROUP_NUM>(kL0Actual)));435 tla::MakeShape(mL1Actual, CeilDiv<MX_SCALE_GROUP_NUM>(kL0Actual)));
402 436 
403 AscendC::WaitFlag<AscendC::HardEvent::M_MTE1>(l0AEventList[l0AListId]);437 AscendC::WaitFlag<AscendC::HardEvent::M_MTE1>(l0AEventList[l0AListId]);
@@ -408,13 +442,14 @@ public:
408 // Locate the current tile on L0B442 // Locate the current tile on L0B
409 auto layoutBInL0 = tla::MakeLayout<ElementB, LayoutTagL0B>(kL0Actual, nL1Actual);443 auto layoutBInL0 = tla::MakeLayout<ElementB, LayoutTagL0B>(kL0Actual, nL1Actual);
410 auto tensorL0B = tla::MakeTensor(l0BTensorList[l0BListId], layoutBInL0, Arch::PositionL0B{});444 auto tensorL0B = tla::MakeTensor(l0BTensorList[l0BListId], layoutBInL0, Arch::PositionL0B{});
445+ 
411 // Locate the current tile of matrix mxScaleB on L1446 // Locate the current tile of matrix mxScaleB on L1
412 auto tensorTileL1B =447 auto tensorTileL1B =
413 GetTile(tensorL1B, tla::MakeCoord(kL0Idx * L0_TILE_K, 0), tla::MakeShape(kL0Actual, nL1Actual));448 GetTile(tensorL1B, tla::MakeCoord(kL0Idx * L0_TILE_K, 0), tla::MakeShape(kL0Actual, nL1Actual));
414 449 
415 // Locate the current tile of matrix mxScaleB on L1450 // Locate the current tile of matrix mxScaleB on L1
416 auto tensorTileL1MxScaleB = GetTile(451 auto tensorTileL1MxScaleB = GetTile(
417- tensorL1MxScaleB, tla::MakeCoord(kL0Idx * L0_TILE_K / MX_SCALE_GROUP_NUM, 0),452+ tensorL1MxScaleB, tla::MakeCoord((l0kOffset + kL0Idx * L0_TILE_K) / MX_SCALE_GROUP_NUM, 0),
418 tla::MakeShape(CeilDiv<MX_SCALE_GROUP_NUM>(kL0Actual), nL1Actual));453 tla::MakeShape(CeilDiv<MX_SCALE_GROUP_NUM>(kL0Actual), nL1Actual));
419 454 
420 // Wait for mmad finished455 // Wait for mmad finished
@@ -470,10 +505,20 @@ public:
470 AscendC::SetFlag<AscendC::HardEvent::M_MTE1>(l0AEventList[l0AListId]);505 AscendC::SetFlag<AscendC::HardEvent::M_MTE1>(l0AEventList[l0AListId]);
471 l0AListId = (l0AListId + 1 < L0A_STAGES) ? (l0AListId + 1) : 0;506 l0AListId = (l0AListId + 1 < L0A_STAGES) ? (l0AListId + 1) : 0;
472 }507 }
508+ 
473 AscendC::SetFlag<AscendC::HardEvent::MTE1_MTE2>(l1AEventList[l1AListId]);509 AscendC::SetFlag<AscendC::HardEvent::MTE1_MTE2>(l1AEventList[l1AListId]);
474 AscendC::SetFlag<AscendC::HardEvent::MTE1_MTE2>(l1BEventList[l1BListId]);510 AscendC::SetFlag<AscendC::HardEvent::MTE1_MTE2>(l1BEventList[l1BListId]);
475- AscendC::CrossCoreSetFlag<AIC_SYNC_AIV_MODE, PIPE_MTE1>(AIV_SYNC_AIC_FLAG);511+ 
476- AscendC::CrossCoreSetFlag<AIC_SYNC_AIV_MODE, PIPE_MTE1>(AIC_SYNC_AIV_FLAG);512+ AscendC::CrossCoreSetFlag<AIC_SYNC_AIV_MODE, PIPE_MTE1>(AIC_SYNC_AIV_FLAG + l1BListId);
513+ AscendC::CrossCoreSetFlag<AIC_SYNC_AIV_MODE, PIPE_MTE1>(AIC_SYNC_AIV_FLAG + FLAG_ID_MAX + l1BListId);
514+ 
515+ l1AListId = (l1AListId + 1 < L1A_STAGES) ? (l1AListId + 1) : 0;
516+ l1BListId = (l1BListId + 1 < L1B_STAGES) ? (l1BListId + 1) : 0;
517+ 
518+ if (((kL1Idx + 1) % L1_SCALE_FACTOR_K == 0) || kL1Idx == kL1Loop - 1) {
519+ l1MxScaleAListId = (l1MxScaleAListId + 1 < L1A_STAGES) ? (l1MxScaleAListId + 1) : 0;
520+ l1MxScaleBListId = (l1MxScaleBListId + 1 < L1B_STAGES) ? (l1MxScaleBListId + 1) : 0;
521+ }
477 }522 }
478 // copy block out523 // copy block out
479 copyL0CToDst.params.scale = 64;524 copyL0CToDst.params.scale = 64;
@@ -559,6 +604,7 @@ protected:
559 604 
560 // Multi-stage tensors list605 // Multi-stage tensors list
561 AscendC::LocalTensor<ElementA> l1ATensorList[L1A_STAGES];606 AscendC::LocalTensor<ElementA> l1ATensorList[L1A_STAGES];
607+ AscendC::LocalTensor<ElementB> l1BTensorList[L1B_STAGES];
562 AscendC::LocalTensor<ElementMxScaleA> l1MxScaleATensorList[L1A_STAGES];608 AscendC::LocalTensor<ElementMxScaleA> l1MxScaleATensorList[L1A_STAGES];
563 AscendC::LocalTensor<ElementMxScaleB> l1MxScaleBTensorList[L1B_STAGES];609 AscendC::LocalTensor<ElementMxScaleB> l1MxScaleBTensorList[L1B_STAGES];
564 AscendC::LocalTensor<ElementL0A> l0ATensorList[L0A_STAGES];610 AscendC::LocalTensor<ElementL0A> l0ATensorList[L0A_STAGES];
@@ -586,9 +632,13 @@ protected:
586 uint32_t l0BListId{0};632 uint32_t l0BListId{0};
587 uint32_t l0CListId{0};633 uint32_t l0CListId{0};
588 634 
635+ uint32_t l1MxScaleAListId{0}; // scale id only used in Mmad3
636+ uint32_t l1MxScaleBListId{0};
637+ 
589 constexpr static uint8_t AIC_SYNC_AIV_MODE = 4;638 constexpr static uint8_t AIC_SYNC_AIV_MODE = 4;
590- static constexpr uint16_t AIV_SYNC_AIC_FLAG = 0;639+ constexpr static uint16_t AIV_SYNC_AIC_FLAG = 6;
591- static constexpr uint16_t AIC_SYNC_AIV_FLAG = 16;640+ constexpr static uint16_t AIC_SYNC_AIV_FLAG = 8;
641+ constexpr static uint16_t FLAG_ID_MAX = 16;
592 642 
593 TileMmad tileMmad;643 TileMmad tileMmad;
594 CopyL1ToL0A copyL1ToL0A;644 CopyL1ToL0A copyL1ToL0A;
Minclude/catlass/gemm/block/block_mx_a8w4_prologue.hpp+332-226
dxdx7月17日

🟡 代码规范：无用 using 类型别名残留，无清理 using _0 = Int<0>; using _16 = Int<16>; using _32 = Int<32>; 仅声明、全文件零使用，属于无效冗余代码，如不使用建议删除。

likedislike
dxdx7月17日

🟡代码规范：注释存在语法、拼写、标点不规范 长注释无分行，句尾缺少句号； 英文句子大小写混乱：sub Vector has to be diabled 单词拼写错误 (diabled→disabled)。

likedislike
@@ -18,38 +18,43 @@
18#include "catlass/gemm/dispatch_policy.hpp"18#include "catlass/gemm/dispatch_policy.hpp"
19#include "catlass/gemm/helper.hpp"19#include "catlass/gemm/helper.hpp"
20#include "catlass/gemm_coord.hpp"20#include "catlass/gemm_coord.hpp"
21+#include "catlass/epilogue/tile/tile_copy.hpp"
21#include "tla/layout.hpp"22#include "tla/layout.hpp"
22#include "tla/tensor.hpp"23#include "tla/tensor.hpp"
23 24 
24namespace Catlass::Gemm::Block {25namespace Catlass::Gemm::Block {
25-template <class ArchTag, class InType_, class OutType_, class TileShapeL1_, class TileCopy_>26+template <class ArchTag, uint32_t L1B_STAGES_, class InType_, class OutType_, class TileShapeL1_, class TileCopy_>
26-struct BlockPrologue<MxA8W4Prologue<ArchTag>, InType_, OutType_, TileShapeL1_, TileCopy_> {27+struct BlockPrologue<MxA8W4Prologue<ArchTag, L1B_STAGES_>, InType_, OutType_, TileShapeL1_, TileCopy_> {
27public:28public:
28- using DispatchPolicy = MxA8W4Prologue<ArchTag>;29+ using DispatchPolicy = MxA8W4Prologue<ArchTag, L1B_STAGES_>;
29 using ElementIn = typename InType_::Element;30 using ElementIn = typename InType_::Element;
30 using ElementOut = typename OutType_::Element;31 using ElementOut = typename OutType_::Element;
31 using LayoutIn = typename InType_::Layout;32 using LayoutIn = typename InType_::Layout;
32 using LayoutOut = typename OutType_::Layout;33 using LayoutOut = typename OutType_::Layout;
33 using TileShapeL1 = TileShapeL1_;34 using TileShapeL1 = TileShapeL1_;
34 using TileCopy = TileCopy_;35 using TileCopy = TileCopy_;
35- using LayoutB = typename TileCopy::LayoutB;36+ using LayoutPrologueB = typename TileCopy::LayoutPrologueB;
36- template <int32_t t>
37- using Int = AscendC::Std::integral_constant<int32_t, t>;
38- using _0 = Int<0>;
39- using _16 = Int<16>;
40- using _32 = Int<32>;
41 37 
42 static constexpr uint32_t L1B_STAGES = DispatchPolicy::L1B_STAGES;38 static constexpr uint32_t L1B_STAGES = DispatchPolicy::L1B_STAGES;
43 39 
44 struct Params {40 struct Params {
45 TileShapeL1 tileShapeL1;41 TileShapeL1 tileShapeL1;
46- LayoutB layoutB;42+ LayoutPrologueB layoutPrologueB;
47- int32_t nUbSize;
48- int32_t kUbSize;
49 bool hasBias;43 bool hasBias;
44+ Catlass::Arch::Resource<ArchTag>& resource;
50 };45 };
51 46 
52- struct VfParamsNormal {47+ struct VfParamsNz {
48+ uint32_t loopKNum;
49+ uint32_t innerLoopNum;
50+ uint32_t loopKDstStride;
51+ uint32_t innerDstStride;
52+ uint32_t nRealSizeAlign;
53+ __ubuf__ ElementIn* weightInUbAddr;
54+ __ubuf__ ElementOut* weightOutUbAddr;
55+ };
56+ 
57+ struct VfParamsNd {
53 uint16_t outExtend;58 uint16_t outExtend;
54 uint16_t innerExtend;59 uint16_t innerExtend;
55 uint32_t dataBlockStride;60 uint32_t dataBlockStride;
@@ -57,49 +62,48 @@ public:
57 int32_t outDimOffset;62 int32_t outDimOffset;
58 uint32_t maskB8Tail0;63 uint32_t maskB8Tail0;
59 uint32_t maskB8Tail1;64 uint32_t maskB8Tail1;
65+ 
66+ int32_t kUbLen;
67+ 
60 __ubuf__ int8_t* weightInUbBaseAddr;68 __ubuf__ int8_t* weightInUbBaseAddr;
61 __ubuf__ ElementOut* weightOutUbAddr;69 __ubuf__ ElementOut* weightOutUbAddr;
62 __ubuf__ ElementOut* weightOutUbAddr1;70 __ubuf__ ElementOut* weightOutUbAddr1;
63 };71 };
64 72 
65- struct VfParamsNz {
66- uint16_t innerExtend;
67- uint32_t innerDstExtend;
68- uint32_t innerSrcExtend;
69- uint32_t shiftLeftSize;
70- uint32_t andMask;
71- __ubuf__ int8_t* weightInUbBaseAddr;
72- __ubuf__ ElementOut* weightOutUbAddr;
73- };
74- 
75 static_assert(73 static_assert(
76 std::is_same_v<LayoutIn, layout::RowMajor> || std::is_same_v<LayoutIn, layout::ColumnMajor> ||74 std::is_same_v<LayoutIn, layout::RowMajor> || std::is_same_v<LayoutIn, layout::ColumnMajor> ||
77- std::is_same_v<LayoutIn, layout::zN>,75+ std::is_same_v<LayoutIn, layout::zN> || std::is_same_v<LayoutIn, layout::Weight4BitnZ>,
78- "Unsupported layout, only can be Rowmajor ColumnMajor or zN");76+ "Unsupported layout, only can be Rowmajor ColumnMajor or zN or nZ");
79 77 
80 CATLASS_DEVICE78 CATLASS_DEVICE
81 BlockPrologue(const Params& params)79 BlockPrologue(const Params& params)
82 {80 {
83- nUbSize_ = params.nUbSize;81+ static constexpr int32_t OFFSET_64 = 64;
84- kUbSize_ = params.kUbSize;82+ vecWeightInLen =
85- nSize_ = tla::get<1>(params.layoutB.shape());83+ (UB_STAGES * (tla::get<1>(params.tileShapeL1) * RoundUp(tla::get<2>(params.tileShapeL1), OFFSET_64))) >>
86- kSize_ = tla::get<0>(params.layoutB.shape());84+ INT4_DTYPE_PARAM;
87- nL1Size_ = tla::get<1>(params.tileShapeL1);85+ vecWeightOutLen = UB_STAGES * (RoundUp(tla::get<1>(params.tileShapeL1), AscendC::BLOCK_CUBE) + 1) *
88- kL1Size_ = tla::get<2>(params.tileShapeL1); // 2 in order to obtain k86+ RoundUp(
89- bL1Size_ = nL1Size_ * RoundUp(kL1Size_, K_ALIGN_SIZE);87+ RoundUp(tla::get<2>(params.tileShapeL1), static_cast<int32_t>(AscendC::ONE_BLK_SIZE)),
90- aL1Size_ = tla::get<0>(params.tileShapeL1) * tla::get<2>(params.tileShapeL1); // 2 in order to obtain k88+ static_cast<int32_t>(K_ALIGN_SIZE));
91- vecWeightInLen_ = (UB_STAGES * (nUbSize_ * RoundUp(kUbSize_, OFFSET_64))) >> INT4_DTYPE_PARAM;
92- vecWeightOutLen_ =
93- UB_STAGES * (RoundUp(nUbSize_, AscendC::BLOCK_CUBE) + 1) *
94- RoundUp(RoundUp(kUbSize_, static_cast<int32_t>(AscendC::ONE_BLK_SIZE)), static_cast<int32_t>(K_ALIGN_SIZE));
95- weightOutUb_ = AscendC::LocalTensor<ElementOut>(AscendC::TPosition::VECCALC, 0, vecWeightOutLen_);
96- uint64_t ubOffset = vecWeightOutLen_ * sizeof(ElementOut);
97- weightInUb_ = AscendC::LocalTensor<ElementIn>(AscendC::TPosition::VECCALC, ubOffset, vecWeightInLen_);
98- l1Local_ = AscendC::LocalTensor<ElementOut>(AscendC::TPosition::B1, 0, ArchTag::L1_SIZE);
99 89 
90+ uint32_t l1Offset = 0;
91+ uint32_t L1B_TILE_SIZE = tla::get<1>(params.tileShapeL1) * tla::get<2>(params.tileShapeL1);
100 for (uint32_t i = 0; i < L1B_STAGES; i++) {92 for (uint32_t i = 0; i < L1B_STAGES; i++) {
101- AscendC::SetFlag<AscendC::HardEvent::V_MTE2>(i);93+ // Assign L1/L0A/L0B space for each stages
102- AscendC::SetFlag<AscendC::HardEvent::MTE3_V>(i);94+ l1BTensorList[i] = params.resource.l1Buf.template GetBufferByByte<ElementOut>(l1Offset);
95+ l1Offset += L1B_TILE_SIZE;
96+ // Assign event ID for each stages
97+ l1BEventList[i] = i;
98+ AscendC::SetFlag<AscendC::HardEvent::V_MTE2>(l1BEventList[i]);
99+ AscendC::SetFlag<AscendC::HardEvent::MTE3_V>(l1BEventList[i]);
100+ }
101+ uint32_t ubOffset = 0;
102+ for (uint32_t i = 0; i < UB_STAGES; i++) {
103+ ubCastInTensor[i] = params.resource.ubBuf.template GetBufferByByte<ElementIn>(ubOffset);
104+ ubOffset += 32 * 1024;
105+ ubCastOutTensor[i] = params.resource.ubBuf.template GetBufferByByte<ElementOut>(ubOffset);
106+ ubOffset += 32 * 1024;
103 }107 }
104 }108 }
105 109 
@@ -107,139 +111,292 @@ public:
107 ~BlockPrologue()111 ~BlockPrologue()
108 {112 {
109 for (uint32_t i = 0; i < L1B_STAGES; i++) {113 for (uint32_t i = 0; i < L1B_STAGES; i++) {
110- AscendC::WaitFlag<AscendC::HardEvent::V_MTE2>(i);114+ AscendC::WaitFlag<AscendC::HardEvent::MTE3_V>(l1BEventList[i]);
111- AscendC::WaitFlag<AscendC::HardEvent::MTE3_V>(i);115+ AscendC::WaitFlag<AscendC::HardEvent::V_MTE2>(l1BEventList[i]);
112 }116 }
113 }117 }
114 118 
115- template <class TensorBIn, class TensorBOut, class ActualBlockShape>119+ template <class TensorBIn, class ActualBlockShape>
116 CATLASS_DEVICE void operator()(120 CATLASS_DEVICE void operator()(
117- const TensorBIn& bGlobal, const TensorBOut& bLocal, const ActualBlockShape& actualBlockShape,121+ const TensorBIn& bGlobal, const ActualBlockShape& actualBlockShape, const Params& params)
118- const Params& params)
119 {122 {
120- nL1Len_ = actualBlockShape.n();123+ uint32_t kSize = tla::get<0>(params.layoutPrologueB.originShape());
121- uint64_t kTileCount = CeilDiv(kSize_, static_cast<uint64_t>(tla::get<2>(params.tileShapeL1)));124+ uint32_t kL1Size = tla::get<2>(params.tileShapeL1);
122- for (uint64_t kLoopIdx = 0; kLoopIdx < kTileCount; kLoopIdx++) {125+ uint32_t kGmLoop = CeilDiv(kSize, static_cast<uint64_t>(tla::get<2>(params.tileShapeL1)));
123- AscendC::CrossCoreWaitFlag<SYNC_MODE, PIPE_MTE3>(AIV_SYNC_AIC_FLAG);126+ 
124- kGmOffset_ = kLoopIdx * kL1Size_;127+ uint32_t kL1TileShape;
125- kL1Len_ = Min(kSize_ - kGmOffset_, kL1Size_);128+ uint32_t kL1TileCoord;
126- auto tensorBlockB = GetTile(129+ uint32_t kL1Offset;
127- bGlobal, tla::MakeCoord(kGmOffset_, 0),130+ for (uint64_t kLoopIdx = 0; kLoopIdx < kGmLoop; kLoopIdx++) {
128- tla::MakeShape(static_cast<uint64_t>(kL1Len_), static_cast<uint64_t>(nL1Len_)));131+ auto l1BLayout = tla::MakeLayout<ElementOut, LayoutOut>(
129- nUbLen_ = nL1Len_;132+ tla::get<2>(params.tileShapeL1), tla::get<1>(params.tileShapeL1));
130- kUbLen_ = kL1Len_;133+ auto tensorL1B = tla::MakeTensor(l1BTensorList[l1BListId], l1BLayout, Arch::PositionL1{});
131- if constexpr (L1B_STAGES == DOUBLE_BUFFER) {134+ kL1TileCoord = kLoopIdx * kL1Size;
132- if (l1BufIdx_ == AscendC::GetSubBlockIdx()) {135+ 
133- ProcessL1NK(tensorBlockB, bLocal);136+ bool disableSubVec = false;
137+ if (kLoopIdx == kGmLoop - 1) {
138+ kL1TileShape = kSize - kLoopIdx * kL1Size;
139+ kL1Offset = RoundUp(kL1TileShape / AscendC::GetSubBlockNum(), 32);
140+ 
141+ // Since kL1Offset is rounded up, kL1TileShape - kL1Offset may be negative.
142+ // Disable the sub-vector in this case.
143+ if (kL1Offset > kL1TileShape && AscendC::GetSubBlockIdx() != 0) {
144+ disableSubVec = true;
145+ }
146+ 
147+ uint32_t klastLoopCoord = kL1TileCoord + kL1Offset * AscendC::GetSubBlockIdx();
148+ uint32_t klastLoopShape = AscendC::GetSubBlockIdx() == 0 ? kL1Offset : kL1TileShape - kL1Offset;
149+ 
150+ auto tileTensorGmB = GetTile(
151+ bGlobal, tla::MakeCoord(klastLoopCoord, 0), tla::MakeShape(klastLoopShape, actualBlockShape.n()));
152+ auto tileTensorL1B = GetTile(
153+ tensorL1B, tla::MakeCoord(kL1Offset * AscendC::GetSubBlockIdx(), 0),
154+ tla::MakeShape(klastLoopShape, actualBlockShape.n()));
155+ 
156+ if constexpr (std::is_same_v<LayoutIn, layout::Weight4BitnZ>) {
157+ ProcessL1Nz(tileTensorL1B, tileTensorGmB, disableSubVec);
158+ } else {
159+ ProcessL1Nd(tileTensorL1B, tileTensorGmB, disableSubVec);
160+ }
161+ } else {
162+ kL1TileShape = kL1Size / AscendC::GetSubBlockNum();
163+ kL1Offset = AscendC::GetSubBlockIdx() * kL1TileShape;
164+ 
165+ auto tileTensorGmB = GetTile(
166+ bGlobal, tla::MakeCoord(kL1TileCoord + kL1Offset, 0),
167+ tla::MakeShape(kL1TileShape, actualBlockShape.n()));
168+ auto tileTensorL1B = GetTile(
169+ tensorL1B, tla::MakeCoord(kL1Offset * AscendC::GetSubBlockIdx(), 0),
170+ tla::MakeShape(kL1TileShape, actualBlockShape.n()));
171+ 
172+ if constexpr (std::is_same_v<LayoutIn, layout::Weight4BitnZ>) {
173+ ProcessL1Nz(tileTensorL1B, tileTensorGmB);
174+ } else {
175+ ProcessL1Nd(tileTensorL1B, tileTensorGmB);
134 }176 }
135- } else if (AscendC::GetSubBlockIdx() == 0) {
136- ProcessL1NK(tensorBlockB, bLocal);
137 }177 }
138- l1BufIdx_ = (l1BufIdx_ + 1) % L1B_STAGES;178+ l1BListId = (l1BListId + 1 < L1B_STAGES) ? (l1BListId + 1) : 0;
139- AscendC::CrossCoreSetFlag<SYNC_MODE, PIPE_MTE3>(AIV_SYNC_AIC_FLAG);
140 }179 }
141 }180 }
142 181 
143- template <class TensorBIn, class TensorBOut>182+ template <class TensorOut, class TensorIn>
144- __aicore__ inline void ProcessL1NK(const TensorBIn& tensorBlockB, const TensorBOut& tensorL1B)183+ __aicore__ inline void ProcessL1Nz(const TensorOut& tensorOut, const TensorIn& tensorIn, bool disableSubVec = false)
145 {184 {
146- int32_t nFactor = CeilDiv(nL1Len_, nUbSize_);185+ AscendC::WaitFlag<AscendC::HardEvent::V_MTE2>(l1BEventList[l1BListId]);
147- int32_t kFactor = CeilDiv(kL1Len_, kUbSize_);186+ uint32_t shapeN;
148- for (int32_t nLoopIdx = 0; nLoopIdx < nFactor; nLoopIdx++) {187+ uint32_t shapeK;
149- nL1Offset_ = nLoopIdx * nUbSize_;188+ 
150- nUbLen_ = Min(nL1Len_ - static_cast<int32_t>(nL1Offset_), nUbSize_);189+ shapeN = tla::get<1, 0>(tensorIn.shape()) * tla::get<1, 1>(tensorIn.shape()); // nZ
151- for (int32_t kLoopIdx = 0; kLoopIdx < kFactor; kLoopIdx++) {190+ shapeK = tla::get<0, 0>(tensorIn.shape()) * tla::get<0, 1>(tensorIn.shape());
152- kL1Offset_ = kLoopIdx * kUbSize_;191+ 
153- kUbLen_ = Min(kL1Len_ - static_cast<int32_t>(kL1Offset_), kUbSize_);192+ auto layoutCastIn = tla::MakeLayout<ElementOut, layout::Weight4BitnZ>(shapeK, shapeN);
154- int64_t l1Offset = (l1BufIdx_ & 0x1) * L1_BUFFER_HALF_SIZE / sizeof(ElementOut) +193+ auto layoutCastOut = tla::MakeLayout<ElementOut, layout::nZ>(shapeK, shapeN);
155- RoundUp(nL1Size_, AscendC::BLOCK_CUBE) * kL1Offset_ +194+ 
156- nL1Offset_ * AscendC::ONE_BLK_SIZE;195+ auto tensorCastIn = tla::MakeTensor(ubCastInTensor[l1BListId], layoutCastIn, Catlass::Arch::PositionUB{});
157- ProcessL1(tensorBlockB, l1Offset, tensorL1B);196+ auto tensorCastOut = tla::MakeTensor(ubCastOutTensor[l1BListId], layoutCastOut, Catlass::Arch::PositionUB{});
158- }197+ 
198+ if (!disableSubVec) {
199+ CopyNzInWeightTensor(tensorCastIn, tensorIn);
159 }200 }
201+ 
202+ AscendC::SetFlag<AscendC::HardEvent::MTE2_V>(l1BEventList[l1BListId]);
203+ AscendC::WaitFlag<AscendC::HardEvent::MTE2_V>(l1BEventList[l1BListId]);
204+ 
205+ AscendC::WaitFlag<AscendC::HardEvent::MTE3_V>(l1BEventList[l1BListId]);
206+ 
207+ if (!disableSubVec) {
208+ AntiQuantComputeNz(tensorCastOut, tensorCastIn);
209+ }
210+ 
211+ AscendC::SetFlag<AscendC::HardEvent::V_MTE3>(l1BEventList[l1BListId]);
212+ AscendC::WaitFlag<AscendC::HardEvent::V_MTE3>(l1BEventList[l1BListId]);
213+ 
214+ AscendC::SetFlag<AscendC::HardEvent::V_MTE2>(l1BEventList[l1BListId]);
215+ 
216+ AscendC::CrossCoreWaitFlag<SYNC_MODE, PIPE_MTE3>(AIC_SYNC_AIV_FLAG + l1BListId);
217+ 
218+ if (!disableSubVec) {
219+ CopyUb2L1(tensorOut, tensorCastOut);
220+ }
221+ 
222+ AscendC::CrossCoreSetFlag<SYNC_MODE, PIPE_MTE3>(AIV_SYNC_AIC_FLAG + l1BListId);
223+ AscendC::SetFlag<AscendC::HardEvent::MTE3_V>(l1BEventList[l1BListId]);
160 }224 }
161 225 
162- template <class TensorBIn, class TensorBOut>226+ template <class TensorOut, class TensorIn>
163- __aicore__ inline void ProcessL1(const TensorBIn& tensorBlockB, int64_t l1Offset, const TensorBOut& tensorL1B)227+ __aicore__ inline void ProcessL1Nd(const TensorOut& tensorOut, const TensorIn& tensorIn, bool disableSubVec = false)
164 {228 {
165- ubBufIdx_ = ubBufIdx_ % L1B_STAGES;229+ AscendC::WaitFlag<AscendC::HardEvent::V_MTE2>(l1BEventList[l1BListId]);
166- AscendC::WaitFlag<AscendC::HardEvent::V_MTE2>(ubBufIdx_);230+ 
167- CopyInTensorWeight(tensorBlockB);231+ uint32_t shapeN = tla::get<1>(tensorIn.shape());
168- AscendC::SetFlag<AscendC::HardEvent::MTE2_V>(ubBufIdx_);232+ uint32_t shapeK = tla::get<0>(tensorIn.shape());
169- AscendC::WaitFlag<AscendC::HardEvent::MTE3_V>(ubBufIdx_);233+ 
170- AscendC::WaitFlag<AscendC::HardEvent::MTE2_V>(ubBufIdx_);234+ auto layoutCastIn = tla::MakeLayout<ElementIn, layout::ColumnMajor>(shapeK, shapeN);
171- AntiQuantCompute();235+ auto layoutCastOut = tla::MakeLayout<ElementOut, layout::nZ>(shapeK, shapeN);
172- AscendC::SetFlag<AscendC::HardEvent::V_MTE3>(ubBufIdx_);236+ 
173- AscendC::SetFlag<AscendC::HardEvent::V_MTE2>(ubBufIdx_);237+ auto tensorCastIn = tla::MakeTensor(ubCastInTensor[l1BListId], layoutCastIn, Catlass::Arch::PositionUB{});
174- AscendC::WaitFlag<AscendC::HardEvent::V_MTE3>(ubBufIdx_);238+ auto tensorCastOut = tla::MakeTensor(ubCastOutTensor[l1BListId], layoutCastOut, Catlass::Arch::PositionUB{});
175- uint64_t weightOutUbOffset = ubBufIdx_ * (vecWeightOutLen_ / sizeof(ElementOut) / L1B_STAGES);239+ 
176- CopyVecOut2L1(l1Offset, weightOutUb_[weightOutUbOffset], tensorL1B);240+ if (!disableSubVec) {
177- AscendC::SetFlag<AscendC::HardEvent::MTE3_V>(ubBufIdx_);241+ CopyNdInWeightTensor(tensorCastIn, tensorIn);
178- ubBufIdx_++;242+ }
243+ 
244+ AscendC::SetFlag<AscendC::HardEvent::MTE2_V>(l1BEventList[l1BListId]);
245+ AscendC::WaitFlag<AscendC::HardEvent::MTE3_V>(l1BEventList[l1BListId]);
246+ AscendC::WaitFlag<AscendC::HardEvent::MTE2_V>(l1BEventList[l1BListId]);
247+ 
248+ if (!disableSubVec) {
249+ AntiQuantComputeNd(tensorCastOut, tensorCastIn);
250+ }
251+ 
252+ AscendC::SetFlag<AscendC::HardEvent::V_MTE3>(l1BEventList[l1BListId]);
253+ AscendC::SetFlag<AscendC::HardEvent::V_MTE2>(l1BEventList[l1BListId]);
254+ AscendC::WaitFlag<AscendC::HardEvent::V_MTE3>(l1BEventList[l1BListId]);
255+ 
256+ AscendC::CrossCoreWaitFlag<SYNC_MODE, PIPE_MTE3>(AIC_SYNC_AIV_FLAG + l1BListId);
257+ 
258+ if (!disableSubVec) {
259+ CopyUb2L1(tensorOut, tensorCastOut);
260+ }
261+ 
262+ AscendC::CrossCoreSetFlag<SYNC_MODE, PIPE_MTE3>(AIV_SYNC_AIC_FLAG + l1BListId);
263+ AscendC::SetFlag<AscendC::HardEvent::MTE3_V>(l1BEventList[l1BListId]);
179 }264 }
180 265 
181- template <class TensorBIn>266+ template <class TensorDst, class TensorSrc>
182- __aicore__ inline void CopyInTensorWeight(const TensorBIn& tensorBlockB)267+ __aicore__ inline void CopyNzInWeightTensor(TensorDst const& dstTensor, TensorSrc const& srcTensor)
268+ {
269+ uint32_t blockCount = tla::get<0, 1>(srcTensor.shape());
270+ uint32_t blockLen = tla::get<0, 1>(dstTensor.stride()) >> 1;
271+ AscendC::DataCopyExtParams repeatParams;
272+ 
273+ repeatParams.blockCount = blockCount;
274+ repeatParams.blockLen = blockLen;
275+ repeatParams.srcStride = (tla::get<0, 1>(srcTensor.stride()) - tla::get<0, 1>(dstTensor.stride())) >> 1;
276+ repeatParams.dstStride = 0;
277+ 
278+ auto dstOffset = dstTensor.layout()(dstTensor.coord());
279+ auto srcOffset = srcTensor.layout()(srcTensor.coord());
280+ AscendC::DataCopyPadExtParams<typename TensorDst::Element> padParams;
281+ AscendC::DataCopyPad(dstTensor.data()[dstOffset], srcTensor.data()[srcOffset], repeatParams, padParams);
282+ }
283+ 
284+ template <class TensorDst, class TensorSrc>
285+ __aicore__ inline void CopyNdInWeightTensor(TensorDst const& dstTensor, TensorSrc const& srcTensor)
183 {286 {
184 AscendC::DataCopyExtParams intriParams;287 AscendC::DataCopyExtParams intriParams;
185 intriParams.dstStride = 0;288 intriParams.dstStride = 0;
186 AscendC::DataCopyPadExtParams<ElementIn> padParams;289 AscendC::DataCopyPadExtParams<ElementIn> padParams;
187- intriParams.blockCount = nUbLen_;290+ intriParams.blockCount = tla::get<1>(srcTensor.shape());
188- intriParams.blockLen = CeilDiv(kUbLen_, 2);291+ intriParams.blockLen = CeilDiv(tla::get<0>(srcTensor.shape()), 2);
189- intriParams.srcStride = CeilDiv(kSize_, 2) - CeilDiv(kUbLen_, 2);292+ intriParams.srcStride =
190- uint64_t weightInOffset = ubBufIdx_ * (vecWeightInLen_ << INT4_DTYPE_PARAM) / L1B_STAGES;293+ CeilDiv(tla::get<1>(srcTensor.stride()), 2) - CeilDiv(tla::get<0>(srcTensor.shape()), 2);
191- auto subTensorBlockB = tensorBlockB(tla::MakeCoord(kL1Offset_, nL1Offset_));294+ auto dstOffset = dstTensor.layout()(dstTensor.coord());
192- DataCopyPad(weightInUb_[weightInOffset], subTensorBlockB, intriParams, padParams);295+ auto srcOffset = srcTensor.layout()(srcTensor.coord());
296+ 
297+ AscendC::DataCopyPad(dstTensor.data()[dstOffset], srcTensor.data()[srcOffset], intriParams, padParams);
193 }298 }
194 299 
195- template <class TensorBOut>300+ template <class TensorCastOut, class TensorCastIn>
196- __aicore__ inline void CopyVecOut2L1(301+ __aicore__ inline void AntiQuantComputeNz(const TensorCastOut& tensorCastOut, const TensorCastIn& tensorCastIn)
197- int64_t l1Offset, const AscendC::LocalTensor<ElementOut>& ubLocal, const TensorBOut& tensorL1B)
198 {302 {
199- AscendC::DataCopyParams params;303+ static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits<ElementIn>::value;
200- params.blockLen = nUbLen_;304+ 
201- params.blockCount = CeilDiv(kUbLen_, static_cast<int32_t>(GROUP_SIZE));305+ VfParamsNz params;
202- params.srcStride = 1 + RoundUp(nUbLen_, AscendC::BLOCK_CUBE) - nUbLen_;306+ params.weightInUbAddr = (__ubuf__ ElementIn*)tensorCastIn.data().GetPhyAddr();
203- params.dstStride = RoundUp(nL1Size_, AscendC::BLOCK_CUBE) - nUbLen_;307+ params.weightOutUbAddr = (__ubuf__ ElementOut*)tensorCastOut.data().GetPhyAddr();
204- DataCopy(l1Local_[l1Offset], ubLocal, params);308+ 
309+ params.loopKNum = tla::get<0, 1>(tensorCastIn.shape());
310+ params.nRealSizeAlign = tla::get<1, 1>(tensorCastIn.shape()) * AscendC::BLOCK_CUBE;
311+ params.innerDstStride = AscendC::GetVecLen();
312+ params.innerLoopNum = (params.nRealSizeAlign * tla::get<0, 0>(tensorCastIn.shape())) /
313+ static_cast<uint64_t>(AscendC::GetVecLen());
314+ params.loopKDstStride = params.innerLoopNum * params.innerDstStride;
315+ 
316+ RegComputeNz<TensorCastOut, TensorCastIn>(params);
205 }317 }
206 318 
207- __aicore__ inline void AntiQuantCompute()319+ template <class TensorCastOut, class TensorCastIn>
320+ __simd_vf__ inline void RegComputeNz(VfParamsNz params)
208 {321 {
209- uint64_t weightOutUbOffset;322+ AscendC::Reg::RegTensor<int8_t> wShrReg;
210- uint64_t weightInUbOffset;323+ AscendC::Reg::RegTensor<int8_t> wShlReg;
211- weightOutUbOffset = ubBufIdx_ * (vecWeightOutLen_ / sizeof(ElementOut) / L1B_STAGES);324+ AscendC::Reg::RegTensor<int8_t> wAndReg;
212- weightInUbOffset = ubBufIdx_ * (vecWeightInLen_ << INT4_DTYPE_PARAM) / L1B_STAGES;325+ AscendC::Reg::RegTensor<int8_t> wLoad;
213- weightInUbBaseAddr_ = (__ubuf__ int8_t*)weightInUb_[weightInUbOffset].GetPhyAddr();326+ AscendC::Reg::RegTensor<int8_t> wShl;
214- weightOutUbAddr_ = (__ubuf__ ElementOut*)weightOutUb_[weightOutUbOffset].GetPhyAddr();327+ AscendC::Reg::RegTensor<int8_t> wShr0;
328+ AscendC::Reg::RegTensor<int8_t> wShr1;
329+ AscendC::Reg::RegTensor<int8_t> wSel;
330+ AscendC::Reg::RegTensor<int8_t> wAnd;
215 331 
216- uint16_t blockStride = RoundUp(nUbLen_, AscendC::BLOCK_CUBE) + 1;332+ AscendC::Reg::MaskReg preg = AscendC::Reg::CreateMask<uint8_t, AscendC::Reg::MaskPattern::ALL>();
217- weightOutUbAddr1_ = weightOutUbAddr_ + VEC_MAX_ELEM_B8 * blockStride;333+ AscendC::Reg::MaskReg pregVsel = AscendC::Reg::CreateMask<uint16_t, AscendC::Reg::MaskPattern::ALL>();
218- AntiQuantComputeNormal();334+ 
335+ AscendC::Reg::Duplicate<int8_t, AscendC::Reg::MaskMergeMode::ZEROING>(wShrReg, E2M1_SHIFT_RIGHT_SIZE, preg);
336+ AscendC::Reg::Duplicate<int8_t, AscendC::Reg::MaskMergeMode::ZEROING>(wShlReg, SHIFT_LEFT_SIZE, preg);
337+ AscendC::Reg::Duplicate<int8_t, AscendC::Reg::MaskMergeMode::ZEROING>(wAndReg, E2M1_AND_MASK, preg);
338+ 
339+ for (uint16_t loopKIdx = 0; loopKIdx < params.loopKNum; ++loopKIdx) {
340+ for (uint16_t innerLoopIdx = 0; innerLoopIdx < params.innerLoopNum; ++innerLoopIdx) {
341+ // DIST_US_B8 load mode expands each packed B4 byte into lane-aligned B8 slots.
342+ // Packed B4 address offset (bytes) = logical element index >> 1.
343+ AscendC::Reg::AddrReg aregWeightB8In = AscendC::Reg::CreateAddrReg<uint8_t>(
344+ loopKIdx, (C0_SIZE_B8 * params.nRealSizeAlign) >> 1, innerLoopIdx, AscendC::GetVecLen() >> 1);
345+ AscendC::Reg::LoadAlign<uint8_t, AscendC::Reg::LoadDist::DIST_US_B8>(
346+ (AscendC::Reg::RegTensor<uint8_t>&)wLoad, (__ubuf__ uint8_t*&)params.weightInUbAddr,
347+ aregWeightB8In);
348+ 
349+ AscendC::Reg::ShiftRight(wShr0, wLoad, wShrReg, preg);
350+ AscendC::Reg::ShiftLeft(wShl, wLoad, wShlReg, preg);
351+ AscendC::Reg::ShiftRight(wShr1, wShl, wShrReg, preg);
352+ AscendC::Reg::Select(wSel, wShr1, wShr0, pregVsel);
353+ AscendC::Reg::And(wAnd, wSel, wAndReg, preg);
354+ 
355+ AscendC::Reg::AddrReg aregWeightB8Out = AscendC::Reg::CreateAddrReg<uint8_t>(
356+ loopKIdx, params.loopKDstStride, innerLoopIdx, params.innerDstStride);
357+ AscendC::Reg::StoreAlign<uint8_t, AscendC::Reg::StoreDist::DIST_NORM_B8>(
358+ (__ubuf__ uint8_t*&)params.weightOutUbAddr, (AscendC::Reg::RegTensor<uint8_t>&)wAnd,
359+ aregWeightB8Out, preg);
360+ }
361+ }
219 }362 }
220 363 
221- __aicore__ inline void AntiQuantComputeNormal()364+ template <class TensorCastOut, class TensorCastIn>
365+ __aicore__ inline void AntiQuantComputeNd(const TensorCastOut& tensorCastOut, const TensorCastIn& tensorCastIn)
222 {366 {
223- VfParamsNormal wParams;367+ uint32_t weightOutUbOffset;
224- wParams.outExtend = static_cast<uint16_t>(nUbLen_);368+ uint32_t weightInUbOffset;
225- wParams.innerExtend = CeilDiv(RoundUp(kUbLen_, UB_ALIGN_SIZE_FOR_4BITS), VECTOR_REG_WIDTH_FOR_4BITS);369+ weightOutUbOffset = l1BListId * (vecWeightOutLen / sizeof(ElementOut) / L1B_STAGES);
226- wParams.dataBlockStride = RoundUp(nUbLen_, AscendC::BLOCK_CUBE) + 1;370+ weightInUbOffset = l1BListId * (vecWeightInLen << INT4_DTYPE_PARAM) / L1B_STAGES;
371+ __ubuf__ int8_t* weightInUbAddr = (__ubuf__ int8_t*)tensorCastIn.data().GetPhyAddr();
372+ __ubuf__ ElementOut* weightOutUbAddr = (__ubuf__ ElementOut*)tensorCastOut.data().GetPhyAddr();
373+ 
374+ int32_t kUbLen = tla::get<0>(tensorCastIn.shape());
375+ int32_t nUbLen = tla::get<1>(tensorCastIn.shape());
376+ uint16_t blockStride = RoundUp(nUbLen, AscendC::BLOCK_CUBE);
377+ __ubuf__ ElementOut* weightOutUbAddr1 = weightOutUbAddr + VEC_MAX_ELEM_B8 * blockStride;
378+ 
379+ VfParamsNd wParams;
380+ wParams.outExtend = static_cast<uint16_t>(nUbLen);
381+ wParams.innerExtend = CeilDiv(RoundUp(kUbLen, K_ALIGN_SIZE), VECTOR_REG_WIDTH_FOR_4BITS);
382+ wParams.dataBlockStride = RoundUp(nUbLen, AscendC::BLOCK_CUBE);
227 wParams.repeatStride = wParams.dataBlockStride * AscendC::BLOCK_CUBE;383 wParams.repeatStride = wParams.dataBlockStride * AscendC::BLOCK_CUBE;
228 wParams.outDimOffset =384 wParams.outDimOffset =
229 AscendC::ONE_BLOCK_SIZE - wParams.innerExtend * wParams.repeatStride * AscendC::ONE_BLOCK_SIZE;385 AscendC::ONE_BLOCK_SIZE - wParams.innerExtend * wParams.repeatStride * AscendC::ONE_BLOCK_SIZE;
230 wParams.maskB8Tail0 =386 wParams.maskB8Tail0 =
231- Min(kUbLen_ % VECTOR_REG_WIDTH_FOR_4BITS, static_cast<int32_t>(AscendC::VECTOR_REG_WIDTH)) +387+ Min(kUbLen % VECTOR_REG_WIDTH_FOR_4BITS, static_cast<int32_t>(AscendC::VECTOR_REG_WIDTH)) +
232- kUbLen_ / VECTOR_REG_WIDTH_FOR_4BITS * AscendC::VECTOR_REG_WIDTH;388+ kUbLen / VECTOR_REG_WIDTH_FOR_4BITS * AscendC::VECTOR_REG_WIDTH;
233 wParams.maskB8Tail1 =389 wParams.maskB8Tail1 =
234- Max(kUbLen_ % VECTOR_REG_WIDTH_FOR_4BITS - static_cast<int32_t>(AscendC::VECTOR_REG_WIDTH), 0) +390+ Max(kUbLen % VECTOR_REG_WIDTH_FOR_4BITS - static_cast<int32_t>(AscendC::VECTOR_REG_WIDTH), 0) +
235- kUbLen_ / VECTOR_REG_WIDTH_FOR_4BITS * AscendC::VECTOR_REG_WIDTH;391+ kUbLen / VECTOR_REG_WIDTH_FOR_4BITS * AscendC::VECTOR_REG_WIDTH;
236- wParams.weightInUbBaseAddr = weightInUbBaseAddr_;392+ wParams.kUbLen = kUbLen;
237- wParams.weightOutUbAddr = weightOutUbAddr_;393+ wParams.weightInUbBaseAddr = weightInUbAddr;
238- wParams.weightOutUbAddr1 = weightOutUbAddr1_;394+ wParams.weightOutUbAddr = weightOutUbAddr;
239- RegCompute(wParams);395+ wParams.weightOutUbAddr1 = weightOutUbAddr1;
396+ RegComputeNd(wParams);
240 }397 }
241 398 
242- __simd_vf__ inline void RegCompute(const VfParamsNormal wParams)399+ __simd_vf__ inline void RegComputeNd(const VfParamsNd wParams)
243 {400 {
244 __ubuf__ ElementOut* weightOutUbAddr = wParams.weightOutUbAddr;401 __ubuf__ ElementOut* weightOutUbAddr = wParams.weightOutUbAddr;
245 __ubuf__ ElementOut* weightOutUbAddr1 = wParams.weightOutUbAddr1;402 __ubuf__ ElementOut* weightOutUbAddr1 = wParams.weightOutUbAddr1;
@@ -261,7 +418,8 @@ public:
261 AscendC::MicroAPI::MaskReg MaskRegB8Tail0 = AscendC::MicroAPI::UpdateMask<uint8_t>(maskWeight0Tmp);418 AscendC::MicroAPI::MaskReg MaskRegB8Tail0 = AscendC::MicroAPI::UpdateMask<uint8_t>(maskWeight0Tmp);
262 AscendC::MicroAPI::MaskReg MaskRegB8Tail1 = AscendC::MicroAPI::UpdateMask<uint8_t>(maskWeight1Tmp);419 AscendC::MicroAPI::MaskReg MaskRegB8Tail1 = AscendC::MicroAPI::UpdateMask<uint8_t>(maskWeight1Tmp);
263 AscendC::MicroAPI::AddrReg aregWeightB8 = AscendC::MicroAPI::CreateAddrReg<uint8_t>(420 AscendC::MicroAPI::AddrReg aregWeightB8 = AscendC::MicroAPI::CreateAddrReg<uint8_t>(
264- outIdx, RoundUp(kUbLen_, static_cast<int32_t>(K_ALIGN_SIZE)) >> 1, repeatIdx, VEC_MAX_ELEM_B8);421+ outIdx, RoundUp(wParams.kUbLen, static_cast<int32_t>(K_ALIGN_SIZE)) >> 1, repeatIdx,
422+ VEC_MAX_ELEM_B8);
265 AscendC::MicroAPI::LoadAlign(wLoad0, (__ubuf__ uint8_t*&)wParams.weightInUbBaseAddr, aregWeightB8);423 AscendC::MicroAPI::LoadAlign(wLoad0, (__ubuf__ uint8_t*&)wParams.weightInUbBaseAddr, aregWeightB8);
266 // 提取E/M424 // 提取E/M
267 AscendC::MicroAPI::ShiftRight(wShr, wLoad0, wdup0, preg); // vr1425 AscendC::MicroAPI::ShiftRight(wShr, wLoad0, wdup0, preg); // vr1
@@ -292,106 +450,54 @@ public:
292 }450 }
293 }451 }
294 452 
295- __aicore__ inline void AntiQuantComputeNKMxNz()453+ template <class TensorDst, class TensorSrc>
454+ __aicore__ inline void CopyUb2L1(TensorDst const& dstTensor, TensorSrc const& srcTensor)
296 {455 {
297- static_assert(456+ static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits<ElementOut>::value;
298- AscendC::Std::is_one_of_v<ElementIn, fp4x2_e2m1_t, fp4x2_e1m2_t>,457+ int64_t brustCount = tla::get<0, 1>(srcTensor.shape());
299- "only support fp4x2_e2m1_t and fp4x2_e1m2_t");458+ int64_t burstLen = tla::get<1, 0>(srcTensor.shape()) * tla::get<1, 1>(srcTensor.shape());
300- VfParamsNz wParams;459+ 
301- wParams.shiftLeftSize =460+ AscendC::DataCopyParams dataCopyParams(
302- AscendC::IsSameType<ElementIn, fp4x2_e2m1_t>::value ? E2M1_SHIFT_LEFT_SIZE : E1M2_SHIFT_LEFT_SIZE;461+ brustCount, burstLen, (tla::get<0, 1>(srcTensor.stride()) / ELE_NUM_PER_C0 - burstLen),
303- wParams.andMask = AscendC::IsSameType<ElementIn, fp4x2_e2m1_t>::value ? E2M1_AND_MASK : E1M2_AND_MASK;462+ (tla::get<0, 1>(dstTensor.stride()) / ELE_NUM_PER_C0 - burstLen));
304- wParams.innerExtend =463+ 
305- CeilDiv(kUbLen_ * RoundUp(nUbLen_, AscendC::BLOCK_CUBE), static_cast<int32_t>(AscendC::VECTOR_REG_WIDTH));464+ auto dstOffset = dstTensor.layout()(dstTensor.coord());
306- wParams.innerDstExtend = AscendC::VECTOR_REG_WIDTH * L1B_STAGES;465+ auto srcOffset = srcTensor.layout()(srcTensor.coord());
307- wParams.innerSrcExtend = AscendC::VECTOR_REG_WIDTH >> 1;466+ 
308- wParams.weightInUbBaseAddr = weightInUbBaseAddr_;467+ AscendC::DataCopy(dstTensor.data()[dstOffset], srcTensor.data()[srcOffset], dataCopyParams);
309- wParams.weightOutUbAddr = weightOutUbAddr_;
310- RegComputeNkNz(wParams);
311 }468 }
312 469 
313- __simd_vf__ inline void RegComputeNkNz(const VfParamsNz wParams)
314- {
315- AscendC::MicroAPI::RegTensor<int8_t> wdup0, wdup1, wdup2, wLoad0, wShl, wShr0, wShr1, wSel0, sAnd0;
316- AscendC::MicroAPI::MaskReg preg = AscendC::MicroAPI::CreateMask<uint8_t, AscendC::MicroAPI::MaskPattern::ALL>();
317- AscendC::MicroAPI::MaskReg pregVsel =
318- AscendC::MicroAPI::CreateMask<uint16_t, AscendC::MicroAPI::MaskPattern::ALL>();
319- AscendC::MicroAPI::Duplicate<int8_t, AscendC::MicroAPI::MaskMergeMode::ZEROING>(
320- wdup0, wParams.shiftLeftSize, preg);
321- AscendC::MicroAPI::Duplicate<int8_t, AscendC::MicroAPI::MaskMergeMode::ZEROING>(wdup1, SHIFT_RIGHT_SIZE, preg);
322- AscendC::MicroAPI::Duplicate<int8_t, AscendC::MicroAPI::MaskMergeMode::ZEROING>(wdup2, wParams.andMask, preg);
323- for (uint16_t repeatIdx = 0; repeatIdx < wParams.innerExtend; ++repeatIdx) {
324- AscendC::MicroAPI::AddrReg aregWeightB8In =
325- AscendC::MicroAPI::CreateAddrReg<uint8_t>(repeatIdx, wParams.innerSrcExtend);
326- AscendC::MicroAPI::AddrReg aregWeightB8Out =
327- AscendC::MicroAPI::CreateAddrReg<uint8_t>(repeatIdx, wParams.innerDstExtend);
328- AscendC::MicroAPI::LoadAlign<uint8_t, AscendC::MicroAPI::LoadDist::DIST_US_B8>(
329- (AscendC::MicroAPI::RegTensor<uint8_t>&)wLoad0, (__ubuf__ uint8_t*&)wParams.weightInUbBaseAddr,
330- aregWeightB8In);
331- AscendC::MicroAPI::ShiftRight(wShr0, wLoad0, wdup0, preg);
332- AscendC::MicroAPI::ShiftLeft(wShl, wLoad0, wdup1, preg);
333- AscendC::MicroAPI::ShiftRight(wShr1, wShl, wdup0, preg);
334- AscendC::MicroAPI::Select(wSel0, wShr1, wShr0, pregVsel);
335- AscendC::MicroAPI::And(sAnd0, wSel0, wdup2, preg);
336- AscendC::MicroAPI::StoreAlign<uint8_t, AscendC::MicroAPI::StoreDist::DIST_NORM_B8>(
337- (__ubuf__ uint8_t*&)wParams.weightOutUbAddr, (AscendC::MicroAPI::RegTensor<uint8_t>&)sAnd0,
338- aregWeightB8Out, preg);
339- }
340- }
341- 
342- static constexpr int64_t DOUBLE_BUFFER = 2;
343- static constexpr uint64_t SYNC_MODE4 = 4;
344- static constexpr uint64_t L1_BUFFER_HALF_SIZE = 256 * 1024;
345 static constexpr uint64_t K_ALIGN_SIZE = 64;470 static constexpr uint64_t K_ALIGN_SIZE = 64;
346- static constexpr uint64_t INT4_DTYPE_PARAM = 1;
347- static constexpr uint64_t BLOCK_NUM_REG = BYTE_PER_VECTOR_FRACTAL / BYTE_PER_BLK;
348- static constexpr uint64_t SINGLE_BUFFER = 1;
349- static constexpr uint64_t GROUP_SIZE = 32;
350 static constexpr int32_t C0_SIZE_B8 = 32;471 static constexpr int32_t C0_SIZE_B8 = 32;
351- static constexpr int32_t UB_ALIGN_SIZE_FOR_4BITS = 64;472+ static constexpr uint64_t INT4_DTYPE_PARAM = 1;
473+ static constexpr int32_t VECTOR_REG_WIDTH_FOR_4BITS = 512;
474+ static constexpr int32_t VEC_MAX_ELEM_B8 = BYTE_PER_VECTOR_FRACTAL / sizeof(ElementOut);
475+ 
352 static constexpr uint32_t DUP_CONFIG_2 = 0x2;476 static constexpr uint32_t DUP_CONFIG_2 = 0x2;
353 static constexpr uint32_t DUP_CONFIG_MODE_1C = 0x1C;477 static constexpr uint32_t DUP_CONFIG_MODE_1C = 0x1C;
354 static constexpr uint32_t DUP_CONFIG_4 = 0x4;478 static constexpr uint32_t DUP_CONFIG_4 = 0x4;
355 static constexpr uint32_t DUP_FLAG_80 = 0x80;479 static constexpr uint32_t DUP_FLAG_80 = 0x80;
356- static constexpr uint32_t E1M2_SHIFT_LEFT_SIZE = 0x3;480+ static constexpr uint32_t E2M1_SHIFT_RIGHT_SIZE = 0x2;
357- static constexpr uint32_t E1M2_AND_MASK = 0x8E;
358- static constexpr uint32_t E2M1_SHIFT_LEFT_SIZE = 0x2;
359 static constexpr uint32_t E2M1_AND_MASK = 0x9C;481 static constexpr uint32_t E2M1_AND_MASK = 0x9C;
360 static constexpr uint32_t SHIFT_RIGHT_SIZE = 0x4;482 static constexpr uint32_t SHIFT_RIGHT_SIZE = 0x4;
361- static constexpr int32_t VEC_MAX_ELEM_B8 = BYTE_PER_VECTOR_FRACTAL / sizeof(ElementOut);483+ static constexpr uint32_t SHIFT_LEFT_SIZE = 0x4;
362- static constexpr int32_t VECTOR_REG_WIDTH_FOR_4BITS = 512;484+ 
363- static constexpr int32_t OFFSET_64 = 64;485+ static constexpr int32_t SYNC_MODE = 4;
364- static constexpr int32_t SYNC_MODE = 4;486+ constexpr static uint16_t AIV_SYNC_AIC_FLAG = 6;
365- static constexpr uint16_t AIV_SYNC_AIC_FLAG = 0;487+ constexpr static uint16_t AIC_SYNC_AIV_FLAG = 8;
366 488 
367- uint64_t nSize_;
368- uint64_t kSize_;
369- int32_t nUbSize_;
370- int32_t kUbSize_;
371- int32_t nUbLen_;
372- int32_t kUbLen_;
373- int32_t nBiasUbLen_ = 0;
374- uint64_t nL1Size_;
375- uint64_t kL1Size_;
376- uint64_t kGmOffset_;
377- int32_t nL1Len_;
378- int32_t kL1Len_;
379- uint64_t aL1Size_;
380- uint64_t bL1Size_;
381- uint64_t vecWeightOutLen_;
382- uint64_t vecWeightInLen_;
383- uint64_t vecBiasLen_;
384- uint64_t ubBufIdx_ = 0;
385- int64_t l1BufIdx_ = 0;
386- uint64_t nL1Offset_ = 0;
387- uint64_t kL1Offset_ = 0;
388 static constexpr int64_t UB_STAGES = L1B_STAGES;489 static constexpr int64_t UB_STAGES = L1B_STAGES;
389- __ubuf__ ElementOut* weightOutUbAddr_;490+ 
390- __ubuf__ ElementOut* weightOutUbAddr1_;491+ uint32_t l1BListId{0};
391- __ubuf__ int8_t* weightInUbBaseAddr_;492+ uint32_t vecWeightOutLen;
392- AscendC::LocalTensor<ElementIn> weightInUb_;493+ uint32_t vecWeightInLen;
393- AscendC::LocalTensor<ElementOut> weightOutUb_;494+ int32_t nUbLen;
394- AscendC::LocalTensor<ElementOut> l1Local_;495+ int32_t kUbLen;
496+ int32_t l1BEventList[L1B_STAGES];
497+ 
498+ AscendC::LocalTensor<ElementIn> ubCastInTensor[UB_STAGES];
499+ AscendC::LocalTensor<ElementOut> ubCastOutTensor[UB_STAGES];
500+ AscendC::LocalTensor<ElementOut> l1BTensorList[L1B_STAGES];
395};501};
396} // namespace Catlass::Gemm::Block502} // namespace Catlass::Gemm::Block
397 503 
Minclude/catlass/gemm/dispatch_policy.hpp+5-3
@@ -420,8 +420,9 @@ struct MmadMx : public MmadBase<ArchTag_, false> {
420};420};
421 421 
422template <422template <
423- class ArchTag_, bool ENABLE_UNIT_FLAG_ = false, uint32_t L0C_STAGES_ = 1, bool ENABLE_L1_RESIDENT_ = false,423+ class ArchTag_, bool ENABLE_UNIT_FLAG_ = false, bool ENABLE_L1_RESIDENT_ = false, uint32_t L1_SCALE_FACTOR_K_ = 1,
424- uint32_t L1A_STAGES_ = 2, uint32_t L1B_STAGES_ = 2, uint32_t L0A_STAGES_ = 2, uint32_t L0B_STAGES_ = 2>424+ uint32_t L0C_STAGES_ = 1, uint32_t L1A_STAGES_ = 2, uint32_t L1B_STAGES_ = 1, uint32_t L0A_STAGES_ = 2,
425+ uint32_t L0B_STAGES_ = 2>
425struct MmadA8W4Mx : public MmadBase<ArchTag_, false> {426struct MmadA8W4Mx : public MmadBase<ArchTag_, false> {
426 static constexpr uint32_t L1A_STAGES = L1A_STAGES_;427 static constexpr uint32_t L1A_STAGES = L1A_STAGES_;
427 static constexpr uint32_t L1B_STAGES = L1B_STAGES_;428 static constexpr uint32_t L1B_STAGES = L1B_STAGES_;
@@ -430,9 +431,10 @@ struct MmadA8W4Mx : public MmadBase<ArchTag_, false> {
430 static constexpr uint32_t L0C_STAGES = L0C_STAGES_;431 static constexpr uint32_t L0C_STAGES = L0C_STAGES_;
431 static constexpr bool ENABLE_UNIT_FLAG = ENABLE_UNIT_FLAG_;432 static constexpr bool ENABLE_UNIT_FLAG = ENABLE_UNIT_FLAG_;
432 static constexpr bool ENABLE_L1_RESIDENT = ENABLE_L1_RESIDENT_;433 static constexpr bool ENABLE_L1_RESIDENT = ENABLE_L1_RESIDENT_;
434+ static constexpr uint32_t L1_SCALE_FACTOR_K = L1_SCALE_FACTOR_K_;
433};435};
434 436 
435-template <class ArchTag_, uint32_t L1B_STAGES_ = 1>437+template <class ArchTag_, uint32_t L1B_STAGES_ = 2>
436struct MxA8W4Prologue : public MmadBase<ArchTag_, false> {438struct MxA8W4Prologue : public MmadBase<ArchTag_, false> {
437 static constexpr uint32_t L1B_STAGES = L1B_STAGES_;439 static constexpr uint32_t L1B_STAGES = L1B_STAGES_;
438};440};
Minclude/catlass/gemm/kernel/a8w4_mx_matmul.hpp+30-45
@@ -43,7 +43,8 @@ public:
43 using ElementC = typename BlockMmad::ElementC;43 using ElementC = typename BlockMmad::ElementC;
44 using LayoutC = typename BlockMmad::LayoutC;44 using LayoutC = typename BlockMmad::LayoutC;
45 using ElementBias = typename BlockMmad::ElementBias;45 using ElementBias = typename BlockMmad::ElementBias;
46- using ElementPrologueB = typename BlockMmad::ElementPrologueB;46+ using ElementPrologueB = typename BlockMmad::TileCopy::ElementPrologueB;
47+ using LayoutPrologueB = typename BlockMmad::TileCopy::LayoutPrologueB;
47 using ElementAccumulator = typename BlockMmad::ElementAccumulator;48 using ElementAccumulator = typename BlockMmad::ElementAccumulator;
48 using LayoutTagL1B = typename BlockMmad::LayoutTagL1B;49 using LayoutTagL1B = typename BlockMmad::LayoutTagL1B;
49 50 
@@ -66,7 +67,7 @@ public:
66 GM_ADDR ptrA;67 GM_ADDR ptrA;
67 LayoutA layoutA;68 LayoutA layoutA;
68 GM_ADDR ptrB;69 GM_ADDR ptrB;
69- LayoutB layoutB;70+ LayoutPrologueB layoutPrologueB;
70 GM_ADDR ptrMxScaleA;71 GM_ADDR ptrMxScaleA;
71 LayoutMxScaleA layoutMxScaleA;72 LayoutMxScaleA layoutMxScaleA;
72 GM_ADDR ptrMxScaleB;73 GM_ADDR ptrMxScaleB;
@@ -82,14 +83,15 @@ public:
82 83 
83 CATLASS_HOST_DEVICE84 CATLASS_HOST_DEVICE
84 Params(85 Params(
85- GemmCoord const& problemShape_, GM_ADDR ptrA_, LayoutA layoutA_, GM_ADDR ptrB_, LayoutB layoutB_,86+ GemmCoord const& problemShape_, GM_ADDR ptrA_, LayoutA layoutA_, GM_ADDR ptrB_,
86- GM_ADDR ptrMxScaleA_, LayoutMxScaleA layoutMxScaleA_, GM_ADDR ptrMxScaleB_, LayoutMxScaleB layoutMxScaleB_,87+ LayoutPrologueB layoutPrologueB_, GM_ADDR ptrMxScaleA_, LayoutMxScaleA layoutMxScaleA_,
87- GM_ADDR ptrC_, LayoutC layoutC_, GM_ADDR ptrBias_ = nullptr)88+ GM_ADDR ptrMxScaleB_, LayoutMxScaleB layoutMxScaleB_, GM_ADDR ptrC_, LayoutC layoutC_,
89+ GM_ADDR ptrBias_ = nullptr)
88 : problemShape(problemShape_),90 : problemShape(problemShape_),
89 ptrA(ptrA_),91 ptrA(ptrA_),
90 layoutA(layoutA_),92 layoutA(layoutA_),
91 ptrB(ptrB_),93 ptrB(ptrB_),
92- layoutB(layoutB_),94+ layoutPrologueB(layoutPrologueB_),
93 ptrMxScaleA(ptrMxScaleA_),95 ptrMxScaleA(ptrMxScaleA_),
94 layoutMxScaleA(layoutMxScaleA_),96 layoutMxScaleA(layoutMxScaleA_),
95 ptrMxScaleB(ptrMxScaleB_),97 ptrMxScaleB(ptrMxScaleB_),
@@ -105,7 +107,7 @@ public:
105 uint8_t* ptrA;107 uint8_t* ptrA;
106 LayoutA layoutA;108 LayoutA layoutA;
107 uint8_t* ptrB;109 uint8_t* ptrB;
108- LayoutB layoutB;110+ LayoutPrologueB layoutPrologueB;
109 uint8_t* ptrMxScaleA;111 uint8_t* ptrMxScaleA;
110 LayoutMxScaleA layoutMxScaleA;112 LayoutMxScaleA layoutMxScaleA;
111 uint8_t* ptrMxScaleB;113 uint8_t* ptrMxScaleB;
@@ -127,32 +129,25 @@ public:
127 129 
128 static Params ToUnderlyingArguments(const Arguments& args, uint8_t* workspace)130 static Params ToUnderlyingArguments(const Arguments& args, uint8_t* workspace)
129 {131 {
130- Params params{args.problemShape, args.ptrA, args.layoutA, args.ptrB,132+ Params params{args.problemShape, args.ptrA, args.layoutA, args.ptrB,
131- args.layoutB, args.ptrMxScaleA, args.layoutMxScaleA, args.ptrMxScaleB,133+ args.layoutPrologueB, args.ptrMxScaleA, args.layoutMxScaleA, args.ptrMxScaleB,
132- args.layoutMxScaleB, args.ptrC, args.layoutC, args.ptrBias};134+ args.layoutMxScaleB, args.ptrC, args.layoutC, args.ptrBias};
133 return params;135 return params;
134 }136 }
135 137 
136 // Methods138 // Methods
137 CATLASS_DEVICE139 CATLASS_DEVICE
138- A8W4MxMatmul(uint32_t l1BufAddrStart = 0)140+ A8W4MxMatmul()
139 {141 {
140 Arch::Resource<ArchTag> resource;142 Arch::Resource<ArchTag> resource;
141 if constexpr (tla::detail::isRowMajor<LayoutC>::value) {143 if constexpr (tla::detail::isRowMajor<LayoutC>::value) {
142 AscendC::SetMMLayoutTransform(true);144 AscendC::SetMMLayoutTransform(true);
143 }145 }
144- uint32_t l1Offset = l1BufAddrStart;
145- for (uint32_t i = 0; i < L1B_STAGES; i++) {
146- // Assign L1/L0A/L0B space for each stages
147- l1BTensorList[i] = resource.l1Buf.template GetBufferByByte<ElementB>(l1Offset);
148- l1Offset += L1B_TILE_SIZE;
149- // Assign event ID for each stages
150- l1BEventList[i] = i;
151- // The event id that needs to be set before the loop
152- }
153 if ASCEND_IS_AIC {146 if ASCEND_IS_AIC {
154- AscendC::CrossCoreSetFlag<AIC_SYNC_AIV_MODE, PIPE_MTE1>(AIV_SYNC_AIC_FLAG);147+ for (uint32_t i = 0; i < L1B_STAGES; i++) {
155- AscendC::CrossCoreSetFlag<AIC_SYNC_AIV_MODE, PIPE_MTE1>(AIC_SYNC_AIV_FLAG);148+ AscendC::CrossCoreSetFlag<AIC_SYNC_AIV_MODE, PIPE_MTE1>(AIC_SYNC_AIV_FLAG + i);
149+ AscendC::CrossCoreSetFlag<AIC_SYNC_AIV_MODE, PIPE_MTE1>(AIC_SYNC_AIV_FLAG + FLAG_ID_MAX + i);
150+ }
156 }151 }
157 }152 }
158 153 
@@ -161,7 +156,9 @@ public:
161 ~A8W4MxMatmul()156 ~A8W4MxMatmul()
162 {157 {
163 if ASCEND_IS_AIV {158 if ASCEND_IS_AIV {
164- AscendC::CrossCoreWaitFlag<AIC_SYNC_AIV_MODE, PIPE_MTE3>(AIV_SYNC_AIC_FLAG);159+ for (uint32_t i = 0; i < L1B_STAGES; i++) {
160+ AscendC::CrossCoreWaitFlag<AIC_SYNC_AIV_MODE, PIPE_MTE3>(AIC_SYNC_AIV_FLAG + i);
161+ }
165 }162 }
166 }163 }
167 164 
@@ -175,8 +172,7 @@ public:
175 BlockScheduler matmulBlockScheduler(params.problemShape, MakeCoord(L1_TILE_M, L1_TILE_N));172 BlockScheduler matmulBlockScheduler(params.problemShape, MakeCoord(L1_TILE_M, L1_TILE_N));
176 uint32_t coreLoops = matmulBlockScheduler.GetCoreLoops();173 uint32_t coreLoops = matmulBlockScheduler.GetCoreLoops();
177 174 
178- Arch::Resource<ArchTag> resource;175+ BlockMmad blockMmad(resource);
179- BlockMmad blockMmad(resource, L1B_TILE_SIZE * L1B_STAGES);
180 176 
181 // Represent the full gm177 // Represent the full gm
182 AscendC::GlobalTensor<ElementA> gmA;178 AscendC::GlobalTensor<ElementA> gmA;
@@ -208,9 +204,6 @@ public:
208 auto tensorC = tla::MakeTensor(gmC, params.layoutC, Arch::PositionGM{});204 auto tensorC = tla::MakeTensor(gmC, params.layoutC, Arch::PositionGM{});
209 auto tensorBias = tla::MakeTensor(gmBias, layoutBias, Arch::PositionGM{});205 auto tensorBias = tla::MakeTensor(gmBias, layoutBias, Arch::PositionGM{});
210 206 
211- // make L1 TensorB
212- auto tensorL1B = tla::MakeTensor(l1BTensorList[l1BListId], L1B_LAYOUT, Arch::PositionL1{});
213- 
214 for (uint32_t loopIdx = AscendC::GetBlockIdx(); loopIdx < coreLoops; loopIdx += AscendC::GetBlockNum()) {207 for (uint32_t loopIdx = AscendC::GetBlockIdx(); loopIdx < coreLoops; loopIdx += AscendC::GetBlockNum()) {
215 GemmCoord blockCoord = matmulBlockScheduler.GetBlockCoord(loopIdx);208 GemmCoord blockCoord = matmulBlockScheduler.GetBlockCoord(loopIdx);
216 GemmCoord actualBlockShape = matmulBlockScheduler.GetActualBlockShape(blockCoord);209 GemmCoord actualBlockShape = matmulBlockScheduler.GetActualBlockShape(blockCoord);
@@ -234,13 +227,12 @@ public:
234 227 
235 // Compute block-scoped matrix multiply-add228 // Compute block-scoped matrix multiply-add
236 if constexpr (std::is_void_v<ElementBias>) {229 if constexpr (std::is_void_v<ElementBias>) {
237- blockMmad(230+ blockMmad(tensorBlockA, tensorBlockC, actualBlockShape, tensorBlockMxScaleA, tensorBlockMxScaleB);
238- tensorBlockA, tensorBlockC, actualBlockShape, tensorL1B, tensorBlockMxScaleA, tensorBlockMxScaleB);
239 } else {231 } else {
240 auto tensorBlockBias = GetTile(232 auto tensorBlockBias = GetTile(
241 tensorBias, tla::MakeCoord(blockCoord.n() * L1_TILE_N), tla::MakeShape(actualBlockShape.n()));233 tensorBias, tla::MakeCoord(blockCoord.n() * L1_TILE_N), tla::MakeShape(actualBlockShape.n()));
242 blockMmad(234 blockMmad(
243- tensorBlockA, tensorBlockC, actualBlockShape, tensorL1B, tensorBlockMxScaleA, tensorBlockMxScaleB,235+ tensorBlockA, tensorBlockC, actualBlockShape, tensorBlockMxScaleA, tensorBlockMxScaleB,
244 tensorBlockBias);236 tensorBlockBias);
245 }237 }
246 }238 }
@@ -250,7 +242,7 @@ public:
250 CATLASS_DEVICE void operator()<AscendC::AIV>(Params const& params)242 CATLASS_DEVICE void operator()<AscendC::AIV>(Params const& params)
251 {243 {
252 using PrologueParams = typename BlockPrologue::Params;244 using PrologueParams = typename BlockPrologue::Params;
253- PrologueParams prologueParams{L1TileShape{}, params.layoutB, L1_TILE_N, L1_TILE_K, false};245+ PrologueParams prologueParams{L1TileShape{}, params.layoutPrologueB, false, resource};
254 246 
255 BlockPrologue blockPrologue(prologueParams);247 BlockPrologue blockPrologue(prologueParams);
256 248 
@@ -258,15 +250,12 @@ public:
258 AscendC::GlobalTensor<ElementPrologueB> gmB;250 AscendC::GlobalTensor<ElementPrologueB> gmB;
259 gmB.SetGlobalBuffer((__gm__ ElementPrologueB*)params.ptrB);251 gmB.SetGlobalBuffer((__gm__ ElementPrologueB*)params.ptrB);
260 252 
261- auto tensorGmB = tla::MakeTensor(gmB, params.layoutB, Arch::PositionGM{});253+ auto tensorGmB = tla::MakeTensor(gmB, params.layoutPrologueB, Arch::PositionGM{});
262 254 
263 BlockScheduler matmulBlockScheduler(params.problemShape, MakeCoord(L1_TILE_M, L1_TILE_N));255 BlockScheduler matmulBlockScheduler(params.problemShape, MakeCoord(L1_TILE_M, L1_TILE_N));
264 // veccore256 // veccore
265 uint32_t coreLoops = matmulBlockScheduler.GetCoreLoops();257 uint32_t coreLoops = matmulBlockScheduler.GetCoreLoops();
266 258 
267- // make L1 TensorB
268- auto tensorL1B = tla::MakeTensor(l1BTensorList[l1BListId], L1B_LAYOUT, Arch::PositionL1{});
269- 
270 for (uint32_t loopIdx = AscendC::GetBlockIdx() / AscendC::GetSubBlockNum(); loopIdx < coreLoops;259 for (uint32_t loopIdx = AscendC::GetBlockIdx() / AscendC::GetSubBlockNum(); loopIdx < coreLoops;
271 loopIdx += AscendC::GetBlockNum()) {260 loopIdx += AscendC::GetBlockNum()) {
272 GemmCoord blockCoord = matmulBlockScheduler.GetBlockCoord(loopIdx);261 GemmCoord blockCoord = matmulBlockScheduler.GetBlockCoord(loopIdx);
@@ -274,20 +263,16 @@ public:
274 auto blockTensorB = GetTile(263 auto blockTensorB = GetTile(
275 tensorGmB, tla::MakeCoord(blockCoord.k() * L1_TILE_K, blockCoord.n() * L1_TILE_N),264 tensorGmB, tla::MakeCoord(blockCoord.k() * L1_TILE_K, blockCoord.n() * L1_TILE_N),
276 tla::MakeShape(actualBlockShape.k(), actualBlockShape.n()));265 tla::MakeShape(actualBlockShape.k(), actualBlockShape.n()));
277- blockPrologue(blockTensorB, tensorL1B, actualBlockShape, prologueParams);266+ blockPrologue(blockTensorB, actualBlockShape, prologueParams);
278 }267 }
279 }268 }
280 269 
281- // L1B TensorList270+ Arch::Resource<ArchTag> resource;
282- AscendC::LocalTensor<ElementB> l1BTensorList[L1B_STAGES];
283- int32_t l1BEventList[L1B_STAGES];
284- 
285- // The id of current stage
286- uint32_t l1BListId{0};
287 271 
288 constexpr static uint8_t AIC_SYNC_AIV_MODE = 4;272 constexpr static uint8_t AIC_SYNC_AIV_MODE = 4;
289- static constexpr uint16_t AIV_SYNC_AIC_FLAG = 0;273+ constexpr static uint16_t AIV_SYNC_AIC_FLAG = 6;
290- static constexpr uint16_t AIC_SYNC_AIV_FLAG = 16;274+ constexpr static uint16_t AIC_SYNC_AIV_FLAG = 8;
275+ constexpr static uint16_t FLAG_ID_MAX = 16;
291 276 
292 // L1BLayout277 // L1BLayout
293 static constexpr auto L1B_LAYOUT =278 static constexpr auto L1B_LAYOUT =
Ainclude/catlass/gemm/kernel/weight_quant_a8w4_grouped_mx_matmul.hpp+375-0
@@ -0,0 +1,375 @@
1+/**
2+ * This program is free software, you can redistribute it and/or modify.
3+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
4+ * This file is a part of the CANN Open Software.
5+ * Licensed under CANN Open Software License Agreement Version 2.0 (the "License").
6+ * Please refer to the License for details. You may not use this file except in compliance with the License.
7+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING
8+ * BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. See LICENSE in the root of
9+ * the software repository for the full text of the License.
10+ */
11+ 
12+#ifndef CATLASS_GEMM_KERNEL_WEIGHT_QUANT_A8W4_GROUPED_MX_MATMUL_HPP
13+#define CATLASS_GEMM_KERNEL_WEIGHT_QUANT_A8W4_GROUPED_MX_MATMUL_HPP
14+ 
15+#include "catlass/arch/resource.hpp"
16+#include "catlass/catlass.hpp"
17+#include "catlass/coord.hpp"
18+#include "catlass/gemm_coord.hpp"
19+#include "catlass/matrix_coord.hpp"
20+#include "tla/layout.hpp"
21+#include "tla/tensor.hpp"
22+ 
23+namespace Catlass::Gemm::Kernel {
24+ 
25+#if (defined(CATLASS_ARCH) && CATLASS_ARCH == 3510)
26+ 
27+// Template for MxMatmul kernel. Compute C = (MxScaleA x A) * (MxScaleB x B) + Bias
28+template <class BlockMmad_, class BlockPrologue_, class BlockEpilogue_, class BlockScheduler_, class ElementGroupList_>
29+class A8W4GroupedMxMatmul {
30+public:
31+ using BlockMmad = BlockMmad_;
32+ using BlockPrologue = BlockPrologue_;
33+ using ArchTag = typename BlockMmad::ArchTag;
34+ using L1TileShape = typename BlockMmad::L1TileShape;
35+ using ElementA = typename BlockMmad::ElementA;
36+ using LayoutA = typename BlockMmad::LayoutA;
37+ using ElementB = typename BlockMmad::ElementB;
38+ using LayoutB = typename BlockMmad::LayoutB;
39+ using ElementMxScaleA = typename BlockMmad::TileCopy::ElementMxScaleA;
40+ using LayoutMxScaleA = typename BlockMmad::TileCopy::LayoutMxScaleA;
41+ using ElementMxScaleB = typename BlockMmad::TileCopy::ElementMxScaleB;
42+ using LayoutMxScaleB = typename BlockMmad::TileCopy::LayoutMxScaleB;
43+ using ElementC = typename BlockMmad::ElementC;
44+ using LayoutC = typename BlockMmad::LayoutC;
45+ using ElementBias = typename BlockMmad::ElementBias;
46+ using ElementPrologueB = typename BlockMmad::TileCopy::ElementPrologueB;
47+ using LayoutPrologueB = typename BlockMmad::TileCopy::LayoutPrologueB;
48+ using ElementAccumulator = typename BlockMmad::ElementAccumulator;
49+ using LayoutTagL1B = typename BlockMmad::LayoutTagL1B;
50+ 
51+ using BlockScheduler = BlockScheduler_;
52+ using ElementGroupList = ElementGroupList_;
53+ 
54+ static constexpr uint32_t L1_TILE_M = tla::get<0>(L1TileShape{});
55+ static constexpr uint32_t L1_TILE_N = tla::get<1>(L1TileShape{});
56+ static constexpr uint32_t L1_TILE_K = tla::get<2>(L1TileShape{});
57+ 
58+ // L1B STAGES
59+ static constexpr uint32_t L1B_STAGES = 2;
60+ 
61+ /// Parameters structure
62+ struct Params {
63+ // Data members
64+ GemmCoord problemShape;
65+ uint32_t problemCount;
66+ GM_ADDR ptrGroupList;
67+ GM_ADDR ptrA;
68+ LayoutA layoutA;
69+ GM_ADDR ptrB;
70+ LayoutPrologueB layoutPrologueB;
71+ GM_ADDR ptrMxScaleA;
72+ LayoutMxScaleA layoutMxScaleA;
73+ GM_ADDR ptrMxScaleB;
74+ LayoutMxScaleB layoutMxScaleB;
75+ GM_ADDR ptrC;
76+ LayoutC layoutC;
77+ GM_ADDR ptrBias;
78+ 
79+ // Methods
80+ CATLASS_HOST_DEVICE
81+ Params()
82+ {}
83+ 
84+ CATLASS_HOST_DEVICE
85+ Params(
86+ GemmCoord const& problemShape_, uint32_t problemCount_, GM_ADDR ptrGroupList_, GM_ADDR ptrA_,
87+ LayoutA layoutA_, GM_ADDR ptrB_, LayoutPrologueB layoutPrologueB_, GM_ADDR ptrMxScaleA_,
88+ LayoutMxScaleA layoutMxScaleA_, GM_ADDR ptrMxScaleB_, LayoutMxScaleB layoutMxScaleB_, GM_ADDR ptrC_,
89+ LayoutC layoutC_, GM_ADDR ptrBias_ = nullptr)
90+ : problemShape(problemShape_),
91+ problemCount(problemCount_),
92+ ptrGroupList(ptrGroupList_),
93+ ptrA(ptrA_),
94+ layoutA(layoutA_),
95+ ptrB(ptrB_),
96+ layoutPrologueB(layoutPrologueB_),
97+ ptrMxScaleA(ptrMxScaleA_),
98+ layoutMxScaleA(layoutMxScaleA_),
99+ ptrMxScaleB(ptrMxScaleB_),
100+ layoutMxScaleB(layoutMxScaleB_),
101+ ptrC(ptrC_),
102+ layoutC(layoutC_),
103+ ptrBias(ptrBias_)
104+ {}
105+ };
106+ 
107+ struct Arguments {
108+ GemmCoord problemShape;
109+ uint32_t problemCount;
110+ uint8_t* ptrGroupList;
111+ uint8_t* ptrA;
112+ LayoutA layoutA;
113+ uint8_t* ptrB;
114+ LayoutPrologueB layoutPrologueB;
115+ uint8_t* ptrMxScaleA;
116+ LayoutMxScaleA layoutMxScaleA;
117+ uint8_t* ptrMxScaleB;
118+ LayoutMxScaleB layoutMxScaleB;
119+ uint8_t* ptrC;
120+ LayoutC layoutC;
121+ uint8_t* ptrBias{nullptr};
122+ };
123+ 
124+ static bool CanImplement(const Arguments& args)
125+ {
126+ return AscendC::Std::is_one_of_v<ElementA, float8_e4m3_t, float8_e5m2_t, float4_e2m1x2_t, float4_e1m2x2_t> &&
127+ AscendC::Std::is_one_of_v<ElementB, float8_e4m3_t, float8_e5m2_t, float4_e2m1x2_t, float4_e1m2x2_t> &&
128+ std::is_same_v<ElementMxScaleA, float8_e8m0_t> && std::is_same_v<ElementMxScaleB, float8_e8m0_t>;
129+ }
130+ 
131+ static size_t GetWorkspaceSize(const Arguments& args)
132+ {
133+ return 0;
134+ }
135+ 
136+ static Params ToUnderlyingArguments(const Arguments& args, uint8_t* workspace)
137+ {
138+ Params params{args.problemShape, args.problemCount, args.ptrGroupList, args.ptrA,
139+ args.layoutA, args.ptrB, args.layoutPrologueB, args.ptrMxScaleA,
140+ args.layoutMxScaleA, args.ptrMxScaleB, args.layoutMxScaleB, args.ptrC,
141+ args.layoutC, args.ptrBias};
142+ return params;
143+ }
144+ 
145+ // Methods
146+ CATLASS_DEVICE
147+ A8W4GroupedMxMatmul()
148+ {
149+ if constexpr (tla::detail::isRowMajor<LayoutC>::value) {
150+ AscendC::SetMMLayoutTransform(true);
151+ }
152+ if ASCEND_IS_AIC {
153+ for (uint32_t i = 0; i < L1B_STAGES; i++) {
154+ AscendC::CrossCoreSetFlag<AIC_SYNC_AIV_MODE, PIPE_MTE1>(AIC_SYNC_AIV_FLAG + i);
155+ AscendC::CrossCoreSetFlag<AIC_SYNC_AIV_MODE, PIPE_MTE1>(AIC_SYNC_AIV_FLAG + FLAG_ID_MAX + i);
156+ }
157+ }
158+ }
159+ 
160+ // Destructor
161+ CATLASS_DEVICE
162+ ~A8W4GroupedMxMatmul()
163+ {
164+ if ASCEND_IS_AIV {
165+ for (uint32_t i = 0; i < L1B_STAGES; i++) {
166+ AscendC::CrossCoreWaitFlag<AIC_SYNC_AIV_MODE, PIPE_MTE3>(AIC_SYNC_AIV_FLAG + i);
167+ }
atomgit-bot
atomgit-botatomgit-bot7月9日

🟡 Medium Priority

在 weight_quant_a8w4_grouped_mx_matmul.hpp 中：

Set 和 Wait 使用了不同的 Pipe（MTE1 vs MTE3）。根据 AscendC 同步机制，SetFlag/WaitFlag 需要在同一 Pipe 上配对才能正确同步。当前 MTE1 上 Set 的 flag 在 MTE3 上 Wait 可能等不到，导致 AIV 侧的 Wait 死等或错过同步信号。

同时，AIC 构造函数还在第 149 行设置了 AIC_SYNC_AIV_FLAG + FLAG_ID_MAX + i 的 flag，但 AIV 析构函数中从未 Wait 对应的 flag（仅 Wait 了 AIC_SYNC_AIV_FLAG + i），造成另一组 flag 设置后无人等待。

建议：确认 CrossCoreSetFlag 与 CrossCoreWaitFlag 的 Pipe 参数是否应一致。如果 MTE1/MTE3 差异是有意为之（不同方向上使用不同 Pipe），请确认 AscendC 文档支持此用法。同时确认 FLAG_ID_MAX 偏移的 flag 是否需要对应 Wait，或移除多余的 Set。

likedislike
不准确？
168+ }
169+ }
170+ 
171+ template <int32_t CORE_TYPE = g_coreType>
172+ CATLASS_DEVICE void operator()(Params const& params);
173+ 
174+ /// Executes one Matmul
175+ template <>
176+ CATLASS_DEVICE void operator()<AscendC::AIC>(Params const& params)
177+ {
178+ BlockMmad blockMmad(resource);
179+ 
180+ uint32_t coreIdx = AscendC::GetBlockIdx();
181+ uint32_t coreNum = AscendC::GetBlockNum();
182+ 
183+ // Represent the full gm
184+ AscendC::GlobalTensor<ElementA> gmA;
185+ gmA.SetGlobalBuffer((__gm__ ElementA*)params.ptrA);
186+ 
187+ AscendC::GlobalTensor<ElementC> gmC;
188+ gmC.SetGlobalBuffer((__gm__ ElementC*)params.ptrC);
189+ 
190+ AscendC::GlobalTensor<ElementGroupList> groupList;
191+ groupList.SetGlobalBuffer((__gm__ ElementGroupList*)params.ptrGroupList);
192+ 
193+ // Setting L2 Cache to Disable, data reads will bypass L2 Cache.
194+ if (CeilDiv(params.problemShape.n(), L1_TILE_N) == 1) {
195+ gmA.SetL2CacheHint(AscendC::CacheMode::CACHE_MODE_DISABLE);
196+ }
197+ 
198+ using GlobalTensorBiasType = std::conditional_t<std::is_void_v<ElementBias>, uint8_t, ElementBias>;
199+ AscendC::GlobalTensor<GlobalTensorBiasType> gmBias;
200+ 
201+ auto layoutBias = tla::MakeLayout(params.problemShape.n());
202+ 
203+ int64_t gmGroupOffsetMxScaleA = 0;
204+ int64_t gmGroupOffsetMxScaleB = 0;
205+ int64_t mxScaleAlignedK =
206+ static_cast<int64_t>(CeilDiv<MX_BASEK_FACTOR>(params.problemShape.k()) * MX_SCALE_COPY_GROUP_NUM);
207+ 
208+ int64_t totalM = 0;
209+ uint32_t startCoreIdx = 0;
210+ 
211+ // Represent the full tensors
212+ auto tensorA = tla::MakeTensor(gmA, params.layoutA, Arch::PositionGM{});
213+ auto tensorC = tla::MakeTensor(gmC, params.layoutC, Arch::PositionGM{});
214+ auto tensorBias = tla::MakeTensor(gmBias, layoutBias, Arch::PositionGM{});
215+ 
216+ for (uint32_t groupIdx = 0; groupIdx < params.problemCount; ++groupIdx) {
217+ uint32_t currentM = groupList.GetValue(groupIdx);
218+ GemmCoord inGroupProblemShape{currentM, params.problemShape.n(), params.problemShape.k()};
219+ 
悠悠的云朵
悠悠的云朵悠悠的云朵7月12日

添加L2 Cache Hit优化： if (CeilDiv(currentM, tileM) == 1) { gmB.SetL2CacheHint(AscendC::CacheMode::CACHE_MODE_DISABLE); } else { gmB.SetL2CacheHint(AscendC::CacheMode::CACHE_MODE_NORMAL); }

likedislike
220+ auto tileM = GetTileM(currentM);
221+ 
222+ auto layoutMxScaleA = tla::MakeMxScaleLayout<ElementMxScaleA, layout::RowMajor, false>(
223+ inGroupProblemShape.m(), CeilDiv<MX_SCALE_GROUP_NUM>(inGroupProblemShape.k()));
224+ 
225+ BlockScheduler matmulBlockScheduler(inGroupProblemShape, MakeCoord(tileM, L1_TILE_N));
226+ uint32_t coreLoops = matmulBlockScheduler.GetCoreLoops();
227+ 
228+ AscendC::GlobalTensor<ElementMxScaleB> gmMxScaleB;
229+ gmMxScaleB.SetGlobalBuffer((__gm__ ElementMxScaleB*)params.ptrMxScaleB + gmGroupOffsetMxScaleB);
230+ 
231+ AscendC::GlobalTensor<ElementMxScaleA> gmMxScaleA;
232+ gmMxScaleA.SetGlobalBuffer((__gm__ ElementMxScaleA*)params.ptrMxScaleA + gmGroupOffsetMxScaleA);
233+ 
234+ uint32_t startLoopIdx;
235+ if (coreIdx < startCoreIdx) {
236+ startLoopIdx = coreIdx + coreNum - startCoreIdx;
237+ } else {
238+ startLoopIdx = coreIdx - startCoreIdx;
239+ }
240+ 
241+ auto tensorMxScaleA = tla::MakeTensor(gmMxScaleA, layoutMxScaleA, Arch::PositionGM{});
242+ auto tensorMxScaleB = tla::MakeTensor(gmMxScaleB, params.layoutMxScaleB, Arch::PositionGM{});
243+ 
244+ for (uint32_t loopIdx = startLoopIdx; loopIdx < coreLoops; loopIdx += coreNum) {
245+ GemmCoord blockCoord = matmulBlockScheduler.GetBlockCoord(loopIdx);
246+ GemmCoord actualBlockShape = matmulBlockScheduler.GetActualBlockShape(blockCoord);
247+ 
248+ auto tensorBlockA = GetTile(
249+ tensorA, tla::MakeCoord(totalM + blockCoord.m() * tileM, blockCoord.k() * L1_TILE_K),
250+ tla::MakeShape(actualBlockShape.m(), actualBlockShape.k()));
251+ 
252+ auto tensorBlockC = GetTile(
253+ tensorC, tla::MakeCoord(totalM + blockCoord.m() * tileM, blockCoord.n() * L1_TILE_N),
254+ tla::MakeShape(actualBlockShape.m(), actualBlockShape.n()));
255+ 
256+ auto tensorBlockMxScaleA = GetTile(
257+ tensorMxScaleA,
258+ tla::MakeCoord(blockCoord.m() * tileM, blockCoord.k() * L1_TILE_K / MX_SCALE_GROUP_NUM),
259+ tla::MakeShape(actualBlockShape.m(), CeilDiv<MX_SCALE_GROUP_NUM>(actualBlockShape.k())));
260+ 
261+ auto tensorBlockMxScaleB = GetTile(
262+ tensorMxScaleB,
263+ tla::MakeCoord(blockCoord.k() * L1_TILE_K / MX_SCALE_GROUP_NUM, blockCoord.n() * L1_TILE_N),
264+ tla::MakeShape(CeilDiv<MX_SCALE_GROUP_NUM>(actualBlockShape.k()), actualBlockShape.n()));
265+ 
266+ if constexpr (std::is_void_v<ElementBias>) {
267+ blockMmad(tensorBlockA, tensorBlockC, actualBlockShape, tensorBlockMxScaleA, tensorBlockMxScaleB);
268+ } else {
269+ auto tensorBlockBias = GetTile(
270+ tensorBias, tla::MakeCoord(blockCoord.n() * L1_TILE_N), tla::MakeShape(actualBlockShape.n()));
271+ blockMmad(
272+ tensorBlockA, tensorBlockC, actualBlockShape, tensorBlockMxScaleA, tensorBlockMxScaleB,
273+ tensorBlockBias);
274+ }
275+ }
276+ 
277+ totalM += inGroupProblemShape.m();
278+ 
279+ gmGroupOffsetMxScaleA += inGroupProblemShape.m() * mxScaleAlignedK;
280+ gmGroupOffsetMxScaleB += mxScaleAlignedK * inGroupProblemShape.n();
281+ 
282+ startCoreIdx = (startCoreIdx + coreLoops) % coreNum;
283+ }
284+ 
285+ if constexpr (BlockMmad::DispatchPolicy::ASYNC) {
286+ blockMmad.template SynchronizeBlock<decltype(tensorC)>();
287+ }
288+ 
289+ AscendC::PipeBarrier<PIPE_ALL>();
290+ }
291+ 
292+ template <>
293+ CATLASS_DEVICE void operator()<AscendC::AIV>(Params const& params)
294+ {
295+ using PrologueParams = typename BlockPrologue::Params;
296+ PrologueParams prologueParams{L1TileShape{}, params.layoutPrologueB, false, resource};
297+ 
298+ BlockPrologue blockPrologue(prologueParams);
299+ 
300+ AscendC::GlobalTensor<ElementGroupList> groupList;
301+ groupList.SetGlobalBuffer((__gm__ ElementGroupList*)params.ptrGroupList);
302+ 
303+ uint32_t coreIdx = AscendC::GetBlockIdx() / AscendC::GetSubBlockNum();
304+ uint32_t coreNum = AscendC::GetBlockNum();
305+ int64_t gmGroupOffsetB = 0;
306+ uint32_t startCoreIdx = 0;
307+ 
308+ for (uint32_t groupIdx = 0; groupIdx < params.problemCount; ++groupIdx) {
309+ uint32_t currentM = groupList.GetValue(groupIdx);
310+ 
311+ // Represent the full gmB
312+ AscendC::GlobalTensor<ElementPrologueB> gmB;
313+ gmB.SetGlobalBuffer((__gm__ ElementPrologueB*)params.ptrB + gmGroupOffsetB);
314+ auto tensorGmB = tla::MakeTensor(gmB, params.layoutPrologueB, Arch::PositionGM{});
315+ 
316+ GemmCoord inGroupProblemShape{currentM, params.problemShape.n(), params.problemShape.k()};
317+ auto tileM = GetTileM(currentM);
318+ BlockScheduler matmulBlockScheduler(inGroupProblemShape, MakeCoord(tileM, L1_TILE_N));
319+ 
320+ if (CeilDiv(currentM, tileM) == 1) {
321+ gmB.SetL2CacheHint(AscendC::CacheMode::CACHE_MODE_DISABLE);
322+ } else {
323+ gmB.SetL2CacheHint(AscendC::CacheMode::CACHE_MODE_NORMAL);
324+ }
325+ 
326+ // veccore
327+ uint32_t coreLoops = matmulBlockScheduler.GetCoreLoops();
328+ 
329+ uint32_t startLoopIdx;
330+ if (coreIdx < startCoreIdx) {
331+ startLoopIdx = coreIdx + coreNum - startCoreIdx;
332+ } else {
333+ startLoopIdx = coreIdx - startCoreIdx;
334+ }
335+ 
336+ for (uint32_t loopIdx = startLoopIdx; loopIdx < coreLoops; loopIdx += AscendC::GetBlockNum()) {
337+ GemmCoord blockCoord = matmulBlockScheduler.GetBlockCoord(loopIdx);
338+ GemmCoord actualBlockShape = matmulBlockScheduler.GetActualBlockShape(blockCoord);
339+ auto blockTensorB = GetTile(
340+ tensorGmB, tla::MakeCoord(blockCoord.k() * L1_TILE_K, blockCoord.n() * L1_TILE_N),
341+ tla::MakeShape(actualBlockShape.k(), actualBlockShape.n()));
342+ blockPrologue(blockTensorB, actualBlockShape, prologueParams);
343+ }
344+ 
345+ gmGroupOffsetB += static_cast<int64_t>(CeilDiv<2>(inGroupProblemShape.k())) * inGroupProblemShape.n();
346+ startCoreIdx = (startCoreIdx + coreLoops) % coreNum;
347+ }
348+ AscendC::PipeBarrier<PIPE_ALL>();
349+ }
350+ 
351+ CATLASS_DEVICE
352+ uint32_t GetTileM(uint32_t mSize)
353+ {
354+ if (mSize <= L1_TILE_M) {
355+ return L1_TILE_M;
356+ } else {
357+ uint32_t tileNum = CeilDiv(mSize, L1_TILE_M);
358+ return CeilDiv((mSize / tileNum), TILE_M_ALIGNMENT) * TILE_M_ALIGNMENT;
359+ }
360+ }
361+ 
362+ Arch::Resource<ArchTag> resource;
363+ 
364+ constexpr static uint8_t AIC_SYNC_AIV_MODE = 4;
365+ constexpr static uint16_t AIV_SYNC_AIC_FLAG = 6;
366+ constexpr static uint16_t AIC_SYNC_AIV_FLAG = 8;
367+ constexpr static uint16_t FLAG_ID_MAX = 16;
368+ constexpr static uint16_t TILE_M_ALIGNMENT = 16;
369+};
370+ 
371+#endif // (defined(CATLASS_ARCH) && CATLASS_ARCH == 3510)
372+ 
373+} // namespace Catlass::Gemm::Kernel
374+ 
375+#endif // CATLASS_GEMM_KERNEL_WEIGHT_QUANT_A8W4_GROUPED_MX_MATMUL_HPP
Minclude/catlass/gemm/tile/tile_copy.hpp+4-3
@@ -527,8 +527,8 @@ template <
527 class LayoutMxScaleB_, class ElementC_, class LayoutTagC, class ElementBias = void, bool ReluEnable_ = false,527 class LayoutMxScaleB_, class ElementC_, class LayoutTagC, class ElementBias = void, bool ReluEnable_ = false,
528 ScaleGranularity DEQUANT_GRANULARITY = ScaleGranularity::NO_QUANT, class L0CCopyMode = CopyToGM>528 ScaleGranularity DEQUANT_GRANULARITY = ScaleGranularity::NO_QUANT, class L0CCopyMode = CopyToGM>
529struct PackedMxA8W4TileCopyTla : public PackedTileCopyTla<529struct PackedMxA8W4TileCopyTla : public PackedTileCopyTla<
530- ArchTag, ElementA_, LayoutTagA, ElementB_, LayoutTagPrologueB, ElementC_,530+ ArchTag, ElementA_, LayoutTagA, ElementB_, LayoutTagB, ElementC_, LayoutTagC,
531- LayoutTagC, ElementBias, ReluEnable_, DEQUANT_GRANULARITY, L0CCopyMode> {531+ ElementBias, ReluEnable_, DEQUANT_GRANULARITY, L0CCopyMode> {
532 using ElementMxScaleA = ElementMxScaleA_;532 using ElementMxScaleA = ElementMxScaleA_;
533 using ElementMxScaleB = ElementMxScaleB_;533 using ElementMxScaleB = ElementMxScaleB_;
534 534 
@@ -540,7 +540,8 @@ struct PackedMxA8W4TileCopyTla : public PackedTileCopyTla<
540 using LayoutL1MxScaleA = detail::TagToLayout_t<ElementMxScaleA, LayoutTagL1MxScaleA>;540 using LayoutL1MxScaleA = detail::TagToLayout_t<ElementMxScaleA, LayoutTagL1MxScaleA>;
541 using LayoutL1MxScaleB = detail::TagToLayout_t<ElementMxScaleB, LayoutTagL1MxScaleB>;541 using LayoutL1MxScaleB = detail::TagToLayout_t<ElementMxScaleB, LayoutTagL1MxScaleB>;
542 542 
543- using LayoutB = detail::TagToLayout_t<ElementPrologueB_, LayoutTagPrologueB>;543+ using ElementPrologueB = ElementPrologueB_;
544+ using LayoutPrologueB = detail::TagToLayout_t<ElementPrologueB, LayoutTagPrologueB>;
544 using LayoutTagL1B = typename helper::L1BTypeSelector<Gemm::GemmType<ElementB_, LayoutTagB>>::L1BType::Layout;545 using LayoutTagL1B = typename helper::L1BTypeSelector<Gemm::GemmType<ElementB_, LayoutTagB>>::L1BType::Layout;
545 using LayoutTagL0B = layout::nZ;546 using LayoutTagL0B = layout::nZ;
546 547 
Minclude/catlass/layout/matrix.hpp+168-0
@@ -1734,6 +1734,174 @@ public:
1734 return stride_[idx];1734 return stride_[idx];
1735 }1735 }
1736 1736 
1737+private:
1738+ /// Origin Shape data member
1739+ OrgShape orgShape_;
1740+ 
1741+ /// Shape data member
1742+ Shape shape_;
1743+ 
1744+ /// Stride data member
1745+ Stride stride_;
1746+};
1747+ 
1748+/// Mapping function for nZ matrices which is col-major inside fractal and row-major between fractal
1749+struct Weight4BitnZ {
1750+public:
1751+ /// Logical rank of tensor
1752+ static constexpr int RANK = 4;
1753+ 
1754+ /// Index type used for coordinates
1755+ using Index = uint32_t;
1756+ 
1757+ /// Long index type used for offsets
1758+ using LongIndex = int64_t;
1759+ 
1760+ /// Logical rank of orgshape
1761+ static constexpr int ORG_SHAPE_RANK = 2;
1762+ 
1763+ /// Logical coordinate
1764+ using OrgShape = Coord<ORG_SHAPE_RANK, Index>;
1765+ 
1766+ /// Logical coordinate
1767+ using Shape = Coord<RANK, Index>;
1768+ 
1769+ /// Stride vector
1770+ using Stride = Coord<RANK, LongIndex>;
1771+ 
1772+public:
1773+ // Methods
1774+ 
1775+ /// Constructor
1776+ CATLASS_HOST_DEVICE constexpr Weight4BitnZ(
1777+ Index orgRows = 0, /// Number of rows of origin matrices
1778+ Index orgCols = 0, /// Number of cols of origin matrices
1779+ Index rowsInFractal = 0, /// Number of rows inside the fractal
1780+ Index rowsByFractal = 0, /// number of rows by the fractal
1781+ Index colsInFractal = 0, /// number of cols inside the fractal
1782+ Index colsByFractal = 0, /// number of cols by the fractal
1783+ LongIndex strideRowsInFractal = 0, /// number of elements between adjacent rows inside the fractal
1784+ LongIndex strideRowsByFractal = 0, /// number of elements between adjacent fractal rows
1785+ LongIndex strideColsInFractal = 0, /// number of elements between adjacent cols inside the fractal
1786+ LongIndex strideColsByFractal = 0) /// number of elements between adjacent fractal cols
1787+ : orgShape_(MakeCoord(orgRows, orgCols)),
1788+ shape_(MakeCoord(rowsInFractal, rowsByFractal, colsInFractal, colsByFractal)),
1789+ stride_(MakeCoord(strideRowsInFractal, strideRowsByFractal, strideColsInFractal, strideColsByFractal))
1790+ {}
1791+ 
1792+ /// Ctor
1793+ CATLASS_HOST_DEVICE constexpr Weight4BitnZ(OrgShape orgShape, Shape shape, Stride stride)
1794+ : orgShape_(orgShape), shape_(shape), stride_(stride)
1795+ {}
1796+ 
1797+ /// Make the layout of a coordinate (row, column)
1798+ template <class Element>
1799+ CATLASS_HOST_DEVICE constexpr static Weight4BitnZ MakeLayout(Index orgRows, Index orgCols)
1800+ {
1801+ constexpr uint32_t ELE_NUM_PER_C0 = 16;
1802+ constexpr uint32_t ELE_NUM_PER_FRACTAL = 512;
1803+ Index rowsRound = RoundUp<ELE_NUM_PER_C0>(orgRows);
1804+ Index colsRound = RoundUp<C0_NUM_PER_FRACTAL>(orgCols);
1805+ return Weight4BitnZ(
1806+ orgRows, orgCols, ELE_NUM_PER_C0, rowsRound / ELE_NUM_PER_C0, C0_NUM_PER_FRACTAL,
1807+ colsRound / C0_NUM_PER_FRACTAL, 1, colsRound * ELE_NUM_PER_C0, ELE_NUM_PER_C0, ELE_NUM_PER_FRACTAL);
1808+ }
1809+ 
1810+ /// Returns the offset of a coordinate in linear memory.
1811+ /// Assumes coordinate has convention (row, column)
1812+ CATLASS_HOST_DEVICE
1813+ LongIndex GetOffset(MatrixCoord const& coord) const
1814+ {
1815+ return LongIndex(coord.row()) / shape_[0] * stride_[1] + LongIndex(coord.column()) / shape_[2] * stride_[3] +
1816+ (LongIndex(coord.row()) % shape_[0]) * stride_[0] + (LongIndex(coord.column()) % shape_[2]) * stride_[2];
1817+ }
1818+ 
1819+ /// Returns the layout of a tile.
1820+ CATLASS_HOST_DEVICE
1821+ Weight4BitnZ GetTileLayout(MatrixCoord const& tileOriShape) const
1822+ {
1823+ auto tileShape = MakeCoord(
1824+ shape(0), CeilDiv(tileOriShape.row(), shape(0)), shape(2), CeilDiv(tileOriShape.column(), shape(2)));
1825+ return Weight4BitnZ(tileOriShape, tileShape, stride());
1826+ }
1827+ 
1828+ /// Returns the origin shape of the layout
1829+ CATLASS_HOST_DEVICE
1830+ typename OrgShape::Index orgShape(int idx) const
1831+ {
1832+ return orgShape_[idx];
1833+ }
1834+ 
1835+ /// Returns the origin shape of the layout
1836+ CATLASS_HOST_DEVICE
1837+ typename OrgShape::Index& orgShape(int idx)
1838+ {
1839+ return orgShape_[idx];
1840+ }
1841+ 
1842+ /// Returns the shape of the layout
1843+ CATLASS_HOST_DEVICE
1844+ Shape shape() const
1845+ {
1846+ return shape_;
1847+ }
1848+ 
1849+ /// Returns the shape of the layout
1850+ CATLASS_HOST_DEVICE
1851+ Shape& shape()
1852+ {
1853+ return shape_;
1854+ }
1855+ 
1856+ /// Returns the shape of the layout
1857+ CATLASS_HOST_DEVICE
1858+ typename Shape::Index shape(int idx) const
1859+ {
1860+ return shape_[idx];
1861+ }
1862+ 
1863+ /// Returns the shape of the layout
1864+ CATLASS_HOST_DEVICE
1865+ typename Shape::Index& shape(int idx)
1866+ {
1867+ return shape_[idx];
1868+ }
1869+ 
1870+ /// Returns the stride of the layout
1871+ CATLASS_HOST_DEVICE
1872+ Stride stride() const
1873+ {
1874+ return stride_;
1875+ }
1876+ 
1877+ /// Returns the stride of the layout
1878+ CATLASS_HOST_DEVICE
1879+ Stride& stride()
1880+ {
1881+ return stride_;
1882+ }
1883+ 
1884+ /// Returns the stride of the layout
1885+ CATLASS_HOST_DEVICE
1886+ typename Stride::Index stride(int idx) const
1887+ {
1888+ return stride_[idx];
1889+ }
1890+ 
1891+ /// Returns the stride of the layout
1892+ CATLASS_HOST_DEVICE
1893+ typename Stride::Index& stride(int idx)
1894+ {
1895+ return stride_[idx];
1896+ }
1897+ 
1898+ /// Returns the length of the layout
1899+ CATLASS_HOST_DEVICE
1900+ LongIndex Capacity() const
1901+ {
1902+ return static_cast<LongIndex>(stride_[1]) * shape_[1];
1903+ }
1904+ 
1737private:1905private:
1738 /// Origin Shape data member1906 /// Origin Shape data member
1739 OrgShape orgShape_;1907 OrgShape orgShape_;
Minclude/tla/layout.hpp+13-1
@@ -522,7 +522,8 @@ CATLASS_HOST_DEVICE constexpr auto MakeLayout(T const& rows, U const& cols)
522 std::is_same_v<LayoutTag, Catlass::layout::ColumnMajor> ||522 std::is_same_v<LayoutTag, Catlass::layout::ColumnMajor> ||
523 std::is_same_v<LayoutTag, Catlass::layout::VectorLayout> ||523 std::is_same_v<LayoutTag, Catlass::layout::VectorLayout> ||
524 std::is_same_v<LayoutTag, Catlass::layout::zN> || std::is_same_v<LayoutTag, Catlass::layout::nZ> ||524 std::is_same_v<LayoutTag, Catlass::layout::zN> || std::is_same_v<LayoutTag, Catlass::layout::nZ> ||
525- std::is_same_v<LayoutTag, Catlass::layout::zZ> || std::is_same_v<LayoutTag, Catlass::layout::L0C>,525+ std::is_same_v<LayoutTag, Catlass::layout::zZ> || std::is_same_v<LayoutTag, Catlass::layout::L0C> ||
526+ std::is_same_v<LayoutTag, Catlass::layout::Weight4BitnZ>,
526 "Unsupported LayoutTag for MakeLayoutFromTag, only support Catlass::layout::RowMajor or"527 "Unsupported LayoutTag for MakeLayoutFromTag, only support Catlass::layout::RowMajor or"
527 "Catlass::layout::ColumnMajor or Catlass::layout::zN or Catlass::layout::nZ or Catlass::layout::zZ or "528 "Catlass::layout::ColumnMajor or Catlass::layout::zN or Catlass::layout::nZ or Catlass::layout::zZ or "
528 "Catlass::layout::L0C");529 "Catlass::layout::L0C");
@@ -581,6 +582,17 @@ CATLASS_HOST_DEVICE constexpr auto MakeLayout(T const& rows, U const& cols)
581 Int<1>{},582 Int<1>{},
582 RoundUp((int64_t)rows, Int<Catlass::C0_NUM_PER_FRACTAL>{}) * Catlass::C0_NUM_PER_FRACTAL)),583 RoundUp((int64_t)rows, Int<Catlass::C0_NUM_PER_FRACTAL>{}) * Catlass::C0_NUM_PER_FRACTAL)),
583 MakeShape(rows, cols));584 MakeShape(rows, cols));
585+ } else if constexpr (std::is_same_v<LayoutTag, Catlass::layout::Weight4BitnZ>) {
586+ constexpr uint32_t ELE_NUM_PER_C0 = 32;
587+ constexpr uint32_t ELE_NUM_PER_FRACTAL = 512;
588+ return MakeLayout(
589+ MakeShape(
590+ MakeShape(Int<ELE_NUM_PER_C0>{}, CeilDiv(rows, Int<ELE_NUM_PER_C0>{})),
591+ MakeShape(Int<Catlass::C0_NUM_PER_FRACTAL>{}, CeilDiv(cols, Int<Catlass::C0_NUM_PER_FRACTAL>{}))),
592+ MakeStride(
593+ MakeStride(Int<1>{}, RoundUp((int64_t)cols, Int<Catlass::C0_NUM_PER_FRACTAL>{}) * ELE_NUM_PER_C0),
594+ MakeStride(Int<ELE_NUM_PER_C0>{}, Int<ELE_NUM_PER_FRACTAL>{})),
595+ MakeShape(rows, cols));
584 } else {596 } else {
585 return MakeLayout(597 return MakeLayout(
586 MakeShape(598 MakeShape(
Mtests/optest/README.md+1-0
@@ -82,6 +82,7 @@ Python API (`torch_catlass.ops.*`)
82- [x] 60_ascend950_grouped_matmul_slice_m (Ascend950)82- [x] 60_ascend950_grouped_matmul_slice_m (Ascend950)
83- [x] 67_ascend950_batched_matmul (Ascend950)83- [x] 67_ascend950_batched_matmul (Ascend950)
84- [x] 73_ascend950_matmul_full_loadA (Ascend950)84- [x] 73_ascend950_matmul_full_loadA (Ascend950)
85+- [x] 74_ascend950_weight_quant_a8w4_grouped_mx_matmul (Ascend950)
85 86 
86### 暂未接入87### 暂未接入
87 88 
Mtests/optest/include/catlass_kernel_jit.h+10-0
@@ -492,6 +492,16 @@ void Ascend950Fp8MxGroupedMatmulFinalizeRoutingNoDeter(
492void Ascend950MatmulFullLoadA(492void Ascend950MatmulFullLoadA(
493 const uint32_t blockNum, aclrtStream stream, const TParams& tParams, const MatmulParams& params);493 const uint32_t blockNum, aclrtStream stream, const TParams& tParams, const MatmulParams& params);
494 494 
495+/**
496+ * @brief Reserved JIT interface for example 74_ascend950_weight_quant_a8w4_grouped_mx_matmul.
497+ *
498+ * Grouped MX A8W4 matmul: C = (MxScaleA * A_fp8) @ (MxScaleB * B_fp4) per group.
499+ * B is the packed FP4 prologue (int8 bytes, Weight4BitnZ layout). Output is FP32.
500+ * ``params.batch`` carries the group (expert) count; ``inputAddr[2]`` is the group list.
501+ */
502+void Ascend950A8W4GroupedMxMatmul(
503+ const uint32_t blockNum, aclrtStream stream, const TParams& tParams, const GroupedMatmulParams& params);
504+ 
495/**505/**
496 * @brief Reserved JIT interface for example 102_dynamic_optimized_matmul.506 * @brief Reserved JIT interface for example 102_dynamic_optimized_matmul.
497 */507 */
Mtests/optest/kernels/59_ascend950_a8w4_mx_matmul/a8w4_mx_matmul_impl.cpp+1-2
@@ -110,8 +110,7 @@ extern "C" void run(uint32_t blockNum, aclrtStream stream, const CatlassKernel::
110 Catlass::Gemm::Tile::ScaleGranularity::PER_TENSOR>;110 Catlass::Gemm::Tile::ScaleGranularity::PER_TENSOR>;
111 111 
112 using BlockMmad = Catlass::Gemm::Block::BlockMmadA8W4Mx<112 using BlockMmad = Catlass::Gemm::Block::BlockMmadA8W4Mx<
113- DispatchPolicyMmad, L1TileShape, L0TileShape, ElementA, ElementB, ElementC, ElementPrologueB, ElementBias,113+ DispatchPolicyMmad, L1TileShape, L0TileShape, ElementA, ElementB, ElementC, ElementBias, TileCopy>;
114- TileCopy>;
115 114 
116 using BlockPrologue = Catlass::Gemm::Block::BlockPrologue<115 using BlockPrologue = Catlass::Gemm::Block::BlockPrologue<
117 DispatchPolicyPrologue, PrologueSrcType, PrologueDstType, L1TileShape, TileCopy>;116 DispatchPolicyPrologue, PrologueSrcType, PrologueDstType, L1TileShape, TileCopy>;
Atests/optest/kernels/74_ascend950_weight_quant_a8w4_grouped_mx_matmul/CMakeLists.txt+15-0
@@ -0,0 +1,15 @@
1+# -----------------------------------------------------------------------------------------------------------
2+# Copyright (c) 2026 Huawei Technologies Co., Ltd.
3+# This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+# CANN Open Software License Agreement Version 2.0 (the "License").
5+# Please refer to the License for details. You may not use this file except in compliance with the License.
6+# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+# See LICENSE in the root of the software repository for the full text of the License.
9+# -----------------------------------------------------------------------------------------------------------
10+ 
11+add_kernel(NAME ascend950_a8w4_grouped_mx_matmul
12+ NPU_ARCH_LIST 3510
13+ KERNEL_TYPE jit
14+ ${CMAKE_CURRENT_SOURCE_DIR}/weight_quant_a8w4_grouped_mx_matmul.cpp
15+ TEMPLATE ${CMAKE_CURRENT_SOURCE_DIR}/weight_quant_a8w4_grouped_mx_matmul_impl.cpp)
Atests/optest/kernels/74_ascend950_weight_quant_a8w4_grouped_mx_matmul/weight_quant_a8w4_grouped_mx_matmul.cpp+35-0
@@ -0,0 +1,35 @@
1+/**
2+ * This program is free software, you can redistribute it and/or modify.
3+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
4+ * This file is a part of the CANN Open Software.
5+ * Licensed under CANN Open Software License Agreement Version 2.0 (the "License").
6+ * Please refer to the License for details. You may not use this file except in compliance with the License.
7+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
8+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. See LICENSE in the root of
9+ * the software repository for the full text of the License.
10+ */
11+ 
12+#include "catlass_kernel.h"
13+#include "jit_compiler.h"
14+#include "jit_macro_generator.h"
15+ 
16+namespace CatlassKernel {
17+ 
18+/**
19+ * @brief example 74_ascend950_weight_quant_a8w4_grouped_mx_matmul:
20+ * Resolve and launch the JIT grouped A8W4 MX matmul implementation.
21+ */
22+extern "C" void Ascend950A8W4GroupedMxMatmul(
23+ const uint32_t blockNum, aclrtStream stream, const TParams& tParams, const GroupedMatmulParams& params)
24+{
25+ auto* entry = JitCompiler::instance().getKernel(
26+ "weight_quant_a8w4_grouped_mx_matmul_impl.cpp",
27+ JitMacroGenerator<TParams>::generate("ascend950_a8w4_grouped_mx_matmul", tParams),
28+ JitKernelType::MIX);
29+ if (entry) {
30+ entry(blockNum, stream, &params);
31+ }
32+ aclrtSynchronizeStream(stream);
33+}
34+ 
35+} // namespace CatlassKernel
Atests/optest/kernels/74_ascend950_weight_quant_a8w4_grouped_mx_matmul/weight_quant_a8w4_grouped_mx_matmul_impl.cpp+140-0
@@ -0,0 +1,140 @@
1+/**
2+ * This program is free software, you can redistribute it and/or modify.
3+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
4+ * This file is a part of the CANN Open Software.
5+ * Licensed under CANN Open Software License Agreement Version 2.0 (the "License").
6+ * Please refer to the License for details. You may not use this file except in compliance with the License.
7+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
8+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. See LICENSE in the root of
9+ * the software repository for the full text of the License.
10+ */
11+ 
12+#ifndef K_MAX_SHAPE_DIM
13+#define K_MAX_SHAPE_DIM 0
14+#endif
15+ 
16+#include <algorithm>
17+#include <cstddef>
18+using std::size_t;
19+ 
20+#include <kernel_operator.h>
21+ 
22+#include "catlass/gemm/kernel/weight_quant_a8w4_grouped_mx_matmul.hpp"
23+#include "catlass/arch/arch.hpp"
24+#include "catlass/catlass.hpp"
25+#include "catlass/gemm/block/block_mmad.hpp"
26+#include "catlass/gemm/block/block_swizzle.hpp"
27+#include "catlass/gemm/dispatch_policy.hpp"
28+#include "catlass/gemm/gemm_type.hpp"
29+#include "catlass/layout/layout.hpp"
30+#include "catlass/status.hpp"
31+#include "tla/layout.hpp"
32+ 
33+#include "catlass_kernel.h"
34+#include "common/kernel_runner.h"
35+ 
36+#ifndef CATLASS_JIT_ELEMENT_A
37+#define CATLASS_JIT_ELEMENT_A float8_e4m3_t
38+#endif
39+#ifndef CATLASS_JIT_ELEMENT_B
40+#define CATLASS_JIT_ELEMENT_B float4_e2m1x2_t
41+#endif
42+#ifndef CATLASS_JIT_ELEMENT_MX_SCALE
43+#define CATLASS_JIT_ELEMENT_MX_SCALE float8_e8m0_t
44+#endif
45+#ifndef CATLASS_JIT_LAYOUT_A
46+#define CATLASS_JIT_LAYOUT_A RowMajor
47+#endif
48+#ifndef CATLASS_JIT_LAYOUT_B
49+#define CATLASS_JIT_LAYOUT_B Weight4BitnZ
50+#endif
51+ 
52+using ElementA = CATLASS_JIT_ELEMENT_A;
53+using ElementPrologueB = CATLASS_JIT_ELEMENT_B;
54+using ElementB = float8_e4m3_t;
55+using ElementMxScale = CATLASS_JIT_ELEMENT_MX_SCALE;
56+// The accumulator type is fixed because the JIT output dtype describes the
57+// tensor interface, not the internal workspace representation.
58+using ElementC = float;
59+using ElementBias = void;
60+using ElementGroupList = int64_t;
61+ 
62+using LayoutTagA = Catlass::layout::CATLASS_JIT_LAYOUT_A;
63+using LayoutTagPrologueB = Catlass::layout::CATLASS_JIT_LAYOUT_B;
64+using LayoutTagMxScaleB = Catlass::layout::ColumnMajor;
65+using LayoutTagC = Catlass::layout::RowMajor;
66+using LayoutTagB = Catlass::layout::nZ;
67+ 
68+using ArchTag = Catlass::Arch::Ascend950;
69+constexpr bool enableUnitFlag = true;
70+static constexpr uint32_t L1_SCALE_FACTOR_K = 16;
71+static constexpr uint32_t L1A_STAGES = 2;
72+static constexpr uint32_t L1B_STAGES = 2;
73+static constexpr uint32_t L0A_STAGES = 2;
74+static constexpr uint32_t L0B_STAGES = 2;
75+static constexpr uint32_t L0C_STAGES = 1;
76+ 
77+using L1TileShape = tla::Shape<tla::Int<256>, tla::Int<256>, tla::Int<256>>;
78+using L0TileShape = tla::Shape<tla::Int<256>, tla::Int<256>, tla::Int<128>>;
79+ 
80+using PrologueSrcType = Catlass::Gemm::GemmType<ElementPrologueB, LayoutTagPrologueB>;
81+using PrologueDstType = Catlass::Gemm::GemmType<ElementB, LayoutTagB>;
82+ 
83+using DispatchPolicyMmad = Catlass::Gemm::MmadA8W4Mx<
84+ ArchTag, enableUnitFlag, false, L1_SCALE_FACTOR_K,
85+ L0C_STAGES, L1A_STAGES, L1B_STAGES, L0A_STAGES, L0B_STAGES>;
86+using DispatchPolicyPrologue = Catlass::Gemm::MxA8W4Prologue<ArchTag, L1B_STAGES>;
87+ 
88+using BlockScheduler = typename Catlass::Gemm::Block::GemmIdentityBlockSwizzle<3, 0>;
89+ 
90+using BlockEpilogue = void;
91+ 
92+extern "C" void run(uint32_t blockNum, aclrtStream stream, const CatlassKernel::MatmulParams* params)
93+{
94+ uint32_t m = params->m;
95+ uint32_t n = params->n;
96+ uint32_t k = params->k;
97+ uint32_t problemCount = params->batch;
98+ uint32_t mxScaleK = CeilDiv<Catlass::MX_SCALE_GROUP_NUM>(k);
99+ 
100+ Catlass::GemmCoord problemShape{m, n, k};
101+ 
102+ uint8_t* deviceA = params->inputAddr[0];
103+ uint8_t* devicePrologueB = params->inputAddr[1];
104+ uint8_t* deviceGroupList = params->inputAddr[2];
105+ uint8_t* deviceMxScaleA = params->inputAddr[3];
106+ uint8_t* deviceMxScaleB = params->inputAddr[4];
107+ uint8_t* deviceC = params->outputAddr[0];
108+ 
109+ auto layoutA = tla::MakeLayout<ElementA, LayoutTagA>(m, k);
110+ auto layoutPrologueB = tla::MakeLayout<ElementPrologueB, LayoutTagPrologueB>(k, n);
111+ auto layoutMxScaleA = tla::MakeMxScaleLayout<ElementMxScale, LayoutTagA, false>(m, mxScaleK);
112+ auto layoutMxScaleB = tla::MakeMxScaleLayout<ElementMxScale, LayoutTagMxScaleB, true>(mxScaleK, n);
113+ auto layoutC = tla::MakeLayout<ElementC, LayoutTagC>(m, n);
114+ 
115+ using TileCopy = Catlass::Gemm::Tile::PackedMxA8W4TileCopyTla<
116+ ArchTag, ElementA, LayoutTagA, ElementPrologueB, LayoutTagPrologueB, ElementB, LayoutTagB, ElementMxScale,
117+ decltype(layoutMxScaleA), ElementMxScale, decltype(layoutMxScaleB), ElementC, LayoutTagC, ElementBias, false,
118+ Catlass::Gemm::Tile::ScaleGranularity::PER_TENSOR>;
119+ 
120+ using BlockMmad = Catlass::Gemm::Block::BlockMmadA8W4Mx<
121+ DispatchPolicyMmad, L1TileShape, L0TileShape, ElementA, ElementB, ElementC, ElementBias, TileCopy>;
122+ 
123+ using BlockPrologue = Catlass::Gemm::Block::BlockPrologue<
124+ DispatchPolicyPrologue, PrologueSrcType, PrologueDstType, L1TileShape, TileCopy>;
125+ 
126+ using MatmulKernel = Catlass::Gemm::Kernel::A8W4GroupedMxMatmul<
127+ BlockMmad, BlockPrologue, BlockEpilogue, BlockScheduler, ElementGroupList>;
128+ 
129+ typename MatmulKernel::Arguments arguments{
130+ problemShape, problemCount, deviceGroupList,
131+ deviceA, layoutA, devicePrologueB, layoutPrologueB, deviceMxScaleA,
132+ layoutMxScaleA, deviceMxScaleB, layoutMxScaleB, deviceC, layoutC, nullptr};
133+ 
134+ uint64_t taskNum64 = static_cast<uint64_t>(CeilDiv(m, tla::get<0>(L1TileShape{}))) *
135+ static_cast<uint64_t>(CeilDiv(n, tla::get<1>(L1TileShape{})));
136+ uint32_t taskNum = static_cast<uint32_t>(std::min(taskNum64, static_cast<uint64_t>(UINT32_MAX)));
137+ uint32_t aicCoreUsed = std::min(blockNum, taskNum);
138+ 
139+ Catlass::RunKernel<MatmulKernel>(arguments, stream, aicCoreUsed);
140+}
Mtests/optest/kernels/CMakeLists.txt+1-0
@@ -274,6 +274,7 @@ add_subdirectory(69_ascend950_tail_multi_core_splitk_matmul)
274add_subdirectory(70_ascend950_flash_attention_chunk_prefill)274add_subdirectory(70_ascend950_flash_attention_chunk_prefill)
275add_subdirectory(71_ascend950_fp8_mx_grouped_matmul_finalize_routing)275add_subdirectory(71_ascend950_fp8_mx_grouped_matmul_finalize_routing)
276add_subdirectory(73_ascend950_matmul_full_loadA)276add_subdirectory(73_ascend950_matmul_full_loadA)
277+add_subdirectory(74_ascend950_weight_quant_a8w4_grouped_mx_matmul)
277add_subdirectory(56_ascend950_basic_conv2d_tla)278add_subdirectory(56_ascend950_basic_conv2d_tla)
278add_subdirectory(72_ascend950_fp8_mx_flash_attention_infer)279add_subdirectory(72_ascend950_fp8_mx_flash_attention_infer)
279add_subdirectory(80_grouped_matmul_slice_m_gelu)280add_subdirectory(80_grouped_matmul_slice_m_gelu)
Mtests/optest/src/catlass_torch.cpp+5-0
@@ -46,6 +46,7 @@
46#include "template/group_gemm.h"46#include "template/group_gemm.h"
47#include "template/mx_grouped_matmul_swiglu_mx_quant.h"47#include "template/mx_grouped_matmul_swiglu_mx_quant.h"
48#include "template/a8w4_mx_matmul.h"48#include "template/a8w4_mx_matmul.h"
49+#include "template/a8w4_grouped_mx_matmul.h"
49#include "template/svd_quant_matmul.h"50#include "template/svd_quant_matmul.h"
50#include "template/conv_bias.h"51#include "template/conv_bias.h"
51 52 
@@ -350,6 +351,10 @@ using Ascend950A8W4MxMatmulOp = A8W4MxMatmulLike<CatlassKernel::Ascend950A8W4MxM
350static auto& ascend950_a8w4_mx_matmul = Ascend950A8W4MxMatmulOp::Run;351static auto& ascend950_a8w4_mx_matmul = Ascend950A8W4MxMatmulOp::Run;
351REGISTER_TORCH_FUNC(ascend950_a8w4_mx_matmul);352REGISTER_TORCH_FUNC(ascend950_a8w4_mx_matmul);
352 353 
354+using Ascend950A8W4GroupedMxMatmulOp = A8W4GroupedMxMatmulLike<CatlassKernel::Ascend950A8W4GroupedMxMatmul>;
355+static auto& ascend950_a8w4_grouped_mx_matmul = Ascend950A8W4GroupedMxMatmulOp::Run;
356+REGISTER_TORCH_FUNC(ascend950_a8w4_grouped_mx_matmul);
357+ 
353using Ascend950Fp8MxGroupedMatmulSliceMSwigluMxQuantOp =358using Ascend950Fp8MxGroupedMatmulSliceMSwigluMxQuantOp =
354 GroupedMxSwigluMxQuantMatmulLike<CatlassKernel::Ascend950Fp8MxGroupedMatmulSliceMSwigluMxQuant>;359 GroupedMxSwigluMxQuantMatmulLike<CatlassKernel::Ascend950Fp8MxGroupedMatmulSliceMSwigluMxQuant>;
355static auto& ascend950_fp8_mx_grouped_matmul_slice_m_swiglu_mx_quant =360static auto& ascend950_fp8_mx_grouped_matmul_slice_m_swiglu_mx_quant =
Atests/optest/src/include/template/a8w4_grouped_mx_matmul.h+160-0
@@ -0,0 +1,160 @@
1+/**
2+ * This program is free software, you can redistribute it and/or modify.
3+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
4+ * This file is a part of the CANN Open Software.
5+ * Licensed under CANN Open Software License Agreement Version 2.0 (the "License").
6+ * Please refer to the License for details. You may not use this file except in compliance with the License.
7+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
8+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. See LICENSE in the root of
9+ * the software repository for the full text of the License.
10+ */
11+ 
12+#ifndef OPTEST_A8W4_GROUPED_MX_MATMUL_H
13+#define OPTEST_A8W4_GROUPED_MX_MATMUL_H
14+ 
15+#include <torch/torch.h>
16+#include <tiling/platform/platform_ascendc.h>
17+ 
18+#include "catlass_kernel_jit.h"
19+#include "common/run_npu_func.h"
20+#include "mx_matmul.h"
21+#include "torch_utils.h"
22+#include "type_utils.hpp"
23+ 
24+namespace CatlassKernelWrapper {
25+ 
26+using A8W4GroupedKernelFn = void (*)(
27+ const uint32_t, aclrtStream, const CatlassKernel::TParams&, const CatlassKernel::GroupedMatmulParams&);
28+ 
29+/**
30+ * @brief Adapter for example 74_ascend950_weight_quant_a8w4_grouped_mx_matmul.
31+ *
32+ * Grouped MX A8W4 matmul: C = (MxScaleA * A_fp8) @ (MxScaleB * B_fp4), sliced on M.
33+ * A is shared across groups (float8_e4m3fn, RowMajor M×K). B is the packed FP4
34+ * prologue (int8 bytes, Weight4BitnZ layout, per group). Scales are float8_e8m0fnu.
35+ * Output is FP32 (M×N) — the kernel accumulates in float per the CType convention.
36+ */
37+template <A8W4GroupedKernelFn KernelFunc>
38+struct A8W4GroupedMxMatmulLike {
39+ using OutputType = at::Tensor;
40+ 
41+ static OutputType Run(
42+ const at::Tensor& mat1, const at::Tensor& mat2,
43+ const at::Tensor& groupList,
44+ const at::Tensor& mx_scale_a, const at::Tensor& mx_scale_b)
45+ {
46+ CatlassKernel::TParams tParams;
47+ CatlassKernel::GroupedMatmulParams params;
48+ 
49+ // ── Validation ──
50+ CheckNpuTensor(mat1, "mat1");
51+ CheckNpuTensor(mat2, "mat2");
52+ CheckNpuTensor(groupList, "groupList");
53+ CheckNpuTensor(mx_scale_a, "mx_scale_a");
54+ CheckNpuTensor(mx_scale_b, "mx_scale_b");
55+ CheckSameDevice(mat1, "mat1", mat2, "mat2");
56+ CheckSameDevice(mat1, "mat1", groupList, "groupList");
57+ CheckSameDevice(mat1, "mat1", mx_scale_a, "mx_scale_a");
58+ CheckSameDevice(mat1, "mat1", mx_scale_b, "mx_scale_b");
59+ CheckMxScaleDType(mx_scale_a, "mx_scale_a");
60+ CheckMxScaleDType(mx_scale_b, "mx_scale_b");
61+ 
62+ TORCH_CHECK(
63+ mat1.scalar_type() == torch::kFloat8_e4m3fn,
64+ "mat1 must have dtype torch.float8_e4m3fn, got ", mat1.scalar_type());
65+ TORCH_CHECK(mat1.dim() == 2, "mat1 must be 2-D with shape (M, K)");
66+ TORCH_CHECK(
67+ mat2.scalar_type() == torch::kInt8,
68+ "mat2 must be int8 packed FP4 prologue bytes (Weight4BitnZ layout), got ", mat2.scalar_type());
69+ TORCH_CHECK(
70+ groupList.scalar_type() == torch::kInt64,
71+ "groupList must be int64, got ", groupList.scalar_type());
72+ TORCH_CHECK(mat1.is_contiguous(), "mat1 must be contiguous");
73+ TORCH_CHECK(mat2.is_contiguous(), "mat2 must be contiguous");
74+ TORCH_CHECK(mx_scale_a.is_contiguous(), "mx_scale_a must be contiguous");
75+ TORCH_CHECK(mx_scale_b.is_contiguous(), "mx_scale_b must be contiguous");
76+ 
77+ const int64_t m = mat1.size(0);
78+ const int64_t k = mat1.size(1);
79+ const uint32_t g = static_cast<uint32_t>(groupList.numel());
80+ 
81+ // N is inferred from mx_scale_b: (G, N, mxScaleAlignedK/2, 2) for the grouped case.
82+ int64_t n = 0;
83+ if (mx_scale_b.dim() == 4) {
84+ TORCH_CHECK(
85+ mx_scale_b.size(0) == static_cast<int64_t>(g),
86+ "mx_scale_b dim 0 must equal group count, got ", mx_scale_b.size(0), " expected ", g);
87+ n = mx_scale_b.size(1);
88+ } else if (mx_scale_b.dim() == 3) {
89+ // Single-group fallback: (N, mxScaleAlignedK/2, 2)
90+ n = mx_scale_b.size(0);
91+ } else {
92+ TORCH_CHECK(false, "mx_scale_b must be 4-D (G, N, mxScaleAlignedK/2, 2) or 3-D (N, mxScaleAlignedK/2, 2)");
93+ }
94+ 
95+ // Packed FP4 bytes for Weight4BitnZ: K padded to 32, N padded to 16, fractal (32, 16).
96+ const int64_t kPadded = ((k + 31) / 32) * 32;
97+ const int64_t nPadded = ((n + 15) / 16) * 16;
98+ const int64_t packedBytesPerGroup = kPadded * nPadded / 2; // 2 FP4 nibbles per byte
99+ const int64_t expectedBytes = packedBytesPerGroup * static_cast<int64_t>(g);
100+ TORCH_CHECK(
101+ mat2.numel() >= expectedBytes,
102+ "mat2 packed bytes insufficient: need at least ", expectedBytes,
103+ " (g=", g, ", kPadded=", kPadded, ", nPadded=", nPadded, "), got ", mat2.numel());
104+ 
105+ // MX scale size validation.
106+ uint32_t mxScaleK = static_cast<uint32_t>((k + 31) / 32); // CeilDiv<32>(k)
107+ uint32_t mxScaleAlignedK = ((mxScaleK + 1) / 2) * 2; // RoundUp<2>(mxScaleK)
108+ const int64_t scaleANumel = m * mxScaleAlignedK;
109+ const int64_t scaleBNumel = static_cast<int64_t>(g) * n * mxScaleAlignedK;
110+ TORCH_CHECK(
111+ mx_scale_a.numel() == scaleANumel,
112+ "mx_scale_a must have ", scaleANumel, " elements (m=", m, ", mxScaleAlignedK=", mxScaleAlignedK,
113+ "), got ", mx_scale_a.numel());
114+ TORCH_CHECK(
115+ mx_scale_b.numel() == scaleBNumel,
116+ "mx_scale_b must have ", scaleBNumel, " elements (g=", g, ", n=", n, ", mxScaleAlignedK=", mxScaleAlignedK,
117+ "), got ", mx_scale_b.numel());
118+ 
119+ // ── TParams (compile-time JIT parameters) ──
120+ tParams.element["A"] = TorchDtypeToAclDtype(mat1.scalar_type());
121+ tParams.element["B"] = test_utils::TypeCast<std::string, aclDataType>("float4_e2m1fn_x2");
122+ tParams.element["C"] = ACL_FLOAT;
123+ tParams.element["MX_SCALE"] = test_utils::TypeCast<std::string, aclDataType>("float8_e8m0fnu");
124+ tParams.transpose["A"] = false;
125+ tParams.transpose["B"] = true;
126+ tParams.transpose["C"] = false;
127+ tParams.useNz["A"] = false;
128+ tParams.useNz["B"] = true; // Weight4BitnZ prologue layout
129+ tParams.useNz["C"] = false;
130+ tParams.l1TileShape = {256, 256, 256};
131+ tParams.l0TileShape = {256, 256, 128};
132+ tParams.swizzle = {3, 0, 0};
133+ 
134+ // ── Params (runtime) ──
135+ params.m = static_cast<uint32_t>(m);
136+ params.n = static_cast<uint32_t>(n);
137+ params.k = static_cast<uint32_t>(k);
138+ params.batch = g;
139+ params.sliceMode = CatlassKernel::GroupedMatmulParams::SliceMode::M;
140+ params.inputAddr.resize(5);
141+ params.inputAddr[0] = static_cast<uint8_t*>(const_cast<void*>(mat1.storage().data()));
142+ params.inputAddr[1] = static_cast<uint8_t*>(const_cast<void*>(mat2.storage().data()));
143+ params.inputAddr[2] = static_cast<uint8_t*>(const_cast<void*>(groupList.storage().data()));
144+ params.inputAddr[3] = static_cast<uint8_t*>(const_cast<void*>(mx_scale_a.storage().data()));
145+ params.inputAddr[4] = static_cast<uint8_t*>(const_cast<void*>(mx_scale_b.storage().data()));
146+ 
147+ OutputType output = GetOutputTensor({params.m, params.n}, torch::kFloat32);
148+ params.outputAddr.resize(1);
149+ params.outputAddr[0] = static_cast<uint8_t*>(const_cast<void*>(output.storage().data()));
150+ 
151+ aclrtStream stream = c10_npu::getCurrentNPUStream().stream(false);
152+ uint32_t aicCoreNum = platform_ascendc::PlatformAscendCManager::GetInstance()->GetCoreNumAic();
153+ RUN_NPU_FUNC(KernelFunc, aicCoreNum, stream, tParams, params);
154+ return output;
155+ }
156+};
157+ 
158+} // namespace CatlassKernelWrapper
159+ 
160+#endif
Mtests/optest/tests/a8w4_golden.py+106-2
@@ -50,10 +50,9 @@ def _case_random_seed(*values) -> int:
50 return seed50 return seed
51 51 
52 52 
53-def _set_case_random_seed(*values) -> int:53+def _set_case_random_seed(*values) -> None:
54 seed = _case_random_seed(*values)54 seed = _case_random_seed(*values)
55 torch.manual_seed(seed)55 torch.manual_seed(seed)
56- return seed
57 56 
58 57 
59def _build_e4m3_lut() -> torch.Tensor:58def _build_e4m3_lut() -> torch.Tensor:
@@ -292,3 +291,108 @@ def compare_a8w4_result(result: torch.Tensor, expected: torch.Tensor, k: int) ->
292 threshold = rtol * torch.clamp(expected_cpu.abs(), min=1.0)291 threshold = rtol * torch.clamp(expected_cpu.abs(), min=1.0)
293 max_diff = diff.max().item()292 max_diff = diff.max().item()
294 return bool((diff <= threshold).all()), max_diff293 return bool((diff <= threshold).all()), max_diff
294+ 
295+ 
296+def _trans_nd2nz(input_data: torch.Tensor) -> torch.Tensor:
297+ """Transform (G, N_pad, K_pad) int8 indices to nZ fractal (G, K_blk, N_blk, 16, 32)."""
298+ g, n_pad, k_pad = input_data.shape
299+ return input_data.reshape(g, n_pad // 16, 16, k_pad // 32, 32).permute(0, 3, 1, 2, 4)
300+ 
301+ 
302+def prepare_a8w4_grouped_mx_inputs(
303+ group_sizes, n: int, k: int, device: str = "npu", trans_a: int = 0, trans_b: int = 1
304+):
305+ """Build grouped A8W4 MX inputs (Weight4BitnZ B layout) for example 74.
306+ 
307+ Mirrors ``examples/74_ascend950_weight_quant_a8w4_grouped_mx_matmul/gen_data.py``
308+ with ``isNz=1`` (the default nZ prologue layout).
309+ 
310+ Returns:
311+ (a_fp8, b_int8, group_list, a_scale, b_scale, expected) where:
312+ - a_fp8: (M, K) float8_e4m3fn, M = sum(group_sizes), shared across groups.
313+ - b_int8: flat int8 packed FP4 bytes in Weight4BitnZ fractal layout.
314+ - group_list: (G,) int64 non-cumsum group sizes.
315+ - a_scale: (M, mxScaleAlignedK/2, 2) float8_e8m0fnu.
316+ - b_scale: (G, N, mxScaleAlignedK/2, 2) float8_e8m0fnu.
317+ - expected: (M, N) float32 dequant reference.
318+ """
319+ group_sizes = tuple(int(s) for s in group_sizes)
320+ if not group_sizes:
321+ raise ValueError("group_sizes must not be empty")
322+ if any(s <= 0 for s in group_sizes):
323+ raise ValueError(f"group_sizes must be positive, got {group_sizes}")
324+ m = sum(group_sizes)
325+ g = len(group_sizes)
326+ _set_case_random_seed(74, m, n, k, trans_a, trans_b)
327+ 
328+ # ── A: shared, fp8 e4m3, quantize along K (axis=1) ──
329+ a_fp8, a_scale_raw, a_deq = _gen_data_fp8_e4m3(m, k, 1)
330+ a_scale = a_scale_raw.reshape(a_scale_raw.shape[0], a_scale_raw.shape[1] // 2, 2).contiguous()
331+ if trans_a == 1:
332+ a_scale = a_scale.permute(1, 0, 2).contiguous()
333+ 
334+ # ── B: per group, fp4 e2m1, quantize along K (axis=0), nZ layout ──
335+ b_indices_list = []
336+ b_scale_list = []
337+ b_deq_list = []
338+ for _ in range(g):
339+ matrix = torch.randn((k, n), dtype=torch.float32)
340+ quantized_vals, scale, dequantized = _quantize_fp4(matrix, "E2M1", axis=0)
341+ if trans_b == 1:
342+ quantized_vals = quantized_vals.t().contiguous()
343+ _, fp4_indices = _quantize_to_fp4_lut(quantized_vals, "E2M1")
344+ b_indices_list.append(fp4_indices) # (N, K) or (K, N), uint8
345+ b_scale_list.append(scale.to(torch.float8_e8m0fnu))
346+ b_deq_list.append(dequantized) # (K, N)
347+ 
348+ b_indices_stacked = torch.stack(b_indices_list, dim=0) # (G, N, K) when trans_b=1
349+ 
350+ # Pad K to 32, N to 16 for the Weight4BitnZ fractal (32, 16).
351+ pad_k = (32 - k % 32) % 32
352+ pad_n = (16 - n % 16) % 16
353+ if pad_k > 0 or pad_n > 0:
354+ b_indices_stacked = torch.nn.functional.pad(b_indices_stacked, (0, pad_k, 0, pad_n), "constant", 0)
355+ 
356+ n_pad = n + pad_n
357+ k_pad = k + pad_k
358+ b_quant_nz = _trans_nd2nz(b_indices_stacked) # (G, k_pad//32, n_pad//16, 16, 32)
359+ b_quant_flat = b_quant_nz.reshape(-1, 32) # (total_rows, 32)
360+ b_packed = _pack_fp4_nibbles(b_quant_flat) # (total_rows, 16) bytes
361+ b_int8 = _clone_int8_storage(b_packed.flatten())
362+ 
363+ # ── B scales: reshape + permute to (G, N, mxScaleAlignedK/2, 2) ──
364+ b_scale_stacked = torch.stack(b_scale_list, dim=0) # (G, num_blocks_padded, N)
365+ b_scale_processed_list = []
366+ for i in range(g):
367+ scale = b_scale_stacked[i] # (num_blocks_padded, N)
368+ scale = scale.reshape(scale.shape[0] // 2, 2, scale.shape[1])
369+ if trans_b == 1:
370+ scale = scale.permute(2, 0, 1).contiguous() # (N, num_blocks_padded/2, 2)
371+ else:
372+ scale = scale.permute(0, 2, 1).contiguous() # (num_blocks_padded/2, N, 2)
373+ b_scale_processed_list.append(scale)
374+ b_scale_processed = torch.stack(b_scale_processed_list, dim=0).contiguous()
375+ 
376+ # ── Group list ──
377+ group_list = torch.tensor(group_sizes, dtype=torch.int64)
378+ 
379+ # ── Expected: per-group dequant matmul concatenated along M ──
380+ c_fp32_list = []
381+ m_offset = 0
382+ for i, group_m in enumerate(group_sizes):
383+ if group_m == 0:
384+ continue
385+ a_group = a_deq[m_offset:m_offset + group_m] # (group_m, K)
386+ b_group = b_deq_list[i] # (K, N)
387+ c_fp32_list.append(a_group @ b_group)
388+ m_offset += group_m
389+ expected = torch.cat(c_fp32_list, dim=0) # (M, N)
390+ 
391+ if device == "npu":
392+ a_fp8 = a_fp8.npu()
393+ b_int8 = b_int8.npu()
394+ group_list = group_list.npu()
395+ a_scale = a_scale.npu()
396+ b_scale_processed = b_scale_processed.npu()
397+ 
398+ return a_fp8, b_int8, group_list, a_scale, b_scale_processed, expected
Mtests/optest/torch_catlass/ops/__init__.py+22-13