昇腾CANN上的FlashAttention工程实战ops-transformer源码拆解最近在昇腾NPU上部署一套大模型推理服务性能瓶颈死死卡在Attention层。翻CANN的算子仓库发现ops-transformer里直接给了FlashAttention的实现省了自己从头造轮子。这篇文章记录我在CANN 8.0环境下把ops-transformer的FlashAttention算子接进推理链路的完整过程顺带拆解它在昇腾达芬奇架构上的工程实现。标准Attention为什么在NPU上跑不动标准Scaled Dot-Product Attention的计算流程Q乘K的转置得到Score矩阵Score除以缩放因子后做Softmax最后乘V得到输出。问题出在中间结果——Score矩阵的大小是seq_len × seq_len序列长度128K时仅这一个中间矩阵在fp16下就要吃掉32GB显存。更麻烦的是这套流程在昇腾NPU上的数据搬运开销远超计算开销QK^T的结果要从Cube Unit Buffer搬出来做SoftmaxVector单元做完再搬回去乘V又是Cube单元。来回折腾带宽全浪费在搬运上了。Tiling Online Softmax原理回顾FlashAttention的核心思路是分块计算Tiling加上在线Softmax。Q、K、V按固定大小的Tile加载到片上缓存每个Tile内独立完成注意力计算。Softmax不做全序列归一化而是维护两个累积变量——当前最大值m和指数和l。每处理完一个新Tile用新Tile的局部最大值更新m再用新旧m的差值校正之前所有Tile的累加结果。这样整条链路中中间结果O始终只占Tile大小的空间显存占用从O(N²)降到O(N)。ops-transformer怎么落在昇腾硬件上ops-transformer的FlashAttention实现把这个算法精确映射到了达芬奇架构的计算单元上。Ascend 910有两套核心计算单元Cube Unit负责大矩阵乘GEMM算力密度高Vector Unit负责逐元素运算灵活但吞吐低。FlashAttention的计算图天然分成两类——QK^T和PV走CubeSoftmax、Scale、Dropout走Vector。关键设计在于Buffer管理。标准FlashAttention论文假设硬件有一块统一的SRAM但达芬奇架构的Cube Unit Buffer和Vector Unit Buffer是两块独立的片上存储。ops-transformer对两块Buffer做了分别管理矩阵乘的中间结果S矩阵的每个Tile驻留在Cube BufferSoftmax的累加状态m、l、O的当前分块放在Vector Buffer。避免了数据在两种Buffer之间反复搬运。这个细节是昇腾实现区别于GPU上CUDA实现的核心差异也是性能能打的关键。另一个工程细节Layout转换。昇腾NPU的矩阵乘对数据布局有要求输入需要从ND行主序转成NZ分块列主序格式才能喂给Cube Unit。ops-transformer在算子入口处自动做了这个转换用户侧无感。实际跑出来的数据基于Llama-70B推理batch_size1fp16精度在Ascend 910上测试了三组序列长度seq_len标准Attention 吞吐FlashAttention 吞吐吞吐提升标准Attention 显存FlashAttention 显存2,0481,680 t/s2,950 t/s75.6%16.2 GB8.4 GB8,1921,180 t/s3,420 t/s189.8%38.6 GB12.1 GB32,768OOM2,760 t/s∞OOM28.3 GB32K序列长度下标准Attention直接OOMFlashAttention还在正常跑。吞吐随序列长度的增长退化也比标准方案温和得多。接入踩坑从ops-transformer拉源码编译后接入PyTorch模型替换原来的attention实现fromops_transformerimportflash_attention# 替换F.scaled_dot_product_attention# 省掉HBM中间结果的反复搬运seq_len32K也不OOMoutflash_attention(q,k,v,scale1.0/math.sqrt(d))踩到的坑CANN 8.0之前的版本只支持fp16bf16支持是在CANN 8.0才补上的。如果你用的CANN版本低于8.0精度在高频推理场景下会掉点先检查版本再查代码。另外ops-transformer要求输入Q/K/V的shape必须是Tile大小的整数倍不对齐时需要做padding这个在文档里有写但容易被忽略。仓库地址https://atomgit.com/cann/ops-transformer
昇腾CANN上的FlashAttention工程实战：ops-transformer源码拆解
发布时间：2026/5/22 3:35:49
昇腾CANN上的FlashAttention工程实战ops-transformer源码拆解最近在昇腾NPU上部署一套大模型推理服务性能瓶颈死死卡在Attention层。翻CANN的算子仓库发现ops-transformer里直接给了FlashAttention的实现省了自己从头造轮子。这篇文章记录我在CANN 8.0环境下把ops-transformer的FlashAttention算子接进推理链路的完整过程顺带拆解它在昇腾达芬奇架构上的工程实现。标准Attention为什么在NPU上跑不动标准Scaled Dot-Product Attention的计算流程Q乘K的转置得到Score矩阵Score除以缩放因子后做Softmax最后乘V得到输出。问题出在中间结果——Score矩阵的大小是seq_len × seq_len序列长度128K时仅这一个中间矩阵在fp16下就要吃掉32GB显存。更麻烦的是这套流程在昇腾NPU上的数据搬运开销远超计算开销QK^T的结果要从Cube Unit Buffer搬出来做SoftmaxVector单元做完再搬回去乘V又是Cube单元。来回折腾带宽全浪费在搬运上了。Tiling Online Softmax原理回顾FlashAttention的核心思路是分块计算Tiling加上在线Softmax。Q、K、V按固定大小的Tile加载到片上缓存每个Tile内独立完成注意力计算。Softmax不做全序列归一化而是维护两个累积变量——当前最大值m和指数和l。每处理完一个新Tile用新Tile的局部最大值更新m再用新旧m的差值校正之前所有Tile的累加结果。这样整条链路中中间结果O始终只占Tile大小的空间显存占用从O(N²)降到O(N)。ops-transformer怎么落在昇腾硬件上ops-transformer的FlashAttention实现把这个算法精确映射到了达芬奇架构的计算单元上。Ascend 910有两套核心计算单元Cube Unit负责大矩阵乘GEMM算力密度高Vector Unit负责逐元素运算灵活但吞吐低。FlashAttention的计算图天然分成两类——QK^T和PV走CubeSoftmax、Scale、Dropout走Vector。关键设计在于Buffer管理。标准FlashAttention论文假设硬件有一块统一的SRAM但达芬奇架构的Cube Unit Buffer和Vector Unit Buffer是两块独立的片上存储。ops-transformer对两块Buffer做了分别管理矩阵乘的中间结果S矩阵的每个Tile驻留在Cube BufferSoftmax的累加状态m、l、O的当前分块放在Vector Buffer。避免了数据在两种Buffer之间反复搬运。这个细节是昇腾实现区别于GPU上CUDA实现的核心差异也是性能能打的关键。另一个工程细节Layout转换。昇腾NPU的矩阵乘对数据布局有要求输入需要从ND行主序转成NZ分块列主序格式才能喂给Cube Unit。ops-transformer在算子入口处自动做了这个转换用户侧无感。实际跑出来的数据基于Llama-70B推理batch_size1fp16精度在Ascend 910上测试了三组序列长度seq_len标准Attention 吞吐FlashAttention 吞吐吞吐提升标准Attention 显存FlashAttention 显存2,0481,680 t/s2,950 t/s75.6%16.2 GB8.4 GB8,1921,180 t/s3,420 t/s189.8%38.6 GB12.1 GB32,768OOM2,760 t/s∞OOM28.3 GB32K序列长度下标准Attention直接OOMFlashAttention还在正常跑。吞吐随序列长度的增长退化也比标准方案温和得多。接入踩坑从ops-transformer拉源码编译后接入PyTorch模型替换原来的attention实现fromops_transformerimportflash_attention# 替换F.scaled_dot_product_attention# 省掉HBM中间结果的反复搬运seq_len32K也不OOMoutflash_attention(q,k,v,scale1.0/math.sqrt(d))踩到的坑CANN 8.0之前的版本只支持fp16bf16支持是在CANN 8.0才补上的。如果你用的CANN版本低于8.0精度在高频推理场景下会掉点先检查版本再查代码。另外ops-transformer要求输入Q/K/V的shape必须是Tile大小的整数倍不对齐时需要做padding这个在文档里有写但容易被忽略。仓库地址https://atomgit.com/cann/ops-transformer
