CANN——FA算子

融合算子是指将多个独立的"小算子"融合成一个"大算子",二者功能等价,但性能内存上"大算子"明显优于"小算子"。 FA(Flash Attention)是大语言模型(Large LanguageModel, LLM)中最核心的一个融合算子。

融合算子的优势:

• 减少计算量:简化计算过程,减少计算量,提高计算效率。

• 减少内存占用:合并中间结果,提高内存利用率。

• 优化数据流:减少数据传输,提高数据处理效率。

• 简化代码实现:减少代码量,提高可读可维性。

功能描述

正向计算公式:

attention_out=Dropout(Softmax(Mask(scale*(pse+query*key'),atte"en_mask)),keep_prob)*value

接口原型

torch_npu.npu_fusion_attention(
	const at::Tensor &query, const at::Tensor &key, const at::Tensor 8&value, int64_t head_num,c10::string_view input_layout,
	const c10::optional<at::Tensor> &pse_opt, const c10:optional<at::Tensor> &padding_mask_opt,const
	c10::optional<at::Tensor>&atten_mask_opt,
	double scale,double keep_prob, int64_t pre_tockens, int64_t next_tockens, t inner_precise
	at::OptionallntArrayRef prefix_opt, at::OptionallntArrayRef actual_seq_qlen, at::OptionallntArrayRef actual_seq_kvlen,
	int64_t sparse_mode, bool gen_mask_parallel, bool sync
)

共7个输出:(Tensor,Tensor,Tensor,Tensor,int,int,int):

1. 第1个输出为Tensor,计算公式的最终输出attention_out

2. 第2个输出为Tensor,Softmax计算的Max中间结果,用于反向计算

3. 第3个输出为Tensor,Softmax计算的Sum中间结果,用于反向计算

4. 第4个输出为Tensor,保留参数,暂未使用。

5. 第5个输出为int,DSA生成dropoutmask中,Philox算法的seed。

6. 第6个输出为int,DSA生成dropoutmask中,Philox算法的offset。

7. 第7个输出为int,DSA生成dropoutmask的长度。