Flash Attention 2 support possible?

[Dampfinchen]

Hello!

So Flash Attention 2 has just been released. https://github.com/Dao-AILab/flash-attention

Apparently, this can also be used to speed up inference and significantly decrease memory consumption for context.

Less memory usage for ctx could definately be useful for llama.cpp as ctx does need a significant amount of memory, regardless if you are using partial or full CUDA GPU offloading.

Right now it just supports recent RTX GPUs, but support for Turing (RTX 2000) is coming soon.

I'm interested to hear your opinions about this. Seems like good stuff.

[drax-xard]

And with llama2 having a native 4k context this comes at the perfect time!

[imrohankataria]

Is this still in progress?

[ggerganov]

There is a WIP branch in #5021 - should already be able to run tests with it on most models