{"id":"flash-attention","name":"optimizing-attention-flash","summary":"フラッシュアテンションでトランスフォーマーの注意力を最適化し、2〜4倍の高速化と10〜20倍のメモリ削減を実現します。","body":"# Flash Attention - Fast Memory-Efficient Attention\n\n## Quick start\n\nFlash Attention provides 2-4x speedup and 10-20x memory reduction for transformer attention through IO-aware tiling and recomputation.\n\n**PyTorch native (easiest, PyTorch 2.2+)**:\n```python\nimport torch\nimport torch.nn.functional as F\n\nq = torch.randn(2, 8, 512, 64, device='cuda', dtype=torch.float16)  # [batch, heads, seq, dim]\nk = torch.randn(2, 8, 512, 64, device='cuda', dtype=torch.float16)\nv = torch.randn(2, 8, 512, 64, device='cuda', dtype=torch.float16)\n\n# Automatically uses Flash Attention if available\nout = F.scaled_dot_product_attention(q, k, v)\n```\n\n**flash-attn library (more features)**:\n```bash\npip install flash-attn --no-build-isolation\n```\n\n```python\nfrom flash_attn import flash_attn_func\n\n# q, k, v: [batch, seqlen, nheads, headdim]\nout = flash_attn_func(q, k, v, dropout_p=0.0, causal=True)\n```\n\n## Common workflows\n\n### Workflow 1: Enable in existing PyTorch model\n\nCopy this checklist:\n\n```\nFlash Attention Integration:\n- [ ] Step 1: Check PyTorch version (≥2.2)\n- [ ] Step 2: Enable Flash Attention backend\n- [ ] Step 3: Verify speedup with profiling\n- [ ] Step 4: Test accuracy matches baseline\n```\n\n**Step 1: Check PyTorch version**\n\n```bash\npython -c \"import torch; print(torch.__version__)\"\n# Should be ≥2.2.0\n```\n\nIf <2.2, upgrade:\n```bash\npip install --upgrade torch\n```\n\n**Step 2: Enable Flash Attention backend**\n\nReplace standard attention:\n```python\n# Before (standard attention)\nattn_weights = torch.softmax(q @ k.transpose(-2, -1) / math.sqrt(d_k), dim=-1)\nout = attn_weights @ v\n\n# After (Flash Attention)\nimport torch.nn.functional as F\nout = F.scaled_dot_product_attention(q, k, v, attn_mask=mask)\n```\n\nForce Flash Attention backend:\n```python\nwith torch.backends.cuda.sdp_kernel(\n    enable_flash=True,\n    enable_math=False,\n    enable_mem_efficient=False\n):\n    out = F.scaled_dot_product_attention(q, k, v)\n```\n\n**Step 3: Verify speedup with profiling**\n\n```python\nimport torch.utils.benchmark as benchmark\n\ndef test_attention(use_flash):\n    q, k, v = [torch.randn(2, 8, 2048, 64, device='cuda', dtype=torch.float16) for _ in range(3)]\n\n    if use_flash:\n        with torch.backends.cuda.sdp_kernel(enable_flash=True):\n            return F.scaled_dot_product_attention(q, k, v)\n    else:\n        attn = (q @ k.transpose(-2, -1) / 8.0).softmax(dim=-1)\n        return attn @ v\n\n# Benchmark\nt_flash = benchmark.Timer(stmt='test_attention(True)', globals=globals())\nt_standard = benchmark.Timer(stmt='test_attention(False)', globals=globals())\n\nprint(f\"Flash: {t_flash.timeit(100).mean:.3f}s\")\nprint(f\"Standard: {t_standard.timeit(100).mean:.3f}s\")\n```\n\nExpected: 2-4x speedup for sequences >512 tokens.\n\n**Step 4: Test accuracy matches baseline**\n\n```python\n# Compare outputs\nq, k, v = [torch.randn(1, 8, 512, 64, device='cuda', dtype=torch.float16) for _ in range(3)]\n\n# Flash Attention\nout_flash = F.scaled_dot_product_attention(q, k, v)\n\n# Standard attention\nattn_weights = torch.softmax(q @ k.transpose(-2, -1) / 8.0, dim=-1)\nout_standard = attn_weights @ v\n\n# Check difference\ndiff = (out_flash - out_standard).abs().max()\nprint(f\"Max difference: {diff:.6f}\")\n# Should be <1e-3 for float16\n```\n\n### Workflow 2: Use flash-attn library for advanced features\n\nFor multi-query attention, sliding window, or H100 FP8.\n\nCopy this checklist:\n\n```\nflash-attn Library Setup:\n- [ ] Step 1: Install flash-attn library\n- [ ] Step 2: Modify attention code\n- [ ] Step 3: Enable advanced features\n- [ ] Step 4: Benchmark performance\n```\n\n**Step 1: Install flash-attn library**\n\n```bash\n# NVIDIA GPUs (CUDA 12.0+)\npip install flash-attn --no-build-isolation\n\n# Verify installation\npython -c \"from flash_attn import flash_attn_func; print('Success')\"\n```\n\n**Step 2: Modify attention code**\n\n```python\nfrom flash_attn import flash_attn_func\n\n# Input: [batch_size, seq_len, num_heads, head_dim]\n# Transpose from [batch, heads, seq, dim] if needed\nq = q.transpose(1, 2)  # [batch, seq, heads, dim]\nk = k.transpose(1, 2)\nv = v.transpose(1, 2)\n\nout = flash_attn_func(\n    q, k, v,\n    dropout_p=0.1,\n    causal=True,  # For autoregressive models\n    window_size=(-1, -1),  # No sliding window\n    softmax_scale=None  # Auto-scale\n)\n\nout = out.transpose(1, 2)  # Back to [batch, heads, seq, dim]\n```\n\n**Step 3: Enable advanced features**\n\nMulti-query attention (shared K/V across heads):\n```python\nfrom flash_attn import flash_attn_func\n\n# q: [batch, seq, num_q_heads, dim]\n# k, v: [batch, seq, num_kv_heads, dim]  # Fewer KV heads\nout = flash_attn_func(q, k, v)  # Automatically handles MQA\n```\n\nSliding window attention (local attention):\n```python\n# Only attend to window of 256 tokens before/after\nout = flash_attn_func(\n    q, k, v,\n    window_size=(256, 256),  # (left, right) window\n    causal=True\n)\n```\n\n**Step 4: Benchmark performance**\n\n```python\nimport torch\nfrom flash_attn import flash_attn_func\nimport time\n\nq, k, v = [torch.randn(4, 4096, 32, 64, device='cuda', dtype=torch.float16) for _ in range(3)]\n\n# Warmup\nfor _ in range(10):\n    _ = flash_attn_func(q, k, v)\n\n# Benchmark\ntorch.cuda.synchronize()\nstart = time.time()\nfor _ in range(100):\n    out = flash_attn_func(q, k, v)\n    torch.cuda.synchronize()\nend = time.time()\n\nprint(f\"Time per iteration: {(end-start)/100*1000:.2f}ms\")\nprint(f\"Memory allocated: {torch.cuda.max_memory_allocated()/1e9:.2f}GB\")\n```\n\n### Workflow 3: H100 FP8 optimization (FlashAttention-3)\n\nFor maximum performance on H100 GPUs.\n\n```\nFP8 Setup:\n- [ ] Step 1: Verify H100 GPU available\n- [ ] Step 2: Install flash-attn with FP8 support\n- [ ] Step 3: Convert inputs to FP8\n- [ ] Step 4: Run with FP8 attention\n```\n\n**Step 1: Verify H100 GPU**\n\n```bash\nnvidia-smi --query-gpu=name --format=csv\n# Should show \"H100\" or \"H800\"\n```\n\n**Step 2: Install flash-attn with FP8 support**\n\n```bash\npip install flash-attn --no-build-isolation\n# FP8 support included for H100\n```\n\n**Step 3: Convert inputs to FP8**\n\n```python\nimport torch\n\nq = torch.randn(2, 4096, 32, 64, device='cuda', dtype=torch.float16)\nk = torch.randn(2, 4096, 32, 64, device='cuda', dtype=torch.float16)\nv = torch.randn(2, 4096, 32, 64, device='cuda', dtype=torch.float16)\n\n# Convert to float8_e4m3 (FP8)\nq_fp8 = q.to(torch.float8_e4m3fn)\nk_fp8 = k.to(torch.float8_e4m3fn)\nv_fp8 = v.to(torch.float8_e4m3fn)\n```\n\n**Step 4: Run with FP8 attention**\n\n```python\nfrom flash_attn import flash_attn_func\n\n# FlashAttention-3 automatically uses FP8 kernels on H100\nout = flash_attn_func(q_fp8, k_fp8, v_fp8)\n# Result: ~1.2 PFLOPS, 1.5-2x faster than FP16\n```\n\n## When to use vs alternatives\n\n**Use Flash Attention when:**\n- Training transformers with sequences >512 tokens\n- Running inference with long context (>2K tokens)\n- GPU memory constrained (OOM with standard attention)\n- Need 2-4x speedup without accuracy loss\n- Using PyTorch 2.2+ or can install flash-attn\n\n**Use alternatives instead:**\n- **Standard attention**: Sequences <256 tokens (overhead not worth it)\n- **xFormers**: Need more attention variants (not just speed)\n- **Memory-efficient attention**: CPU inference (Flash Attention needs GPU)\n\n## Common issues\n\n**Issue: ImportError: cannot import flash_attn**\n\nInstall with no-build-isolation flag:\n```bash\npip install flash-attn --no-build-isolation\n```\n\nOr install CUDA toolkit first:\n```bash\nconda install cuda -c nvidia\npip install flash-attn --no-build-isolation\n```\n\n**Issue: Slower than expected (no speedup)**\n\nFlash Attention benefits increase with sequence length:\n- <512 tokens: Minimal speedup (10-20%)\n- 512-2K tokens: 2-3x speedup\n- >2K tokens: 3-4x speedup\n\nCheck sequence length is sufficient.\n\n**Issue: RuntimeError: CUDA error**\n\nVerify GPU supports Flash Attention:\n```python\nimport torch\nprint(torch.cuda.get_device_capability())\n# Should be ≥(7, 5) for Turing+\n```\n\nFlash Attention requires:\n- Ampere (A100, A10): ✅ Full support\n- Turing (T4): ✅ Supported\n- Volta (V100): ❌ Not supported\n\n**Issue: Accuracy degradation**\n\nCheck dtype is float16 or bfloat16 (not float32):\n```python\nq = q.to(torch.float16)  # Or torch.bfloat16\n```\n\nFlash Attention uses float16/bfloat16 for speed. Float32 not supported.\n\n## Advanced topics\n\n**Integration with HuggingFace Transformers**: See [references/transformers-integration.md](references/transformers-integration.md) for enabling Flash Attention in BERT, GPT, Llama models.\n\n**Performance benchmarks**: See [references/benchmarks.md](references/benchmarks.md) for detailed speed and memory comparisons across GPUs and sequence lengths.\n\n**Algorithm details**: See [references/algorithm.md](references/algorithm.md) for tiling strategy, recomputation, and IO complexity analysis.\n\n**Advanced features**: See [references/advanced-features.md](references/advanced-features.md) for rotary embeddings, ALiBi, paged KV cache, and custom attention masks.\n\n## Hardware requirements\n\n- **GPU**: NVIDIA Ampere+ (A100, A10, A30) or AMD MI200+\n- **VRAM**: Same as standard attention (Flash Attention doesn't increase memory)\n- **CUDA**: 12.0+ (11.8 minimum)\n- **PyTorch**: 2.2+ for native support\n\n**Not supported**: V100 (Volta), CPU inference\n\n## Resources\n\n- Paper: \"FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness\" (NeurIPS 2022)\n- Paper: \"FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning\" (ICLR 2024)\n- Blog: https://tridao.me/blog/2024/flash3/\n- GitHub: https://github.com/Dao-AILab/flash-attention\n- PyTorch docs: https://pytorch.org/docs/stable/generated/torch.nn.functional.scaled_dot_product_attention.html","author":"@Orchestra-Research","ownerProfile":null,"authorContacts":null,"sourceUrl":"https://github.com/Orchestra-Research/AI-Research-SKILLs/tree/main/10-optimization/flash-attention","license":"MIT","category":"writing","lang":"en","tokens":2755,"stars":0,"calls30d":1,"claimed":false,"visibility":"public","origin":"crawler","version":"0.1.0","createdAt":"2026-08-22","updatedAt":"2026-08-22","files":[{"path":"references/benchmarks.md","size":7129,"sha256":"264ff3060af824c32f7a6e3c0e0208c38646c728b2292fb35ede923a38d60b6b"},{"path":"references/transformers-integration.md","size":7427,"sha256":"89ebe8e20756ee218103a1144d4682e960a0faa07870c2588530bc8d26c441a1"}],"requires":{"mcp":[],"tools":[]},"safety":{"flags":[],"scannedAt":"2026-08-22","hasScripts":false,"networkEndpoints":["pytorch.org","tridao.me"]}}