刘孟玄 Blog

「离开世界之前一切都是过程」

FlashAttention 中的 Safe Softmax

从数值稳定性到分块重缩放

本文面向仅熟悉标准 Softmax的读者，用逐行推导的方式拆解 FlashAttention 为解决数值溢出与显存爆炸而引入的 Safe Softmax 及其分块实现。 1. 标准 Softmax 的数值陷阱给定向量 $x=[x_1,x_2,\dots,x_N]$，标准 Softmax 为 $\displaystyle \frac{e^{x_i}}{\sum_{j=1}^{N} ...

Posted by Liu Mengxuan on August 28, 2025