输入输出列表
列表中的shape仅为通用场景下的shape。特殊场景的shape见功能列表。
参数 |
维度 |
数据类型 |
格式 |
cpu/npu |
描述 |
使用场景 |
|---|---|---|---|---|---|---|
query |
[num_tokens, num_head, head_size] |
float16/bf16/int8 |
ND |
npu |
各batch的query在num_tokens轴合并；当开启全量化功能时仅支持int8。 |
基础场景 |
keyCache |
|
float16/bf16/int8 |
|
npu |
cache好的key；当开启反量化融合功能和全量化功能时，只支持int8。 |
基础场景 |
valueCache |
|
float16/bf16/int8 |
|
npu |
cache好的value；当开启反量化融合功能和全量化功能时，只支持int8。当开启mla合并输出kvcache功能时不传入。 |
基础场景 |
blockTables |
[num_tokens, max_num_blocks_per_query] |
int32 |
ND |
npu |
每个query的kvcache的block table，第一维是token索引，第二维表示block索引。 |
基础场景 |
contextLens |
[batch] |
int32 |
ND |
cpu |
每个query对应的key/value的token数量。 |
基础场景 |
mask |
较复杂，见功能列表mask类型 |
float16/bf16 |
|
npu |
mask场景 |
|
batchRunStatus |
[batch] |
int32 |
ND |
cpu |
当开启控制可计算batch功能时需要传此tensor以作为哪些batch参与计算的标志位。 |
控制可计算batch场景 |
kDescale |
[k_head_num*head_size] |
int64/float |
ND |
npu |
当量化类型为反量化时，步长tensor；数据类型需与vDescale保持一致。 |
反量化融合场景或全量化场景 |
kOffset |
[k_head_num*head_size] |
int32 |
ND |
npu |
当量化类型为反量化，且“hasQuantOffset”为true时，输入k的偏移量。 |
非对称反量化场景 |
vDescale |
[v_head_num*head_size] |
int64/float |
ND |
npu |
当量化类型为反量化时，步长tensor。 |
反量化融合场景或全量化场景 |
vOffset |
[v_head_num*head_size] |
int32 |
ND |
npu |
当量化类型为反量化，且“hasQuantOffset”为true时，输入v的偏移量。 |
非对称反量化场景 |
qSeqLens |
[batch] |
int32 |
ND |
cpu |
当开启并行解码功能时需要传此tensor，每个batch对应的seqLen。 |
并行解码场景 |
razorOffset |
[num_blocks, block_size] |
float |
ND |
npu |
当开启Razor Rope功能时需要传此tensor。 |
多头自适应压缩（rope场景） |
pScale |
[q_head_num] |
float |
ND |
npu |
当开启离线全量化时需要传此tensor，传入attention计算中间值p_score的量化scale。 |
离线全量化场景 |
logN |
[batch] |
float |
ND |
npu |
当scale类型为logN时需要传此tensor，为各batch对应的缩放系数。 |
LogN场景 |
attnOut |
[num_tokens, num_head, head_size_v] |
float16/bf16 |
ND |
npu |
经过计算输出的query。 |
基础场景 |
上表中，contextLens和qSeqLens较为特殊，其虽为输入tensor，在使用时却是当做param使用，故需要有一个bind tensor的操作，使其从device侧提取出来绑定到host侧。contextLens在
上表中，使用场景表示各个输入tensor所对应的场景，其中基础场景是指该tensor为必须项，PA所有情况均有此tensor，其中包含5个输入tensor和1个输出tensor。而属于其他使用场景的输入tensor则是在对应场景下才能够传此输入tensor，非对应场景下不能传此输入tensor。以下为不同场景的定义：
- mask场景指“maskType”不为UNDEFINED时。
- 反量化融合场景指“quantType”为TYPE_DEQUANT_FUSION时。
- 全量化场景指“quantType”为TYPE_QUANT_QKV_OFFLINE（离线全量化场景）或TYPE_QUANT_QKV_ONLINE（在线全量化场景）时。
- 非对称反量化场景指“quantType”为TYPE_DEQUANT_FUSION且“hasQuantOffset”为True时。
- 并行解码场景指“calcType”为CALC_TYPE_SPEC时。
- 多头自适应压缩（rope场景）指“compressType”为COMPRESS_TYPE_KVHEAD_ROPE时。
- 离线全量化场景指“quantType”为TYPE_QUANT_QKV_OFFLINE时。
- LogN场景指“scaleType”为SCALE_TYPE_LOGN时。