--- id: wiki-2026-0508-flash-attention title: Flash Attention category: 10_Wiki/Topics status: verified canonical_id: self aliases: [FlashAttention, FA2, FA3, IO-aware attention, Tri Dao, online softmax] duplicate_of: none source_trust_level: A confidence_score: 0.98 verification_status: applied tags: [transformer, attention, gpu, optimization, flash-attention, memory-efficient] raw_sources: [] last_reinforced: 2026-05-10 github_commit: pending tech_stack: language: CUDA / PyTorch framework: flash-attn / xformers / vLLM --- # Flash Attention ## 매 한 줄 > **"매 attention 의 IO-aware tile-based exact algorithm"**. Tri Dao 2022 (FA1), 2023 (FA2), 2024 (FA3). 매 quadratic memory 의 fix — 매 O(N²) → 매 O(N) memory. 매 modern transformer 의 standard. 매 vLLM, xformers, native PyTorch. ## 매 핵심 ### 매 problem (vanilla) - **Standard attention**: 매 O(N²) memory (매 N×N attention matrix). - **HBM bandwidth**: 매 bottleneck (>FLOPS). - **Long context**: 매 OOM. ### 매 solution (Flash) - **Tile** Q, K, V into blocks. - **Online softmax**: 매 incremental, no full matrix. - **SRAM compute**: 매 fast on-chip. - **Recomputation**: 매 backward 의 의 의 trade compute for memory. ### 매 versions - **FA1** (2022): 매 baseline. - **FA2** (2023): 매 better parallelism, 2x faster. - **FA3** (2024): 매 H100-optimized, async. ### 매 응용 1. **All transformer training**. 2. **Long-context** (100K+). 3. **Inference** (vLLM, TGI). 4. **Multi-query / GQA**. 5. **Sparse / sliding window**. ## 💻 패턴 ### PyTorch native (FA built-in) ```python import torch import torch.nn.functional as F # 매 PyTorch 2.0+ scaled_dot_product_attention auto-uses Flash if eligible out = F.scaled_dot_product_attention(q, k, v, is_causal=True) ``` ### flash-attn (Tri Dao package) ```python from flash_attn import flash_attn_func, flash_attn_varlen_func # 매 standard out = flash_attn_func(q, k, v, dropout_p=0.0, causal=True) # 매 variable length (no padding waste) out = flash_attn_varlen_func(q, k, v, cu_seqlens_q, cu_seqlens_k, max_seqlen_q, max_seqlen_k, causal=True) ``` ### xformers ```python from xformers.ops import memory_efficient_attention out = memory_efficient_attention(q, k, v, attn_bias=causal_mask) ``` ### vLLM (paged attention serving) ```python from vllm import LLM, SamplingParams llm = LLM(model='meta-llama/Llama-3-8B', dtype='bfloat16') # 매 internally uses Flash + paged attention outputs = llm.generate(['Hello'], SamplingParams(max_tokens=100)) ``` ### Manual Flash-style (educational, simplified) ```python def flash_attention_simple(Q, K, V, block_size=64): """매 simplified — actual implementation 의 CUDA.""" N = Q.shape[1] O = torch.zeros_like(Q) L = torch.zeros(Q.shape[:2]) # 매 max M = torch.full(Q.shape[:2], float('-inf')) # 매 normalize for j in range(0, N, block_size): Kj = K[:, j:j+block_size] Vj = V[:, j:j+block_size] for i in range(0, N, block_size): Qi = Q[:, i:i+block_size] Sij = Qi @ Kj.transpose(-1, -2) Mij = Sij.max(dim=-1, keepdim=True).values Mi_new = torch.maximum(M[:, i:i+block_size, None], Mij) Pij = torch.exp(Sij - Mi_new) # 매 online normalization scale = torch.exp(M[:, i:i+block_size, None] - Mi_new) O[:, i:i+block_size] = O[:, i:i+block_size] * scale + Pij @ Vj M[:, i:i+block_size] = Mi_new.squeeze(-1) return O / L # 매 simplified ``` ### Sliding window (Mistral-style) ```python from flash_attn import flash_attn_func out = flash_attn_func(q, k, v, window_size=(window_left, 0), causal=True) ``` ### Grouped Query Attention (GQA) ```python class GQA(nn.Module): def __init__(self, dim, n_heads, n_kv_heads): super().__init__() self.q = nn.Linear(dim, n_heads * head_dim) self.k = nn.Linear(dim, n_kv_heads * head_dim) self.v = nn.Linear(dim, n_kv_heads * head_dim) def forward(self, x): q = self.q(x).view(...) k = self.k(x).view(...).repeat_interleave(n_heads // n_kv_heads, dim=2) v = self.v(x).view(...).repeat_interleave(n_heads // n_kv_heads, dim=2) return flash_attn_func(q, k, v, causal=True) ``` ### KV cache (inference) ```python # 매 paged attention class PagedKVCache: def __init__(self, n_layers, max_seqs, block_size=16): self.blocks = {} # 매 logical block → physical self.block_size = block_size def append(self, seq_id, k_block, v_block): physical = self.allocate_block() self.blocks[(seq_id, len(self.blocks))] = physical # 매 → flash_attn_with_kvcache ``` ### Backward (recomputation) ```python # 매 forward 의 small statistics + recompute on backward # 매 native to flash_attn — automatic out = flash_attn_func(q, k, v).backward() ``` ### Compile + Flash ```python # 매 PyTorch 2.x compile 의 fuse model = torch.compile(model) # 매 internally uses sdpa (Flash if available) ``` ### Detect Flash availability ```python def has_flash(): try: from flash_attn import flash_attn_func return torch.cuda.is_available() and torch.cuda.get_device_capability()[0] >= 8 except ImportError: return False ``` ### H100 / FA3 ```python # 매 fa3 (2024) — H100 hopper async from flash_attn_interface import flash_attn_func # 매 same API, 1.5-2x faster on H100 ``` ### Mask custom (block-sparse) ```python # 매 매 custom mask 의 efficient 의 X # 매 fully sparse (e.g., longformer global+local) → flash-attn variants from flash_attn.flash_attn_triton import flash_attn_func out = flash_attn_func(q, k, v, custom_block_mask) ``` ### vLLM serving ```bash python -m vllm.entrypoints.openai.api_server \ --model meta-llama/Llama-3-8B \ --dtype bfloat16 \ --max-model-len 32768 ``` ## 매 결정 기준 | 상황 | Approach | |---|---| | Default training | PyTorch sdpa (auto) | | Long context | flash_attn_varlen | | Production serving | vLLM (paged) | | Custom mask | xformers / flash-attn variants | | H100 | FA3 | | Mobile / non-CUDA | Use math fallback | **기본값**: 매 PyTorch sdpa + 매 vLLM serving + 매 GQA + 매 paged KV cache + 매 H100 FA3. ## 🔗 Graph - 부모: [[Transformer]] · [[Attention-Mechanism]] - 변형: [[Paged-Attention]] · [[Sliding-Window]] · [[GQA]] - 응용: [[LLM_Optimization_and_Deployment_Strategies|vLLM]] · [[Long-Context]] - Adjacent: [[LLM_Optimization_and_Deployment_Strategies|Quantization]] · [[Foundation-Models]] ## 🤖 LLM 활용 **언제**: 매 모든 transformer training/inference. **언제 X**: 매 non-CUDA (mobile). ## ❌ 안티패턴 - **Manual attention loop**: 매 slow. - **Pad to max in batch**: 매 use varlen. - **No KV cache**: 매 inference quadratic. - **Old non-Flash 의 prod**: 매 cost ↑. ## 🧪 검증 / 중복 - Verified (Dao 2022/2023/2024 FA papers, vLLM Kwon 2023). - 신뢰도 A. ## 🕓 Changelog | 날짜 | 변경 | |---|---| | 2026-04-20 | Auto-reinforced | | 2026-05-08 | Phase 1 | | 2026-05-10 | Manual cleanup — algorithm + 매 PyTorch / flash-attn / vLLM / GQA / FA3 code |