MlaPreprocess

English | 简体中文

Product Support

Product	Supported
Ascend 950PR&950DT products	×
Atlas A3 products	√
Atlas A2 products	√
Atlas 200I/500 A2 inference products	×
Atlas inference products	×
Atlas training products	×

Function Description

• Operator function: Computes the preprocessing of Multi-Head Latent Attention in inference scenarios. The main computation process is as follows:

  • First, the input $x$ is processed by RmsNormQuant and then multiplied by $W^{DQKV}$ for downsampling, and the result is split into path 1 and path 2.
  • Path 1 is processed by RmsNormQuant, multiplied by $W^{UQ}$, and then split into path 3 and path 4.
  • Path 3 is multiplied by $W^{uk}$ and outputs $q^N$.
  • Path 4 passes through the rotary position embedding and outputs $q^R$.
  • Path 2 is split into path 5 and path 6.
  • Path 5 passes through RmsNorm and is written into the cache to obtain $k^N$.
  • Path 6 passes through the rotary position embedding and is written into another cache to obtain $k^R$.

• Computation formulas:

  RmsNormQuant formula

  $$\text{RMS}(x)=\sqrt{\frac{1}{N}\sum _{i=1}^N{x}_i^2+ϵ}$$

  $$\text{RmsNorm}(x)=\gamma \cdot \frac{x_i}{\text{RMS}(x)}$$

  $$RmsNormQuant(x)=(RmsNorm(x)+bias)\ast deqScale$$

  Query computation formulas, including the W^{DQKV} matrix multiplication, W^{UK} matrix multiplication, RmsNormQuant, and ROPE rotary position embedding processing

  $$q^N=RmsNormQuant(x)\cdot W^{DQKV}\cdot W^{UK}$$

  $$q^R=ROPE(x^Q)$$

  Key computation formulas, including RmsNorm and ROPE. The computation results are stored in the cache.

  $$k^N=Cache(RmsNorm(RmsNormQuant(x)))$$

  $$k^R=Cache(ROPE(RmsNormQuant(x)))$$

  When both cos and sin of the ACLNN interface are passed as null pointers, the Q-RoPE and K-RoPE mathematical operations are skipped, and the outputs and caches are still written with the unrotated Q-rope and K-rope components in the original layout. When both are non-null tensors, the original RoPE behavior is retained. Passing only one null pointer is prohibited.

Parameter Description

Parameter name	Input/Output/Attribute	Description	Data type	Data format
input	Input	Device-side aclTensor, the x used to compute the Query and Key, with shape [tokenNum,hiddenSize]	FLOAT16, BFLOAT16	ND
gamma0	Input	Device-side aclTensor, the γ parameter in the first RmsNorm computation, with shape [hiddenSize]	FLOAT16, BFLOAT16	ND
beta0	Input	Device-side aclTensor, the β parameter in the first RmsNorm computation, with shape [hiddenSize]	FLOAT16, BFLOAT16	ND
quantScale0	Input	Device-side aclTensor, the quantization scale parameter in the first RmsNorm formula, with shape [1]	FLOAT16, BFLOAT16	ND
quantOffset0	Input	Device-side aclTensor, the quantization offset parameter in the first RmsNorm formula, with shape [1]	INT8	ND
wdqkv	Input	Device-side aclTensor, the dimension reduction matrix that first performs matrix multiplication with the input, with shape [qLoraDim + keyTotalDim,hiddenSize]	INT8, FLOAT16, BFLOAT16	NZ
deScale0	Input	Device-side aclTensor, the coefficient in the dimension reduction matrix that first performs matrix multiplication with the input, with shape [qLoraDim + keyTotalDim]. INT64 is supported when the dtype of input is FLOAT16, and FLOAT is supported when the input is BFLOAT16	INT64, FLOAT	ND
bias0	Input	Device-side aclTensor, the coefficient in the dimension reduction matrix that first performs matrix multiplication with the input, with shape [qLoraDim + keyTotalDim]. Passing an empty tensor is supported; do not pass this parameter when quantMode is 1 or 3	INT32	ND
gamma1	Input	Device-side aclTensor, the γ parameter in the second RmsNorm computation, with shape [qLoraDim]	FLOAT16, BFLOAT16	ND
beta1	Input	Device-side aclTensor, the β parameter in the second RmsNorm computation, with shape [qLoraDim]	FLOAT16, BFLOAT16	ND
quantScale1	Input	Device-side aclTensor, the quantization scale parameter in the second RmsNorm formula, with shape [1]. Pass this parameter only when quantMode is 0	FLOAT16, BFLOAT16	ND
quantOffset1	Input	Device-side aclTensor, the quantization offset parameter in the second RmsNorm formula, with shape [1]. Pass this parameter only when quantMode is 0	INT8	ND
wuq	Input	Device-side aclTensor, the weight matrix, with shape [headNum * (qNoRopeDim + qRopeDim),qLoraDim]	INT8, FLOAT16, BFLOAT16	NZ
deScale1	Input	Device-side aclTensor, the coefficient involved in the wuq matrix multiplication, with shape [headNum * (qNoRopeDim + qRopeDim)]. INT64 is supported when the dtype of input is FLOAT16, and FLOAT is supported when the input is BFLOAT16	INT64, FLOAT	ND
bias1	Input	Device-side aclTensor, the coefficient involved in the wuq matrix multiplication, with shape [headNum * (qNoRopeDim + qRopeDim)]. Do not pass this parameter when quantMode is 1 or 3	INT32	ND
gamma2	Input	Device-side aclTensor, the γ parameter involved in the RmsNormAndreshapeAndCache computation, with shape [512].	FLOAT16, BFLOAT16	ND
cos	Input	Device-side aclTensor, the cosine parameter matrix used to compute the rotary position embedding. When RoPE is enabled, the shape is [tokenNum,64], where tokenNum equals dim 0 of input; RoPE is disabled when both cos and sin are passed as null pointers.	FLOAT16, BFLOAT16	ND
sin	Input	Device-side aclTensor, the sine parameter matrix used to compute the rotary position embedding. When RoPE is enabled, the shape is [tokenNum,64], where tokenNum equals dim 0 of input; RoPE is disabled when both cos and sin are passed as null pointers.	FLOAT16, BFLOAT16	ND
wuk	Input	Device-side aclTensor, the upsampling weight used to compute the Key, with shape [headNum,qNoRopeDim,512].	FLOAT16, BFLOAT16	ND
kvCache	Input	Device-side aclTensor, the same tensor as the output kvCacheOut. The input format varies with cacheMode. When cacheMode is 0: shape [blockNum,blockSize,1,576] When cacheMode is 1: shape [blockNum,blockSize,1,512] When cacheMode is 2: shape [blockNum,16,block_size,32] When cacheMode is 3: shape [blockNum,32,block_size,16]	When cacheMode is 0: same as input When cacheMode is 1: same as input When cacheMode is 2: INT8 When cacheMode is 3: same as input	ND ND NZ NZ
kvCacheRope	Input	Device-side aclTensor. Required; a null pointer is not supported. This is the same tensor as the output krCacheOut. When it does not participate in the computation, a placeholder tensor must still be passed and its content is not read. When cacheMode is 0: does not participate in the computation. Placeholder shape [blockNum,blockSize,1,64], data format ND When cacheMode is 1: shape [blockNum,blockSize,1,64] When cacheMode is 2 or 3: shape [blockNum, 4 ,block_size, 16]	Same as input	ND NZ
slotmapping	Input	Device-side aclTensor, the indices used to store kv_cache and kr_cache, with shape [tokenNum]	INT32	ND
ctkvScale	Input	Device-side aclTensor. Required; a null pointer is not supported. Coefficient involved in the output quantization processing, shape [1]. It participates in the computation only when cacheMode is 2. For any other cacheMode, pass a placeholder of the same shape; its content is not read	FLOAT16, BFLOAT16	ND
qNopeScale	Input	Device-side aclTensor. Required; a null pointer is not supported. Coefficient involved in the output quantization processing, shape [1]. It participates in the computation only when cacheMode is 2. For any other cacheMode, pass a placeholder of the same shape; its content is not read	FLOAT16, BFLOAT16	ND
wdqDim	Input	Indicates the dim size split after the matmul. Reserved parameter; currently only 1536 is supported	int64_t	-
qRopeDim	Input	Indicates the dim size of q passed to rope. Reserved parameter; currently only 64 is supported.	int64_t	-
kRopeDim	Input	Indicates the dim size of k passed to rope. Reserved parameter; currently only 64 is supported.	int64_t	-
epsilon	Input	Indicates the value added to the denominator to prevent division by 0	float	-
qRotaryCoeff	Input	Indicates the q rotation coefficient. Reserved parameter; currently only 2 is supported	int64_t	-
kRotaryCoeff	Input	Indicates the k rotation coefficient. Reserved parameter; currently only 2 is supported	int64_t	-
transposeWdq	Input	Indicates whether wdq is transposed. Reserved parameter; currently only true is supported	bool	-
transposeWuq	Input	Indicates whether wuq is transposed. Reserved parameter; currently only true is supported	bool	-
transposeWuk	Input	Indicates whether wuk is transposed. Reserved parameter; currently only true is supported	bool	-
cacheMode	Input	Indicates the specified cache type. Value range: [0, 3] 0: kcache and q are output after concatenation 1: The output kvCacheOut is split into kvCacheOut and krCacheOut, and qOut is split into qOut and qRopeOut 2: krope and ctkv are converted to the NZ format for output, and ctkv and qnope are quantized to the int8 type through per_head static symmetric quantization 3: krope and ctkv are converted to the NZ format for output	int64_t	-
quantMode	Input	Indicates the specified RmsNorm quantization type. Value range: [0, 3] 0: per_tensor static asymmetric quantization, the default quantization type 1: per_token dynamic symmetric quantization, not implemented 2: per_token dynamic asymmetric quantization, not implemented 3: No quantization, floating-point output, not implemented	int64_t	-
doRmsNorm	Input	Indicates whether to perform the RmsNormQuant operation on the input. false indicates that no operation is performed, and true indicates that the operation is performed. Reserved parameter; currently only true is supported	bool	-
wdkvSplitCount	Input	Indicates the number of wdkv splits. [1-3] is supported, indicating no split, a split into 2 dimension reduction matrices, and a split into 3 dimension reduction matrices respectively. Reserved parameter; currently only 1 is supported	int64_t	-
qOut	Output	Indicates the output tensor of the Query, corresponding to the output on the right side of the computation flow graph after NOPE and matrix multiplication. The shape and dtype vary with cacheMode When cacheMode is 0: shape [tokenNum, headNum, 576] When cacheMode is 1 or 3: shape [tokenNum, headNum, 512] When cacheMode is 2: shape [tokenNum, headNum, 512]	When cacheMode is 0: same as input When cacheMode is 1 or 3: same as input When cacheMode is 2: INT8	ND
kvCacheOut	Output	Indicates the output of the Key after ReshapeAndCache. The shape and dtype vary with cacheMode When cacheMode is 0: shape [blockNum, blockSize, 1, 576] When cacheMode is 1: shape [blockNum, blockSize, 1, 512] When cacheMode is 2: shape [blockNum, 16, block_size, 32] When cacheMode is 3: shape [blockNum, 32, block_size, 16]	When cacheMode is 0: same as input When cacheMode is 1: same as input When cacheMode is 2: INT8 When cacheMode is 3: same as input	ND ND NZ NZ
qRopeOut	Output	Indicates the output of the Query after rotary position embedding. Required; a null pointer is not supported. When it does not participate in the computation, a placeholder tensor must still be passed and its content is not written When cacheMode is 0: does not participate in the computation. Placeholder shape [tokenNum, headNum, 64] When cacheMode is 1 or 3: shape [tokenNum, headNum, 64] When cacheMode is 2: shape [tokenNum, headNum, 64]	When cacheMode is 0: same as input When cacheMode is 1 or 3: same as input When cacheMode is 2: same as input	ND ND ND
krCacheOut	Output	Indicates the output of the Key after ROPE and ReshapeAndCache. Required; a null pointer is not supported. When it does not participate in the computation, a placeholder tensor must still be passed and its content is not written When cacheMode is 0: does not participate in the computation. Placeholder shape [blockNum, blockSize, 1, 64] When cacheMode is 1: shape [blockNum, blockSize, 1, 64] When cacheMode is 2 or 3: shape [blockNum, 4 ,block_size, 16]	When cacheMode is 0: same as input When cacheMode is 1: same as input When cacheMode is 2 or 3: same as input	ND ND NZ

Constraint Description

• Meanings and constraints of shape format fields
  • tokenNum: tokenNum indicates the batch size of input samples. Value range: 0~256
  • hiddenSize: hiddenSize indicates the size of the hidden layer. The value is fixed to the range 2048~10240 and must be a multiple of 256
  • headNum: indicates the number of heads. Value range: 1~128
  • blockNum: the number of blocks in the PagedAttention scenario. The value is fixed to 192
  • blockSize: the block size in the PagedAttention scenario. The value is fixed to 128
  • qloraDim: indicates the LoRA input dimension of the Q matrix. Value range: 32~4096, a multiple of 32
  • keyTotalDim: indicates the total dimension of the Key part. The value is fixed to 576 (512 main dimension + 64 rope dimension)
  • qRopeDim: indicates the dimension of the rotary embedding part in the Q matrix. The value is fixed to 64
  • qNoRopeDim: indicates the dimension of the non-rotary embedding part in the Q matrix. Value range: 16~256, a multiple of 16
• RoPE mode constraint
  • The Rotary Position Embedding (RoPE) operation in the mla_preprocess operator uses the half mode. The interleave mode is not yet supported.