Matmul Supporting the Pertoken Quantization Mode

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

The linear operation supports the pertoken quantization mode.

Formula

$\text{[math]}$

Parameter Configuration

Member	Value Range
transposeA	false/true
transposeB	false/true
hasBias	false
outDataType	ACL_FLOAT16/ACL_BF16
enAccum	false
matmulType	MATMUL_UNDEFINED
quantMode	PER_TOKEN

Input

Parameter	Dimension	Data Type	Format	Description
x	[m,k]/[batch,m,k]	int8	ND	Matrix A for matrix multiplication.
weight	[k,n]/[batch,k,n]	int8	ND	Weight of matrix B for matrix multiplication.
deqScale	[n]	float	ND	Dequantization step The actually occupied memory must be 32-byte aligned.
perTokenScale	[m]	float	ND	perToken dequantization step.

Output

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

Description

The following table lists the combinations of input and output attributes. The combinations that are not listed in the table are not supported.

Figure 1 Combinations of input and output attributes

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 in Single-operator, see the following parameter construction part.

// Parameter construction
atb::infer::LinearParam param;
param.transposeA = false;
param.transposeB = false;
param.hasBias = false;
param.outDataType = ACL_FLOAT16;
param.enAccum = false;
param.matmulType = MATMUL_UNDEFINED;
param.quantMode = PER_TOKEN;

# Example
>>> x
tensor([[1, 2],
        [3, 4]])
>>> weight
tensor([[1, 2, 3],
        [4, 5, 6]])
>>> deqScale
tensor([1, 2, 3])
>>> perTokenScale
tensor([1, 2])
>>> output
tensor([[9, 24, 45],
        [38, 104, 198]])
# 9 = (1 * 1 + 2 * 4) * 1 * 1
# 24 = (1 * 2 + 2 * 5) * 2 * 1
# 45 = (1 * 3 + 2 * 6) * 3 * 1
# 38 = (3 * 1 + 4 * 4) * 1 * 2
# 104 = (3 * 2 + 4 * 5) * 2 * 2
# 198 = (3 * 3 + 4 * 6) * 3 * 2

Parent topic: Functions