GroupedMatmulInplaceAddOperation

Applicable Products

Hardware Model	Supported
Atlas 350 accelerator cards	x
Atlas A3 inference products / Atlas A3 training products	√
Atlas A2 training products / Atlas A2 inference products	√
Atlas training products	x
Atlas inference products	x
Atlas 200I/500 A2 inference products	x

Description

Fuses Groupedmatmul and add operators to implement the group matrix multiplication and InplaceAdd functions.

Computational Graph

Performs group matrix multiplication and accumulates in-place at the output address. ki in dimensions of each group matrix multiplications is different, and x_i/weight_i may be concatenated on k_i. The mathematical expression of the operator is $\text{[math]}$ .

g indicates the number of groups, and m_i/k_i/n_i indicates the corresponding shape.

Figure 1 Computation diagram

Definition

      
           struct GroupedMatmulInplaceAddParam {
    bool transposeA = false;
    bool transposeB = false;
    uint8_t rsv[22] = {0};
};

Parameters

Member	Type	Default Value	Value Range	Description
transposeA	bool	False	false/true	Whether matrix A is transposed
transposeB	bool	False	false	Whether matrix B is transposed. Currently, only False is supported.
rsv[22]	uint8_t	{0}	-	Reserved

Input

Parameter	Dimension	Data Type	Format	Mandatory	Description
x	[m, k]	float16/bf16	ND	Yes	Matrix A for matrix multiplication. The value range of m is (0, 65535].
weight	[k, n]	float16/bf16	ND	Yes	Weight of matrix B for matrix multiplication. The value range of n is (0, 65535]. The data type is the same as that of x.
groupList	[numGroup]	int64	ND	Yes	1D tensor, indicating the number of matrix multiplication groups. The maximum length is 128.
out	[numGroup * m, n] or [m, numGroup * n]	float	ND	Yes	Input and output tensors, supporting in-place write. The maximum value of numGroup is 128.

Output

Parameter	Dimension	Data Type	Format	Mandatory	Description
out	[numGroup * m, n] or [m, numGroup * n]	float	ND	Yes	Output.

Constraints

Currently, transposeB can only be set to False. The value of transposeA and the dimensions of matrix A and matrix B must meet the dimension constraints of matrix multiplication.
The output tensor is the same as the last input tensor.
groupList is a 1D tensor whose length is less than or equal to 128, indicating the index of A/B on the k axis. The next element in groupList must be greater than or equal to the previous element, and the size of the last element must be equal to k.
For example, if the dimension of matrix A is [5, 10], the dimension of matrix B is [10, 4], the dimension of groupList is a 1D tensor, numGroup is 3, and the values are [3, 6, 10], it indicates that k is divided into three groups: [0, 3), [3, 6), and [6, 10), and matrix multiplication is performed by group.

Parent topic: atb/infer_op_params.h