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] when transB is zero (the default), 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]]

When transB is non-zero, weights has shape [G, N, K] and each selected weight matrix is transposed before multiplication:

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

M is the number of input rows, G is the number of groups, K is the contraction dimension, and N is the output dimension. Only the last two axes of weights are logically transposed; the group axis and the layouts of input, output, and bias are unchanged.

If bias is present, it must have shape [G, N], and bias[group_indices[m, i]] is added to each corresponding result. If M or S is zero, the output is an empty tensor of shape [M, S, N]. 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.

Quantized weights can be dequantized with DequantizeLinear before this operator. The dequantized weights must have the shape specified by transB.

Attributes

• transB - INT (default is 0):

  Whether to transpose the last two axes of weights before multiplication. If zero, weights has shape [G, K, N]; if non-zero, weights has shape [G, N, K].

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] if transB is zero, or [G, N, K] if transB is non-zero.

• 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

_transposed_weights_nonzero

import numpy as np
import onnx

input = np.array([[1, 2, -1], [3, 0, 4]], dtype=np.float32)
weights = np.arange(18, dtype=np.float32).reshape(3, 2, 3)
group_indices = np.array([[2, 0], [1, 2]], dtype=np.int64)
output = grouped_matmul(input, weights, group_indices, transB=1)
for transB in (-1, 2):
    node = onnx.helper.make_node(
        "GroupedMatMul",
        ["input", "weights", "group_indices"],
        ["output"],
        transB=transB,
    )
    expect(
        node,
        inputs=[input, weights, group_indices],
        outputs=[output],
        name=f"test_groupedmatmul_transB_nonzero_{transB}",
        opset_imports=[onnx.helper.make_opsetid("", 29)],
    )

_transposed_weights

import numpy as np
import onnx

input = np.array([[1, 2, -1], [3, 0, 4]], dtype=np.float32)
weights = np.arange(18, dtype=np.float32).reshape(3, 2, 3)
group_indices = np.array([[2, 0, 2], [1, 2, 0]], dtype=np.int64)
bias = np.array([[0.5, 1], [1.5, -2], [3, 0.25]], dtype=np.float32)
for bias_input in (None, "", "bias"):
    input_names = ["input", "weights", "group_indices"]
    values = [input, weights, group_indices]
    if bias_input is not None:
        input_names.append(bias_input)
    if bias_input:
        values.append(bias)
    node = onnx.helper.make_node(
        "GroupedMatMul", input_names, ["output"], transB=1
    )
    output = grouped_matmul(
        input, weights, group_indices, bias if bias_input else None, transB=1
    )
    suffix = {None: "", "": "_explicit_empty_bias", "bias": "_with_bias"}[
        bias_input
    ]
    expect(
        node,
        inputs=values,
        outputs=[output],
        name=f"test_groupedmatmul_transB{suffix}",
        opset_imports=[onnx.helper.make_opsetid("", 29)],
    )

_transposed_weights_empty

import numpy as np
import onnx

weights = np.arange(24, dtype=np.float32).reshape(2, 4, 3)
for m, s in ((0, 2), (2, 0)):
    input = np.empty((m, 3), dtype=np.float32)
    group_indices = np.empty((m, s), dtype=np.int32)
    node = onnx.helper.make_node(
        "GroupedMatMul",
        ["input", "weights", "group_indices"],
        ["output"],
        transB=1,
    )
    output = grouped_matmul(input, weights, group_indices, transB=1)
    expect(
        node,
        inputs=[input, weights, group_indices],
        outputs=[output],
        name=f"test_groupedmatmul_transB_empty_{m}_{s}",
        opset_imports=[onnx.helper.make_opsetid("", 29)],
    )

_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"],
    transB=0,
)
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)],
)