面试题：为什么 Attention 是 O(n²)？Linear Attention 和 FlashAttention 分别解决了什么问题？

2026年8月13日 · 阅读需 9 分钟

这是「每日一题」专栏的第四篇。每天一道面试题，后端 + AI 混合路线，从原理到代码一次讲清。

题目

大模型面试里有一道题越来越高频：

为什么标准 Self-Attention 的复杂度是 O(n²)？Linear Attention 和 FlashAttention 分别是怎么优化的？它俩解决的是同一个问题吗？

这道题考察的不是背结论，而是你对「注意力机制的数学本质」和「GPU 显存与 IO 层级」的双重理解。很多人张嘴就是 "FlashAttention 把复杂度降到了 O(n)"——这是错的：FlashAttention 的计算复杂度还是 O(n²)，它省的是 IO，不是计算量。

先给结论，再拆开讲。

一句话答案

方法	复杂度	本质	是否精确等价
标准 Attention	O(n²·d)	显式构造 n×n 注意力矩阵	—
Linear Attention	O(n·d²)	核函数分解 + 改变结合顺序	否（近似）
FlashAttention	O(n²·d)，但 IO 大幅下降	分块 + online softmax + kernel fusion	是（逐位一致）

三个关键点：

标准 Attention 的 n² 来自 QKᵀ：n 个 token 两两打分，产生 n×n 的注意力矩阵，这是 softmax 全局归一化在数学上绕不开的（除非换掉 softmax）。
Linear Attention 是真降复杂度，代价是换掉 softmax、用核函数近似，精度和表达力有损。
FlashAttention 没降计算复杂度，它只是把「读写 HBM」的次数降了几个量级，让 GPU 真正算满。

一、标准 Attention 为什么是 O(n²)？

回顾缩放点积注意力（上一篇《Transformer 的 Attention 机制为什么用缩放点积》讲过数学推导）：

scores = Q @ K.T * scale        # [n, n]  ← n² 的来源
attn   = softmax(scores)        # [n, n]
output = attn @ V               # [n, d]

拆开看每一步的复杂度（n = 序列长度，d = 每个 head 的维度）：

QKᵀ：[n, d] × [d, n] → [n, n]，复杂度 O(n²d)。
softmax：对 [n, n] 每一行做归一化，复杂度 O(n²)。
attn @ V：[n, n] × [n, d] → [n, d]，复杂度 O(n²d)。

三步里，第 1、2 步都要物化一个 n×n 的矩阵。序列长度 n 一旦变大（长文档、多轮对话、图像 patch、视频帧），这个 n² 会瞬间爆炸：

n = 2048 时，注意力矩阵约 2048×2048 ≈ 420 万个元素（单头单层）。
n = 32768 时，约 10.7 亿个元素，再乘上 batch、head、层数，显存直接爆掉。

为什么绕不开 n²？ 因为 softmax 的分母要对「所有 key」求和，导致每个 query 都必须和所有 key 交互一次——这就是全局注意力的代价。想优化，只有两条路：改注意力形式（Linear Attention 的思路），或别把 n×n 矩阵写进显存（FlashAttention 的思路）。

二、Linear Attention：把复杂度真正降到 O(n)

核心思想：换掉 softmax，用核函数分解

标准注意力：

attn(Q, K, V) = softmax(QKᵀ / √d) · V

softmax 里的 exp(qᵢ·kⱼ) 本质上是把点积映射成一个「正数权重」。这里藏着一个关键观察：

只要能找到某个特征映射 φ，使得 φ(qᵢ)ᵀ φ(kⱼ) ≈ exp(qᵢ·kⱼ)（即把相似度写成核函数形式），就能把注意力写成「先聚合 key/value，再点乘 query」的形式，从而交换矩阵乘法的结合顺序。

把 softmax 换成一般的相似度 sim(q, k)（比如线性核 q·k，或 φ(q)ᵀφ(k)）：

attnᵢ = Σⱼ sim(qᵢ, kⱼ) · vⱼ / Σⱼ sim(qᵢ, kⱼ)

当 sim(qᵢ, kⱼ) = φ(qᵢ)ᵀ φ(kⱼ) 时，利用矩阵乘法的结合律：

Σⱼ φ(qᵢ)ᵀ φ(kⱼ) vⱼ = φ(qᵢ)ᵀ · ( Σⱼ φ(kⱼ) vⱼᵀ )

关键：右边括号里 Σⱼ φ(kⱼ) vⱼᵀ 是一个 d×d 的矩阵，和 n 无关！可以先一次性算出来，再和每个 query 相乘。于是复杂度从 O(n²d) 降到了 O(n·d²)。当 n 远大于 d 时（长序列场景），这是质的提升。

代码

def linear_attention(Q, K, V, eps=1e-6):
    # φ：特征映射，保证非负即可，常用 elu(x)+1 或 relu
    Q = torch.nn.functional.elu(Q) + 1
    K = torch.nn.functional.elu(K) + 1

    # 先聚合 key 侧，绕开 [n, n] 注意力矩阵（O(n·d²) 的关键）
    KV = K.transpose(-2, -1) @ V            # [b, d, d]
    K_sum = K.sum(dim=-2, keepdim=True)     # [b, 1, d] 归一化分母

    num = Q @ KV                            # [b, n, d]
    den = Q @ K_sum.transpose(-2, -1)       # [b, n, 1]
    return num / (den + eps)

代价：降复杂度不是免费的

Linear Attention 的坑也很明确，面试要主动说出来：

精度损失：换掉 softmax 后，注意力不再是「全局归一化的概率分布」，表达力下降，长文本上可能需要更多层/更大模型补偿。
d² 项放大：head 维度 d 一旦变大，O(n·d²) 里的 d² 会变成新瓶颈，所以 linear attention 通常配小 head 维度。
训练不稳定：φ 的输出可能很大，需要数值稳定处理（如上文的分母归一化）。
工程生态不成熟：FlashAttention 对标准 softmax 的 kernel 高度优化，linear attention 的算子还不够成熟，实际收益取决于实现。

代表工作：Transformer 原班人马的 Performer（随机特征映射近似 softmax）、Linformer（低秩投影），以及 RWKV / RetNet / Mamba 这条线性注意力与状态空间模型路线。

三、FlashAttention：复杂度没降，但把 IO 打穿了

先说结论：FlashAttention 不是数学变体

很多人以为 FlashAttention 是「更快的注意力近似」，错。FlashAttention 在数学上和标准 softmax attention 完全等价（结果逐位一致），它优化的是「在 GPU 上怎么算」，不是「算什么」。

GPU 的内存层级：为什么「读」比「算」更慢

GPU 有两级关键内存：

HBM（高带宽显存，如 H100 的 80GB）：大，但相对慢，带宽约 3.35 TB/s。
SRAM（片上共享内存）：小（每个 SM 只有几百 KB），但极快，带宽是 HBM 的 10 倍以上。

标准 attention 的实现是「算子分开」的：先算 QKᵀ 写回 HBM，再读出来做 softmax 写回 HBM，再读出来乘 V 写回 HBM。这个过程中，n×n 的注意力矩阵被反复从 HBM 读写，而 GPU 的算力大部分时间在等数据——瓶颈是 IO，不是 FLOPS。

三个优化组合拳

分块（Tiling）：把 Q、K、V 切成小块，每次只把一小块装进 SRAM，避免整个 n×n 矩阵落进 HBM。
Online Softmax：softmax 需要先求全行最大值再归一化，但可以「边算边更新」——维护 running max 和 running sum，最后再统一 rescale，无需看到完整一行即可算，且数值稳定。
Kernel Fusion：把「QKᵀ → softmax → ×V」融合成一个 CUDA kernel，中间结果全程留在 SRAM，只把最终输出写回 HBM。

伪代码：

# FlashAttention：分块 + online softmax，结果与标准 attention 逐位一致
for i in blocks(Q):                     # 遍历 Q 的分块
    m_i = full(block_rows, -inf)        # running max
    l_i = zeros(block_rows)             # running sum(exp)
    o_i = zeros(block_rows, d)          # 输出累积

    for j in blocks(K, V):              # 每次只把 K_j、V_j 一小块 load 进 SRAM
        s = Q_i @ K_j.T * scale         # [B_r, B_c] 小块，只留在 SRAM，不写 HBM
        m_ij = s.max(dim=-1)            # 当前块的行最大值
        m_new = maximum(m_i, m_ij)      # 更新 running max
        p = exp(s - m_ij)               # 以块内 max 做数值稳定的 exp
        l_ij = p.sum(dim=-1)            # 当前块的 exp 求和

        # 用 exp(max 差值) 把旧累积 rescale 到新尺度，再累加当前块
        l_new = exp(m_i - m_new) * l_i + exp(m_ij - m_new) * l_ij
        o_i = exp(m_i - m_new)[..., None] * o_i + (exp(m_ij - m_new)[..., None] * p) @ V_j
        m_i, l_i = m_new, l_new

    o_i = o_i / l_i[..., None]
    write o_i to HBM                   # 每块只写回一次

收益

省显存：不物化 n×n 矩阵，显存从 O(n²) 降到 O(n)，这是它能训长序列（32k、128k）的根本原因。
更快：把 attention 从「IO 瓶颈」变成「算力瓶颈」，实测 2~4 倍加速，且序列越长越明显。
精确：结果和标准实现逐位一致，可以无缝替换，不改变模型行为。

后续还有 FlashAttention-2（减少非矩阵乘开销、更好的并行策略）、FlashDecoding（推理场景长序列解码优化）。

四、一张表分清三者的关系

维度	标准 Attention	Linear Attention	FlashAttention
计算复杂度	O(n²d)	O(nd²)	O(n²d)（不变）
显存占用	O(n²)	O(n)	O(n)
是否精确	是（基准）	否（近似）	是（逐位一致）
优化维度	—	算法（换 softmax）	工程（IO）
长序列适用	差	好（n 大 d 小时）	好
典型代表	原始 Transformer	Performer / RWKV / Mamba	FlashAttention-2 / FlashDecoding

五、面试话术（怎么答这道题）

按这个节奏答，基本能拿满分：

先分清两类优化：「这道题问的是两个不同维度的优化。Linear Attention 是算法层面的，把 O(n²) 降到 O(n)；FlashAttention 是工程/IO 层面的，复杂度没变，但显存和速度大幅提升。」
讲 Linear Attention：核函数分解 + 结合律交换，复杂度 O(n·d²)，代价是精度和表达力下降。
讲 FlashAttention：GPU 内存层级（HBM 慢、SRAM 快）→ 标准实现反复读写 HBM → 分块 + online softmax + fusion → 结果精确等价、显存 O(n)、加速 2~4 倍。
收尾拔高：「两者不冲突，生产里常一起用——线性注意力改模型结构，FlashAttention 优化训练/推理的算子和显存。」

总结

问题	答案
n² 哪来的？	QKᵀ 构造 n×n 注意力矩阵，softmax 又依赖全局归一化
Linear Attention 怎么降？	核函数分解 + 改变矩阵结合顺序，O(n²d) → O(n·d²)
FlashAttention 降复杂度了吗？	没有，还是 O(n²d)，它优化的是 HBM 与 SRAM 之间的 IO
两者能一起用吗？	能，一个是算法一个是工程，正交

一道题把注意力机制的「数学本质」和「GPU 工程」两条线都串起来了。

明日预告：后端题 —「消息队列怎么保证消息不丢？」

题目​

一句话答案​

一、标准 Attention 为什么是 O(n²)？​

二、Linear Attention：把复杂度真正降到 O(n)​

核心思想：换掉 softmax，用核函数分解​

代码​

代价：降复杂度不是免费的​

三、FlashAttention：复杂度没降，但把 IO 打穿了​

先说结论：FlashAttention 不是数学变体​

GPU 的内存层级：为什么「读」比「算」更慢​

三个优化组合拳​

收益​

四、一张表分清三者的关系​

五、面试话术（怎么答这道题）​

总结​

题目