cd /news/artificial-intelligence/how-a-baseten-engineer-traced-7-year… · home topics artificial-intelligence article
[ARTICLE · art-81272] src=dev.to ↗ pub= topic=artificial-intelligence verified=true sentiment=· neutral

How a Baseten Engineer Traced 7 Years of Attention Mechanism Evolution -- From GPT-2 to Kimi K3, in Runable PyTorch

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.

read7 min views1 publishedJul 31, 2026

Last week, a Baseten inference engineer who goes by @waterloo_intern published a technical blog post titled "22,580: From GPT-2 to Kimi K3, Explained." It hit 2.4 million views in days.

He 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.

I devoured his post, then cross-checked the key claims against 5 original papers. Here's the full picture.

In 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.

But this isn't a "throw more compute at it" story. It's a story about how we store, update, and retrieve memory.

class Block(nn.Module):
    def forward(self, x):
        x = x + self.attn(self.ln_1(x))
        x = x + self.mlp(self.ln_2(x))
        return x

Every 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.

That's why KV Cache was invented.

Simple idea: cache the already-computed keys and values. For the next token, new Q only needs one dot product against the cached K.

Problem 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.

The bottleneck isn't compute. It's memory bandwidth. This is the key to understanding every improvement that follows.

Can we compress O(N²D) into O(ND²)?

The idea: replace softmax with a feature map.

attention = softmax(QKᵀ / √d) × V

Q' = ELU(Q) + 1
K' = ELU(K) + 1
output = (Q'(K')ᵀ × V') / (Q'(K')ᵀ × ones)

Now we can compute K'^T × V'

first — 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.

The cost? ELU+1 is an approximation of the softmax kernel. Expressiveness drops. But on long-context tasks, this tradeoff is often worth it.

A common confusion point. Ali's post mentions "in 2020, there was no FlashAttention" — but FlashAttention and linear attention solve fundamentally different problems.

FlashAttention (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²).

Linear attention redefined the algorithm itself — replacing softmax with a feature map, going from O(N²D) to O(ND²).

One optimizes IO. One changes the algorithm. Orthogonal.

Linear attention has a fatal flaw: it can only add, never update.

Every new token piles onto the state: S = S + K'^T × V'

. Information only grows. It's like a notebook where you can only write — never erase or correct.

DeltaNet (Songlin Yang et al., NeurIPS 2024) fixes this.

The 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.

v_old = k @ S_old           # read current value at key
delta = v_new - v_old       # compute the delta
S_new = S_old + k^T @ delta # write only the difference

If v_new == v_old

, delta is zero — nothing changes. If completely different, delta equals v_new — equivalent to an overwrite. Everything in between is a smooth interpolation.

This is the delta rule: precise memory updates instead of blunt accumulation.

DeltaNet's state update is strictly sequential — each step depends on the previous S. It looks impossible to parallelize. But it is.

The 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.

Complexity: 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.

DeltaNet can precisely edit individual key-value pairs. But what about forgetting at scale?

Imagine 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.

DeltaNet can overwrite specific entries but can't bulk-decay global memory. Mamba-2 (Dao & Gu, ICML 2024) can:

cache = α × S_old + S_new

α 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.

Gated DeltaNet (Songlin Yang et al., ICLR 2025, NVIDIA) merges DeltaNet's delta updates with Mamba-2's gating:

S_new = α × S_old + k^T @ delta   # decay first, then precise write

α=1 is pure DeltaNet. α=0 is a memory wipe. Everything in between both forgets and updates.

The 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.

Kimi Linear (Moonshot AI, arXiv 2510.26692, October 2025) — the predecessor to K3 — refined the idea further: from scalar gating to per-dimension gating.

Gated 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.

Key data from the paper:

Metric Full MLA Kimi Linear
KV Cache 100% -75%
1M-context decode throughput Baseline 6x
Short-context performance Baseline Outperforms
RL scaling Baseline Outperforms

This 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.

I 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.

K3's technical report (arXiv 2607.24653, July 2026) confirms:

K3 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.

Ali's conclusion is sharper than anything I could write:

From GPT-2 to Kimi K3, each generation's core improvement wasn't "more parameters" — it was redesigning

how memory is accessed: how you store, how you forget, how you retrieve.

The evolution in one table:

Stage Representative Memory Mechanism Problem Solved
1 GPT-2 + KV Cache Full cache, O(N) growth Eliminate recomputation
2 Linear Attention Fixed state, O(D²) Cache doesn't grow with N
3 DeltaNet Delta updates, editable Can't update → can update
4 Gated DeltaNet Gating + delta, forgettable Can't forget → can forget
5 KDA / Kimi K3 Per-dim gating + hybrid Linear surpasses full attention

Parameters grew 22,580×. But if that's all you see, you missed the entire story.

If you're doing long-context work (code review, document analysis, multi-turn agents), attention architecture directly impacts your cost and output quality.

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.

Long ≠ expensive. K3's 1M-token context window isn't "brute-forced." 75% of work goes through the linear path.

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.

── more in #artificial-intelligence 4 stories · sorted by recency
── more on @baseten 3 stories trending now
sponsored brought to you by zahid.host 4,200+ EU-deployed projects
reading about agents? ship yours in a single git push.

Run your AI side-project on zahid.host

EU-based hosting, git-push deploys, automatic HTTPS, no cold starts. Free tier with a custom domain — perfect for shipping the agent you just read about.

$git push zahid main
Live at https://your-agent.zahid.host
Get free account → Pricing
from €0/mo · no card required
LIVE [news/how-a-baseten-engine…] indexed:0 read:7min 2026-07-31 ·