The MLA decode speedup hiding in your model card

Bhabani's Substack

SubscribeSign in

The MLA decode speedup hiding in your model card<br>Naive MLA is an order of magnitude slower than MHA at long context. The fix is one matrix identity and almost nobody writes it the first time.

Bhabani Nayak<br>Aug 11, 2026

Share

I had a benchmarking puzzle.<br>I’d implemented Multi-Head Latent Attention from the obvious read of the description: cache a low-rank latent, reconstruct full keys and values at decode time, run attention. Same architecture, smaller cache. Cache savings should buy faster decode. Right?<br>Wrong. At 32K context with batch=16, my MLA was decoding at 5,000 tokens per second . My MHA at the same setup was doing 42,000 .<br>MLA had a smaller cache and was 8× slower .<br>That ratio bothered me for two days. The architecture is supposed to be a Pareto improvement that’s the whole pitch. The numbers said the opposite. Either the paper was wrong about something fundamental, or I’d written the wrong implementation. The latter is overwhelmingly more likely when you’re debugging your own code, so I went looking.<br>I found a math identity that, applied at model-load time, makes MLA decode at 71,000 tokens per second with the same cache. Thirteen times faster. Same numerics, same outputs. One change.<br>This is that identity. And the reason most first MLA implementations leave a ~30× speedup on the table.<br>(Part two of the KV-cache economics series. Part one, nine attention variants on one Pareto plot, found MLA quality-dominated by GQA at 30M scale. This post is about the other axis: why MLA’s latency reputation is an implementation artifact.)

Nine attention variants, one Pareto plot<br>Bhabani Nayak<br>Jul 28

Read full story

The naive form, written out

Here’s MLA’s decode step the way you write it the first time:<br>def naive_mla_decode_step (q, k_lat, v_lat, W_k_up, W_v_up):<br>"""One decode step. q: (B,H,1,d). k_lat,v_lat: (B,H,T,r). W_*_up: (H,r,d)."""<br># 1. Up-project the cached latents to full key and value.<br>K = torch.einsum("bhtr,hrd->bhtd", k_lat, W_k_up) # (B, H, T, d)<br>V = torch.einsum("bhtr,hrd->bhtd", v_lat, W_v_up) # (B, H, T, d)<br># 2. Standard attention.<br>scores = q @ K.transpose(-2, -1) / sqrt(d) # (B, H, 1, T)<br>probs = scores.softmax(-1)<br>return probs @ V # (B, H, 1, d)Two big matmuls before attention even starts. Each one runs over the entire cache , every cached token, every head. The dominant FLOP cost is O(B · H · T · r · d_head), proportional to T · r · d_head.<br>That T is the killer. For every new decode step, you’re redoing work over the full cached past. Doubling the context doubles the per-step cost. This is why naive MLA’s latency curve goes from 0.15 ms at 1K context to 3 ms at 32K, a 20× increase for a 32× context increase. The cost is in the up-projection, not the attention.<br>The identity

Standard scaled-dot-product attention computes, for each cached position t:<br>score[t] = q · K[t]ᵀIn MLA, K[t] = K_lat[t] · W_k_up. Substituting:<br>score[t] = q · (K_lat[t] · W_k_up)ᵀ<br>= q · W_k_upᵀ · K_lat[t]ᵀ<br>= (q · W_k_upᵀ) · K_lat[t]ᵀThat last regrouping is the whole insight. The product q · W_k_upᵀ doesn’t depend on t. Compute it once per decode step , outside the loop over cached tokens and then run attention against the cached latent directly.<br>Define:<br>q' = q · W_k_upᵀ # shape: (B, H, 1, r)And attention becomes:<br>score[t] = q' · K_lat[t]ᵀq' is small. K_lat is what we already had cached. The full key reconstruction never happens.<br>For V it’s almost the same trick, applied after attention:<br>out_lat = probs · V_lat # in latent space, shape: (B, H, 1, r)<br>out = out_lat · W_v_up # project back to head dim, shape: (B, H, 1, d_head)You compute attention’s weighted sum in latent space , then project once at the end. V is also never reconstructed.<br>The absorbed form

Here’s the decode step rewritten:<br>def absorbed_mla_decode_step(q, k_lat, v_lat, W_k_up, W_v_up):<br>"""One decode step. Same I/O as naive — different math, dramatically faster."""<br># 1. Absorb W_k_up into Q. O(H * r * d_head). NOT a function of T.<br>q_abs = torch.einsum("bhtd,hrd->bhtr", q, W_k_up) # (B, H, 1, r)<br># 2. Attention runs against the latent K directly.<br>scores = q_abs @ k_lat.transpose(-2, -1) / sqrt(d) # (B, H, 1, T)<br>probs = scores.softmax(-1)<br># 3. Aggregate in latent space, then project out.<br>out_lat = probs @ v_lat # (B, H, 1, r)<br>return torch.einsum("bhtr,hrd->bhtd", out_lat, W_v_up)What changed:

The hot path lost a factor of d_head. With d_head = 128 (production scale), that’s the difference between “barely usable” and “competitive.”<br>What the numbers actually look like

I ran both forms on an H100 across a context-length sweep, batch=16, in bf16.

Naive vs absorbed MLA decode throughput, with MHA as reference. Left: my model's dimensions. Right: Llama-7B-class dimensions. Two things to notice: the absorbed line sits 13–29× above the naive line at long context, and at model scale it crosses above MHA past ~8k. Once you're bandwidth-bound, the smaller latent cache reads...