FlashAttention Forward Pass 流程演示

不再读写庞大的 O(N²) 注意力矩阵 S 和 P。将输入切分为 Block，
通过 Tiling (分块) 和 Algebraic Aggregation (代数聚合)，完全在快速 SRAM 中计算后直接写回 HBM。

阶段 0: 初始状态

初始化...

HBM (High Bandwidth Memory - 慢速/大容量)

Q₁

Q₂

Q₃

Q₄

K₁

K₂

K₃

K₄

V₁

V₂

V₃

V₄

↓ 输出与统计信息 (被初始化为 0 / -∞) ↓

m, l

-∞, 0

SRAM (On-Chip Cache - 极速/小容量)

外层循环 (常驻)

Empty

内层循环 (流转)

Empty

ALU (算术逻辑单元) 实时计算

S_ij = Q_i K_j^T

rowmax/sum: m̃_ij, l̃_ij, P̃_ij

m^new_i = max(m_i, m̃_ij)

O_i ← diag(l^new_i)^-1 [ O_i + P̃_ij V_j ]

阶段总结剖析：代数聚合 (Algebraic Aggregation) 的魔法

传统注意力痛点： 需要完整的 S = QK^T 来计算 Softmax 的分母，因此必须把 O(N²) 的庞大矩阵暂存到慢速 HBM 中，导致 IO 瓶颈。
FlashAttention 的解法： Softmax 的分母本质是行的指数和 (rowsum)。我们在 HBM 中只维护 O(N) 级别的统计量：当前的最大值 m_i 和指数和 l_i。
当内层计算出一个新的 S_ij 块时：
1. 计算当前块的局部最大值 m̃_ij 和局部和 l̃_ij。
2. 通过公式 m_i^new = max(m_i, m̃_ij) 获取全局新最大值。
3. 使用缩放因子 e^{m_i - m_i^new} 修正历史结果 O_i 和历史和 l_i。
结果： 算完一块就丢掉一块 S_ij，绝不写回 HBM，完美融合了局部计算与全局 Softmax。