Efficient Attention & FlashAttention: 解释在线 softmax（online softmax）如何支持分块计算。

Core Concept: 逐块处理时维护 running max 与 running sum，用 e^{m_old−m_new} 重缩放已累积结果，实现与全量 softmax 等价。

数学机理：<strong>问题</strong>——softmax 需要<strong>行最大值</strong>（为防溢出）与<strong>行和</strong>（为归一化），但分块计算时无法一次看到整行；若分块独立算 softmax，得到的归一化常数不一致（各块用自己的 max/sum）。<strong>在线 softmax（Milakov & Gimelshein 2018）</strong> 用<strong>流式更新</strong>解决：处理第 j 块时，(1) 计算该块的最大值，更新 running max：m^new=max(m_old, max_j z_j)；(2) <strong>重缩放</strong>已累积的和与输出：把之前的 ℓ 与 O 乘以 e^{m_old−m_new}（因为 max 变大后，之前的指数项需相应缩小以保持等价）；(3) 累积新块的贡献：ℓ^new=e^{m_old−m_new}·ℓ_old+Σ_j e^{z_j−m_new}，O^new=e^{m_old−m_new}·O_old+Σ_j e^{z_j−m_new}v_j。最终 O/ℓ 即精确的 softmax 加权和。<strong>等价性的关键</strong>——softmax 对<strong>整体平移不变</strong>（softmax(z+c)=softmax(z)），故可'延迟'归一化：先按块累积<strong>未归一化</strong>的加权 V 与指数和，过程中用 running max 保证不溢出，最后统一除以 ℓ。<strong>数值精度</strong>——在线 softmax 与全量 softmax 的差异仅来自浮点舍入（不同求和顺序），在 FP32 累加下可忽略；这正是 Flash Attention 在 kernel 内用 FP32 累加 ℓ 与 O 的原因。

深度剖析与工程权衡：① <strong>'平移不变性'是等价性的基础</strong>——理解这一点是掌握在线 softmax 的关键；面试中若被要求证明，应从 softmax 的定义出发说明'整体减常数不改变结果'。② <strong>与 Flash Attention 的关系</strong>——在线 softmax 是 Flash Attention 的<strong>核心组件</strong>；Flash 把'分块 + 在线 softmax + 不物化'组合起来实现 IO 优化。③ <strong>反向传播的重计算</strong>——Flash 的反向不保存 L×L 矩阵，而是用保存的 m、ℓ（每行两个标量）<strong>重算</strong> S 与 P；这是显存 O(L) 的另一半原因（保存的只有 O(L) 统计量而非 O(L²) 矩阵）。④ <strong>与流式/长上下文的联系</strong>——在线 softmax 使'流式处理超长序列'成为可能（不需要一次看到整行）；这与 SSM 的'递归状态'思想有相通之处。⑤ <strong>数值细节</strong>——重缩放因子 e^{m_old−m_new} 在 m_new 远大于 m_old 时趋 0（旧贡献被丢弃，符合'新块有更大值则旧块贡献很小'的直觉）；实现中需注意 FP16 下该因子可能下溢（故用 FP32 累积）。⑥ <strong>面试要点</strong>——被问'分块怎么做 softmax'，应写出 running max/sum 的递推式并解释<strong>重缩放</strong>的作用；能指出'等价性依赖平移不变性'与'FP32 累加保证精度'是深度理解的标志。