Per-token Quantization for Matmul

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

The linear operation supports per-token quantization.

Formula

$\text{[math]}$

Parameters

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 stride. The actually occupied memory must be 32-byte aligned.
perTokenScale	[m]	float	ND	Per-token dequantization stride.

Output

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

Specifications

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