Comparing transformer variants.

Every few weeks someone changes how a language model is wired inside and announces that the new version is more efficient. Checking the claim sounds easy: train both, see which does better. The catch is that you have to decide what stays the same between the two runs (the number of adjustable weights, the amount of arithmetic you pay for, or the hours on the clock), and that choice alone can pick the winner before either run starts.

It matters any time you change a model you train or serve. A change that looks like a wash on a training chart can let your Spark answer several times as many people at once. A change that looks like a clear win can turn out to be one run that got a luckier random start and a better-tuned setting than the other. Research labs like this question in interviews for exactly that reason: it separates people who repeat results from people who can check them.

You bring the GPT you wrote from scratch, Why does the KV cache exist? for the notes a model keeps about everything it has already read, and Why ~17 tokens per second? for why writing each new word is slowed by moving bytes, not by arithmetic. The same counting carries straight into Training on many GPUs, and the luck question gets its full treatment in Is that improvement real?.

Today

• Count, from first principles, how many arithmetic steps (FLOPs) one pass through a transformer costs, including the two pieces the famous shortcut of “2 steps per weight” leaves out, and check your count against PyTorch’s own counter.
• Explain why training costs about three times a plain pass (the 6ND rule: 6 steps for every weight, for every token it trains on), and how to measure what share of the chip’s top speed your training really uses (MFU).
• Say what “a fair comparison” means (same number of weights, same amount of arithmetic, or same hours on the clock) and build an example where the winner flips depending on which one you pick.
• On the Spark, give three attention designs that differ only in how many heads share one set of stored keys and values (MHA, GQA, MQA) the same arithmetic budget across several model sizes, each with its own tuned learning rate and more than one random start.
• Separate “cheaper to train” from “cheaper to serve”, measure both, and write a one-page answer to “which is more compute efficient, and for what?”

The main thing you need is the GPT you’ve already written from scratch. Read the KV cache lesson first for why writing each new word means re-reading everything the model has stored, and Why ~17 tokens per second? helps with the serving half. Budget an evening to read and write the code, one night of Spark time for the sweep, and an evening to measure and write up.

The Problem

Three messages land in a team channel on the same afternoon.

“Swapped MHA for MQA in the 125M config. Val loss is 0.04 worse. MQA is out.”

“GQA matches MHA and has 11% fewer parameters, so it’s more efficient.”

“Ran decode on both: GQA is 3x faster. Ship it.”

Each of these is a real measurement. None of them answers the question. The first one used the learning rate tuned for MHA on both models, ran one seed, and stopped both after the same number of steps. The second one is counting parameters, but nobody is paying for parameters; they are paying for GPU hours. The third one measured decoding at batch size 32 with long prompts, which is exactly where GQA shines; at batch size 1 with short prompts the gap is closer to 1.2x.

“Compare the compute efficiency of two transformer variants” sounds like one experiment. It is really three decisions you have to make before you run anything: what is held equal, what is measured, and which use (training or serving) you care about. Get those wrong and you can make either variant win. This lesson is about making them on purpose.

The Idea

The currency: a FLOP

A FLOPFLOPOne floating-point operation: a single add or a single multiply. “FLOPs” (lowercase s) is a count of them. “FLOP/s” is a rate: how many per second. is one add or one multiply on floating-point numbers. Counting them is the standard way to say how much arithmetic a model costs, because it doesn’t depend on which GPU you used or how good your kernels are. Time depends on both. FLOPs depend on neither.

Almost all the arithmetic in a transformer is matrix multiplication, and matrix multiplication has one rule worth tattooing somewhere: multiplying an (m × k) matrix by a (k × n) matrix costs 2·m·k·n FLOPs. There are m·n numbers in the output. Each one is a dot product of length k, which is k multiplies and (about) k adds. That’s where the 2 comes from.

Now look at it from one token’s point of view.

That last line is the whole trick. Every weight matrix in the model gets touched exactly once per token in the forward pass, and each touch is a multiply and an add. So if the model has N parameters sitting in weight matrices, the forward pass costs about 2N FLOPs per token. This is the famous rule, and it is a good first approximation. It is also wrong in two specific places, and at small scale both matter.

Where N comes from

Take one standard block of a GPT with model width d (d_model) and an MLP that expands to 4d.

• Attention projections: the Q, K, V and output matrices are each d × d. That’s 4d² parameters.
• MLP: up-projection d × 4d and down-projection 4d × d. That’s 8d².

So each layer holds about 12d² parameters, and with L layers, N ≈ 12·L·d² (ignoring the tiny LayerNorm and bias vectors). For GPT-2 small, L = 12 and d = 768: 12 × 12 × 768² ≈ 85M. That matches the usual “85M non-embedding parameters” for the 124M model; the other ~39M are the token and position embeddings.

People usually define N as the non-embedding parametersnon-embedding parametersThe parameters in the transformer blocks, not counting the token embedding table (vocab × d) or position embeddings. Kaplan et al. found loss scales more cleanly with this count., because an embedding lookup is not a matmul. Fetching row 4,182 of a table costs no arithmetic at all.

The two things 2N leaves out

1. The attention itself. The projections are matmuls with weights. But inside attention there are two more matmuls that use no weights at all: the scores Q·Kᵀ, and the mixing of values, softmax(scores)·V. Their cost depends on how many tokens are in the context, not on N.

Count it for one token at position t, in one layer. Its query (length d once you line up all the heads side by side) gets dotted with t keys: 2·t·d FLOPs. Then the t attention weights mix t value vectors of length d: another 2·t·d. So 4·t·d per layer per token.

Averaged over a sequence of length T, t depends on how you count:

• If the kernel computes the full T × T square and then masks the future out (the simple implementation, and how PyTorch’s FLOP counter counts it), every token pays as if t = T: 4·L·T·d per token.
• If you only count the half of the square a causal model actually needs (each position sees, on average, T/2 earlier ones), it is 2·L·T·d per token. This is the term in Kaplan et al.’s formula, written 2·n_layer·n_ctx·d_attn. Fused attention kernels like FlashAttention do skip most of the masked-out blocks, so this is closer to the work actually done.

Neither is “the” right answer. Pick one, say which, and use the same one for both variants.

How big is it next to 2N? Divide: 4·L·T·d over 2·12·L·d² is T / (6d) (full counting) or T / (12d) (causal counting). With d = 768 and T = 1024, full counting gives 22% extra. With a tiny model, d = 256 and T = 2048, it is 133%: the attention costs more than all the weights combined.

2. The output head. The last step turns each token’s d-vector into scores over the whole vocabulary: a d × V matmul, 2·d·V FLOPs per token. With weight tying, those weights are the embedding table, so they’re excluded from N, yet this matmul is very real. For GPT-2 small: 2 × 768 × 50,304 ≈ 77M FLOPs per token, against 2N ≈ 170M. At large scale it’s a rounding error. At the sizes you can train on one Spark, it’s a quarter of your compute.

Now play with it. Start at GPT-2 small. Push sequence length up until attention overtakes the MLP. Then drop d_model to 256 and watch the output head eat the bar. Then toggle MQA: notice which segment changes, and which doesn’t.

That last point is the first surprise about GQA and MQA: they barely change training FLOPs. In GPT-2 small, going from 12 KV heads to 1 trims the K and V projections from 2d² to about d²/6 per layer, which is a 9% cut in forward FLOPs. The expensive things, the MLP, the scores, the head, don’t move. Whatever these designs are good for, it isn’t mainly training arithmetic.

Training: why 6N

Training does a forward pass and then a backward pass. The backward pass costs about twice the forward, and the reason is mechanical.

For a layer that computes Y = X·W, the backward pass gets handed dY (how the loss changes with the output) and has to produce two things:

• dX = dY·Wᵀ, to pass the gradient further back to earlier layers. Same shape of matmul as the forward: 2 FLOPs per weight per token.
• dW = Xᵀ·dY, the gradient for the weights themselves, which the optimizer needs. Another matmul of the same size: 2 more.

So training costs about 6N FLOPs per token, and over a run of D tokens, C ≈ 6·N·D. This is the number scaling-law papers put on their x-axis. The attention and head terms triple the same way, so at small scale the honest version is C ≈ 3 × (forward FLOPs per token) × D, with your full forward count. I checked the 3x on a Mac too: FlopCounterMode around forward plus backward of the GPT-2 config reported exactly 3.0 times the forward.

Two refinements you’ll see:

• If you use activation checkpointingactivation checkpointingThrowing away some activations during the forward pass to save memory, then recomputing them during the backward pass. Costs roughly one extra forward pass., the hardware does about 8N per token, not 6N. The extra 2N is real work but not useful work, so it is not counted as “model FLOPs”.
• The optimizer step (Adam’s few operations per parameter, once per step, not per token) is tiny next to 6N × tokens per batch, and is usually ignored.

MFU: how much of the GPU you’re actually using

Once you can count FLOPs, you can ask how close to the hardware’s limit you are. MFUMFUModel FLOPs Utilization. The model FLOPs your training actually completes per second (6N plus attention, times tokens per second), divided by the hardware’s peak FLOP/s. Defined in Google’s PaLM paper, 2022. is:

MFU = (model FLOPs per token × tokens per second) / peak FLOP/s of the hardware

The PaLM paper introduced it because the older measure counted recomputation as useful work, which flattered runs that recomputed a lot. Well-tuned large runs report MFU in the 40-60% range; PaLM 540B reported 46%. A small model on one GPU is usually much lower, because its matmuls are too skinny to fill the chip, and Python and kernel-launch overhead eat a bigger share of each step.

For the Spark you’ll want a realistic denominator. The headline “1 petaFLOP” on its spec sheet is FP4 with sparsity, which your bf16 training never touches. Dense bf16 is several times lower. Rather than trust a number, you’ll measure it: time a large bf16 matmul and call that your “achievable peak” (Task 0). Then MFU against that tells you how good your code is, not how optimistic the brochure was.

MFU is not part of the comparison itself; it’s a sanity check on it. If one variant runs at 15% MFU and the other at 35%, you’ll see a big wall-clock difference that has nothing to do with the architecture and everything to do with kernels.

What “fair” means

Now the heart of it. You have two variants, A and B. You want to know which is “more efficient”. There are three things you could hold equal:

And the most common thing people actually hold equal is a fourth one, training steps, which is equal-parameters-and-tokens when the batch is the same, with the same trap.

The default currency for “compute efficiency” is equal FLOPs, and the reason is that compute is what you pay for and what you can redistribute. With a fixed budget you can always choose a bigger model trained for fewer tokens, or a smaller one for more. A fair FLOP comparison lets each variant make that choice for itself.

Here’s how the answer flips. Suppose B does 1.5x the FLOPs per parameter per token (it reuses its weights more, the way a weight-shared looped block or a longer context window does) and gets 1.3x as much out of each parameter. At equal parameters and equal tokens, B wins, obviously: it’s using 50% more compute. At equal FLOPs, B has to train on two-thirds as many tokens, or be a smaller model, and that can cost more than its per-parameter advantage gains.

isoFLOP curves: the method

The Chinchilla paper gives a clean recipe for comparing at equal compute without guessing the right model size. It’s called an isoFLOP curveisoFLOP curveFix a compute budget C. Train several model sizes, each for exactly C FLOPs (smaller models see more tokens). Plot final loss against model size. The curve is U-shaped; its bottom is the best size for that budget. (“iso” is Greek for “equal”):

1. Fix a compute budget C.
2. Pick 4 or 5 model sizes N around your guess of the best size. For each, set the number of training tokens to D = C / (FLOPs per token), so every run costs exactly C.
3. Train each to the end of its schedule. Record final validation loss.
4. Plot loss against log N. You get a U: too small a model is limited by its capacity, too big a model didn’t see enough data. The bottom of the U is the compute-optimalcompute-optimalThe model size (and matching token count) that reaches the lowest loss for a given compute budget. point for that C.
5. Repeat at two or three budgets, each about 3x the last.

Do that for both variants and you get two sets of U’s. The comparison is between the bottoms: at each budget, which variant’s best run is lower? Connect the bottoms across budgets and you have each variant’s compute frontiercompute frontierThe best loss a design reaches at each compute budget, using its best size. Plotted against log compute it’s close to a straight line (a power law).. If B’s frontier sits below A’s, B is more compute efficient, and you can say by how much in a currency anyone understands: “B reaches A’s loss with 0.7x the compute”, a compute multipliercompute multiplierHow much compute variant A would need to match variant B’s loss at budget C, divided by C. Above 1 means B is better. Lets you compare improvements of different kinds on one scale. of about 1.4x.

The multiple budgets aren’t decoration. A variant can win at small scale and lose at large (or the reverse) if its frontier has a different slope. Two points can’t tell you a slope reliably; three can at least warn you.

Things to try:

• Costly + strong. The square sits below A’s dot (B wins at equal params), but B’s own dot sits above A’s (B loses at equal FLOPs). The compute multiplier is below 1: B is a worse use of money.
• Cheap + weak. The reverse. This is roughly the MQA story: a little worse per parameter, a little cheaper per token.
• Identical, then nudge B’s FLOPs to 1.1x. A 10% cost increase with no quality change shows up as a compute multiplier near 0.9, and a loss gap at the dots of only about 0.01 to 0.03, depending on the budget. That’s the size of effect you’re hunting for, and it’s about the size of seed noise. Hold that thought.
• Move the budget slider. The compute multiplier doesn’t budge: in this toy it is exactly (quality per param) ÷ (FLOPs per param), because a B that’s m times better per parameter and k times costlier per token behaves just like A with m/k times the compute. So the verdict never flips with scale. Real variants don’t have constant multipliers, and theirs can.

The variants: MHA, GQA, MQA

Here are the designs you’ll actually compare. All three keep the same number of query heads. They differ in how many key and value heads those queries share.

MHAMHAMulti-head attention, the original design: every query head has its own key head and value head. is the original. MQAMQAMulti-query attention (Shazeer, 2019): all query heads share a single key head and a single value head. came from Noam Shazeer in 2019, who noticed that cached decoding is limited by reading keys and values from memory (the point of the KV cache lesson), so sharing one K and V across all heads makes each decode step read far less. The price was a little quality. GQAGQAGrouped-query attention (Ainslie et al., 2023): query heads are split into groups; each group shares one key head and one value head. MHA and MQA are its two extremes., from Joshua Ainslie and colleagues in 2023, is the middle: groups of query heads share a KV head. They reported quality close to MHA and decode speed close to MQA, and showed an existing MHA model can be converted by averaging its KV heads within each group and then training a little more (“uptraining”, about 5% of the original pretraining compute). Llama 2 70B and every Llama 3 model use GQA.

In code the change is tiny. The K and V projections output n_kv_head × head_dim instead of n_head × head_dim, and before attention each KV head is broadcast to the query heads in its group. Recent PyTorch does the broadcast for you: F.scaled_dot_product_attention(q, k, v, is_causal=True, enable_gqa=True) accepts k and v with fewer heads than q (I used this on PyTorch 2.14; check your version’s docs).

Now put these through the fairness machinery.

Training. At the same width and depth, GQA and MQA have fewer parameters and slightly fewer FLOPs per token (Figure 002). So “same config” is neither equal parameters nor equal FLOPs. Papers that compare at equal parameters usually widen the MLP to give back what attention lost. An isoFLOP sweep side-steps the question: each variant gets the same compute and its best size. The expected outcome, from the papers and from how little attention projections matter, is that GQA lands on top of MHA within noise, and MQA is close but perhaps slightly worse. “Perhaps slightly” is exactly the kind of claim this lesson is teaching you to check rather than repeat.

Inference is a different world. In the bandwidth lesson you saw that each decode step has to read all the weights from memory, and in the KV cache lesson that it also reads every cached key and value for every sequence in the batch. Arithmetic is not the bottleneck; bytes are. So the right model for decode time is:

time per decode step ≈ (weight bytes + batch × context × KV bytes per token) / memory bandwidth
KV bytes per token   = 2 (K and V) × layers × n_kv_head × head_dim × bytes per number

The weights term doesn’t care about batch size. The KV term grows with batch and with context. So:

• Batch 1, short context: weights dominate. Shrinking the KV cache barely matters. GQA and MQA are a little faster mostly because their K and V weights are smaller.
• Batch 32, long context: the cache dominates. Cutting it 4x or 16x cuts the step time a lot, and lets you fit far more sequences in memory at once.

So “which is more efficient?” has at least two honest answers, and a good results doc gives both: per training FLOP, they’re roughly tied (measure it); per served token at realistic batch sizes, GQA and MQA win by a lot, and the win grows with batch and context. That separation, training efficiency vs inference efficiency, is why every modern open model uses GQA: you pay almost nothing at training time for a big win at serving time, and serving runs for much longer than training.

Other variants worth the same treatment

The method is general. Three others you’ll meet, each with its own fairness wrinkle:

• SwiGLU vs GELU MLP (Shazeer, 2020, “GLU Variants Improve Transformer”). A SwiGLUSwiGLUAn MLP with three weight matrices instead of two: the input goes through two up-projections, one is passed through SiLU and multiplied elementwise with the other, then a down-projection. Used in Llama, PaLM and most recent models. MLP has three matrices instead of two. To compare fairly, people shrink its hidden size from 4d to 8d/3, which makes parameters and FLOPs equal (3 × d × 8d/3 = 8d²; I checked the counts match exactly for d = 768). At the same hidden size, SwiGLU gets 50% more MLP compute and of course looks better.
• RoPE vs learned positions. You’ve implemented RoPERoPERotary position embedding: positions are encoded by rotating each query and key vector by an angle that depends on its position, so the dot product only depends on how far apart two tokens are.. It has no parameters and negligible FLOPs, while learned position embeddings add a T × d table (no FLOPs, just a lookup). So for this pair, equal-params and equal-FLOPs nearly coincide, and the comparison is about loss and about behaviour at lengths longer than training.
• Pre-norm vs post-norm. Putting LayerNorm before each sublayer (as GPT-2 does) instead of after (as the original Transformer did) changes no FLOPs at all, but post-norm needs a learning-rate warmup and a lower peak LR to train stably (Xiong et al., 2020). Here the whole comparison lives or dies on LR tuning.

Noise: is the gap real?

Train the same config twice with different seedseedThe number that initialises the random number generator: it decides the initial weights and the order of the data. Same seed, same run (nearly); different seed, slightly different final model.s and the final validation loss won’t match. At the sizes you’ll train, the spread is often around 0.005 to 0.02 in loss (nats per token). That’s a rough range; you’ll measure yours. The differences you’re looking for between MHA and GQA are the same size or smaller.

So one run per variant tells you almost nothing. You need either several seeds per point, so you can compare a difference in means to the spread, or a margin so large that noise can’t explain it. Is that improvement real? covers the statistics properly. The working rule here: run at least 2-3 seeds at each frontier point, report mean ± standard deviation, and don’t call a difference smaller than about twice the seed spread.

One mercy: the isoFLOP U is fitted through 4-5 points, and fitting a curve averages away some noise. The bottom of a fitted parabola is more trustworthy than any single run.

Confounders: the comparison you didn’t mean to run

A confounderconfounderSomething other than the thing you’re testing that differs between the two runs and affects the result. It makes you credit (or blame) the variant for what the setup did. is anything other than the architecture that differs between your runs. The usual suspects, in roughly the order they bite:

1. Learning rate tuned for one variant. You tuned MHA, then reused its LR for MQA. Different designs often have different best LRs. Tune each one, at least coarsely, at your smallest budget, and check the best isn’t at the edge of your grid.
2. Schedule length not matched to run length. The Kaplan/Chinchilla lesson. Every run’s cosine decay must end when that run ends.
3. Different tokens. Different data order is fine (that’s seed noise); different data, tokenizer or validation set is not.
4. Different parameter counting. Counting embeddings for one and not the other, or forgetting that tied weights are counted once.
5. Different FLOP counting. Using 2N for one variant whose real cost is dominated by attention.
6. Equal steps when you meant equal FLOPs. A variant that’s 10% cheaper per token should get 10% more tokens.
7. Speed from software, not design. In wall-clock comparisons: one variant hits a fast fused kernel and the other falls back to a slow path. Check MFU for both.

Narang et al. (2021) re-tested dozens of published transformer modifications in one shared codebase. Most of them didn’t help. The ones that did were mostly small and computationally cheap, the GLU variants among them. Treat that as a base rate: when a variant looks like a big win, look for the confounder first.

Build It

Everything runs on the Spark except the FLOP counter check, which runs anywhere. Layout:

archcmp/
  model.py        # configurable GPT: MHA/GQA/MQA, gelu/swiglu
  flops.py        # analytic counter + check against torch
  train.py        # one run: config, budget C, lr, seed -> final val loss (JSON)
  sweep.py        # LR sweep, then isoFLOP sweep
  plot.py         # isoFLOP U's and frontiers
  decode_bench.py # KV cache size, decode tok/s at batch 1 and 32
  results/        # one JSON per run
  RESULTS.md      # the write-up

# peak.py: time a big bf16 matmul and report TFLOP/s
import torch, time
n = 8192
a = torch.randn(n, n, device="cuda", dtype=torch.bfloat16)
b = torch.randn(n, n, device="cuda", dtype=torch.bfloat16)
for _ in range(3):
    a @ b
torch.cuda.synchronize()
iters = 20
t = time.perf_counter()
for _ in range(iters):
    a @ b
torch.cuda.synchronize()
dt = (time.perf_counter() - t) / iters
print(f"{2 * n**3 / dt / 1e12:.1f} TFLOP/s achievable bf16")

# model.py (shape only)
@dataclass
class Cfg:
    vocab: int = 16384
    ctx: int = 512
    n_layer: int = 6
    d: int = 384
    n_head: int = 12        # head_dim = d // n_head
    n_kv_head: int = 12     # 12 = MHA, 3 = GQA (groups of 4), 1 = MQA
    mlp: str = "gelu"       # or "swiglu" with hidden 8d/3
    tie_weights: bool = True

class Attention(nn.Module):
    def __init__(self, c: Cfg): ...
        # q: d -> n_head * hd ; kv: d -> 2 * n_kv_head * hd ; o: d -> d
    def forward(self, x, kv_cache=None): ...
        # returns y, and the updated cache when kv_cache is given.
        # The cache stores the small, un-broadcast (B, n_kv_head, T, hd) K and V.

class GPT(nn.Module):
    def __init__(self, c: Cfg): ...
    def forward(self, idx, targets=None, kv_cache=None): ...
    def num_params(self, non_embedding: bool = True) -> int: ...

# flops.py
def forward_flops_per_token(c: Cfg, T: int, causal_half: bool = False) -> dict[str, int]:
    """Return {'proj': ..., 'attn': ..., 'mlp': ..., 'head': ...}, forward FLOPs per token
    at sequence length T. 'attn' is 4*L*T*d (full square) or 2*L*T*d (causal_half)."""
    ...

def train_flops_per_token(c: Cfg, T: int) -> int:
    """3x the forward total."""
    ...

def check(c: Cfg, T: int = 512) -> None:
    """Run one forward (and one forward+backward) under FlopCounterMode, print torch's total
    next to yours and the ratio. It should be 1.000."""
    ...

from torch.utils.flop_counter import FlopCounterMode
with FlopCounterMode(display=False) as fc:
    model(x)
print(fc.get_total_flops())

from torch.nn.attention import sdpa_kernel, SDPBackend
with sdpa_kernel(SDPBackend.MATH), FlopCounterMode(display=False) as fc:
    model(x)

# plot.py (shape only)
def load(results_dir) -> list[dict]: ...
def fit_bottom(Ns, losses) -> tuple[float, float]:
    """Fit loss = a*(log N)^2 + b*log N + c through one budget's points;
    return (N at the minimum, loss at the minimum)."""
    ...
def compute_multiplier(frontier_a, frontier_b) -> list[float]:
    """Fit log(loss - E) ~ log C for A's frontier (or just interpolate in log C),
    and for each of B's points find the C_A where A reaches B's loss. Return C_A / C."""
    ...

# decode_bench.py (shape only)
def kv_bytes(cache) -> int:
    """Sum of .nbytes over every cached K and V tensor. Compare with the formula."""
    ...

@torch.no_grad()
def bench(model, batch: int, prompt: int = 512, gen: int = 512) -> dict:
    """Prefill `prompt` random tokens for `batch` sequences, then decode `gen` tokens one at a
    time with the cache. Return {'kv_GiB': ..., 'decode_tok_s': batch*gen/seconds, 'peak_mem_GiB': ...}."""
    ...

## Answer
<one paragraph: training efficiency verdict with compute multiplier ± noise;
 inference verdict with tok/s ratios at batch 1 and 32; the recommendation.>

## Setup
data, tokenizer, shapes, budgets, seeds, LR per variant, FLOP counting convention.

## Training: isoFLOP results
plot + table. Seed spread. What is and isn't distinguishable.

## Inference
KV bytes per token, decode tok/s at batch 1 and 32, vs the bandwidth floor.

## Caveats
what scale this was, what might change at 100x, what you didn't control.

Measure

FLOP counter check (any machine):

For reference, my Mac run on a GPT-2-small shape (12 layers, d = 768, T = 1024, vocab 50,304, math attention): MHA 291.7 GFLOP, GQA with 4 KV heads 272.4, MQA 265.2, all ratios 1.0000, and fwd+bwd exactly 3.0x. Without forcing the math path, torch reported 253.1, 233.7 and 226.5: exactly 38.7 GFLOP short each time, which is the attention term.

LR sweep:

isoFLOP bottoms (fitted, mean ± sd over seeds at the nearest points):

Throughput and MFU (training, per shape): tokens/s and MFU against your measured peak. Also write down the same shape’s MFU for each variant: if they differ by more than a few points, your wall-clock comparison is a kernel comparison.

Inference (24 layers, d = 1024, head 64, context 1024 at the end):

The floors are from the formula with the spec-sheet 273 GB/s and are approximate; measured numbers should be lower.

What the numbers should look like, and why:

• GQA vs MHA in training: indistinguishable, most likely. The difference is a small cut in a fifth of the FLOPs. If your GQA bottom is within a seed sd or two of MHA’s, the honest statement is “no detectable difference at these budgets”, not “GQA is 0.3% better”.
• MQA vs MHA in training: close, possibly slightly worse. If you see a gap, check it survives seeds and per-variant LR before believing it.
• The best N shifts. For the same C, GQA’s and MQA’s best runs should sit at a slightly larger shape (they’re cheaper per token at a given shape). This is the equal-FLOP comparison doing its job.
• Decode, batch 1: a modest win for GQA/MQA, partly from smaller weights. If you’re launch-bound, maybe nothing at all.
• Decode, batch 32: a large win, roughly tracking the bandwidth floors’ ratios. If it doesn’t, check that your cache really is stored un-broadcast (a repeat_interleave into the cache gives back all the memory you saved).

Break It

1. Tune the LR for MHA only. Rerun MQA’s U at the middle budget with MHA’s best LR instead of its own. If the gap between the variants changes by more than the seed spread, you’ve just measured how much of a “result” a confounder can manufacture. This is the most common way variant comparisons go wrong.

2. Compare at equal steps, with the wrong SwiGLU. Add a SwiGLU variant with hidden size 4d (not 8d/3). Train it and GELU for the same number of steps on the same shape. SwiGLU “wins” comfortably. Now put both through one isoFLOP budget, counting its real FLOPs (50% more in the MLP). Most of the win shrinks. Then fix the hidden size to 8d/3 and compare again. Whatever is left is the actual design effect.

3. Trust one seed. Run five seeds of one config and look at the spread. Then go back to your seed-0-only isoFLOP points and ask how many of the orderings between variants would survive. Usually, at these budgets, very few of the small ones do.

4. Count FLOPs with 2N at long context. Rerun the smallest shape at sequence length 2048 and set its steps with 6N instead of your full counter. With d = 256, the attention term alone is more than all the weight FLOPs, so the run silently costs about twice its budget, and a variant whose attention is cheaper looks bad. Then remember Task 2: if your FLOP counter can’t see the fused attention kernel, it makes exactly this mistake for you.

Explain It Back

Go Deeper

• “Scaling Laws for Neural Language Models”, Jared Kaplan et al. (OpenAI), 2020. Table 1 is the FLOP accounting this lesson derives, and the paper is where “compare at equal compute” became standard practice.
• “Training Compute-Optimal Large Language Models”, Jordan Hoffmann et al. (DeepMind), 2022. The Chinchilla paper. Read the isoFLOP method (their “Approach 2”) and Appendix F’s FLOP counting, which counts both attention matmuls in full.
• “Fast Transformer Decoding: One Write-Head is All You Need”, Noam Shazeer, 2019. Introduces MQA. Short, and its argument about memory traffic during decoding is the inference half of this lesson.
• “GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints”, Joshua Ainslie et al., 2023. The quality/speed tradeoff between MHA and MQA, measured, plus how to convert an existing model.
• “GLU Variants Improve Transformer”, Noam Shazeer, 2020. Two pages of results, and a model of how to keep params and FLOPs equal when the variant changes the shape (the 8d/3 trick).
• “PaLM: Scaling Language Modeling with Pathways”, Aakanksha Chowdhery et al. (Google), 2022. Defines MFU, with the exact FLOPs-per-token formula including attention, in an appendix.
• “Do Transformer Modifications Transfer Across Implementations and Applications?”, Sharan Narang et al. (Google), 2021. Dozens of modifications re-tested in one codebase; most didn’t help. The best argument for careful baselines.
• “Transformer Math 101”, Quentin Anthony, Stella Biderman and Hailey Schoelkopf, EleutherAI blog, 2023. The practical cheat sheet: 6ND, memory, and how to sanity-check a training run’s throughput.

Connections

This builds directly on Why does the KV cache exist? (what the cache stores and why decoding reads it all) and Why ~17 tokens per second? (decode is bytes, not FLOPs). The training side leans on Train a small LM and Tokenizers for the data and model, and the noise question is the whole subject of Is that improvement real?.

It sets up Training on many GPUs, where the same FLOP accounting tells you how long a run should take and MFU tells you how badly communication is hurting, and Why training needs so much memory, the other half of what a variant costs.