The majority of the optimizations described here also apply to multi-GPU setups!

FlashAttention-2

FlashAttention-2 is experimental and may change considerably in future versions.

FlashAttention-2 is a faster and more efficient implementation of the standard attention mechanism that can significantly speedup inference by:

additionally parallelizing the attention computation over sequence length
partitioning the work between GPU threads to reduce communication and shared memory reads/writes between them