Input and Output

query

[ntokens, qHiddenSize]

float16/bf16

ND

npu

Query matrix.

Basic

key

[ntokens, kHiddenSize]

float16/bf16

ND

npu

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

Basic

value

[ntokens, vHiddenSize]

float16/bf16

ND

npu

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

Basic

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

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

attentionMask

See mask

-

-

-

-

Mask

tokenOffset

[batch]

int32/uint32

ND

cpu

Token offset after calculation.

Basic

seqLen

[batch]

int32/uint32

ND

cpu

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

Basic

layerId

[1]

int32/uint32

ND

npu

KV of the cache for calculation.

Basic

batchStatus

[batch]

int32/uint32

ND

cpu

Batch to be calculated when calculable batch control is enabled.

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

output

[nTokens, headNum * vHeadSize]

float16/bf16

ND

npu

Output.

Basic