人工智能深度学习机器学习【免费下载链接】onnxOpen standard for machine learning interoperability项目地址https://gitcode.com/gh_mirrors/onn/onnx点击查看免费下载导读本文深入解读 ONNX 官方 RFCdocs/proposals/0010-GroupedMatMulOpProposal.md状态为 under discussion提出的标准算子GroupedMatMul它把一批 token 向量乘以一组专家expert权重矩阵、每个 token 选择若干专家这一 Mixture-of-ExpertsMoE前馈层的核心计算封装为紧凑、可融合的 ONNX 函数算子。读完本文你将掌握GroupedMatMul的输入输出语义、函数分解式的规范定义、MoE 前馈层的标准建图方式、形状推断规则、量化QDQ组合方式、边界行为与测试用例以及它与 PyTorchgrouped_mm、cuBLAS grouped-GEMM、ONNX Runtimecom.microsoft.MoE等已有实现的对应关系。背景为什么需要 GroupedMatMulMoEMixture-of-Experts前馈层是 Mixtral、DeepSeek、Grok、Switch Transformer 等大语言模型的核心组件。其核心计算为给定一批M个 token 向量和num_groups个专家权重矩阵由 router 为每个 token 选择一个或多个专家矩阵进行乘法。这种被称为grouped matrix multiplication / grouped GEMM的模式理论上可以用现有标准 ONNX 算子表达但朴素的未融合实现非常低效、不实用朴素分解不切实际自然的做法是Gather权重 →Expandtoken → 批量MatMul这会物化出完整的[M×k, K, N]权重切片和[M×k, K]的 token 副本。对真实的 MoE 层而言这些中间张量动辄数 GB导致分解方案在实践中不可行。融合内核的优势融合的 grouped-GEMM 内核每个专家权重矩阵只需处理一次与选择它的 token 数量无关且每个 token 行在其k个专家间复用、无需拷贝。将GroupedMatMul加入 ONNX 标准ai.onnx域即可为 MoE 模型提供更紧凑、更高效的表示。提案同时指出ONNX Runtime 目前没有独立的 grouped-matmul contrib op而是提供更大的融合算子com.microsoft.MoE以及量化版本com.microsoft.QMoE来实现整个 MoE 前馈层本提案与 onnx/onnx#7902 相关文中仅作背景陈述外部链接不展开。算子语义Guide-level ExplanationGroupedMatMul的输入inputtoken 矩阵形状[M, K]weightsG个专家权重矩阵的堆叠形状[G, K, N]group_indices形状[M, k]为每个 token 选择其要乘的k个专家group索引可选bias逐组偏置形状[G, N]。输出形状为[M, k, N]每个 token 与它选中的k个专家分别相乘含可选偏置后的逐专家结果。任何对这些逐专家结果的加权组合例如 MoE 层中 router 的 top-k加权求和都应当由周围的Mul/ReduceSum标准算子表达下文给出示例。对于稠密单专家场景令k 1即可。关键设计算子内部复用每个 token 行跨其k个选中专家因此不需要Expandtoken 批次——这是与朴素分解的本质区别。典型用法MoE 前馈层一个标准的 top-k MoE FFN含两个投影可以映射为两个GroupedMatMul算子且无需Expandtoken 批次# Notation: B batch, S sequence length, H hidden dim, F FFN inner dim # E num_experts, k experts_per_token scores Softmax(MatMul(hidden, router_W)) # [B, S, E] values, indices TopK(scores, k) # [B, S, k] h Reshape(hidden, [B*S, H]) idx Reshape(indices, [B*S, k]) val Reshape(values, [B*S, k]) # --- Up projection: per-expert output --- # output shape: [B*S, k, F] h_up GroupedMatMul(h, expert_gate_W, idx) # expert_gate_bias (optional) h_up SiLU(h_up) # Reshape for down projection: treat each (token, expert-slot) pair as a row. h_flat Reshape(h_up, [B*S*k, F]) idx2 Reshape(idx, [B*S*k, 1]) # each flat row selects one expert # --- Down projection: per-expert output, then router-weighted sum over the k slots --- d GroupedMatMul(h_flat, expert_down_W, idx2) # [B*S*k, 1, H] d Reshape(d, [B*S, k, H]) # regroup the k slots out ReduceSum(d * Unsqueeze(val, -1), axis1) # [B*S, H] weighted sum over k out Reshape(out, [B, S, H])两个投影都只用GroupedMatMul完成分组矩阵乘法下投影中的 top-krouter 加权求和用MulReduceSum在重组后的k个槽位上显式表达。之所以不把该加权和融进算子是因为这里的k个专家拥有各自不同的逐槽位输入上投影激活后的输出而非共享同一输入行详见后文设计抉择。参考级规范函数分解单一规范性定义GroupedMatMul的语义被定义为一个 ONNX函数function其含义严格等于下述标准 ONNX 算子的组合。由于可选bias输入的存在会改变生成的图函数体是上下文相关的context-dependent——通过SetContextDependentFunctionBodyBuilder构建与CenterCropPad等既有算子的风格一致。以函数定义语义有一个重要附带收益ONNX 的参考求值器onnx.reference.ReferenceEvaluator在未注册专用 Python 内核时会自动通过函数体执行算子因此不需要单独的参考实现文件——一份分解同时充当规范、文档与参考实现。这与仓库中onnx/reference/op_run.py的函数执行机制一致自定义函数在执行期实例化、复用同一实现。下文符号约定M、K、N为形状维度G weights.shape[0]专家数k group_indices.shape[1]每 token 选择的专家数。函数分解Function Decompositionidx_flat Reshapeallowzero 1(group_indices, [M*k]) W_sel Gather(weights, idx_flat, axis0) # [M*k, K, N] — duplicates weights! X Reshapeallowzero 1(Expand(Unsqueeze(input, 1), [M, k, K]), [M*k, 1, K]) # [M*k, 1, K] — copies tokens! r Reshapeallowzero 1(MatMul(X, W_sel), [M, k, N]) # [M, k, N] # If bias is present, add the per-group bias to each selected expert result: bias_sel Reshapeallowzero 1(Gather(bias, idx_flat, axis0), [M, k, N]) # [M, k, N] r r bias_sel # (only when bias present) output r # [M, k, N]上下文相关的函数体只产生两种情形输入 3bias存在时包含 bias 加法行否则省略。关于效率的重要说明上述分解会物化完整的[M*k, K, N]权重切片和[M*k, K]的 token 副本真实 MoE 层中这些中间量达 GB 级——这正是GroupedMatMul需要作为融合算子存在的原因它定义的是计算结果是什么而运行时被期望融合执行而非照搬朴素分解。分解对结果是规范性的对策略不是。边界情况Edge Cases分解已对以下特殊情况给出良定义行为下表记录之情形行为k 0未选中任何专家输出为空张量形状[M, 0, N]真实模型不应出现。k 1每 token 一个专家输出形状[M, 1, N]。G 1所有索引为 0等价于MatMul(input, weights[0])含可选 bias。M 0零 token 输入输出形状[0, k, N]无需计算。索引越界非法输入实现必须报错。算子规范Operator Specification名称与域字段值NameGroupedMatMulDomainai.onnx标准域Opset version下一个可用 opset如 27Since version随该 opset 新增需要说明截至本仓库当前版本VERSION_NUMBER为 1.24.0该算子仍处于提案阶段尚未在onnx/defs/下的 schema 与onnx/backend/test/case/node/测试目录中落地实现——文中规范均以提案文档为准。输入索引名称类型必选形状说明0inputT必选[M, K]行主序 token 矩阵。M个 tokenK为收缩hidden维。1weightsT必选[G, K, N]G个专家权重矩阵的堆叠每个为K × N。所有专家共享相同的K与N。2group_indicestensor(int64)必选[M, k]每个 token 每个槽位slot的组专家索引。每个 token 选择k个专家取值必须在[0, G)内。稠密单专家场景用k1。3biasT可选[G, N]逐组偏置向量加到每个专家的结果上。注意点G weights.shape[0]组 / 专家数。若调用方有[B, M, K]形状的批处理输入应先用Reshape将批维度并入M多数后端中这通常是零拷贝、仅元数据的视图构造。weights与bias对所有 token 相同即它们是模型参数而非逐 token 的。输出索引名称类型形状说明0outputT[M, k, N]逐专家结果每个 token 与它选中的k个专家分别相乘的结果含可选 bias。类型约束约束类型Ttensor(float)、tensor(float16)、tensor(bfloat16)group_indices恒为tensor(int64)。量化QDQ组合方式GroupedMatMul可以直接与 ONNX 标准的量化/反量化QDQ表示组合。量化激活、权重与可选 bias 在算子之前反量化其输出可在之后量化DequantizeLinear(input_q) ─┐ ├─ GroupedMatMul ─ QuantizeLinear (optional) DequantizeLinear(weights_q) ─┘对于形状[G, K, N]的专家权重沿axis1的分块DequantizeLinear支持形状为[G, ceil(K/B), N]的 scale这允许每个专家、每个 K 块、每个输出通道各自独立的 scale。逐专家、逐输出通道量化是其特例[G, 1, N]单个 K 大小的块。GroupedMatMul的 schema不需要任何量化参数或量化张量类型。为高效执行运行时应当识别并融合周围的 QDQ 模式避免物化完整的反量化专家权重张量。若未来需要显式整数累加、requantization 或打包权重格式应另立QLinearGroupedMatMul风格的算子单独考虑。属性无。所有配置都通过输入表达遵循 ONNX 在取值可能动态时偏好输入而非属性的一般原则。形状推断规则Shape Inference设M input.shape[0]K input.shape[1]G weights.shape[0]N weights.shape[2]k group_indices.shape[1]校验检查违反即报错input.rank 2weights.rank 3group_indices.rank 2且group_indices.shape[0] Mweights.shape[1] K收缩维一致若存在biasbias.shape [G, N]输出形状output.shape [M, k, N]测试用例提案计划将用例放入onnx/backend/test/case/node/groupedmatmul.py该目录下已有各算子的 node 级后端测试如onnx/backend/test/case/node/abs.py、add.py等新增测试将遵循同样的组织方式。测试 1 —— 稠密k1无 bias# 4 tokens, K3, G2 groups, N2, k1 input [[1, 0, -1], [0, 1, 2], [1, 1, 0], [0, 0, 1]] # shape [4, 3] weights [[[1, 0], [0, 1], [-1, 0]], [[0, 1], [1, 0], [ 0, 1]]] # shape [2, 3, 2] group_indices [[0], [1], [0], [1]] # shape [4, 1] # Expected output shape [4, 1, 2]: # token 0 - group 0: [1,0,-1] [[1,0],[0,1],[-1,0]] [101, 000] [2, 0] # token 1 - group 1: [0,1, 2] [[0,1],[1,0],[ 0,1]] [010, 002] [1, 2] # token 2 - group 0: [1,1, 0] [[1,0],[0,1],[-1,0]] [100, 010] [1, 1] # token 3 - group 1: [0,0, 1] [[0,1],[1,0],[ 0,1]] [000, 001] [0, 1] output [[[2, 0]], [[1, 2]], [[1, 1]], [[0, 1]]] # shape [4, 1, 2]测试 2 —— Top-kk2含 biasM, k, K, G, N 2, 2, 2, 3, 2 input [[1.0, 0.0], [0.0, 1.0]] # [2, 2] weights [[[1,0],[0,1]], [[0,1],[1,0]], [[1,1],[0,0]]] # [3,2,2] group_indices [[0, 1], [2, 0]] # [2, 2] bias [[0.1, 0.2], [0.3, 0.0], [0.5, 0.5]] # [3, 2] # token 0, slot 0 - g0: [1,0] [[1,0],[0,1]] [0.1,0.2] [1.1, 0.2] # token 0, slot 1 - g1: [1,0] [[0,1],[1,0]] [0.3,0.0] [0.3, 1.0] # token 1, slot 0 - g2: [0,1] [[1,1],[0,0]] [0.5,0.5] [0.5, 0.5] # token 1, slot 1 - g0: [0,1] [[1,0],[0,1]] [0.1,0.2] [0.1, 1.2] output [[[1.1, 0.2], [0.3, 1.0]], [[0.5, 0.5], [0.1, 1.2]]] # [2, 2, 2]测试 3 —— 空组某个专家未被使用# Group 1 receives no tokens. M, k, K, G, N 4, 1, 2, 3, 2 group_indices [[0], [0], [2], [2]] # group 1 unused # weights[1] is never accessed; output is well-defined.测试 4 —— 单组退化到 MatMulG 1 group_indices [[0], [0], [0]] # all tokens - group 0 # output MatMul(input, weights[0]) (reshaped from [M,1,N] to [M,1,N])设计抉择Rationale and Alternatives为什么这样设计将GroupedMatMul表达为上下文相关的 ONNX 函数可以获得单一规范性定义同时充当参考实现保持算子与周围图router、激活、reshape、加权求和的可组合性让运行时自行融合计算而不规定具体融合策略。将算子限定在分组矩阵乘法本身含可选的逐组bias使其小而通用MoE 层使用的 router 加权求和交给标准Mul/ReduceSum必要时可由后续独立算子解决。不做的后果没有GroupedMatMulMoE 模型要么用朴素的Gather/Expand/MatMul分解导出会物化 GB 级中间量、不实用要么依赖 ONNX Runtimecom.microsoft.MoE这类厂商私有 contrib op——无法在 ONNX 生态内移植。提案还评估了以下备选方案k1 的特殊形状当前方案要求group_indices是[M, k]的二维张量。备选方案是允许k1特例使用[M]的一维张量。显式批维本算子将M个 token 展平为单一维度实际使用中批处理时 token 可能是多维的如[Batch, Sequence]。可考虑支持额外批维以避免多余 Reshape甚至让M支持任意维数——但后者复杂度收益不成比例未采纳。更大粒度的融合算子实现中可能做更激进的融合如把下投影、激活、上投影整个融合ONNX Runtime 的 MoE contrib op 正是如此。其缺点是激活函数因模型而异、且新模型不断引入新激活导致 contrib op 需要持续更新以支持新激活目前没有好的解决方案。决策本算子不融合激活、不融合 combine。激活保持为独立 ONNX 算子以保证图跨众多 MoE 路由变体可组合。router 加权求和combine同样用标准Mul/ReduceSum表达而非折入GroupedMatMul标准 MoE 下投影中k个专家拥有各自不同的逐槽位输入加权求和的归约分组与 matmul 自身的k分组不同折入会迫使两种分组解耦并固定规范结果布局——规范与形状推断复杂度都会增加而这类融合更适合由未来专门算子干净地捕获。group_indicesvsgroup_offsets备选方案是使用已排序的 token 缓冲与整数 offsets更贴近 cuBLAS grouped-GEMM API。决策采用group_indices无序索引。索引天然与TopK/Gather组合且不要求调用方预排序 token 批次运行时内部自行排序。堆叠三维权重[G, K, N]vs 变长矩阵列表备选方案是每个专家可能具有不同K/N的权重张量序列异构专家一般情形。决策堆叠[G, K, N]。所有专家共享K、N是已部署 MoE 模型中的压倒性常见情形单一权重张量更简单。实现方仍可差异化处理单一张量内的不同专家切片例如应对整个张量过大无法同时装入内存的情况。已有实现参考Prior Art框架 / 库APIPyTorchtorch.nn.functional.grouped_mmPyTorch ≥ 2.5PyTorchtorch._grouped_mm/torch.ops.aten.mm_group内部JAXjax.lax.dot_general分组批处理cuBLAScublasGemmBatchedEx/cublasGemmGroupedBatchedExCUTLASSGroupedGemm内核OpenVINOGroupConvolution卷积类比ONNX Runtimecom.microsoft.MoE/com.microsoft.QMoEcontrib_ops——更大的融合 MoE 层而非独立 grouped matmulPyTorch 2.5 新增的torch.nn.functional.grouped_mm与本提案语义直接对应# PyTorch grouped_mm — same semantics out torch.nn.functional.grouped_mm(input, weight, offsNone) # offs are contiguous group offsets; our design uses indices instead # (see group_indices vs. group_offsets under Rationale and alternatives)未决问题Unresolved Questionsk 1时是否允许形状为[M]的一维group_indices还是始终要求二维[M, k]见上文k1 的特殊形状。是否引入显式批维例如[B, S, K]的多维 token 输入以省去周围Reshape还是保持展平的[M, K]形式见上文显式批维。未来展望Future Possibilities带归约的融合下投影可设计独立算子表示 MoE 下投影——自身k为 1每个展平的(token, slot)行选单个专家的GroupedMatMul与后续按不同组大小k 1每个 token 的k个专家进行的 router 加权ReduceSum融合。这将融合当前提案以显式Mul/ReduceSum表达的加权和代价是解耦 matmul 分组与归约分组需M*k M2*k2的总数约束并定义规范结果布局。该设计推迟到后续提案以免复杂化GroupedMatMul。异构专家通过权重张量序列支持专家K/N各异的一般 grouped-GEMM 情形。更多融合步骤若生态在稳定激活集上收敛可做更激进融合激活、上/下投影类似部分运行时 contrib op。group_offsets变体未来可为更自然映射到 cuBLAS grouped-GEMM API 的后端补充排序缓冲 / offsets 接口。从源码看函数算子的落地机制虽然GroupedMatMul尚未合入本仓库VERSION_NUMBER为 1.24.0onnx/defs/与onnx/backend/test/case/node/中尚无对应实现其选用的上下文相关函数机制在仓库中已有成熟支撑可作为理解其落地路径的依据机制入口schema.h 中OpSchema::SetContextDependentFunctionBodyBuilder(ContextDependentFunctionBodyBuilder, int opset_version)注册按 opset 版本的函数体构建器BuildContextDependentFunction在需要时实例化函数体HasContextDependentFunction()用于判断 schema 是否带有上下文相关函数。既有范例提案明确点名CenterCropPad为其风格参照其 schema 定义见 tensor/defs.cc——先声明输入/输出/类型约束与形状推断函数再以 lambda 形式的SetContextDependentFunctionBodyBuilder依据输入是否存在来构造不同函数体与GroupedMatMul按bias是否出现生成两种函数体完全同构。此外Gelu、Clip、LayerNormalization等算子math/defs.cc、nn/defs.cc也使用同一 API说明该机制是 ONNX 定义函数型算子的标准路径。参考求值器op_run.py 中函数在调用期实例化、且函数实现无论何时被调用都相同维度共享由调用方与函数体之间的绑定处理印证了提案所称无需单独参考实现文件、ReferenceEvaluator自动经函数体执行的可行性。测试落点新增 node 级测试将置于onnx/backend/test/case/node/与abs.py、add.py等同级由后端测试框架自动收集执行用于验证分解语义与预期输出完全一致。综上GroupedMatMul提案以函数分解即规范的方式为 ONNX 生态提供了一份紧凑、可融合、可移植的 MoE 分组矩阵乘标准表示其语义、形状推断、量化组合与测试用例已在本文完整呈现可据此在后续 opset 中按既有机制落地实现。赞分享人工智能深度学习机器学习【免费下载链接】onnxOpen standard for machine learning interoperability项目地址https://gitcode.com/gh_mirrors/onn/onnx点击查看免费下载相关推荐GroupedMatmul 算子深度指南MoE 分组矩阵乘的接口调用、量化模式与源码实现解析GroupedMatmul 算子深度指南MoE 分组矩阵乘的接口调用、量化模式与源码实现解析 导读 本文以 CANN ops transformer 仓库中算子库人工智能深度学习Ascend解读 ATB GroupedMatmulWithRouting 算子MoE 场景下带路由的分组矩阵乘实现解读 ATB GroupedMatmulWithRouting 算子MoE 场景下带路由的分组矩阵乘实现 导读 GroupedMatmulWithRoutin人工智能大模型算子库深度学习CANNAscendascend-transformer-boost 算子源码导读GroupedMatmulWithRouting带路由分组矩阵乘ascend transformer boost 算子源码导读GroupedMatmulWithRouting带路由分组矩阵乘 本文是 CANN asce人工智能大模型算子库深度学习CANNAscend上一篇终极指南Hunyuan3D-2合作伙伴生态系统与集成方案详解下一篇解决Anki跨端表格显示难题从数据渲染到响应式适配全方案创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考 ONNX GroupedMatMul 算子提案解读：MoE 前馈层分组矩阵乘的标准表示
人工智能深度学习机器学习【免费下载链接】onnxOpen standard for machine learning interoperability项目地址https://gitcode.com/gh_mirrors/onn/onnx点击查看免费下载导读本文深入解读 ONNX 官方 RFCdocs/proposals/0010-GroupedMatMulOpProposal.md状态为 under discussion提出的标准算子GroupedMatMul它把一批 token 向量乘以一组专家expert权重矩阵、每个 token 选择若干专家这一 Mixture-of-ExpertsMoE前馈层的核心计算封装为紧凑、可融合的 ONNX 函数算子。读完本文你将掌握GroupedMatMul的输入输出语义、函数分解式的规范定义、MoE 前馈层的标准建图方式、形状推断规则、量化QDQ组合方式、边界行为与测试用例以及它与 PyTorchgrouped_mm、cuBLAS grouped-GEMM、ONNX Runtimecom.microsoft.MoE等已有实现的对应关系。背景为什么需要 GroupedMatMulMoEMixture-of-Experts前馈层是 Mixtral、DeepSeek、Grok、Switch Transformer 等大语言模型的核心组件。其核心计算为给定一批M个 token 向量和num_groups个专家权重矩阵由 router 为每个 token 选择一个或多个专家矩阵进行乘法。这种被称为grouped matrix multiplication / grouped GEMM的模式理论上可以用现有标准 ONNX 算子表达但朴素的未融合实现非常低效、不实用朴素分解不切实际自然的做法是Gather权重 →Expandtoken → 批量MatMul这会物化出完整的[M×k, K, N]权重切片和[M×k, K]的 token 副本。对真实的 MoE 层而言这些中间张量动辄数 GB导致分解方案在实践中不可行。融合内核的优势融合的 grouped-GEMM 内核每个专家权重矩阵只需处理一次与选择它的 token 数量无关且每个 token 行在其k个专家间复用、无需拷贝。将GroupedMatMul加入 ONNX 标准ai.onnx域即可为 MoE 模型提供更紧凑、更高效的表示。提案同时指出ONNX Runtime 目前没有独立的 grouped-matmul contrib op而是提供更大的融合算子com.microsoft.MoE以及量化版本com.microsoft.QMoE来实现整个 MoE 前馈层本提案与 onnx/onnx#7902 相关文中仅作背景陈述外部链接不展开。算子语义Guide-level ExplanationGroupedMatMul的输入inputtoken 矩阵形状[M, K]weightsG个专家权重矩阵的堆叠形状[G, K, N]group_indices形状[M, k]为每个 token 选择其要乘的k个专家group索引可选bias逐组偏置形状[G, N]。输出形状为[M, k, N]每个 token 与它选中的k个专家分别相乘含可选偏置后的逐专家结果。任何对这些逐专家结果的加权组合例如 MoE 层中 router 的 top-k加权求和都应当由周围的Mul/ReduceSum标准算子表达下文给出示例。对于稠密单专家场景令k 1即可。关键设计算子内部复用每个 token 行跨其k个选中专家因此不需要Expandtoken 批次——这是与朴素分解的本质区别。典型用法MoE 前馈层一个标准的 top-k MoE FFN含两个投影可以映射为两个GroupedMatMul算子且无需Expandtoken 批次# Notation: B batch, S sequence length, H hidden dim, F FFN inner dim # E num_experts, k experts_per_token scores Softmax(MatMul(hidden, router_W)) # [B, S, E] values, indices TopK(scores, k) # [B, S, k] h Reshape(hidden, [B*S, H]) idx Reshape(indices, [B*S, k]) val Reshape(values, [B*S, k]) # --- Up projection: per-expert output --- # output shape: [B*S, k, F] h_up GroupedMatMul(h, expert_gate_W, idx) # expert_gate_bias (optional) h_up SiLU(h_up) # Reshape for down projection: treat each (token, expert-slot) pair as a row. h_flat Reshape(h_up, [B*S*k, F]) idx2 Reshape(idx, [B*S*k, 1]) # each flat row selects one expert # --- Down projection: per-expert output, then router-weighted sum over the k slots --- d GroupedMatMul(h_flat, expert_down_W, idx2) # [B*S*k, 1, H] d Reshape(d, [B*S, k, H]) # regroup the k slots out ReduceSum(d * Unsqueeze(val, -1), axis1) # [B*S, H] weighted sum over k out Reshape(out, [B, S, H])两个投影都只用GroupedMatMul完成分组矩阵乘法下投影中的 top-krouter 加权求和用MulReduceSum在重组后的k个槽位上显式表达。之所以不把该加权和融进算子是因为这里的k个专家拥有各自不同的逐槽位输入上投影激活后的输出而非共享同一输入行详见后文设计抉择。参考级规范函数分解单一规范性定义GroupedMatMul的语义被定义为一个 ONNX函数function其含义严格等于下述标准 ONNX 算子的组合。由于可选bias输入的存在会改变生成的图函数体是上下文相关的context-dependent——通过SetContextDependentFunctionBodyBuilder构建与CenterCropPad等既有算子的风格一致。以函数定义语义有一个重要附带收益ONNX 的参考求值器onnx.reference.ReferenceEvaluator在未注册专用 Python 内核时会自动通过函数体执行算子因此不需要单独的参考实现文件——一份分解同时充当规范、文档与参考实现。这与仓库中onnx/reference/op_run.py的函数执行机制一致自定义函数在执行期实例化、复用同一实现。下文符号约定M、K、N为形状维度G weights.shape[0]专家数k group_indices.shape[1]每 token 选择的专家数。函数分解Function Decompositionidx_flat Reshapeallowzero 1(group_indices, [M*k]) W_sel Gather(weights, idx_flat, axis0) # [M*k, K, N] — duplicates weights! X Reshapeallowzero 1(Expand(Unsqueeze(input, 1), [M, k, K]), [M*k, 1, K]) # [M*k, 1, K] — copies tokens! r Reshapeallowzero 1(MatMul(X, W_sel), [M, k, N]) # [M, k, N] # If bias is present, add the per-group bias to each selected expert result: bias_sel Reshapeallowzero 1(Gather(bias, idx_flat, axis0), [M, k, N]) # [M, k, N] r r bias_sel # (only when bias present) output r # [M, k, N]上下文相关的函数体只产生两种情形输入 3bias存在时包含 bias 加法行否则省略。关于效率的重要说明上述分解会物化完整的[M*k, K, N]权重切片和[M*k, K]的 token 副本真实 MoE 层中这些中间量达 GB 级——这正是GroupedMatMul需要作为融合算子存在的原因它定义的是计算结果是什么而运行时被期望融合执行而非照搬朴素分解。分解对结果是规范性的对策略不是。边界情况Edge Cases分解已对以下特殊情况给出良定义行为下表记录之情形行为k 0未选中任何专家输出为空张量形状[M, 0, N]真实模型不应出现。k 1每 token 一个专家输出形状[M, 1, N]。G 1所有索引为 0等价于MatMul(input, weights[0])含可选 bias。M 0零 token 输入输出形状[0, k, N]无需计算。索引越界非法输入实现必须报错。算子规范Operator Specification名称与域字段值NameGroupedMatMulDomainai.onnx标准域Opset version下一个可用 opset如 27Since version随该 opset 新增需要说明截至本仓库当前版本VERSION_NUMBER为 1.24.0该算子仍处于提案阶段尚未在onnx/defs/下的 schema 与onnx/backend/test/case/node/测试目录中落地实现——文中规范均以提案文档为准。输入索引名称类型必选形状说明0inputT必选[M, K]行主序 token 矩阵。M个 tokenK为收缩hidden维。1weightsT必选[G, K, N]G个专家权重矩阵的堆叠每个为K × N。所有专家共享相同的K与N。2group_indicestensor(int64)必选[M, k]每个 token 每个槽位slot的组专家索引。每个 token 选择k个专家取值必须在[0, G)内。稠密单专家场景用k1。3biasT可选[G, N]逐组偏置向量加到每个专家的结果上。注意点G weights.shape[0]组 / 专家数。若调用方有[B, M, K]形状的批处理输入应先用Reshape将批维度并入M多数后端中这通常是零拷贝、仅元数据的视图构造。weights与bias对所有 token 相同即它们是模型参数而非逐 token 的。输出索引名称类型形状说明0outputT[M, k, N]逐专家结果每个 token 与它选中的k个专家分别相乘的结果含可选 bias。类型约束约束类型Ttensor(float)、tensor(float16)、tensor(bfloat16)group_indices恒为tensor(int64)。量化QDQ组合方式GroupedMatMul可以直接与 ONNX 标准的量化/反量化QDQ表示组合。量化激活、权重与可选 bias 在算子之前反量化其输出可在之后量化DequantizeLinear(input_q) ─┐ ├─ GroupedMatMul ─ QuantizeLinear (optional) DequantizeLinear(weights_q) ─┘对于形状[G, K, N]的专家权重沿axis1的分块DequantizeLinear支持形状为[G, ceil(K/B), N]的 scale这允许每个专家、每个 K 块、每个输出通道各自独立的 scale。逐专家、逐输出通道量化是其特例[G, 1, N]单个 K 大小的块。GroupedMatMul的 schema不需要任何量化参数或量化张量类型。为高效执行运行时应当识别并融合周围的 QDQ 模式避免物化完整的反量化专家权重张量。若未来需要显式整数累加、requantization 或打包权重格式应另立QLinearGroupedMatMul风格的算子单独考虑。属性无。所有配置都通过输入表达遵循 ONNX 在取值可能动态时偏好输入而非属性的一般原则。形状推断规则Shape Inference设M input.shape[0]K input.shape[1]G weights.shape[0]N weights.shape[2]k group_indices.shape[1]校验检查违反即报错input.rank 2weights.rank 3group_indices.rank 2且group_indices.shape[0] Mweights.shape[1] K收缩维一致若存在biasbias.shape [G, N]输出形状output.shape [M, k, N]测试用例提案计划将用例放入onnx/backend/test/case/node/groupedmatmul.py该目录下已有各算子的 node 级后端测试如onnx/backend/test/case/node/abs.py、add.py等新增测试将遵循同样的组织方式。测试 1 —— 稠密k1无 bias# 4 tokens, K3, G2 groups, N2, k1 input [[1, 0, -1], [0, 1, 2], [1, 1, 0], [0, 0, 1]] # shape [4, 3] weights [[[1, 0], [0, 1], [-1, 0]], [[0, 1], [1, 0], [ 0, 1]]] # shape [2, 3, 2] group_indices [[0], [1], [0], [1]] # shape [4, 1] # Expected output shape [4, 1, 2]: # token 0 - group 0: [1,0,-1] [[1,0],[0,1],[-1,0]] [101, 000] [2, 0] # token 1 - group 1: [0,1, 2] [[0,1],[1,0],[ 0,1]] [010, 002] [1, 2] # token 2 - group 0: [1,1, 0] [[1,0],[0,1],[-1,0]] [100, 010] [1, 1] # token 3 - group 1: [0,0, 1] [[0,1],[1,0],[ 0,1]] [000, 001] [0, 1] output [[[2, 0]], [[1, 2]], [[1, 1]], [[0, 1]]] # shape [4, 1, 2]测试 2 —— Top-kk2含 biasM, k, K, G, N 2, 2, 2, 3, 2 input [[1.0, 0.0], [0.0, 1.0]] # [2, 2] weights [[[1,0],[0,1]], [[0,1],[1,0]], [[1,1],[0,0]]] # [3,2,2] group_indices [[0, 1], [2, 0]] # [2, 2] bias [[0.1, 0.2], [0.3, 0.0], [0.5, 0.5]] # [3, 2] # token 0, slot 0 - g0: [1,0] [[1,0],[0,1]] [0.1,0.2] [1.1, 0.2] # token 0, slot 1 - g1: [1,0] [[0,1],[1,0]] [0.3,0.0] [0.3, 1.0] # token 1, slot 0 - g2: [0,1] [[1,1],[0,0]] [0.5,0.5] [0.5, 0.5] # token 1, slot 1 - g0: [0,1] [[1,0],[0,1]] [0.1,0.2] [0.1, 1.2] output [[[1.1, 0.2], [0.3, 1.0]], [[0.5, 0.5], [0.1, 1.2]]] # [2, 2, 2]测试 3 —— 空组某个专家未被使用# Group 1 receives no tokens. M, k, K, G, N 4, 1, 2, 3, 2 group_indices [[0], [0], [2], [2]] # group 1 unused # weights[1] is never accessed; output is well-defined.测试 4 —— 单组退化到 MatMulG 1 group_indices [[0], [0], [0]] # all tokens - group 0 # output MatMul(input, weights[0]) (reshaped from [M,1,N] to [M,1,N])设计抉择Rationale and Alternatives为什么这样设计将GroupedMatMul表达为上下文相关的 ONNX 函数可以获得单一规范性定义同时充当参考实现保持算子与周围图router、激活、reshape、加权求和的可组合性让运行时自行融合计算而不规定具体融合策略。将算子限定在分组矩阵乘法本身含可选的逐组bias使其小而通用MoE 层使用的 router 加权求和交给标准Mul/ReduceSum必要时可由后续独立算子解决。不做的后果没有GroupedMatMulMoE 模型要么用朴素的Gather/Expand/MatMul分解导出会物化 GB 级中间量、不实用要么依赖 ONNX Runtimecom.microsoft.MoE这类厂商私有 contrib op——无法在 ONNX 生态内移植。提案还评估了以下备选方案k1 的特殊形状当前方案要求group_indices是[M, k]的二维张量。备选方案是允许k1特例使用[M]的一维张量。显式批维本算子将M个 token 展平为单一维度实际使用中批处理时 token 可能是多维的如[Batch, Sequence]。可考虑支持额外批维以避免多余 Reshape甚至让M支持任意维数——但后者复杂度收益不成比例未采纳。更大粒度的融合算子实现中可能做更激进的融合如把下投影、激活、上投影整个融合ONNX Runtime 的 MoE contrib op 正是如此。其缺点是激活函数因模型而异、且新模型不断引入新激活导致 contrib op 需要持续更新以支持新激活目前没有好的解决方案。决策本算子不融合激活、不融合 combine。激活保持为独立 ONNX 算子以保证图跨众多 MoE 路由变体可组合。router 加权求和combine同样用标准Mul/ReduceSum表达而非折入GroupedMatMul标准 MoE 下投影中k个专家拥有各自不同的逐槽位输入加权求和的归约分组与 matmul 自身的k分组不同折入会迫使两种分组解耦并固定规范结果布局——规范与形状推断复杂度都会增加而这类融合更适合由未来专门算子干净地捕获。group_indicesvsgroup_offsets备选方案是使用已排序的 token 缓冲与整数 offsets更贴近 cuBLAS grouped-GEMM API。决策采用group_indices无序索引。索引天然与TopK/Gather组合且不要求调用方预排序 token 批次运行时内部自行排序。堆叠三维权重[G, K, N]vs 变长矩阵列表备选方案是每个专家可能具有不同K/N的权重张量序列异构专家一般情形。决策堆叠[G, K, N]。所有专家共享K、N是已部署 MoE 模型中的压倒性常见情形单一权重张量更简单。实现方仍可差异化处理单一张量内的不同专家切片例如应对整个张量过大无法同时装入内存的情况。已有实现参考Prior Art框架 / 库APIPyTorchtorch.nn.functional.grouped_mmPyTorch ≥ 2.5PyTorchtorch._grouped_mm/torch.ops.aten.mm_group内部JAXjax.lax.dot_general分组批处理cuBLAScublasGemmBatchedEx/cublasGemmGroupedBatchedExCUTLASSGroupedGemm内核OpenVINOGroupConvolution卷积类比ONNX Runtimecom.microsoft.MoE/com.microsoft.QMoEcontrib_ops——更大的融合 MoE 层而非独立 grouped matmulPyTorch 2.5 新增的torch.nn.functional.grouped_mm与本提案语义直接对应# PyTorch grouped_mm — same semantics out torch.nn.functional.grouped_mm(input, weight, offsNone) # offs are contiguous group offsets; our design uses indices instead # (see group_indices vs. group_offsets under Rationale and alternatives)未决问题Unresolved Questionsk 1时是否允许形状为[M]的一维group_indices还是始终要求二维[M, k]见上文k1 的特殊形状。是否引入显式批维例如[B, S, K]的多维 token 输入以省去周围Reshape还是保持展平的[M, K]形式见上文显式批维。未来展望Future Possibilities带归约的融合下投影可设计独立算子表示 MoE 下投影——自身k为 1每个展平的(token, slot)行选单个专家的GroupedMatMul与后续按不同组大小k 1每个 token 的k个专家进行的 router 加权ReduceSum融合。这将融合当前提案以显式Mul/ReduceSum表达的加权和代价是解耦 matmul 分组与归约分组需M*k M2*k2的总数约束并定义规范结果布局。该设计推迟到后续提案以免复杂化GroupedMatMul。异构专家通过权重张量序列支持专家K/N各异的一般 grouped-GEMM 情形。更多融合步骤若生态在稳定激活集上收敛可做更激进融合激活、上/下投影类似部分运行时 contrib op。group_offsets变体未来可为更自然映射到 cuBLAS grouped-GEMM API 的后端补充排序缓冲 / offsets 接口。从源码看函数算子的落地机制虽然GroupedMatMul尚未合入本仓库VERSION_NUMBER为 1.24.0onnx/defs/与onnx/backend/test/case/node/中尚无对应实现其选用的上下文相关函数机制在仓库中已有成熟支撑可作为理解其落地路径的依据机制入口schema.h 中OpSchema::SetContextDependentFunctionBodyBuilder(ContextDependentFunctionBodyBuilder, int opset_version)注册按 opset 版本的函数体构建器BuildContextDependentFunction在需要时实例化函数体HasContextDependentFunction()用于判断 schema 是否带有上下文相关函数。既有范例提案明确点名CenterCropPad为其风格参照其 schema 定义见 tensor/defs.cc——先声明输入/输出/类型约束与形状推断函数再以 lambda 形式的SetContextDependentFunctionBodyBuilder依据输入是否存在来构造不同函数体与GroupedMatMul按bias是否出现生成两种函数体完全同构。此外Gelu、Clip、LayerNormalization等算子math/defs.cc、nn/defs.cc也使用同一 API说明该机制是 ONNX 定义函数型算子的标准路径。参考求值器op_run.py 中函数在调用期实例化、且函数实现无论何时被调用都相同维度共享由调用方与函数体之间的绑定处理印证了提案所称无需单独参考实现文件、ReferenceEvaluator自动经函数体执行的可行性。测试落点新增 node 级测试将置于onnx/backend/test/case/node/与abs.py、add.py等同级由后端测试框架自动收集执行用于验证分解语义与预期输出完全一致。综上GroupedMatMul提案以函数分解即规范的方式为 ONNX 生态提供了一份紧凑、可融合、可移植的 MoE 分组矩阵乘标准表示其语义、形状推断、量化组合与测试用例已在本文完整呈现可据此在后续 opset 中按既有机制落地实现。赞分享人工智能深度学习机器学习【免费下载链接】onnxOpen standard for machine learning interoperability项目地址https://gitcode.com/gh_mirrors/onn/onnx点击查看免费下载相关推荐GroupedMatmul 算子深度指南MoE 分组矩阵乘的接口调用、量化模式与源码实现解析GroupedMatmul 算子深度指南MoE 分组矩阵乘的接口调用、量化模式与源码实现解析 导读 本文以 CANN ops transformer 仓库中算子库人工智能深度学习Ascend解读 ATB GroupedMatmulWithRouting 算子MoE 场景下带路由的分组矩阵乘实现解读 ATB GroupedMatmulWithRouting 算子MoE 场景下带路由的分组矩阵乘实现 导读 GroupedMatmulWithRoutin人工智能大模型算子库深度学习CANNAscendascend-transformer-boost 算子源码导读GroupedMatmulWithRouting带路由分组矩阵乘ascend transformer boost 算子源码导读GroupedMatmulWithRouting带路由分组矩阵乘 本文是 CANN asce人工智能大模型算子库深度学习CANNAscend上一篇终极指南Hunyuan3D-2合作伙伴生态系统与集成方案详解下一篇解决Anki跨端表格显示难题从数据渲染到响应式适配全方案创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考