当前位置：首页 > news >正文

面试官：FlashAttention 的实现原理与内存优化方式？为什么能做到 O(N²) attention 的显存线性化？

news 2026/7/24 0:31:11

如果你最近刷到过“FlashAttention”，那你一定见过那句经典介绍：“它让传统 O(N²) 的 Attention，显存占用变成 O(N)。”

很多人平时也都用FlashAttention，但是很少有人能够讲清楚其中的原理。
今天我们就拆开讲清楚：

为什么普通 Attention 显存爆炸；
FlashAttention 究竟改了什么；
为什么它能在保持 O(N²) 计算量的同时，让显存线性化。

一、普通 Attention 的计算与内存瓶颈

标准的自注意力（Self-Attention）计算如下：

A t t e n t i o n ( Q , K , V ) = s o f t m a x ( Q K T d k ) V Attention(Q, K, V) = softmax\left(\frac{QK^T}{\sqrt{d_k}}\right)VAttention(Q,K,V)=softmax(dkQKT)V

假设输入序列长度为 N，特征维度为 d。

那么计算步骤：

计算相似度矩阵

S = QKᵀ → [N, N]

归一化

A = softmax(S)

加权求和

O = A * V

显存问题出在哪？

关键在于那一步S = QKᵀ。
它是一个N×N 的矩阵，会直接占据 O(N²) 的显存。

举个例子：
假设 N=4096，单精度浮点数 4 字节：

4096² × 4B ≈ 64 MB

而在多头 attention、batch 堆叠后，这个数会直接上百 MB。

再加上中间 softmax 的缓存与梯度，整个过程几乎炸显存。

二、FlashAttention 的核心思想

FlashAttention 的核心不是改公式，而是改计算顺序。论文题目里那句关键话非常准确：“An IO-aware exact attention algorithm.”

也就是说：

数学上结果一模一样；
但计算顺序被重排，
以最小化显存访问和缓存中间矩阵为目标。

普通实现流程：

QKᵀ → Softmax → Dropout → (Softmax * V)

问题是：

每一步都需要完整的 [N, N] 矩阵；
每层都要读写显存（global memory）；
Softmax 的数值稳定性还要额外缓存max与sum。

这些中间值不是算力瓶颈，而是IO 瓶颈。
GPU 大部分时间都在“搬运数据”，而不是“算”。

三、FlashAttention 的关键优化

FlashAttention 的思路非常巧妙：把 Attention 计算拆成小块（tiles），每次只在显存中保留局部块，并在块级别完成 softmax 的归一化与累加。

分块计算 QKᵀ

把 Q 和 K 按块划分：

Q = [Q₁, Q₂, ..., Q_M] K = [K₁, K₂, ..., K_M]

对于每个 query 块 Qᵢ：

依次读取每个 key 块 Kⱼ；
计算局部相似度矩阵 Sᵢⱼ = QᵢKⱼᵀ；
同时在寄存器中保留该块的最大值与和。

这样只需要存储一个 tile 的中间矩阵（比如 64×64），不会生成完整的 [N, N] 矩阵。

块内 Softmax 的数值稳定处理

为了保持数值精度，FlashAttention 在块内维护：

当前最大值mᵢ；
累积和lᵢ。

公式如下：

m i ( j ) = m a x ( m i ( j − 1 ) , m a x ( S i j ) ) l i ( j ) = e x p ( m i ( j − 1 ) − m i ( j ) ) ∗ l i ( j − 1 ) + s u m ( e x p ( S i j − m i ( j ) ) ) m_i^{(j)} = max(m_i^{(j-1)}, max(S_{ij})) l_i^{(j)} = exp(m_i^{(j-1)} - m_i^{(j)}) * l_i^{(j-1)} + sum(exp(S_{ij} - m_i^{(j)}))mi(j)=max(mi(j−1),max(Sij))li(j)=exp(mi(j−1)−mi(j))∗li(j−1)+sum(exp(Sij−mi(j)))