Where the time goes: bandwidth, batching and memory interview questions
10 questions with model answers
- junior 1
- mid 3
- senior 6
10 where the time goes: bandwidth, batching and memory questions asked in AI engineer, ML engineer and applied AI engineer interviews, each with a model answer you can say out loud in a minute or two: 1 junior, 3 mid-level and 6 senior. The lesson Where the Time Goes: Memory, Compute and Batching teaches the ideas behind them.Updated
Model answer
Compute-bound work is limited by arithmetic; bandwidth-bound work by how fast data moves from memory. Arithmetic intensity, operations per byte read, decides which. Decode is bandwidth-bound: each step reads every weight for one token.
Scroll sideways to see the whole diagram
Every forward pass has to bring the model's weights from memory to the arithmetic units: about 498 MB for GPT-2 small in 32-bit floats.
Text version
- Every forward pass has to bring the model's weights from memory to the arithmetic units: about 498 MB for GPT-2 small in 32-bit floats.
- A decode step does that whole read to produce one token for one sequence, plus that sequence's KV cache. About 2 operations per weight, 4 bytes per weight: half an operation per byte.
- Decode 8 sequences together and the same read serves all 8. Only the KV caches grow with the batch.
- Prefill pushes every prompt token through in one pass, so one read of the weights serves all 960.
- Operations per byte decides the limit. Few, and the chip waits on memory; many, and it waits on arithmetic. The next pages measure both on one laptop.
GPT-2 small (124,439,808 weights, float32); KV cache from the Module 01 formula at a 128-token context.
GPT-2 small in 4-byte floats reads 497.8 MB per token, about half an operation per byte.
Model answer
Measure both ceilings, a large memory copy and a large matrix multiply, and place each phase on a roofline. On our test laptop's CPU, a batch-1 decode step moved 507 MB in 9.1 ms: 55.7 GB/s, 46% of the copy ceiling, at about 0.5 operations per byte, far left of the ridge at 16. A 960-token prefill reached 51% of the matrix-multiply ceiling. For decode, more bandwidth or fewer bytes helps; more FLOP/s doesn't.
Model answer
Each step reads the weights once and uses them for every sequence; only each sequence's KV cache is extra. On our test laptop's Apple GPU, a step took 6.3 ms for one sequence and 9.3 ms for 8, so throughput rose from 159 to 860 tokens/s.
Limits: the longer step slows each user, and every conversation's KV cache needs memory.
Model answer
Weights plus KV cache per conversation times conversations must fit. Weights take about 26 GB. Each cache is 2 × 40 layers × 5,120 × 2,048 tokens × 2 bytes = 1.68 GB, so the remaining 54 GB holds 32, matching AI Engineering's example. Real servers fit fewer: activations, the framework and fragmentation take memory too.
The cache is linear in tokens and bytes per number, so cap conversation length or store the cache in fewer bytes.
Model answer
Continuous batching. Static batching waits for a full batch and dynamic caps the wait, but both finish a batch together, so a short reply waits for its longest neighbour. Continuous batching, Orca's iteration-level scheduling, returns each reply when done and refills its seat.
Scroll sideways to see the whole diagram
Two requests are decoded together: A needs 8 more tokens, B needs 64.
Text version
- Two requests are decoded together: A needs 8 more tokens, B needs 64.
- In request-level batching the batch finishes together. A's reply is ready after 8 steps but waits 56 more for B.
- In continuous batching A is returned as soon as its last token is generated.
- Its seat is filled by a waiting request at the very next step, so the batch stays full.
- Orca calls this iteration-level scheduling: the batch is decided again before every step.
After AI Engineering, Ch. 9, "Batching", and Orca (Yu et al., OSDI 2022). Illustrative lengths.
Simulated from our measured Apple GPU steps (200 requests, at most 8 at once), continuous gave 343 tokens/s and a 22 ms median first token; static batches gave 299 tokens/s and a 5.3 s median wait.
Model answer
Write the three costs down, then compare times.
# d model, f hidden (often 4d), T tokens, bf16
flops = 2 * (2 * T * d * f) # two matmuls
w_bytes = 2 * d * f * 2 # both weights
act_bytes = T * (2 * d + f) * 2 # in, out, hidden
# tensor parallel over n GPUs: one all-reduce of T*d
comm = 2 * (n - 1) / n * T * d * 2 # bytes per GPU
t_compute = flops / n / peak_flops
t_hbm = (w_bytes / n + act_bytes) / hbm_bw
t_net = comm / link_bw
Scroll sideways to see the whole diagram
Every forward pass has to bring the model's weights from memory to the arithmetic units: about 498 MB for GPT-2 small in 32-bit floats.
Text version
- Every forward pass has to bring the model's weights from memory to the arithmetic units: about 498 MB for GPT-2 small in 32-bit floats.
- A decode step does that whole read to produce one token for one sequence, plus that sequence's KV cache. About 2 operations per weight, 4 bytes per weight: half an operation per byte.
- Decode 8 sequences together and the same read serves all 8. Only the KV caches grow with the batch.
- Prefill pushes every prompt token through in one pass, so one read of the weights serves all 960.
- Operations per byte decides the limit. Few, and the chip waits on memory; many, and it waits on arithmetic. The next pages measure both on one laptop.
GPT-2 small (124,439,808 weights, float32); KV cache from the Module 01 formula at a 128-token context.
The largest time binds. With few tokens, as in decode, intensity is roughly T FLOPs per byte, below a GPU's FLOPs-to-bandwidth ratio, so HBM binds; big prefill batches turn compute-bound. Sharding divides compute and weights by n but not the all-reduce, so over slow inter-node links the network binds.
Model answer
A compiler turns the model graph into fewer, faster kernels. Kernel fusion matters most: chaining elementwise ops, norms or attention into one kernel keeps intermediates in registers or shared memory instead of round-tripping HBM, and most such ops are memory-bound. Others: constant folding, layout changes to what tensor cores want, lowering to vendor libraries, autotuning tile sizes per GPU, static memory planning, and CUDA graphs against launch overhead. Fitting the hardware means sizing tiles to SM resources and keeping matmuls in tensor-core shapes and precisions.
Model answer
In eager PyTorch, each op dispatches to a kernel, often from cuBLAS or cuDNN, and the CPU launches it asynchronously on a CUDA stream. The GPU runs it as thread blocks across its SMs, reading HBM. A compiler like torch.compile instead captures a graph, fuses ops and generates Triton kernels; CUDA graphs replay launches with less CPU overhead. Small ops end up launch- or memory-bound, large matmuls compute-bound on tensor cores.
Model answer
First prove it: a profiler shows each step waiting on the next batch, and GPU utilisation dips in a sawtooth. A quick test is to repeat one cached batch; if throughput jumps, the input pipeline is the bottleneck.
Then move work out of the hot path. Tokenise, filter and dedupe offline, writing packed, pre-tokenised sequences into large shards instead of many small files. At train time, stream shards with parallel workers, prefetching and pinned memory, shuffle by shard order plus a buffer, and make the loader resumable from a saved position.
Model answer
Compare each kernel's arithmetic intensity with the GPU's ratio of compute to bandwidth. Large matmuls are compute-bound; softmax, norms, elementwise ops and small-batch decode are memory-bound, so fuse them.
Naive attention writes the full sequence-by-sequence score matrix to HBM: quadratic memory and traffic. FlashAttention tiles it through on-chip SRAM instead.
MFU counts only the model's required FLOPs against peak; HFU also counts recomputation, so activation checkpointing raises HFU, not MFU.
Profile: timeline gaps for data loading, launch overhead and communication; per-kernel achieved bandwidth versus tensor-core utilization.