【FlashAttention 在昇腾 CANN 上的实现原理与性能优化】

FlashAttention 在昇腾 CANN 上的实现原理与性能优化

上周帮同事调一个 7B 模型的推理流水线，latency 卡在 attention 算子上出不来。他说 GPU 上跑得好好的，换到昇腾 NPU 上就慢了一截。我看了一眼算子实现------直接把整块 attention 扔给 CANN 的 ops-transformer 仓库跑了事，连 tiling 策略都没调。这不行。FlashAttention 在昇腾 CANN 上的行为跟 GPU 完全不是一套逻辑，不搞清楚 Cube 和 Vector 怎么分工、数据怎么搬，性能根本出不来。

先说结论：FlashAttention 的核心加速点不是某个黑魔法，而是把 O(N²) 的显存访问变成 O(N)。昇腾 NPU 上实现这事儿，硬件架构差异大，但思路可以迁移，关键在于理解 CANN 的 tiling 是怎么跟达芬奇架构的 Cube/Vector 两级计算单元配合的。

Attention 的数学：一个矩阵乘法的故事

标准 Attention 公式：

code复制

复制代码

Q, K, V ∈ R^(N×d)

S = Q × K^T // (N, N) 注意力分数矩阵
P = softmax(S) // (N, N) 概率分布
O = P × V // (N, d) 输出

问题出在哪？S 和 P 都是 N×N 的矩阵。N=8192、d=128 的时候，S 一个矩阵就 256MB。backward 还得把 P 存下来，显存直接炸。GPU 上原始实现读写显存 ~O(N²d)，FlashAttention 把它压到 ~O(N²d/M)，M 是 SRAM 大小。

FlashAttention 的做法很直白：不存完整的 S 和 P，分块算。每算一块 S，立刻做 softmax 归一化（需要在线更新），算出对应的 O 子块，中间结果不落回 HBM。这就是所谓的 "tiling" 或 "fused kernel"。

昇腾 NPU vs GPU：架构差异决定一切

GPU 的流式多处理器（SM）是通用计算单元，SRAM 共享，算 attention 时一个 kernel 搞定 Q/K/V 的搬运、矩阵乘、softmax、再矩阵乘。

昇腾达芬奇架构不一样。它有两类计算单元：

单元	擅长	Attention 中的角色
Cube Unit	矩阵乘法（MAC 阵列）	Q×K^T、P×V
Vector Unit	逐元素运算	softmax、scale、mask、归一化

GPU 是一个 SM 干所有事；昇腾是 Cube 干矩阵乘、Vector 干逐元素，中间数据要通过 L1 缓存传递。这意味着 FlashAttention 的 fused kernel 在昇腾上不是"一个 kernel"，而是 Cube 和 Vector 交替执行的流水线。

打个比方：GPU 是一个厨师又炒菜又调味；昇腾是两个厨师------一个专门翻炒（Cube），一个专门调味（Vector），菜得在两人之间递来递去。

CANN 的 Tiling 策略：分块不只是数学问题

FlashAttention 分块在数学上不难，但在昇腾上 tilng 要同时满足：

Cube 的 MAC 阵列利用率：分块大小得让矩阵乘法填满 Cube 的计算单元，否则算力浪费
L1 缓存容量：Q、K、V 的 tile 要能放进 L1，不然数据溢出到 L2 甚至 HBM，延迟暴涨
Vector 吞吐：softmax 归一化在 Vector 上跑，tile 太小 Vector 吃不饱；tile 太大 L1 装不下

ops-transformer 仓库里 FlashAttention 算子的 tiling 策略大致是这样的：

code复制

复制代码

// tiling 参数推算（简化版）
tile_q = L1_capacity / (3 * d * sizeof(float16)) // Q 的一行能放多少
tile_k = Cube_MAC_rows // Cube 阵列行数对齐
tile_v = tile_k // V 跟 K 同步

// 实际运行时根据 seq_len、head_dim 动态调整
if seq_len > 4096:
 tile_q = min(tile_q, 128) // 长 seq 时压小 tile，避免 L1 溢出

关键点：tile 大小不是拍脑袋定的，是 CANN 编译器根据硬件参数和输入 shape 自动推算的。这也是为什么同一个模型不同 seq_len 下性能差异大------tiling 不一样。

执行流：Cube 和 Vector 怎么跳双人舞

完整的 FlashAttention 在昇腾上的执行流：

code复制

复制代码

┌─────────────────────────────────────────────┐
│ HBM → L1: 搬入 Q_tile, K_tile │
│ ↓ │
│ Cube: S_tile = Q_tile × K_tile^T │
│ ↓ │
│ L1 缓存: S_tile 暂存 │
│ ↓ │
│ Vector: S_tile *= scale │
│ S_tile += mask (如需) │
│ P_tile = softmax(S_tile) │
│ 在线更新: O_tile = O_tile + P_tile × V_tile │
│ ↓ │
│ 循环下一块 K/V │
│ ↓ │
│ L1 → HBM: 写出 O_tile │
└─────────────────────────────────────────────┘

这里有个容易踩的坑：softmax 的在线归一化。标准 FlashAttention 论文里用的是 log-sum-exp trick，每次新来一块 K/V，要更新全局的 max 和 sum。在昇腾上这事儿由 Vector 完成，但 Vector 得等 Cube 算完 S_tile 才能开始，这就产生了一个同步点。

GPU 上这个同步在同一个 SM 内，延迟几乎为零。昇腾上 Cube → Vector 的数据传递要走 L1，虽然也在片上，但延迟比 GPU 的共享内存高。这是架构差异带来的固有开销。

伪代码：用 Ascend C 的视角看 FlashAttention

cpp复制

复制代码

// Ascend C 视角的 FlashAttention 伪代码（简化版）
// 只展示核心循环，省略边界处理

for (int q_idx = 0; q_idx < num_q_tiles; q_idx++) {
 // 从 HBM 搬 Q 的第 q_idx 个 tile 到 L1
 local_q = LoadTile(Q, q_idx * tile_q, tile_q);
 local_o = Zeros(tile_q, d); // 输出先清零
 running_max = NegInf(tile_q); // 在线 softmax 的最大值
 running_sum = Zeros(tile_q); // 在线 softmax 的累加和

 for (int kv_idx = 0; kv_idx < num_kv_tiles; kv_idx++) {
 // 搬 K/V tile
 local_k = LoadTile(K, kv_idx * tile_k, tile_k);
 local_v = LoadTile(V, kv_idx * tile_k, tile_k);

 // Cube: 矩阵乘
 s_tile = MatMul(local_q, local_k.T()); // (tile_q, tile_k)

 // Vector: scale + mask + softmax
 s_tile = s_tile * (1.0 / sqrt(d)); // scale
 if (causal) ApplyCausalMask(s_tile, q_idx, kv_idx);

 // 在线 softmax 更新------这里是关键
 new_max = Max(s_tile, axis=-1);
 correction = Exp(running_max - new_max); // 旧值要缩放
 s_tile = Exp(s_tile - new_max);
 new_sum = running_sum * correction + Sum(s_tile, axis=-1);

 // Cube: P_tile × V_tile
 pv = MatMul(s_tile, local_v); // (tile_q, d)

 // Vector: 累加到 O
 local_o = local_o * correction + pv;
 running_max = new_max;
 running_sum = new_sum;
 }

 // 最终归一化
 local_o = local_o / running_sum;
 StoreTile(O, local_o, q_idx * tile_q); // 写回 HBM
}

注意几个工程细节：

correction = Exp(running_max - new_max)：每次 max 更新，之前累加的 O 和 sum 都要缩放。这步在 Vector 上做，多了一轮逐元素运算
causal mask 在 decoder-only 模型里必须有，mask 本身不费算力，但它让一半的 S_tile 变成 -inf，Cube 算了白算------这是 causal attention 在昇腾上的固有浪费
整个内层循环里 Cube 和 Vector 交替执行，中间有隐式同步

内存搬运：最容易被忽视的性能杀手

光看算力，Cube 的 MAC 阵列吞吐不比 GPU 的 Tensor Core 差。但实际性能往往被内存带宽卡住。

昇腾上的内存层次：

code复制

复制代码

HBM (几十GB, 带宽 ~1.2TB/s)
 ↓ DMA 搬运
L2 缓存 (几十MB)
 ↓ 自动/手动预取
L1 缓存 (几百KB~1MB) ← Cube 和 Vector 共享
 ↓
Cube 寄存器 / Vector 寄存器

FlashAttention 的核心收益就是减少 HBM 访问次数。但昇腾上还有一层：L1 到 Cube/Vector 寄存器的搬运。这个搬运 CANN 编译器会自动调度，但如果 tiling 不合理（tile 太小），L1 和计算单元之间的带宽也会成瓶颈。

实测数据（7B 模型，seq_len=4096，Ascend 910）：

配置	attention 吞吐 (tokens/s)	HBM 带宽利用率
未融合（标准 attention）	1,420	35%
FlashAttention（默认 tiling）	3,260	62%
FlashAttention（手动 tiling 调优）	3,810	71%

手动 tiling 调优指的是通过 ops-transformer 的 tiling 参数接口，针对特定 shape 调整 tile_q 和 tile_k。默认 tiling 是通用策略，不保证对所有 shape 最优。

性能瓶颈分析：三看三调

调 FlashAttention 性能，我总结为"三看"：

一看 Cube 利用率。如果 MAC 阵列填充率低于 70%，说明 tile 太小。调大 tile_k（K 的分块行数）通常最直接，但要注意 L1 容量上限。

二看 Vector 空闲比。如果 Vector 大量时间在等 Cube，说明矩阵乘是瓶颈，可以考虑把 scale 和 mask 合并到 Cube 输出后立即执行，减少 Cube→Vector 的切换次数。CANN 的 graph-autofusion 算子自动融合框架就是干这个的------把 scale+mask+softmax 融合成一个 Vector 子图，减少中间数据落 L1 的次数。

三看 HBM 带宽。如果带宽利用率低于 50%，数据搬运会拖后腿。检查 Q/K/V 的数据布局是否连续（NHWC vs NCHW），不连续的布局会让 DMA 搬运效率暴跌。

一个真实场景：我们在做 7B 模型推理时，batch=1、seq=2048 的场景下 FlashAttention 只用了 Cube 算力的 40%。原因？tile_q=32，MAC 阵列 16×16 只填了一半。把 tile_q 调到 64，吞吐直接涨了 22%。

工程实践：从 ops-transformer 入手

ops-transformer 仓库里 FlashAttention 算子的调用方式：

python复制

复制代码

import torch_npu # 昇腾 PyTorch 适配

# 最简调用------走 ATB 加速库的融合算子
from torch_npu.contrib import transfer_to_npu

# 标准 PyTorch 写法，torch_npu 自动路由到 CANN FlashAttention
output = torch.nn.functional.scaled_dot_product_attention(q, k, v)

如果想更细粒度地控制 tiling：

python复制

复制代码

# 通过 ATB 的参数接口设置 tiling
import ascend_transformer_boost as atb

# 针对长 seq 场景的手动 tiling
flash_attn_op = atb.FlashAttention(
 tile_q=64, # Q 分块大小
 tile_kv=128, # K/V 分块大小
 is_causal=True,
)
output = flash_attn_op(q, k, v)

踩坑提醒：tile_q 和 tile_kv 不是越大越好。L1 容量有限，tile 太大会触发溢出到 L2，延迟反而增加。建议从默认值开始，用 CANN 的 Profiling 工具看 Cube 利用率，再逐步调大。

刚从 CUDA 转 CANN 的朋友，最容易犯的错就是把 GPU 上的 kernel 优化思路直接搬过来。昇腾的 Cube+Vector 双单元架构决定了：优化重心不是减少计算量，而是让 Cube 和 Vector 的流水线跑满。关注 tiling、关注 L1 缓存命中率、关注 Cube 和 Vector 之间的数据搬运效率，这三件事做好了，FlashAttention 在昇腾上的性能不会比 GPU 差。

下一步建议：

用 torch_npu 跑一遍你的模型，对比标准 attention 和 FlashAttention 的吞吐差异
打开 CANN Profiling，看 Cube 利用率和 Vector 空闲比
如果 Cube 利用率低于 70%，调整 tiling 参数重跑
仓库源码和完整 API 文档：https://atomgit.com/cann/ops-transformer