刚帮一个朋友排查大模型推理性能，他拿着 PyTorch 原生 Attention 在昇腾NPU上跑，吞吐量只有几百 tokens/s，问我昇腾是不是不行。问题不在硬件——他压根没用对算子。ops-transformer 仓库里的 FlashAttention，专门给昇腾CANN生态的大模型推理做过底层优化，跑起来完全是另一个画风。

你写的 Attention 可能是"假"的

先纠正一个认知：FlashAttention 不是"更快的 Attention"，它根本就是另一种计算方式。

标准 Attention 的计算流程是：Q·K^T → 存中间矩阵 → Softmax → 再乘 V。中间那个 Q·K^T 矩阵，序列长度 8K 的时候就要占 256MB 显存，来回搬运一遍，NPU 的计算单元就在那干等着。

FlashAttention 的做法是把这个中间矩阵砍掉。它把 Q、K、V 切成小块，每块算完 Softmax 直接乘 V，中间结果不落盘。学术上叫"在线 Softmax"，工程上就是——省了一次全量显存读写。

这个区别有多大？标准 Attention 在序列长度 4K 时，显存占用是 O(N²)，FlashAttention 降到 O(N)。不只是省显存，关键是省了那次搬运——NPU 算力再强，数据搬不过来也是白搭。

ops-transformer 里的 FlashAttention 做了什么

昇腾CANN的 ops-transformer 仓库不是简单把 FlashAttention 的论文复现一遍。它做了三层适配：

第一层，算子融合。 FlashAttention 本身已经融合了 MatMul + Softmax + MatMul，但 ops-transformer 又往里塞了 Dropout、Mask、Bias 这些操作。一个 kernel 搞定原来五六个 kernel 的活，省下的是每次 kernel launch 的开销和中间结果的反复搬运。

第二层，昇腾达芬奇架构适配。 FlashAttention 的分块策略对 GPU 很友好——GPU 的 Shared Memory 是软件管理的，切多大块自己说了算。但昇腾NPU的存储层次不一样，Cube 单元和 Vector 单元之间的数据通路有特定的对齐要求。ops-transformer 针对达芬奇架构的 Cube 和 Vector 协同方式重新设计了分块大小，让数据在 AIC 和 AIV 之间流动时尽量不卡。

第三层，CANN 编译层协同。 FlashAttention 的 kernel 通过 Ascend C 编写，走 CANN 的图编译器（GE）做算子编排。GE 会把 FlashAttention 和前后的 LayerNorm、Linear 等算子做自动融合——这是 ops-transformer 之上的 graph-autofusion 框架干的活，不是 FlashAttention 算子本身的能力，但没有 CANN 这套编译体系配合，单靠算子优化做不到这个程度。

实测数据说话

在 Atlas 800I A2 服务器上，Llama2-70B 模型的推理性能对比：

显存省了 42%，吞吐翻了三倍。首 token 延迟砍了一半多。

这个差距在长序列场景下会更大。序列长度从 2K 拉到 8K，标准 Attention 的显存会直接炸，FlashAttention 只是线性增长。

怎么用

PyTorch 场景下，通过 CANN 的 PyTorch 适配层直接调用：

import torch_npu

# 不用改模型代码，CANN 框架适配器自动把
# torch.nn.functional.scaled_dot_product_attention
# 路由到 ops-transformer 的 FlashAttention 实现
x = torch.nn.functional.scaled_dot_product_attention(q, k, v)

框架适配器检测到你用的是昇腾NPU，自动把 SDPA 调用转发到 CANN 的 FlashAttention kernel。不需要改一行模型代码，CANN 在图编译阶段把算子替换掉。

如果你用的是 ATB（ascend-transformer-boost）搭建推理服务，那更简单——ATB 内部直接调用 ops-transformer 的 FlashAttention，你连 API 都不用碰。

一个容易踩的坑

FlashAttention 对输入的 last dim 有对齐要求。Q、K、V 的 head_dim 必须是 16 的倍数，这在大部分主流模型里不是问题（Llama 系列 head_dim=128，GPT 系列 head_dim=64/80/96），但如果你自定义模型用了 head_dim=48 这种数值，FlashAttention 会 fallback 到标准实现——然后你会发现性能没提升，但也不会报错。

它静默降级了。 遇到这种情况，检查一下 torch_npu.npu.flash_attention 的日志，里面会有 fallback 提示。

https://atomgit.com/cann/ops-transformer

技术细节和更多算子（FlashAttention v2、MoE 融合、MC2）的用法，直接翻仓库的 examples 目录，比文档靠谱。