flash-attention-npu：面向昇腾NPU的高性能FlashAttention算子库

flash-attention-npu是什么？

在大模型训练与推理中，Attention计算往往是最耗时、最占显存的环节。随着模型参数规模和上下文长度不断增长，传统注意力计算O(n2)的显存复杂度，成为制约大模型效率的核心瓶颈。FlashAttention通过分块计算、IO感知优化等手段，在不改变计算结果的前提下大幅降低显存占用、提升计算速度，如今已成为大模型基础设施的关键组件。业界主流开源实现Dao-AILab/flash-attention基于GPU架构进行了深度优化，已持续演进多个版本，被广泛应用于大模型训练与推理场景。

flash-attention-npu正是在昇腾NPU上的原生FlashAttention实现：面向昇腾NPU深度优化的FlashAttention实现，提供与业界主流实现一致的接口与用法。用户只需改动一行import，即可将原有注意力计算代码平滑运行在昇腾NPU上，并获得高性能加速。

用户无感切换，开发便捷扩展

flash-attention-npu的设计目标明确：对上层用户，接口一致、无感切换；对开发者，架构清晰、易于扩展。

算子基于昇腾异构计算架构CANN（Compute Architecture for Neural Networks）与昇腾算子模板库CATLASS（CANN Templates for Linear Algebra Subroutines）构建。CATLASS提供了面向昇腾NPU的高性能计算模板，使算子能够复用底层优化能力，避免重复的指令级适配工作；模板化设计也让计算组件解耦复用，方便多版本算子协同演进和新特性快速落地。

整体代码采用四层分层架构——Python接口层、Host侧处理层、Kernel侧内核层、CATLASS模板层。上层接口保持统一稳定，底层优化独立迭代，上层用户无需关心内部实现即可享受性能提升，开发者也能基于分层架构便捷地扩展和优化算子。整体架构如下图所示：

产研联创，助力昇腾AI生态繁荣

flash-attention-npu由中山大学国家超算广州中心沈明华课题组与昇腾技术团队联合研发。双方围绕算子架构设计、工程实现、性能调优与版本演进等全方位开展协同研发，以“学术前沿 + 产业落地”双轮驱动，共同完善昇腾AI生态，同时致力于将flash-attention-npu打造成昇腾AI生态中高性能FlashAttention的标准实现，持续完善算子功能与性能，为大模型计算提供高效、易用的Attention基础算子，让每一位使用昇腾NPU的伙伴们都能获得基于昇腾NPU原生加速的高效注意力计算体验。

当前核心能力

flash-attention-npu已具备覆盖主流大模型场景的能力，从硬件适配、算法版本到训练推理全流程，为昇腾NPU提供开箱即用的注意力加速方案：

多架构适配：适配昇腾910、950系列产品，支持多平台高性能计算

多版本支持：覆盖FlashAttention 2/3/4三代版本

全流程覆盖：覆盖训练（前向 + 反向传播）与推理（KV缓存等）两大场景

核心特性：支持MHA/MQA/GQA注意力架构、BSND、BNSD和TND形状、分页KV缓存、FlashDecoding、因果注意力和滑动窗口注意力等

精度支持：覆盖FP16 / BF16主流精度，FP8量化计算即将支持

快速上手

flash-attention-npu的使用门槛极低——安装一行命令、接口与主流实现完全一致、改动一行import即可完成迁移。以下是快速上手指南：

• 一键安装flash-attn-npu包

Plain Textpip install flash-attn-npu --no-build-isolation 

• 便捷使用flash-attn-npu算子

算子接口与业界主流FlashAttention完全对齐——原有代码只需将 import 语句替换为flash_attn_npu，即可在昇腾NPU上运行。以下是带KV Cache的推理场景用法示例，核心调用仅需一行：

Pythonfrom flash_attn_npu import flash_attn_with_kvcache  # 算子接入class LlamaAttentionLayer(nn.Module):    """    截取LlamaAttentionLayer的核心计算部分, 作为用法示例。    """    def __init__(self, config):        super().__init__()        # 步骤1：初始化参数配置        self.config = config        self.hidden_size = config.hidden_size        self.num_attention_heads = config.num_attention_heads        self.num_key_value_heads = config.num_key_value_heads        self.head_dim = self.hidden_size // self.num_attention_heads        self.causal = config.causal        self.window_size = config.window_size        self.num_splits = config.num_splits        self.k_cache = None        self.v_cache = None        # 步骤2：实例化计算所需 Matmul 算子        self.q_proj = nn.Linear(self.hidden_size, self.num_attention_heads * self.head_dim,                                 bias=config.attention_bias)        self.k_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim,                                 bias=config.attention_bias)        self.v_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim,                                 bias=config.attention_bias)        self.o_proj = nn.Linear(self.num_attention_heads * self.head_dim, self.hidden_size,                                 bias=config.attention_bias)    def forward(self, hidden_states):        # 步骤1：先通过 Matmul 算子计算 Q/K/V 张量        B, Sq, _ = hidden_states.shape        q = self.q_proj(hidden_states).view(B, Sq, self.num_attention_heads, self.head_dim)        k = self.k_proj(hidden_states).view(B, Sq, self.num_key_value_heads, self.head_dim)        v = self.v_proj(hidden_states).view(B, Sq, self.num_key_value_heads, self.head_dim)        # 步骤2：调用 flash_attn_npu 算子计算注意力结果        out = flash_attn_with_kvcache(q, self.k_cache, self.v_cache, k, v,                               cache_seqlens=self.cache_seqlens, causal=self.causal,                               window_size=self.window_size, num_splits=self.num_splits)        self.cache_seqlens += k.shape[1]        # 步骤3：通过 Matmul 算子计算出最终结果        out = out.reshape(B, Sq, self.num_attention_heads * self.head_dim)        out = self.o_proj(out)        return out

新版发布与后续计划
新版发布

近期，v0.5.0版本即将发布，面向昇腾910、950系列产品，从算子能力、性能、接口兼容、工程稳定性迭代，对齐主流FlashAttention，优化业务使用体验。该版本核心更新包括：

算子功能扩展：FA3/FA4接口支持变长序列、元数据接口、支持softcap，多平台扩展高head dim场景，优化确定性反向训练功能。

算子性能优化：优化LSE初始化、拷贝与释放，消除跨任务等待；优化昇腾950 FD合并、splitfuse流水线，精简Host冗余逻辑，增加变长输入校验，约束workspace。

接口兼容性与一致性增强：修复接口差异，统一不同硬件版本元数据行为，修正SWA、分页短KV调度窗口，规范化滑动窗口参数，保证跨版本一致性。

健壮性与工程适配改进：适配CANN 9.2，修复指令异常、头文件检测、DMA等待等bug；优化CI流水线，提升代码合入质量与回归效率，确保算子库在多硬件场景下的可靠交付。

后续计划

展望未来，flash-attention-npu将围绕三个方向持续演进：一是持续完善FlashAttention功能特性，依托CATLASS以及CATLASS DSL，探索适配昇腾硬件层级的新范式，孵化昇腾专属的高效注意力变种算子；二是持续优化算子性能与计算稳定性，提升代码质量和算子性能；完善CI/CD流水线，保障用户安装与使用的体验；三是建设完善的使用文档与典型场景示例，帮助用户更顺畅地上手和落地。我们由衷欢迎广大社区伙伴参与讨论贡献，共同推动昇腾AI生态的繁荣。

源码地址 & 诚邀体验

flash-attention-npu已在GitHub正式开源。无论是在昇腾NPU上训练大模型，还是做推理部署，flash-attention-npu都能带来更高效的注意力计算体验。欢迎各开发者一起使用或提Issue，一起助力昇腾AI基础软件生态的持续发展。

• 源码地址：https://github.com/MinghuasLab/flash-attention-npu

• Issue反馈：https://github.com/MinghuasLab/flash-attention-npu/issues

• 贡献代码：https://github.com/MinghuasLab/flash-attention-npu/pulls

fa-npu开源代码仓二维码

◆往期推荐◆

昇腾支持MiniMax M3.1 Flash，释放多款产品推理潜能