PagedAttentionOperation

功能

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

约束

batch取值范围为(0, 1000]。
要求block_size * head_size ≤ 128 * 128，且block_size为16的整数倍，推荐block_size = 128。
blockTables中元素的值须在[0, num_blocks)之间。
head_size取值范围为(0, 256]。
在Atlas 推理系列产品（配置Ascend 310P AI处理器）上，query、keyValue、valueCache、maskTensor输入的数据类型不支持bf16。
Atlas 推理系列产品（配置Ascend 310P AI处理器）上max_seq_len16对齐，且维度维度描述中的除法均为ceil div。

定义

struct PagedAttentionParam {
    int32_t headNum = 0;
    float qkScale = 1.0; // scale after qkbmm
    int32_t kvHeadNum = 0;
    enum MaskType : int {
        UNDEFINED = 0,
        MASK_TYPE_NORM,
        MASK_TYPE_ALIBI
    };
    MaskType maskType = UNDEFINED;
    bool batchRunStatusEnable = false;
    enum QuantType : int {
        TYPE_QUANT_UNDEFINED = 0,
        TYPE_DEQUANT_FUSION
    };
    QuantType quantType = TYPE_QUANT_UNDEFINED;
    bool hasQuantOffset = false;
};

成员

成员名称	描述
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。
maskType	遮蔽类型。Atlas 推理系列产品（配置Ascend 310P AI处理器）不支持alibi mask压缩。
batchRunStatusEnable	是否开启动态batch。Atlas 推理系列产品（配置Ascend 310P AI处理器）不支持动态batch功能。
quantType	量化类型。Atlas 推理系列产品（配置Ascend 310P AI处理器）不支持反量化融合。
hasQuantOffset	在量化情况下，是否需要输入偏移tensor

输入

参数	维度	数据类型	格式	描述
query	[num_tokens, num_heads, head_size]	float16/bf16	ND	各batch的tensor以num_tokens轴合并。
keyCache	Atlas 800I A2推理产品：[num_blocks, block_size, k_head_num, head_size] Atlas 推理系列产品（配置Ascend 310P AI处理器）：[num_blocks, num_tokens *num_heads / 16 ,block_size, 16]	float16/bf16	Atlas 800I A2推理产品：ND Atlas 推理系列产品（配置Ascend 310P AI处理器）：NZ	已有的key，待与当前key拼接。
valueCache	Atlas 800I A2推理产品：[num_blocks, block_size, v_head_num, head_size] Atlas 推理系列产品（配置Ascend 310P AI处理器）：[num_blocks, num_tokens *num_heads / 16 ,block_size, 16]	float16/bf16	Atlas 800I A2推理产品：ND Atlas 推理系列产品（配置Ascend 310P AI处理器）：NZ	已有的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数量。
maskTensor	Atlas 800I A2推理产品：[batch, num_heads, 1, max_seqlen] Atlas 推理系列产品（配置Ascend 310P AI处理器）：[batch, num_heads, max_seqlen/16, 16, 16]	float16/bf16	Atlas 800I A2推理产品：ND Atlas 推理系列产品（配置Ascend 310P AI处理器）：NZ	Atlas 800I A2推理产品：当maskType不为UNDEFINED时输入。
kDeScale	[num_heads * head_size]	int64	ND	当量化类型为反量化时，步长tensor
kOffset	[num_heads * head_size]	int32	ND	当量化类型为反量化，且hasQuantOffset为true时，输入k的偏移量
vDeScale	[num_heads * head_size]	int64	ND	当量化类型为反量化时，步长tensor
vOffset	[num_heads * head_size]	int32	ND	当量化类型为反量化，且hasQuantOffset为true时，输入v的偏移量

输出

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

父主题： atb/infer_op_params.h