calcType为PREFIX_ENCODER时输入输出
参数 |
维度 |
数据类型 |
格式 |
cpu/npu |
描述 |
使用场景 |
---|---|---|---|---|---|---|
query |
[nTokens, qHiddenSize] 或[nTokens, headNum, headSize] |
float16/bf16 |
ND |
npu |
query矩阵。
|
基础场景 |
key |
[numBlocks, blockSize, kvHiddenSize] 或[numBlocks, blockSize, kvHeadNum, headSize] |
float16/bf16 |
ND |
npu |
key矩阵。
|
基础场景 |
value |
[numBlocks, blockSize, kvHiddenSize] 或[numBlocks, blockSize, kvHeadNum, 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 |
|
float16/bf16 |
ND |
npu |
|
基础场景 |
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 |
[nTokens, qHiddenSize] 或[nTokens, headNum, headSize] |
float16/bf16 |
ND |
npu |
输出。维度和query一致。 |
基础场景 |