Inference-aware model design

Choices you make before training (how many key-value heads, what precision the cache uses, whether a draft model can ride along) decide how much memory each token costs and how fast the model can answer. This stage is about doing that arithmetic first.

Source: Suraj Sharma's roadmap post, stage 8 · saved in sources/08-suraj-ai-research-engineer.md

Going deeper: everything below is the full version, with some maths. On a first read, skim the headings and the "Why it matters for robotics" box, and come back for the rest later.

1. What this stage covers

• Learn: GQA, MQA, KV-cache quantisation, and what makes a model work well with speculative decoding.
• Practice: convert a standard transformer to GQA, measure the KV-cache memory saved, and benchmark latency and throughput in vLLM.
• Why (the post's claim): a model too slow or expensive to serve is useless, and architecture sets the inference costs. Largely true; see section 4.

2. Concepts, explained

First, the KV-cache arithmetic

In a decoder, each layer computes keys and values for every token. Generating token t+1 needs the keys and values of tokens 1…t, so servers keep them instead of recomputing. The size is exact and worth memorising:

KV bytes = 2 × L × n_kv × d_head × T × B × bytes_per_value
           (K and V) (layers) (KV heads) (head dim) (tokens) (batch)

Llama-2-7B shape (L=32, n_kv=32, d_head=128), fp16:
  per token = 2 × 32 × 32 × 128 × 2 B = 512 KiB
  4,096 tokens = 2 GiB per sequence

Llama-3-8B shape (L=32, n_kv=8, d_head=128), fp16:
  per token = 128 KiB  →  8,192 tokens = 1 GiB per sequence

Why it matters for speed: decoding one token at batch 1 is memory-bandwidth bound. The GPU must stream all weights plus the cache from memory for very little arithmetic. A rough floor is (weight bytes + KV bytes) ÷ bandwidth. For a 7B fp16 model (about 14 GB of weights) on a GPU with 1 TB/s, that's about 14 ms per token before the cache. At short contexts the weights dominate. At long contexts or large batches the cache overtakes them, because weights are read once per step for the whole batch while each sequence brings its own cache. So shrinking the cache buys bigger batches (throughput) and longer contexts, and only buys per-token latency once the cache is a large share of the bytes read.

Multi-query attention (MQA)

Noam Shazeer's 2019 paper "Fast Transformer Decoding: One Write-Head is All You Need" kept many query heads but shared a single key head and value head across all of them. The cache shrinks by a factor of n_heads (32× for a 32-head model). The paper reports much faster decoding and only minor quality loss. In practice MQA can hurt quality and training stability at scale, and with tensor parallelism the single KV head has to be duplicated on every GPU, which gives back some of the saving.

Grouped-query attention (GQA)

GQA (Ainslie, Lee-Thorp, de Jong, Zemlyanskiy, Lebrón and Sanghai, 2023) is the middle ground: split the query heads into G groups, and give each group one shared KV head. G = n_heads is ordinary multi-head attention (MHA); G = 1 is MQA.

MHA:  n_q = 32, n_kv = 32   cache × 1
GQA:  n_q = 32, n_kv = 8    cache × 1/4    (each KV head serves 4 query heads)
MQA:  n_q = 32, n_kv = 1    cache × 1/32

attention for query head h uses K, V of KV head  ⌊h / (n_q / n_kv)⌋

The paper's second contribution is uptraining: you don't need to retrain from scratch. Mean-pool the key and value projection weights of the heads in each group into one head, then continue pre-training for about 5% of the original compute. Uptrained GQA came close to MHA quality at a speed close to MQA. Eight KV heads has become a common choice: Llama 2 70B, Llama 3 and Mistral 7B all use it.

A further step is to compress the cache rather than share it. DeepSeek's Multi-head Latent Attention (MLA) caches one small latent vector per token and reconstructs keys and values from it. Module 05 explains MLA and how Kimi K3 mixes it with linear-attention layers whose state doesn't grow at all; see the length section there.

KV-cache quantisation

Weights are fixed and can be quantised offline with care. The cache is created on the fly, so its quantisation must be cheap and must cope with whatever activations appear. Two findings shape the methods:

• Keys have outlier channels. A few channels of the key vectors carry much larger values than the rest, consistently across tokens. Quantising keys per channel isolates them. Values have no such pattern and are quantised per token. This is the core finding of KIVI (Liu et al., 2024), a tuning-free 2-bit scheme that also keeps the most recent tokens in full precision until a group fills.
• Rotary embeddings mix channels. KVQuant (Hooper et al., 2024) quantises keys before RoPE is applied, uses non-uniform per-layer data types, and stores outliers separately in a sparse format. It reports under 0.1 perplexity loss at 3 bits on Llama-family models.

q = round((x − min) / scale),  scale = (max − min) / (2^bits − 1)
keys:   min/max/scale per channel (over a group of tokens)
values: min/max/scale per token   (over its channels)
memory ≈ bits/16 of fp16, plus a small overhead for the scales

In production the common, safe step is FP8 KV cache, which vLLM supports through --kv-cache-dtype. It halves the cache with usually small quality impact. Going to 4 or 2 bits needs the tricks above and careful evaluation, especially on long-context retrieval tasks, where small errors in old keys can make the model attend to the wrong place.

Speculative-decoding compatibility

Speculative decoding (Leviathan, Kalman and Matias, 2022) lets a cheap draft propose several tokens and the large model verify them in one forward pass, with an acceptance rule that leaves the output distribution exactly unchanged. Module 05 covers the mechanism and how K3 trained its draft for acceptance rate (serving section). The design questions for this stage are what make a model easy to speculate with:

• Same tokenizer. A separate draft model must share the target's vocabulary. Plan a small sibling model at training time.
• Built-in drafts. Extra prediction heads (Medusa, EAGLE) or multi-token prediction trained alongside the main model give a draft for free, but only if they're designed in.
• Cheap verification. Verifying k tokens is a short prefill that reads the whole cache. GQA and KV quantisation make that cheaper too.
• Memory for two caches. The draft has its own KV cache, which counts against the same budget.
• Predictable outputs. Acceptance is higher on predictable text (code, structured output) and lower at high sampling temperature, so the gain depends on the workload.

3. The practice project

(a) Laptop/Colab version: MHA vs GQA vs MQA on a tiny model

Train three tiny character-level GPTs that differ only in KV heads, then measure quality, cache size and decode speed. This runs on a laptop CPU in a few hours or a Colab GPU in under one.

1. Start from nanoGPT (or your Stage 01 transformer) on Tiny Shakespeare: 6 layers, 8 query heads, embedding 256, context 256.
2. Swap the attention for the GQA module below. Train three models with n_kv = 8, 2 and 1 for the same number of steps and seed.
3. Measure quality: validation loss for each. Expect GQA-2 to be close to MHA and MQA a little worse, but check rather than assume; at this tiny scale the differences may be within seed noise, so run two seeds.
4. Measure memory: compute KV bytes with the formula, then confirm by summing the sizes of the cached tensors during generation. They should agree exactly.
5. Measure speed: tokens per second when generating 2,000 tokens at batch 1 and batch 32. Expect little gain at batch 1 on a tiny model, and a clearer gain at larger batch and context.
6. Uptraining exercise: take the trained MHA model, mean-pool its KV heads into 2 groups, measure the loss jump, then fine-tune for 5% of the original steps and measure again.

import torch, torch.nn as nn, torch.nn.functional as F

class GQA(nn.Module):
    def __init__(s, d, n_q, n_kv):
        super().__init__()
        s.n_q, s.n_kv, s.hd = n_q, n_kv, d // n_q
        s.q = nn.Linear(d, n_q * s.hd, bias=False)
        s.kv = nn.Linear(d, 2 * n_kv * s.hd, bias=False)
        s.o = nn.Linear(d, d, bias=False)

    def forward(s, x, cache=None):                     # x: [B, T, d]
        B, T, _ = x.shape
        q = s.q(x).view(B, T, s.n_q, s.hd).transpose(1, 2)
        k, v = s.kv(x).view(B, T, 2, s.n_kv, s.hd).permute(2, 0, 3, 1, 4)
        if cache is not None:                          # [B, n_kv, T_past, hd]
            k, v = torch.cat([cache[0], k], 2), torch.cat([cache[1], v], 2)
        new_cache = (k, v)                             # the part GQA shrinks
        r = s.n_q // s.n_kv
        k, v = k.repeat_interleave(r, 1), v.repeat_interleave(r, 1)
        y = F.scaled_dot_product_attention(q, k, v, is_causal=cache is None)
        return s.o(y.transpose(1, 2).reshape(B, T, -1)), new_cache

def pool_heads(W, n_q, n_kv):                          # W: [n_q*hd, d] → [n_kv*hd, d]
    hd = W.shape[0] // n_q
    return W.view(n_kv, n_q // n_kv, hd, -1).mean(1).reshape(n_kv * hd, -1)

Note that repeat_interleave copies data at compute time; only the cache is smaller. Real kernels (and recent PyTorch's enable_gqa option) index the shared heads without copying.

(b) The full version

The post asks for a GQA conversion benchmarked in vLLM. A realistic plan: take an open MHA checkpoint of about 1B parameters, convert it to 4 or 8 KV heads by mean-pooling, uptrain on a few billion tokens of similar data, and serve both versions in vLLM on one 24–80 GB GPU. Uptraining a 1B model on a few billion tokens is a multi-GPU-day job, so budget roughly tens to low hundreds of dollars of rented compute. Then:

• Record the KV-cache capacity vLLM reports at startup for each model. With the same memory setting, the GQA model should hold n_q/n_kv times more tokens.
• Benchmark with vLLM's benchmarking tools across concurrency levels and prompt lengths. Report time to first token, inter-token latency and output tokens per second.
• Repeat with --kv-cache-dtype fp8 on both.
• Evaluate quality on the same tasks before and after, including at least one long-context retrieval test.

(c) Pitfalls

• Measuring VRAM with nvidia-smi. vLLM pre-allocates a fixed share of GPU memory, so total usage looks the same for both models. Compare KV-cache token capacity instead.
• Expecting latency wins at batch 1, short prompts. There the weights dominate the bytes read. GQA's gains show up as larger batches, longer contexts and higher throughput.
• "The exact drop". The cache saving is exactly n_q/n_kv by construction. The interesting measurements are the quality cost and the end-to-end throughput change.
• Skipping uptraining. Mean-pooled heads without further training lose noticeable quality. The GQA paper's result depends on uptraining.
• Tensor parallelism. With more GPUs than KV heads, KV heads are replicated, and the saving per GPU shrinks.

4. Reality check

The post says a model too slow or expensive to serve is useless, and that architecture dictates inference physics.

• True: serving cost is a design constraint, not an afterthought. GQA became near-universal in open models because it changes how many users fit on a GPU, and Module 05 shows Kimi K3's attention design was shaped by serving.
• Overstated: architecture is one lever among several. Quantisation of weights, batching, paged memory, prefix caching, speculative decoding and the kernel stack often matter as much, and they can be added after training. "Useless" is too strong: slow, accurate models are routinely distilled into fast ones (Module 04).
• Missing: the cache is not the bottleneck for every workload. Short chats at small batch are bound by weights; that's why weight quantisation and smaller models come first for edge and robot deployment.

5. In this course

• Module 05 · Serving: MLA, the paged KV pool shared by two cache types, and speculative decoding with a draft trained for acceptance rate. Read this before or after this page; it's the production version of these ideas.
• Module 05 · Length: linear-attention layers whose state doesn't grow with context.
• Module 04 · Latency is a safety number: how milliseconds turn into metres on a moving robot.
• Module 02 · TensorTonic: KV cache, quantisation and speculative decoding practice problems.
• 30u30 Day 13 · The Transformer: multi-head attention, the starting point for GQA.
• Stage 07 · Multimodal fusion: why multi-camera robot models produce so many tokens to cache.

6. Why it matters for robotics

7. Self-check

8. Courses & content

Start with the inference chapters of How To Scale Your Model: they teach the bandwidth-versus-compute arithmetic this whole stage rests on, with worked examples. Then use the vLLM docs for the benchmark.

Full course map for all 12 stages →

9. Papers & docs

• Fast Transformer Decoding: One Write-Head is All You Need (MQA) (Shazeer, 2019)
• GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints (Ainslie et al., 2023)
• Efficient Memory Management for LLM Serving with PagedAttention (vLLM) (Kwon et al., 2023)
• KIVI: A Tuning-Free Asymmetric 2bit Quantization for KV Cache (Liu et al., 2024)
• KVQuant: Towards 10 Million Context Length LLM Inference with KV Cache Quantization (Hooper et al., 2024)
• Fast Inference from Transformers via Speculative Decoding (Leviathan, Kalman, Matias, 2022)
• π0: A Vision-Language-Action Flow Model for General Robot Control (Black et al., 2024)
• Docs: vLLM documentation (quantised KV cache, speculative decoding, benchmarking)