Commits · 184b992dcb2a0890adaa19eb9b541c3e4f9d2a08 · gaoqiong / flash-attention

27 Jul, 2023 1 commit

Kirthi Shankar Sivamani authored Jul 27, 2023



* Add RNG state to kernel launch params
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Save seed and offset for backward
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Single thread write to global mem
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* compute_dq_dk_dv_1colblock get seed and offset from launch params
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* compute_dq_dk_dv_1rowblock get seed and offset from launch params
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Change forward c++ APIs to save RNG state for backward
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Change backward c++ APIs to set RNG state for bprop launcher
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Bug fixes
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Python side API changes
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Bug fix; only save seeds instead of full offset
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Account for 3D grid size
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

---------
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

a03f6f8e

21 Jul, 2023 1 commit
- Fix using dO stride for O, which can cause memory error in bwd · 9ee0ff1d
  Tri Dao authored Jul 20, 2023
  
  9ee0ff1d
19 Jul, 2023 1 commit

Fix compile error on MSVC · 538d570c

danthe3rd authored Jul 19, 2023

See also: https://stackoverflow.com/questions/55136414/constexpr-variable-captured-inside-lambda-loses-its-constexpr-ness

538d570c

17 Jul, 2023 1 commit
- FlashAttention-2 release · 4f285b35
  Tri Dao authored Jul 17, 2023
  
  4f285b35
26 Apr, 2023 1 commit
- [Docs] Clearer error message for bwd d > 64, bump to v1.0.4 · ad113948
  Tri Dao authored Apr 26, 2023
  
  ad113948
15 Apr, 2023 1 commit
- only 1 thread writes to global mem in fprop · 45567a25
  Kirthi Shankar Sivamani authored Apr 15, 2023
```
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
```
  45567a25
13 Apr, 2023 1 commit
- Handle FlashAttnQKVPackedSplitFunc by making rng_state optional in backward · 7d25a4ec
  Kirthi Shankar Sivamani authored Apr 13, 2023
```
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
```
  7d25a4ec
12 Apr, 2023 1 commit
- Support CUDA graph capture · 31018c5f
  Kirthi Shankar Sivamani authored Apr 12, 2023
```
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
```
  31018c5f
15 Mar, 2023 1 commit
- Support H100 · 1b18f1b7
  Tri Dao authored Mar 15, 2023
  
  1b18f1b7
25 Jan, 2023 1 commit
- [FA] Remove unused variable rng_engine_inputs · 6b4a4821
  Tri Dao authored Jan 25, 2023
  
  6b4a4821
06 Jan, 2023 1 commit
- [Compilation] Change BOOL_SWITCH to fix Windows compilation · a1f49a2b
  Tri Dao authored Jan 06, 2023
```
Follow xFormers's DISTPATCH_BOOL. Haven't tested it on Windows.
```
  a1f49a2b
06 Dec, 2022 1 commit
- Simplify BOOL_SWITCH macro to fix compiling error on gcc 7 · 8a2ece89
  Tri Dao authored Dec 06, 2022
  
  8a2ece89
26 Nov, 2022 2 commits
- Fix typo in comments · 9bc63d1e
  Tri Dao authored Nov 25, 2022
  
  9bc63d1e
- Speed up compilation by splitting into separate .cu files · d95ee1a9
  Tri Dao authored Nov 25, 2022
  
  d95ee1a9
09 Nov, 2022 1 commit
- Fix out-of-bound memory read · 6998e0ec
  Tri Dao authored Nov 09, 2022
  
  6998e0ec
05 Nov, 2022 1 commit
- Parallelize CUDA bwd along seqlen_k instead of seqlen_q · 55778193
  Tri Dao authored Nov 05, 2022
```
This is faster since we only need to do atomic adds on dq, instead of atomic
adds on both dk and dv.
```
  55778193
25 Oct, 2022 1 commit
- Get rid of o_rows_are_valid since we don't have headdim=16 anymore · c422fee3
  Tri Dao authored Oct 24, 2022
  
  c422fee3
24 Oct, 2022 2 commits
- Support all head dims that are multiples of 8, up to 128 · 46fd2a20
  Tri Dao authored Oct 24, 2022
  
  46fd2a20
- Cast q.get_device() to char to avoid compiler warning (narrowing) · 97e13de2
  Tri Dao authored Oct 24, 2022
  
  97e13de2
23 Oct, 2022 2 commits
- Attempt to use atomicCAS to replace atomicAdd(bfloat16) · 9e92a1f2
  Tri Dao authored Oct 23, 2022
  
  9e92a1f2
- Split bwd on the seqlen_q dimension · a5a8806d
  Tri Dao authored Oct 23, 2022
  
  a5a8806d
22 Oct, 2022 1 commit
- Don't need to run configure for the forward pass · 871db479
  Tri Dao authored Oct 21, 2022
  
  871db479
21 Oct, 2022 3 commits
- Use block_size=128 for headdim=128 on SM80 · 7fc39832
  Tri Dao authored Oct 21, 2022
```
Previously we were using block_size=256.
```
  7fc39832
- Split fwd on the seqlen_q dimension · a44f48df
  Tri Dao authored Oct 21, 2022
  
  a44f48df
- Rework dropout to decouple forward and backward · 1aa6d7d9
  Tri Dao authored Oct 18, 2022
```
They don't have to have the same block size, number of threads, etc.
```
  1aa6d7d9
17 Oct, 2022 1 commit
- fix typo in function mha_fwd · ff07250e
  YangShu authored Oct 17, 2022
```
as title.
```
  ff07250e
16 Oct, 2022 1 commit
- Fix #54: set device for multi-GPU case · 52fb4b72
  Tri Dao authored Oct 16, 2022
  
  52fb4b72
14 Oct, 2022 1 commit
- Implement attention kernel that splits the batch into two · 5badfb78
  Tri Dao authored Oct 13, 2022
  
  5badfb78
05 Oct, 2022 2 commits
- Fixed switch statement, thanks @yocabon · 2211db5f
  Eric Engelhart authored Oct 04, 2022
  
  2211db5f
- Replace BOOL_SWITCH with FP16_SWITCH to work around MSVC bug with constexpr variables and templates · 9d7fd5b6
  Eric Engelhart authored Oct 03, 2022
  
  9d7fd5b6
12 Sep, 2022 1 commit
- Use block_size=128 for d=128 on SM86 to avoid exceeding smem limit · 8166063a
  Tri Dao authored Sep 12, 2022
  
  8166063a
11 Jul, 2022 1 commit
- Don't nest BOOL_SWITCH to work around gcc 7 bug · bc2c2102
  Tri Dao authored Jul 11, 2022
  
  bc2c2102
10 Jul, 2022 4 commits
- Implement for bf16 · de19de7a
  Tri Dao authored Jul 09, 2022
  
  de19de7a
- Refactor gemm_cl to template on either __half or __nv_bfloat16 · 6a77a6da
  Tri Dao authored Jul 08, 2022
  
  6a77a6da
- Refactor to template on __half, implement bf16 util functions · e518a4b3
  Tri Dao authored Jul 08, 2022
  
  e518a4b3
- Fix Illegal Memory Access bug in fwd when d=16 · 2dc1b205
  Tri Dao authored Jul 09, 2022
  
  2dc1b205
04 Jul, 2022 3 commits
- Apply dropout scaling to dQ and dK instead of to V (in bwd) · 5b838a8b
  Tri Dao authored Jun 29, 2022
```
Theoretically this might have lower numerical error since the scaling is in
fp32 instead of fp16 (not sure, I haven't thought too carefully about it).
However, in practice, the numerical errors seem about the same.
```
  5b838a8b
- Do P * dP (pointwise) in the bwd in fp32 instead of fp16 · a5559a0e
  Tri Dao authored Jul 03, 2022
  
  a5559a0e
- Implement cross attention · 6c3a8c65
  Tri Dao authored Jun 30, 2022
  
  6c3a8c65
30 Jun, 2022 1 commit
- Support batch size > 64K by swapping grid.x and grid.y · f66603cb
  Tri Dao authored Jun 29, 2022
  
  f66603cb