SelfAttentionOperation

功能

KVCache + KVCache + Muls + FlashAttention。

图1 SelfAttentionOperation

约束

tokenOffset ≥ seqLen。
attentionMask的维度：
- 支持[batch, maxSeqLen, maxSeqLen]，此时需要保证batch的维度和其他输入batch维度的一致性。
- 支持[maxSeqLen, maxSeqLen]，此时所有batch的mask相同。
- 支持[batch, 1, maxSeqLen]，此时所有batch的mask不同，mask为向量。
SelfAttentionOperation内包含枚举CalcType。
- 当calcType为PA_ENCODER时，为pagedAttention相应的全量阶段。
- 当calcType为ENCODER、DECODER时，为flashAttention。

在calcType = PA_ENCODER下，Atlas 800I A2推理产品上的q，k，v可传二维[nTokens, hiddenSize]或四维[batch, seq_len, head_num, head_size]。
当开启高精度功能且“maskType”为“NORM”或“NORM_COMPRESS”时，“mask”值需传1。

对于Atlas 推理系列产品上，0 < batch <= 2000。
如需使用GQA模式，需满足以下条件：
- headNum > kvHeadNum
- headNum % kvHeadNum == 0
- head_size取值范围为（0, 256]。
Atlas 推理系列产品上的mask大小必须使用真实的max_seq_len，q_seq_len必须等于kv_seq_len。

关于tensor维度中的nTokens：
- Atlas 800I A2推理产品上为各batch上seq_len之和。
- Atlas 推理系列产品上：PA_ENCODER下为所有batch的seq_len之和向上对齐到16的整数倍，其余情况下为所有batch上的seq_len先向上对齐到16的整数倍，再求和。
- 开启压缩mask功能且maskType为MASK_TYPE_ALIBI_COMPRESS时，mask的维度：在Atlas 800I A2推理产品上为[head_num, seqlen, 128]或[256, 256]。Atlas 推理系列产品上为[head_num, 128//16, maxSeqlen, 16]或[1, 256//16, 256, 16]。
- 开启压缩mask功能且maskType为MASK_TYPE_NORM_COMPRESS时，mask的维度：在Atlas 800I A2推理产品上为[128， 128]。Atlas 推理系列产品上为[1, 128//16, 128, 16]。
关于选取何种mask：
- 在全量阶段，依据原生模型的配置，决定是NORM还是ALIBI。
- 在增量阶段，如果原模型非alibi，且mask没有做padding，可以使用UNDEFINED格式的mask，即不传mask；否则依据原生模型的配置，决定是NORM还是ALIBI。
- 在长序列场景下，可以考虑压缩mask。

定义

struct SelfAttentionParam {
    enum CalcType : int {
        UNDEFINED = 0,
        ENCODER, // encoder for flashAttention
        DECODER,  // decoder for flashAttention
        PA_ENCODER // encoder for pagedAttention
    };
    enum KernelType : int {
        KERNELTYPE_DEFAULT = 0, // i:fp16, bmm:fp16, o:fp16
        KERNELTYPE_HIGH_PRECISION // i:fp16, bmm:fp32, o:fp16
    };
    enum ClampType : int {
        CLAMP_TYPE_UNDEFINED = 0,
        CLAMP_TYPE_MIN_MAX
    };
    enum MaskType : int {
        MASK_TYPE_UNDEFINED = 0,
        MASK_TYPE_NORM,
        MASK_TYPE_ALIBI,
        MASK_TYPE_NORM_COMPRESS,
        MASK_TYPE_ALIBI_COMPRESS,
        MASK_TYPE_ALIBI_COMPRESS_SQRT,
        MASK_TYPE_ALIBI_COMPRESS_LEFT_ALIGN
    };
    int32_t headNum = 0;
    int32_t kvHeadNum = 0;
    float qScale = 1;  // qtensor scale before qkbmm
    float qkScale = 1; // scale after qkbmm
    bool batchRunStatusEnable = false;
    uint32_t isTriuMask = 0;
    CalcType calcType = UNDEFINED;
    KernelType kernelType = KERNELTYPE_DEFAULT;
    ClampType clampType = CLAMP_TYPE_UNDEFINED;
    float clampMin = 0;
    float clampMax = 0;
    MaskType maskType = MASK_TYPE_UNDEFINED;
};

成员

成员名称	描述
CalcType	计算类型。 UNDEFINED：同时使用flashAttention的decoder和encoder。 ENCODER：使用FlashAttention的encoder。 DECODER：使用FlashAttention的decoder。 PA_ENCODER：使用PageAttention的encoder。
KernelType	算子内部计算类型。 KERNELTYPE_DEFAULT ：输入fp16，batch matmul使用fp16，输出fp16。 KERNELTYPE_HIGH_PRECISION ：输入fp16，batch matmul使用fp32，输出fp16。
ClampType	缩放类型。 CLAMP_TYPE_UNDEFINED：不进行缩放。 CLAMP_TYPE_MIN_MAX：进行缩放，同时指定最大最小值。
MaskType	AttentionMask类型。 MASK_TYPE_UNDEFINED：默认值，全0mask。 MASK_TYPE_NORM：倒三角mask。 MASK_TYPE_ALIBI：Alibi Mask。 MASK_TYPE_NORM_COMPRESS：倒三角压缩mask。 MASK_TYPE_ALIBI_COMPRESS：alibi压缩mask。 MASK_TYPE_ALIBI_COMPRESS_SQRT：alibi压缩开平方mask。 MASK_TYPE_ALIBI_COMPRESS_LEFT_ALIGN：alibi压缩mask左对齐。当前仅支持Atlas 800I A2推理产品。
headNum	query头大小。 headNum需大于或等于0。
kvHeadNum	kv头数量，该值需要用户根据使用的模型实际情况传入。 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值相同。
qScale	query缩放系数。
qkScale	算子tor值，在Q*K^T后乘。
batchRunStatusEnable	是否开启动态batch。
isTriuMask	是否开启倒三角优化，只有mask为倒三角时才能开启优化。如果为1，使用倒三角优化性能优化。如果为0，不使用倒三角优化性能优化。
calcType	计算类型。CalcType类型枚举值。
kernelType	内核精度类型。KernelType类型枚举值。
clampType	缩放类型。ClampType类型枚举值。
clampMin	缩放功能最小值。
clampMax	缩放功能最大值。
maskType	Mask类型，MaskType类型枚举值。

函数输入输出描述

当在Atlas 推理系列产品上运行时，cacheK，cacheV的格式为NZ格式，相应的维度为[layer, batch, hiddenSize/16, maxSeqLen, 16]，且maxSeqLen应与16对齐。

当在Atlas 推理系列产品上运行时，mask的格式可以为NZ格式，相应的维度为[batch, kvMaxSeqLen / 16, qMaxSeqLen, 16], [1, kvMaxSeqLen / 16, qMaxSeqLen, 16], [batch * head, kvMaxSeqLen / 16, qMaxSeqLen, 16], [head, kvMaxSeqLen / 16, qMaxSeqLen, 16]，且kvMaxSeqLen，qMaxSeqLen应与16对齐。

以上维度说明中，涉及除法的均为ceil div。

参数	维度	数据类型	格式	设备	描述
query	[nTokens, qHiddenSize]	float16/bfloat16	ND	npu	query矩阵。
key	[nTokens, hiddenSize]	float16/bfloat16	ND	npu	key矩阵。
value	[nTokens, hiddenSize]	float16/bfloat16	ND	npu	value矩阵。
cacheK	[layer, batch, maxSeqLen, hiddenSize] 当开启动态batch功能时，cacheK、cacheV的shape为[batch, maxSeqLen, hiddenSize]。	float16/bfloat16	ND/NZ	npu	存储之前所有的k，本次执行时将key刷新到cacheK上。
cacheV	[layer, batch, maxSeqLen, hiddenSize] 当开启动态batch功能时，cacheK、cacheV的shape为[batch, maxSeqLen, hiddenSize]。	float16/bfloat16	ND/NZ	npu	存储之前所有的v，本次执行时将value刷新到cacheV上。
attentionMask	[maxSeqLen, maxSeqLen] [batch, maxSeqLen, maxSeqLen] [batch, 1, maxSeqLen] [batch, headNum, maxSeqLen, maxSeqLen]	float16/bfloat16	ND/NZ	npu	所有batch相同，方阵。 batch不同时的方阵。 batch不同时的向量。 alibi场景。
tokenOffset	[batch]	int32/uint32	ND	cpu	计算完成后的token偏移。
seqLen	[batch]	int32/uint32	ND	cpu	等于1时，为增量或全量。大于1时，为全量。
layerId	[1]	int32/uint32	ND	npu	取cache的kv中哪一个kv进行计算。
batchStatus	[batch]	int32/uint32	ND	cpu	开启动态batch功能时，通过标志位控制具体需要运算的batch。
out	[nTokens, qHiddenSize]	float16/bfloat16	ND	npu	输出。

函数输入输出描述（calcType = PA_ENCODER）

参数	维度	数据类型	格式	设备	描述
query	[nTokens, head_num, head_size]	float16/bfloat16	ND	npu	query矩阵。
key	[nTokens, head_num, head_size]	float16/bfloat16	ND	npu	key矩阵。
value	[nTokens, head_num, head_size]	float16/bfloat16	ND	npu	value矩阵。
mask	与FlashAttention相同。	float16/bfloat16	ND/NZ	npu	与FlashAttention相同，当maskType为undefined时不传此tensor。
seqLen	[batch]	int32	ND	cpu	等于1时，为增量或全量。大于1时，为全量。
slopes	[head_num]	Atlas 800I A2推理产品：float16 Atlas 推理系列产品：float	ND	npu	当maskType为alibi压缩，且mask为[256,256]时需传入此tensor，为alibi mask每个head的系数。
output	[nTokens, head_num, head_size]	float16/bfloat16	ND	npu	输出。

父主题： atb/infer_op_params.h