{"slug": "how-a-baseten-engineer-traced-7-years-of-attention-mechanism-evolution-from-gpt", "title": "How a Baseten Engineer Traced 7 Years of Attention Mechanism Evolution -- From GPT-2 to Kimi K3, in Runable PyTorch", "summary": "A Baseten inference engineer known as @waterloo_intern published a technical blog post titled '22,580: From GPT-2 to Kimi K3, Explained,' which has garnered 2.4 million views. The post provides runnable PyTorch code tracing the evolution of attention mechanisms from GPT-2 to Kimi K3, covering innovations like KV cache, linear attention, and DeltaNet. The engineer explains how memory bandwidth became the bottleneck and how architectural changes addressed it, culminating in Moonshot AI's Kimi K3 with 2.8T parameters.", "body_md": "Last week, a Baseten inference engineer who goes by [@waterloo_intern](https://x.com/waterloo_intern/status/2081762065392541951) published a technical blog post titled **\"22,580: From GPT-2 to Kimi K3, Explained.\"** It hit 2.4 million views in days.\n\nHe didn't write a press release. He wrote **runnable PyTorch code** — starting from GPT-2's attention block, stepping through every architectural change, explaining one problem and one cost per iteration. It's the best transformer lineage explanation I've seen.\n\nI devoured his post, then cross-checked the key claims against 5 original papers. Here's the full picture.\n\nIn February 2019, OpenAI released GPT-2 — 124M parameters. Seven years later, Moonshot AI open-sourced Kimi K3 — 2.8T parameters. You could fit **22,580 GPT-2s inside one Kimi K3**.\n\nBut this isn't a \"throw more compute at it\" story. It's a story about **how we store, update, and retrieve memory**.\n\n``` python\nclass Block(nn.Module):\n    def forward(self, x):\n        x = x + self.attn(self.ln_1(x))\n        x = x + self.mlp(self.ln_2(x))\n        return x\n```\n\nEvery time the model generates a new token, it recomputes Q, K, V projections for **all** historical tokens, then runs an O(N²) softmax attention. K and V from tokens 1 through N-1? Thrown away. Token N+1 arrives? Recompute everything.\n\nThat's why KV Cache was invented.\n\nSimple idea: cache the already-computed keys and values. For the next token, new Q only needs one dot product against the cached K.\n\nProblem solved — but a new one created. KV cache grows **linearly** with sequence length. At 1M tokens × d_model × layers, that's dozens of GB of VRAM. Every decoding step reads all of it from HBM.\n\n**The bottleneck isn't compute. It's memory bandwidth.** This is the key to understanding every improvement that follows.\n\nCan we compress O(N²D) into O(ND²)?\n\nThe idea: replace softmax with a feature map.\n\n```\n# Standard softmax (must materialize N×N first)\nattention = softmax(QKᵀ / √d) × V\n\n# Linear attention (fold K and V first)\nQ' = ELU(Q) + 1\nK' = ELU(K) + 1\noutput = (Q'(K')ᵀ × V') / (Q'(K')ᵀ × ones)\n```\n\nNow we can compute `K'^T × V'`\n\nfirst — a fixed-size D×D matrix — then multiply by Q'. Historical KV information gets \"folded\" into a constant-size state matrix. Cache no longer grows with N.\n\nThe cost? ELU+1 is an approximation of the softmax kernel. Expressiveness drops. But on long-context tasks, this tradeoff is often worth it.\n\nA common confusion point. Ali's post mentions \"in 2020, there was no FlashAttention\" — but FlashAttention and linear attention solve fundamentally different problems.\n\nFlashAttention (Tri Dao, NeurIPS 2022) didn't change the attention **algorithm**. It optimized the GPU **IO pattern** — tiling the N×N matrix so it never fully lands in HBM. It makes softmax faster, but it's still O(N²).\n\nLinear attention **redefined the algorithm itself** — replacing softmax with a feature map, going from O(N²D) to O(ND²).\n\nOne optimizes IO. One changes the algorithm. Orthogonal.\n\nLinear attention has a fatal flaw: it can only **add**, never **update**.\n\nEvery new token piles onto the state: `S = S + K'^T × V'`\n\n. Information only grows. It's like a notebook where you can only write — never erase or correct.\n\nDeltaNet (Songlin Yang et al., NeurIPS 2024) fixes this.\n\nThe core idea traces back to Schmidhuber's \"Fast Weight Programmers\" from the 1990s, later formalized by Schlag et al. (ICML 2021) linking linear attention to fast weights. DeltaNet's approach: before writing, read what this key position currently stores.\n\n```\nv_old = k @ S_old           # read current value at key\ndelta = v_new - v_old       # compute the delta\nS_new = S_old + k^T @ delta # write only the difference\n```\n\nIf `v_new == v_old`\n\n, delta is zero — nothing changes. If completely different, delta equals v_new — equivalent to an overwrite. Everything in between is a smooth interpolation.\n\nThis is the **delta rule**: precise memory updates instead of blunt accumulation.\n\nDeltaNet's state update is strictly sequential — each step depends on the previous S. It looks impossible to parallelize. But it is.\n\nThe method: split the sequence into chunks of size C. Within each chunk, use normal masked attention (GPU-parallel). Between chunks, use the state matrix + one matmul (Q @ S). A Householder transformation reparameterizes the delta updates so all deltas within a chunk can be computed at once.\n\nComplexity: fixed cost 2LD² (state maintenance) + variable cost 2LCD (intra-chunk attention). Bigger C = more variable cost but better GPU efficiency. In practice, C=64 or 128 works best — FLOPs aren't the only metric; tensor core utilization matters too.\n\nDeltaNet can precisely edit individual key-value pairs. But what about **forgetting at scale**?\n\nImagine reading through all documents about Topic A in a 1M-token context, then switching to Topic B. The model should ideally \"forget\" Topic A to free up capacity for B.\n\nDeltaNet can overwrite specific entries but can't bulk-decay global memory. Mamba-2 (Dao & Gu, ICML 2024) can:\n\n```\ncache = α × S_old + S_new\n```\n\nα is a gating value between 0 and 1, uniformly decaying all old memories. This is the core of Mamba-2's \"State Space Duality\" theory — softmax attention and SSMs are mathematically the same thing expressed differently. Mamba-2 unified them through gating.\n\nGated DeltaNet (Songlin Yang et al., ICLR 2025, NVIDIA) merges DeltaNet's delta updates with Mamba-2's gating:\n\n```\nS_new = α × S_old + k^T @ delta   # decay first, then precise write\n```\n\nα=1 is pure DeltaNet. α=0 is a memory wipe. Everything in between both forgets and updates.\n\nThe crucial mechanism: a token written at timestep x, read at x+t, has been through t rounds of cumulative α decay (αₓ × αₓ₊₁ × … × αₓ₊ₜ). **Information written at different times forgets at different speeds** — recent items decay little, distant ones may be fully gone.\n\nKimi Linear (Moonshot AI, arXiv 2510.26692, October 2025) — the predecessor to K3 — refined the idea further: **from scalar gating to per-dimension gating**.\n\nGated DeltaNet uses α as a single scalar. One number controlling forgetting speed for all memory dimensions. KDA turns α into a **vector** (or matrix). Each dimension independently controls its own forgetting rate. Which concepts to retain, which dimensions to decay — the model learns it.\n\nKey data from the paper:\n\n| Metric | Full MLA | Kimi Linear |\n|---|---|---|\n| KV Cache | 100% | -75% |\n| 1M-context decode throughput | Baseline | 6x |\n| Short-context performance | Baseline | Outperforms |\n| RL scaling | Baseline | Outperforms |\n\nThis is the **first time linear attention outperforms full attention under fair comparisons across all scenarios.** Not just long context — short context too. And that \"-75% KV cache\" translates directly to inference cost savings.\n\nI verified this claim directly against the paper's abstract: *\"for the first time, outperforms full attention under fair comparisons across various scenarios.\"* Not marketing. Core conclusion.\n\nK3's technical report (arXiv 2607.24653, July 2026) confirms:\n\nK3 isn't a pure linear-attention model. It's a **hybrid system**: KDA handles the bulk processing (cheap, fast, fixed-state), while periodic softmax layers do precision retrieval — recovering details that linear compression might lose. Exactly the tradeoff Ali analyzed: linear attention + periodic softmax retrieval.\n\nAli's conclusion is sharper than anything I could write:\n\nFrom GPT-2 to Kimi K3, each generation's core improvement wasn't \"more parameters\" — it was redesigning\n\nhow memory is accessed: how you store, how you forget, how you retrieve.\n\nThe evolution in one table:\n\n| Stage | Representative | Memory Mechanism | Problem Solved |\n|---|---|---|---|\n| 1 | GPT-2 + KV Cache | Full cache, O(N) growth | Eliminate recomputation |\n| 2 | Linear Attention | Fixed state, O(D²) | Cache doesn't grow with N |\n| 3 | DeltaNet | Delta updates, editable | Can't update → can update |\n| 4 | Gated DeltaNet | Gating + delta, forgettable | Can't forget → can forget |\n| 5 | KDA / Kimi K3 | Per-dim gating + hybrid | Linear surpasses full attention |\n\nParameters grew 22,580×. But if that's all you see, you missed the entire story.\n\nIf you're doing long-context work (code review, document analysis, multi-turn agents), attention architecture directly impacts your cost and output quality.\n\n**Cost isn't fixed.** Same 1M-token context: KDA's KV cache is only 25% of full attention's. Lower memory bandwidth pressure, significantly better latency and throughput.\n\n**Long ≠ expensive.** K3's 1M-token context window isn't \"brute-forced.\" 75% of work goes through the linear path.\n\n**Hybrid is the trend.** Linear won't replace softmax — they'll complement each other. Precision retrieval via softmax, bulk processing via linear. This paradigm will spread.", "url": "https://wpnews.pro/news/how-a-baseten-engineer-traced-7-years-of-attention-mechanism-evolution-from-gpt", "canonical_source": "https://dev.to/cdragon123code/how-a-baseten-engineer-traced-7-years-of-attention-mechanism-evolution-from-gpt-2-to-kimi-k3-in-pl7", "published_at": "2026-07-31 03:23:15+00:00", "updated_at": "2026-07-31 03:59:19.379501+00:00", "lang": "en", "topics": ["artificial-intelligence", "machine-learning", "large-language-models", "ai-research", "developer-tools"], "entities": ["Baseten", "OpenAI", "Moonshot AI", "Kimi K3", "GPT-2", "Tri Dao", "Songlin Yang", "Schmidhuber"], "alternates": {"html": "https://wpnews.pro/news/how-a-baseten-engineer-traced-7-years-of-attention-mechanism-evolution-from-gpt", "markdown": "https://wpnews.pro/news/how-a-baseten-engineer-traced-7-years-of-attention-mechanism-evolution-from-gpt.md", "text": "https://wpnews.pro/news/how-a-baseten-engineer-traced-7-years-of-attention-mechanism-evolution-from-gpt.txt", "jsonld": "https://wpnews.pro/news/how-a-baseten-engineer-traced-7-years-of-attention-mechanism-evolution-from-gpt.jsonld"}}