输入输出列表

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

ND/NZ

cpu/npu

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

基础场景

cacheV

[layerNum, batch, maxSeqLen, vhiddenSize]

float16/bf16

ND/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]

float16/bf16

ND/NZ

npu

四种shape分别对应：

1.所有batch相同，方阵。

2.batch不同时的方阵。

3. batch不同时的向量。

4. alibi场景。

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

输出。

基础场景

query

[nTokens, head_num, head_size]

float16/bf16

开启量化功能时为int8

ND

npu

query矩阵, nTokens情况较复杂，见规格约束第一项。

基础场景

key

[nTokens, head_num, head_size]

float16/bf16

开启量化功能时为int8

ND

npu

key矩阵。

基础场景

value

[nTokens, head_num, head_size]

float16/bf16

开启量化功能时为int8

ND

npu

value矩阵。

基础场景

mask

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

开启mask压缩功能时与FA有所不同，见压缩mask

float16/bf16

ND/NZ

npu

四种shape分别对应：

1.所有batch相同，方阵。

2. batch不同时的方阵。

3. batch不同时的向量。

4. alibi场景。

当maskType为undefined时不传此tensor。

mask场景

seqLen

[batch]

int32

ND

cpu

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

基础场景

slopes

[head_num]

• Atlas 800I A2 推理产品/Atlas A2 训练系列产品：float
• Atlas 推理系列产品: float16

ND

npu

当maskType为alibi压缩时需传入此tensor，为alibi mask每个head的系数。

alibi压缩mask场景

output

[nTokens, head_num, head_size]

float16/bf16

ND

npu

输出。

基础场景