Transformer Inference Optimization Strategies

The third post covered the architecture of a Transformer. This one is about running it fast. Generating text is dominated by moving bytes, not by arithmetic, and almost every optimisation in this post is a way of moving fewer bytes or moving them less often: the KV cache, FlashAttention and quantization.

The fourth post covers four topics:

• Why inference efficiency is so critical
• How the KV cache speeds up inference
• FlashAttention's low-level implementation and limitations
• Quantization and other model compression techniques

Every question has a diagram, and every formula that matters is written out with a worked example. The numbers use an A100-class GPU (312 TFLOP/s in FP16, about 2 TB/s of HBM bandwidth) and Llama-2-7B unless said otherwise.

#Why is inference efficiency so critical?

#Where are the main bottlenecks in LLM inference: computation or memory?

LLM inference is mostly memory-bound, not compute-bound. A GPU has far more arithmetic throughput than the memory system can feed.

The Roofline model makes this precise. It describes a kernel by its operational intensity I, the FLOPs it performs per byte it moves from memory:

P(I) = min( P_peak , BW × I )          where  I = FLOPs / bytes moved

ridge point:  I* = P_peak / BW = 312 TFLOP/s ÷ 2 TB/s ≈ 156 FLOPs per byte

Below the ridge point, speed is set by bandwidth (BW × I). Above it, speed is capped by the arithmetic peak.

Why decoding sits so far to the left:

• To produce one token, every weight is read from HBM once. An FP16 weight is 2 bytes, and it takes part in one multiply-add, which is 2 FLOPs. So I ≈ 2 FLOPs ÷ 2 bytes = 1.
• If B sequences are decoded together, each weight is still read once but used B times, so I ≈ B. At batch 1 the GPU runs at about 0.6% of its peak.
• Prefill processes the whole prompt as one big matrix multiplication, so its intensity is in the hundreds and it is compute-bound.

Here is the same fact in time units, for one decode step of a 7B FP16 model:

memory time  = 14 GB ÷ 2 TB/s            ≈ 7 ms
compute time = 14 GFLOP ÷ 312 TFLOP/s   ≈ 0.045 ms        (2 FLOPs per parameter)

This is why the optimisations in this post target bytes, not FLOPs:

#How is the latency per token related to sequence length?

Generation has two phases with different costs:

• Prefill. The prompt of Lp tokens is processed in parallel. Attention here is quadratic in Lp, but it happens once per request.
• Decode. Tokens come out one at a time. Each step is paid again for every generated token, so it dominates the total time.

With a KV cache, each decode step needs two things from HBM: the weights (a fixed cost) and the KV cache of every earlier token (a cost that grows with the context). That gives the idealised model:

t_token ≈ ( W_bytes + B × T × c_kv ) ÷ BW

W_bytes = weight bytes (14 GB for 7B in FP16)       B = sequences in the batch
T       = tokens already in the cache               c_kv = KV bytes per token (0.5 MiB for Llama-2-7B)

So latency per token is linear in the context length, on top of a constant floor set by the weights:

• For short contexts, the weights dominate and latency is almost flat.
• At batch 1 the two costs meet at T ≈ 14 GB ÷ 0.5 MiB ≈ 27,000 tokens. With a larger batch B, the KV term grows B times and takes over much earlier.
• The attention arithmetic itself is also linear per step, O(T · d), but it is cheap next to the bytes it has to read.

#Why does inference memory usage grow in proportion to prompt length?

Every token the model has seen leaves a key and a value vector in every layer, and those vectors must stay in memory to be attended to later. Prompt tokens add their entries during prefill, and every generated token adds one more.

• KV cache memory grows linearly with prompt plus generated length, and also linearly with the batch size.
• The weights are a fixed cost. The KV cache is the part that changes with traffic.
• On long-context workloads (document summarisation, long chats) the KV cache, not the weights, is usually what runs the GPU out of VRAM.

#How the KV cache speeds up inference

#What are the dimensions of the Key/Value cache? How is it stored?

Dimensions. For each layer, the cache holds one K tensor and one V tensor:

[batch, num_kv_heads, seq_len, head_dim]

• num_kv_heads equals num_heads in standard multi-head attention (MHA). It is smaller in grouped-query attention (GQA), and 1 in multi-query attention (MQA).
• head_dim = d_model / num_heads.
• seq_len grows by one with each generated token. Prefill writes the whole prompt at once.
• Some implementations put seq_len before num_kv_heads. The layout varies, but the content is the same.

The attention variant decides how many KV heads there are, and so how big the cache is:

How it is stored. There are three common strategies:

Other ways to shrink the cache:

• Lower precision. Store K and V in FP8 or INT8 (even INT4) instead of FP16.
• Fewer or smaller entries. GQA and MQA (above), sliding-window attention, or evicting old tokens.
• MLA (multi-head latent attention, used by DeepSeek) caches a small compressed latent vector per token instead of full keys and values.

#How does attention for a prefix prompt avoid redundant computation?

Without a cache, producing the next token means running the model over the whole sequence again, because the new token's query must be compared with the keys and values of every earlier token. Those keys and values did not change since the last step, so recomputing them is pure waste.

The KV cache removes the waste:

1. Prefill. The prompt goes through the model once. In every layer the keys and values of every prompt token are computed and stored.
2. Decode. Each step feeds in only the newest token. Its K and V are computed, appended to the cache, and its query attends to the whole cache.

In the attention score matrix this means that only one row is new at each step:

The cost per decode step, for a context of T tokens and model width d:

The cache is a trade: it buys a large amount of compute for 2 × L × H_kv × T × d × bytes of memory (next question). It also changes which side of the roofline decode lives on, because each step now reads a long cache instead of recomputing it.

# Hugging Face: the cache is explicit
out = model(input_ids, use_cache=True)             # prefill: processes the prompt, fills the cache
past = out.past_key_values                          # K and V of every layer
next_id = out.logits[:, -1].argmax(-1, keepdim=True)

out = model(next_id, past_key_values=past, use_cache=True)  # decode: one token in, cache grows by one

The same idea works across requests. If two requests start with the same system prompt, the keys and values of that prefix are identical, so they can be computed once and shared. This is prefix caching:

• Blocks are identified by a hash of their tokens and everything before them, so equal prefixes find equal blocks.
• A request that hits the cache skips the prefill work for those tokens, which lowers the time to first token.

#How large can the cache get? Can it cause out-of-memory? Does it need to be manually cleared?

How large. The size is a product of seven factors:

KV bytes = 2 × L × B × H_kv × T × d × b

2    one K and one V            L    number of layers          B   batch size
H_kv number of KV heads         T    tokens in the cache       d   head dimension
b    bytes per element (2 for FP16/BF16, 1 for FP8/INT8)

For full multi-head attention H_kv × d = d_model, so the formula simplifies to 2 × L × B × T × d_model × b. For GQA or MQA, use the real number of KV heads, not the number of query heads.

def kv_cache_bytes(n_layers, batch, seq_len, n_kv_heads, head_dim, bytes_per_elem=2):
    return 2 * n_layers * batch * seq_len * n_kv_heads * head_dim * bytes_per_elem

kv_cache_bytes(32, 1, 4096, 32, 128, 2)   # Llama-2-7B, FP16: 2,147,483,648 bytes = 2 GiB
kv_cache_bytes(32, 1, 4096, 8, 128, 2)    # Llama-3-8B, 8 KV heads: 536,870,912 bytes = 512 MiB

Per token that is 512 KiB for Llama-2-7B and 128 KiB for Llama-3-8B. Both grow linearly with the batch size and the length, so a batch of 32 sequences of 2048 tokens needs 32 × 2048 × 512 KiB = 32 GiB for Llama-2-7B, on top of the 13 GiB of weights.

Can it cause an out-of-memory error? Yes, and it is the usual cause in serving. Each new token appends K and V to every layer, so the cache grows steadily, and with many concurrent long requests it can exhaust VRAM. That caps the maximum context length and the maximum batch size, and with them the throughput.

Does it need to be cleared manually? No. Inference engines manage it for you. When a request ends (an [EOS] token, the maximum length, or a cancellation), the blocks it held go back to a pool and are reused by other requests.

• With PagedAttention, blocks are small and need not be contiguous, so the cache is allocated on demand instead of reserving the maximum length up front.
• If the pool runs out, the scheduler stops admitting new requests and, if needed, preempts a running one, which frees its blocks and recomputes it later. Post 05 follows this pool in detail.
• In your own code with a Hugging Face model, the cache lives inside the past_key_values object, so it is freed when you drop the reference (and call torch.cuda.empty_cache() only if you need the memory back for another process).

#FlashAttention's low-level implementation and limitations

#How does FlashAttention fuse memory access with computation?

Standard attention runs as three separate steps. Each one writes its result to HBM, and the next one reads it back:

S = Q · Kᵀ / √d        # N × N scores        -> written to HBM
P = softmax(S)         # N × N probabilities -> read from and written to HBM
O = P · V              # output              -> read P from HBM

The N × N matrices S and P are the expensive part: they are large, and they travel across the slow memory twice each. FlashAttention computes the exact same result but never builds them in HBM.

It relies on two ideas:

• Tiling. Split Q, K and V into blocks sized to fit in the GPU's fast on-chip SRAM. For one block of queries Q_i, sweep across the blocks K_j, V_j, doing all the work on the tile that is currently in SRAM.
• Kernel fusion. The matmul, the scaling, the softmax and the second matmul are one GPU kernel, so intermediate values stay in SRAM and registers. This also removes the launch overhead of several kernels.

Why this matters is the memory hierarchy: SRAM is roughly an order of magnitude faster than HBM, but tiny.

The hard part: softmax across tiles. A softmax needs the maximum and the sum of a whole row, but a tile only sees part of the row. FlashAttention keeps running statistics per row, the running max m and the running sum l, and rescales its partial results whenever the max changes (online softmax):

m_new = max( m_old , rowmax(S_ij) )
l_new = e^(m_old − m_new) · l_old + rowsum( e^(S_ij − m_new) )
O_new = e^(m_old − m_new) · O_old + e^(S_ij − m_new) · V_j
final:  O = O_new / l_new

A worked example, with the row of scores [1, 3, 2, 4] processed in two blocks:

• Block 1 gives m = 3 and l = e^(1−3) + e^(3−3) = 1.135.
• Block 2 raises the max to 4. The old sum shrinks by e^(3−4) = 0.368, and l = 0.368 × 1.135 + e^(2−4) + e^(4−4) = 1.553.
• A single pass over all four scores gives l = e^(−3) + e^(−1) + e^(−2) + 1 = 1.553. Identical.

In PyTorch you rarely call it directly. scaled_dot_product_attention picks a fused backend when the hardware and dtype allow it:

import torch.nn.functional as F
from torch.nn.attention import sdpa_kernel, SDPBackend

# q, k, v: (batch, heads, seq_len, head_dim) in FP16 or BF16 on a CUDA GPU
with sdpa_kernel(SDPBackend.FLASH_ATTENTION):          # force the FlashAttention backend
    out = F.scaled_dot_product_attention(q, k, v, is_causal=True)

#What are the differences between FlashAttention v1 and v2? Which devices are supported?

Both versions are exact and IO-aware. v2 reorganises the work so the GPU is kept busier:

Supported devices. The official kernels target NVIDIA GPUs with Tensor Cores:

Support changes quickly, so check the project README for the current list before depending on it.

#Why is FlashAttention especially effective for long sequences?

The N × N matrix is the problem. It grows with the square of the length, in memory and in HBM traffic:

standard attention:  HBM accesses  Θ( N·d + N² )        extra memory  O(N²)
FlashAttention:      HBM accesses  Θ( N²·d² / M )       extra memory  O(N)       (M = SRAM size)

For d much smaller than √M (typical values: d = 64–128, M about 100 KB), the second line is several times smaller. The O(N) extra memory is one number per row (the log-sum-exp), not a matrix.

The size of the matrix being avoided, per head and per sequence in FP16:

• At 32k tokens a single head's score matrix is 2 GiB, and a model has dozens of heads per layer. Standard attention simply cannot hold it.
• Less HBM traffic means the attention step is faster too, and the GPU spends more time on arithmetic and less waiting on memory.
• This is what makes long-context training and inference practical on a single GPU.

#Can it be used for inference, or only for training?

Both. The two cases use the same idea in a different way:

Decode still reads the whole KV cache from HBM, so it stays memory-bound. FlashAttention makes that read efficient and parallel, but it does not remove it.

#Quantization and other model compression techniques

Quantization stores numbers with fewer bits. The weights of a 7B-parameter model take 28 GB in FP32, 14 GB in FP16, 7 GB in INT8 and 3.5 GB in INT4 (the chart in post 02). Because decoding is memory-bound, fewer bytes per weight means faster tokens as well as smaller models.

#What are the principles of INT8 / FP16 / GPTQ? Which ones impact precision?

First, how the formats spend their bits:

FP16 and BF16 (16-bit floats). A float keeps a sign, an exponent (range) and a mantissa (precision):

FP16:  value = (−1)^s × 2^(e − 15) × (1 + m / 1024)     5 exponent bits, 10 mantissa bits
       largest value 65,504     smallest normal value 2^−14 ≈ 6.1 × 10⁻⁵     about 3–4 decimal digits
BF16:  same 8 exponent bits as FP32 (range up to 3.4 × 10³⁸), only 7 mantissa bits (about 2–3 digits)

• Floats are spaced logarithmically: small numbers are represented very finely and big ones coarsely. That suits weights and activations, which span a wide range.
• Neural networks tolerate the lost digits well. FP16 can overflow because of its small range, which is why BF16 is often preferred.
• FP8 goes further on Hopper and Ada GPUs. E4M3 (4 exponent bits, 3 mantissa bits, a maximum of 448) is used for weights and activations, and E5M2 (more range, up to 57,344) mostly for gradients.

INT8 (8-bit integers). An integer grid is spaced evenly, so we need a scale to map real values onto it. There are two schemes:

Symmetric:    s = α / 127                     q = clamp( round(x / s), −127, 127 )      x̂ = s · q
Asymmetric:   s = (max − min) / 255           z = round(−min / s)
              q = clamp( round(x / s) + z, 0, 255 )                                      x̂ = s · (q − z)

rounding error per value ≤ s / 2                 typical error (uniform) = s / √12

Example with α = 1.27: s = 0.01, and x = 0.347 becomes q = round(34.7) = 35, which dequantizes to x̂ = 0.35. The error, 0.003, is below s / 2 = 0.005.

def quantize_int8(x):
    s = x.abs().max() / 127                                  # one scale for the whole tensor
    q = torch.clamp(torch.round(x / s), -127, 127).to(torch.int8)
    return q, s                                              # dequantize with q.float() * s

• The scale can be per tensor, per output channel, or per group of weights (say 128 values). Finer scales cost a little memory but track the range better.
• Outliers are the real enemy. One huge value inflates α, which inflates s, and every other value loses resolution. Activations have more outliers than weights, so quantizing both (W8A8) is harder than quantizing only the weights. Two well-known fixes:
  • SmoothQuant divides each activation channel by a factor and multiplies the matching weights by the same factor. The product is unchanged, but the outliers move from the activations into the weights, which handle them better.
  • AWQ (activation-aware weight quantization) finds the small share of weight channels that matter most, judged by the size of their activations, and scales them up before quantizing so that they keep more precision.

GPTQ (post-training weight quantization). GPTQ (Generative Pre-trained Transformer Quantization) compresses the weights of a finished model to 4 or 3 bits without retraining. Rounding each weight independently adds error. GPTQ rounds the weights of a layer one column at a time and pushes each column's rounding error onto the columns that are still in full precision, so the layer's output stays close to the original:

per layer:   minimise  ‖ W·X − Ŵ·X ‖²      with Ŵ on a 4-bit grid       X = calibration activations
Hessian:     H = 2 · X · Xᵀ
for each column q:   δ = ( w_q − quant(w_q) ) / [H⁻¹]_qq
                     w_rest ← w_rest − δ · [H⁻¹]_(q, rest)

• It needs only a small unlabelled calibration set (a few hundred samples) to measure how much each weight matters.
• The result is weight-only quantization: the weights are stored in 4 bits, then dequantized on the fly for the matrix multiply.

Which ones impact precision? All three lose some, in different ways:

The loss depends more on the number of bits and the method than on the format family. Fewer bits and cruder scales lose more.

#If quantized models can be used for inference, can they also be used for training?

Yes, with a special technique: Quantization-Aware Training (QAT). It simulates quantization during the forward pass so the model learns weights that survive it.

PauseData

• Fake quantization. In the forward pass the weights are quantized and immediately dequantized back to full precision: ŵ = s · clamp(round(w / s), q_min, q_max). The layer computes with ŵ, so it sees the rounding noise.
• The gradient problem. round() is flat almost everywhere, so its true derivative is 0 and no gradient would reach w.
• Straight-through estimator (STE). In the backward pass we pretend rounding is the identity, inside the clip range:

forward:    ŵ = s · clamp( round(w / s), q_min, q_max )
backward:   ∂ŵ/∂w ≈ 1    if  q_min ≤ w/s ≤ q_max          (0 outside the range)

class FakeQuant(torch.autograd.Function):
    @staticmethod
    def forward(ctx, w, s):
        return torch.clamp(torch.round(w / s), -127, 127) * s   # quantize, then dequantize

    @staticmethod
    def backward(ctx, grad):
        return grad, None     # straight-through: the gradient passes unchanged

• Master weights stay in FP32. Updates are tiny, and low-bit values would swallow them. Only the forward pass sees the quantized values.

How this applies to each format:

For fine-tuning large models there is a practical shortcut, QLoRA: keep the base model frozen in 4-bit, and train only small low-rank adapters in BF16. The quantized weights are never updated, so there is no gradient problem.

#How are quantized weights executed on the GPU? Are there dedicated kernels?

Yes. Modern GPUs have Tensor Cores, units that multiply small matrix tiles in hardware. For a tile they compute:

D = A × B + C        A, B: low-precision inputs      C, D: higher-precision accumulators
INT8:   A, B in INT8,  C, D in INT32           FP16/BF16:   A, B in FP16 or BF16,  C, D in FP32

There are two common execution paths:

W8A8:          Y ≈ s_x · s_w · ( X_q · W_qᵀ )      integer matmul, INT32 accumulation
weight-only:   Y ≈ X · ( s · (W_q − z) )ᵀ          dequantize W in registers, then an FP16 matmul

• Why INT32 accumulation. Each product of two INT8 values is at most 127 × 127 = 16,129. A sum over K = 4096 terms is at most about 6.6 × 10⁷, far below INT32's limit of about 2.1 × 10⁹. The result cannot overflow.
• Why weight-only is still faster. Decoding is memory-bound, so reading 4-bit weights moves 4× fewer bytes than FP16. The dequantize step is fused into the kernel, so the FP16 weights never touch HBM.
• Compute-bound phases (prefill) gain from the faster INT8 Tensor Core throughput itself. Memory-bound phases gain from the smaller footprint.

The software stack that exposes this:

#How are operations with large dynamic range (e.g., LayerNorm, softmax) handled?

They are kept in higher precision. Matrix multiplications hold nearly all the FLOPs and tolerate low precision. LayerNorm and softmax hold almost none, but they contain sums of squares and exponentials, which are very sensitive to rounding and overflow.

LayerNorm normalises each token's feature vector of size d:

μ  = (1/d) Σ x_i
σ² = (1/d) Σ (x_i − μ)²
y_i = γ_i · (x_i − μ) / √(σ² + ε) + β_i

• The sum of squares overflows quickly in low precision. In FP16, x² > 65,504 once |x| > 256, and LLM activations can have outliers that large.
• Computing the variance as E[x²] − μ² can lose all its digits to cancellation.
• So the mean and variance are computed in FP32, even when the input and output are INT8 or FP16.

Softmax turns scores into probabilities:

softmax(z)_j = e^(z_j) / Σ_k e^(z_k)  =  e^(z_j − z_max) / Σ_k e^(z_k − z_max)

The exponential grows so fast that it overflows low-precision formats at modest inputs. In FP16, e^12 = 162,755 is already larger than 65,504. The fix is the second form: subtract the maximum first, so the largest exponent is 0 and everything is at most 1.

PauseData

def softmax(z):
    z = z - z.max()              # the largest exponent becomes 0, so exp() never overflows
    e = torch.exp(z)
    return e / e.sum()           # in practice, accumulate the sum in FP32

The strategies, in order of how often they are used:

#Summary

Decoding is memory-bound. Each step streams the weights and the KV cache from HBM, so the time per token is roughly the bytes moved divided by the bandwidth, and every optimisation here moves fewer bytes or reuses them more. The KV cache stores every earlier token's keys and values (2 × L × B × H_kv × T × d × b bytes) so that each step computes only the new token. It grows linearly with the context and the batch, and GQA, MQA, paging and lower precision shrink it. FlashAttention computes exact attention in tiles held in SRAM, with an online softmax, and never materialises the N × N matrix, which helps training, prefill and, with Flash-Decoding, decode. Quantization trades bits for memory and speed: floats spend their bits on range and precision, integers need a scale, GPTQ compensates for rounding error, and QAT trains through fake quantization with the straight-through estimator. On the GPU, low-precision Tensor Core matmuls accumulate in higher precision, while LayerNorm and softmax stay in FP32.

#Concepts at a glance

#Test yourself

Next in the series: 05 · vLLM Framework and Inference Engine Design.