输入输出列表

参数

维度

数据类型

格式

cpu/npu

描述

使用场景

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 800I A2 推理产品/Atlas 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 800I A2 推理产品/Atlas 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 800I A2 推理产品/Atlas 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

输出。

基础场景

当参数“calcType”为PA_ENCODER时，输入输出如下：

参数

维度

数据类型

格式

cpu/npu

描述

使用场景

query

[nTokens, headNum, headSize]

float16/bf16

开启量化功能时为int8

ND

npu

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

基础场景

key

[nTokens, headNum, headSize]

float16/bf16

开启量化功能时为int8

ND

npu

key矩阵。

基础场景

value

[nTokens, headNum, headSize]

float16/bf16

开启量化功能时为int8

ND

npu

value矩阵。当mlaVHeadSize > 0时不传此tensor。

基础场景

mask

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

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

float16/bf16

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

Atlas 推理系列产品：NZ

npu

四种shape分别对应：

1.所有batch相同，方阵。

2. batch不同时的方阵。

3. decoder场景下，q_seqlen为1时的mask向量。

4. alibi场景。

当maskType为undefined时不传此tensor。

mask场景

seqLen

[batch]

int32

ND

cpu

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

基础场景

slopes

[headNum]
  • Atlas 800I A2 推理产品/Atlas A2 训练系列产品和Atlas A3 推理系列产品/Atlas A3 训练系列产品：float
  • Atlas 推理系列产品: float16

ND

npu

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

alibi压缩mask场景

logN

全量阶段为[maxSeqLen]

float

ND

npu

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

LogN场景

output

[nTokens, headNum, headSize]

float16/bf16

ND

npu

输出。

基础场景

当参数“calcType”为PREFIX_ENCODER时，输入输出如下：

参数

维度

数据类型

格式

cpu/npu

描述

使用场景

query

[batch*qSeqLen, qHiddenSize] 或[batch*qSeqLen, headNum, headSize]

float16/bf16

ND

npu

query矩阵。

  • query的shape和key、value的不需要一致。
  • PrefixCacheLen的计算由kvseqlen减seqLen得到，由于实现方案限制，需要让prefixCacheLen 128对齐。
  • headSize需要16对齐。

基础场景

key

[numBlocks, blockSize, kvHiddenSize] 或[numBlocks, blockSize, headNum, headSize]

float16/bf16

ND

npu

key矩阵。

  • key，value的HeadSize <= 128。
  • blockSize必须设为128。
  • BlockSize × HeadSize <= 128 × 128。

基础场景

value

[numBlocks, blockSize, kvHiddenSize] 或[numBlocks, blockSize, headNum, headSize]

float16/bf16

ND

npu

value矩阵。key，value的shape必须一致。

基础场景

blockTables

[batch, maxBlockNum]

int32

ND

npu

每个query的kvcache的block table，第一维是token索引，第二维表示block索引。

blocktable 中item的value < maxBlockNum

max(kvSeqLen) / blockSize <= maxblockNum

基础场景

mask
  1. [256, 256]
  2. [head, seqLen, 128]
  3. [128, 128]

float16/bf16

ND

npu

仅支持压缩mask。

  • MASK_TYPE_ALIBI_COMPRESS支持维度1和2。
  • MASK_TYPE_ALIBI_COMPRESS_SQRT支持维度1和2。
  • MASK_TYPE_NORM_COMPRESS支持维度3。

基础场景

seqLen

[batch]

int32/uint32

ND

cpu

query对应的每个batch的序列长度。

基础场景

kvSeqLen

[batch]

int32/uint32

ND

cpu

key, value对应的每个batch的序列长度。

基础场景

slopes

[head]

float32

ND

npu

即alibi coefficient，当使用alibiMask时传入tensor时，为alibiMask的每个head的系数。

alibi压缩mask场景

output

[batch*qSeqLen, qHiddenSize] 或[batch*qSeqLen, headNum, headSize]

float16/bf16

ND

npu

输出。维度和query一致。

基础场景
父主题： SelfAttentionOperation