输入输出

query

[ntokens, qHiddenSize]

float16/bf16

ND

npu

query矩阵。

基础场景

key

[ntokens, kHiddenSize]

float16/bf16

ND

npu

key矩阵，当“kvcacheCfg”配置为K_BYPASS_V_BYPASS，不传入。

基础场景

value

[ntokens, vHiddenSize]

float16/bf16

ND

npu

value矩阵，当“kvcacheCfg”配置为K_BYPASS_V_BYPASS，不传入。

基础场景

cacheK

[layerNum, batch, maxSeqLen, kHiddenSize]

float16/bf16

Atlas A2 训练系列产品/Atlas 800I A2 推理产品和Atlas A3 推理系列产品/Atlas A3 训练系列产品：ND

Atlas 推理系列产品：NZ

cpu/npu

• NPU：存储之前所有的k，本次执行时将key刷新到cacheK上 。
• CPU：输入为已经准备好的cacheK，输入时根据分成batch个tensor作为std::vector<tensor>传入，此时layer维度要求为1。

基础场景

cacheV

[layerNum, batch, maxSeqLen, vhiddenSize]

float16/bf16

Atlas A2 训练系列产品/Atlas 800I A2 推理产品和Atlas A3 推理系列产品/Atlas A3 训练系列产品：ND

Atlas 推理系列产品：NZ

npu/cpu

• NPU：存储之前所有的v，本次执行时将value刷新到cacheV上。
• CPU：输入为已经准备好的cacheV，输入时根据分成batch个tensor作为std::vector<tensor>传入，此时layer维度要求为1。

基础场景

attentionMask

[maxSeqLen, maxSeqLen]/[batch, maxSeqLen, maxSeqLen]/[batch, 1, maxSeqLen]/[batch, headNum, maxSeqLen, maxSeqLen]/[512,512]

float16/bf16

Atlas A2 训练系列产品/Atlas 800I A2 推理产品和Atlas A3 推理系列产品/Atlas A3 训练系列产品：ND

Atlas 推理系列产品：ND/NZ

npu

五种shape分别对应：

1.所有batch相同或SWA norm mask，方阵。

2.batch不同时的方阵。

3. batch不同时的向量。

4. alibi场景。

5.SWA compress mask。

mask场景

tokenOffset

[batch]

int32/uint32

ND

cpu

计算完成后的token偏移。

基础场景

seqLen

[batch]

int32/uint32

ND

cpu

等于1时，为增量或全量；大于1时，为全量。

基础场景

layerId

[1]

int32/uint32

ND

npu

取cache的kv中哪一个kv进行计算。

基础场景

batchStatus

[batch]

int32/uint32

ND

cpu

开启控制可计算batch功能时，通过标志位控制具体需要运算的batch。

控制可计算batch场景

logN

• 全量阶段为[maxSeqLen]
• 增量阶段为[batch]

float

ND

npu

全量阶段为长度maxSeqLen的logN序列，batch内每条请求根据自己的序列长度seqlen从该向量中取值；增量阶段为各batch新增请求对应的logN；当scale类型为logN时需要传此tensor。

LogN场景

output

[nTokens, vHiddenSize]

float16/bf16

ND

npu

输出。

基础场景