基于昇腾进行Glm5.1模型swa算子接入优化

卡路里45465

14人浏览 · 2026-07-08 14:22:53

卡路里45465 · 2026-07-08 14:22:53 发布

作者：昇腾实战派
知识地图：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枚举定义如下：

enum class SparseMode {
    NO_MASK = 0,
    ALL_MASK,
    LEFT_UP_CAUSAL,
    RIGHT_DOWN_CAUSAL,
    BAND,
    PREFIX,
    PREFIX_COMPRESS,
    RIGHT_DOWN_CAUSAL_BAND,
    BAND_LEFT_UP_CAUSAL
};

其中，模式0表示无mask，模式4（BAND）表示使用pre_tockens与next_tockens控制窗口的滑动窗口mask。本实现仅需支持这两种模式，无需关注底层AscendC算子的具体实现。

2. 编译安装

这个需求需要编译 op-plugin 这个库。

先拉取CANN镜像

docker pull  swr.cn-south-1.myhuaweicloud.com/ascendhub/cann:8.5.2-910b-ubuntu22.04-py3.11

启动容器

NAME=cann8.5-test
DEVICES="0,1,2,3,4,5,6,7"
IMAGE="swr.cn-south-1.myhuaweicloud.com/ascendhub/cann:8.5.2-910b-ubuntu22.04-py3.11"
docker run -itd -u 0  --ipc=host  --privileged \
-e VLLM_USE_MODELSCOPE=True -e PYTORCH_NPU_ALLOC_CONF=max_split_size_mb:256 \
-e ASCEND_RT_VISIBLE_DEVICES=$DEVICES \
--name $NAME \
--net=host \
--device /dev/davinci_manager \
--device /dev/devmm_svm \
--device /dev/hisi_hdc \
--shm-size=1200g \
-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/ \
--privileged=true \
-it $IMAGE bash

下载对应OpPlugin版本分支代码，进入插件根目录，以v2.7.1为例。

git clone --branch 7.3.0 https://gitcode.com/ascend/op-plugin.git
cd op-plugin

python安装依赖（torch以2.7.1为例子）：

pip install torch==2.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的分支名。

bash ci/build.sh --python=3.11 --pytorch=v2.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包

pip3 install --upgrade dist/torch_npu-{torch_npu_version}-{Python_version}-{arch}.whl

3. 代码修改

找到#if VERSIONBETWEEN(V2R2, VERSIONNEWEST)块，修改

npu_fusion_attention改名为swa_forward

npu_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大于16K，sparse_mode = 4

int64_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::SmallVector<int64_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};
    }