llama.cpp/ggml-cuda
Georgi Gerganov 9c67c2773d
ggml : add Flash Attention (#5021)
* ggml : add ggml_flash_attn_ext API

* ggml : fix GQA support in ggml_flash_attn_ext

* ggml : online attention (CPU)

* metal : initial implementation

* metal : f16 precision

* metal : reduce branches

* metal : specialize for head size

* wip : 8 rows per simd group

* wip : 4 rows per simd group

* wip : template for rows per warp

* metal : parallelize across KV size

* metal : parallel reduce across heads

* metal : efficient flash_attn_f16 implementation

* metal : avoid redundant loads of the attention

* metal : scale and mask in matrix form

* metal : fix comment

* llama : avoid ggml_cast, use F32 query

* metal : add parallel reduce version (disabled)

* metal : move output into local memory + optimize

- the result from each simdgroup now stays in the registers
- significantly reduced SRAM usage
- more efficient skipping of -INF blocks
- avoid simdgroup barrier in hot loop
- add comments

* metal : add tests, fix scaling, support C > 32

* metal : improve precision

* ggml : fix f16 mad

* metal : minor

* metal : support Q > 8

* tests : add ATTN tests

* metal : disable buffer allocation logs

* tests : more

* metal : faster inner loop for C == 32

* metal : fix array initialization

* tests : ifdef

* ggml : switch to padded F16 mask for ggml_soft_max, ggml_flash_attn_ext

* ggml : fix ggml_soft_max mask requirement

* cuda : fix soft_max to use correct mask size

* cuda : add flash_attn kernel (wip)

* metal : optimize softmax for C > 32

* metal : optimize softmax

* tests : minor fix

* cuda : avoid zeroing fragments

* tests : update dims

* cuda : fix __hisinf() result check

* cuda : avoid warp_reduce for smax

* cuda : use int instead of int64_t

Noticeably improves performance (thanks to Johannes)

* cuda : make loops use the same loop values

Thanks Johannes again for the tip

* cuda : unroll some of the loops

* cuda : avoid __hisinf branches

* cuda : use half2 in softmax

* cuda : switch to 1 warp for bs > 16

* cuda : speed-up reduce part of the kernel

* cuda : unroll Q*K^T loop

* cuda : fix -INF block check

* cuda : simplify softmax

* cuda : fix matrix names

* cuda : minor

* llama : adapt to F16 KQ_pos

* llama : adapt new models to F16 KQ_mask

* ggml : fix F16 store (ARM NEON)

* llama : fix type of KQ_mask and KQ_pos

* ggml : fix CPU soft_max

* tests : add hs=256

* cuda : fix build

* metal : improve perf via smaller int registers

* cuda : adapt soft_max to F16 mask and pos

* CUDA: faster FlashAttention, kernel for bs == 1

* 16 cols for Phi-2

* no vec for hs, no hs==256 ncols==32 for Volta

* adjust kernel selection logic

* 4 warps, 256 stride for all D

* no ncols == 64

* Multiple parallel blocks for batch size 1

* fix compile warnings

* fix excessive KQ_b loads

* fix cmake build

* fix KV cache padding, NaN from INFINITY (#6438)

* llama : flash_attn cparam + fix defrag

* server: support flash_attn param

* server: bench: enable flash_attn param

* CUDA: refactor host code, dyn. par. blocks

* fix flash_attn_vec_f16 race condition

* flush softmax exp below threshold to 0

* store temp KQ in registers

* Calculate KQ as FP32 if KQV has GGML_PREC_F32

* Add __hgt2_mask implementation for CUDA 11

* fix KQ FP32 precision fpr parallel_blocks > 1

* llama-bench : add -fa,--flash-attn arg

* metal : add BS=1 kernel for flash attention (#6508)

* metal : add BS=1 kernel for flash attention (wip)

* metal : support more than 1 warps

* metal : opts

* metal : opt

* metal : switch to parallel reduce

* metal : reduce registers

* metal : simplify

* metal : initial FA vec kernel

* metal : use F32 attention accumulators

* batched-bench : add fattn arg

* llama : simplify llama_build_kv_store

ggml-ci

* llama : adapt build_olmo to changes

* ggml : fix arm fp16 store on windows

* metal : clean-up

* metal : clean-up kernel code

* metal : minor

* tests : remove benchmarks

ggml-ci

* ggml : fix avx512 const correctness

ggml-ci

* ggml : fix soft_max with bias on CPU

ggml-ci

* common : print --flash-attn in help

* ggml : fix num dimensions in ggml_flash_attn_ext

* llama : force disable flash attention for incompatible models

* ggml : ggml_soft_max support F16/F32 mask/pos

ggml-ci

* cuda : uint -> uint32_t

* cuda : "constexpr dim3" -> "const dim3"

ggml-ci

* cuda : try to fix __hgt2_mask

ggml-ci

* ggml : add TODO's for F16/F32 mask/pos support in other backends

* llama : replace bool need_kq_pos with use_alibi

* llama : prep ALiBi support for BERT models

ggml-ci

* llama : fix n_batch requirements

ggml-ci

* cont

* server : add help for --flash-attn arg

* llama : disable FA for AMD

* tests : remove TMP_ATTN_BENCH

ggml-ci

* llama : support save/load state with FA enabled

ggml-ci

* ci : add CUDA save-load-state tests

ggml-ci

* llama : llama_kv_cache_clear zeroes data + fix save-load seq

ggml-ci

* llama : fix copy-paste errors, add TODO

* llama : disallow incompatible states

* llama : update llama_state_get_size after v_trans field

* metal : remove tmp log

* llama : add static reminder for llama_state_get_size

* metal : fix max nsg

ggml-ci

* ci : fix arg order

ggml-ci

---------

Co-authored-by: Johannes Gäßler <johannesg@5d6.de>
Co-authored-by: Pierrick HYMBERT <pierrick.hymbert@gmail.com>
2024-04-30 12:16:08 +03:00
..
acc.cu cuda : refactor into multiple files (#6269) 2024-03-25 13:50:23 +01:00
acc.cuh cuda : refactor into multiple files (#6269) 2024-03-25 13:50:23 +01:00
alibi.cu cuda : refactor into multiple files (#6269) 2024-03-25 13:50:23 +01:00
alibi.cuh cuda : refactor into multiple files (#6269) 2024-03-25 13:50:23 +01:00
arange.cu cuda : refactor into multiple files (#6269) 2024-03-25 13:50:23 +01:00
arange.cuh cuda : refactor into multiple files (#6269) 2024-03-25 13:50:23 +01:00
argsort.cu ggml : mul_mat_id use the same tensor for all the experts (#6387) 2024-04-03 16:07:05 +03:00
argsort.cuh cuda : refactor into multiple files (#6269) 2024-03-25 13:50:23 +01:00
binbcast.cu ggml : group all experts in a single ggml_mul_mat_id (#6505) 2024-04-18 15:18:48 +02:00
binbcast.cuh cuda : refactor into multiple files (#6269) 2024-03-25 13:50:23 +01:00
clamp.cu cuda : refactor into multiple files (#6269) 2024-03-25 13:50:23 +01:00
clamp.cuh cuda : refactor into multiple files (#6269) 2024-03-25 13:50:23 +01:00
common.cuh ggml : add Flash Attention (#5021) 2024-04-30 12:16:08 +03:00
concat.cu cuda : refactor into multiple files (#6269) 2024-03-25 13:50:23 +01:00
concat.cuh cuda : refactor into multiple files (#6269) 2024-03-25 13:50:23 +01:00
convert.cu Fix more int overflow during quant (PPL/CUDA). (#6563) 2024-04-29 00:38:44 +02:00
convert.cuh llama : add Command R Plus support (#6491) 2024-04-09 11:16:13 +03:00
cpy.cu cuda : refactor into multiple files (#6269) 2024-03-25 13:50:23 +01:00
cpy.cuh cuda : refactor into multiple files (#6269) 2024-03-25 13:50:23 +01:00
dequantize.cuh llama : add Command R Plus support (#6491) 2024-04-09 11:16:13 +03:00
diagmask.cu cuda : refactor into multiple files (#6269) 2024-03-25 13:50:23 +01:00
diagmask.cuh cuda : refactor into multiple files (#6269) 2024-03-25 13:50:23 +01:00
dmmv.cu llama : add Command R Plus support (#6491) 2024-04-09 11:16:13 +03:00
dmmv.cuh sync : ggml (#6351) 2024-03-29 17:45:46 +02:00
fattn.cu ggml : add Flash Attention (#5021) 2024-04-30 12:16:08 +03:00
fattn.cuh ggml : add Flash Attention (#5021) 2024-04-30 12:16:08 +03:00
getrows.cu cuda : refactor into multiple files (#6269) 2024-03-25 13:50:23 +01:00
getrows.cuh cuda : refactor into multiple files (#6269) 2024-03-25 13:50:23 +01:00
im2col.cu cuda : refactor into multiple files (#6269) 2024-03-25 13:50:23 +01:00
im2col.cuh cuda : refactor into multiple files (#6269) 2024-03-25 13:50:23 +01:00
mmq.cu cuda : refactor into multiple files (#6269) 2024-03-25 13:50:23 +01:00
mmq.cuh cuda : refactor into multiple files (#6269) 2024-03-25 13:50:23 +01:00
mmvq.cu IQ1_M: 1.75 bpw quantization (#6302) 2024-03-26 15:21:27 +01:00
mmvq.cuh cuda : refactor into multiple files (#6269) 2024-03-25 13:50:23 +01:00
norm.cu cuda : refactor into multiple files (#6269) 2024-03-25 13:50:23 +01:00
norm.cuh cuda : refactor into multiple files (#6269) 2024-03-25 13:50:23 +01:00
pad.cu cuda : refactor into multiple files (#6269) 2024-03-25 13:50:23 +01:00
pad.cuh cuda : refactor into multiple files (#6269) 2024-03-25 13:50:23 +01:00
pool2d.cu cuda : refactor into multiple files (#6269) 2024-03-25 13:50:23 +01:00
pool2d.cuh cuda : refactor into multiple files (#6269) 2024-03-25 13:50:23 +01:00
quantize.cu llama : add Command R Plus support (#6491) 2024-04-09 11:16:13 +03:00
quantize.cuh llama : add Command R Plus support (#6491) 2024-04-09 11:16:13 +03:00
rope.cu cuda : refactor into multiple files (#6269) 2024-03-25 13:50:23 +01:00
rope.cuh cuda : refactor into multiple files (#6269) 2024-03-25 13:50:23 +01:00
scale.cu cuda : refactor into multiple files (#6269) 2024-03-25 13:50:23 +01:00
scale.cuh cuda : refactor into multiple files (#6269) 2024-03-25 13:50:23 +01:00
softmax.cu ggml : add Flash Attention (#5021) 2024-04-30 12:16:08 +03:00
softmax.cuh cuda : refactor into multiple files (#6269) 2024-03-25 13:50:23 +01:00
sumrows.cu cuda : refactor into multiple files (#6269) 2024-03-25 13:50:23 +01:00
sumrows.cuh cuda : refactor into multiple files (#6269) 2024-03-25 13:50:23 +01:00
tsembd.cu cuda : refactor into multiple files (#6269) 2024-03-25 13:50:23 +01:00
tsembd.cuh cuda : refactor into multiple files (#6269) 2024-03-25 13:50:23 +01:00
unary.cu cuda : refactor into multiple files (#6269) 2024-03-25 13:50:23 +01:00
unary.cuh cuda : refactor into multiple files (#6269) 2024-03-25 13:50:23 +01:00
upscale.cu cuda : refactor into multiple files (#6269) 2024-03-25 13:50:23 +01:00
upscale.cuh cuda : refactor into multiple files (#6269) 2024-03-25 13:50:23 +01:00
vecdotq.cuh IQ1_M: 1.75 bpw quantization (#6302) 2024-03-26 15:21:27 +01:00