PagedAttentionOperation
功能
PagedAttention。
一个Q有多个token,一个token对应多个KV的token,以token0为例,block_table代表其对应的KV的block_id,-1代表截止, 所以第二行和第四行为其目标block,context_lens则表示KV有多少个token,则代表仅有block_id为(3,4,5,9,10)是需要与Q进行计算的。
约束
- 因硬件限制,embedding * blockSize ≤ 128 * 128。
- block_size %16 == 0 推荐 block_size = 128,blockTables中元素的值须在[0, num_blocks)之间,head_size 范围(0, 256] 。
- query、keyCache、valueCache、maskTensor四个入参在Atlas 推理系列产品上 只支持float16。
- Atlas 推理系列产品上 0<batch<=2000。
- 关于选取何种mask,如果原模型非alibi,且mask没有做padding,可以使用UNDEFINED格式的mask,即不传mask;否则依据原生模型的配置,决定是NORM还是ALIBI。
定义
struct PagedAttentionParam { enum MaskType : int { UNDEFINED = 0, MASK_TYPE_NORM, MASK_TYPE_ALIBI, MASK_TYPE_SPEC }; enum QuantType : int { TYPE_QUANT_UNDEFINED = 0, TYPE_DEQUANT_FUSION }; enum CompressType : int{ COMPRESS_TYPE_UNDEFINED = 0, COMPRESS_TYPE_KVHEAD }; enum CalcType : int { CALC_TYPE_UNDEFINED = 0, CALC_TYPE_SPEC }; int32_t headNum = 0; float qkScale = 1.0; // scale after qkbmm int32_t kvHeadNum = 0; MaskType maskType = UNDEFINED; bool batchRunStatusEnable = false; QuantType quantType = TYPE_QUANT_UNDEFINED; bool hasQuantOffset = false; CompressType compressType = COMPRESS_TYPE_UNDEFINED; CalcType calcType = CALC_TYPE_UNDEFINED; };
成员
成员名称 |
描述 |
---|---|
maskType |
MaskType的类型值。。
|
quantType |
量化类型。
|
compressType |
压缩类型。
|
calcType |
CalcType的类型值。
|
headNum |
query 头大小。 |
qkScale |
算子tor值, 在Q*K^T后乘。 |
kvHeadNum |
kv头数量。 |
maskType |
mask类型。 |
batchRunStatusEnable |
是否开启动态batch。 |
quantType |
量化类型。 |
hasQuantOffset |
开启量化功能后是否使用offset。 |
compressType |
压缩方式。 |
calcType |
计算类型。 |
输入
参数 |
维度 |
数据类型 |
格式 |
设备 |
描述 |
---|---|---|---|---|---|
query |
[num_tokens, num_heads, head_size] |
float16/bfloat16 |
ND |
npu |
各batch的query在num_tokens轴合并。 |
keyCache |
Atlas 800I A2推理产品:[num_blocks, block_size, k_head_num, head_size] Atlas 推理系列产品:[num_blocks, num_tokens *num_heads / 16 ,block_size, 16] |
float16/bfloat16/int8 |
Atlas 800I A2推理产品:ND Atlas 推理系列产品:NZ |
npu |
已cache的key,当开启反量化融合功能时,只支持int8。 |
valueCache |
Atlas 800I A2推理产品:[num_blocks, block_size, v_head_num, head_size] Atlas 推理系列产品:[num_blocks, num_tokens *num_heads / 16 ,block_size, 16] |
float16/bfloat16/int8 |
Atlas 800I A2推理产品:ND Atlas 推理系列产品:NZ |
npu |
已cache的value,当开启反量化融合功能时,只支持int8。 |
blockTables |
[num_tokens, max_num_blocks_per_query] |
int32 |
ND |
npu |
每个query的kvcache的block table,第一维是token索引,第二维表示block索引。 |
contextLens |
[num_tokens] |
int32 |
ND |
cpu |
每个query对应的key/value的token数量(token数量 // block_size == block_num + 1)。 |
mask |
见表1 |
float16/bfloat16 |
Atlas 800I A2推理产品:ND Atlas 推理系列产品:NZ |
npu |
Attention mask。 Atlas 800I A2推理产品:当maskType不为UNDEFINED时输入。 |
batchRunStatus |
[batch] |
int32 |
ND |
cpu |
当开启动态batch功能时需要传此tensor以作为哪些batch参与计算的标志位。 |
kDescale |
[k_head_num*head_size] |
int64/fp32 |
ND |
npu |
当量化类型为反量化时,步长tensor。数据类型需与vDescale保持一致。 |
kOffset |
[k_head_num*head_size] |
int32 |
ND |
npu |
当量化类型为反量化,且hasQuantOffset为true时,输入k的偏移量。 |
vDescale |
[v_head_num*head_size] |
int64/fp32 |
ND |
npu |
当量化类型为反量化时,步长tensor。 |
vOffset |
[v_head_num*head_size] |
int32 |
ND |
npu |
当量化类型为反量化,且hasQuantOffset为true时,输入v的偏移量。 |
maskType |
硬件类型 |
维度 |
---|---|---|
UNDEFINED |
不传mask |
不传mask |
MASK_TYPE_NORM |
Atlas 800I A2推理产品 |
[batch, 1, max_seq_len] 或 [1, max_seq_len] 或 [max_seq_len, max_seq_len] |
Atlas 推理系列产品 |
[batch, max_seq_len / 16, 16, 16] 或 [1, max_seq_len / 16, 16, 16] |
|
MASK_TYPE_ALIBI |
Atlas 800I A2推理产品 |
[batch, num_head, 1, max_seq_len] 或 [num_head, 1, max_seq_len] |
Atlas 推理系列产品 |
[batch * num_head, max_seq_len / 16, 16, 16] 或 [num_head, max_seq_len / 16, 16, 16] |
|
MASK_TYPE_SPEC |
Atlas 800I A2推理产品 |
[num_tokens, max_seq_len] |
Atlas 推理系列产品 |
[1, max_seq_len / 16, num_tokens, 16] |
输出
参数 |
维度 |
数据类型 |
格式 |
设备 |
描述 |
---|---|---|---|---|---|
attnOut |
[num_tokens, num_head, head_size] |
float16/bfloat16 |
ND |
npu |
经过计算输出的query。 |