想象一家小有名气的餐厅「Flash 小馆」，主厨有个怪癖：他从不给菜品拍照片。每道菜做好端上桌，配料、火候、用量就全凭脑子里的两行账——本桌的「最高分」和「总分」。客人如果说不满意，主厨不慌不忙，照着小账本把菜重新做一遍，边做边估摸每种调料到底放多放少。这套「不存档、现做现算」的做法，正是 Flash Attention 反向传播的精髓。 大模型的注意力机制，本质上就是让每个词（顾客）看看其他所有词（菜品）「有多投缘」，再把相关内容按投缘程度加权汇总。标准做法会把所有「投缘分数」存成一张 N×N 的大表，反向传播时再翻出来用；而 Flash Attention 只在计算时临时用一用，转头就扔，只留下 m（最高分）和 l（总分）两行台账。当误差信号回流时，它按台账把注意力权重「重建」出来，再分块给 Q、K、V 三位调味师算账。这就是重计算、分块与数值稳定三者的精妙合奏。

岗位关联

读懂它，你就能把注意力训练的显存占用从 O(N²) 降到 O(N)，是手写高性能注意力 kernel、训练超长上下文大模型的必修课。

目录
  1. 可视化演示
  2. 显存困境
  3. 前向台账
  4. 反向总览
  5. Softmax梯度
  6. QKV梯度
  7. 数值稳定
  8. 算力对比
  9. AI 应用场景
  10. Python 演示
  11. 常见误区
  12. 面试题

可视化演示

展示FCFS、SJF、Round Robin三种CPU调度算法的执行过程，通过甘特图对比不同调度策略的等待时间和周转时间。

本节配有交互动效（进入视口后自动加载，约 3KB）。

FCFS先来先服务调度 甘特图展示进程执行顺序，每个色块代表一个进程的CPU占用时段 下方圆点表示进程队列，高亮为正在执行的进程 FCFS按到达顺序调度，简单但可能导致短作业等待 等待时间 = 开始时间 - 到达时间 FCFS是非抢占式调度，按进程到达顺序依次执行

显存困境

传统后厨的规矩是「做一道、拍一张」。每道菜出锅，立刻有人举着相机拍证件照，照片按桌号、菜名归档进仓库。客人 N 位，菜品 N 道，两两之间的「合拍程度」都要留档，于是仓库里堆了 N² 张照片。当 N 从几千涨到几万甚至几十万，仓库立刻爆炸——这就是标准 Attention 反向传播的隐忧：它必须保存完整的注意力矩阵 P。

为什么反向传播非要这些照片？因为误差信号回流时，需要知道「当初这桌顾客给每道菜投了多少关注」，才能把梯度公平地送回 Q、K、V。于是 P 矩阵不仅前向要用，反向还要原封不动地从显存里再读一遍。对一张 8192 长度的序列，P 矩阵就有 8192×8192 个元素，仅这一张「照片」就要占掉上百兆字节的显存。

如果不想让仓库爆掉，就不能按老规矩存照片。Flash 主厨决定：照片不拍了，只记两行台账。

P=softmax ⁣(QK⊤d)∈RN×N\displaystyle P=\mathrm{softmax}\!\left(\frac{QK^{\top}}{\sqrt{d}}\right)\in\mathbb{R}^{N\times N}
  1. 打分矩阵 S先让 Q 与 K 做内积，衡量每个顾客和每道菜的「投缘分」，除以 √d 防止维度变大方差过大。
    Sij=Qi⋅Kjd\displaystyle S_{ij}=\frac{Q_i\cdot K_j}{\sqrt{d}}
  2. softmax 归一化对每一行做 softmax，把投缘分变成「好感占比」，总和为 1。
    Pij=eSij∑keSik\displaystyle P_{ij}=\frac{e^{S_{ij}}}{\sum_{k}e^{S_{ik}}}
  3. 存储规模结果是一个 N×N 的稠密矩阵，显存随序列长度平方级增长。
    MP=O(N2)\displaystyle M_{P}=O(N^{2})

前向台账

Flash 主厨发现，客人真正在意的是「哪道菜最合口味」和「所有菜的总分值」。于是每到一桌，他只在账本上记两行数：m_i 表示本桌目前的最高评分，l_i 表示把分数换算成「好感指数」后的总分。至于每道菜具体得了多少好感分？算完就忘。

这正是 FlashAttention 前向在做的事：把序列切成小块，逐块计算评分 S，用当前块的最大值更新 m，用 e^{S-m} 的和更新 l，同时把输出 O 的累积器更新掉。P 矩阵本身不落显存，只在片上缓存里短暂停留。反向传播需要的全部「档案」，就只有 m 和 l 这两行。

有了这两行台账，反向传播时就能像按菜谱重做一样，把 P 完完整整地重建出来。

\displaystyle m_i=\max_{j}s_{ij},\qquadl_i=\sum_{j}\exp(s_{ij}-m_i),\qquadP_{ij}=\exp(s_{ij}-m_i)
  1. 行最大值 m记录本桌评分的最高分，它是后续数值稳定的保险丝。
    mi=max⁡jsij\displaystyle m_i=\max_{j}s_{ij}
  2. 未归一化权重 P每个评分减去最高分后再取指数，得分再高也不会突破 e^0=1。
    Pij=exp⁡(sij−mi)\displaystyle P_{ij}=\exp(s_{ij}-m_i)
  3. 行和 l把这一行的所有指数求和，反向传播时要靠它把 P 还原成真正的 softmax 概率。
    li=∑jPij\displaystyle l_i=\sum_{j}P_{ij}

反向总览

现在客人投诉了（误差 dO 传入）。Flash 主厨的应对分两轮：第一轮，按台账把菜重新炒一遍——重算每个分块的 P；第二轮，一边重炒一边给三位调味师记「用量偏差」——Q 管顾客口味，K 管菜品特征，V 管菜品内容。

反向传播遵循链式法则：从输出梯度 dO 出发，先还原 P（经由 softmax 梯度 dS），再沿着 S=QK^T/√d 这条线把梯度分别推到 Q、K，同时把 dV 直接从 P^T dO 得到。整个过程不需要翻任何照片，只需要 Q、K、V 本身和两份台账。

三位调味师各拿一张「用量偏差单」，但这单子不是从照片里找出来的，而是现场重算出来的。

∂L∂Q=dS K⋅1d,∂L∂K=dS⊤Q⋅1d,∂L∂V=P^⊤dO\displaystyle \frac{\partial L}{\partial Q}=\mathrm{d}S\,K\cdot\frac{1}{\sqrt{d}},\qquad\frac{\partial L}{\partial K}=\mathrm{d}S^{\top}Q\cdot\frac{1}{\sqrt{d}},\qquad\frac{\partial L}{\partial V}=\hat P^{\top}\mathrm{d}O
  1. 重算概率 P̂用台账 m、l 和分块评分 S，重建归一化注意力权重。
    P^ij=exp⁡(sij−mi)li\displaystyle \hat P_{ij}=\frac{\exp(s_{ij}-m_i)}{l_i}
  2. dQ 回推把 dS 与 K 块相乘，再乘缩放因子，得到 Q 的梯度片段。
    dQi=dSijKj⋅1d\displaystyle \mathrm{d}Q_i=\mathrm{d}S_{ij}K_j\cdot\frac{1}{\sqrt{d}}
  3. dK 回推把 dS 转置后与 Q 块相乘，得到 K 的梯度片段。
    dKj=dSij⊤Qi⋅1d\displaystyle \mathrm{d}K_j=\mathrm{d}S_{ij}^{\top}Q_i\cdot\frac{1}{\sqrt{d}}
  4. dV 直达dO 经由 P̂ 的转置直接回流到 V，不需要经过 softmax 梯度。
    dVj=P^ij⊤dOi\displaystyle \mathrm{d}V_j=\hat P_{ij}^{\top}\mathrm{d}O_i

Softmax梯度

要算账，先要知道「这一桌的注意力分布误差」是如何经过 softmax 被放大的。softmax 的本质是归一化，归一化有个特点：某一项变高，其他项必然被压低。这种「此消彼长」的连带效应，正是反向传播中最容易算漏的部分——你不能只盯着某一项自己的梯度，还要算上它对整行归一化的挤压力。

为此 FlashAttention 反向先算一个行辅助量 D_i，它代表「输出误差与输出本身的点积」。原始的 dP=dO V^T 只考虑了「输出直接依赖 P」这一条路径，而 D_i 捕捉了「分母 l_i 随 P 变化」的间接路径。两者结合，才能得到对 S 的真实梯度 dS。

这步就像先尝一口整桌菜的平均口味，再判断每一道菜该加盐还是加糖，避免把「归一化的连带影响」算错。

Di=∑jdOij∘Oij,dSij=P^ij((dOi⋅Vj)−Di)\displaystyle D_i=\sum_{j}\mathrm{d}O_{ij}\circ O_{ij},\qquad\mathrm{d}S_{ij}=\hat P_{ij}\left((\mathrm{d}O_i\cdot V_j)-D_i\right)
  1. 行辅助量 D把输出梯度 dO 与输出 O 逐元素相乘再按行求和，衡量归一化分母的连带影响。
    Di=∑jdOij∘Oij\displaystyle D_i=\sum_{j}\mathrm{d}O_{ij}\circ O_{ij}
  2. 原始梯度 dPdO 与 V 做矩阵乘，得到「输出直接依赖 P」的梯度。
    dPij=dOi⋅Vj\displaystyle \mathrm{d}P_{ij}=\mathrm{d}O_i\cdot V_j
  3. 修正与加权用 D 修正 dP，再乘上归一化权重 P̂，得到对打分 S 的完整梯度。
    dSij=P^ij(dPij−Di)\displaystyle \mathrm{d}S_{ij}=\hat P_{ij}\left(\mathrm{d}P_{ij}-D_i\right)

QKV梯度

有了 dS，就能给三位调味师各写一张偏差单。主厨在小小的灶台（SRAM）上一次只能同时处理几个分块，所以他的办法是：外层循环固定一个 K/V 块，内层循环扫描所有 Q 块，把当前块的梯度累加到对应的全局账页上。

分块计算的本质是「流式累加」：每个 (i,j) 块算出的 dV、dQ、dK 片段，立即加到全局梯度矩阵的对应位置。因为每个块只依赖重算出的 P 块，不需要看到全表，所以片上内存始终只放一个小分块，显存里只需要维护最终的梯度累积器。

就这样，一桌一桌地算，最后三张完整的「用量偏差单」就诞生了。

dVj=dVj+P^ij⊤dOi,dQi=dQi+dSijKj1d,dKj=dKj+dSij⊤Qi1d\displaystyle \mathrm{d}V_j=\mathrm{d}V_j+\hat P_{ij}^{\top}\mathrm{d}O_i,\qquad\mathrm{d}Q_i=\mathrm{d}Q_i+\mathrm{d}S_{ij}K_j\frac{1}{\sqrt{d}},\qquad\mathrm{d}K_j=\mathrm{d}K_j+\mathrm{d}S_{ij}^{\top}Q_i\frac{1}{\sqrt{d}}
  1. dV 累加每个 Q 块都会向同一个 K/V 块的 dV 贡献一份梯度，因此用加等号累加。
    dVj+=P^ij⊤dOi\displaystyle \mathrm{d}V_j \mathrel{+}=\hat P_{ij}^{\top}\mathrm{d}O_i
  2. dQ 累加遍历所有 K 块时，同一个 Q 块的 dQ 从多个 K 块收集梯度。
    dQi+=dSijKj⋅1d\displaystyle \mathrm{d}Q_i \mathrel{+}=\mathrm{d}S_{ij}K_j\cdot\frac{1}{\sqrt{d}}
  3. dK 累加同理，同一个 K 块从多个 Q 块收集梯度。
    dKj+=dSij⊤Qi⋅1d\displaystyle \mathrm{d}K_j \mathrel{+}=\mathrm{d}S_{ij}^{\top}Q_i\cdot\frac{1}{\sqrt{d}}

数值稳定

台账上那个「最高分 m」还有一个救命用途：防爆表。假设某道菜的好感分高达 1000，直接算 e^{1000}，任何浮点数都得爆成 inf，整个反向传播立刻坍塌成一堆 NaN。这不是理论风险——长序列训练中，早期未收敛的模型完全可能产生很大的 logits。

FlashAttention 的解决方式无比朴素：所有指数都减去本行最大值 m_i，让指数输入落在 (-∞,0] 区间，e^0=1 封顶，既不会溢出，精度也足够。关键在于：反向重算 P 时，必须使用前向记录的那一份 m，否则重算出的 P 和当初不一致，梯度就会失真。

所以那份台账不只是省显存，更是「安全重算」的保险丝。

P^ij=exp⁡(sij−mi)li,mi=max⁡jsij\displaystyle \hat P_{ij}=\frac{\exp(s_{ij}-m_i)}{l_i},\qquad m_i=\max_{j}s_{ij}
  1. 重算指数反向分块时，用当前块评分 s_{ij} 减去台账中的 m_i 再做指数。
    pij=exp⁡(sij−mi)\displaystyle p_{ij}=\exp(s_{ij}-m_i)
  2. 归一化除以台账中的 l_i，还原成与标准 softmax 完全一致的权重。
    P^ij=pijli\displaystyle \hat P_{ij}=\frac{p_{ij}}{l_i}
  3. 范围保证因为 m_i 是本行最大分，s_{ij}-m_i≤0，exp 结果必然落在 (0,1]。
    exp⁡(sij−mi)∈(0,1]\displaystyle \exp(s_{ij}-m_i)\in(0,1]

算力对比

有人会问：重新做一遍菜，岂不是多花时间？是的，FLOPs 确实增加了一些——反向要多算一次注意力得分和指数，整体训练浮点运算大约增加三分之一。但省下的显存和 HBM 带宽是决定性的，因为显存带宽往往是训练真正的瓶颈。

标准 attention 需要把 N² 的 P 写入显存再读回；Flash 只把 O、m、l 写回，反向也是块内重算。显存从 O(N²) 降到 O(N)，意味着序列长度可以成倍增长。对大模型训练而言，带宽就是生命线：少搬一次数据，往往比少算几个乘法更关键。

主厨用「多炒一遍菜」换来了「不用建仓库」，在大模型动辄百万 token 的今天，这笔账非常划算。

\displaystyle M_{\text{standard}}=O(N^{2}),\qquad M_{\text{flash}}=O(N);\qquadT_{\text{flash}}\approx 1.33\times T_{\text{standard}}
  1. 标准显存完整 P 矩阵让显存随 N 平方增长，序列一长就触顶。
    Mstandard=O(N2)\displaystyle M_{\text{standard}}=O(N^{2})
  2. Flash 显存只存 O、m、l 和梯度累积器，随 N 线性增长。
    Mflash=O(N)\displaystyle M_{\text{flash}}=O(N)
  3. 时间代价重算 P 使总 FLOPs 增加约三分之一，但不搬大矩阵反而快。
    Tflash≈1.33×Tstandard\displaystyle T_{\text{flash}}\approx 1.33\times T_{\text{standard}}

AI 应用场景

大模型训练工程师的显存急救

阶段：训练优化

解决的问题：训练 128K 上下文时，标准 Attention 的 N² P 矩阵会瞬间耗尽 HBM。FlashAttention 反向传播通过重算 P、只保留 O/m/l，把注意力显存从 O(N²) 降到 O(N)，让超长序列训练在单卡上成为可能。

具体案例：一位工程师在训练 7B 模型时，上下文长度一直卡在 4K。改用基于 FlashAttention 反向传播的 kernel 后，长序列训练不再需要保存完整注意力矩阵，配合激活检查点，单卡即可扩展到 128K 上下文。

CUDA/Triton kernel 开发者实现自定义注意力

阶段：推理与部署优化

解决的问题：手写融合注意力 kernel 时，最大的难点不是前向，而是反向如何在分块框架下重建 softmax 中间结果。理解 m/l 台账和 dS 修正项后，开发者才能写出既省显存又数值稳定的反向 kernel。

具体案例：开发者在 Triton 中实现 Flash Attention 反向，每个 kernel 只加载一个 Q 块和 K/V 块，内部重算 P、累积 dQ/dK/dV，最终训练吞吐比标准 PyTorch 实现提升约 2-3 倍。

Python 演示

python79 行
import numpy as np  def flash_fwd(Q, K, V, block=2):    # 分块前向: 只记录 m/l 台账, 不保存完整 P    N, d = Q.shape    O = np.zeros_like(Q)    m = np.full(N, -1e9)    l = np.zeros(N)    scale = 1.0 / np.sqrt(d)    for i in range(0, N, block):        for j in range(0, N, block):            s = (Q[i:i + block] @ K[j:j + block].T) * scale            m_new = np.maximum(m[i:i + block], s.max(axis=1))            p = np.exp(s - m_new[:, None])            corr = np.exp(m[i:i + block] - m_new)            l[i:i + block] = l[i:i + block] * corr + p.sum(axis=1)            O[i:i + block] = O[i:i + block] * corr[:, None] + p @ V[j:j + block]            m[i:i + block] = m_new    return O / l[:, None], m, l  def flash_bwd(Q, K, V, dO, block=2):    # 反向: 分块重算 P, 累加 dQ/dK/dV    N, d = Q.shape    O, m, l = flash_fwd(Q, K, V, block)    D = (dO * O).sum(axis=1)    dQ = np.zeros_like(Q)    dK = np.zeros_like(K)    dV = np.zeros_like(V)    scale = 1.0 / np.sqrt(d)    for i in range(0, N, block):        for j in range(0, N, block):            s = (Q[i:i + block] @ K[j:j + block].T) * scale            p = np.exp(s - m[i:i + block, None])            p = p / l[i:i + block, None]            dV[j:j + block] += p.T @ dO[i:i + block]            dp = dO[i:i + block] @ V[j:j + block].T            ds = p * (dp - D[i:i + block, None])            dQ[i:i + block] += ds @ K[j:j + block] * scale            dK[j:j + block] += ds.T @ Q[i:i + block] * scale    return dQ, dK, dV  def standard_bwd(Q, K, V, dO):    # 标准反向: 保存完整 softmax, 用于对比验证    d = Q.shape[1]    scale = 1.0 / np.sqrt(d)    S = (Q @ K.T) * scale    m = S.max(axis=1, keepdims=True)    P = np.exp(S - m)    l = P.sum(axis=1, keepdims=True)    P = P / l    O = P @ V    D = (dO * O).sum(axis=1, keepdims=True)    dP = dO @ V.T    dS = P * (dP - D)    dQ = dS @ K * scale    dK = dS.T @ Q * scale    dV = P.T @ dO    return dQ, dK, dV  def main():    rng = np.random.default_rng(0)    N, d = 6, 8    Q = rng.normal(size=(N, d))    K = rng.normal(size=(N, d))    V = rng.normal(size=(N, d))    dO = rng.normal(size=(N, d))    dQ, dK, dV = flash_bwd(Q, K, V, dO, block=2)    gQ, gK, gV = standard_bwd(Q, K, V, dO)    print('dQ match:', np.allclose(dQ, gQ, atol=1e-9))    print('dK match:', np.allclose(dK, gK, atol=1e-9))    print('dV match:', np.allclose(dV, gV, atol=1e-9))  if __name__ == '__main__':    main()
目的：用 NumPy 实现 FlashAttention 前向的 m/l 台账记录与反向的分块重算，并和标准 Attention 反向梯度对比，验证公式正确性。
  1. O, m, l = flash_fwd(Q, K, V, block)前向分块扫描，只返回输出和两行台账，P 不落显存。
  2. D = (dO * O).sum(axis=1)计算行辅助量 D，捕获 softmax 归一化的连带效应。
  3. p = np.exp(s - m[i:i + block, None])反向分块中用台账 m 重算未归一化权重，防止指数溢出。
  4. ds = p * (dp - D[i:i + block, None])用 D 修正原始梯度，得到对打分 S 的真实梯度。
  5. dQ[i:i + block] += ds @ K[j:j + block] * scale分块流式累加，每个 Q 块从所有 K 块收集梯度。
  6. print('dQ match:', np.allclose(dQ, gQ, atol=1e-9))与标准反向对比，验证分块重算的梯度完全一致。

常见误区

误区反向传播计算梯度时，softmax 的权重 P 必须完整保存在显存里，不然没法算 dV 和 dS。

正解FlashAttention 证明这是「照片依赖」。反向只需按 Q、K 和台账 m/l 重算每个分块的 P，算完即弃。就像主厨凭菜谱重做一遍菜，不需要当初的菜品照片。显存里始终只有一个小分块和梯度累积器，整体占用从 O(N²) 降到 O(N)。

为何常见：标准 Attention 的实现把 P 矩阵显式存下来，让人误以为这是数学上的必然，其实只是「存储换方便」的实现选择。

误区反向重算 P 要重新做一遍矩阵乘和 exp，太浪费了，不如直接存 P。

正解重算确实让 FLOPs 增加约三分之一，但大模型训练的瓶颈往往是 HBM 带宽而不是算力。省掉 N² 矩阵的写入和读回，节省的带宽远大于多算的浮点运算，实际训练反而更快。这是典型的「用计算换访存」。

为何常见：只看 FLOPs 数字会得出错误结论，忽略了显存带宽这个真正的瓶颈。

误区m 和 l 只是前向为了数值稳定和归一化临时用的，反向传播中没它们什么事。

正解反向重算 P 时，必须用前向保存的 m 减掉再取 exp，再用 l 归一化。如果不用同一份 m，重算出的 P 会与实际计算路径不一致，梯度就会出错；l 更是把未归一化权重还原成 softmax 概率的关键。

为何常见：前向推导中 m/l 看起来只是中间变量，容易忽略它们跨过了前向与反向的整个时间边界。

误区为了省显存而重算 P，得到的梯度应该和标准 Attention 有偏差，训练效果会打折。

正解只要重算时使用相同的 m/l 台账，反向的 dS、dQ、dK、dV 在数学上与标准 Attention 完全一致，不存在近似误差。分块只是改变计算顺序和存储位置，不改变数值结果。

为何常见：把「重计算」和「近似估计」混为一谈，其实重算是精确复现同一个中间量。

面试题

Q1. FlashAttention 反向传播时，需要重建以下哪个中间矩阵？

单选题 · easy

  • A.完整 QK^T 打分矩阵
  • B.归一化注意力权重 P← 正确答案
  • C.嵌入层输出向量
  • D.前馈网络权重矩阵

答案：B。反向传播算 dV 和 dS 都需要注意力权重 P。FlashAttention 不保存 P，而是用 m/l 台账重算出来。QK^T 打分矩阵是中间态，不需要整体重建。

Q2. FlashAttention 前向为反向留下了哪两条台账记录？

单选题 · easy

  • A.行最大值 m 与行和 l← 正确答案
  • B.列最小值与列和
  • C.完整注意力矩阵
  • D.梯度累积器数值

答案：A。m 是本行评分最大值，l 是指数求和。反向重算 P 时用 m 防溢出、用 l 做归一化，两者缺一不可。

Q3. 反向重算注意力权重时，为什么要先用 m_i 减去评分？

单选题 · mid

  • A.防止 exp 溢出为无穷大← 正确答案
  • B.加快矩阵乘法速度
  • C.减少梯度的随机噪声
  • D.让分块可以完全并行

答案：A。评分可能很大，直接 exp 会爆成 inf。减去本行最大值 m_i 后指数输入落在 (-∞,0]，exp 结果不超过 1，从而保证数值稳定。

Q4. 反向传播中 D_i = sum(dO_i * O_i) 的作用是什么？

单选题 · mid

  • A.计算输出误差的自相关
  • B.捕获 softmax 归一化的连带效应← 正确答案
  • C.缩放学习率
  • D.更新 V 的梯度

答案：B。softmax 的归一化分母 l 随每个 P 变化，D_i 捕捉这条间接路径。修正 dP 后才能真正还原对 S 的梯度，避免只算直接路径而漏算连带效应。

Q5. FlashAttention 反向计算 dV 时，用哪个矩阵乘 dO？

单选题 · mid

  • A.Q 与 K 的内积矩阵
  • B.归一化注意力权重 P^T← 正确答案
  • C.V 自身的协方差矩阵
  • D.恒等单位矩阵 I

答案：B。O = P̂V，所以 dV = P̂^T dO。P̂ 是反向重算并归一化后的注意力权重，这一路不需要经过 softmax 梯度。

Q6. 相比标准 Attention，FlashAttention 反向因重计算导致训练 FLOPs 大约如何变化？

单选题 · hard

  • A.几乎不增加
  • B.减少约三分之一
  • C.增加约三分之一← 正确答案
  • D.增加整整一倍

答案：C。反向要多算一次注意力得分和 exp，整体训练浮点运算约增加三分之一。但因为省去 N² 矩阵的显存读写，实际训练往往更快。

Q7. FlashAttention 能训练超长序列的根本原因是什么？

单选题 · hard

  • A.使用混合精度训练
  • B.不保存完整 P 矩阵← 正确答案
  • C.采用梯度累积策略
  • D.减小了 batch size

答案：B。注意力显存从 O(N²) 降到 O(N) 的关键就是不再保存 P 矩阵，反向时按 m/l 台账重算。混合精度和梯度累积只是辅助手段，不是本质原因。

Q8. FlashAttention 反向传播需要哪些输入数据？

多选题 · mid

  • A.Q、K、V 矩阵← 正确答案
  • B.输出梯度 dO← 正确答案
  • C.台账 m 与 l← 正确答案
  • D.完整 P 矩阵

答案：ABC。反向重算 P 需要 Q、K 和 m，归一化需要 l，dO 是误差来源。完整 P 矩阵恰恰是 FlashAttention 刻意不保存的东西。

Q9. 关于 FlashAttention 反向传播的正确说法有哪些？

多选题 · mid

  • A.P 矩阵在反向时重算← 正确答案
  • B.显存占用依然 O(N²)
  • C.使用分块流式累加梯度← 正确答案
  • D.必须保存 S 矩阵

答案：AC。反向按分块重算 P，并把 dQ/dK/dV 片段累加到全局梯度矩阵，显存降为 O(N)。S 和 P 都不需要保存，只需台账 m/l。

Q10. 实现 FlashAttention 反向 kernel 时，数值稳定的关键处理包括哪些？

多选题 · extreme

  • A.前向记录行最大值 m← 正确答案
  • B.反向重算时再次减 m← 正确答案
  • C.用 l 归一化注意力权重← 正确答案
  • D.将 exp 结果存入 float16

答案：ABC。m 贯穿前向与反向，保证重算的 P 与真实路径一致；l 负责把未归一化权重还原成 softmax 概率。float16 存储不是数值稳定的手段，反而可能丢精度。