GroupedMatMul¶

GroupedMatMul - 29¶

Version¶

name: GroupedMatMul (GitHub)
domain: main
since_version: 29
function: True
support_level: SupportType.COMMON
shape inference: True

This version of the operator has been available since version 29.

Summary¶

GroupedMatMul multiplies each row of a token matrix by one or more selected group (expert) weight matrices. Given input of shape [M, K], weights of shape [G, K, N], and group_indices of shape [M, S], where S is the number of selections per token, the output has shape [M, S, N] and is defined by:

output[m, i] = input[m] @ weights[group_indices[m, i]]

If bias is present, bias[group_indices[m, i]] is added to each corresponding result. Every value in group_indices must be in the range [0, G). Values outside this range, including negative indices, are invalid input; the behavior is undefined for such inputs. In particular, negative indices do not select groups by counting backwards from G.

This operator is represented as a context-dependent function, which also provides its reference implementation. Its decomposition uses Gather, Expand, and MatMul to define the result, while runtimes are expected to use a fused grouped-matrix-multiplication implementation to avoid materializing the expanded intermediate tensors, including selected weights of shape [M * S, K, N].

Weight-only quantized experts are intended to be expressed by applying DequantizeLinear to quantized weights (with the appropriate scales and zero points) and passing the result to GroupedMatMul. Runtimes may fuse this composition to avoid materializing the dequantized [G, K, N] weights.

Inputs¶

Between 3 and 4 inputs.

input (heterogeneous) - T:

Token matrix with shape [M, K].
weights (heterogeneous) - T:

Stack of G group weight matrices with shape [G, K, N].
group_indices (heterogeneous) - Tind:

Group index for each token and selection, with shape [M, S]. Values outside [0, G) are invalid input.
bias (optional, heterogeneous) - T:

Optional per-group bias with shape [G, N].

Outputs¶

output (heterogeneous) - T:

Per-group matrix multiplication results with shape [M, S, N].

Type Constraints¶

T in ( tensor(bfloat16), tensor(double), tensor(float), tensor(float16) ):

Constrain input, weights, bias, and output to floating-point tensors.
Tind in ( tensor(int32), tensor(int64) ):

Constrain group indices to int32 or int64 tensors.

Examples¶

_additional_types¶

import numpy as np
import onnx

for dtype in (np.float16, np.float64):
    for index_dtype in (np.int32, np.int64):
        node = onnx.helper.make_node(
            "GroupedMatMul",
            inputs=["input", "weights", "group_indices", "bias"],
            outputs=["output"],
        )
        input = np.array([[1, 2], [3, 4]], dtype=dtype)
        weights = np.array([[[1, 0], [0, 1]], [[0, 1], [1, 0]]], dtype=dtype)
        group_indices = np.array([[0, 1], [1, 0]], dtype=index_dtype)
        bias = np.array([[0.5, 1], [1.5, 2]], dtype=dtype)
        output = grouped_matmul(input, weights, group_indices, bias)

        expect(
            node,
            inputs=[input, weights, group_indices, bias],
            outputs=[output],
            name=f"test_groupedmatmul_{np.dtype(dtype).name}_{np.dtype(index_dtype).name}",
            opset_imports=[onnx.helper.make_opsetid("", 29)],
        )

default¶

import numpy as np
import onnx

node = onnx.helper.make_node(
    "GroupedMatMul",
    inputs=["input", "weights", "group_indices"],
    outputs=["output"],
)
input = np.array(
    [[1, 0, -1], [0, 1, 2], [1, 1, 0], [0, 0, 1]], dtype=np.float32
)
weights = np.array(
    [[[1, 0], [0, 1], [-1, 0]], [[0, 1], [1, 0], [0, 1]]],
    dtype=np.float32,
)
group_indices = np.array([[0], [1], [0], [1]], dtype=np.int64)
output = grouped_matmul(input, weights, group_indices)

expect(
    node,
    inputs=[input, weights, group_indices],
    outputs=[output],
    name="test_groupedmatmul",
    opset_imports=[onnx.helper.make_opsetid("", 29)],
)

_without_bias_explicit_empty_input¶

import numpy as np
import onnx

node = onnx.helper.make_node(
    "GroupedMatMul",
    inputs=["input", "weights", "group_indices", ""],
    outputs=["output"],
)
input = np.array(
    [[1, 0, -1], [0, 1, 2], [1, 1, 0], [0, 0, 1]], dtype=np.float32
)
weights = np.array(
    [[[1, 0], [0, 1], [-1, 0]], [[0, 1], [1, 0], [0, 1]]],
    dtype=np.float32,
)
group_indices = np.array([[0], [1], [0], [1]], dtype=np.int64)
output = grouped_matmul(input, weights, group_indices)

expect(
    node,
    inputs=[input, weights, group_indices],
    outputs=[output],
    name="test_groupedmatmul_without_bias_explicit_empty_input",
    opset_imports=[onnx.helper.make_opsetid("", 29)],
)

_with_bias¶

import numpy as np
import onnx

node = onnx.helper.make_node(
    "GroupedMatMul",
    inputs=["input", "weights", "group_indices", "bias"],
    outputs=["output"],
)
input = np.array([[1, 0], [0, 1]], dtype=np.float32)
weights = np.array(
    [[[1, 0], [0, 1]], [[0, 1], [1, 0]], [[1, 1], [0, 0]]],
    dtype=np.float32,
)
group_indices = np.array([[0, 1], [2, 0]], dtype=np.int64)
bias = np.array([[0.1, 0.2], [0.3, 0.0], [0.5, 0.5]], dtype=np.float32)
output = grouped_matmul(input, weights, group_indices, bias)

expect(
    node,
    inputs=[input, weights, group_indices, bias],
    outputs=[output],
    name="test_groupedmatmul_with_bias",
    opset_imports=[onnx.helper.make_opsetid("", 29)],
)

_with_unused_group¶

import numpy as np
import onnx

node = onnx.helper.make_node(
    "GroupedMatMul",
    inputs=["input", "weights", "group_indices"],
    outputs=["output"],
)
input = np.array([[1, 2], [3, 4], [5, 6], [7, 8]], dtype=np.float32)
weights = np.arange(12, dtype=np.float32).reshape(3, 2, 2)
group_indices = np.array([[0], [0], [2], [2]], dtype=np.int64)
output = grouped_matmul(input, weights, group_indices)

expect(
    node,
    inputs=[input, weights, group_indices],
    outputs=[output],
    name="test_groupedmatmul_with_unused_group",
    opset_imports=[onnx.helper.make_opsetid("", 29)],
)

_single_group¶

import numpy as np
import onnx

node = onnx.helper.make_node(
    "GroupedMatMul",
    inputs=["input", "weights", "group_indices"],
    outputs=["output"],
)
input = np.array([[1, 2], [3, 4], [5, 6]], dtype=np.float32)
weights = np.array([[[1, 2, 3], [4, 5, 6]]], dtype=np.float32)
group_indices = np.zeros((3, 1), dtype=np.int64)
output = grouped_matmul(input, weights, group_indices)

expect(
    node,
    inputs=[input, weights, group_indices],
    outputs=[output],
    name="test_groupedmatmul_single_group",
    opset_imports=[onnx.helper.make_opsetid("", 29)],
)

_zero_selections¶

import numpy as np
import onnx

node = onnx.helper.make_node(
    "GroupedMatMul",
    inputs=["input", "weights", "group_indices"],
    outputs=["output"],
)
input = np.arange(6, dtype=np.float32).reshape(2, 3)
weights = np.arange(24, dtype=np.float32).reshape(2, 3, 4)
group_indices = np.empty((2, 0), dtype=np.int64)
output = grouped_matmul(input, weights, group_indices)

expect(
    node,
    inputs=[input, weights, group_indices],
    outputs=[output],
    name="test_groupedmatmul_zero_selections",
    opset_imports=[onnx.helper.make_opsetid("", 29)],
)

_zero_tokens¶

import numpy as np
import onnx

node = onnx.helper.make_node(
    "GroupedMatMul",
    inputs=["input", "weights", "group_indices"],
    outputs=["output"],
)
input = np.empty((0, 3), dtype=np.float32)
weights = np.arange(24, dtype=np.float32).reshape(2, 3, 4)
group_indices = np.empty((0, 2), dtype=np.int64)
output = grouped_matmul(input, weights, group_indices)

expect(
    node,
    inputs=[input, weights, group_indices],
    outputs=[output],
    name="test_groupedmatmul_zero_tokens",
    opset_imports=[onnx.helper.make_opsetid("", 29)],
)