[Feature Request] Add GroupedMatMul op for efficient MoE support

[Feature Request] Add GroupedMatMul op for efficient MoE support

Motivation

Mixture-of-Experts (MoE) models are becoming standard (Gemma 4, DeepSeek-V2/V3, Mixtral, Qwen-MoE, DBRX). Currently, MoE layers can be represented using standard ONNX ops (Gather + MatMul + TopK), but this decomposition misses significant optimization opportunities.

The key value of this op is not the GEMM speedup itself (measured at ~1.2-1.8x for grouped vs batched GEMM), but the semantic signal it provides to runtimes. Without a dedicated op, runtimes see a soup of Gather + MatMul + Reshape ops and cannot reliably identify "this is an MoE layer — please apply end-to-end fusion (permute + grouped FFN + unpermute + weighted sum in one kernel)." With GroupedMatMul, the intent is explicit, enabling runtimes to trigger full MoE kernel fusion where the real gains are.

PyTorch already has torch.nn.functional.grouped_mm (docs) for this exact use case. ONNX needs a corresponding op to enable clean export without performance regression.

Why Not a Full MoE Op?

We considered a dedicated MoELayer op (routing + FFN + weighted sum in one op) for maximum fusion potential. However, modern MoE architectures are too diverse for a single op:

• Gemma 4: 128 experts + 1 shared expert + top-8 + GeGLU + Per-Layer Embeddings
• DeepSeek-V2/V3: 160 experts + shared experts + top-6
• Mixtral: 8 experts + top-2 + SwiGLU
• Expert Choice routing: experts select tokens (inverted)
• Soft MoE: all experts participate with continuous weights

A dedicated MoE op would need to constantly grow new attributes/inputs to keep up — it becomes "GemmaLayer" or "DeepSeekLayer" rather than a general building block. Even ORT's existing com.microsoft.MoE contrib op (which has SwiGLU fusion, fc3, sparse mixer) cannot fully represent Gemma 4.

GroupedMatMul is the correct level of abstraction — general enough to survive architecture changes, specific enough to enable optimizations. All MoE variants share one thing: grouped matrix multiplication by expert index.

Formal Spec: GroupedMatMul

Op Name

GroupedMatMul

Summary

Grouped matrix multiplication where each token uses one or more weight matrices selected by group index. Core computation for Mixture-of-Experts (MoE) layers.

Inputs

Outputs

Semantics

def GroupedMatMul(input, weights, group_indices, bias=None):
    B, M, K = input.shape
    num_groups, _, N = weights.shape
    
    if group_indices.ndim == 2:
        # Single group per token: [B, M] → [B, M, N]
        output = zeros([B, M, N])
        for b in range(B):
            for i in range(M):
                g = group_indices[b, i]
                output[b, i] = input[b, i] @ weights[g]
                if bias is not None:
                    output[b, i] += bias[g]
    else:
        # Multiple groups per token: [B, M, k] → [B, M, k, N]
        k = group_indices.shape[2]
        output = zeros([B, M, k, N])
        for b in range(B):
            for i in range(M):
                for j in range(k):
                    g = group_indices[b, i, j]
                    output[b, i, j] = input[b, i] @ weights[g]
                    if bias is not None:
                        output[b, i, j] += bias[g]
    
    return output

Type Constraints

• T: tensor(float16), tensor(bfloat16), tensor(float32), tensor(float64)
• FP8/INT8: compose with existing DequantizeLinear/QuantizeLinear ops. Not baked into this op.

Edge Cases

• Empty groups: Valid. If no token maps to group g, weights[g] is unused.
• Single sample: Use B=1.
• All tokens same group: Degenerates to standard batched MatMul.
• k=1 with 3D indices: Equivalent to 2D indices (output has extra dim of size 1).

MoE Usage Patterns

Top-k Routing (Direct — No Reshape Needed for First Layer)

With 3D group_indices, the first GroupedMatMul call needs no Reshape/Expand:

# Router selects top-k experts
scores = Softmax(MatMul(hidden, router_W))            # [B, M, E]
values, indices = TopK(scores, k=2)                    # [B, M, 2]

# First FFN layer: 3D indices, no reshape needed
h = GroupedMatMul(hidden, expert_W1, indices)           # [B, M, 2, intermediate]
h = SiLU(h)

# Second FFN layer: flatten for chaining (one reshape)
h_flat = Reshape(h, [B, M*2, intermediate])            # [B, M*k, intermediate]
indices_flat = Reshape(indices, [B, M*2])              # [B, M*k]
out = GroupedMatMul(h_flat, expert_W2, indices_flat)    # [B, M*k, hidden]

# Reshape back and weighted sum
out = Reshape(out, [B, M, 2, hidden])                  # [B, M, k, hidden]
output = ReduceSum(out * Unsqueeze(values, -1), axis=2) # [B, M, hidden]

Gemma 4 (128 experts, top-8, shared expert, GeGLU)

# Shared expert (always active) — standard MatMul
shared_out = MatMul(hidden, shared_W1)
shared_out = GELU(shared_out)
shared_out = MatMul(shared_out, shared_W2)

# Routed experts — GroupedMatMul with 3D indices
scores = Softmax(MatMul(hidden, router_W))              # [B, M, 128]
values, indices = TopK(scores, k=8)                     # [B, M, 8]

# GeGLU: two GroupedMatMul calls with 3D indices (no reshape for up/gate)
gate = GroupedMatMul(hidden, expert_W1, indices)         # [B, M, 8, intermediate]
gate = GELU(gate)
up = GroupedMatMul(hidden, expert_W3, indices)           # [B, M, 8, intermediate]
h = gate * up                                            # element-wise

# Down projection (flatten for chaining — one reshape)
h_flat = Reshape(h, [B, M*8, intermediate])
indices_flat = Reshape(indices, [B, M*8])
out = GroupedMatMul(h_flat, expert_W2, indices_flat)     # [B, M*8, hidden]

# Reshape + weighted sum + add shared
out = Reshape(out, [B, M, 8, hidden])
routed_out = ReduceSum(out * Unsqueeze(values, -1), axis=2)  # [B, M, hidden]
output = shared_out + routed_out

Expert Choice Routing

# Each expert selects its top-N tokens (inverted routing)
selected_input = GatherElements(hidden, expert_selections)
flat_indices = ...  # assign each selected token to its expert (2D)
output = GroupedMatMul(selected_input, weights, flat_indices)
# Scatter results back to original positions

Design Decisions

Performance Impact

Without fusion, MoE layers decomposed into standard ops suffer from:

1. Multiple kernel launches (Router → TopK → Gather → MatMul → Scatter)
2. Redundant memory reads/writes for intermediate results
3. Inability to do token permutation + grouped GEMM co-optimization

Performance references:

• NVIDIA cuBLAS grouped GEMM API achieves ~1.2x over naive batched GEMM loops (source)
• Triton grouped GEMM achieves ~1.5-1.8x over cuBLAS for small N (source)
• Full MoE layer fusion (routing + permute + GEMM + unpermute) can yield larger gains by eliminating kernel launch overhead and intermediate memory allocation
• The semantic signal from GroupedMatMul enables runtimes to identify MoE patterns and apply end-to-end fusion

Optimization Notes for Runtime Implementers

• Reorder tokens by group → perform one batched GEMM per group
• Use CUTLASS grouped GEMM or similar fused kernels
• Support variable group sizes (different groups may have different token counts)
• Consider overlapping communication with computation for expert parallelism
• When consecutive GroupedMatMul nodes with same group_indices are detected, fuse the entire FFN block
• 3D indices hint at top-k pattern — runtime can optimize the implicit broadcast internally

Related Work

• PyTorch: torch.nn.functional.grouped_mm (PyTorch 2.11+)
• PyTorch blog: Triton Grouped GEMM kernel for MoE
• ONNX Runtime: com.microsoft.MoE contrib op (non-standard, ORT-only)
• vLLM: fused MoE kernels
• CUTLASS: grouped GEMM API

cc @justinchuby