Skip to content

Commit 4509cab

Browse files
GushanFallGushanFall
committed
docs: flash_attention_backward
1 parent 76a3894 commit 4509cab

File tree

1 file changed

+176
-0
lines changed
  • infiniop/ops/flash_attention_backward

1 file changed

+176
-0
lines changed

‎infiniop/ops/flash_attention_backward/README.md‎

Lines changed: 176 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,176 @@
1+
# `FlashAttentionBackward`
2+
3+
`FlashAttentionBackward` 是算子 `FlashAttention` 的反向传播。
4+
5+
与正向传播相同，先进行分块。$K$ 和 $V$ 在外循环中逐块加载，而 $Q$ 在内循环中逐块加载。
6+
7+
1. 外循环中，每次循环需要先初始化当前块的 $\mathbf{d}\mathbf{K}_j$ 和 $\mathbf{d}\mathbf{V}_j$ 为 0；
8+
2. 内循环中，按以下顺序计算
9+
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}$
11+
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}$
12+
4. $\mathbf{dP}_i^{(j)}=\mathbf{dO}_i\mathbf{V}_j^\top\in\mathbb{R}^{B_r\times B_c}$
13+
5. $\mathbf{dQ}_i\leftarrow\mathbf{dQ}_i+\mathbf{dS}_i^{(j)}\mathbf{K}_j$
14+
6. $\mathbf{dK}_j\leftarrow\mathbf{dK}_j+\mathbf{dS}_i^{(j)\top}\mathbf{Q}_i\in\mathbb{R}^{B_c\times d}$
15+
16+
## 接口
17+
18+
### 计算
19+
20+
```c
21+
infiniStatus_t infiniopFlashAttentionBackward(
22+
infiniopFlashAttentionBackwardDescriptor_t desc,
23+
void *workspace,
24+
size_t workspace_size,
25+
void *grad_q,
26+
void *grad_k,
27+
void *grad_v,
28+
const void *q,
29+
const void *k,
30+
const void *v,
31+
const void *grad_out,
32+
const void *mask,
33+
void *stream
34+
);
35+
```
36+
37+
<div style="background-color: lightblue; padding: 1px;"> 参数： </div>
38+
39+
- `desc`:
40+
已使用 `infiniopCreateFlashAttentionBackwardDescriptor()` 初始化的算子描述符；
41+
- `workspace`:
42+
指向算子计算所需的额外工作空间；
43+
- `workspace_size`:
44+
`workspace` 的大小，单位：字节；
45+
- `grad_q`:
46+
查询（Query）梯度计算结果地址。张量限制见[创建算子描述](#创建算子描述)部分。
47+
- `grad_k`:
48+
键（Key）梯度计算结果地址。张量限制见[创建算子描述](#创建算子描述)部分。
49+
- `grad_v`:
50+
值（Value）梯度计算结果地址。张量限制见[创建算子描述](#创建算子描述)部分。
51+
- `q`:
52+
查询（Query）张量数据指针。张量限制见[创建算子描述](#创建算子描述)部分。
53+
- `k`:
54+
键（Key）张量数据指针。张量限制见[创建算子描述](#创建算子描述)部分。
55+
- `v`:
56+
值（Value）张量数据指针。张量限制见[创建算子描述](#创建算子描述)部分。
57+
- `grad_out`
58+
注意力梯度张量数据指针。张量限制见[创建算子描述](#创建算子描述)部分。
59+
- `mask`:
60+
注意力掩码的数据指针，可选参数。取值为 `0` 时表示保留对应位置的元素（参与计算）；取值为 `-inf` （负无穷）时表示屏蔽对应位置的元素（即跳过，不参与计算）。张量限制见[创建算子描述](#创建算子描述)部分。
61+
- `stream`:
62+
计算流/队列。
63+
64+
<div style="background-color: lightblue; padding: 1px;"> 返回值：</div>
65+
66+
- [`INFINI_STATUS_SUCCESS`], [`INFINI_STATUS_NULL_POINTER`], [`INFINI_STATUS_INSUFFICIENT_WORKSPACE`], [`INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED`], [`INFINI_STATUS_INTERNAL_ERROR`].
67+
68+
### 创建算子描述
69+
70+
```c
71+
infiniStatus_t infiniopCreateFlashAttentionBackwardDescriptor(
72+
infiniopHandle_t handle,
73+
infiniopFlashAttentionBackwardDescriptor_t *desc_ptr,
74+
infiniopTensorDescriptor_t grad_q_desc,
75+
infiniopTensorDescriptor_t grad_k_desc,
76+
infiniopTensorDescriptor_t grad_v_desc,
77+
infiniopTensorDescriptor_t q_desc,
78+
infiniopTensorDescriptor_t k_desc,
79+
infiniopTensorDescriptor_t v_desc,
80+
infiniopTensorDescriptor_t grad_out_desc,
81+
infiniopTensorDescriptor_t mask_desc,
82+
infiniopAttentionMaskType_t mask_type
83+
);
84+
```
85+
86+
<div style="background-color: lightblue; padding: 1px;"> 参数：</div>
87+
88+
- `handle`:
89+
`infiniopHandle_t`类型的硬件控柄。详见 [`InfiniopHandle_t`]。
90+
- `desc_ptr`:
91+
`infiniopFlashAttentionBackwardDescriptor_t` 指针，指向将被初始化的算子描述符地址。
92+
- `grad_q_desc` - { dT | ((batch_size,) seq_len_q, num_heads_q, head_dim) | ($\ldots, 1$)}:
93+
算子计算参数 `grad_q` 的张量描述，四维或者三维，最后一维连续。
94+
- `grad_k_desc` - { dT | ((batch_size,) seq_len_kv, num_heads_kv, head_dim) | ($\ldots, 1$)}:
95+
算子计算参数 `grad_k` 的张量描述，四维或者三维，最后一维连续。
96+
- `grad_v_desc` - { dT | ((batch_size,) seq_len_kv, num_heads_kv, head_dim) | ($\ldots, 1$)}:
97+
算子计算参数 `grad_v` 的张量描述，四维或者三维，最后一维连续。
98+
- `q_desc` - { dT | ((batch_size,) seq_len_q, num_heads_q, head_dim) | ($\ldots, 1$)}:
99+
算子计算参数 `q` 的张量描述，形状与 `grad_q_desc` 一致，最后一维连续。
100+
- `k_desc` - { dT | ((batch_size,) seq_len_kv, num_heads_kv, head_dim) | ($\ldots, 1$)}:
101+
算子计算参数 `k` 的张量描述，形状与 `grad_k_desc` 一致，最后一维连续。
102+
- `v_desc` - { dT | ((batch_size,) seq_len_kv, num_heads_kv, head_dim) | ($\ldots, 1$)}:
103+
算子计算参数 `v` 的张量描述，形状与 `grad_v_desc` 一致，最后一维连续。
104+
- `grad_out` - { dT | ((batch_size,) seq_len_q, num_heads_q, head_dim) | ($\ldots, 1$)}:
105+
算子计算参数 `grad_out` 的张量描述，形状与 `grad_q_desc` 一致，最后一维连续。
106+
- `mask_desc` - { dM | (seq_len_q, seq_len_kv) | (~)}:
107+
算子计算参数 `mask` 的张量描述，当 `mask_type=INFINIOP_ATTENTION_MASK_TYPE_FULL` 时，`mask` 不可为空，其余情况 `mask` 可为`nullptr`。
108+
- `mask_type` - `infiniopAttentionMaskType_t`:
109+
注意力类型参数，有三种类型可选，详细见参数限制。
110+
111+
参数限制：
112+
113+
- `dT`: `Float16`, `Float32` 或 `BFloat16`。
114+
- `dM`: `Flaot32`。
115+
- `seq_len_q` 与 `seq_len_kv` 可以不同。
116+
- `num_heads_q` 与 `num_heads_kv` 可以不同，但需满足前者是后者的整数倍（非0整数）。
117+
- 当 $N_q/N_{kv}=1$ 时，即为 MQA (multi-query attention)
118+
- 当 $N_q/N_{kv}>1$ 时，即为 GQA (grouped-query attention)
119+
- `mask_type` 的三种类型：
120+
- `INFINIOP_ATTENTION_MASK_TYPE_NONE=0`: 不使用注意力掩码，忽略 `mask` 取值；
121+
- `INFINIOP_ATTENTION_MASK_TYPE_FULL=1`: 使用完整 mask 矩阵，此时 `mask` 不可为空；
122+
- `INFINIOP_ATTENTION_MASK_TYPE_CAUSAL=2`: 使用标准因果掩码，对应以左上顶点划分的下三角场景，忽略 `mask` 取值；
123+
124+
<div style="background-color: lightblue; padding: 1px;"> 返回值：</div>
125+
126+
- [`INFINI_STATUS_SUCCESS`], [`INFINI_STATUS_BAD_PARAM`], [`INFINI_STATUS_BAD_TENSOR_SHAPE`], [`INFINI_STATUS_BAD_TENSOR_DTYPE`], [`INFINI_STATUS_BAD_TENSOR_STRIDES`], [`INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED`].
127+
128+
### 计算额外工作空间
129+
130+
```c
131+
infiniStatus_t infiniopGetFlashAttentionBackwardWorkspaceSize(
132+
infiniopFlashAttentionBackwardDescriptor_t desc,
133+
size_t *size
134+
);
135+
```
136+
137+
<div style="background-color: lightblue; padding: 1px;"> 参数：</div>
138+
139+
- `desc`:
140+
已使用 `infiniopCreateFlashAttentionBackwardDescriptor()` 初始化的算子描述符；
141+
- `size`:
142+
额外空间大小的计算结果的写入地址；
143+
144+
<div style="background-color: lightblue; padding: 1px;"> 返回值：</div>
145+
146+
- [`INFINI_STATUS_SUCCESS`], [`INFINI_STATUS_BAD_PARAM`], [`INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED`].
147+
148+
### 销毁算子描述符
149+
150+
```c
151+
infiniStatus_t infiniopDestoryFlashAttentionBackwardDescriptor(
152+
infiniopFlashAttentionBackwardDescriptor_t desc
153+
);
154+
```
155+
156+
<div style="background-color: lightblue; padding: 1px;"> 参数： </div>
157+
158+
- `desc`:
159+
输入。 待销毁的算子描述符；
160+
161+
<div style="background-color: lightblue; padding: 1px;"> 返回值： </div>
162+
163+
- [`INFINI_STATUS_SUCCESS`], [`INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED`].
164+
165+
<!-- 链接 -->
166+
[`InfiniopHandle_t`]: /infiniop/handle/README.md
167+
168+
[`INFINI_STATUS_SUCCESS`]: /common/status/README.md#INFINI_STATUS_SUCCESS
169+
[`INFINI_STATUS_BAD_PARAM`]: /common/status/README.md#INFINI_STATUS_BAD_PARAM
170+
[`INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED`]: /common/status/README.md#INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED
171+
[`INFINI_STATUS_BAD_TENSOR_SHAPE`]: /common/status/README.md#INFINI_STATUS_BAD_TENSOR_SHAPE
172+
[`INFINI_STATUS_BAD_TENSOR_DTYPE`]: /common/status/README.md#INFINI_STATUS_BAD_TENSOR_DTYPE
173+
[`INFINI_STATUS_BAD_TENSOR_STRIDES`]: /common/status/README.md#INFINI_STATUS_BAD_TENSOR_STRIDES
174+
[`INFINI_STATUS_NULL_POINTER`]:/common/status/README.md#INFINI_STATUS_NULL_POINTER
175+
[`INFINI_STATUS_INSUFFICIENT_WORKSPACE`]:/common/status/README.md#INFINI_STATUS_INSUFFICIENT_WORKSPACE
176+
[`INFINI_STATUS_INTERNAL_ERROR`]:/common/status/README.md#INFINI_STATUS_INTERNAL_ERROR

0 commit comments

Comments
 (0)