LM10-T6 — Inference Efficiency and Architecture Trade-Offs

Learning objectives

After this lesson, you should be able to:

• Build a memory and speed budget from weights, cache, batch and bandwidth.
• Explain exact attention kernels, speculative decoding, mixture of experts and sliding windows in terms of the cost they remove.
• Choose levers for a given constraint and state what to evaluate for each.
• Recognise levers that trade quality and those that do not.

The idea: four costs, many levers

A deployment is limited by one of four quantities. Weight memory: the model must fit. Cache memory: the contexts of all concurrent users must fit. Bandwidth: every decode step streams weights and cache. Arithmetic: prefill, and decode at large batch, need compute. Each technique in this module reduces one of them:

The second-to-last column is the one to read first. Levers that do not change the output (batching, exact kernels, speculative decoding with the proper rule) are safe to try; they cost engineering effort but no quality. Levers that change the output need the evaluation methods of Module 9 before you rely on them.

Exact attention kernels

Standard attention writes the score matrix (n × n per head) to memory, reads it back for the softmax, writes the probabilities and reads them again to multiply with the values. For long sequences these round trips dominate. FlashAttention (Dao et al., 2022) computes exact attention in tiles that stay in fast on-chip memory, keeping running softmax statistics, so the full score matrix is never stored. The output is the same as standard attention up to rounding; the memory for scores grows linearly, not quadratically, with length. For one layer with 32 heads in 16 bits, the stored score matrix would be 0.27 GB at 2,048 tokens, 4.3 GB at 8,192 and 68.7 GB at 32,768: the quadratic growth that exact kernels avoid.

Speculative decoding

Decode is memory-bound, so verifying several tokens in one pass costs almost the same as producing one. A small draft model proposes k tokens cheaply; the large model scores all k in a single pass and accepts the longest prefix that it agrees with, using an acceptance rule that keeps the output distribution identical to the large model’s own. If each proposed token is accepted with probability a (assumed independent), the expected number of tokens produced per large-model pass is (1 − a^k+1)/(1 − a). With k = 4: 1.94 tokens at a = 0.5, 2.77 at a = 0.7 and 4.10 at a = 0.9. The gain is real only if the draft is much cheaper than the large model and agrees with it often; measure the acceptance rate on your own traffic.

Exactness is a property to verify

The mathematics: a serving budget

Let the model have P parameters stored at b_w bytes each, L layers, H_kv key-value heads of width d_h, cache precision b_c; serve B requests of n tokens; memory bandwidth BW, usable arithmetic rate F.

M ≤ memory

the feasibility condition; leave margin for activations, fragmentation and the framework (10 to 20%).

t_step

memory term dominates for small B and short n; the arithmetic term takes over at large B.

B

raising B raises total throughput until the arithmetic term or the memory limit binds.

Worked numbers (P = 7·10⁹, L = 32, d_h = 128, BW = 10¹² B/s, F = 4·10¹³ FLOP/s, 80 GB GPU, B = 16, n = 8,192):

Reading the table: grouped-query attention is what makes the workload fit; quantisation then doubles the speed because the traffic halves; batching multiplies the total throughput while costing each request some speed (compare B = 1 with B = 16 at the same precision: 66 versus 513 tokens/s in total, 66 versus 32 per request). Batch 64 no longer fits at 16 bits.

A worked example: the budget calculator

This is the calculator behind the table, plus the expected-speedup formula for speculative decoding. Copy it and replace the constants with your own model, hardware and workload.

P, L, HD = 7e9, 32, 128
BW, FU, MEM = 1e12, 4e13, 80e9                  # bytes/s, usable FLOP/s, GPU memory in bytes

def plan(weight_bytes, nkv, cache_bytes, batch, ctx):
    weights = P * weight_bytes
    cache = 2 * L * nkv * HD * ctx * batch * cache_bytes
    mem = weights + cache
    step = max((weights + cache) / BW, 2 * P * batch / FU)
    return {'gb': round(mem / 1e9, 1), 'fits': mem <= MEM, 'step_ms': round(step * 1000, 1),
            'per_req': round(1 / step, 1), 'total': round(batch / step)}

for args in [(2, 32, 2, 16, 8192), (2, 8, 2, 16, 8192), (1, 8, 1, 16, 8192),
             (0.5, 8, 1, 16, 8192), (2, 8, 2, 64, 8192), (2, 8, 2, 1, 8192)]:
    print(args, plan(*args))
# (2, 32, 2, 16, 8192) {'gb': 82.7, 'fits': False, 'step_ms': 82.7, 'per_req': 12.1, 'total': 193}
# (2, 8, 2, 16, 8192) {'gb': 31.2, 'fits': True, 'step_ms': 31.2, 'per_req': 32.1, 'total': 513}
# (1, 8, 1, 16, 8192) {'gb': 15.6, 'fits': True, 'step_ms': 15.6, 'per_req': 64.1, 'total': 1026}
# (0.5, 8, 1, 16, 8192) {'gb': 12.1, 'fits': True, 'step_ms': 12.1, 'per_req': 82.7, 'total': 1323}
# (2, 8, 2, 64, 8192) {'gb': 82.7, 'fits': False, 'step_ms': 82.7, 'per_req': 12.1, 'total': 774}
# (2, 8, 2, 1, 8192) {'gb': 15.1, 'fits': True, 'step_ms': 15.1, 'per_req': 66.3, 'total': 66}

def speculative_tokens(accept, k):
    return (1 - accept ** (k + 1)) / (1 - accept)

print([round(speculative_tokens(a, 4), 2) for a in (0.5, 0.7, 0.9)])   # [1.94, 2.77, 4.1]
print([round(32 * n * n * 2 / 1e9, 2) for n in (2048, 8192, 32768)])   # [0.27, 4.29, 68.72]

Rows that do not fit still print a step time; ignore it, as the configuration is infeasible. The fifth row is a reminder that the calculator’s time estimate does not know about the memory limit, so always read fits first.

A decision procedure

1. State the workload,Model, concurrency, context length, latency targets (time to first token, tokens per second per user), hardware.
2. Compute the baseline,Run the budget at the highest precision and the architecture you have. Note whether it fails on memory, speed or both.
3. Apply output-preserving levers first,Batching policy, exact attention kernels, prefix caching, speculative decoding. They cost no quality.
4. Then quality-affecting levers,Quantisation of weights and cache, fewer key-value heads (a model choice), smaller models. Evaluate each against your task set with confidence intervals (Module 9).
5. Re-measure on the real system,Budgets are estimates good to a factor of about two. Measure time to first token, tokens per second and percentiles at the target load.

Engineering implications and limitations

Compute before you build. A ten-line budget eliminates configurations that cannot work and tells you which lever matters; it is far cheaper than discovering an out-of-memory error in staging.

Leave headroom. Activations, fragmentation, the runtime and spikes need 10 to 20% of memory. A configuration that fits at 79 GB of 80 will fail in production.

Treat quantisation as a model change. Run the same evaluation set (Module 9) for the original and the quantised model, with intervals, including long-form and long-context cases where errors accumulate.

Log acceptance rate and cache use. For speculative decoding the acceptance rate on real traffic decides the speed-up; for the cache, occupancy and eviction rates tell you when capacity is the problem.

Do not stack levers blindly. Combinations interact: a 4-bit model with an 8-bit cache and a draft model may each be fine alone and poor together. Test the combination you ship.

Prefer simple wins. Shorter prompts, shorter outputs and a smaller model that passes your evaluation save more than most clever tricks.

For robots. On-board inference has a hard memory ceiling and a deadline. Start from the memory (weights plus cache for the longest context you allow) and the deadline, pick the smallest model that passes the task evaluation, and keep a safe fallback behaviour. The Physical AI and VLA tracks on this site cover the robot side of such systems.

Common misconceptions

Engineering takeaways

1. Memory M = P·b_w + 2·L·H_kv·d_h·n·B·b_c and step time = max(memory traffic / BW, 2·P·B / F) give a serving budget before any code is deployed.
2. In the 7B example, multi-head attention needs 82.7 GB for 16 × 8,192 tokens; grouped-query attention brings it to 31.2 GB and 8-bit weights and cache to 15.6 GB, doubling the decode speed.
3. Levers that preserve the output (batching, exact attention kernels, prefix caching, speculative decoding) come first; levers that change it need evaluation.
4. Speculative decoding yields (1 − a^(k+1))/(1 − a) tokens per large-model pass; it is worth it only with a cheap, well-aligned draft model.
5. Estimates are good to about a factor of two; confirm with measurements at the real load.

References and further reading

The budget numbers were produced by the code in this lesson with round illustrative hardware figures, not the specification of a particular device. The expected-token formula for speculative decoding assumes independent acceptance with a constant probability, as in the cited analysis.

• Dao, T. et al. (2022). FlashAttention: Fast and memory-efficient exact attention with IO-awareness. arXiv:2205.14135.
• Leviathan, Y., Kalman, M. and Matias, Y. (2023). Fast inference from transformers via speculative decoding. arXiv:2211.17192.
• Ainslie, J. et al. (2023). GQA: Training generalized multi-query transformer models from multi-head checkpoints. arXiv:2305.13245.
• Shazeer, N. et al. (2017). Outrageously large neural networks: The sparsely-gated mixture-of-experts layer. arXiv:1701.06538.
• Dao et al. (2022), FlashAttention (research paper)

Knowledge check

Three conceptual questions. Write an answer, then reveal the explanation. Your answers stay in your browser.

What’s next

Next, Module 11: Building LLM Applications, where the model becomes one component of a system with retrieval, tools and safeguards.

Related content

Nothing below is required to finish this lesson.

• Before
  LM10-T5: KV Caching, Prefill and Decode

  Cache keys and values to avoid recomputation and distinguish compute-bound prefill from memory-bound decode.

• Lab
  LM10-L1: KV-Cache Memory and Behaviour

  Compute KV-cache size and count the operations with and without a cache.

• Lab
  LM10-L2: Compare Attention Variants

  Count parameters, cache size and compute for multi-head, multi-query and grouped-query attention.