Intermediate LevelAI, LLMs, Embeddings & MCP
How does FlashAttention (FlashAttention-2/3) accelerate Transformer training and inference?
❓ Interview Prompt: Explain memory hierarchy (HBM vs SRAM), tiling, online softmax, and I/O-aware exact attention computation.
Standard Attention computes the full $N \times N$ attention matrix and writes it to GPU High Bandwidth Memory (HBM). For long sequences, HBM memory read/write bandwidth becomes the main bottleneck. FlashAttention is an I/O-aware exact attention algorithm.
Core Technical Innovations:
1
Tiling: Splitting Query, Key, and Value matrices into small blocks that fit entirely inside ultra-fast GPU SRAM on-chip cache (19TB/s bandwidth vs 2TB/s HBM).
2
Online Softmax: Computes exact Softmax scaling incrementally block-by-block without storing the giant $N \times N$ attention matrix in HBM.
Speedup:
Provides 2x to 4x faster training/inference while reducing memory usage from $O(N^2)$ to $O(N)$ without any mathematical approximation!
Code Implementation Example
python
# FlashAttention usage in PyTorch / HuggingFace (Python)
import torch
# Standard PyTorch 2.0+ automatically uses FlashAttention-2 scaled_dot_product_attention
output = torch.nn.functional.scaled_dot_product_attention(
query, key, value,
is_causal=True # Enables GPU SRAM tiling kernel automatically!
)💡Invoking FlashAttention-2 kernel via PyTorch scaled_dot_product_attention.
🎯 Interview Key Takeaway
FlashAttention uses GPU SRAM tiling and online softmax to compute exact attention without writing $N \times N$ matrices to HBM.
#FlashAttention#GPU#CUDA#SRAM#Intermediate
Back to all AI, LLMs, Embeddings & MCP Questions→