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 scenario

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 scenario

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 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 blocktable < maxBlockNum

max(kvSeqLen) / blockSize <= maxblockNum

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

output

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

float16/bf16

ND

npu

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

Basic scenario