File tree Expand file tree Collapse file tree 1 file changed +2
-2
lines changed
infiniop/ops/flash_attention_backward Expand file tree Collapse file tree 1 file changed +2
-2
lines changed Original file line number Diff line number Diff line change 22
33` FlashAttentionBackward ` 是算子 ` FlashAttention ` 的反向传播。
44
5- 与正向传播相同,先进行分块。$K$ 和 $V$ 在外循环中逐块加载,而 $Q$ 在内循环中逐块加载。
5+ 与正向传播相同,先进行分块。 $K$ 和 $V$ 在外循环中逐块加载,而 $Q$ 在内循环中逐块加载。
66
771 . 外循环中,每次循环需要先初始化当前块的 $\mathbf{d}\mathbf{K}_ j$ 和 $\mathbf{d}\mathbf{V}_ j$ 为 0;
882 . 内循环中,按以下顺序计算
99 1 . $\mathbf{S}_ i^{(j)}=\mathbf{Q}_ i\mathbf{K}_ j^T\in\mathbb{R}^{B_r\times B_c}$
10- 2 . $\mathbf{P}_ i^{(j)}=\exp(\mathbf{S}_ {ij}-L_i)\in\mathbb{R}^{B_r\times B_c}$
10+ 2 . $\mathbf{P}_ i^{(j)}=\exp(\mathbf{S} _ {ij}-L_i)\in\mathbb{R}^{B_r\times B_c}$
1111 3 . $\mathbf{d}\mathbf{V}_ j\leftarrow\mathbf{d}\mathbf{V}_ j+(\mathbf{P}_ i^{(j)})^\top\mathbf{d}\mathbf{O}_ i\in\mathbb{R}^{B_c\times d}$
1212 4 . $\mathbf{dP}_ i^{(j)}=\mathbf{dO}_ i\mathbf{V}_ j^\top\in\mathbb{R}^{B_r\times B_c}$
1313 5 . $\mathbf{dQ}_ i\leftarrow\mathbf{dQ}_ i+\mathbf{dS}_ i^{(j)}\mathbf{K}_ j$
You can’t perform that action at this time.
0 commit comments