FlashAttentionScore

1
2

res = torch.nn.functional.scaled_dot_product_attention(query, key, value, atten_mask=attention_mask,
                                     dropout_p=0.0, is_causal=False)

替换为：

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14

if atten_mask.dtype == torch.bool:
    atten_mask_npu = torch.logical_not(attention_mask.bool()).to(device) // atten_mask需要取反
else:
    atten_mask_npu = 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=atten_mask_npu,
                       scale=1.0 / math.sqrt(query.shape[-1]),
                       pre_tockens=2147483647,
                       next_tockens=2147483647,
                       keep_prob=1
                   )[0]

1
2

res = torch.nn.functional.scaled_dot_product_attention(query, key, value, atten_mask=None,
                                     dropout_p=0.0, is_causal=True)

替换为：

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11

atten_mask_npu= torch.triu(torch.ones([2048, 2048]), diagonal=1).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=atten_mask_npu,
                       scale=1.0 / math.sqrt(query.shape[-1]),
                       pre_tockens=2147483647,
                       next_tockens=2147483647,
                       keep_prob=1,
                       sparse_mode=2 )[0]

• 接口参数对应表格：

  表3 接口参数替换

  gpu参数名称	npu参数名称	说明
  q	q	query；gpu、npu参数名一致且含义一致。
  k	k	key；gpu、npu参数名一致且含义一致 。
  v	v	value；gpu、npu参数名一致且含义一致。
  dropout	keep_prob	keep_prob= 1 - dropout；gpu接口中dropout代表数据需要忽略的概率，npu接口中keep_prob代表数据需要保留的概率。
  softmax_scale	scale	对QKT 的缩放系数；gpu、npu参数名不一致，含义一致；npu接口参数的默认值为1.0，适配时查阅原实现的scale值。
  causal	atten_mask	gpu接口参数causal=true时，npu接口需要传入下三角形式的atten mask；gpu接口参数causal=false时，npu接口不需要传入atten mask。
  -	head_num	npu接口新增参数，表示query的头数。
  -	layout	npu接口新增参数，表示qkv的layout，gpu接口默认layout为" BSND"。
  -	sparse_mode	npu接口新增参数，表示稀疏计算模式。sparse_mode=2表示leftUp causal，sparse_mode=3表示rightDown causal；gpu接口在FA2.0版本及之前，causal场景默认是leftUp，在2.0版本之后，默认是rightDown。

• 接口参数替换实例：
  不使能causal时，模型中替换代码：

  1	out= flash_attn_func(q, k, v, dropout_p=0.0, softmax_scale=scale, causal=False)

  替换为：

  1 2 3

  head_num = q.shape[2] out = torch_npu.npu_fusion_attention(q, k, v, head_num, "BSND", keep_prob=1.0, scale=scale)[0]

  使能causal时，模型中替换代码：

  1	out= flash_attn_func(q, k, v, dropout_p=0.0, softmax_scale=scale, causal=True)

  替换为：

  1 2 3 4

  atten_mask_npu = torch.triu(torch.ones([2048, 2048]), diagonal=1).bool().to(device) head_num = q.shape[2] out = torch_npu.npu_fusion_attention(q, k, v, head_num, "BSND", keep_prob=1.0, scale=scale, atten_mask=atten_mask_npu, sparse_mode=3)[0]

  

  当替换flash-attention为2.0或之前版本时，应设置sparse_mode=2；当替换2.1或之后版本时，应设置sparse_mode=3。

• 接口参数对应表格：

  表4 接口参数替换

  gpu参数名称	npu参数名称	说明
  q	q	query；gpu、npu参数名一致且含义一致。
  k	k	key；gpu、npu参数名一致且含义一致 。
  v	v	value；gpu、npu参数名一致且含义一致。
  dropout	keep_prob	keep_prob= 1 - dropout；gpu接口中dropout代表数据需要忽略的概率，npu接口中keep_prob代表数据需要保留的概率。
  softmax_scale	scale	对QKT 的缩放系数；gpu、npu参数名不一致，含义一致；npu接口参数的默认值为1.0，适配时查阅原实现的scale值。
  causal	atten_mask	gpu接口参数causal=true时，npu接口需要传入下三角形式的atten mask；gpu接口参数causal=false时，npu接口不需要传入atten mask。
  cu_seqlens_q	actual_seq_qlen	query 序列的累积长度；gpu、npu参数名不一致，含义一致；npu接口该参数需要转换为host侧的list格式。
  cu_seqlens_k	ctual_seq_kvlen	key、value序列的累积长度 ；gpu、npu参数名不一致，含义一致；npu接口该参数需要转换为host侧的list格式。
  max_seqlen_q	-	npu无需配置该参数； gpu接口中表示query序列的最大长度，npu在接口内部计算。
  max_seqlen_k	-	npu无需配置该参数； gpu接口中表示key、value序列的最大长度，npu在接口内部计算。
  -	head_num	npu接口新增参数，表示query的头数。
  -	layout	npu接口新增参数，表示qkv的layout，gpu接口默认layout为" BSND"。
  -	sparse_mode	npu接口新增参数，表示稀疏计算模式。sparse_mode=2表示leftUp causal，sparse_mode=3表示rightDown causal；gpu接口在FA2.0版本及之前，causal场景默认是leftUp，在2.0版本之后，默认是rightDown。

• 接口参数替换实例：

  不使能causal时，原调用接口代码：

  1 2 3 4 5 6 7 8 9 10 11 12

  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 )

  替换为：

  1 2 3 4 5 6 7 8 9 10

  head_num = q.shape[1] output = torch_npu.npu_fusion_attention( q, k, v, head_num, pse=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()))[0]

  使能causal时，GPU调用接口代码：

  1 2 3 4 5 6 7 8 9 10 11 12

  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 )

  替换为：

  1 2 3 4 5 6 7 8 9 10 11 12 13

  atten_mask_npu = torch.triu(torch.ones([2048, 2048]), diagonal=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()), sparse_mode=3)[0]

  

  当替换flash-attention为2.0或之前版本时，应设置sparse_mode=2；当替换2.1或之后版本时，应设置sparse_mode=3。