Training Stability & Mixed Precision: 解释 attention 的数值稳定性问题与 Flash Attention 的处理。

Core Concept: attention 需对 L 个元素做 softmax，直接物化会 O(L²) 显存且 FP16 下累加误差大；Flash 用在线 softmax + 分块。

数学机理：<strong>数值问题</strong>有两层。(1) <strong>softmax 上溢</strong>——attention 分数 z=q·k/√d 若直接算 exp(z) 会溢出；标准做法是减去行最大值 exp(z−max)。但朴素实现需要<strong>先物化完整的 L×L 分数矩阵</strong>才能求 max，显存 O(L²)——L=8192 时单头即 64M 元素、FP16 下 128 MB，多层多头后不可行。(2) <strong>累加精度</strong>——softmax 的归一化要对 L 个 exp 求和，FP16 逐元素累加会累积误差。<strong>Flash Attention</strong> 的核心是<strong>在线 softmax（online / streaming softmax）</strong>：把 key/value 按块（block）处理，每处理一块就更新 <strong>running max m</strong> 与 <strong>running sum ℓ</strong>：先算 m^new=max(m, max_j z_j)，再把已累积的 ℓ 按 e^{m−m^new} 重新缩放、加上新块的贡献。这样<strong>无需保存完整分数矩阵</strong>（显存 O(L)），且数学上与全量 softmax <strong>完全等价</strong>（因为 softmax 对 max 的平移不变性）。同时它在 kernel 内用 <strong>FP32 累加</strong> ℓ 与输出，保证精度。<strong>额外收益</strong>：不物化 L×L 矩阵意味着大幅减少 HBM 读写（attention 是 memory-bound 的），故 Flash Attention 不仅省显存，还<strong>更快</strong>（2~4 倍）。

深度剖析与工程权衡：① <strong>等价性证明要点</strong>——softmax 的输出对分数整体平移不变，故可'延迟'归一化：先按块累积未归一化的加权 V 与指数和，最后统一除以 ℓ；过程中用 running max 保证指数不溢出。这是'分块 + 在线'能等价的关键。② <strong>反向的重计算</strong>——Flash Attention 的反向不保存中间矩阵，而是用保存的 m、ℓ 重新计算（类似检查点思想）；这是它显存 O(L) 的另一半原因。③ <strong>与 GQA/MQA 的关系</strong>——减少 KV head 数可降低 KV cache 显存，与 Flash Attention 的'计算时显存'优化互补（前者省推理显存、后者省训练显存）。④ <strong>长上下文的组合拳</strong>——Flash Attention（省激活）+ GQA（省 KV cache）+ RoPE 插值/位置外推 + 稀疏/滑窗注意力，是长上下文训练的标准组合。⑤ <strong>其他实现</strong>——xformers 的 memory-efficient attention、以及 FlashAttention-2/3（进一步优化并行与 warp 调度）都是同一思想的工程演进。⑥ <strong>面试要点</strong>——被问'attention 的数值稳定'，应能写出<strong>在线 softmax 的递推式</strong>并解释'为何与全量等价'；同时指出 Flash Attention 的收益是'显存 + 速度'双重（因 memory-bound），这比只说'省显存'更完整。