Optimizing Attention Flash

# Flash Attention - Fast Memory-Efficient Attention ## Quick start Flash Attention provides 2-4x speedup and 10-20x memory reduction for transformer attention through IO-aware tiling and recomputation. **PyTorch native (easiest, PyTorch 2.2+)**: ```python import torch import torch.nn.functional as F q = torch.randn(2, 8, 512, 64, device='cuda', dtype=torch.float16) # [batch, heads, seq, dim] k = torch.randn(2, 8, 512, 64, device='cuda', dtype=torch.float16) v = torch.randn(2, 8, 512, 64, device='cuda', dtype=torch.float16) # Automatically uses Flash Attention if available out = F.scaled_dot_product_attention(q, k, v) ``` **flash-attn library (more features)**: ```bash pip install flash-attn --no-build-isolation ``` ```python from flash_attn import flash_attn_func # q, k, v: [batch, seqlen, nheads, headdim] out = flash_attn_func(q, k, v, dropout_p=0.0, causal=True) ``` ## Common workflows ### Workflow 1: Enable in existing PyTorch model Copy this checklist: ``` Flash Attention Integration: - [ ] Step 1: Check PyTorch version (≥2.2) - [ ] Step 2: Enable Flash Attention backend - [ ] Step 3: Verify speedup with profiling - [ ] Step 4: Test accuracy matches baseline ``` **Step 1: Check PyTorch version** ```bash python -c "import torch; print(torch.__version__)" # Should be ≥2.2.0 ``` If 512 tokens. **Step 4: Test accuracy matches baseline** ```python # Compare outputs q, k, v = [torch.randn(1, 8, 512, 64, device='cuda', dtype=torch.float16) for _ in range(3)] # Flash Attention out_flash = F.scaled_dot_product_attention(q, k, v) # Standard attention attn_weights = torch.softmax(q @ k.transpose(-2, -1) / 8.0, dim=-1) out_standard = attn_weights @ v # Check difference diff = (out_flash - out_standard).abs().max() print(f"Max difference: {diff:.6f}") # Should be 512 tokens - Running inference with long context (>2K tokens) - GPU memory constrained (OOM with standard attention) - Need 2-4x speedup without accuracy loss - Using PyTorch 2.2+ or can install flash-attn **Use alternatives instead:** - **Standard attention**: Sequences <256 tokens (overhead not worth it) - **xFormers**: Need more attention variants (not just speed) - **Memory-efficient attention**: CPU inference (Flash Attention needs GPU) ## Common issues **Issue: ImportError: cannot import flash_attn** Install with no-build-isolation flag: ```bash pip install flash-attn --no-build-isolation ``` Or install CUDA toolkit first: ```bash conda install cuda -c nvidia pip install flash-attn --no-build-isolation ``` **Issue: Slower than expected (no speedup)** Flash Attention benefits increase with sequence length: - 2K tokens: 3-4x speedup Check sequence length is sufficient. **Issue: RuntimeError: CUDA error** Verify GPU supports Flash Attention: ```python import torch print(torch.cuda.get_device_capability()) # Should be ≥(7, 5) for Turing+ ``` Flash Attention requires: - Ampere (A100, A10): ✅ Full support - Turing (T4): ✅ Supported - Volta (V100): ❌ Not supported **Issue: Accuracy degradation** Check dtype is float16 or bfloat16 (not float32): ```python q = q.to(torch.float16) # Or torch.bfloat16 ``` Flash Attention uses float16/bfloat16 for speed. Float32 not supported. ## Advanced topics **Integration with HuggingFace Transformers**: See [references/transformers-integration.md](references/transformers-integration.md) for enabling Flash Attention in BERT, GPT, Llama models. **Performance benchmarks**: See [references/benchmarks.md](references/benchmarks.md) for detailed speed and memory comparisons across GPUs and sequence lengths. **Algorithm details**: See [references/algorithm.md](references/algorithm.md) for tiling strategy, recomputation, and IO complexity analysis. **Advanced features**: See [references/advanced-features.md](references/advanced-features.md) for rotary embeddings, ALiBi, paged KV cache, and custom attention masks. ## Hardware requirements - **GPU**: NVIDIA Ampere+ (A100, A10, A30) or AMD MI200+ - **VRAM**: Same as standard attention (Flash Attention doesn't increase memory) - **CUDA**: 12.0+ (11.8 minimum) - **PyTorch**: 2.2+ for native support **Not supported**: V100 (Volta), CPU inference ## Resources - Paper: "FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness" (NeurIPS 2022) - Paper: "FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning" (ICLR 2024) - Blog: https://tridao.me/blog/2024/flash3/ - GitHub: https://github.com/Dao-AILab/flash-attention - PyTorch docs: https://pytorch.org/docs/stable/generated/torch.nn.functional.scaled_dot_product_attention.html