Arc 5: Decoding & the Real Inference Algorithm

Why One New Token Means One New Row

8 min readbuilds on Training View vs. Inference View, Self-Attention: Q, K, V from First Principles, Causal Masking and the Autoregressive Constraint and Prefill vs. Decode: The Two Phases of Inference

Inspect OneNewRow

Post 5.1 estimated the arithmetic intensity of batch-one decode. Here we'll trace the matrix operations behind that estimate.

During training or prefill, attention computes the full lower triangle. Each cached decode step computes one new row at the bottom.

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:

Q = X W_Q \in \mathbb{R}^{n \times d_k}, \quad K = X W_K \in \mathbb{R}^{n \times d_k}

The attention score matrix is:

S = Q K^\top \in \mathbb{R}^{n \times n}

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 mask future-position logits with negative infinity before softmax, giving the upper triangle zero probability so no token can attend to future positions. But the computation still produces $n(n+1)/2$ entries in that lower triangle.

Then apply the scale and causal mask: $A=\operatorname{softmax}(S/\sqrt{d_k}+M)$ . Multiply $A$ 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.

Training and prefill use these multi-row matrix products, which give the GPU more opportunities to reuse weights than a single-row decode step.

Now: generate one 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$ keys through the current position, including its own newly appended key. So $K$ is $(L, d_k)$ after the append.

The attention computation becomes:

q_t K^\top = (1, d_k) \times (d_k, L) = (1, L)

That's a matrix product with one output row, often described as a matrix–vector operation. It's a single dot product swept across $L$ cached keys, producing an $L$ -dimensional score vector. After softmax, you multiply by $V$ :

(1, L) \times (L, d_v) = (1, d_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 "Process next token," watch what happens: one new row appears at the bottom of the matrix, highlighted in blue. The previous rows do not need recomputing. The cache stores their keys and values, not the score rows drawn here.

Attention Matrix — Decode Steps

The grid is a history of attention computations. The cache stores K and V activations, not these attention rows. Prefill logits select the first continuation token; processing that token creates the next row and predicts the following token.

previously computed rowsmasked (future)

4 tokens prefilled — click "Process next token"

Each decode step adds exactly one new row to the attention matrix. Watch the dimension panel: Q is always (1, d_k) during decode.

The KV cache retains the keys and values needed for the next row.

The one new row, and what it costs

At short contexts and small batches, reading the dense model's weights can dominate a decode step. Attention still reads $O(L)$ cached keys and values per layer, so its bandwidth cost becomes substantial at long contexts or larger batches.

In batch-one decode, 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 decode step, one weight matrix

(1, d_in) · (d_in, d_out) = (1, d_out)

Inspect ONR WeightSweep

Follow the orange sweep across W: each column contributes one output coordinate. The counter uses one read per FP16 weight and one multiply-add per weight.

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, with a low bandwidth-limited ceiling in the simplified batch-one FP16 model. Large prefill and training matmuls can reach much higher utilization.

What does this look like in code?

This implementation computes the new query’s scores, applies softmax, and sums the cached values.

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)   -- keys through current position (already appended)
    V_cache: (L, d_v)   -- values through current position (already appended)
    """
    # 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

The leading dimension of 1 identifies the single query position. The cached K and V tensors retain the full context length.

What this changes for inference

The following techniques change different costs in this decode step: repeated prefix computation, weight traffic, or the number of useful output rows per weight read.

What changes around the one-row decode step?

Ordinary batch-one decode processes one new query row. These techniques change different costs; they do not all increase arithmetic intensity.

KV caching: Avoid recomputing old positions.
Stores K/V; each new query still reads the relevant history.
Batching: Apply shared weights to more requests in one iteration.
More weight reuse; each request still has its own attention work.
Weight quantization: Move fewer weight bytes.
Metadata, conversion work and supported kernels affect the gain.
Speculative decoding: Verify several proposed positions together.
More work per target pass; acceptance and draft cost determine latency.

Compare which techniques avoid computation and which reuse or reduce weight traffic.

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 and cache-read traffic for avoiding recomputation of earlier positions.

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: FP16-to-INT4 packing reduces raw weight bytes by a factor of four. Scale metadata, packing overhead, and dequantization work affect the realized reduction. In the ideal weights-only model the arithmetic intensity rises from 1 to 4 FLOPs/byte; this is not a guaranteed fourfold speedup.

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. Which traffic dominates depends on context length, batch size, head sharing, and the model.

What's next

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.
Dao, T., Haziza, D., Massa, F., Sizov, G. (2023). Flash-Decoding for long-context inference. Splits keys and values along the sequence dimension, computes partial attention in parallel, and combines the results.
Hong, K., Dai, G., et al. (2023). FlashDecoding++: Faster Large Language Model Inference on GPUs. A separate system studying asynchronous softmax, flat GEMM optimization, and adaptive dataflow.