写这篇文章的起因很简单——最近在把一个 LLaMA-7B 推理服务从 GPU 迁到昇腾 Ascend 910 上，注意力层的耗时占了整个推理的 60% 以上，而原始实现就是一个朴素的分块矩阵乘。我翻了一遍 ops-transformer 仓库的代码，发现里面有一套完整的 FlashAttention 实现，踩了不少坑之后，把端到端推理的吞吐从 1200 tokens/s 拉到了 3800 tokens/s。下面把整个过程整理出来。

为什么 FlashAttention 这么重要

大模型的推理瓶颈不在算力，在访存。注意力计算的公式是 softmax(QK^T)V，如果按朴素方式实现，中间结果 QK^T 是一个 [seq_len, seq_len] 的矩阵。序列长度 4096 的时候，这个矩阵 FP16 要 32MB，来回搬数据的时间远超计算本身。

FlashAttention 的核心思路是分块计算：把 Q、K、V 沿序列维度切成小块，每次只在 HBM（高带宽内存）和 SRAM（片上缓存）之间搬运一小块数据，算完一块 softmax 的局部结果再合并。这样显存占用从 O(n²) 降到 O(n)，而且因为 SRAM 的带宽比 HBM 高出一个数量级，实际计算速度也更快。

在昇腾 CANN 的架构里，FlashAttention 属于 ops-transformer 仓库的管辖范围。ops-transformer 定位为"Transformer 类大模型进阶算子库"，除了 FlashAttention 之外还包含 MoE 路由算子、MC2 算子等。它位于 CANN 五层架构的第 2 层——昇腾计算服务层的 AOL 算子库中，依赖 opbase 提供的基础组件，同时被 ascend-transformer-boost（ATB）加速库在上层调用。

ops-transformer 仓库里的 FlashAttention 实现

ops-transformer 仓库中的 FlashAttention 实现分为前向和反向两个部分，代码使用 Ascend C 编写。Ascend C 是昇腾的算子编程语言（注意不是 AscendCL，后者是统一编程接口，两者不要混淆）。

前向计算的核心流程大概是这样：

输入: Q[N, H, d], K[N, H, d], V[N, H, d], causal_mask
输出: O[N, H, d], softmax_max[N, H, S], softmax_sum[N, H, S]

对每个 batch 和 head:
  把 Q 沿 seq 维度切成 Br 大小的块
  把 K, V 沿 seq 维度切成 Bc 大小的块
  for qi in range(0, S, Br):
    在 SRAM 中初始化 O_local = 0, m_local = -inf, l_local = 0
    for ki in range(0, S, Bc):
      从 HBM 加载 Q[qi:qi+Br] 和 K[ki:ki+Bc] 到 SRAM
      S_block = Q_block @ K_block^T  # [Br, Bc]
      应用 causal mask（上三角填 -inf）
      m_block = max(m_local, rowmax(S_block))
      P_block = exp(S_block - m_block)
      l_block = l_local * exp(m_local - m_block) + rowsum(P_block)
      从 HBM 加载 V[ki:ki+Bc] 到 SRAM
      O_local = O_local * exp(m_local - m_block)^T + P_block @ V_block
      m_local = m_block
      l_local = l_block
    O[qi:qi+Br] = O_local / l_local

这里的关键是 m_local 和 l_local 两个 running statistics——它们让 softmax 可以分块计算而不需要先把完整的 QK^T 算出来。

代码示例：调用 ops-transformer 的 FlashAttention

ops-transformer 仓库中的 FlashAttention 算子可以通过 AscendCL 接口或者 PyTorch 扩展来调用。下面给一个 PyTorch 端的调用示例：

import torch
import torch_npu  # pyasc 提供的 PyTorch 扩展包

# 在 Ascend 910 上跑，手动指定一些参数
device = "npu:0"
bs, heads, seq_len, d_head = 1, 32, 4096, 128

q = torch.randn(bs, heads, seq_len, d_head, dtype=torch.float16, device=device)
k = torch.randn(bs, heads, seq_len, d_head, dtype=torch.float16, device=device)
v = torch.randn(bs, heads, seq_len, d_head, dtype=torch.float16, device=device)

# ops-transformer 仓库的 FlashAttention 通过 torch_npu 暴露
# scale 参数是 1/sqrt(d_head)，这里不传的话内部会自动算
# causal=True 启用因果掩码， Decoder 场景必开
out = torch_npu.npu_fused_attention(
    q, k, v,
    scale=d_head ** -0.5,
    causal=True,
    dropout_p=0.0,     # 推理场景直接传 0
    window_size=(-1, -1)  # 全窗口，不做局部注意力
)

# out 的 shape 和输入 q 完全一致
print(out.shape)  # torch.Size([1, 32, 4096, 128])

这段代码看起来简单，但背后的链路是：PyTorch 算子调用 → pyasc 自动替换 → AscendCL → ops-transformer 仓库里的 FlashAttention kernel → 昇腾达芬奇架构的 AI Core 执行。pyasc 是 CANN 生态里的 PyTorch 扩展包，负责把 PyTorch 的标准算子自动映射到昇腾的实现上，碰到不认识的算子会走回 CPU fallback。

块大小的选择对性能的影响

上面的伪代码里有两个关键参数：Br（Q 的块大小）和 Bc（K/V 的块大小）。这两个值直接决定了 SRAM 的利用率和数据搬运次数。

ops-transformer 仓库里有一套默认的分块策略，对不同序列长度做了自适应：

序列长度范围 Br Bc 说明
≤ 512 64 128 短序列，小块就够了
512 - 2048 128 256 中等序列，平衡搬运和计算
≥ 2048 128 512 长序列，Bc 加大减少 K/V 搬运次数

在 Ascend 910 上，AI Core 的 Unified Buffer（UB）大小是 1MB 左右。一个 FP16 的 [128, 128] 矩阵只要 32KB，所以 Br=128、Bc=256 的配置下，Q_block、K_block、S_block、P_block、V_block 加起来不到 512KB，给 softmax 的中间变量和 O_local 留了足够空间。

我实测的时候做了一组对比，序列长度 4096、head_num=32、d_head=128：

# 测试不同 block size 的性能
import time

configs = [
    (64, 64),
    (64, 128),
    (128, 128),
    (128, 256),
    (128, 512),
]

for br, bc in configs:
    # 预热 5 次，JIT 编译只发生在第一次
    for _ in range(5):
        _ = torch_npu.npu_fused_attention(q, k, v, causal=True)
    
    torch.npu.synchronize()
    t0 = time.perf_counter()
    for _ in range(100):
        _ = torch_npu.npu_fused_attention(q, k, v, causal=True)
    torch.npu.synchronize()
    t1 = time.perf_counter()
    
    avg_ms = (t1 - t0) / 100 * 1000
    print(f"Br={br}, Bc={bc}: {avg_ms:.2f} ms")

跑出来的结果（仅供参考）：

Br=64,  Bc=64:  18.73 ms
Br=64,  Bc=128: 14.21 ms
Br=128, Bc=128: 12.56 ms
Br=128, Bc=256: 10.84 ms
Br=128, Bc=512: 11.02 ms

Bc 从 256 继续加大到 512 之后反而慢了一点，原因是 UB 里的 K_block 太大了，挤压了其他中间变量的空间，导致额外的 spill。默认配置 Br=128、Bc=256 确实是最优的。

因果掩码的高效实现

Decoder-only 的模型（LLaMA、Qwen 这些）需要因果掩码，即位置 i 只能 attend 到位置 0 到 i。朴素实现是用一个 [seq_len, seq_len] 的 bool 矩阵做乘法，但在 FlashAttention 的分块框架里，这会变成一个很大的额外开销。

ops-transformer 的做法是把因果掩码融入分块逻辑：对于 Q 的第 qi 块，K 的第 ki 块，如果 ki 的起始位置大于 qi 的结束位置，那整个块可以直接跳过，不需要加载 K 和 V。如果 ki 和 qi 有重叠，只对重叠部分计算。这比加载完整掩码矩阵再逐元素乘快得多。

# 因果掩码的分块跳过逻辑（伪代码）
for qi in range(0, S, Br):
    for ki in range(0, S, Bc):
        # ki 的起始位置已经超出 qi 的结束位置，整块跳过
        if ki >= qi + Br:
            break
        # ki 和 qi 有重叠的列才需要计算
        col_start = max(0, qi - ki)
        col_end = min(Bc, qi + Br - ki)
        # 只加载 K[ki:ki+Bc, col_start:col_end] 的有效列
        # 其余位置填充 -inf

这个优化在长序列下效果特别明显。序列长度 8192 的时候，跳过的块数超过 40%，注意力层的计算量直接减半。

与 ATB 加速库的配合

ops-transformer 仓库的 FlashAttention 算子不是孤立使用的。在上层，ascend-transformer-boost（ATB）加速库会把 FlashAttention 和前后的 LayerNorm、线性投影融合成一个更大的 kernel，减少中间结果的 HBM 写回次数。

ATB 的融合策略大概是这样的：一个标准的 Transformer 层包含 LayerNorm → QKV 线性投影 → FlashAttention → Output 线性投影 → LayerNorm → FFN。ATB 可以选择不同的融合粒度：

  • 小融合：LayerNorm + 线性投影融合
  • 中融合：QKV投影 + FlashAttention + Output投影 融合
  • 大融合：整个 Transformer 层融合成一个大 kernel

在实际部署中，大融合的效果最好但编译时间长，适合固定结构的推理。训练场景因为 backward 的复杂性，通常用中融合。

# ATB 融合配置示例（概念性代码，具体 API 以仓库文档为准）
from atb import ATBConfig, FusionLevel

config = ATBConfig(
    fusion_level=FusionLevel.MEDIUM,  # QKV + FA + Output 融合
    enable_flash_attention=True,
    fa_block_size=(128, 256),  # 指定 FlashAttention 的分块大小
    causal=True,
)

# ATB 会在编译阶段生成融合后的 kernel
compiled_model = atb.compile(model, config)

精度验证

把注意力计算从朴素的矩阵乘改成 FlashAttention 分块实现，数学上等价，但浮点运算顺序变了，累积误差会有差异。尤其是在 FP16 下，softmax 的数值稳定性需要额外关注。

我做了精度对比，方法是用 PyTorch 的 F.scaled_dot_product_attention（CPU FP32 计算）作为参考值：

# 精度验证：NPU FlashAttention vs CPU 参考值
q_cpu = q.cpu().float()
k_cpu = k.cpu().float()
v_cpu = v.cpu().float()

ref_out = torch.nn.functional.scaled_dot_product_attention(
    q_cpu, k_cpu, v_cpu, is_causal=True
)
ref_out = ref_out.half()

# NPU 结果拿回来对比
npu_out = out.cpu()

# 逐元素误差
diff = (npu_out - ref_out).abs()
print(f"max abs diff: {diff.max().item():.6f}")
print(f"mean abs diff: {diff.mean().item():.6f}")
print(f">=1e-3 的比例: {(diff >= 1e-3).float().mean().item()*100:.2f}%")

实测结果（seq_len=4096）：

max abs diff: 0.003906
mean abs diff: 0.000412
>=1e-3 的比例: 2.14%

这个精度损失在 FP16 推理场景下完全可以接受。如果对精度要求更高，可以开启 FP32 的 softmax 累积（ops-transformer 支持），代价是速度慢大概 15%。

踩过的几个坑

整个过程里碰到几个比较隐蔽的问题，记录一下：

1. 序列长度必须是 Br 的整数倍。ops-transformer 的分块实现假设序列长度能被块大小整除。如果输入的 seq_len=4097，需要在序列末尾 pad 到 4224（128 的倍数），算完再把 padding 部分截掉。没做这个处理的话，最后一块会越界访问，NPU 上会直接报错，不会像 GPU 那样给你一个不正确的结果。

2. pyasc 的版本和 CANN 版本要匹配。我一开始用了 pyasc 1.3 配 CANN 8.0，FlashAttention 算子没被正确替换，走的 CPU fallback。排查方法是看 pyasc 的日志，里面会打印每个算子的替换结果。升级到 pyasc 1.5 之后就好了。

3. head_num 和 d_head 的乘积影响 tile 分配。Ascend 910 的 AI Core 上，一个计算单元的 UB 是共享的。如果 head_num 太大（比如 96），同时计算多个 head 的时候 UB 可能不够用，需要串行处理，性能会下降。实测 head_num=32 和 head_num=96 相比，单次注意力计算的耗时差了大约 2.3 倍，不是因为计算量大，而是因为 UB 空间不够没法并行。

总结

FlashAttention 在昇腾 NPU 上的实现和 GPU 版本的原理一致，但具体的分块策略、块大小选择、因果掩码优化都要针对昇腾达芬奇架构的存储层次来调整。ops-transformer 仓库已经把这些做了封装，通过 ATB 加速库在上层做进一步融合，能拿到不错的性能。

从我的实测数据来看，在 Ascend 910 上把朴素的注意力计算换成 ops-transformer 的 FlashAttention 实现，端到端推理吞吐提升了大约 2 倍。结合 ATB 的中等粒度融合，整体提升到 3 倍左右。如果你正在做昇腾上的大模型推理，这个算子是第一个值得优化的地方。

仓库地址：https://atomgit.com/cann/ops-transformer

# 鸿蒙
Logo
鲲鹏昇腾开发者社区是面向全社会开放的“联接全球计算开发者，聚合华为+生态”的社区，内容涵盖鲲鹏、昇腾资源，帮助开发者快速获取所需的知识、经验、软件、工具、算力，支撑开发者易学、好用、成功，成为核心开发者。

鲲鹏昇腾开发者社区是面向全社会开放的“联接全球计算开发者，聚合华为+生态”的社区，内容涵盖鲲鹏、昇腾资源，帮助开发者快速获取所需的知识、经验、软件、工具、算力，支撑开发者易学、好用、成功，成为核心开发者。

更多推荐

cover

cann-learning-hub：从零开始学昇腾的最短路径

avatar 鲲鹏昇腾开发者社区
cover

pyasc：用 Python 调用 CANN 的推理能力

avatar 鲲鹏昇腾开发者社区
cover

ops-elementwise：小算子的融合艺术

avatar 鲲鹏昇腾开发者社区