​作者​昇腾实战派​知识地图​https://blog.csdn.net/Lumos_Lovegood/article/details/161601003背景概述在昇腾AI处理器的模型推理与训练场景中FlashAttention算子npu_fusion_attention是Transformer类模型的核心组件。该算子同时封装了前向与反向计算逻辑但在某些业务场景下我们需要将前向与反向逻辑解耦以便更灵活地控制计算流程。此外针对长序列场景如序列长度超过16K需要引入滑动窗口注意力Sliding Window Attention机制以降低计算复杂度。本文以Atlas 800I A2设备为例基于Ascend op-plugin仓库详细介绍了如何将FlashAttention算子拆分为独立的前向与反向函数并实现支持稀疏模式SparseMode为0和4的滑动窗口注意力算子。1. 需求分析1.1 代码梳理原始FlashAttention算子的实现位于以下文件中https://gitcode.com/Ascend/op-plugin/tree/7.3.0/op_plugin/ops/opapi/FlashAttentionKernelNpuOpApi.cpp该文件包含三个核心函数其作用如下表所示符号作用npu_fusion_attention前向函数内部调用aclnnFlashAttentionScore或aclnnFlashAttentionVarLenScore返回(attention_score, softmax_max, softmax_sum, softmax_out, seed, offset, numels)npu_fusion_attention_grad“高层”反向函数根据seed/offset/numels等参数生成dropout mask再调用底层反向函数npu_fusion_attention_backward“底层”反向函数调用aclnnFlashAttentionScoreGrad/UnpaddingScoreGrad注意npu_fusion_attention_grad末尾会调用npu_fusion_attention_backward。代码文件中存在三对上述函数通过#if VERSION_BETWEEN宏区分不同版本。1.2 目标实现路径我们需要将torch_npu.npu_fusion_attention()算子入口拆分为独立的前向和反向函数具体步骤如下克隆op-plugin代码仓库新建代码文件SlideWindowsAttentionKernelNpuOpApi.cpp添加swa_forward和swa_backward函数将SparseMode限制为0和4两种模式具体含义参考后表在op_plugin/config/op_plugin_functions.yaml中添加对应的函数声明编写CI测试用例验证函数功能1.3 SparseMode说明SparseMode枚举定义如下enumclassSparseMode{NO_MASK0,ALL_MASK,LEFT_UP_CAUSAL,RIGHT_DOWN_CAUSAL,BAND,PREFIX,PREFIX_COMPRESS,RIGHT_DOWN_CAUSAL_BAND,BAND_LEFT_UP_CAUSAL};其中模式0表示无mask模式4BAND表示使用pre_tockens与next_tockens控制窗口的滑动窗口mask。本实现仅需支持这两种模式无需关注底层AscendC算子的具体实现。2. 编译安装这个需求需要编译 op-plugin 这个库。先拉取CANN镜像dockerpull swr.cn-south-1.myhuaweicloud.com/ascendhub/cann:8.5.2-910b-ubuntu22.04-py3.11启动容器NAMEcann8.5-testDEVICES0,1,2,3,4,5,6,7IMAGEswr.cn-south-1.myhuaweicloud.com/ascendhub/cann:8.5.2-910b-ubuntu22.04-py3.11dockerrun-itd-u0--ipchost--privileged\-eVLLM_USE_MODELSCOPETrue-ePYTORCH_NPU_ALLOC_CONFmax_split_size_mb:256\-eASCEND_RT_VISIBLE_DEVICES$DEVICES\--name$NAME\--nethost\--device/dev/davinci_manager\--device/dev/devmm_svm\--device/dev/hisi_hdc\--shm-size1200g\-v/usr/local/dcmi:/usr/local/dcmi\-v/usr/local/Ascend/driver/tools/hccn_tool:/usr/local/Ascend/driver/tools/hccn_tool\-v/usr/local/bin/npu-smi:/usr/local/bin/npu-smi\-v/usr/local/bin/ais_bench:/usr/local/bin/ais_bench\-v/usr/local/Ascend/driver/lib64/:/usr/local/Ascend/driver/lib64/\-v/usr/local/Ascend/driver/version.info:/usr/local/Ascend/driver/version.info\-v/etc/ascend_install.info:/etc/ascend_install.info\-v/etc/hccn.conf:/etc/hccn.conf\-v/root/.cache:/root/.cache\-v/home/m30071636/:/home/m30071636/\--privilegedtrue\-it$IMAGEbash下载对应OpPlugin版本分支代码进入插件根目录以v2.7.1为例。gitclone--branch7.3.0 https://gitcode.com/ascend/op-plugin.gitcdop-pluginpython安装依赖torch以2.7.1为例子pipinstalltorch2.7.1 pyyaml numpy attrs decorator psutil scipy执行编译构建当前支持torch_npu 2.6.0/2.7.1/2.8.0/2.9.0版本下述命令中v2.7.1-7.3.0表示匹配OpPlugin仓7.3.0版本的PyTorchv2.7.1的分支名。bashci/build.sh--python3.11--pytorchv2.7.1-7.3.0编译好之后会显示.... adding torch_npu.egg-info/PKG-INFO adding torch_npu.egg-info/SOURCES.txt adding torch_npu.egg-info/dependency_links.txt adding torch_npu.egg-info/entry_points.txt adding torch_npu.egg-info/requires.txt adding torch_npu.egg-info/top_level.txt adding torch_npu-2.7.1.post3.dist-info/METADATA adding torch_npu-2.7.1.post3.dist-info/WHEEL adding torch_npu-2.7.1.post3.dist-info/entry_points.txt adding torch_npu-2.7.1.post3.dist-info/top_level.txt adding torch_npu-2.7.1.post3.dist-info/RECORD removing build/bdist.linux-aarch64/wheel完成编译后安装dist目录下生成的插件torch_npu包pip3install--upgradedist/torch_npu-{torch_npu_version}-{Python_version}-{arch}.whl3. 代码修改找到#if VERSIONBETWEEN(V2R2, VERSIONNEWEST)块修改npu_fusion_attention改名为swa_forwardnpu_fusion_attention_grad改名为swa_backward另外sparse_mode只能是0和4加上两个函数内加上校验TORCH_CHECK(sparse_mode 0 || sparse_mode 4, The sparse_mode value should be 0 or 4, but got , sparse_mode, OPS_ERROR(ErrCode::PARAM));再加上特殊业务需求T大于16Ksparse_mode 4int64_t B 0; int64_t S0 0; // S for query int64_t S1 0; // S for key value int64_t N_local 0; // N for npu_fusion_attention int64_t D 0; int64_t H 0; int64_t T 0; int64_t D2 0; // D2 for value head-dim c10::SmallVectorint64_t atten_score_shape; if (input_layout_str TND) { T query.size(0); if (T 16 * 1024) { sparse_mode 4; // set sparse_mode to 4 when sequence length is greater than 16K in TND layout, which means using high-precision kernel } N_local query.size(1); D query.size(THIRD_ELEMENT); D2 value.size(THIRD_ELEMENT); atten_score_shape {T, N_local, D2}; }