Input and Output When calcType Is PREFIX_ENCODER

query

[nTokens, qHiddenSize] or [nTokens, headNum, headSize]

float16/bf16

ND

npu

Query matrix

• The shape of query does not need to be the same as that of key or value.
• PrefixCacheLen is calculated by subtracting seqLen from kvseqlen. Due to implementation restrictions, prefixCacheLen must be 128-aligned.
• headSize must be 16-aligned.

Basic scenario

key

[numBlocks, blockSize, kvHiddenSize] or [numBlocks, blockSize, kvHeadNum, headSize]

float16/bf16

ND

npu

Key matrix

• HeadSize of key and value <= 128
• blockSize must be set to 128.
• BlockSize × HeadSize <= 128 × 128.

Basic scenario

value

[numBlocks, blockSize, kvHiddenSize] or [numBlocks, blockSize, kvHeadNum, headSize]

float16/bf16

ND

npu

Value matrix The shapes of the key and value must be the same.

Basic scenario

blockTables

[batch, maxBlockNum]

int32

ND

npu

Block table of kvcache of each query. The first dimension is the token index, and the second dimension indicates the block index.

Value of the item in blockTables < maxBlockNum

max(kvSeqLen) / blockSize <= maxblockNum

Basic scenario

mask

• [256, 256]
• [head, seqLen, 128]
• [128, 128]

float16/bf16

ND

npu

• MASK_TYPE_ALIBI_COMPRESS supports dimensions 1 and 2.
• MASK_TYPE_ALIBI_COMPRESS_SQRT supports dimensions 1 and 2.
• MASK_TYPE_NORM_COMPRESS supports dimension 3.

Basic scenario

seqLen

[batch]

int32/uint32

ND

cpu

Sequence length of each batch corresponding to query.

Basic scenario

kvSeqLen

[batch]

int32/uint32

ND

cpu

Sequence length of each batch corresponding to key and value.

Basic scenario

slopes

[head]

float32

ND

npu

ALiBi coefficient. When alibiMask is used and a tensor is passed, the value is the coefficient of each head of alibiMask.

Alibi compression mask

output

[nTokens, qHiddenSize] or [nTokens, headNum, headSize]

float16/bf16

ND

npu

Output. The dimension is the same as that of query.

Basic scenario