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 of Restrictions.

Basic scenario

key

[nTokens, headNum, headSize]

float16/bf16

int8 when quantization is enabled

ND

npu

Key matrix

Basic scenario

value

[nTokens, headNum, headSize]

float16/bf16

int8 when quantization is enabled

ND

npu

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

Basic scenario

mask

1. [maxSeqLen, maxSeqLen]
2. [batch, maxSeqLen, maxSeqLen]
3. [batch, 1, maxSeqLen]
4. [batch, headNum, maxSeqLen, maxSeqLen]
5. [headNum, maxSeqLen, maxSeqLen]
6. [512, 512]
7. [2048, 2048]

Different from FA when mask compression is enabled. For details, see Compression Mask.

float16/bf16

Atlas 350 accelerator card, Atlas A2 training products/Atlas A2 inference products, and Atlas A3 inference products/Atlas A3 training products: ND

Atlas inference products and Atlas training products: NZ

npu

The four shapes are described as follows:

1. Square matrix with the same batch

2. Matrix with different batches.

3. Mask vector when q_seqlen is 1.

4. Square matrix with different batches in the ALiBi scenario.

5. Square matrix with the same batch in the ALiBi scenario.

6. SWA compression matrix.

7. Atlas 350 accelerator card: square matrix when maskType is MASK_TYPE_NORM or MASK_TYPE_NORM_COMPRESS.

This tensor is not passed when maskType is undefined.

Mask scenario

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 scenario

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 compression 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 scenario

output

[nTokens, headNum, headSize]

float16/bf16

ND

npu

Output.

Basic scenario