Einstein Multiplication Enabling for Matmul

Applicable Products

Hardware Model	Supported or Not
Atlas 350 accelerator card	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

Enables Einstein multiplication for matrix multiplication (matmul).

Formula

Two input tensors and are used for matrix multiplication, and the output tensor is .

$\text{[math]}$

Parameter Configuration

Member	Value Range
transposeA	false
transposeB	false/true
hasBias	false
outDataType	ACL_DT_UNDEFINED
enAccum	false
matmulType	MATMUL_EIN_SUM
quantMode	QUANT_UNDEFINED

Input

Parameter	Dimension	Data Type	Format	Description
x	[m, batch, k]	float16/bf16	ND	Matrix A for matrix multiplication.
weight	ND: [batch, k, n] NZ: [batch, n/16, k, 16]	float16/bf16	ND/NZ	Weight of matrix B for matrix multiplication. When there are four dimensions, the values of k and n are integer multiples of 16.

Parameter

Dimension

Data Type

Format

Description

[m, batch, k]

float16/bf16

Matrix A for matrix multiplication.

weight

ND: [batch, k, n]

NZ: [batch, n/16, k, 16]

float16/bf16

ND/NZ

Weight of matrix B for matrix multiplication.

When there are four dimensions, the values of k and n are integer multiples of 16.

Output

Parameter	Dimension	Data Type	Format	Description
output	[m, batch, n]	float16/bf16	ND	Matrix multiplication result.

OP Usage and Typical Scenarios

For details about how to use OP, see the usage process in Operator Usage Guide (ATB C++ APIs). For details about how to construct the Operation parameter as instructed in Single-operator, see the parameter construction in the following scenarios.

// Parameter construction
atb::infer::LinearParam param;
param.transposeA = false;
param.transposeB = false;
param.hasBias = false;
param.outDataType = ACL_DT_UNDEFINED;
param.enAccum = false;
param.matmulType = MATMUL_EIN_SUM;
param.quantMode = QUANT_UNDEFINED;

# Example
>>> x
tensor([[[1, 2]],
        [[3, 4]]])
>>> weight
tensor([[[1, 2, 3],
         [4, 5, 6]]])
>>> output
tensor([[[9, 12, 15],
       [19, 26, 33]]])
# 9 = 1 * 1 + 2 * 4
# 12 = 1 * 2 + 2 * 5
# 15 = 1 * 3 + 2 * 6 
# 19 = 3 * 1 + 4 * 4
# 26 = 3 * 2 + 4 * 5 
# 33 = 3 * 3 + 4 * 6

Parent topic: Functions