Decode Compute Math, in depth: a per-operator FLOP and byte ledger for one LLM decode step

Three counting rules

Three counting rules cover almost everything in a transformer forward pass.

1. A matrix-vector product costs 2 FLOPs per weight. Multiplying a token's activation vector of length d_in by a weight matrix d_in x d_out is one multiply and one add per weight. For B tokens it is 2 * d_in * d_out * B FLOPs.
2. The weight bytes are paid once per step, not once per token. A good kernel streams each weight tile once and applies it to every token in the batch, so a weight GEMM moves d_in * d_out * bytes_per_weight bytes regardless of B.
3. Arithmetic intensity is FLOPs divided by bytes. A GPU can only use its tensor cores fully if the work supplies at least peak_FLOPs / bandwidth FLOPs per byte, the ridge point. For an H100 SXM, 989 TFLOP/s dense BF16 over 3.35 TB/s of HBM3 gives about 295 FLOP/byte.

For BF16 weights (2 bytes each), a decode weight GEMM therefore has intensity 2 * B / 2 = B: at batch 1, one FLOP per byte, about 300 times below the ridge. The interesting operators are those where rule 2 fails.

The ledger for one step

Llama 3 8B has 32 layers, a hidden size of 4096, 32 query heads and 8 key/value heads of dimension 128 (grouped-query attention with groups of four), a SwiGLU MLP with intermediate size 14336, and a 128,256-token vocabulary. Per layer and per step, the operators are:

Two things are left out on purpose. The embedding lookup reads only B rows of a 1 GB table, so it is negligible. Norms, rotary embeddings, residual adds and sampling are tiny in FLOPs and bytes, but each is a kernel launch unless fused, which is one reason measured steps sit above these floors.

Attention is the odd one out. Each sequence has its own KV cache, so there is nothing to share across the batch. For one sequence at context L, a layer reads 2 x 8 x 128 x L values (K and V) at 2 bytes each, and computes QK^T and the weighted sum of V for 32 query heads: 4 x 32 x 128 x L FLOPs. Divide and the intensity is exactly 4 FLOP/byte, which is the GQA group size (32 query heads sharing 8 KV heads) for a 2-byte cache. It does not depend on B or L. Doubling the batch doubles both the FLOPs and the bytes.

A calculator and its output

The calculator below implements the ledger; swap in your own model's config.

# Per-operator FLOP and byte ledger for ONE decode step of a dense GQA transformer.
from dataclasses import dataclass

@dataclass
class Model:
    layers: int = 32        # Llama 3 8B
    d: int = 4096           # hidden size
    heads: int = 32
    kv_heads: int = 8
    head_dim: int = 128
    ffn: int = 14336        # SwiGLU intermediate size
    vocab: int = 128256
    w_bytes: float = 2.0    # BF16 weights
    kv_bytes: float = 2.0   # BF16 KV cache

def step_ledger(m, batch, ctx):
    """FLOPs and HBM bytes for one decode step: `batch` sequences, each with `ctx` cached tokens."""
    kv = m.kv_heads * m.head_dim
    weights = {                     # matrix shapes (in, out), read once per step
        "qkv_proj": m.d * (m.d + 2 * kv),
        "o_proj":   m.d * m.d,
        "mlp":      3 * m.d * m.ffn,  # gate, up, down
    }
    rows = {}
    for name, n in weights.items():
        rows[name] = (2 * n * batch * m.layers,                 # 2 FLOPs per weight per token
                      n * m.w_bytes * m.layers)                 # bytes: independent of batch
    # Attention: every sequence reads its own K and V; nothing is shared across the batch.
    attn_flops = 4 * m.heads * m.head_dim * ctx * batch * m.layers   # QK^T and PV
    attn_bytes = 2 * kv * ctx * m.kv_bytes * batch * m.layers         # K and V
    rows["attention"] = (attn_flops, attn_bytes)
    rows["lm_head"] = (2 * m.d * m.vocab * batch, m.d * m.vocab * m.w_bytes)
    return rows

def report(m, batch, ctx, peak=989e12, bw=3.35e12):
    rows = step_ledger(m, batch, ctx)
    tf = sum(f for f, _ in rows.values()); tb = sum(b for _, b in rows.values())
    print(f"batch={batch} ctx={ctx}")
    for k, (f, b) in rows.items():
        print(f"  {k:10s} {f/1e9:9.1f} GFLOP {b/1e9:7.2f} GB  intensity {f/b:6.1f}")
    print(f"  total      {tf/1e9:9.1f} GFLOP {tb/1e9:7.2f} GB  intensity {tf/tb:6.1f}")
    print(f"  floors: memory {tb/bw*1e3:.2f} ms, compute {tf/peak*1e3:.2f} ms")

Running report(Model(), 1, 2048) and report(Model(), 64, 2048) prints these numbers (GB here are 10^9 bytes):

Three readings. First, the MLP is three quarters of weight bytes. Second, at batch 1 the step reads 15.3 GB and does 16 GFLOP, so the memory floor at full bandwidth is 4.6 ms while the compute floor is 0.02 ms. Third, at batch 64 the weight bytes have not moved but attention bytes went from 0.27 GB to 17.2 GB and are now the largest single item.

Why batching plateaus

To turn bytes and FLOPs into time, apply a roofline per operator: each operator takes max(flops / achieved_flops, bytes / achieved_bw), and the step is the sum. Using 60 percent of peak compute and 80 percent of peak bandwidth as achieved rates (assumptions; measure your own with a microbenchmark), the sweep at 2,048 tokens of context per sequence looks like this:

def step_time(m, batch, ctx, peak=989e12 * 0.6, bw=3.35e12 * 0.8):
    rows = step_ledger(m, batch, ctx)
    return sum(max(f / peak, b / bw) for f, b in rows.values())

for b in [1, 8, 32, 64, 128]:
    t = step_time(Model(), b, 2048)
    print(b, round(t * 1e3, 2), "ms", round(b / t), "tok/s")

From batch 1 to 8, throughput rises about sevenfold for a 12 percent increase in step time: the weights were being read anyway. From 64 to 128, doubling the batch buys only 30 percent more, because attention bytes scale linearly with batch. Batch 256 at this context would need about 84 GB for weights plus cache, which does not fit on an 80 GB card at all, so memory capacity ends the curve before compute does. Every weight GEMM is still far below the ridge at batch 128: the tensor cores idle even at the plateau.

There is a clean crossover rule hiding in the ledger. The weights read per step total 15.0 GB (excluding the embedding table) and each token of context costs 131,072 bytes (128 KiB) of KV across all layers. KV bytes exceed weight bytes once the sum of context lengths across the batch passes about 114,500 tokens. Below that, batching is nearly free. Above it, you are paying for every additional sequence in full. Cache sizing is covered in KV cache sizing.

Long context changes the mix

Now hold the total cache constant and change its shape. Batch 8 at 32,768 tokens holds exactly as many cached tokens as batch 128 at 2,048: the step reads the same 49.4 GB and takes the same 18.4 ms. But it produces 8 tokens per step instead of 128, so throughput drops from about 6,900 to about 430 tokens per second. Total intensity falls to 5.2 FLOP/byte because attention, at intensity 4, now dominates the mix.

Long context is expensive in a way per-token pricing hides. The lever is attention bytes per cached token:

• Fewer KV heads. GQA with 8 KV heads already cut the cache 4x relative to full multi-head attention with 32. Intensity is the group size, so more sharing also raises intensity.
• A smaller cache dtype. An FP8 KV cache halves attention bytes and doubles attention intensity to 8. In the calculator, batch 8 at 32k drops from 18.4 ms to 12.0 ms per step; batch 64 at 2k drops from 12.0 ms to 8.8 ms. Evaluate quality first.
• Reading less of it. Sliding windows, eviction and sparse attention shrink L directly but change model behaviour.

FlashAttention-style kernels do not change these byte counts in decode. They help by avoiding intermediate writes and splitting long sequences across thread blocks; see FlashAttention for the mechanics.

Weight quantisation in the ledger

Weight-only quantisation such as W4A16 stores weights in 4 bits and dequantises them to 16 bits inside the GEMM kernel. In the ledger it divides weight bytes by four and leaves FLOPs unchanged. With the same roofline, batch 1 drops from 5.7 ms to about 1.5 ms, close to the full 4x. At batch 64 and 2k context it drops only from 12.0 ms to 8.0 ms, because 17 GB of BF16 attention bytes are untouched.

Two caveats. The calculator ignores dequantisation work. It hides under memory time at small batch, and W4A16 kernels typically lose their edge over BF16 as batch grows; where depends on kernel and GPU. Second, for a batch job dominated by KV reads, KV quantisation is the better lever.

Mixture of experts: the union problem

In a mixture-of-experts layer, each token is routed to k of E experts. At batch 1 the step reads only k experts' weights. At larger batches, different tokens pick different experts, and the step reads the union. Under uniform, independent routing, an assumption that real routers violate, the expected number of distinct experts touched is:

def expected_experts(E, k, B):
    # each token picks k of E; P(a given expert is untouched by one token) = 1 - k/E
    return E * (1 - (1 - k / E) ** B)

# Mixtral-style 8 experts, top-2:   B=1 -> 2.0, B=4 -> 5.5, B=8 -> 7.2, B=16 -> 7.9
# 128 experts, top-8:               B=1 -> 8,   B=8 -> 51.6, B=32 -> 111.8, B=64 -> 125.9

The consequence is that MoE decode is cheap only at small batch. With 8 experts and top-2 routing, by batch 16 almost every expert is read on every step, so weight bytes approach those of a dense model with the full parameter count, while each expert sees only about B x k / E tokens and its GEMM intensity is correspondingly low. Real routing is skewed, shrinking the union but creating hot experts; see expert parallelism.

Measuring against the ledger

Turning the ledger into measurements: the two metrics to compute from a profile are model bandwidth utilisation (MBU), the ledger's bytes per step divided by measured step time, divided by peak bandwidth; and model FLOPs utilisation (MFU), the same with FLOPs. Healthy decode shows high MBU and low MFU. If MBU is low too, time is going somewhere the ledger does not count:

• Kernel launch and Python overhead. Hundreds of kernels per step add up at batch 1; CUDA Graphs replay the step with one launch.
• Tensor-parallel communication. Splitting across N GPUs divides weight and KV bytes per GPU by N but adds two all-reduces of B x d activations per layer. These are latency-bound at small batch.

A common failure: benchmarking at batch 1 with short prompts, then deploying into batch 64 at 8k tokens, where attention dominates and the chosen weight-side optimisation barely helps. Compute the ledger for your production mix of batch and context.

What to do next

1. Plug your model's config into the calculator: layers, hidden size, KV heads, head dimension, MLP size and vocabulary. Check the per-token KV bytes against your serving engine's reported cache usage.
2. From production logs, compute total cached tokens per step and compare it with the crossover (weight bytes over KV bytes per token).
3. Profile one decode step and compute MBU. Below about 60 percent, look for launches, unfused kernels and communication before buying hardware.
4. If you are below the crossover, pursue weight-side levers: quantised weights, bigger batches, speculative decoding.
5. If you are above it, pursue cache-side levers: an FP8 KV cache, prefix sharing, shorter retained context, and admission limits on total cached tokens.
6. For MoE models, compute expected experts touched at your batch size and route-skew measurements before assuming the active parameter count sets your cost.