FlashAttention 分块运算可视化
重置
← 上一步
下一步 →
自动播放
步骤 0 / 25
初始化：设置 N=12, d=4, Br=3, Bc=2。SRAM 大小为 M。在 HBM 中初始化 O=0, ℓ=0, m=-∞。
Q 矩阵 (N×d)
K 矩阵 (N×d)
O 输出 (N×d)
V 矩阵 (N×d)
m 向量（行最大值，按 Br=3 分组）
ℓ 向量（行和，按 Br=3 分组）
注意力分数 S 矩阵 (N×N)，每个色块 = Br×Bc = 3×2，显示计算值
未计算
当前块
已完成
on-chip SRAM 计算
Qi
3×4
×
Kjᵀ
4×2
=
Sij
3×2
→
P̃ij
exp(S−m̃)
×
Vj
2×4
→
Oi
3×4
FlashAttention 前向传播公式
①
S
ij
= Q
i
· K
j
T
②
m̃
ij
= rowmax(S
ij
),
P̃
ij
= exp(S
ij
− m̃
ij
),
ℓ̃
ij
= rowsum(P̃
ij
)
③
m
new
= max(m
i
, m̃
ij
),
ℓ
new
= e
m
i
−m
new
·ℓ
i
+ e
m̃
ij
−m
new
·ℓ̃
ij
④
O
i
← diag(ℓ
new
)
−1
· (diag(ℓ
i
)·e
m
i
−m
new
·O
i
+ e
m̃
ij
−m
new
·P̃
ij
·V
j
)