Input and Output When calcType Is PA_ENCODER

query

[nTokens, headNum, headSize]

float16/bf16

int8 when quantization is enabled

ND

npu

Query matrix. For details about nTokens, see the first item in Constraints.

Basic

key

[nTokens, headNum, headSize]

float16/bf16

int8 when quantization is enabled

ND

npu

Key matrix.

Basic

value

[nTokens, headNum, headSize]

float16/bf16

int8 when quantization is enabled

ND

npu

Value matrix. This tensor is not passed when mlaVHeadSize > 0.

Basic

mask

See mask

-

-

-

-

Mask

seqlen

[batch] / [2, batch]

int32/uint32

ND

cpu

1: incremental or full mode. > 1: full mode.

If the shape is [batch], it indicates the sequence length of each batch. The sequence length of each batch is the same for query, cacheK, and cacheV.

If the shape is [2, batch], seqlen[0] indicates the sequence length of query, and seqlen[1] indicates the sequence length of cacheK and cacheV.

Basic

slopes

[headNum]

• Atlas A2 training products/Atlas A2 inference products and Atlas A3 inference products/Atlas A3 training products: float
• Atlas inference products and Atlas training products: float16

ND

npu

This tensor is required when maskType is set to alibi. It is the coefficient of each head of the ALiBi mask.

ALiBi-compressed mask

logN

Full phase: [maxSeqLen]

float

ND

npu

In the full phase, the value is the logN sequence of the length maxSeqLen. Each request in a batch obtains a value from this vector based on seqlen. This tensor is required when the scale type is logN.

LogN

output

[nTokens, headNum, headSize]

float16/bf16

ND

npu

Output.

Basic