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 |
npu |
四种shape分别对应： 1.所有batch相同，方阵。 2. batch不同时的方阵。 3. q_seqlen为1时的mask向量。 4. alibi场景。 当maskType为undefined时不传此tensor。 |
mask场景 |
|
seqlen |
[batch] / [2, batch] |
int32/uint32 |
ND |
cpu |
值等于1时，为增量或全量；大于1时，为全量。 若shape为[batch] ，代表每个batch的序列长度，query，cacheK，cacheV相同。 若shape为[2，batch]，seqlen[0]代表query的序列长度，seqlen[1]代表cacheK，cacheV的序列长度。 |
基础场景 |
slopes |
[headNum] |
|
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 |
输出。 |
基础场景 |