Accelerating transformer blocks
PyTorch 2.0 includes a scaled dot-product attention function as part of torch.nn.functional. This function encompasses several implementations that can be applied depending on the inputs and the hardware in use. Before PyTorch 2.0, you had to search for third-party implementations and install separate packages in order to take advantage of memory optimized algorithms, such as FlashAttention. The available implementations are:
* FlashAttention, from the official FlashAttention project. 
* Memory-Efficient Attention, from the xFormers project.
* A native C++ implementation suitable for non-CUDA devices or when high-precision is required.