Add Flash Decoding

[zfang]

Feature request

See https://pytorch.org/blog/flash-decoding/#:~:text=Flash%2DDecoding%20works%20in%203,exp%20of%20the%20attention%20values.

Motivation

Flash decoding further improves attention mechanism compared to FlashAttention V2 on long context

Your contribution

None

[ssmi153]

+1!

It looks like this might be included in FlashAttention v2.2. It's not clear from the blog whether any inference code needs to be changed to see the benefits of this.

[taishan1994]

+1

[OlivierDehaene]

It's not clear if it is superior to paged attention: all the tests I saw were vs native Transformers which we know is not optimised.
The kernel fusing is nice though. I will make some tests and report back if we want to have this or not.

[dongs0104]

@OlivierDehaene PA announce to PagedAttention V2 implements a similar idea to boost the performance when the batch size or the number of attention heads per GPU is small.

[OlivierDehaene]

#1183 instead.