Tri Dao
0c01568daf
Only run backward test for d=128 on A100
2022-10-04 18:06:08 -07:00
Tri Dao
8166063a55
Use block_size=128 for d=128 on SM86 to avoid exceeding smem limit
2022-09-12 14:21:29 -07:00
Tri Dao
13403e8115
Relax assert to allow both bf16 and fp16
2022-09-11 12:09:43 -07:00
Tri Dao
64f42cd057
Change license from Apache 2.0 to BSD
2022-09-09 12:07:35 -07:00
Tri Dao
04fb198523
Merge pull request #43 from eric-tc-wong/patch-1
...
Update flash_attention.py
2022-09-06 14:37:31 -07:00
eric-tc-wong
b410d14f28
Update flash_attention.py
...
Recasting query and key after rotary_emb()
2022-09-06 17:29:49 -04:00
Tri Dao
19d1261025
Add back need_weights in FlashMHA
2022-08-09 10:14:10 -07:00
Tri Dao
6cc7342575
Support index_first_axis with more than 2 dimensions
2022-08-05 09:48:16 -07:00
Tri Dao
713ea302d7
Allow headdim 128 in FlashMHA interface
2022-08-05 09:47:22 -07:00
Tri Dao
2ed471ecc4
Add tests for numerical error
2022-07-22 17:54:09 -04:00
Tri Dao
42f54d8840
Edit mention of Triton implementation
...
Phil Tillet suggests calling it "experimental".
2022-07-11 17:02:29 -07:00
Tri Dao
4577151ff8
Link to Triton implementation
2022-07-11 16:01:43 -07:00
Tri Dao
bc2c210254
Don't nest BOOL_SWITCH to work around gcc 7 bug
2022-07-11 10:28:46 -07:00
Tri Dao
d1fc80a3bb
Link to IEEE Spectrum article on MLPerf
2022-07-10 12:11:46 -07:00
Tri Dao
1bbebccc0a
Edit README to mention bf16 support
2022-07-09 23:34:29 -07:00
Tri Dao
de19de7ab1
Implement for bf16
2022-07-09 23:31:56 -07:00
Tri Dao
6a77a6da10
Refactor gemm_cl to template on either __half or __nv_bfloat16
2022-07-09 23:18:26 -07:00
Tri Dao
e518a4b327
Refactor to template on __half, implement bf16 util functions
2022-07-09 23:18:26 -07:00
Tri Dao
2dc1b205f6
Fix Illegal Memory Access bug in fwd when d=16
2022-07-09 23:17:14 -07:00
Tri Dao
5b838a8bef
Apply dropout scaling to dQ and dK instead of to V (in bwd)
...
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.
2022-07-03 17:53:37 -07:00
Tri Dao
a5559a0e75
Do P * dP (pointwise) in the bwd in fp32 instead of fp16
2022-07-03 17:52:05 -07:00
Tri Dao
6c3a8c65af
Implement cross attention
2022-07-03 17:48:12 -07:00
Tri Dao
01947bc93b
Merge pull request #18 from gahdritz/main
...
Slightly improve installation process
2022-07-02 23:33:14 -07:00
Gustaf
af4a9ce024
Add missing __init__.py
2022-07-03 02:04:55 -04:00
Gustaf
440e9c49f2
Add einops installation to setup.py
2022-07-03 02:04:24 -04:00
Tri Dao
f66603cb6f
Support batch size > 64K by swapping grid.x and grid.y
2022-06-29 23:16:24 -07:00
Tri Dao
450b64fe44
Add README section on issues
2022-06-27 13:50:16 -07:00
Tri Dao
c0daa62eaa
Add type check (fp16) in the forward pass
2022-06-26 11:41:30 -07:00
Tri Dao
ea38d3d261
Fix race condition in backward pass (smem_dq)
2022-06-25 18:02:30 -07:00
Tri Dao
eeca63a72a
Bug fix: wrong smem_o write pointer for d=16
2022-06-25 15:18:33 -07:00
Dan Fu
765741c1ee
More explanation
2022-06-14 11:55:14 -07:00
Dan Fu
2d5b2483b8
Speedup graph for A100, d128
2022-06-14 11:54:16 -07:00
Tri Dao
5d07483bbc
Refactor Gmem code to store q, k, v pointers separately
2022-06-12 16:37:32 -07:00
Tri Dao
d3e6440958
Implement bwd for head dim 128
2022-06-11 17:52:36 -07:00
Tri Dao
0d854692c6
Implement fwd for head dim 128
2022-06-11 17:52:36 -07:00
Dan Fu
0a398dfc37
Broken link
2022-06-04 17:28:45 -07:00
Dan Fu
bd60750e0b
T4
2022-06-04 17:27:51 -07:00
Tri Dao
321c57d07d
Set block size of SM75 fwd to 256 if there's no dropout
...
This speeds up the fwd by 1.5x.
2022-06-04 16:51:28 -07:00
Tri Dao
f2d8d4104e
Edit README: support Turing (SM75)
2022-06-04 16:06:48 -07:00
Tri Dao
d380e87fb6
Don't use Smem_dp_sum in backward pass
...
To reduce smem usage for SM75
2022-06-04 16:01:36 -07:00
Tri Dao
b17c6fe235
Reduce smem usage for Q and dO in the backward pass
...
From 4KB per buffer to 2KB per buffer. This saves us 8KB of smem (each Q and dO
have 2 buffers)
2022-06-03 16:59:11 -07:00
Tri Dao
2712aa4c8d
Support Turing mma instructions
2022-06-03 16:58:44 -07:00
Tri Dao
050873327e
Remove softmax fp16 max
2022-06-02 14:09:46 -07:00
Tri Dao
14dc326e59
Use Cutlass gemm as WarpMma
2022-06-02 10:33:32 -07:00
Tri Dao
e78e7c9553
Remove old backward
2022-06-02 10:13:44 -07:00
Tri Dao
512c98ee05
Add Cutlass as submodule
2022-06-02 09:54:16 -07:00
Dan Fu
ad6c694bb3
3090 speedup
2022-06-01 20:07:00 -07:00
Tri Dao
5a61cb7729
Rename src -> flash_attn
2022-06-01 18:50:26 -07:00
Tri Dao
c41479d66d
Support SM86 GPUs
2022-06-01 18:49:47 -07:00
Dan Fu
4b7cfb5f45
Citation
2022-05-30 13:29:04 -07:00
