Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

WIP: Flash Attention implementation (forward + backward) #1

Closed
wants to merge 72 commits into from
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
72 commits
Select commit Hold shift + click to select a range
f7bcfb0
cuda: add flash attention + test
FSSRepo Jan 17, 2024
e53de28
fix compilation
FSSRepo Jan 18, 2024
a1c004e
ggml : add ggml_flash_attn_ext API
ggerganov Jan 18, 2024
fa7ebcc
ggml : fix GQA support in ggml_flash_attn_ext
ggerganov Jan 19, 2024
09db1a7
Merge branch 'gg/flash-attn' of https://github.com/ggerganov/llama.cp…
FSSRepo Jan 19, 2024
fded2e6
apply suggestions
FSSRepo Jan 20, 2024
c3cdfff
Merge branch 'master' into gg/flash-attn
ggerganov Jan 20, 2024
a9681fe
ggml : online attention (CPU)
ggerganov Jan 20, 2024
1173f49
metal : initial implementation
ggerganov Jan 20, 2024
528da75
metal : f16 precision
ggerganov Jan 21, 2024
52ae085
metal : reduce branches
ggerganov Jan 21, 2024
b973258
metal : specialize for head size
ggerganov Jan 21, 2024
8cde449
wip : 8 rows per simd group
ggerganov Jan 21, 2024
f31955f
wip : 4 rows per simd group
ggerganov Jan 21, 2024
a4b6341
wip : template for rows per warp
ggerganov Jan 21, 2024
77d08f3
metal : parallelize across KV size
ggerganov Jan 21, 2024
17720fa
metal : parallel reduce across heads
ggerganov Jan 21, 2024
a689b02
Merge branch 'gg/flash-attn' of https://github.com/ggerganov/llama.cp…
FSSRepo Jan 23, 2024
6374bc5
cuda: port metal version flash_attn_ext
FSSRepo Jan 23, 2024
6416821
fix equivalent fp16 math functions, compiler error 'undefined'
FSSRepo Jan 24, 2024
972c2ad
use half2 instead half4
FSSRepo Jan 24, 2024
0fc36d8
match to metal impl
FSSRepo Jan 24, 2024
1446a12
metal : efficient flash_attn_f16 implementation
ggerganov Jan 23, 2024
d917746
metal : avoid redundant loads of the attention
ggerganov Jan 25, 2024
432ad04
metal : scale and mask in matrix form
ggerganov Jan 25, 2024
40ea8cd
metal : fix comment
ggerganov Jan 25, 2024
78da338
Merge branch 'gg/flash-attn' of https://github.com/ggerganov/llama.cp…
FSSRepo Jan 25, 2024
f9ca5dc
llama : avoid ggml_cast, use F32 query
ggerganov Jan 25, 2024
6e7cb0e
update implementation
FSSRepo Jan 25, 2024
6fea843
metal : add parallel reduce version (disabled)
ggerganov Jan 25, 2024
0a481fe
integrate tensor cores
FSSRepo Jan 27, 2024
7cea973
Merge branch 'gg/flash-attn' of https://github.com/ggerganov/llama.cp…
FSSRepo Jan 27, 2024
2455a8d
update impl
FSSRepo Jan 27, 2024
b3dd7d9
Merge branch 'master' into gg/flash-attn
ggerganov Jan 28, 2024
77f6976
metal : move output into local memory + optimize
ggerganov Jan 28, 2024
ecc466a
metal : add tests, fix scaling, support C > 32
ggerganov Jan 28, 2024
3a428a1
metal : improve precision
ggerganov Jan 28, 2024
8612864
ggml : fix f16 mad
ggerganov Jan 28, 2024
0ad44ba
Merge branch 'master' into gg/flash-attn
ggerganov Jan 28, 2024
134c81c
metal : minor
ggerganov Jan 28, 2024
1db22d7
metal : support Q > 8
ggerganov Jan 28, 2024
4794821
tests : add ATTN tests
ggerganov Jan 29, 2024
abeaf0d
metal : disable buffer allocation logs
ggerganov Jan 29, 2024
c6c1132
tests : more
ggerganov Jan 29, 2024
5fcb9c1
metal : faster inner loop for C == 32
ggerganov Jan 29, 2024
a1d5a12
fix compiler error
FSSRepo Jan 29, 2024
7980178
Merge branch 'gg/flash-attn' of https://github.com/ggerganov/llama.cp…
FSSRepo Jan 29, 2024
d073e4f
metal : fix array initialization
ggerganov Jan 30, 2024
78df552
tests : ifdef
ggerganov Jan 30, 2024
3d03bcb
Merge branch 'master' into gg/flash-attn
ggerganov Jan 30, 2024
3b0f74b
latest kernel update, wrong values
FSSRepo Jan 30, 2024
2ddc9bb
Merge branch 'master' into gg/flash-attn
ggerganov Jan 31, 2024
b1479df
fix kernel
FSSRepo Jan 31, 2024
8ad92dc
ggml : switch to padded F16 mask for ggml_soft_max, ggml_flash_attn_ext
ggerganov Jan 31, 2024
0afe47f
fix naive implementation
FSSRepo Jan 31, 2024
3df0b8d
Merge branch 'gg/flash-attn' of https://github.com/ggerganov/llama.cp…
FSSRepo Jan 31, 2024
fd878f7
cuda: mask as fp16
FSSRepo Jan 31, 2024
71b69aa
cuda : fix flash_attn kernel to produce same results as CPU
ggerganov Feb 1, 2024
2c04bee
cuda : avoid extra QxQ matrix in shared memory
ggerganov Feb 1, 2024
9a5c2a1
cuda : switch to F16 scalars + tune warps for RTX 2060
ggerganov Feb 1, 2024
ac26f27
cuda : increase C to 128 for better performance
ggerganov Feb 1, 2024
43f7156
Merge pull request #3 from ggerganov/flash-attn-cuda
FSSRepo Feb 1, 2024
9240a84
fix mask nullptr
FSSRepo Feb 1, 2024
8d7a606
don't require LLAMA_CUDA_F16 to compile
FSSRepo Feb 1, 2024
19e0b8e
#ifdef -> #if + fix check -inf
FSSRepo Feb 1, 2024
cae985c
cmake: remove unused changes
FSSRepo Feb 1, 2024
53621e3
refactor flash_attn function + improve tests
FSSRepo Feb 1, 2024
674d5ac
unroll 2 loops, int64_t -> int, 309 µs
JohannesGaessler Feb 3, 2024
8b51ab4
Merge pull request #4 from Pints-App/jg/flash-attn-cuda
FSSRepo Feb 3, 2024
a1f9ffe
bring optimizations from gg/flash-attn
FSSRepo Feb 3, 2024
ba7699d
Merge branch 'flash-attn-cuda' of https://github.com/Pints-App/llama.…
FSSRepo Feb 3, 2024
f659f57
fix merge conflicts
FSSRepo Feb 3, 2024
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion examples/batched-bench/batched-bench.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -104,7 +104,7 @@ int main(int argc, char ** argv) {

ctx_params.seed = 1234;
ctx_params.n_ctx = n_kv_max;
ctx_params.n_batch = 512;
ctx_params.n_batch = 2048;
ctx_params.mul_mat_q = mmq;

ctx_params.n_threads = params.n_threads;
Expand Down
Loading