CANN/cann-bench: GroupedMatmul算子API描述

GroupedMatmul 算子 API 描述【免费下载链接】cann-bench评测AI在处理CANN领域代码任务的能力涵盖算子生成、算子优化等领域支撑模型选型、训练效果评估统一量化评估标准识别Agent能力短板构建CANN领域评测平台推动AI能力在CANN领域的持续演进。项目地址: https://gitcode.com/cann/cann-bench1. 算子简介分组矩阵乘法算子。激活x沿M轴合并为单 tensor[M, K]权重weight按 expert 维堆叠为[E, K, N]或[E, N, K]通过group_listcumsum 语义描述每组 token 在M轴上的边界。主要应用场景MoEMixture of Experts模型中多专家的并行矩阵运算多头注意力机制中的分组线性变换批量处理同一 K、N 维度、按 M 切分多组的矩阵乘法算子特征难度等级L3Contraction输入 / 输出容器与level4/grouped_matmul_swiglu_quant对齐x单 tensor weight堆叠 group_list切组所有组的K与N在 expert 维上一致仅M各组可不同2. 算子定义数学公式$$ \begin{aligned} \text{对每个专家 } g \in [0, E),\ \text{根据 } group_list\ (\text{cumsum}) \text{ 取属于该组的 token 行 } rows_g: \ y[rows_g] x[rows_g] \cdot weight[g] \ (\ bias[g]) \end{aligned} $$其中xshape 为[M, K]所有组沿M轴合并weightshape 为[E, K, N]transpose_weightfalse或[E, N, K]transpose_weighttruebiasshape 为[E, N]可选group_list长度Ecumsum 语义最后一个值等于Myshape 为[M, N]group_list 语义group_list [c_0, c_1, ..., c_{E-1}]表示前 g1 个 expert 累计接收c_g个 token。每组的行范围为[c_{g-1}, c_g)约定c_{-1} 0。允许c_g c_{g-1}表示该 expert 为空组跳过计算。3. 接口规范算子原型cann_bench.grouped_matmul( Tensor x, Tensor weight, Tensor? bias, int[] group_list, int split_item 0, bool transpose_weight False, ) - Tensor y输入参数参数类型Shape描述xTensor[M, K]激活矩阵所有组沿M合并weightTensor[E, K, N]或[E, N, K]专家权重按 expert 维堆叠biasTensor?可选[E, N]偏置group_listList[int]长度E累计 token 数cumsum最后值等于Msplit_itemint—输出切分模式详见下transpose_weightbool—是否转置权重输出参数Shapedtype描述y[M, N]split_item2/3或 List[Tensor] 长度 E0/1与 x 相同每组[m_i, N]数据类型输入 (x) dtype输入 (weight) dtype输入 (bias) dtype输出 (y) dtypefloat16float16float16float16bfloat16bfloat16bfloat16bfloat16float32float32float32float32split_item 取值split_item输出形式说明0 / 1List[Tensor]长度 E按group_list把y切回每组[m_i, N]空组返回长度 0 的[0, N]张量2 / 3单 Tensor[M, N]直接返回沿 M 轴拼接好的结果规则与约束K 一致性x.shape[1]必须等于weight中 K 维weight.shape[1]或weight.shape[2]由transpose_weight决定N 一致性所有 expert 共享同一Ngroup_list长度等于E严格非递减最后一个值等于M空组允许group_list[i] group_list[i-1]对应 experti不参与计算transpose_weightfalseweight 形状[E, K, N]每片weight[g]直接参与 matmultrueweight 形状[E, N, K]每片weight[g]需 transpose 最后两维后参与 matmul维度限制每维大小在 32 字节对齐后应小于 int32 最大值支持范围输入 tensor 各维度与参数的支持范围维度 / 参数范围备注Eexpert / 组数1 ~ 128cases.csv 实测 1 ~ 8等于weight.shape[0]与group_list长度Mtoken 总数 / x 行数1 ~ 16384cases.csv 实测 64 ~ 7001含非 2 幂奇数 1009 / 1023 / 1031 / 2047 / 7001等于group_list[-1]Kcontraction 维1 ~ 8192cases.csv 实测 32 ~ 1040含非 2 幂奇数 127 / 511、非 2 幂 1023 / 1040N输出维1 ~ 8192cases.csv 实测 32 ~ 2048含非 2 幂奇数 255 / 1023、非 2 幂 1056m_i每组 token 数0 ~Mcases.csv 实测 0 ~ 1751允许m_i 0表示空组case_7 含 idx2,5 两个空组group_list长度E单调非递减末值 Mcumsum 语义相邻相等表示空组元素类型 int64split_item{0, 1, 2, 3}cases.csv 实测 0 / 2 / 30/1 →List[Tensor]长度E2/3 → 单 Tensor[M, N]transpose_weight{false, true}cases.csv 实测均覆盖false → weight[E, K, N]true → weight[E, N, K]需对最后两维 transposebias可选None 或[E, N]cases.csv 实测含 bias / 无 bias 均覆盖提供时 dtype 与 x 一致或 fp32约束x.shape[1] K且weight中对应维weight.shape[1]若 twfalse否则weight.shape[2]必须等于Kweight.shape[0] E所有 expert 共享同一Ngroup_list[-1] Mgroup_list[i] group_list[i-1]每维大小在 32 字节对齐后应小于 int32 最大值4. 精度要求采用生态算子精度标准进行验证。误差指标平均相对误差MERE采样点中相对误差平均值$$ \text{MERE} \text{avg}(\frac{\text{abs}(actual - golden)}{\text{abs}(golden)\text{1e-7}}) $$最大相对误差MARE采样点中相对误差最大值$$ \text{MARE} \max(\frac{\text{abs}(actual - golden)}{\text{abs}(golden)\text{1e-7}}) $$通过标准数据类型FLOAT16BFLOAT16FLOAT32HiFLOAT32FLOAT8 E4M3FLOAT8 E5M2通过阈值(Threshold)2^-102^-72^-132^-112^-32^-2当平均相对误差 MERE Threshold最大相对误差 MARE 10 * Threshold 时判定为通过。5. 标准 Golden 代码import torch from typing import List, Optional, Union GroupedMatmul 算子 Torch Golden 参考实现 分组矩阵乘法算子x 沿 M 轴合并、weight 按 expert 维堆叠。 公式对每个专家 g ∈ [0, E)根据 group_listcumsum取属于该组的 token 行 rows_g y[rows_g] x[rows_g] weight[g] ( bias[g]) def grouped_matmul( x: torch.Tensor, weight: torch.Tensor, bias: Optional[torch.Tensor] None, group_listNone, split_item: int 0, transpose_weight: bool False, ) - Union[torch.Tensor, List[torch.Tensor]]: M, K x.shape E weight.shape[0] if transpose_weight: N weight.shape[1] else: N weight.shape[2] if isinstance(group_list, torch.Tensor): ends group_list.to(torch.int64).tolist() else: ends list(group_list) starts [0] ends[:-1] y torch.zeros((M, N), dtypex.dtype, devicex.device) x_f x.float() for g in range(E): s, e starts[g], ends[g] if s e: continue w_g weight[g].float() if transpose_weight: mm torch.matmul(x_f[s:e], w_g.transpose(-2, -1)) else: mm torch.matmul(x_f[s:e], w_g) if bias is not None: mm mm bias[g].float().unsqueeze(0) y[s:e] mm.to(x.dtype) if split_item in (0, 1): return [y[starts[g]:ends[g]] for g in range(E)] return y6. 额外信息算子调用示例import torch import cann_bench M, K, N, E 128, 256, 512, 4 x torch.randn(M, K, dtypetorch.float16, devicenpu) weight torch.randn(E, K, N, dtypetorch.float16, devicenpu) bias torch.randn(E, N, dtypetorch.float16, devicenpu) group_list [32, 64, 96, 128] # cumsum 语义最后值 M # split_item2: 输出单 tensor [M, N] y cann_bench.grouped_matmul(x, weight, bias, group_list, split_item2, transpose_weightFalse) # y shape: [128, 512] # split_item0: 输出 List[Tensor] 长度 E y_list cann_bench.grouped_matmul(x, weight, bias, group_list, split_item0, transpose_weightFalse) # y_list [Tensor[32,512], Tensor[32,512], Tensor[32,512], Tensor[32,512]] # transpose_weightTrue: weight 形状 [E, N, K]需 transpose 后 matmul weight_t torch.randn(E, N, K, dtypetorch.float16, devicenpu) y cann_bench.grouped_matmul(x, weight_t, bias, group_list, split_item2, transpose_weightTrue)【免费下载链接】cann-bench评测AI在处理CANN领域代码任务的能力涵盖算子生成、算子优化等领域支撑模型选型、训练效果评估统一量化评估标准识别Agent能力短板构建CANN领域评测平台推动AI能力在CANN领域的持续演进。项目地址: https://gitcode.com/cann/cann-bench创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考