LLM serving and inference optimization

Prefill processes the entire input prompt in parallel in a single forward pass, making it compute-bound (high arithmetic intensity, GPU FLOPs saturated). Decode generates one token at a time, each step depending on the prior, making it memory-bandwidth-bound because it reads the full model weights and KV cache to produce just one token. They have opposite bottlenecks, so optimizations differ: prefill benefits from compute efficiency, decode from memory/bandwidth and batching.

The KV cache stores the key and value projections of all previously processed tokens for each attention layer so they need not be recomputed at every decode step. Its size is 2 (K and V) x num_layers x num_kv_heads x head_dim x seq_len x batch_size x dtype_bytes, so it scales linearly with both the number of tokens and the number of concurrent sequences. It often dominates GPU memory during serving and limits achievable batch size.

Naive serving pre-allocates a contiguous KV buffer sized to max sequence length per request, causing internal fragmentation and reservation waste. PagedAttention borrows OS virtual-memory paging: it splits the KV cache into fixed-size blocks stored non-contiguously, mapped via a block table per sequence. This nearly eliminates fragmentation (under ~4% waste), enables higher batch sizes, and allows block sharing/copy-on-write for prefix reuse and beam search.

Static batching waits to assemble a fixed batch and runs all requests until the longest finishes, wasting GPU cycles on padded/finished sequences. Continuous batching operates at the iteration (token) level: as soon as a sequence finishes, it is evicted and a waiting request is admitted into the batch slot mid-flight. This keeps the GPU saturated, dramatically raising throughput and reducing average latency under mixed-length workloads.

Speculative decoding uses a small, fast draft model to propose K candidate tokens, then the large target model verifies them in a single parallel forward pass. Tokens are accepted as long as they match the target's distribution (via a rejection-sampling criterion that preserves the exact target output distribution); on the first mismatch, the target's corrected token is used and drafting resumes. It accelerates decode because verifying K tokens costs roughly one target forward pass, exploiting the memory-bound nature of decode without changing output quality.

It uses modified rejection sampling: a drafted token x is accepted with probability min(1, p_target(x)/p_draft(x)); if rejected, a token is resampled from the normalized residual distribution proportional to max(0, p_target - p_draft). This procedure is mathematically equivalent to sampling directly from p_target, so the marginal distribution of emitted tokens is identical to the target model's. Thus it is a lossless speedup (for sampling), trading extra compute for latency without quality loss.

TTFT (Time To First Token) measures the delay from request arrival until the first output token, dominated by queueing plus prefill of the prompt. TPOT (Time Per Output Token, a.k.a. inter-token latency, ITL) measures the average time between subsequent tokens during decode. End-to-end latency is roughly TTFT + (num_output_tokens - 1) x TPOT; interactive/streaming UX cares about both, while throughput-oriented batch jobs care mainly about tokens/sec.

In MHA every query head has its own key and value heads, so KV cache scales with num_heads. GQA partitions query heads into G groups that each share one K/V head, cutting KV cache and decode memory bandwidth by a factor of num_heads/G with minimal quality loss. Multi-query attention (MQA) is the extreme case G=1 (a single shared K/V head), giving the smallest cache but slightly more quality degradation.

Standard attention materializes the full N x N score matrix in HBM, making it memory-bandwidth-bound and O(N^2) in memory. FlashAttention is an IO-aware, tiled kernel that computes softmax-weighted attention block by block in fast SRAM using online softmax, never writing the full attention matrix to HBM. This reduces HBM reads/writes, gives wall-clock speedups and linear memory in sequence length, while producing numerically equivalent results.

Tensor parallelism (TP) splits individual layers' weight matrices across GPUs, requiring an all-reduce per layer; it adds high-bandwidth communication every forward step so it is best within a single NVLink-connected node. Pipeline parallelism (PP) places different layer groups (stages) on different GPUs and passes activations between them; communication is lighter but it introduces pipeline bubbles and is better across nodes. Serving often combines TP within a node and PP across nodes.

During decode each step processes a single token (per sequence) but must read all model weights (and KV cache) from HBM, so arithmetic intensity is very low and the GPU's compute units idle waiting on memory. Because the same weights are reused, batching many concurrent sequences amortizes those weight reads across more tokens, increasing arithmetic intensity and throughput nearly for free until the KV cache or memory bandwidth saturates. This is why throughput scales strongly with batch size in the memory-bound regime.

Weight-only INT4 stores weights at 4 bits and dequantizes them on the fly to FP16 for matmul, cutting model memory ~4x and, crucially, reducing HBM weight-read traffic which speeds up the memory-bound decode phase. It primarily helps single-stream/low-batch latency and lets larger models fit on smaller GPUs; at very large batch sizes (compute-bound) the dequant overhead can erode gains. Quality impact is small when using calibration-aware methods like AWQ/GPTQ.

Prefix caching stores the KV cache computed for a shared prompt prefix (e.g., a long system prompt or few-shot examples) and reuses it across requests that share that prefix, skipping recomputation of prefill for the shared tokens. The benefit is largest when many requests share long common prefixes, dramatically reducing TTFT and prefill compute. Systems like vLLM implement it via hash-based block sharing with copy-on-write.

Throughput keeps rising while the workload stays in the memory-bound region (arithmetic intensity below the machine's ridge point), because larger batches raise intensity by reusing weight reads. Once batching pushes intensity past the ridge point into the compute-bound region, throughput plateaus at peak FLOPs and further batching mainly adds latency and KV memory pressure. The crossover depends on hardware FLOPs:bandwidth ratio, model size, and sequence/KV sizes.

Chunked prefill splits a long prompt's prefill into smaller token chunks that are interleaved with ongoing decode steps in the same batched iteration. Without it, a large prefill monopolizes a GPU iteration and stalls decode for other users, spiking their inter-token latency (a head-of-line blocking / latency-throughput tradeoff). By mixing prefill chunks with decodes, it smooths TPOT for active streams while still making prefill progress, balancing TTFT and TPOT.

Prefill is compute-bound and bursty while decode is memory-bandwidth-bound and long-running, so colocating them causes interference (prefill spikes hurt decode TPOT). Disaggregation runs prefill on one pool and decode on another, transferring the KV cache between them; each pool can be independently sized, batched, and parallelized for its bottleneck. This yields more predictable per-phase latency (TTFT vs TPOT SLOs) and better resource utilization at scale, at the cost of KV transfer overhead.

Larger batches amortize weight reads and raise tokens/sec (throughput) but each individual request waits longer in the batch and shares compute, increasing per-request latency (TTFT/TPOT). Small batches give low latency but underutilize the GPU and low throughput. Serving systems pick a batching policy (often continuous batching with a max batch/token budget and SLO-aware admission) to hit latency targets while maximizing throughput.

For one token, KV bytes = 2 (K and V) x num_layers x num_kv_heads x head_dim x bytes_per_element. With GQA, num_kv_heads < num_attention_heads reduces this. Multiply by sequence length and batch size for total. For example, a 32-layer model with num_kv_heads x head_dim = 1024 in FP16 uses 2 x 32 x 1024 x 2 = 131,072 bytes (~128 KB) per token.

kv_bytes_per_token = 2 * num_layers * num_kv_heads * head_dim * dtype_bytes
total = kv_bytes_per_token * seq_len * batch_size

Weight quantization shrinks model weights (memory + decode bandwidth), activation quantization compresses intermediate tensors to enable low-precision matmul (e.g., FP8/INT8 compute), and KV cache quantization stores K/V at lower precision. KV cache quantization most directly relaxes the batch-size/context-length ceiling because the KV cache is what grows with concurrency and sequence length and often dominates memory during serving. The others mainly cut model footprint or compute.

Decode launches many tiny kernels per step; CPU-side kernel launch overhead can dominate at low batch sizes. CUDA graphs capture the sequence of kernel launches once and replay them as a single graph, eliminating per-launch CPU overhead and reducing TPOT. The constraint is that the captured graph requires static shapes/tensor addresses, so serving systems capture graphs for a set of fixed batch sizes (padding to the nearest captured size).

FP8 (e.g., E4M3/E5M2) doubles tensor-core throughput and halves memory traffic versus FP16, accelerating both compute-bound prefill and memory-bound decode while fitting larger models. The key risk is its very limited dynamic range/precision, which can cause accuracy loss without careful per-tensor/per-channel scaling and calibration of sensitive layers. Mixed strategies (FP8 matmuls with higher-precision accumulation and selective layers) mitigate this.

A token budget is the maximum number of tokens the engine will process in one model iteration, summing decode tokens (one per active sequence) plus any prefill tokens admitted that step. The scheduler admits new requests and prefill chunks only up to this budget so each iteration's compute/latency stays bounded. It is the main knob for balancing TTFT (admitting prefills) against TPOT (keeping decode iterations fast).

The acceptance rate is the fraction of draft-proposed tokens the target model accepts; higher acceptance means more tokens emitted per target forward pass. Speedup grows with both acceptance rate and the draft length K, but is bounded by draft model cost and the target verification cost; a poorly aligned draft model (low acceptance) yields little or negative speedup. Optimal K and draft choice depend on the acceptance distribution and the draft/target cost ratio.

Classic speculative decoding requires a separate smaller draft model. Medusa instead adds multiple lightweight decoding heads on top of the target model itself to predict several future tokens in parallel, and EAGLE predicts at the feature level using the target's own hidden states; a tree of candidates is then verified by the target. This avoids hosting/aligning a second model and improves acceptance by sharing the target's representations, at the cost of extra trained heads.

Data parallelism runs independent full-model replicas behind a load balancer to increase total throughput/QPS; each replica must fit the whole model in its GPUs. Model parallelism (tensor/pipeline) shards a single model across GPUs to fit a model too large for one device or to lower single-request latency, with inter-GPU communication cost. Production stacks typically combine them: shard a large model within a replica, then run many replicas for throughput.

Sampling parameters shape the probability distribution but the dominant inference cost is the forward pass; the sampling step (sorting/filtering logits) is cheap relative to the model compute. However, decoding strategy affects total tokens generated: higher temperature or certain stop behaviors can lengthen outputs, and beam search multiplies KV cache and compute by the beam width. So cost scales mainly with number of tokens and concurrent sequences, not the sampling math itself.

Round-robin ignores per-replica state, so it can send a request to a replica whose KV cache or batch is saturated while sending another's shared-prefix request to a replica lacking that cached prefix. KV-cache-aware (prefix-affinity) routing sends requests with a shared prefix to the replica that already holds that prefix's KV blocks, boosting prefix-cache hit rate and cutting TTFT; least-loaded routing balances queue depth to avoid hotspots. Together they raise effective throughput and reduce tail latency.

Max model length caps the tokens (prompt + generation) per sequence, which directly bounds the worst-case KV cache per request and thus how many concurrent sequences fit in GPU memory. Capacity planning sizes the KV memory pool = (total GPU mem - weights - activations) / per-token KV, then divides by expected sequence length to estimate max concurrency. Overstating max len reserves memory and lowers achievable batch size; right-sizing it increases throughput.

Issuing a separate request or blocking call per token adds enormous overhead and breaks batching. Production serving uses a single streaming connection (SSE or gRPC/HTTP2 streaming) where the engine pushes tokens as they are generated from a continuously batched scheduler, decoupling generation from transport. This keeps the GPU batch full while delivering low TTFT and steady inter-token streaming to the client.

Constrained decoding restricts the next-token distribution to tokens allowed by a grammar/regex/JSON schema, guaranteeing structurally valid output by masking disallowed logits each step. Its main runtime cost is computing the allowed token mask per step (building/walking a finite-state machine over the vocabulary); precompiled FSMs or token-prefix automata amortize this so overhead is small. It improves reliability of structured outputs without changing the model.

Prefill attention computes a full L x L score matrix per layer, so its FLOPs grow ~O(L^2) with prompt length L (dominating long-prompt TTFT), even though the per-token feed-forward cost is linear. The KV cache, by contrast, only stores K and V per token, so it grows O(L) in memory. This is why long prompts blow up prefill latency faster than they blow up memory, motivating optimizations like chunked prefill and attention kernels.

It hurts when the GPU is already compute-bound, e.g., large batch sizes or long prefills, because the extra parallel verification work competes for FLOPs that are already saturated, and the draft model adds overhead with little memory-bandwidth slack to exploit. It also hurts with low acceptance rates (misaligned draft) where most drafted tokens are rejected, wasting the verification pass. It shines in latency-sensitive, low-batch, memory-bound single-stream decode.