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

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

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

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 item in blockTables < maxBlockNum

max(kvSeqLen) / blockSize <= maxblockNum

Basic

mask

See mask

-

-

-

-

Basic

seqLen

[batch]

int32/uint32

ND

cpu

Sequence length of each batch corresponding to query.

Basic

kvSeqLen

[batch]

int32/uint32

ND

cpu

Sequence length of each batch corresponding to key and value.

Basic

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-compressed mask

output

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

float16/bf16

ND

npu

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

Basic