MLA Input and Output

qNope

[tokenNum, headNum, 512]

float16/bf16/int8

ND

NPU

Input. Query without positional encoding.

If cacheMode is set to 2, the value is of the int8 data type.

qRope

[tokenNum, headNum, 64]

float16/bf16

ND

NPU

Input. Query with rotational positional encoding.

ctKV

[blockNum, blockSize, kvHeadNum, 512]

• If cacheMode is set to 2:

  [blockNum, kvHeadNum*512/32, blockSize, 32]

• If cacheMode is set to 3:

  [blockNum, kvHeadNum*512/16, blockSize, 16]

float16/bf16/int8

ND/NZ

NPU

Input. ctkv without positional encoding.

• If cacheMode is set to 2, the value is of the int8 data type and the format is NZ.
• If cacheMode is set to 3: the value is of the float16/bf16 data type and the format is NZ.

kRope

[blockNum, blockSize, kvHeadNum, 64]

If cacheMode is set to 2 or 3:

[blockNum, kvHeadNum*64 / 16 ,blockSize, 16]

float16/bf16

ND/NZ

NPU

Input. k with rotational positional encoding.

If cacheMode is set to 2 or 3, NZ is used.

block_tables

[batch, maxBlockNumPerQuery]

int32

ND

NPU

Input. kvcache block mapping table of each query.

kvSeqlen

[batch]

int32

ND

CPU

Input. Context length (kvseqlen) corresponding to each query.

mask

• MASK_TYPE_SPEC:

  [tokenNum (combined axis), maxSeqLen]

• MASK_TYPE_MASK_FREE:

  [125 + 2 * qSeqLen, 128]

  [batch, 125 + 2 * qSeqLen, 128]

float16/bf16

ND

NPU

Input. Attention mask. This parameter is passed when maskType is not 0.

qseqlen

[batch]

int32

ND

CPU

Input. seqLen corresponding to each batch. This parameter is passed when calcType is 1. The value range is [1, 4].

qkDescale

[headNum]

float

ND

NPU

Input. This parameter is passed when cacheMode is 2.

pvDescale

[headNum]

float

ND

NPU

Input. This parameter is passed when cacheMode is 2.

maskUseStatus

[batch]

int32

ND

CPU

This parameter can be used in the MTP1 scenario where the Spec is used together with the Ring. That is, this parameter is passed when calcType is 3. This parameter indicates whether each mask is used.

How to enable this function:

maskUseStatus = MASK_USE_STATUS_TYPE_BATCH_MASK

Internal padding value:

1: The mask is used.

0: No mask is used.

MTP1: qseqlen[i] is 2.

attenOut

[tokenNum, headNum, 512]

float16/bf16

ND

NPU

Attention output.

Ise

[tokenNum, headNum, 1]

float16/bf16/float

ND

NPU

Ise output.

• This tensor is required when calcType is set to CALC_TYPE_RING or CALC_TYPE_SPEC_AND_RING.
• If calcType is set to CALC_TYPE_RING, MaskType can only be set to UNDEFINED.
• If calcType is set to CALC_TYPE_SPEC_AND_RING:
  • MaskType can be set to MASK_TYPE_SPEC and cacheMode can be set to KROPE_CTKV.
  • MaskType can be set to UNDEFINED and cacheMode can be set to INT8_NZCACHE.
• If cacheMode is set to INT8_NZCACHE, dtype is float.