Input and Output When calcType Is PREFIX_ENCODER and maskType Is MASK_TYPE_CAUSAL_MASK

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, headNum, 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, headNum, 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 block table < maxBlockNum

max(kvSeqLen) / blockSize <= maxblockNum

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

output

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

float16/bf16

ND

npu

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

Basic