PagedAttentionOperation

attention计算，支持k, v内存不连续存放。

struct PagedAttentionParam {
    int32_t headNum = 0;
    float qkScale = 1.0; // scale after qkbmm
    int32_t kvHeadNum = 0;
};

成员名称	描述
headNum	attention多头的个数，需要大于0。
qkScale	qk相乘后的缩放系数，需根据实际模型参数需求传入。
kvHeadNum	该值需要用户根据使用的模型实际情况传入。 kvHeadNum = 0时，keyCache的k_head_num，valueCache的v_head_num与query的num_heads一致，均为num_heads的数值。 kvHeadNum != 0时，keyCache的k_head_num， valueCache的v_head_num与kvHeadNum值相同。 kvHeadNum需要满足大于或等于0，当kvHeadNum大于0时, 须满足kvHeadNum mod headNum为0。

成员名称

描述

headNum

attention多头的个数，需要大于0。

qkScale

qk相乘后的缩放系数，需根据实际模型参数需求传入。

kvHeadNum

该值需要用户根据使用的模型实际情况传入。

kvHeadNum = 0时，keyCache的k_head_num，valueCache的v_head_num与query的num_heads一致，均为num_heads的数值。
kvHeadNum != 0时，keyCache的k_head_num， valueCache的v_head_num与kvHeadNum值相同。
kvHeadNum需要满足大于或等于0，当kvHeadNum大于0时, 须满足kvHeadNum mod headNum为0。

参数	维度	数据类型	格式	描述
query	[num_tokens, num_heads, head_size]	float16	ND	各batch的tensor以num_tokens轴合并。
keyCache	[num_blocks, block_size, k_head_num, head_size]	float16	ND	已有的key，待与当前key拼接。
valueCache	[num_blocks, block_size, v_head_num, head_size]	float16	ND	已有的value，待与当前value拼接。
blockTables	[num_tokens, max_num_blocks_per_query]	int32	ND	每个query的kvcache的block table，第一维是token索引，第二维表示block索引。
contextLens	[num_tokens]	int32	ND	每个query对应的key/value的token数量。

参数	维度	数据类型	格式
output	[num_tokens, num_heads, head_size]	float16	ND

父主题： atb/infer_op_params.h