FlashAttentionScore
算子基础信息
算子名称 |
FlashAttentionScore |
---|---|
torch_npu api接口 |
torch_npu.npu_fusion_attention |
支持的torch_npu版本 |
1.11.0, 2.1.0, 2.2.0 |
支持的芯片类型 |
Atlas A2 训练系列产品 |
支持的数据类型 |
float16, bfloat16 |
算子名称 |
FlashAttentionScore |
---|---|
torch_npu api接口 |
torch.nn.functional.scaled_dot_product_attention |
支持的torch_npu版本 |
2.1.0, 2.2.0 |
支持的芯片类型 |
Atlas A2 训练系列产品 |
支持的数据类型 |
float16, bfloat16 |
算子IR及torch_npu接口参数
REG_OP(FlashAttentionScore) .INPUT(query, TensorType({DT_FLOAT16, DT_BF16})) .INPUT(key, TensorType({DT_FLOAT16, DT_BF16})) .INPUT(value, TensorType({DT_FLOAT16, DT_BF16})) .OPTIONAL_INPUT(real_shift, TensorType({DT_FLOAT16, DT_BF16})) .OPTIONAL_INPUT(drop_mask, TensorType({DT_UINT8})) .OPTIONAL_INPUT(padding_mask, TensorType({DT_FLOAT16, DT_BF16})) .OPTIONAL_INPUT(atten_mask, TensorType({DT_BOOL, DT_UINT8})) .OPTIONAL_INPUT(prefix, TensorType({DT_INT64})) .OPTIONAL_INPUT(actual_seq_qlen, TensorType({DT_INT64})) .OPTIONAL_INPUT(actual_seq_kvlen, TensorType({DT_INT64})) .OUTPUT(softmax_max, TensorType({DT_FLOAT32})) .OUTPUT(softmax_sum, TensorType({DT_FLOAT32})) .OUTPUT(softmax_out, TensorType({DT_FLOAT16, DT_BF16})) .OUTPUT(attention_out, TensorType({DT_FLOAT16, DT_BF16})) .ATTR(scale_value, Float, 1.0) .ATTR(keep_prob, Float, 1.0) .ATTR(pre_tockens, Int, 2147483647) .ATTR(next_tockens, Int, 2147483647) .REQUIRED_ATTR(head_num, Int) .REQUIRED_ATTR(input_layout, String) .ATTR(inner_precise, Int, 0) .ATTR(sparse_mode, Int, 0) .OP_END_FACTORY_REG(FlashAttentionScore)
torch_npu.npu_fusion_attention(Tensor query, Tensor key, Tensor value, int head_num, str input_layout, Tensor? pse=None, Tensor? padding_mask=None, Tensor? atten_mask=None, float scale=1., float keep_prob=1., int pre_tockens=2147483647, int next_tockens=2147483647, int inner_precise=0, int[]? prefix=None, int[]? actual_seq_qlen=None, int[]? actual_seq_kvlen=None, int sparse_mode=0, bool gen_mask_parallel=True, bool sync=False) -> (Tensor, Tensor, Tensor, Tensor, int, int, int)
torch_npu接口中的问号表示这个输入参数是可选的。
实现“Transformer Attention Score”的融合计算,实现的计算公式如下:
- query:Device侧的Tensor,数据类型支持FLOAT16、BFLOAT16,数据格式支持ND。综合约束请见约束说明。
- key:Device侧的Tensor,数据类型支持FLOAT16、BFLOAT16,数据格式支持ND。综合约束请见约束说明。
- value:Device侧的Tensor,数据类型支持FLOAT16、BFLOAT16,数据格式支持ND。综合约束请见约束说明。
- head_num:Host侧的int,代表head个数,数据类型支持INT64。综合约束请见约束说明。
- input_layout:Host侧的string,代表输入query、key、value的数据排布格式,支持BSH、SBH、BSND、BNSD、TND(actual_seq_qlen/actual_seq_kvlen需传值);后续章节如无特殊说明,S表示query或key、value的sequence length,Sq表示query的sequence length,Skv表示key、value的sequence length,SS表示Sq*Skv。
- pse:Device侧的Tensor,可选参数,表示位置编码。数据类型支持FLOAT16、BFLOAT16,数据格式支持ND。非varlen场景支持四维输入,包含BNSS格式、BN1Skv格式、1NSS格式。如果非varlen场景Sq大于1024或varlen场景、每个batch的Sq与Skv等长且是sparse_mode为0、2、3的下三角掩码场景,可使能alibi位置编码压缩,此时只需要输入原始PSE最后1024行进行内存优化,即alibi_compress = ori_pse[:, :, -1024:, :],参数每个batch不相同时,输入BNHSkv(H=1024),每个batch相同时,输入1NHSkv(H=1024)。
- padding_mask:Device侧的Tensor,暂不支持该参数。
- atten_mask:Device侧的Tensor,可选参数,取值为1代表该位不参与计算(不生效),为0代表该位参与计算,数据类型支持BOOL、UINT8,数据格式支持ND格式,输入shape类型支持BNSS格式、B1SS格式、11SS格式、SS格式。varlen场景只支持SS格式,SS分别是maxSq和maxSkv。综合约束请见约束说明。
- scale:Host侧的double,可选参数,代表缩放系数,作为计算流中Muls的scalar值,数据类型支持DOUBLE,默认值为1。
- keep_prob:Host侧的double,可选参数,代表Dropout中1的比例,数据类型支持DOUBLE,默认值为1,表示全部保留。
- pre_tockens:Host侧的int,用于稀疏计算的参数,可选参数,数据类型支持INT64,默认值为2147483647。综合约束请见约束说明。
- next_tockens:Host侧的int,用于稀疏计算的参数,可选参数,数据类型支持INT64,默认值为2147483647。next_tockens和pre_tockens取值与atten_mask的关系请参见sparse_mode参数,参数取值与atten_mask分布不一致会导致精度问题。综合约束请见约束说明。
- inner_precise:Host侧的int,用于提升精度,数据类型支持INT64,默认值为0。
说明:当前0、1为保留配置值,2为使能无效行计算,其功能是避免在计算过程中存在整行mask进而导致精度有损失,但是该配置会导致性能下降。
如果算子可判断出存在无效行场景,会自动使能无效行计算,例如sparse_mode为3,Sq > Skv场景。
- prefix:Host侧的int array,可选参数,代表prefix稀疏计算场景每个Batch的N值。数据类型支持INT64,数据格式支持ND。综合约束请见约束说明。
- actual_seq_qlen:Host侧的int array,可选参数,varlen场景时需要传入此参数。表示query每个S的累加和长度,数据类型支持INT64,数据格式支持ND。综合约束请见约束说明。
比如真正的S长度列表为:2 2 2 2 2
则actual_seq_qlen传:2 4 6 8 10
- actual_seq_kvlen:Host侧的int array,可选参数,varlen场景时需要传入此参数。表示key/value每个S的累加和长度。数据类型支持INT64,数据格式支持ND。综合约束请见约束说明。
比如真正的S长度列表为:2 2 2 2 2
则actual_seq_kvlen传:2 4 6 8 10
- sparse_mode:Host侧的int,表示sparse的模式,可选参数。数据类型支持:INT64,默认值为0,支持配置值为0、1、2、3、4、5、6、7、8。当整网的atten_mask都相同且shape小于2048*2048时,建议使用defaultMask模式,来减少内存使用量。综合约束请见约束说明。
表3 sparse_mode不同取值场景说明 sparse_mode
含义
备注
0
defaultMask模式。
-
1
allMask模式。
-
2
leftUpCausal模式。
-
3
rightDownCausal模式。
-
4
band模式。
-
5
prefix非压缩模式。
varlen场景不支持。
6
prefix压缩模式。
-
7
varlen外切场景,rightDownCausal模式。
仅varlen场景支持。
8
varlen外切场景,leftUpCausal模式。
仅varlen场景支持。
atten_mask的工作原理为,在Mask为True的位置遮蔽query(Q)与key(K)的转置矩阵乘积的值,示意如下:
QKT矩阵在atten_mask为True的位置会被遮蔽,效果如下:
说明:下图中的蓝色表示保留该值,atten_mask中,应该配置为False;阴影表示遮蔽该值,atten_mask中应配置为True。- sparse_mode为0时,代表defaultMask模式。
- 不传mask:如果atten_mask未传入则不做mask操作,atten_mask取值为None,忽略pre_tockens和next_tockens取值。Masked QKT矩阵示意如下:
- next_tockens取值为0,pre_tockens大于等于Sq,表示causal场景sparse,atten_mask应传入下三角矩阵,此时pre_tockens和next_tockens之间的部分需要计算,Masked QKT矩阵示意如下:
atten_mask应传入下三角矩阵,示意如下:
- pre_tockens小于Sq,next_tockens小于Skv,且都大于等于0,表示band场景,此时pre_tockens和next_tockens之间的部分需要计算。Masked QKT矩阵示意如下:
atten_mask应传入band形状矩阵,示意如下:
- next_tockens为负数,以pre_tockens=9,next_tockens=-3为例,pre_tockens和next_tockens之间的部分需要计算。Masked QKT示意如下:
说明:next_tockens为负数时,pre_tockens取值必须大于等于next_tockens的绝对值,且next_tockens的绝对值小于Skv。
- pre_tockens为负数,以next_tockens=7,pre_tockens=-3为例,pre_tockens和next_tockens之间的部分需要计算。Masked QKT示意如下:
说明:pre_tockens为负数时,next_tockens取值必须大于等于pre_tockens的绝对值,且pre_tockens的绝对值小于Sq。
- 不传mask:如果atten_mask未传入则不做mask操作,atten_mask取值为None,忽略pre_tockens和next_tockens取值。Masked QKT矩阵示意如下:
- sparse_mode为1时,代表allMask,即传入完整的atten_mask矩阵。
该场景下忽略next_tockens、pre_tockens取值,Masked QKT矩阵示意如下:
- sparse_mode为2时,代表leftUpCausal模式的mask,对应以左上顶点划分的下三角场景(参数起点为左上角)。该场景下忽略pre_tockens、next_tockens取值,Masked QKT矩阵示意如下:
传入的atten_mask为优化后的压缩下三角矩阵(2048*2048),压缩下三角矩阵示意(下同):
- sparse_mode为3时,代表rightDownCausal模式的mask,对应以右下顶点划分的下三角场景(参数起点为右下角)。该场景下忽略pre_tockens、next_tockens取值。atten_mask为优化后的压缩下三角矩阵(2048*2048),Masked QKT矩阵示意如下:
- sparse_mode为4时,代表band场景,即计算pre_tockens和next_tockens之间的部分,参数起点为右下角,pre_tockens和next_tockens之间需要有交集。atten_mask为优化后的压缩下三角矩阵(2048*2048)。Masked QKT矩阵示意如下:
- sparse_mode为5时,代表prefix非压缩场景,即在rightDownCasual的基础上,左侧加上一个长为Sq,宽为N的矩阵,N的值由可选输入prefix获取,例如下图中表示batch=2场景下prefix传入数组[4,5],每个batch轴的N值可以不一样,参数起点为左上角。
该场景下忽略pre_tockens、next_tockens取值,atten_mask矩阵数据格式须为BNSS或B1SS,Masked QKT矩阵示意如下:
atten_mask应传入矩阵示意如下:
- sparse_mode为6时,代表prefix压缩场景,即prefix场景时,attenMask为优化后的压缩下三角+矩形的矩阵(3072*2048):其中上半部分[2048,2048]的下三角矩阵,下半部分为[1024,2048]的矩形矩阵,矩形矩阵左半部分全0,右半部分全1,atten_mask应传入矩阵示意如下。该场景下忽略pre_tockens、next_tockens取值。
- sparse_mode为7时,表示varlen且为长序列外切场景(即长序列在模型脚本中进行多卡切query的sequence length);用户需要确保外切前为使用sparse_mode 3的场景;当前mode下用户需要设置pre_tockens和next_tockens(起点为右下顶点),且需要保证参数正确,否则会存在精度问题。
Masked QKT矩阵示意如下,在第二个batch对query进行切分,key和value不切分,4x6的mask矩阵被切分成2x6和2x6的mask,分别在卡1和卡2上计算:
- 卡1的最后一块mask为band类型的mask,配置pre_tockens=6(保证大于等于最后一个Skv),next_tockens=-2,actual_seq_qlen应传入{3,5},actual_seq_kvlen应传入{3,9}。
- 卡2的mask类型切分后不变,sparse_mode为3,actual_seq_qlen应传入{2,7,11},actual_seq_kvlen应传入{6,11,15}。
- 如果配置sparse_mode=7,但实际只存在一个batch,用户需按照band模式的要求来配置参数;sparse_mode=7时,用户需要输入2048x2048的下三角mask作为该融合算子的输入。
- 基于sparse_mode=3进行外切产生的band模式的sparse的参数应符合以下条件:
- pre_tockens >= last_Skv。
- next_tockens <= 0。
- 当前模式下不支持可选输入pse。
- sparse_mode为8时,表示varlen且为长序列外切场景;用户需要确保外切前为使用sparse_mode 2的场景;当前mode下用户需要设置pre_tockens和next_tockens(起点为右下顶点),且需要保证参数正确,否则会存在精度问题。
Masked QKT矩阵示意如下,在第二个batch对query进行切分,key和value不切分,5x4的mask矩阵被切分成2x4和3x4的mask,分别在卡1和卡2上计算:
- 卡1的mask类型切分后不变,sparse_mode为2,actual_seq_qlen应传入{3,5},actual_seq_kvlen应传入{3,7}。
- 卡2的第一块mask为band类型的mask,配置pre_tockens=4(保证大于等于第一个Skv),next_tockens=1,actual_seq_qlen应传入{3,8,12},actual_seq_kvlen应传入{4,9,13}。
- 如果配置sparse_mode=8,但实际只存在一个batch,用户需按照band模式的要求来配置参数;sparse_mode=8时,用户需要输入2048x2048的下三角mask作为该融合算子的输入。
- 基于sparse_mode=2进行外切产生的band模式的sparse的参数应符合以下条件:
- pre_tockens >= first_Skv。
- next_tockens范围无约束,根据实际情况进行配置。
- 当前模式下不支持可选输入pse。
- sparse_mode为0时,代表defaultMask模式。
- gen_mask_parallel:debug参数,DSA生成dropout随机数向量mask的控制开关,默认值为True:同AICORE计算并行,False:同AICORE计算串行。
- sync:debug参数,DSA生成dropout随机数向量mask的控制开关,默认值为False:dropout mask异步生成,True:dropout mask同步生成。
输出说明:
共7个输出
- 第1个输出为Tensor,计算公式的最终输出y,数据类型支持:FLOAT16、BFLOAT16。
- 第2个输出为Tensor,Softmax 计算的Max中间结果,用于反向计算,数据类型支持:FLOAT。
- 第3个输出为Tensor,Softmax计算的Sum中间结果,用于反向计算,数据类型支持:FLOAT。
- 第4个输出为Tensor,保留参数,暂未使用。
- 第5个输出为int,DSA生成dropoutmask中,Philox算法的seed。
- 第6个输出为int,DSA生成dropoutmask中,Philox算法的offset。
- 第7个输出为int,DSA生成dropoutmask的长度。
约束说明:
- 输入query、key、value的B:batchsize必须相等,取值范围1~2K。
- 输入query、key、value、pse的数据类型必须一致。
- 输入query、key、value的input_layout必须一致。
- 支持输入query的N和key/value的N不相等,但必须成比例关系,即Nq/Nkv必须是非0整数,Nq取值范围1~256。当Nq/Nkv > 1时,即为GQA(grouped-query attention);当Nkv=1时,即为MQA(multi-query attention)。本文如无特殊说明,N表示的是Nq。
- 输入key/value的shape必须一致。
- 输入query、key、value的S:sequence length,取值范围1~512K。
- 输入query、key、value的D:head dim,取值范围1~512。
- varlen场景T(B*S)取值范围1~512K。
- keep_prob的取值范围为(0, 1] 。
- sparse_mode为1、2、3、4、5、6、7、8时,应传入对应正确的atten_mask,否则将导致计算结果错误。当atten_mask输入为None时,sparse_mode,pre_tockens,next_tockens参数不生效,固定为全计算。
- sparse_mode配置为1、2、3、5、6时,用户配置的pre_tockens、next_tockens不会生效。
- sparse_mode配置为0、4时,须保证atten_mask与pre_tockens、next_tockens的范围一致。
- prefix稀疏计算场景B不大于32,varlen场景不支持非压缩prefix,即不支持sparse_mode=5;当Sq>Skv时,prefix的N值取值范围[0, Skv],当Sq<=Skv时,prefix的N值取值范围[Skv-Sq, Skv]。
- sparse_mode=7或者8时,不支持可选输入pse。
- varlen场景:
- atten_mask输入不支持补pad,即atten_mask中不能存在某一行全1的场景。
- pse仅支持alibi场景,即shape为BNHSkv(H=1024)、1NHSkv(H=1024)两种场景。
- 支持actual_seq_qlen中某个Batch上的S长度为0;如果存在S为0的情况,不支持pse输入, 假设真实的S长度为[2,2,0,2,2],则传入的actual_seq_qlen为[2,4,4,6,8];不支持某个batch中Sq不为0,但是Skv为0的场景。
模型中替换代码
当前GPU模式下,调用FA算子的方式有多种,torch调用FA的接口scaled_dot_product_attention,通过flash-attention库中的flash_attn_func、flash_attn_varlen_func等接口调用。NPU模式下除了已经适配的sdpa接口,其余模式需要通过torch_npu接口——npu_fusion_attention接口实现调用。两者之间的适配可能涉及一些脚本迁移工作,以下通过范例说明接口适配方式。
torch原生接口
scaled_dot_product_attention:
- 不使能is_causal时,原调用接口代码:
res = torch.nn.functional.scaled_dot_product_attention(query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False)
替换为:attn_mask_npu = torch.logical_not(attention_mask.bool()).to(device) head_num = query.shape[1] res = torch_npu.npu_fusion_attention( query, key, value, head_num, input_layout="BNSD", pse=None, atten_mask=attn_mask_npu, scale=1.0 / math.sqrt(query.shape[-1]), pre_tockens=2147483647, next_tockens=2147483647, keep_prob=1, inner_precise=0 )[0]
- 使能is_causal时,原调用接口代码:
res = torch.nn.functional.scaled_dot_product_attention(query, key, value, attn_mask=None, dropout_p=0.0, is_causal=True)
替换为:
atten_mask = torch.triu(torch.ones([2048, 2048]), diagonal=1) head_num = query.shape[1] res = torch_npu.npu_fusion_attention( query, key, value, head_num, input_layout="BNSD", pse=None, atten_mask=attn_mask_npu, scale=1.0 / math.sqrt(query.shape[-1]), pre_tockens=2147483647, next_tockens=2147483647, keep_prob=1, inner_precise=0, sparse_mode=2 )[0]
- flash_attn_func
替换为:
head_num = q.shape[2] out = torch_npu.npu_fusion_attention(q, k, v, head_num, "BSND", keep_prob=1.0, scale=scale)
- flash_attn_varlen_func
不使能causal时,原调用接口代码:
out = flash_attn_varlen_func( q, k, v, cu_seqlens_q, cu_seqlens_k, max_seqlen_q, max_seqlen_k, dropout_p=0.0, softmax_scale=None, causal=False )
替换为:head_num = q.shape[1] output = torch_npu.npu_fusion_attention( q, k, v, head_num, pse=None, padding_mask=None, atten_mask=None, scale=1.0 / math.sqrt(q.shape[-1]), keep_prob=1, input_layout="TND", actual_seq_qlen=tuple(cu_seqlens_q[1:].cpu().numpy().tolist()), actual_seq_kvlen=tuple(cu_seqlens_k[1:].cpu().numpy().tolist()), pre_tockens=2147483647, next_tockens=2147483647, inner_precise=inner_precise, sparse_mode=0)[0]
使能causal时,GPU调用接口代码:out = flash_attn_varlen_func( q, k, v, cu_seqlens_q, cu_seqlens_k, max_seqlen_q, max_seqlen_k, dropout_p=0.0, softmax_scale=None, causal=True )
替换为:atten_mask_npu= torch.from_numpy(np.triu(np.ones([max_seqlen_q, max_seqlen_k]), k=1)).bool().to(device) head_num = q.shape[1] output = torch_npu.npu_fusion_attention( q, k, v, head_num, pse=None, padding_mask=None, atten_mask=atten_mask_npu, scale=1.0 / math.sqrt(q.shape[-1]), keep_prob=1, input_layout="TND", actual_seq_qlen=tuple(cu_seqlens_q[1:].cpu().numpy().tolist()), actual_seq_kvlen=tuple(cu_seqlens_k[1:].cpu().numpy().tolist()), pre_tockens=2147483647, next_tockens=0, inner_precise=inner_precise)[0]
xFormers库
memory_efficient_attention
def __init__(self, attention_op: Optional[Callable] = None): self.attention_op = attention_op def head_to_batch_dim(self, tensor): tensor = tensor.reshape(batch_size * head_size, seq_len, dim // head_size) return tensor ... query = attn.head_to_batch_dim(query).contiguous() key = attn.head_to_batch_dim(key).contiguous() value = attn.head_to_batch_dim(value).contiguous() hidden_states = xformers.ops.memory_efficient_attention( query, key, value, attn_bias=attention_mask, op=self.attention_op, scale=attn.scale )
def head_to_batch_dim(self, tensor, out_dim=3): head_size = self.heads batch_size, seq_len, dim = tensor.shape tensor = torch_npu.npu_confusion_transpose(tensor, [0, 2, 1, 3], (batch_size, seq_len, head_size, dim // head_size), False) if out_dim == 3: tensor = tensor.reshape(batch_size * head_size, seq_len, dim // head_size) return tensor ... query = attn.head_to_batch_dim(query, out_dim=4) key = attn.head_to_batch_dim(key, out_dim=4) value = attn.head_to_batch_dim(value, out_dim=4) hidden_states = torch_npu.npu_fusion_attention( query, key, value, heads, input_layout="BNSD", pse=None, atten_mask=attention_mask, scale=scale, pre_tockens=2147483647, next_tockens=2147483647, keep_prob=1., sync=False, inner_precise=0, )[0]
算子计算逻辑
def tforward(q, k, v, drop_mask, atten_mask, pse, scale, keep_prob): if pse is None or len(pse.shape) == 0: qk = torch.matmul(q, k.permute(0, 1, 3, 2)).mul(scale) else: qk = (torch.matmul(q, k.permute(0, 1, 3, 2)) + pse).mul(scale) if atten_mask is None or len(atten_mask.shape) == 0: qk = qk else: qk = qk + atten_mask * torch.finfo(torch.float32).min softmax_res, x_max, x_sum = tsoftmax(qk) if drop_mask is None or len(drop_mask.shape) == 0: drop_res = softmax_res else: drop_res = softmax_res * drop_mask * (1.0 / (keep_prob)) y = torch.matmul(drop_res, v) return y, softmax_res
已支持模型典型case
如下case均包含fp16、bf16。
id |
BNSD |
---|---|
1 |
[1, 8, 4096, 4096] |
2 |
[4, 32, 2048, 2048] |
3 |
[8, 16, 512, 2048] |
4 |
[8, 16, 512, 1024] |
5 |
[8, 16, 512, 1536] |
6 |
[8, 16, 512, 512] |
7 |
[8, 16, 512, 256] |
8 |
[4, 4, 2048, 2048] |
torch原生接口调用FA算子使用限制
接口和参数说明:
使用限制:
- 参数输入符合规格:
- 输入query、key、value的N:batch size,当前只支持[N,head_num, S(L), E(Ev)]的排布方式,取值范围1~2K
- 输入query的head num和key/value的head num必须成比例关系,即Nq/Nkv必须是非0整数,取值范围1~256
- 输入query的L:Target sequence length,取值范围1~512K
- 输入key、value的S:Source sequence length,取值范围1~512K
- 输入query、key、value的E:Embedding dimension of the query and key,取值范围1~512
- 输入value的Ev:Embedding dimension of the value,必须与E相等
- 输入attn_mask:当前支持[N, 1, L, S]、[N, head_num, L, S]、[1, 1, L, S]、[L, S]共4种排布方式
- 在使能is_causal计算时,attn_mask必须为None;不使能is_causal时,若attn_mask输入有效数据,输入数据类型必须是Bool类型
- 与原接口除了规格限制之外存在差异点
- NPU用DSA硬件实现,算法在DSA引擎固化存在跟GPU算法实现差异,导致dropout功能和GPU结果不一致
- 当前接口支持输入query的head num和key/value的head num不等长,而原生PyTorch接口不支持
- 输入query、key、value的数据类型fp16、bf16并且使能requires_grad时,执行FA算子
- 输入query、key、value的数据类型fp16、bf16并且使能requires_grad时,若参数超出规格要求范围,接口会报错