Efficient AI & Systems

01Intuition

You are always waiting on one of two things: the GPU's math units, or data moving through memory. Knowing which one tells you which fix helps.

A GPU has two separate hard limits:

• Compute throughput. How many floating-point operations it can execute per second.
• Memory bandwidth. How many bytes per second it can move between HBM (the GPU's main memory) and the chip's actual arithmetic units.

Every piece of work you hand the GPU is bound by whichever limit it hits first. That gives you two regimes:

• Compute-bound. The data arrives long before the arithmetic on it finishes. More bandwidth wouldn't help. You need faster or fewer FLOPs.
• Memory-bound. The arithmetic finishes almost instantly, then the chip sits idle waiting for the next chunk of data. More compute units wouldn't help at all. You need to move less data, or move it faster.

The number that tells you which regime you're in is arithmetic intensity: how many FLOPs of work you get per byte of data moved. Lots of math per byte fetched means compute-bound. A little math per byte means memory-bound.

Training a large model with big batched matrix multiplications tends to sit on the compute-bound side. You reuse the same weights across many examples in the batch, so each byte you load does a lot of work.

Making a model smaller often means storing each weight in fewer bits. Here a million random weights are squeezed from 32 bits to 8 and the damage is measured.

import numpy as np

rng = np.random.default_rng(0)
w = rng.normal(0, 0.02, 1_000_000).astype(np.float32)   # fake weights

scale = np.abs(w).max() / 127                 # largest weight maps to 127
q = np.round(w / scale).astype(np.int8)       # 1 byte per weight
back = q.astype(np.float32) * scale

print(f"float32: {w.nbytes / 1e6:.1f} MB   int8: {q.nbytes / 1e6:.1f} MB")
print(f"mean absolute error per weight: {np.abs(w - back).mean():.6f}")
print(f"typical weight size:            {np.abs(w).mean():.6f}")

02Timeline

03Architecture: the modern serving stack, end to end

None of the pieces below are exotic on their own. The insight is how they stack. Follow one request through:

• It gets folded into a running batch instead of waiting in a queue.
• Its KV cache lives in non-contiguous pages instead of one reserved block.
• Its attention computation avoids ever writing the full score matrix to slow memory.
• Optionally, a small model proposes tokens that the big model verifies in bulk rather than generating them one at a time.

04Training-time efficiency: splitting the model across devices

A frontier model's parameters, gradients, and optimizer state don't fit on one GPU. Even if they did, training on one GPU would take years.

The four strategies below answer different questions about what to split. Real training runs at scale use several simultaneously, because each one alone runs into a wall.

• Data parallelism is the simplest. Replicate the entire model on every GPU, split the training batch across them, run forward and backward independently, then average the gradients before the optimizer step.

  It scales beautifully as long as the model itself fits in one GPU's memory. The communication cost is just one gradient sync per step. It does nothing to help when the model is too big for a single device, which is the situation for every modern frontier LLM.

• Tensor parallelism splits individual matrix multiplications across devices. Split a weight matrix by columns, and each GPU computes part of the output, exchanging partial results mid-layer.

  You need this when a single layer's weights are too large for one GPU, full stop. Data parallelism offers no way around it. The cost: GPUs now communicate within every single layer, which demands very fast interconnects. Devices doing tensor parallelism are usually on the same physical node.

• Pipeline parallelism splits layers across devices instead. GPU 1 holds layers 1–8, GPU 2 holds layers 9–16, and activations flow through like an assembly line.

  It needs far less bandwidth than tensor parallelism, since only activations cross stage boundaries. That makes it viable over slower links between nodes. The catch is the pipeline bubble: early stages sit idle waiting for the first microbatch to reach later stages, and later stages sit idle waiting for the first to start. You need many microbatches in flight to keep the idle fraction small.

• Expert parallelism is specific to mixture-of-experts models. Different experts, meaning separate feed-forward sub-networks, live on different devices. A router sends each token to only the handful of experts it needs, over the network, mid-forward-pass.

  This scales total parameter count far beyond what any single accelerator could hold, while keeping compute per token roughly fixed. The extra capacity costs communication and leaves compute cost alone.

Why combine them. At real scale, no single strategy is sufficient on its own. A frontier training run typically uses tensor parallelism within a node, pipeline parallelism across nodes, and data parallelism across many copies of that whole pipeline. Expert parallelism sometimes layers on top for MoE architectures.

Each strategy solves a different constraint. Stacking them is how you turn a few thousand GPUs into one coherent training job.

05ZeRO / FSDP: sharding what data parallelism used to duplicate

Plain data parallelism has a wasteful property. Every single replica holds a full copy of three things: the model's parameters, its gradients, and its optimizer state. For Adam, that optimizer state is itself roughly twice the size of the parameters, since it stores momentum and variance terms for every weight.

For a large model, that redundancy alone can make data parallelism impossible before you even get to compute.

ZeRO (Zero Redundancy Optimizer), and its widely used implementation in PyTorch's FSDP (Fully Sharded Data Parallel), attacks exactly that redundancy. Instead of every device holding a full copy, each of those three things is sharded: split into pieces across the data-parallel devices.

Each device permanently holds only its own shard. When a layer's full parameters are needed for a forward or backward computation, the devices briefly gather the missing pieces from each other, use them, then discard them again rather than keeping a persistent full copy.

What this buys you is significant. You can train models far larger than any single GPU's memory could hold, using ordinary data-parallel-style training, without hand-designing a tensor-parallel or pipeline-parallel split of the architecture.

It doesn't replace tensor or pipeline parallelism at the largest scales, and you still combine them, but it removes a huge amount of unnecessary memory duplication that would otherwise force you into more complex parallelism sooner than necessary. It's usually the first lever pulled before reaching for more architecturally invasive splitting strategies.

06Mixed precision: doing arithmetic in fewer bits

Numbers in deep learning don't need to be stored or computed at the same precision throughout. Using fewer bits is a direct, compounding win. Less memory to store and move, which matters enormously given the memory-bandwidth story above. And faster arithmetic, since modern GPUs execute lower-precision matrix multiplications at higher throughput than full FP32.

• FP32 (32-bit float) is the traditional default. High precision, high range, expensive in both memory and compute.
• FP16 (16-bit float) halves both. But its exponent range is narrow, so very large or very small values, common in gradients, can overflow or underflow to zero during training. That historically required loss scaling: multiply the loss by a constant before backprop, then divide it back out, keeping small gradient values representable.
• BF16 (bfloat16) also uses 16 bits but allocates them differently. It keeps FP32's wide exponent range and sacrifices mantissa precision instead. That wider range is why BF16 tends to be numerically stable for training without needing loss scaling at all. It can represent the same enormous dynamic range as FP32, just less precisely, and range turns out to matter far more than precision for training stability.
• FP8 pushes further still. Mostly used for inference, and increasingly for training the least precision-sensitive parts of the computation. More precision traded for more speed and memory savings.

In practice the precision is mixed and not uniform. Keep a master copy of the weights in FP32, or at least higher precision, for the optimizer's accumulation step where small updates need to register. Run the forward and backward passes themselves in FP16 or BF16.

You get the speed and memory benefit of low precision where it barely costs you anything, in the bulk of the matrix multiplications. And you protect the one part of training that is sensitive to precision loss: the accumulation of many small weight updates over time.

07Inference-time efficiency: where most interview questions land

Training happens once, while inference happens continuously, at whatever your product's traffic is, for the entire lifetime of the model in production. That's exactly why the serving stack gets so much systems attention, and why it's the part of this page you should know coldest.

KV cache growth (the direct callback to page 08)

Page 08 already covered the mechanism. During autoregressive generation, you cache every past token's key and value vectors so you don't recompute them from scratch at every new step.

The systems consequence is what matters here. Cache size grows linearly with (tokens generated) × (layers) × (heads) × (head dimension), per request. For long conversations or long contexts this routinely becomes the dominant consumer of GPU memory during serving, often larger than the model's own weights.

Every technique below responds to the fact that the cache is large, grows unpredictably per request, and has to be managed without wasting the GPU memory it lives in.

# Llama-3-8B: 32 layers, 8 KV heads after GQA, head dim 128, fp16.
L, KVH, HD, BY = 32, 8, 128, 2
W_GB = 8.03e9 * BY / 1e9              # fp16 weights: 16.1 GB, and fixed

kv_tok = L * KVH * HD * 2 * BY        # x2 because K and V are both cached
print("KV per token per sequence:", kv_tok / 1024, "KB")
print("weights", round(W_GB, 1), "GB on an 80 GB card\n")

print(f"{'ctx':>7} {'batch 1':>9} {'batch 32':>10}")
for ctx in (2048, 8192, 32768, 131072):
    one = kv_tok * ctx / 1e9
    b32 = one * 32
    tag = ""
    if b32 > W_GB:
        tag = "1 seq > weights" if one > W_GB else "b32 > weights"
    if b32 > 80 - W_GB:
        tag += "  OOM"
    print(f"{ctx:>7} {one:8.1f}G {b32:9.1f}G  {tag}")

mha = L * 32 * HD * 2 * BY * 32768 / 1e9   # the same model without GQA
print(f"\n32k ctx, 1 seq:  GQA {kv_tok*32768/1e9:.1f} GB"
      f"   MHA {mha:.1f} GB")

PagedAttention

Naive KV cache implementations reserve one contiguous block of memory per request, sized for the maximum sequence length it might reach. Most requests never get that long, so most of that reserved memory sits wasted. Worse: because it's reserved contiguously, that wasted memory can't even be reused for a different request in the meantime.

This is the fragmentation problem operating systems solved decades ago with virtual memory paging. PagedAttention, introduced with vLLM, applies the same idea to KV caches:

• Memory is divided into fixed-size blocks, or pages.
• A request's cache lives across possibly non-contiguous blocks.
• A block table tracks which physical blocks belong to which logical positions in each sequence.

That eliminates fragmentation and lets memory be allocated exactly as needed rather than reserved upfront. As a bonus, it makes sharing cache blocks across requests cheap, which is useful when many requests share a common prefix like a system prompt.

Continuous batching

Static batching waits for an entire fixed-size batch to finish generating before starting the next one. But requests finish at wildly different lengths: a one-sentence answer versus a long explanation. So the GPU spends much of its time processing a batch that's mostly finished requests waiting on one straggler, wasting the capacity freed up by everything that already completed.

Continuous batching, also called in-flight batching, removes that lockstep constraint. At every generation step, finished requests are evicted and new ones from the queue slot in immediately. Batch composition changes continuously rather than being fixed for a whole batch's lifetime.

This keeps the GPU close to fully utilized almost all the time, instead of oscillating between busy and mostly-idle. It's one of the largest practical throughput wins in modern serving, precisely because the earlier intuition holds: more real work per byte of weights moved is what pushes you toward compute-bound instead of memory-bound.

FlashAttention

Also introduced in page 08, and worth restating precisely here since this is where it belongs in depth.

FlashAttention computes the exact same attention output as the standard formula. It changes nothing about the math. What it changes is the memory-access pattern.

Standard implementations compute the full $n \times n$ attention score matrix and write it out to HBM, the GPU's relatively slow main memory, before applying softmax and multiplying by V.

FlashAttention instead processes attention in small tiles that fit in SRAM, the GPU's much faster and much smaller on-chip memory. During the backward pass it recomputes what it needs on the fly rather than storing the large intermediate matrix at all.

It's IO-aware exact attention: the same result with dramatically less data shuffled between fast and slow memory. Given that memory bandwidth is usually the actual bottleneck, that translates directly into real wall-clock speedups without any approximation.

Speculative decoding

Autoregressive generation is inherently sequential. Token $t{+}1$ needs token $t$ to exist first, so you can't parallelize across the generation steps of a single sequence the way you can across a batch.

Speculative decoding works around this indirectly:

• A small, cheap draft model proposes several tokens in a row.
• The large target model checks all of them in a single parallel forward pass. Verification parallelizes even though generation doesn't.
• It accepts the prefix of proposed tokens matching what it would have generated anyway, then rejects and corrects from the first mismatch onward.

08The one number that decides everything

• FLOPs total floating-point operations the computation performs. The numerator: how much math you're asking the chip to do.
• Bytes moved total bytes read from and written to memory to perform that computation. The denominator: how much data has to physically travel.
• Threshold every GPU has a ridge point, its peak compute throughput divided by its peak memory bandwidth, in FLOPs/byte. Above that ridge point you're compute-bound: the chip's math units are the limiting factor. Below it you're memory-bound: the chip finishes the math and then waits on data.
• Practical read operations with a lot of reuse per byte loaded, like the big batched matrix multiplications in most of training, sit above the ridge point. Operations that load a lot of data to do comparatively little math with it, like single-token autoregressive decoding steps, sit well below it. That's the formal version of "LLM inference is memory-bound."

# Llama-3-8B shapes, fp16 weights, decoding on one H100.
P = 8.03e9                      # parameters
L, KVH, HD = 32, 8, 128         # layers, KV heads (GQA), head dim
CTX = 4096                      # tokens already sitting in the cache

w_bytes = P * 2                 # every weight is read on every step
kv_tok = L * KVH * HD * 2 * 2   # K and V, fp16, per token, per seq
ridge = 989e12 / 3.35e12        # peak FLOP/s divided by peak bytes/s

print("weights", w_bytes / 1e9, "GB   KV per token", kv_tok / 1024, "KB")
print("ridge point", round(ridge, 1), "FLOPs/byte\n")

for B in (1, 8, 32, 256):
    moved = w_bytes + kv_tok * CTX * B   # weights plus the whole cache
    flops = 2 * P * B                    # 2 FLOPs per param per token
    ai = flops / moved
    side = "compute" if ai > ridge else "MEMORY"
    print(f"batch {B:4d}  {moved/1e9:6.1f} GB moved  AI {ai:6.2f}"
          f"  {side}-bound")

pre = 2 * P * 2048 / (w_bytes + kv_tok * 2048)   # a 2048-token prefill
print(f"\nprefill 2048   AI {pre:6.2f}  compute-bound")

Keep this conceptual and don't memorize exact ridge-point numbers for specific chips. An interviewer is checking whether you understand that this ratio, and neither raw FLOPs nor raw bandwidth in isolation, determines your bottleneck.

09Why memory bandwidth, not FLOPs, dominates LLM inference

Modern GPUs are built with tensor cores: specialized hardware units purpose-built for exactly one operation. A multiply-accumulate on small matrix tiles, extremely fast, in parallel, in a single instruction.

Most deep learning compute is matrix multiplication, which is why they exist. Every linear layer, every attention projection, every feed-forward block reduces to it. Building dedicated silicon for that one operation buys enormous throughput gains over general-purpose arithmetic units.

But a tensor core can only run as fast as data arrives at it. That constraint bites during generation.

During training, or during the prefill step of inference where the initial prompt is processed before generation starts, you run the full sequence through the model at once. A large batch of tokens gets multiplied against the same loaded weights, so each weight you fetch gets reused across many tokens' worth of arithmetic. High arithmetic intensity, comfortably compute-bound.

During autoregressive decoding, the picture inverts. Generating tokens one at a time, one per sequence per step, each forward pass pulls the entire set of model weights through memory. Those weights can be tens or hundreds of gigabytes, and the growing KV cache comes along too. All of that, to compute output for typically just one new token per sequence.

You pay the full memory cost of moving those weights and get only a sliver of arithmetic per byte moved. That's low arithmetic intensity, and it is why single-request autoregressive generation is memory-bandwidth-bound and not compute-bound. The tensor cores finish their multiply-accumulate work almost instantly, then sit idle waiting for the next chunk of weights to arrive from HBM.

That is why batching multiple requests together helps so dramatically. Batching reuses the same loaded weights across many requests' tokens in one pass. That raises the useful compute extracted per byte moved, pushing the operation back toward the compute-bound side of the ridge point, where the expensive tensor cores are earning their keep instead of idling.

10Tradeoffs, and what breaks

Three failure modes worth naming explicitly, because they show up in real deployments and in interview follow-ups:

• Quantization accuracy loss compounds. Drop weight precision, then also use speculative decoding, then also run a distilled model. Each is individually reasonable. Together they can stack into a noticeably degraded model, if nobody tracks the combined effect as well as each optimization in isolation.
• MoE routing imbalance is a hard serving problem. If the router sends a disproportionate share of tokens to a handful of popular experts, those devices become bottlenecks while others sit underused. The expert-parallel overhead of shipping tokens to remote experts can dominate when routing isn't reasonably balanced.
• Speculative decoding can backfire. If the draft model's proposals are frequently wrong, you pay for its forward passes and get little acceptance benefit back. It helps exactly when draft and target agree often, and can slow things down when they don't.

11Compression: quantization, distillation, pruning

Quantization reduces the numeric precision used to store and compute with weights, and sometimes activations. Represent weights in int8 or int4 instead of FP16. Two ways to get there:

• Post-training quantization (PTQ) takes an already-trained model and quantizes it afterward. A calibration step runs a small representative dataset through the model to pick good scaling factors for mapping continuous weight values into the reduced integer range.
• Quantization-aware training (QAT) instead simulates quantization during training itself, so weights adapt to tolerate the reduced precision. Generally more accurate at very low bit-widths, at the cost of needing to retrain or fine-tune rather than just post-process a finished checkpoint.

import torch
torch.manual_seed(0)

W = torch.randn(512, 512) * 0.02  # a plausible weight matrix
W[7] *= 40                        # one outlier channel -- real LLMs
                                  # have them, and they set the scale

def fake_quant(w, per_channel):
    if per_channel:
        amax = w.abs().amax(dim=1, keepdim=True)  # one scale per row
    else:
        amax = w.abs().max()                      # one scale for all
    s = amax / 127                                # symmetric int8
    q = torch.clamp((w / s).round(), -127, 127).to(torch.int8)
    return q, q.float() * s                       # dequantised back

for name, pc in (("per-tensor", False), ("per-channel", True)):
    q, deq = fake_quant(W, pc)
    err = (W - deq).abs()
    codes = q[0].unique().numel()   # distinct levels a normal row uses
    print(f"{name:12s} q {tuple(q.shape)} {q.dtype}")
    print(f"  max err {err.max():.5f}  mean err {err.mean():.6f}"
          f"  row 0 uses {codes:3d}/255 codes")

print("fp16", W.numel() * 2 / 1e6, "MB  ->  int8", W.numel() / 1e6, "MB")

Distillation trains a smaller "student" model to match the behavior of a larger "teacher", typically on the teacher's output distribution rather than only hard labels. The result is a smaller, cheaper model, not a compressed version of the same one.

Pruning and sparsity remove weights, or entire structural components like attention heads and neurons, that contribute little to the output. Two flavours:

• Unstructured sparsity removes individual scattered weights. Hard to accelerate on standard hardware without specialized kernels.
• Structured sparsity removes aligned blocks. Easier to get speedups from.

These three techniques complement each other and don't compete. A real deployment might run a distilled model that's also quantized and has some structural sparsity applied, stacking savings from three different angles at once.

12MoE serving, edge inference, and deployment tooling

Serving a mixture-of-experts model adds two problems that dense models don't have:

• Routing imbalance. Some experts get far more tokens than others. The router's choices are learned and not guaranteed to be even, so this is a pure load-balancing problem. An overloaded expert becomes a serving-time bottleneck no matter how fast the underlying hardware is.
• Expert-parallel communication overhead. When experts live on different devices, routed tokens physically travel over the network to reach their assigned expert, and results travel back. A dense model, with everything colocated, doesn't pay that cost.

Both are why MoE serving infrastructure is meaningfully more complex than dense-model serving, even though MoE's whole appeal is cheaper compute per token at a given parameter count.

Edge and mobile inference operates under a different constraint entirely. Not "how do I maximize throughput across a data-center's worth of GPUs" but "how do I fit and run anything useful inside a phone's tight memory, power, and thermal budget."

That usually means small models to begin with, aggressive quantization (often int4 or lower), and hardware-specific compiled kernels rather than a general-purpose serving framework. The entire optimization target shifts from throughput-at-scale to fitting inside a hard, fixed resource ceiling.

ONNX and TensorRT-style deployment tools sit at a different layer again: graph compilation and kernel fusion.

A trained model is a graph of many small operations. A matrix multiply, then an activation function, then a normalization, and so on. Naive execution runs each as a separate kernel launch, each with its own overhead and its own round trip to memory for intermediate results.

Compiling the graph lets the tool fuse adjacent operations into single kernels, skipping the memory round trip between them, and pick optimized kernel implementations for the specific target hardware. Similar in spirit to what FlashAttention does for attention, but applied generally across an entire model's computation graph.

13Profiling: find the bottleneck before you fix it

The first diagnostic step, always, before applying any optimization from this page: figure out whether the workload is compute-bound or memory-bound.

Profiling tools report GPU utilization, memory bandwidth utilization, and time spent in each kernel. Read the result two ways:

• Compute utilization high, bandwidth low. You're compute-bound. Look at reducing FLOPs (smaller model, fewer layers, algorithmic changes) or using faster numeric formats.
• Bandwidth near its ceiling, compute units idle. You're memory-bound. Look at reducing bytes moved (quantization, better batching, KV cache management) and leave alone the arithmetic that was never the bottleneck.

Applying a compute-side fix to a memory-bound problem, or the reverse, is a common and completely avoidable mistake. Profile first, optimize second.

14Build this

The roofline turns "memory-bound" from a phrase you repeat into a place on a graph. Your GPU has one, and you can measure it yourself in an afternoon.

Where this runs in production

Say you're standing up serving infrastructure for a chat product. Three decisions, in order:

• The KV cache budget per GPU sets your concurrency ceiling. Work out your model's per-token cache size: layers × heads × head_dim × 2 for K and V × precision in bytes.

  Multiply by your target max context length and expected concurrent conversations. That's how much GPU memory is spoken for by cache alone, before you've served a single request beyond what fits.

• Continuous batching determines how efficiently you use that concurrency. With static batching you leave a large fraction of GPU time on the table waiting for stragglers. With continuous batching you keep the batch close to full almost all the time.
• PagedAttention lets you hit the budget. No memory wasted on over-reserved, under-used contiguous blocks, and cheap cache sharing for a common system prompt across many concurrent users.

Together these determine realized throughput: tokens served per second per GPU. Combine that with your GPU's hourly cost and you get cost-per-token.

That's the number that decides whether the product is economically viable at the traffic you expect. Not whether any single request feels fast in isolation.

15Interview questions

16Go deeper

●Now write it yourself

Reading the derivation and being able to produce it are different skills. These are Deep-ML problems that exercise what this page covers — each one is checked against real test cases, not multiple choice.

Matched to this page from Deep-ML's catalogue of 1,380 problems. More at deep-ml.com, and Where to practise covers the other platforms and what each one trains.

17Related pages

Other techniques for this problem

A scoped slice of the full Technique Map — every technique this page covers, grouped by what it solves.