FlashAttention 昇腾优化：从 O(N²) 到 O(N) 的显存革命

本文介绍了FlashAttention在昇腾NPU上的优化实践，通过改变显存访问模式显著提升推理性能。传统Attention计算中的QK^T矩阵（seq=4096时占32MB显存）导致HBM频繁读写，吞吐仅18 tokens/s。FlashAttention采用分块计算策略，将中间结果保留在L1缓存而非HBM，使显存占用从1.2GB降至180MB，吞吐提升7倍至147 tokens/s。关键优化包

高级c

99人浏览 · 2026-05-22 13:29:34

高级c · 2026-05-22 13:29:34 发布

前言

7B 模型推理 seq=4096 时，Attention 的 KV Cache 占 1.2GB 显存。batch=8 就爆显存，吞吐只有 18 tokens/s。用 FlashAttention 后显存降到 180MB，吞吐涨到 147 tokens/s。涨了 7 倍多，不是算力变强了，是 HBM 访问次数从 O(N²) 压到了 O(N)。

很多人以为 FlashAttention 只是"矩阵乘优化"，其实它的核心是显存访问模式的根本改变：不存 QK^T 这个 N×N 的大矩阵，改成分块算、边算边用，中间结果走 L1 不落 HBM。

Attention 的 O(N²) 显存问题

标准 Attention 的计算公式：

Q = x @ Wq
K = x @ Wk
V = x @ Wv
S = Q @ K^T / sqrt(d)
P = softmax(S)
O = P @ V

问题出在 S（QK^T）这个矩阵。seq=4096 时，S 的大小是 4096×4096=16M 个元素。FP16 下占 32MB 显存。看起来不大，但：

Prefill 阶段：每个 batch 都要算一次 S，batch=8 就是 256MB
Decode 阶段：每生成一个 token 都要读整个 KV Cache，seq=4096 时每个 token 读 57KB×4096=233MB

更严重的是：S 矩阵要写回 HBM，softmax 之后再读出来算 P×V。两次 HBM 读写，延迟爆炸。

实测数据（Qwen2.5-7B，910B 单卡，FP16，seq=4096）：

阶段	HBM 访问量	耗时占比
QKV 投影	38MB	12%
QK^T + Softmax	256MB	35%
P×V	233MB	28%
其他	-	25%

Attention 阶段占 63% 的时间，其中 80% 是 HBM 读写。

工程经验：7B 模型推理时，HBM 带宽利用率只有 35%。原因？S 矩阵写 HBM 再读出来，白白浪费带宽。开了 FlashAttention 后带宽利用率拉到 82%，吞吐直接翻倍。

FlashAttention 的 Tiling 策略

FlashAttention 的核心思路：不存 S 矩阵，改成分块算。

把 Q 分成 tile_q×d 的小块，K 分成 tile_k×d 的小块，V 分成 tile_k×d 的小块。每次只算一个 tile_q 的 Q 和一个 tile_k 的 K/V：

for each tile_q in Q:
    for each tile_k in K, V:
        S_tile = Q_tile @ K_tile^T
        P_tile = softmax(S_tile)
        O_tile += P_tile @ V_tile

关键点：S_tile 和 P_tile 走 L1 缓存，不落 HBM。

为什么能省显存？

标准 Attention：存整个 S 矩阵（N×N），显存 O(N²)
FlashAttention：只存 tile_q×tile_k 的 S_tile，显存 O(tile_q×tile_k)

seq=4096 时，标准 Attention 存 16M 元素；FlashAttention 用 tile_q=64、tile_k=64，只存 4K 元素。差了 4000 倍。

为什么能省 HBM 访问？

标准 Attention：Q→HBM，K→HBM，S→HBM，P→HBM，V→HBM，O→HBM，6 次大矩阵读写
FlashAttention：Q_tile 走 L1，K_tile 走 L1，S_tile 走 L1，P_tile 走 L1，V_tile 走 L1，只有 O 最后写一次 HBM

HBM 访问从 6 次降到 1 次。

昇腾 NPU 的内存层次

FlashAttention 在昇腾上的实现，必须理解内存层次：

HBM（高带宽内存）
  ↓ 带宽 1.2TB/s，延迟 200ns
L1 缓存（每个 AI Core 独立）
  ↓ 容量 1MB，带宽 ~10TB/s，延迟 10ns
L0A / L0B（Cube Unit 的输入缓冲区）
  ↓ 容量各 64KB
Cube Unit（矩阵乘单元）
  └─ MAC 阵列 16×16

关键限制：L1 只有 1MB。

FlashAttention 的 tile_q×tile_k×dtype 必须小于 L1 容量，否则溢出到 HBM，性能暴跌。

计算 tile 上限（FP16）：

S_tile = tile_q × tile_k × 2 bytes
P_tile = tile_q × tile_k × 2 bytes
Q_tile = tile_q × d × 2 bytes
K_tile = tile_k × d × 2 bytes
V_tile = tile_k × d × 2 bytes
O_tile = tile_q × d × 2 bytes

总和 < 1MB

假设 d=3584（Qwen2.5-7B），解出 tile_q × tile_k < 16K

实际选择：tile_q=64, tile_k=64, block_size=64（tile_q×tile_k=4096，留足空间给其他 buffer）

Cube/Vector 双缓冲流水线

昇腾的 Cube Unit 算矩阵乘，Vector Unit 算逐元素运算（scale、mask、softmax）。FlashAttention 要同时用 Cube 和 Vector，关键是流水线设计。

标准实现（无流水线）：

Cube: 算 Q×K^T → 等 Vector 算 softmax
Vector: 等 Cube 算完 → 算 softmax → 等 Cube 算 P×V
Cube: 等 Vector 算完 → 算 P×V

Cube 和 Vector 互相等，空转 50% 时间。

双缓冲流水线：

Cube: 算 Q1×K1^T → 算 Q1×K2^T → 算 P1×V1 → 算 P1×V2 → ...
Vector:           等 Q1×K1^T → 算 softmax1 → 等 Q1×K2^T → 算 softmax2 → ...

关键：Cube 算下一个 tile 时，Vector 在算上一个 tile 的 softmax。Cube 不等 Vector，Vector 不等 Cube。

实测交叠率：75%（Cube 和 Vector 同时工作的时间占比）

工程经验：Qwen2.5-7B 在 910B 上，开双缓冲流水线后吞吐从 67 tokens/s 涨到 89 tokens/s（+33%）。不开流水线，Cube 等 Vector 占 40% 时间。

tile_q / tile_k / block_size 参数调优

这三个参数直接决定性能。

tile_q：Q 的分块大小

太小（<32）：MAC 阵列填不满，吞吐低
太大（>128）：L1 装不下，溢出到 HBM，性能暴跌
最优值：64（填满 MAC 阵列，同时 L1 不溢出）

tile_k：K/V 的分块大小

太小（<32）：循环次数多，调度开销大
太大（>128）：L1 装不下
最优值：64

block_size：softmax 的分块大小（用于 Online Softmax）

影响 softmax 的数值稳定性
通常等于 tile_k

实测数据（Qwen2.5-7B，910B 单卡，seq=2048）：

tile_q	tile_k	吞吐 (tokens/s)	L1 溢出
32	32	72	无
64	64	89	无
128	64	85	轻微
128	128	61	严重

tile_q=64、tile_k=64 最优。tile_q=128 时 L1 开始溢出，性能反而掉。

工程经验：tile_q 调优有个坑——不同 seq 的最优 tile_q 不同。seq<1024 时 tile_q=32 最快，seq>2048 时 tile_q=64 最快。我们做了动态选择：seq<1024 用 tile_q=32，seq>=1024 用 tile_q=64。

与 GPU FlashAttention 的差异

很多人以为"FlashAttention 是通用的，GPU 能跑 NPU 也能跑"，其实差异很大。

维度	GPU（NVIDIA）	NPU（昇腾）
L2 缓存	40-50MB（全局共享）	无（只有 L1）
L1 缓存	128KB/SM	1MB/AI Core
执行单元	SM（统一）	Cube + Vector（分离）
最大 tile_q	128	64（L1 限制）
流水线	硬件自动调度	软件显式编排