Why One New Token Means One New Row

The last two posts kept saying "decode is memory-bound" like it was an obvious fact. Post 5.1 has the napkin math and the hardware specs to back it up, but neither post really traced the matrix operations that make it true. The explanation is simple once you stare at the shapes.

The whole thing comes down to a single observation about the attention matrix. During training (or prefill), you compute the full lower triangle. During generation, you add exactly one new row at the bottom. That's the entire decode step, as far as attention is concerned.

Once this picture clicks, it changes how you think about autoregressive generation. So let me build it up piece by piece.

Prerequisites: You should be comfortable with attention as a matrix operation and causal masking. Post 5.1 has the arithmetic intensity and roofline background, and prefill vs. decode sets the stage for why decode is its own phase.

Full-sequence attention

Let's start from what we built in Post 4.1. You have a sequence of $n$ tokens. Each gets projected into a query and a key, both of dimension $d_k$. Stack them up:

The attention score matrix is:

Every row is a query, every column is a key, and every entry $S_{ij}$ is the dot product between token $i$'s query and token $j$'s key. With causal masking (Post 4.3), we zero out the upper triangle so no token can attend to future positions. But the computation still produces $n(n+1)/2$ entries in that lower triangle.

Then you multiply the softmaxed scores by $V \in \mathbb{R}^{n \times d_v}$ to get the output: $(n, n) \times (n, d_v) = (n, d_v)$. Every position gets a contextualized output vector.

This is what happens during training and during the prefill phase of inference. Big, fat matrix multiplies, the kind GPUs are built for, and hardware utilization stays high.

Now: generate one token

So what happens when we shift from processing a full sequence to generating a single new token?

At decode step $t$, you have exactly one new token. It produces a single query vector: $q_t$ of shape $(1, d_k)$. But it needs to attend to all $L$ prior keys (prompt tokens plus previously generated tokens). So $K$ is $(L, d_k)$, sitting in the cache.

The attention computation becomes:

That's one row of scores, not a matrix multiply in any meaningful sense. It's a single dot product swept across $L$ cached keys, producing an $L$-dimensional score vector. After softmax, you multiply by $V$:

You get one output vector.

The picture

Click through the visualization below. It starts with a 4-token prefill (the full lower triangle already computed). Each time you hit "Next token," watch what happens: one new row appears at the bottom of the matrix, highlighted in blue. The previous rows don't change. They're cached.

The side panel makes the key point. During prefill, Q is $(n, d_k)$ and the score matrix is $(n, n)$. During decode, Q is $(1, d_k)$ and the score matrix is $(1, L)$. The shapes collapsed from a matrix to a vector.

And the previous rows? They sit in the KV cache, which is the subject of the next post. You never recompute them. You just grow the matrix by one row at a time.

The one new row, and what it costs

So the attention part of each decode step is cheap. One row, $L$ dot products, done. But here's the thing: attention isn't even the bottleneck. The bottleneck is everything else in the forward pass. The weight-matrix multiplications that happen at every layer.

Here's the picture I want you to carry away from this whole post, and honestly from the whole arc. At every decode step, the model produces one new row of activations with shape $(1, d_\text{model})$. That single row has to pass through every layer: attention projections, attention output, feed-forward up-projection, feed-forward down-projection, output head. Each of those is a matmul of a $(1, d_\text{in})$ row against a $(d_\text{in}, d_\text{out})$ weight matrix, producing another $(1, d_\text{out})$ row. One row in, one row out, and every weight in the matrix has to be read from memory to make it happen.

That's the handle: the one new row. I'll reuse this phrase across the next several posts. Any time you see it, it means: during decode, every layer does a vector-matrix multiply (not a matrix-matrix multiply), and the weight matrix has to be dragged from HBM to produce a single row of output.

This is where the arithmetic intensity number from Post 5.1 comes from. Read every weight once, do one multiply-add per weight, and the ratio of FLOPs to bytes lands at roughly 1. That's what puts decode at the far-left end of the roofline plot, running at under 1% of the GPU's peak compute, while prefill and training sit up against the compute ceiling.

What does this look like in code?

The actual decode-step attention is surprisingly minimal.

import torch
 
def decode_step_attention(q, K_cache, V_cache):
    """
    q:       (1, d_k)   -- the new token's query
    K_cache: (L, d_k)   -- all past keys (from cache)
    V_cache: (L, d_v)   -- all past values (from cache)
    """
    # Score vector: (1, d_k) @ (d_k, L) -> (1, L)
    scores = q @ K_cache.T / (q.shape[-1] ** 0.5)
 
    # Softmax over L keys
    weights = torch.softmax(scores, dim=-1)  # (1, L)
 
    # Weighted sum: (1, L) @ (L, d_v) -> (1, d_v)
    output = weights @ V_cache  # (1, d_v)
 
    return output

Every dimension with a leading 1 is telling you the same thing: this is a vector operation, not a matrix operation. And the weight matrices ($W_Q$, $W_K$, $W_V$, $W_O$, plus the FFN weights) that are potentially billions of parameters? They all get read from memory just to produce that single vector.

Why this single fact is load-bearing

I keep coming back to the one-new-row picture because so many techniques in the inference stack are direct responses to it. Once you see them that way, most of Arc 5 and Arc 6 reduces to a single question: how do you push the ratio of FLOPs to bytes up from 1:1?

The KV cache (Post 5.4): If you didn't cache the keys and values from previous steps, you'd have to recompute every prior token's projections at every step. The cache trades memory capacity for memory bandwidth, storing $O(L)$ key-value pairs so you only do the new token's work.

Speculative decoding (Post 5.6): If each decode step wastes most of the GPU's compute (because it's memory-bound), why not use a small draft model to guess several tokens, then verify the guesses in one batched forward pass of the big model? You spend roughly the same bandwidth but potentially produce multiple accepted tokens.

Continuous batching (Post 5.7): If you're going to read the model's entire weight set anyway, process multiple users' single-token decode steps at once. A batch of 32 users means you do 32 tokens' worth of math for the bandwidth cost of reading the weights once. Arithmetic intensity goes from 1 to 32 FLOPs/byte.

Quantization: Shrink the weights from fp16 to int4 and you cut the data volume by 4x. The FLOP count doesn't change, but the bytes you have to move do, so arithmetic intensity jumps from 1 to 4. You're still memory-bound, but the effective ceiling is 4x higher.

All of these exist because of the same root cause: at decode time, the attention matrix grows by one row, the model reads its entire weight set, and the result is a single vector.

Misconceptions

"The whole attention matrix is recomputed at each step." No. Only the new row is computed. The previous rows are exactly what they were before, and their key-value data sits in the cache. This is why the KV cache works: past attention patterns don't change when you add a new token (under causal masking, no existing token can attend to the new one).

"Decode is slow because attention is O(n^2)." During decode, the attention computation for the new token is $O(L)$, not $O(L^2)$. You compute one row of $L$ dot products. The quadratic cost shows up during prefill (or training), where you compute the full $n \times n$ matrix. The decode bottleneck is weight-reading, not attention.

What's next

I've been leaning on the cache this whole post without ever saying what it is. The next post covers the KV cache from first principles, deriving it by noticing exactly which attention computations you'd be redoing, and then watching the memory it costs grow with every token you generate.

Additional reading (and watching)

• Pope, R., Douglas, S., Chowdhery, A., et al. (2022). Efficiently Scaling Transformer Inference. arXiv:2211.05102.

• Kim, S., Hooper, C., Wattanawong, T., et al. (2023). Full Stack Optimization of Transformer Inference: a Survey. arXiv:2302.14017.

• Dao, T., Fu, D., Ermon, S., Rudra, A., Ré, C. (2022). FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness. NeurIPS 2022. The IO-aware kernel that decode-specific variants (FlashDecoding, FlashDecoding++) build on.

• Ainslie, J., Lee-Thorp, J., et al. (2023). GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints. EMNLP 2023. Introduces grouped-query attention as a knob between MHA and MQA.

• Hong, C., Dao, T., et al. (2023). FlashDecoding++: Faster Large Language Model Inference on GPUs. arXiv:2311.01282. Optimizes the single-query decode kernel by parallelizing across the KV sequence dimension.