Tracing 7 years of attention mechanism evolution from GPT-2 to Kimi K3 in runnable PyTorch

Short summary

A Baseten engineer traced 7 years of attention mechanism evolution from GPT-2 (124M params) to Kimi K3 (2.8T params) using runnable PyTorch code. The progression covers KV cache for memory bandwidth, linear attention to compress O(N²) to O(ND²), FlashAttention for IO optimization, and DeltaNet's delta rule for precise memory updates instead of blunt accumulation. Each step addresses a specific bottleneck in how models store, update, and retrieve memory.

• •KV cache solved recomputation but created memory bandwidth bottlenecks at long contexts
• •Linear attention replaced softmax with feature maps to compress O(N²D) to O(ND²), trading expressiveness for efficiency
• •DeltaNet introduced the delta rule for precise memory updates, enabling add/overwrite/correct instead of accumulation-only

Generated with AI, which can make mistakes.