想象你是全校最用功的图书管理员。学校有个巨大的中央仓库（HBM），里面存着所有书籍，容量惊人，但每次取一趟书都要穿越半个校园、搭乘电梯、过三道安检。而你面前那张小办公桌（SRAM），地方虽小，却伸手可及，拿放几乎零成本。如果你需要一个“注意力机制”——反复参考整座仓库的资料来完成一页重点摘要——你会怎么做？传统做法是：每写一行总结，就搬一摞书上桌，翻完又搬回去；下一步骤再搬一趟。来回折腾，一天下来精疲力竭，真正写字的力气没花多少，时间全耗在路上。 Flash Attention 的思路恰好相反：把资料合理地切成小份，一次只搬一小捆到桌面，在桌面上完成全部处理和记录，直到不再需要为止。它没有改变你要算的东西，却把“路上的时间”省到了极致。这就是分块（Tiling）、在线 Softmax（Online Softmax）、重计算（Recomputation）、Kernel 融合这些招数共同完成的奇迹。本篇文章，就像跟着这位图书管理员走一遍“省力改革”的全过程。

岗位关联

掌握 Flash Attention 的底层逻辑，是你成为「IO 感知型」AI 系统工程师的分水岭——从此你不再只关心算得快不快，更关心数据搬得少不少。

目录
  1. 可视化演示
  2. 注意力内存瓶颈
  3. 分块计算思想
  4. Online Softmax
  5. IO复杂度分析
  6. 重计算策略
  7. Kernel融合优化
  8. AI 应用场景
  9. Python 演示
  10. 常见误区
  11. 面试题

可视化演示

动态演示Flash Attentio的核心机制，通过开始/执行/监控/完成四种模式展示关键概念与交互流程。

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

Flash Attentio：GPU内 GPU内存层级与HBM/SRA Online Softmax的 核心机制可视化 Flash Attentio整体架构与关键组件

注意力内存瓶颈

想象你受命整理一份全校重点文献关联表，N 本书两两之间都要判断“相关程度”。传统做法是：先在草稿区画一张 N×N 的大表格，把每个格子填满，然后搬回仓库储存；下一步要据此汇总每本书的“精华页码”，又得把整张表重新搬出来。N 一旦超过几千，这张大表格本身就能撑爆桌面——于是大量时间不是花在“判断相关程度”上，而是花在“搬表格进进出出”上了。

这个“N×N 的大表格”就是 Transformer 里的注意力矩阵。标准实现中，注意力分数矩阵 S = QKᵀ/√d 先写入 HBM，Softmax 后的概率矩阵 P 再写回 HBM，最后 P 与 V 相乘又要重新读取 P。每一步都在 HBM 与 SRAM 之间穿行。当序列长度 N 达到 8K、32K 甚至 128K 时，N×N 矩阵的存储和访存就是训练中最大的时间黑洞。

标准注意力的数学形式如下：

这四步的中间产物——N×N 大表格——被反复存入仓库又取出。Flash Attention 的核心念头就是：能不能让这张大表格根本不落到仓库？也就是让注意力计算全程发生在桌面（SRAM）上。这正是下一节要讲的分块策略。

Attention(Q,K,V)=softmax(QK⊤dk)V\displaystyle \text{Attention}(Q,K,V)=\text{softmax}\left(\frac{QK^{\top}}{\sqrt{d_k}}\right)V
  1. 计算 QKᵀ将查询矩阵 Q（N×d）与键矩阵 K 的转置（d×N）相乘，得到 N×N 的注意力分数矩阵，每个元素表示一个查询与一个键的相似度。
    QK⊤∈RN×NQK^{\top}\in\mathbb{R}^{N\times N}
  2. 缩放分数除以 √d_k，避免高维点积数值过大而使 Softmax 进入饱和区，保证梯度稳定。
    QK⊤dk\frac{QK^{\top}}{\sqrt{d_k}}
  3. Softmax 归一化对每一行做指数归一化，使每行权重之和为 1，得到注意力权重分布。
    softmax(x)i=exi−m∑jexj−m,m=max⁡jxj\text{softmax}(x)_i=\frac{e^{x_i-m}}{\sum_{j}e^{x_j-m}},\quad m=\max_j x_j
  4. 乘以 V 加权聚合用归一化权重对值矩阵 V 的每行做加权平均，得到最终输出。
    Attention(Q,K,V)=softmax(QK⊤dk)V\text{Attention}(Q,K,V)=\text{softmax}\left(\frac{QK^{\top}}{\sqrt{d_k}}\right)V

分块计算思想

既然仓库太远，那就使用“小推车策略”。每次只从仓库取一小捆“同批次”的资料放上小推车，在推车上完成资料的标注与汇总，把成果浓缩成几个数字带回桌面角落。下一批资料到来时，小推车依然地方充足，桌面永远不会被 N×N 的大表格塞满。

分块（Tiling）正是这样：将 Q、K、V 按行切成大小为 B 的块。外层循环遍历 Q 的每个块 Q_i，内层循环遍历 K、V 的每个块 K_j、V_j。块内计算出的 B×B 注意力分数矩阵被限制在 SRAM 中，从不写入 HBM。每个输出块 O_i 的数学形式如下：

每个小块都在桌面内完成矩阵乘、指数计算与加权求和，只有分子分母两个“浓缩数字”随推车返回。仓库（HBM）的取货次数被压到最低，桌面始终保持整洁。但这里有个关键问题：遍历过程中 m_i 是未知的，怎么保证累加结果的正确性？答案就在下一节的 Online Softmax。

Oi=∑jeSij−miVj∑jeSij−mi\displaystyle O_i=\frac{\sum_{j}e^{S_{ij}-m_i}V_j}{\sum_{j}e^{S_{ij}-m_i}}
  1. 块内分数矩阵对第 i 个 Q 块与第 j 个 K 块做矩阵乘法，得到 B×B 的局部注意力分数。
    Sij=QiKj⊤dkS_{ij}=\frac{Q_iK_j^{\top}}{\sqrt{d_k}}
  2. 全局行最大值m_i 是第 i 行所有块中的最大分数，需要在遍历全部 K、V 块的过程中动态维护，无法一次得到。
    mi=max⁡j, blockSijm_i=\max_{j,\ \text{block}}S_{ij}
  3. 分子累加指数化后的分数与对应的 V 块相乘并累加，分子逐块更新。
    num←num+eSij−miVj\text{num}\leftarrow\text{num}+e^{S_{ij}-m_i}V_j
  4. 分母累加同时累加所有指数项，作为最终的归一化系数。
    den←den+∑reSij−mi\text{den}\leftarrow\text{den}+\sum_{r}e^{S_{ij}-m_i}

Online Softmax

假设你登山时没有全局地图，却要统计“每个海拔高度的植被覆盖比例”。你一边爬一边在便签上写下当前见过的最高海拔 m 与对应的累计覆盖量 l。每见一座新山峰，先比较海拔：若更高，就把便签上旧的累计量按比例缩小；若不高，新山峰的贡献就按当前最高海拔来算。这样从头到尾只需一遍，就能得到准确的比例分布。

标准 Softmax 必须完整看到一整行数据才能确定最大值 m，因此要读取两遍。Online Softmax（在线 Softmax）只需一遍：每当一个新块到达，用局部最大值更新全局最大值，并用指数修正因子将过去累积的分子分母同时修正到新的尺度上。其更新规则如下：

两张便签 m 和 l 从第一座山峰一直带到山顶，从未回头看过已走过的路。这就是“在线”的含义：数据流过一遍，统计结果自动正确。数值稳定性也由此保证——每次指数运算都以当前最大值为中心，绝不会出现 e 的几千次方这种溢出灾难。

\displaystyle m_{\text{new}}=\max(m_{\text{old}},\;m_{\text{block}}),\qquadl_{\text{new}}=l_{\text{old}}\cdot e^{m_{\text{old}}-m_{\text{new}}}+l_{\text{block}}\cdot e^{m_{\text{block}}-m_{\text{new}}}
  1. 更新当前最高峰比较新块的行最大值与全局最大值，取较大者作为新的 m。
    mnew=max⁡(mold,  mblock)m_{\text{new}}=\max(m_{\text{old}},\;m_{\text{block}})
  2. 修正旧统计量若全局最大值上升，旧块累计的指数和必须乘以衰减因子 e^{m_old-m_new}，否则会高估旧块的贡献。
    lold⋅emold−mnewl_{\text{old}}\cdot e^{m_{\text{old}}-m_{\text{new}}}
  3. 加入新块贡献新块的指数和也按与最新最大值的差值对齐后累加进分母。
    lblock⋅emblock−mnewl_{\text{block}}\cdot e^{m_{\text{block}}-m_{\text{new}}}
  4. 分子同步修正输出的加权和（分子）也要乘以同样的衰减因子，保证分子分母尺度一致，最终归一化结果不变。
    O←O⋅emold−mnew+eSblock−mnewVblockO\leftarrow O\cdot e^{m_{\text{old}}-m_{\text{new}}}+e^{S_{\text{block}}-m_{\text{new}}}V_{\text{block}}

IO复杂度分析

效率的关键在“路”不在“脑”。图书管理员再聪明，如果取一次书要半小时，一天也只能干几趟活。芯片世界也一样：计算单元内部的运算（SRAM 内）快如闪电，而 HBM 访存则像是“跨校区取书”，慢上几个数量级。

GPU 有显式的内存层级：HBM 容量大（几十 GB）但带宽远低于片上 SRAM；SRAM 容量小（几十到几百 KB）但带宽极高。标准注意力的问题在于：N×N 矩阵要多次在 HBM 中写读，让 HBM 成为瓶颈。Flash Attention 的核心优化目标，就是尽量减少 HBM 访问次数。

把“搬整张 N×N 表格”变成“每次只用小推车带一小块数据”，路上的时间省下一个数量级。这就是为什么 Flash Attention 能“算得差不多，却快好几倍”。注意：这里的 O(N) 是相对“避免 N×N 矩阵落盘”的访存量级而言，实际论文复杂度还包含 SRAM 容量因子，但核心洞察不变——减少搬运，就是减少时间。

IOstandard=O(N2),IOflash=O(N)\displaystyle \text{IO}_{\text{standard}}=O(N^2),\qquad\text{IO}_{\text{flash}}=O(N)
  1. 标准实现 O(N²)注意力分数矩阵与概率矩阵都是 N×N，写一次、读一次、再读一次，HBM 访问量与 N² 成正比。
    IOstandard=O(N2)\text{IO}_{\text{standard}}=O(N^2)
  2. Flash 实现 O(N)N×N 矩阵从不落 HBM，HBM 只与 Q、K、V、O 这些 N×d 规模的矩阵交互，总访问量随 N 线性增长。
    IOflash=O(N)\text{IO}_{\text{flash}}=O(N)
  3. 更严谨的表达式论文给出精确复杂度 O(N²d²/M)，其中 M 为 SRAM 容量。当块大小与 M 适配时，常数因子远小于标准实现，在长序列场景效果尤其显著。
    IOflash=O ⁣(N2d2M)\text{IO}_{\text{flash}}=O\!\left(\frac{N^2d^2}{M}\right)

重计算策略

有些同学复习时保留每一张中间草稿纸，宿舍堆得连落脚的地方都没有；学霸则只记结论，需要中间步骤时在头脑里快速重推一遍。Flash Attention 选择了学霸路线——不囤草稿纸。

深度学习反向传播需要前向的中间结果（注意力概率 P 和分数 S）来求梯度。标准实现把这些中间结果存入 HBM，付出 O(N²) 显存代价。Flash Attention 前向时不保存 P 和 S，反向需要时在 SRAM 中按分块方式重算。显存占用对比如下：

不囤草稿纸，需要时在桌面（SRAM）上快速重推。省下的宿舍空间（HBM 显存）足以装下更长的序列，而“重推”的代价远小于再跑一趟仓库。这是典型的“以算换存”策略——但因为在 SRAM 里算，这笔交换非常划算。

Memorystandard=O(N2+Nd),Memoryflash=O(Nd)\displaystyle \text{Memory}_{\text{standard}}=O(N^2+Nd),\qquad\text{Memory}_{\text{flash}}=O(Nd)
  1. 标准显存占用N×N 的注意力分数矩阵与概率矩阵占主导，显存随序列长度平方级增长。
    Memorystandard=O(N2+Nd)\text{Memory}_{\text{standard}}=O(N^2+Nd)
  2. Flash 显存占用只保存 Q、K、V 与输出 O，都是 N×d 规模，显存线性增长。
    Memoryflash=O(Nd)\text{Memory}_{\text{flash}}=O(Nd)
  3. 重算代价反向时重新按分块方式计算 P_ij 与 S_ij，因在 SRAM 内完成，重算速度远快于从 HBM 读取同样数据。
    Pij=softmax(QiKj⊤dk)P_{ij}=\text{softmax}\left(\frac{Q_iK_j^{\top}}{\sqrt{d_k}}\right)

Kernel融合优化

普通职工做三件事要分别跑三趟仓库；技术娴熟的工匠则把“查资料、写批注、贴标签”合并成一次取书后连续完成的动作。Kernel 融合正是这样的工匠精神：把多个 GPU Kernel 合并成一个，让数据在片上完成全流程。

标准实现中，QKᵀ、Softmax、乘 V 是三个独立 GPU Kernel，每个 Kernel 启动都要从 HBM 读取和写回中间结果。Flash Attention 将其融合为一个 Kernel：从 HBM 读入 Q_i、K_j、V_j 后，在 SRAM 内连续执行矩阵乘、在线 Softmax 更新与 V 加权，整个块处理完毕才把输出写回 HBM。融合 Kernel 的单块更新流程如下：

三件小事合并成一次取书、一次完成的动作。GPU 的 Kernel 启动次数从三次降为一次，中间结果不再经历仓库往返。再配合重计算策略，前后向传播都能以“尽量少打扰仓库”的方式进行——这就是 Flash Attention 能在真实硬件上跑出接近理论峰值性能的原因。

Qi,Kj,Vj→SRAM内  Sij  →  mnew,lnew  →  Oinew\displaystyle Q_i,K_j,V_j\xrightarrow{\text{SRAM内}}\;S_{ij}\;\to\;m_{\text{new}},l_{\text{new}}\;\to\;O_i^{\text{new}}
  1. 读入块数据一次从 HBM 读入 Q 块、K 块、V 块，放入 SRAM。
    Qi,  Kj,  Vj→SRAMQ_i,\;K_j,\;V_j\to \text{SRAM}
  2. 块内矩阵乘计算 B×B 的注意力分数矩阵 S_ij。
    Sij=QiKj⊤/dkS_{ij}=Q_iK_j^{\top}/\sqrt{d_k}
  3. 在线更新统计量用 Online Softmax 同时更新最大值 m 与分母 l。
    mnew,  lnew←OnlineSoftmax(m,l,Sij)m_{\text{new}},\;l_{\text{new}}\leftarrow\text{OnlineSoftmax}(m,l,S_{ij})
  4. 更新输出并写回输出块 O_i 归一化后写回 HBM，只有这一趟“回家”访问。
    Oinew=numli,Oi→HBMO_i^{\text{new}}=\frac{\text{num}}{l_i},\quad O_i\to\text{HBM}

AI 应用场景

长序列大模型预训练

阶段：模型训练（前向传播 + 反向传播）

解决的问题：标准注意力在序列长度 8K 以上时显存随 N² 爆炸，训练吞吐被 HBM 带宽锁死。传统做法只能通过梯度检查点等妥协手段勉强支撑，无法高效利用硬件。

具体案例：NVIDIA Megatron-LM 与 Hugging Face 集成 FlashAttention 后，GPT 类模型训练序列从 2K 扩展到 128K，训练吞吐提升约 2-3 倍，显存占用下降一个数量级，使“超长上下文”成为大模型标配能力。

LLM 推理加速引擎

阶段：模型推理（自回归解码阶段）

解决的问题：推理时 KV Cache 随序列长度线性膨胀，访存密集导致解码延迟居高不下。传统解码每个 token 都要读取全部 KV Cache，HBM 带宽成为吞吐瓶颈。

具体案例：vLLM 等推理框架采用 PagedAttention（Flash Attention 的分块思想变体）将 KV Cache 分块管理，配合融合 Kernel，推理吞吐提升数倍，并支持更长上下文下的高效并发服务。

长文档理解与智能问答

阶段：应用推理 / 数据处理

解决的问题：法律合同、医学书籍等长文档动辄数十万字，单次送入模型需要极大显存与等待时间，传统方案只能截断或分片，损失全局信息。

具体案例：某金融科技公司用 Flash Attention 在单张 A100 上处理 10 万 token 的招股说明书，实现全文级问答与风险条款定位，耗时从分钟级降到秒级。

Python 演示

python57 行
"""简化版 Flash Attention: 分块计算 + Online Softmax"""import mathimport numpy as np  def flash_attention(Q, K, V, block_size=32):    """在 CPU 上模拟 Flash Attention 的核心流程"""    N, d = Q.shape    O = np.zeros((N, d))    l = np.zeros(N)                # 分母累加器    m = np.full(N, -np.inf)        # 行最大值累加器     for i in range(0, N, block_size):        Qi = Q[i:i + block_size]        nb = Qi.shape[0]        Oi = np.zeros((nb, d))        li = np.zeros(nb)        mi = np.full(nb, -np.inf)         for j in range(0, N, block_size):            Kj = K[j:j + block_size]            Vj = V[j:j + block_size]            Sij = Qi @ Kj.T / math.sqrt(d)      # 块内 QK^T            m_ij = Sij.max(axis=1)              # 块内行最大值            m_new = np.maximum(mi, m_ij)        # 在线更新全局最大值            alpha = np.exp(mi - m_new)          # 旧统计修正因子            Pij = np.exp(Sij - m_new[:, None])  # 块内 softmax 分子            Oi = Oi * alpha[:, None] + Pij @ Vj            li = li * alpha + Pij.sum(axis=1)            mi = m_new         O[i:i + block_size] = Oi / li[:, None]    return O  def standard_attention(Q, K, V):    """标准注意力，作为正确性对照"""    d = Q.shape[-1]    S = Q @ K.T / math.sqrt(d)    P = np.exp(S - S.max(axis=-1, keepdims=True))    P = P / P.sum(axis=-1, keepdims=True)    return P @ V  if __name__ == "__main__":    np.random.seed(0)    N, d = 128, 16    Q = np.random.randn(N, d)    K = np.random.randn(N, d)    V = np.random.randn(N, d)     out_flash = flash_attention(Q, K, V, block_size=32)    out_std = standard_attention(Q, K, V)    err = np.abs(out_flash - out_std).max()    print(f"最大误差: {err:.2e}")    assert err < 1e-6, "Flash Attention 与标准实现不一致!"    print("验证通过: 分块 + Online Softmax 结果正确")
目的：用 NumPy 在 CPU 上完整实现 Flash Attention 的两大核心——分块遍历与 Online Softmax，并用标准注意力验证数值一致性。
  1. m = np.full(N, -np.inf)用负无穷初始化行最大值，保证第一个块到达时直接接管统计。
  2. Sij = Qi @ Kj.T / math.sqrt(d)块内矩阵乘法，只计算 B×B 的局部注意力分数。
  3. m_new = np.maximum(mi, m_ij)在线更新全局最大值，这是 Online Softmax 的“便签”。
  4. alpha = np.exp(mi - m_new)计算旧统计量的修正因子，最大值变化时按比例缩放历史积累。
  5. Oi = Oi * alpha[:, None] + Pij @ Vj分子累加：旧输出乘以修正因子，再加上新块的加权 V。
  6. O[i:i + block_size] = Oi / li[:, None]所有块处理完后一次归一化，得到最终输出。

常见误区

误区以为 Flash Attention 通过简化矩阵乘法让 FLOPs 降低，所以训练变快。

正解Flash Attention 完全保留了标准注意力的数学结果，算术运算量几乎没有减少。加速的真正来源是 HBM 访存减少——不再把 N×N 中间矩阵写进内存又读出来。就像图书管理员省下的不是翻阅时间，而是路上搬运时间。

为何常见：初看“分块”二字，容易联想到降采样或近似计算，误以为它对矩阵乘做了减法。

误区以为把每块的局部 Softmax 结果简单拼接在一起，就等于全局 Softmax。

正解每块的局部 Softmax 分母基于局部最大值，无法直接合并。必须在合并时用全局最大值进行修正——动态维护 m 与 l，这正是 Online Softmax 的核心意义。忽略这一层，结果会系统性偏大。

为何常见：“分而治之”的朴素直觉让人以为子结果天然可拼接，没意识到 Softmax 分母的全局依赖性。

误区以为块越小，SRAM 压力越小，Flash Attention 效果越好。

正解块太小意味着内层循环迭代次数增多，K、V 块会被重复加载更多次，HBM 访问反而上升；同时 GPU 并行粒度不足，Kernel 启动与同步开销放大。块大小必须权衡 SRAM 容量与重复访存代价，通常取 64~128 行。

为何常见：只看到“减少 HBM 访问”这一面，忽略了分块粒度本身会影响访问模式和并行效率。

误区反向传播要重新计算前向结果，训练速度一定会减半。

正解重算发生在 SRAM 内，带宽极高，而且重算的是小块矩阵乘，比从 HBM 读取海量中间矩阵快得多。实际训练中，Flash Attention 配合重计算往往比“省算力但费访存”的标准实现更快。

为何常见：低估了 SRAM 的带宽优势，把“算两次”直接等同于“慢两倍”。

误区以为只要加大 SRAM，就能彻底消除 HBM 访问，Flash Attention 就没有存在必要了。

正解SRAM 在芯片上极其昂贵，容量通常在几十到几百 KB 级别。Flash Attention 的设计目标是在给定 SRAM 规模下，用分块和在线统计把访存最小化，而不是清零。硬件约束永远存在，IO 感知的算法设计才有价值。

为何常见：把“缓存越大越好”的朴素想法直接套到芯片架构上，忽略了片上存储的物理成本与散热极限。

面试题

Q1. Flash Attention 主要解决什么问题？

单选题 · easy

  • A.减少乘法的计算次数
  • B.减少HBM与SRAM的搬运← 正确答案
  • C.提高浮点运算的精度
  • D.替代所有的归一化层

答案：B。Flash Attention 的数学结果与标准注意力完全一致，FLOPs 并没有减少。它的核心贡献是减少 GPU HBM 与 SRAM 之间的数据搬运——就像图书馆管理员省下的是搬书路程而非阅读时间。B 准确描述了这一点。

Q2. 标准注意力中 N×N 的注意力分数矩阵主要存放在哪里？

单选题 · easy

  • A.SRAM 片上缓存
  • B.寄存器文件
  • C.HBM 高带宽显存← 正确答案
  • D.本地磁盘仓库

答案：C。标准实现中，N×N 的 S 和 P 矩阵太大，必须存放在容量大但访问慢的 HBM（显存）中。SRAM 太小放不下整个矩阵。选择 C 是因为中间矩阵的读写只能交给大容量 HBM，这也正是访存瓶颈的来源。

Q3. Online Softmax 的核心作用是什么？

单选题 · mid

  • A.分块时动态维护行最大值和指数和← 正确答案
  • B.把softmax替换成线性函数
  • C.减少需要计算的行数
  • D.让矩阵乘法结果直接为整数

答案：A。Online Softmax 在遍历分块过程中，用更新公式动态维护每行的最大值 m 和指数和 l，保证不预先扫描整行也能得到正确 softmax。它没有改变计算行数，也不是近似替代，只是把两遍扫描压缩成一遍。

Q4. Flash Attention 能加速训练的原因包括哪些？

多选题 · mid

  • A.避免N×N矩阵反复读写HBM← 正确答案
  • B.利用SRAM完成块内计算← 正确答案
  • C.减少乘法运算量
  • D.使用更低精度的数据类型

答案：AB。Flash Attention 的加速来自 IO 优化——N×N 中间矩阵不再落 HBM（A），块内计算全部在 SRAM 完成（B）。它没有减少 FLOPs（C 错），也不依赖低精度（D 错）。类比：省的是搬书时间，不是看书时间。

Q5. 反向传播时 Flash Attention 通常采用哪种策略来节省显存？

单选题 · mid

  • A.保存所有中间注意力矩阵
  • B.压缩中间矩阵后进行存储
  • C.不存中间矩阵反向时重算← 正确答案
  • D.反向时只保存Q矩阵

答案：C。Flash Attention 采用重计算策略：前向时不保存 N×N 的注意力矩阵，反向传播需要梯度时，在 SRAM 里重新按分块计算它们。SRAM 中的重算速度远快于 HBM 读取，既省显存又不拖慢训练。

Q6. 下面哪个组件容量小但速度极快、位于 GPU 芯片内部？

单选题 · easy

  • A.HBM 高带宽显存
  • B.SRAM 片上高速缓存← 正确答案
  • C.磁盘存储阵列
  • D.网络传输通道

答案：B。SRAM 位于 GPU 芯片内部，容量小（几十到几百 KB）但带宽极高，是 Flash Attention 中“桌面工作台”的硬件原型。HBM 虽然叫“高带宽”，但容量大、访问延迟远高于 SRAM，对应文中的“大仓库”。

Q7. 关于 Flash Attention 的块大小选择，正确的说法有哪些？

多选题 · hard

  • A.块须能全部放进SRAM← 正确答案
  • B.块小则KV重复加载越多← 正确答案
  • C.块大则SRAM要求越高← 正确答案
  • D.块大小与头维度无关

答案：ABC。块必须能放入 SRAM（A）。块越小，内层循环要遍历更多 KV 块，重复加载次数越多（B）。块越大，需要 SRAM 一次性容纳的数据越多（C）。块大小与注意力头维度 d 强相关——B×d 的数据量直接决定 SRAM 占用（D 错）。

Q8. 标准 Softmax 与 Online Softmax 的关键差异是什么？

单选题 · mid

  • A.标准方法结果更精确
  • B.在线无需整行最大值即可← 正确答案
  • C.在线方法计算量翻倍
  • D.标准方法支持流式计算

答案：B。标准 Softmax 必须先扫描整行找到最大值 m 才能计算指数和，因此至少访问两遍数据。Online Softmax 通过动态更新 m 和 l，到达一个元素就能立刻累积贡献，一遍扫描即可完成。两者的数值结果等价。

Q9. Kernel 融合优化带来的收益有哪些？

多选题 · hard

  • A.减少GPU Kernel启动次数← 正确答案
  • B.中间数据不落HBM省访存← 正确答案
  • C.消除N×N矩阵持久化存储← 正确答案
  • D.增加了模型参数量

答案：ABC。融合 Kernel 把 QKᵀ、Softmax、×V 三个算子合并为一个，Kernel 启动次数从三次降为一次（A）。中间结果在 SRAM 内流转，不写回 HBM（B），N×N 矩阵不产生持久化存储（C）。参数量不因 Kernel 融合而改变（D 错）。

Q10. 论文中 Flash Attention 的 HBM 访问复杂度的严谨表达式是？

单选题 · extreme

  • A.O(N²d²/M)的量级← 正确答案
  • B.O(N³)的量级
  • C.O(1)常数量级
  • D.O(N²logN)量级

答案：A。Flash Attention 论文给出的 HBM 访问复杂度是 O(N²d²/M)，其中 M 为 SRAM 容量。直观上，块对数约为 (N/B)²，每块访问量约 M，而 B≈M/d，整理即得此式。虽然形式上仍是 N² 量级，但常数因子 d²/M 远小于标准实现的 N² 大矩阵搬运。