InfraLens

Concept explanation: decoder-only block data flow

For input X: (B,S,D), the block projects Q=XWq, K=XWk, and V=XWv, splits them across H heads with Dh=D/H, and computes attention scores shaped (B,H,S,S). The output returns to (B,S,D) before the FFN and residual path.

Parameter count mental model

Self-attention projections are roughly 4D^2 per layer: Q, K, V, and output projection. The FFN is usually 2 * D * Dff; if Dff=4D, that is about 8D^2. Many model cards summarize the layer as mD^2, with embeddings adding V*D.

For a rough estimate, plug values such as D=4096,L=32,V=32000,m=8 into the calculator below, then compare the result with the memory ledger rather than treating parameter count as total memory.

Training memory vs inference memory

Do not mix these into one vague "memory" number. A useful memory ledger separates weights, gradients, optimizer states, activations, temporary buffers, communication buckets, and KV Cache.

Attention activations grow with sequence length, while KV Cache grows with retained context during serving. That difference is why training and inference memory fixes are not interchangeable.

The numbers are intentionally rough. They are good for interview reasoning and sanity checks, not for replacing profiler traces or framework memory summaries.

Memory hierarchy mental model

Start from where bytes live and how often they move. A kernel can have plenty of math available and still be slow if it repeatedly reads from HBM or writes large intermediates.

Register

Per-thread storage, fastest but limited; excessive use can reduce occupancy.

Shared memory

Block-local scratch space used for tiling, reuse, reductions, and avoiding repeated HBM reads.

L2 / cache

On-GPU cache layer that helps reuse across memory operations but is still much slower than registers.

HBM / global memory

Large device memory with high bandwidth, but expensive enough that layout and access pattern dominate many kernels.

Key mechanisms

FlashAttention mental model

Naive attention materializes the (B,H,S,S) score matrix. FlashAttention-style kernels compute attention in blocks with online softmax so the full score matrix does not need to be written to HBM.

Data parallel baseline

DDP keeps a full model replica on each rank, feeds each rank a different data shard, and all-reduces gradient buckets during backward so all replicas apply the same update.

DDP code path to recognize

The user code looks local, but DDP installs autograd hooks. When gradients are produced, those hooks bucket the gradients and launch collectives through the process group.

dist.init_process_group(backend="nccl")
local_rank = int(os.environ["LOCAL_RANK"])
torch.cuda.set_device(local_rank)

model = DDP(build_model().cuda(), device_ids=[local_rank])
sampler = DistributedSampler(train_dataset, shuffle=True)

for epoch in range(num_epochs):
    sampler.set_epoch(epoch)
    for batch in loader:
        loss = model(**move_to_cuda(batch)).loss
        loss.backward()       # autograd hooks trigger gradient bucket all-reduce
        optimizer.step()
        optimizer.zero_grad(set_to_none=True)

Sharding training state

ZeRO and FSDP reduce per-GPU memory by sharding optimizer states, gradients, and parameters at different levels. That saves memory but adds all-gather, reduce-scatter, checkpoint, and runtime peak-memory considerations.

Parallelism families

Collective input/output intuition

AllReduce

Every rank contributes a tensor and every rank receives the reduced result. DDP gradient sync is the common example.

AllGather

Each rank starts with one shard and receives all shards, often used before a sharded parameter is needed.

ReduceScatter

Ranks reduce inputs and scatter the reduced shards, commonly paired with all-gather in sharded training.

AllToAll

Each rank sends a different slice to every other rank. Expert parallelism and token routing often pay this cost.

nvidia-smi topo -m How to read it

Topology labels tell you whether traffic stays on fast GPU links or crosses PCIe, host bridges, NUMA boundaries, or the system interconnect. Use them to place ranks, not just to describe hardware.

Diagnosing the bottleneck

First decide whether the wait is compute, memory, communication, or orchestration. Then tie the diagnosis to a trace, counter, or controlled experiment.

Tool split

Roofline mental model

Roofline reasoning asks whether arithmetic intensity is high enough to use available compute. If it is not, optimizing memory movement may matter more than adding more math throughput.

torchrun --nproc-per-node=8 train.py --config config.yaml
nsys profile --trace=cuda,nvtx,osrt --output=step_report python train.py
ncu --set full --target-processes all python kernel_bench.py
# Serving comparison: keep prompt/output length distributions fixed before comparing engines.

Serving loop mental model

A request first enters a scheduler or queue, then prefill computes the prompt and writes KV Cache. Decode repeatedly reads that cache, appends one token step, and streams output while the scheduler admits more work.

Prefill vs Decode

KV Cache memory scales with batch * layers * kv_heads * head_dim * sequence_length * 2(K,V) * bytes. That is why request length distribution and cache management are first-order serving concerns.

request arrives
  -> scheduler / queue
  -> prefill builds KV Cache and determines TTFT
  -> decode loop reads KV Cache one step at a time
  -> stream output tokens and track TPOT, P95/P99

PagedAttention, Continuous batching and Prefix cache

PagedAttention-style cache management reduces waste from variable-length requests. Continuous batching improves GPU utilization by mixing requests at different decode positions, but it can increase queueing complexity and tail-latency risk.

Quantization and speculative decoding

Quantization can reduce memory and bandwidth pressure when kernels support the target dtype and accuracy remains acceptable. Speculative decoding helps only when acceptance rate, draft-model cost, and scheduler overhead make the extra machinery worthwhile.

The most common mistake is naming a technique without naming the state object or resource it changes.

Full practice page: AI Infra Interview Practice

• What does AI infrastructure optimize beyond model accuracy?
• How would you debug a slow ML training job?
• What signals belong in ML observability?
• What makes LLM infrastructure different from ordinary model serving?

Concept connections: RL runtime · multimodal runtime · systems runtime · performance modeling

Read the related handbook section first, then use the lab page and starter file to connect the concept to concrete variables, shapes, APIs, and interview-ready explanations.

Internal pages

External references