Input and Output

query

[ntokens, qHiddenSize]

float16/bf16

ND

npu

Query matrix.

Basic scenario

key

[ntokens, kHiddenSize]

float16/bf16

ND

npu

Key matrix. This parameter is not passed when kvcacheCfg is set to K_BYPASS_V_BYPASS.

Basic scenario

value

[ntokens, vHiddenSize]

float16/bf16

ND

npu

Value matrix. This parameter is not passed when kvcacheCfg is set to K_BYPASS_V_BYPASS.

Basic scenario

cacheK

[layerNum, batch, maxSeqLen, kHiddenSize]

float16/bf16

Atlas A2 training products/Atlas A2 inference products and Atlas A3 inference products/Atlas A3 training products: ND

Atlas inference products: NZ

cpu/npu

• NPU: stores all previous keys. During execution, keys are updated to cacheK.
• CPU: The input is the prepared cacheK. The input is divided into batch tensors as the input std::vector<tensor>. In this case, the layer dimension must be 1.

Basic scenario

cacheV

[layerNum, batch, maxSeqLen, vhiddenSize]

float16/bf16

Atlas A2 training products/Atlas A2 inference products and Atlas A3 inference products/Atlas A3 training products: ND

Atlas inference products: NZ

npu/cpu

• NPU: stores all previous values. During execution, values are updated to cacheV.
• CPU: The input is the prepared cacheV. The input is divided into batch tensors as the input std::vector<tensor>. In this case, the layer dimension must be 1.

Basic scenario

attentionMask

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

float16/bf16

Atlas A2 training products/Atlas A2 inference products and Atlas A3 inference products/Atlas A3 training products: ND

Atlas inference products: ND/NZ

npu

The six shapes are described as follows:

1. Same batch or SWA norm mask, square matrix.

2. Square matrix with different batches.

3. Vectors with different batches.

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

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

6. SWA compress mask.

Mask scenario

tokenOffset

[batch]

int32/uint32

ND

cpu

Token offset after calculation.

Basic scenario

seqLen

[batch]

int32/uint32

ND

cpu

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

Basic scenario

layerId

[1]

int32/uint32

ND

npu

KV of the cache for calculation.

Basic scenario

batchStatus

[batch]

int32/uint32

ND

cpu

Batch to be calculated when the calculable batch function is enabled.

Scenario of calculable batch control

logN

• Full phase: [maxSeqLen]
• Incremental phase: [batch]

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. In the incremental phase, the value is the logN corresponding to the new request in each batch. This tensor is required when the scale type is logN.

LogN scenario

output

[nTokens, headNum * vHeadSize]

float16/bf16

ND

npu

Output.

Basic scenario