optimizing-attention-flash

Use torch.nn.functional.scaled_dot_product_attention in PyTorch 2.2+, which automatically dispatches to Flash Attention when available. You can force the backend with torch.backends.cuda.sdp_kernel(enable_flash=True) and pass float16 tensors on CUDA.

Pass attn_implementation="flash_attention_2" to AutoModel.from_pretrained along with torch_dtype=torch.float16. This works for supported architectures like Llama, Mistral, Falcon, and Qwen with transformers 4.36+ and flash-attn installed.

PyTorch SDPA is the easiest option built into PyTorch 2.2+ and covers standard attention. The flash-attn library adds advanced features like multi-query attention, sliding window attention, and FP8 support on H100 GPUs.

Flash Attention requires NVIDIA Ampere or newer GPUs such as A100, A10, and H100, plus Turing GPUs like T4. V100 (Volta) and CPU inference are not supported, and CUDA 11.8 or higher is required.

Speedups depend on sequence length: under 512 tokens gains are only 10-20%, while 2K+ tokens yield 3-4x. Also confirm inputs are float16 or bfloat16, since float32 is not supported by Flash Attention kernels.

Avoid it for sequences under 256 tokens where overhead outweighs gains, for CPU inference, or when you need attention variants beyond speed optimization, where xFormers or memory-efficient attention may fit better.