Flash attention

[minhthuc2502]

I currently work on the integration for flash attention. It based on the kernel developed in original repo.

For the current version, I don't see any improvement in performance compared with the standard MHA. Tested with GPU A100.

Something to consider here:

• Implement flash attention with kvcache as fused.
• Accept loss in perf by additional transpose operation for qkv to align the same shape between standard and flash attention

Update:

• With long prompt, we have a good improvement in performance.
• The size of package increases after the integration of feature
• Compiling time increases due to the heavy use of template (but it improves the perf in runtime)

[BBC-Esq]

Super excited about this. Let me know if you need someone with Windows and an RTX 4090 to test.

[Purfview]

For the current version, I don't see any improvement in performance

For performance you probably should switch to cuDNN 9.x

From cuDNN 9.0.0 release notes:

FP16 and BF16 fused flash attention engine performance has been significantly improved for NVIDIA GPUs:
    Speed-up of up to 50% over cuDNN 8.9.7 on Hopper GPUs.
    Speed-up of up to 100% over cuDNN 8.9.7 on Ampere GPUs.

[minhthuc2502]

For performance you probably should switch to cuDNN 9.x

I think the cuDNN is currently used only for conv ops, so it wouldn't affect the perf of the flash-attention. By the way, I see an improvement when trying with long length of input. I will do some more tests and improve the time of compilation.