在线 Softmax 与 FlashAttention

需要会做矩阵乘法，并知道朴素 attention 的分数矩阵来自 $QK^T$。

本课是“在线 Softmax 与 FlashAttention”单元的入口。先证明 tile 计算不会改写 $QK^T$，再学习在线 softmax 如何在不保存完整分数矩阵时得到同一输出。

忽略缩放因子 $1/\sqrt{d_k}$，注意力分数为

若 $Q,K\in\mathbb{R}^{N\times d}$，则 $A\in\mathbb{R}^{N\times N}$。

朴素实现通常把这个完整 $N\times N$ 中间量写到高带宽显存 HBM，再读取它做 softmax 和后续乘 $V$。

当序列长度 $N$ 增大时，这个中间量的元素数按 $N^2$ 增长。

FlashAttention 的核心工程目标不是修改 attention 定义，而是减少这个中间量在 HBM 中的读写和物化。

矩阵乘法定义给出

这里 $Q_{i,:}$ 是 Q 的第 $i$ 行，$K_{j,:}^T$ 是 K 的第 $j$ 行转置成的列向量。

展开后

这条公式非常重要：任何分块方案只要让每个 $(i,j)$ 最终累加完全相同的 $d$ 个乘积，结果就不会改变。

所以正确性检查不需要迷信“FlashAttention”这个名字，只要回到每个输出元素的求和项即可。

把 Q 按 token 行切成两个块：

K 原本也按 token 行切：

转置后就成为列块：

因此

每个 $Q_iK_j^T$ 都是完整注意力矩阵中的一个矩形 tile。

课程用具体数字逐块计算，发现每个 tile 与朴素结果中相同位置的小块完全一致。

原因并不神秘。

设第 $i$ 个 Q 块覆盖行集合 $I_i$，第 $j$ 个 K 块覆盖行集合 $J_j$。

块乘法 $Q_iK_j^T$ 中任意元素仍是

这与朴素矩阵中 $A_{uv}$ 的定义一字不差。

分块改变的是：

分块没有改变的是：

因此，在相同浮点精度与舍入顺序假设下，分块结果与朴素结果相同；实际浮点实现可能有极小舍入差异，但算法语义等价。

HBM 容量大但访问代价高，片上 SRAM 容量小但带宽高。

FlashAttention 让线程块负责一个或若干 tile：

课程在结尾指出一个关键困难：softmax 是逐行的，但每个线程块一次只看到一行的一部分。

这意味着证明 $QK^T$ 分块等价还不够，还要证明局部 softmax 状态可以正确合并。

对一行分数 $s$，稳定 softmax 通常使用全局最大值

和归一化分母

若当前只看到一个分块 $s^{(b)}$，可先得到局部最大值 $m_b$ 与局部分母 $\ell_b$。

把旧状态与新块合并时，先更新

再把二者缩放到同一基准：

输出的加权和也做同样重标定，最终再除以 $\ell_{\text{new}}$。

因此完整 FlashAttention 正确性包含两层：

本条视频完整说明了第一层，并在结尾引出第二层；它不是对整个 FlashAttention 前向算法的全部证明。

注意力在忽略缩放、掩码和多头维度后，可先写成

视频用一个称为 SlimAttention 的方案说明：把 $Q$ 沿 token 行切成 tile，但在当前计算层次让 $K^\top$ 保持覆盖完整 key 轴。

问题是，分开计算每个 query tile，最后还能否严格得到原来的 $A$？

本课先证明线性代数等价，再讨论这一分块对 Softmax 与内存层次意味着什么。

设

则

第 $i,j$ 个元素为

图中选择一行 $q_i$ 和 $K^\top$ 的一列，也就是原 $K$ 的一行 $k_j$。

二者内积填回 $A_{ij}$。

只要分块后仍计算同一对 $q_i,k_j$，该元素的值就不会改变。

把 $Q$ 按行分成

其中 $Q_1$、$Q_2$ 分别包含若干完整 query 行。

这里没有沿特征维 $d$ 切断单个 token 向量。

每个 query tile 仍包含完整的 $d$ 维向量，因此可以与每个 key 做完整内积。

矩阵乘法对纵向拼接满足

左侧是一次计算完整 $QK^\top$。

右侧是分别计算两个行块，再按原顺序纵向拼接。

两者逐元素相同，不依赖近似，也不依赖浮点外的特殊假设。

在真实浮点执行中，不同 kernel 的归约顺序可能造成末位舍入差异；这里的“等价”指相同数学表达式，不承诺 bitwise identical。

$Q_1$、$Q_2$ 是沿行轴切分，所以各自产物 shape 为

它们都覆盖完整 key 列，只覆盖不同 query 行。

因此必须沿行轴拼接。

若错误地沿列轴拼接，shape 和 $A$ 的 query/key 语义都会改变。

判断拼接方向时，应始终追踪“被切的是哪一根轴”。

假设

将 $Q$ 切成两个 $2\times4$ tile：

则

按行拼接得到 $4\times4$，与完整乘法 shape 一致。

视频板书中的数字例子还逐项验证了对应元素相同。

视频随后给出内存示意：Q、K 位于较慢的主存层，执行单元把当前 Q tile 与所需 K 数据搬入更快的片上存储，再计算一个 A 的行块。

一个线程或线程块负责哪个 tile、K 是否一次完整驻留，都取决于硬件与实现。

数学证明只要求每个所需内积最终被计算，不能从证明本身推出某种固定调度一定最快。

注意力 Softmax 沿 key 轴逐行计算：

若一个 Q tile 与完整 $K^\top$ 相乘，所得行块覆盖每一行的全部 $N_k$ 个 key 分数。

于是该 tile 内可以直接得到每行的全局最大值 $m_i$ 与完整分母。

不必跨多个 K tile 逐次修正同一行的 Softmax 状态。

FlashAttention 通常同时对 query 轴与 key 轴分块。

单个分数 tile 只看到部分 key，因此没有一整行的全局最大值与分母。

它需要维护在线 Softmax 状态，在新 K/V tile 到来时重标定旧的部分结果。

视频中的 SlimAttention 行分块则让一个 Q tile 在该层次覆盖完整 K 轴，所以行归一化更直接。

但完整 K 是否能放进快速存储、是否需要在更低层继续切块，是实现层问题。

沿 Q 行分块永远保持矩阵乘法的数学正确性。

性能却取决于：

因此不能从“线性代数保证正确”直接跳到“所有长序列都更快”。

视频画面用 DRAM/SRAM 表示慢速与快速存储，并口述 CPU cache。

在 CPU 上更常见的具体层级是 DRAM 与 L1/L2/L3 cache；在 GPU 讨论中常见 HBM、shared memory、register。

它们共同表达“分层存储与块复用”这个概念，但容量、编程模型和搬运机制不同。

课程稿保留视频的抽象，不把板书中的 SRAM 标签当成所有 CPU 实现的字面硬件接口。

“SlimAttention”并不是一个仅凭名字就能唯一确定所有分块、缓存与 Softmax 细节的通用规范。

本课只解释视频所示的“Q 行分块、当前层次完整 K”路径。

阅读具体论文或代码时，还要核对：

对向量 $x=(x_1,\ldots,x_N)$，Softmax 定义为

直接实现需要先得到分母，才能写出每个 $p_i$。

安全 Softmax 还要先求全局最大值，避免指数溢出，于是看起来需要三次遍历。

在线 Softmax 的关键是：最大值变化时，把此前累计的分母转换到新的指数基准，从而在同一遍流式扫描中同时维护最大值与分母。

以视频中的

为例。

第一遍计算

第二遍再写出

在不知道 $S$ 之前，不能得到最终归一化结果。

因此朴素的独立数组实现至少有“归约 + 写出”两个逻辑阶段。

若某个 $x_i$ 很大，$e^{x_i}$ 可能超出浮点数可表示范围。

即使最终比例本应有限，中间指数也可能先变成正无穷，导致不定比值或 NaN。

Softmax 对整体平移不变：

所以可以选择 $c=\max_j x_j$。

视频示例的全局最大值为 5。

减去它后得到

此时最大的指数是 $e^0=1$，其他指数都位于 $(0,1]$，显著降低上溢风险。

画面也证实了 ASR 中漏掉的负号：移位向量是“[-4,-3,-2,0]”，不是“[4,-3,-2,0]”。

安全 Softmax 可按以下方式实现：

第二遍依赖第一遍的 $m$，第三遍又依赖第二遍的 $\ell$。

若把三步机械实现为三次读取输入，就会产生三遍数据访问。

在线 Softmax 处理到前 $t$ 个元素时维护

和

$m_t$ 是当前前缀最大值。

$\ell_t$ 不是原始指数和，而是所有已见元素相对于当前最大值 $m_t$ 的指数和。

这两个量足以表示已经扫描部分的安全归一化信息。

读入 $x_t$ 后，先更新

旧分母 $\ell_{t-1}$ 是以 $m_{t-1}$ 为基准。

要换成新基准 $m_t$，必须乘 $e^{m_{t-1}-m_t}$，于是

若新元素没有刷新最大值，旧分母系数为 1。

若新元素更大，旧的所有指数项会统一缩小到新基准。

设分块 a、b 分别维护

先取

再把两个局部分母都变换到基准 $m$：

这就是板书中用两个局部向量信息更新全局 max 与 sum 的公式。

它允许树形归约、线程块局部归约或流式扫描，不要求先单独遍历完整向量求最大值。

左块：

右块：

合并最大值为 $m=5$，所以

这正好等于直接对“[-4,-3,-2,0]”求指数和。

在线更新把“求最大值”和“求相对指数和”融合到第一次扫描：

第二遍再根据最终 $m,\ell$ 写出

所以独立 Softmax 数组实现从三次逻辑遍历变成两次。

少掉的是单独求最大值的那一遍，不是消除最终归一化。

在线状态始终以“当前已见最大值”为指数基准。

新最大值只会不减，所有指数的自变量都满足

因此不会为少一次遍历而恢复到直接计算巨大指数的不安全做法。

有限精度下，不同归约树仍会产生微小舍入差异，但算法维持了 max-shift 的稳定原则。

视频用 for 循环与 HBM 到 SRAM 搬运解释收益。

这是很好的数据流直觉，但工程上还要区分：

若安全 Softmax 已由高度融合 kernel 实现，不能机械地把“三个公式阶段”换算成“三次完整 HBM 往返”。

在线公式提供可融合性，真实性能仍需测量。

FlashAttention 不把完整分数行存到 HBM 后再跑独立 Softmax。

它在遍历 score tiles 时维护每行的运行最大值、分母和输出累加器。

新 tile 改变最大值时，不仅要重标定 $\ell$，还要重标定此前的输出部分和。

因此在线 Softmax 是 FlashAttention 分块等价性的核心组件，但实际 kernel 比本课的独立向量两遍算法更进一步融合了 $PV$。

视频把在线 Softmax 追溯到 NVIDIA 研究者 2018 年的工作，并说明它被 FlashAttention 使用。

本课聚焦 max/sum 合并公式，不把所有现代 kernel 的具体访存次数都等同于这份板书。

FlashAttention 不同版本、硬件后端和编译器可能采用不同 tile、warp 归约与流水方式，但要保持精确结果，都必须等价维护相应归一化状态。

Decoder-only Transformer 要保证位置 $i$ 不能看到未来位置 $j>i$。

传统写法先计算完整分数矩阵，再在上三角加 $-\infty$。

FlashAttention 按 tile 计算，不必真的先算出那些注定被掩码为零的完整块。

核心判断是：

忽略多头与缩放因子，原始分数为

因果掩码 $M$ 定义为

掩码后分数为

第 $i$ 行只有当前及历史 key 是有限值。

逐行 Softmax 为

对未来位置 $j>i$，

所以

随后输出

中，这些项都是 $0\cdot V_j$，不会贡献任何值。

把 query 行与 key 列都分块。

若某个块中的每一对 $(i,j)$ 都满足 $j>i$，该块所有掩码分数都是 $-\infty$。

整个 Softmax 权重 tile 都为零，进而

所以不必先计算分数再发现它为零。

运行时可以根据 tile 的位置关系预先判定并跳过。

对齐分块后，可以把 tile 分为三类：

只有第二类可以整块省略。

“省略上三角计算”不等于对角线边界块里完全没有 mask 判断。

作为编者补充，设 query tile 行区间为半开区间

key tile 列区间为

若

则最小 key 下标也大于最大 query 下标，整块完全位于未来，可直接跳过。

若区间跨过对角线，就不能用整块零替代。

视频以一个 query tile $Q_2$ 为例。

线程块把 $Q_2$ 与所需的 $K_1$ 搬入片上 SRAM，计算局部分数

这里的 tile 可以包含多个 token，尺寸由片上容量、数据类型和 kernel 配置共同决定。

$A_{21}$ 在片上立即参与 Softmax 状态更新与 $V_1$ 加权。

使用完成后即可丢弃，不必物化到 HBM。

这与“跳过上三角”是两项相关但不同的优化：

当前块得到的部分结果类似

但 $Q_2$ 还可能关注 $K_2,K_3,\ldots$ 中的合法位置。

因此它必须与后续 K/V tile 的贡献累加。

更重要的是，每个分数块只有局部最大值和分母。

新块改变运行最大值时，旧输出累加器也要按在线 Softmax 比例重标定，才能与完整一行 Softmax 等价。

考虑最早的一组 query，例如 $Q_1$，与很靠后的 $K_N$。

若二者对应的整个分数块都位于因果上三角，那么该块最终全为零。

此时可以同时省去：

收益不只是少做乘法，也包括少搬数据。

有人会担心：不计算这些项，Softmax 分母是否会改变？

不会。

被掩码项的指数本来就是

把零项从求和中删除不改变分母，也不改变最大值，因为合法行至少包含当前位置本身。

因此跳过是代数上的精确消元，不是近似稀疏化。

视频 ASR 把 “Decoder-only Transformer” 识别得不清楚，不能据此把课程理解成单 token decode。

在处理完整 causal 序列的训练或 Prefill 中，$Q$ 和 $K$ 都覆盖许多位置，确实存在大块上三角区域。

而标准自回归单 token decode 只计算最新 query，它可以关注所有已有 key，通常没有同样的完整上三角矩阵可跳。

因此该优化的显著场景主要是多 query 的 causal attention。

数学上用 $-\infty$ 表示 mask 最清楚。

kernel 中可能：

只要最终等价于无效位置的 Softmax 权重为零，具体表示可以不同。

视频以 FlashAttention v2 讲解外 Q tile 与因果上三角跳过。

这一“完全掩码 tile 不加载、不计算”的原则并不应理解为只有 v2 才能采用。

FlashAttention 不同版本、不同后端以及 PyTorch SDPA/Triton 等实现，可能有不同分块大小、对角线处理和调度策略。

判断某版本是否实际跳过哪些 tile，应查看目标 kernel 与运行配置。

FlashAttention 计算的数学目标仍是

它不把完整注意力分数矩阵写回 HBM，而是让局部分数 tile 在片上被消费。

当 key 轴也被切块时，一个 query tile 的最终输出依赖所有合法 K/V tiles。

循环顺序决定：究竟让当前 query/output tile 留在片上，还是让当前 K/V tile 留在片上。

FlashAttention-2 选择外层 Q、内层 K/V，视频用“Q/O 定、K/V 动”概括其访存收益。

完整 Q、K、V 与最终输出 O 通常位于 HBM。

分块计算会把当前需要的小块搬到寄存器或共享内存等片上存储。

完整分数矩阵

则不需要持久化到 HBM。

对 query block $Q_i$ 与 key block $K_j$，先计算

局部 $S_{ij}$ 参与在线 Softmax，再与 $V_j$ 相乘。

一旦它对在线统计量和输出累加器的贡献已经吸收，原始分数 tile 就可以释放。

因此循环顺序主要影响 Q、K、V、O 及在线状态怎样搬运，而不是 A 如何写回。

视频选 $Q_2$ 作为例子。

把它搬入片上存储后，与第一个 K block 计算

再读取 $V_1$，形成第一份局部输出贡献。

之后依次处理

内层循环结束时，当前 query 行已经看过全部合法 key。

设第一个分数块的局部最大值和分母为 $m_1,\ell_1$，第二块为 $m_2,\ell_2$。

若分别对两个块归一化再直接相加，会让每块权重各自和为 1，结果不等于完整行 Softmax。

必须维护运行状态

旧输出部分和也要按相同基准重标定。

对 $Q_2$，第一块产生 $S_{21}V_1$，第二块产生 $S_{22}V_2$。

更准确的在线更新形式可写为

其中 $\tilde O$ 是尚未除以最终分母的加权和。

等所有 K/V blocks 处理完，再除以最终 $\ell$ 得到 $O_2$。

$O_2$ 不是一次 $S_{21}V_1$ 就完成的。

它与在线最大值 $m_2$、分母 $\ell_2$ 一起构成跨 K/V blocks 的运行状态。

如果每处理一个 K/V block 都把这组状态写回 HBM，下一轮又读回来，会产生重复访存。

更好的做法是让负责 $Q_2$ 的线程块在内层循环期间保留这些状态。

FlashAttention-2 的抽象循环可以写成：

这就是“Q/O 定，K/V 动”。

这里的“SRAM”是视频的简化表达；具体 kernel 可能把不同状态分配在 shared memory 与 registers。

视频指出 v1 的高层伪代码先遍历 K/V blocks，再在内层遍历 Q blocks。

在这种组织下，一个 K/V tile 可以被多个 query tiles 复用。

但对某个 $Q_i,O_i$ 而言，每个外层 K/V block 到来时，都可能需要：

于是 Q 与输出状态可能随 K/V 外层循环反复搬运。

交换为外层 query 后，一个线程块从头到尾负责一个 $Q_i/O_i$ tile。

它只需在开始时读取 $Q_i$，在结束时写回一次最终 $O_i$。

在线 Softmax 状态也可在片上存活整个内层循环。

因此减少了 Q、O 以及 $m,\ell$ 的 HBM 往返。

对固定 query block，完整输出是对所有 key blocks 的贡献归约。

只要在线 Softmax 的重标定正确，按 $K_1,K_2,\ldots$ 的顺序逐块合并，与一次计算完整 key 轴等价。

循环重排改变数据驻留与并行任务划分，不改变最终

有限精度下归约顺序可能带来小幅舍入差异，但不是算法近似。

外层遍历多个 Q tiles 时，同一个 K/V tile 会被不同线程块读取。

所以 v2 不是让所有张量都只读一次，而是选择减少更昂贵或更频繁的 Q/O 状态往返，并换取更好的 query 维并行。

真实性能取决于：

循环顺序是 I/O 权衡，不是无条件消除所有重复加载。

外层 query tiles 彼此输出独立，可以由不同线程块并行处理。

每个线程块只需维护自己的 $O_i,m_i,\ell_i$。

这为较长序列提供 query 维的并行任务。

FlashAttention-2 的论文与实现还包含减少非矩阵乘 FLOPs、改善线程块/warp 工作划分等优化。

视频本课只聚焦循环交换与 Q/O 搬运，不应把 v2 的全部收益都归因于这一点。

在 causal attention 中，固定 $Q_i$ 后，内层只需遍历合法 K/V blocks。

严格位于未来的 blocks 可以直接跳过；对角线边界块执行块内 mask。

因此外 Q 内 K 的结构也便于为每个 query tile 确定自己的 key 范围。

但具体 kernel 如何计算边界、是否使用 split-K 或不同调度，随实现变化。

本课按视频解释 FlashAttention v1/v2 论文级高层循环差异。

现代 FlashAttention、框架集成和不同 GPU 架构可能采用多套 kernel；短序列、decode、varlen、causal、GQA 等路径不一定共享同一伪代码。

判断当前软件版本的真实数据驻留和循环次序，应查看对应 kernel、编译参数与 profiler，而不是只凭版本名。

FlashAttention 经常被概括为把注意力显存从平方级降到线性级。

这句话容易被误解成“注意力不再计算 $N^2$ 个 query-key 关系”。

视频抓住了更准确的瓶颈：

数学输出仍是精确的 dense attention，主要改变的是内存复杂度与数据搬运。

对单个 attention head，设：

则

视频用长上下文 $N\approx10^4$、head dimension $d\approx128$ 作量级示例。

具体模型的 $d$ 可以不同，但长序列时通常关注 $N\gg d$ 的情形。

转置后

因此

第一根 N 轴枚举 query token，第二根 N 轴枚举 key token。

每个元素

表示一对 token 的注意力分数。

若传统实现把 A 写回 HBM，它需要保存

个元素。

而单个 Q、K、V 或 O 只包含

个元素。

两者比值为

当 $N\gg d$ 时，A 很快成为更大的中间张量。

单头 Q 的元素数为

A 的元素数为

A 约为单个 Q 的 78 倍。

若仅为说明量级，按 FP16/BF16 每元素 2 字节，单头 A 约 200 MB。

实际总量还会乘 batch、head 数，并受 dtype、是否保存 logits/probabilities、mask 与框架实现影响。

这个数字只是 shape 估算，不是任意模型运行时的固定显存值。

把 A 写回 HBM 后，后续 Softmax 和 $PV$ 还要把它读回来。

因此物化 A 同时带来：

许多 attention kernel 在实际硬件上受内存 I/O 限制，而不只是算术吞吐限制。

FlashAttention 把 Q 沿 query token 行切成 tiles，把 $K^\top$ 沿 key token 列切成 tiles。

在视频的简化画法里，每个 token 向量的完整 $d$ 维保留在 tile 中。

一个 $B_r\times d$ 的 Q tile 与 $d\times B_c$ 的 $K^\top$ tile 相乘，得到

$B_r,B_c$ 根据片上容量和 kernel 设计选择。

每个注意力分数需要完整收缩

若沿 $d$ 切分，也不是数学上不可能，但必须再对各段部分内积做归约。

视频强调“token 完整性”，是在当前 tile 解释中把完整 head dimension 放进计算块，避免再引入 contraction 轴的跨块合并。

不能把它理解为任何 kernel 都绝不会在硬件微操作层拆分 $d$。

局部分数 $S_{\text{tile}}$ 产生后，立即用于更新：

它不写回 HBM。

随后下一个 K/V tile 覆盖另一段 key 轴，在线 Softmax 公式把新旧局部状态精确合并。

所有 key blocks 处理完后，只写回最终 O tile。

分块只改变求值顺序。

对每个 query 行，FlashAttention 维护全局等价状态

以及未归一化输出和

每来一个新 tile，若最大值变化，就重标定旧 $\ell_i$ 与 $\tilde O_i$。

最终

与完整矩阵 Softmax 相同。

不物化 A 后，主要持久 attention 张量 Q、K、V、O 都是 $N\times d$ 量级。

所以 attention 的额外持久存储可写作

当 $d$ 对所研究的序列长度 N 视为固定模型参数时，它关于 N 是线性的。

更严谨的说法是“从 $O(N^2)$ 中间存储降到 $O(Nd)$”，而不是无条件把两个变量都省略成 $O(N)$。

Dense attention 仍需覆盖所有合法 query-key 对。

$QK^\top$ 的算术量级仍是

$PV$ 也有同阶计算。

FlashAttention 的主要贡献是 I/O-aware：减少 HBM 流量与中间存储，从而让相同 dense attention 计算在硬件上更高效。

它没有把稠密注意力变成线性 attention 算法。

因果掩码只保留下三角，合法 query-key 对约为

这能把常数约减半，并允许跳过完全上三角 tiles。

但渐近计算量仍是 $O(N^2d)$。

因此“跳过上三角”和“不物化完整 A”应分别理解：前者减少无效计算，后者改变中间存储与 I/O。

训练需要反向传播通过 Softmax。

传统自动微分常保存大型中间激活。

FlashAttention 的反向可以利用前向保存的较小统计量和输出，重新计算局部分数与概率，而不是保存完整 A/P。

这用额外重计算换取更低内存。

具体保存哪些统计量、dropout RNG 如何恢复，会随算法版本与框架实现变化。

模型总显存还包括：

FlashAttention 消除的是 attention 中最关键的 $N\times N$ 中间物化。

它不保证整个训练或推理进程的全部显存只剩 $O(Nd)$。

视频用“FlashAttention 之前的 $N^2$ 注意力”概括传统 materialize-then-softmax 路径。

现代框架即使不显式标为 FlashAttention，也可能通过 memory-efficient SDPA、Triton kernel 或其他融合实现避免完整 A。

FlashAttention v1/v2/v3 与不同硬件后端的 tile、并行和反向策略也不同。

判断当前运行是否真的物化 $N\times N$ 中间量，应检查所选 backend、shape 支持条件和 profiler，而不是只看 API 名称。

上一课已经解释了完整注意力分数矩阵为何是 $N\times N$。

本课进一步问一个更量化的问题：

答案不是固定的“慢若干倍”。

差距随 N 增大而继续放大。

要看清这个趋势，先固定单个 attention head 的记号：

其中：

传统实现先计算

因为

所以

若把 A 或由它得到的 Softmax 概率写回 HBM，单头就是 $N^2$ 个元素。

视频在 00:14 开始把这一路径称为平方存储。

FlashAttention 仍计算精确的 dense attention。

不同点是它按块消费分数、执行在线 Softmax 并累加输出，不让完整 A 成为 HBM 中的持久中间量。

因此主要持久张量仍是 $N\times d$ 量级。

视频在 00:24 先画了 N 很小时的情况。

若 N 与 d 同量级，$N^2$ 和 $Nd$ 的数值不会相差太远。

例如 N=10、d=8：

只看这个规模，完整分数矩阵似乎“不算大”。

这正是容易形成误判的地方：

把 A 与一个 $N\times d$ 张量比较：

因此，只要 d 固定，N 每扩大一倍，这个相对倍数也扩大一倍。

若把 Q、K、V、O 四个 $N\times d$ 张量合计作为参照，则比值为

这两个比值口径不同：

不能把二者混用。

视频在 00:59 取 N=100000，且在 01:07 说明 d 可取 128 作示意。

一个 $N\times d$ 张量包含

个元素。

完整 A 则包含

个元素。

A 相对单个 $N\times d$ 张量大

倍。

相对四个 Q、K、V、O 的元素数总和，仍约大

倍。

若只作元素字节的理想化估算，FP16 的 A 单头就需要约

这个数字只用于说明平方项规模；真实实现还受 batch、head 数、mask、重计算、融合策略和数据类型影响。

视频在 00:50 明确指出 FlashAttention 不写回 A。

在 01:27 又用长序列图重申这一点。

所以“线性”更准确地说是：

它不意味着：

标准 dense attention 的核心算术量仍约为 $O(N^2d)$。

判断下列说法：

答案依次是：错、对、错。

错。

相对倍数含有 $N/d$，会随 N 变化。

更严格的 shape 口径是 $O(Nd)$。

只有把 d 视为固定模型常数时，才简称“关于 N 线性”。

FlashAttention 的主要贡献是 IO-aware 分块与不物化平方中间量。

标准稠密 attention 的算术关系仍是平方级。

完整训练还包括参数、优化器状态、其他层激活、梯度、临时 workspace 等。

本课只比较 attention 路径中的关键中间量。

真实 kernel 可能重计算、融合、使用不同累加精度，也可能受布局与对齐影响。

元素数分析用于判断渐近主项，不替代 profiler。

普通安全 Softmax 对一行 logits $x_1,\ldots,x_n$ 先取

再计算

概率为

在线 Softmax 的关键是：不必一次看到整行，仍能维护出同一个 m 与 $\ell$。

视频在 00:06 特别说明，把这看成“动态规划”主要是帮助理解局部状态如何维护全局结果。

它并不是经典的二维 DP 表格。

视频在 00:12 给出数量结论：

对每一行而言，在线 Softmax 的状态是：

这里的 $\ell$ 是标量。

m 也是标量。

它们都是“每行一个”，不是整个 attention 矩阵只保存一个。

设旧状态覆盖集合 A：

新块 B 的局部状态为

合并后的最大值是

但 $\ell_A$ 与 $\ell_B$ 使用了不同指数基准，不能直接相加。

先统一到新的 m：

视频从 00:45 开始解释这个分母合并。

每个缩放因子都不大于 1，因此保留了安全 Softmax 的数值稳定性。

视频用 $[1,2]$ 与 $[3,5]$ 两块演示。

第一块局部最大值是 2，第二块局部最大值是 5。

于是局部指数分别是：

和

第二块的 $e^0$ 与第一块的 $e^0$ 并不代表相同的全局权重。

因为它们分别以 5 和 2 为参考。

把第一块改写到全局最大值 5 的基准，需要整体乘

这就是“局部信息维护全局信息”的核心。

attention 输出一行是

若 V 也分块流入 SRAM，kernel 不能只维护 m 与 $\ell$。

它还要维护与 V 维度相同的输出累加状态。

视频在 01:08 转入矩阵分块，在 01:47 说明 SRAM 只能容纳部分 V。

因此，“三个量”应理解为三类每行状态：

第三项不是一个标量。

最清楚的记法是定义块 B 的未归一化输出分子

旧累加器为

统一到新的全局最大值 m 后：

最终归一化输出是

视频从 02:28 开始强调局部输出还需要按最大值差缩放，并在 02:39 补上最终除以全局分母。

有些讲解或实现把第三个状态直接记成已归一化输出 O。

设

则更新式要写成

这与未归一化记法完全等价。

但不能把已归一化 O 直接乘指数因子后相加，再忘掉 $\ell_A$、$\ell_B$。

三类状态是理解前向数学合并所需的最小概念集合。

具体 kernel 还可能保存：

FlashAttention-1、FlashAttention-2、FlashAttention-3 及不同框架封装的具体寄存器、SRAM 与 HBM 布局并不相同。

所以本课的“2 个/3 个”是算法状态分类，不是某个版本源码中变量名的机械计数。