Interview: Inference Optimization

For autoregressive generation, the GPU must load the entire model's weights (and the growing KV cache) from memory for every single token generated, but only does a relatively small amount of compute per token (batch size of one token per step, per sequence). This means time is dominated by moving data from GPU memory to compute units, not by the FLOPs themselves - which is why techniques that reduce memory movement (quantization, better batching, smaller KV cache) give outsized speedups.

During autoregressive generation, each new token's attention needs the Key and Value vectors of all previous tokens. Without caching, you'd recompute K/V for the entire sequence at every new token (quadratic cost). The KV cache stores previously computed K/V vectors so each new token only computes its own K/V and reuses the rest, making generation linear instead of quadratic in sequence length - at the cost of memory that grows with sequence length and batch size.

Quantization represents model weights (and sometimes activations) in lower precision (8-bit or 4-bit integers instead of 16/32-bit floats), shrinking memory footprint and increasing memory-bandwidth-bound throughput, with some quality loss. GPTQ quantizes layer-by-layer to minimize output error introduced by quantization using a calibration dataset. AWQ (Activation-aware Weight Quantization) identifies and preserves the small subset of weights most important for activations at full precision, quantizing the rest more aggressively, generally giving better quality than naive quantization at the same bit-width.

Static batching waits for a fixed batch of requests, processes them together, and only starts the next batch once every sequence in the current one has finished - wasting GPU cycles on requests that finished early while waiting for the longest one. Continuous batching (used by vLLM, TGI) dynamically adds new requests into the batch and removes completed ones at each generation step, keeping the GPU consistently full and dramatically improving throughput for real-world traffic with variable-length outputs.

Naively, KV caches are allocated as large contiguous blocks sized for the maximum possible sequence length, wasting significant GPU memory on unused space (internal fragmentation) and making memory management for many concurrent sequences inefficient. PagedAttention borrows the idea of OS virtual memory paging - the KV cache is split into small fixed-size blocks that are allocated on demand and can be non-contiguous, dramatically reducing wasted memory and allowing far higher batch sizes/throughput for the same GPU memory.

A small, fast "draft" model generates several candidate tokens ahead, and the large target model verifies them all in a single forward pass (which is cheap because verification is parallelizable, unlike generation). Correct draft tokens are accepted for free; the first incorrect one is corrected by the target model. Because verification is much cheaper than generating token-by-token with the large model, this can significantly speed up generation with mathematically identical output distribution to using the large model alone.

Distillation trains a smaller "student" model to mimic a larger "teacher" model's outputs (or output distributions), transferring capability into a genuinely smaller architecture with fewer parameters. Quantization keeps the same architecture and parameter count but represents each weight with fewer bits. They're complementary - you can quantize a distilled model for further gains.

Prefill is the initial pass that processes the entire input prompt at once (compute-bound, highly parallel, good GPU utilization). Decode is the token-by-token generation phase (memory-bandwidth-bound, low utilization per step). Because these phases have very different resource profiles, some high-throughput serving architectures separate them onto different GPU pools (prefill-decode disaggregation) so each can be optimized/scaled independently rather than sharing hardware inefficiently.

• Streaming - send tokens to the user as they're generated instead of waiting for the full response, improving perceived latency (time-to-first-token) even if total generation time is unchanged.
• Prompt caching - cache the KV state of a static/shared prefix (e.g. system prompt) across requests so it isn't recomputed every time.
• Smaller/distilled models for latency-sensitive sub-tasks, reserving the largest model for what actually needs it.
• Parallelizing independent sub-calls instead of chaining them sequentially where possible.

FlashAttention is an IO-aware attention implementation that tiles computation to reduce reads/writes between GPU HBM and SRAM. It produces the same output as standard attention but faster and with lower memory, enabling longer sequences and larger batches. FlashAttention-2/3 extend this for modern GPUs and is standard in training and inference stacks.

Storing KV cache in lower precision (INT8, FP8) cuts memory during decode, allowing more concurrent users per GPU. Quality impact is usually small compared to weight quantization. It pairs well with GQA/MQA since the cache is already smaller.