Inference Optimization

During autoregressive generation, each new token attends to all previous tokens. Without caching, keys and values for all previous tokens are recomputed at every step, giving O(N^2) total compute. The KV cache stores K and V tensors for all previous tokens so only the new token's K/V are computed per step, reducing compute to O(N) over the full sequence. Memory cost: 2 × L × H_kv × d_head × seq_len × dtype_bytes. GQA and MQA directly reduce this memory footprint.

Speculative decoding (Leviathan et al., 2023) uses a small draft model to generate K candidate tokens cheaply, then the large target model verifies all K tokens in a single forward pass. Tokens accepted by the target model are kept; the first rejected token is resampled from the corrected distribution. Since the target model runs in parallel over K tokens (not sequentially), accepted tokens are essentially free. Speedup depends on acceptance rate; with a good draft model, 2-3x throughput improvement is typical without changing output distribution.

Traditional static batching waits until an entire batch finishes before adding new requests, leaving GPUs idle when fast sequences finish while long ones run. Continuous batching (Orca) slots new requests into the batch at token-generation boundaries — when a sequence finishes (generates EOS), the slot is immediately filled with the next waiting request. This keeps the GPU continuously active, maximizing throughput. vLLM, TGI, and TensorRT-LLM all implement continuous batching as the default serving strategy.

PagedAttention (vLLM) manages KV cache memory using fixed-size pages (blocks) analogous to virtual memory paging in OS. Each request's KV cache grows dynamically into non-contiguous physical memory blocks, eliminating internal fragmentation from pre-allocating maximum-length buffers. Blocks can be shared across requests with common prefixes (prefix caching) using copy-on-write. This allows vLLM to serve 3-4x more concurrent requests than systems pre-allocating contiguous max-length KV buffers.

INT8 PTQ converts float16/32 weights and activations to 8-bit integers after training, halving memory and enabling INT8 GEMM kernels (2x faster than FP16 on tensor cores). The main challenge is activation quantization: unlike weights (which can be statically calibrated), activations have dynamic ranges that vary per input and layer. LLM.int8() (Dettmers et al.) uses mixed-precision decomposition — detecting and preserving large-magnitude 'outlier' activation dimensions in FP16 while quantizing the rest to INT8.

from transformers import BitsAndBytesConfig
config = BitsAndBytesConfig(load_in_8bit=True)
model = AutoModelForCausalLM.from_pretrained(name, quantization_config=config)

GPTQ (Frantar et al., 2022) is a one-shot post-training quantization method using second-order (Hessian-based) information. It quantizes weights layer by layer: for each row of a weight matrix, it minimizes the increase in layer output error (using a calibration dataset) by compensating unquantized weights as each weight is quantized. This greedy compensation (Optimal Brain Quantization framework) produces INT4 models with quality nearly matching FP16 at 4-bit, far better than naive round-to-nearest quantization. GPTQ is available via AutoGPTQ.

AWQ (Lin et al., 2023) improves INT4 weight quantization by observing that not all weights are equally important — weights corresponding to large-activation input channels cause more quantization error. AWQ scales up salient weight channels (those with large input activations) before quantization to give them higher effective precision, then scales the activations down correspondingly. This is mathematically equivalent (scale-invariant) but places more quantization budget on important weights. AWQ achieves better quality than GPTQ at 4-bit with less calibration overhead.

GGUF (GPT-Generated Unified Format) is a binary file format for storing quantized LLM weights, supporting Q4_K_M, Q5_K_M, Q8_0, and other formats using k-quants (mix of 4-bit and 6-bit for different tensor groups). llama.cpp loads GGUF models and runs inference purely on CPU (or partially on GPU with offloading) using optimized SIMD matrix multiplication. This enables running 7B–70B models on consumer hardware. Mixed quantization applies higher bits to more sensitive layers (embeddings, last layer) and lower bits to bulk FFN layers.

# Run quantized model:
# ./llama-cli -m model-Q4_K_M.gguf -p 'Hello' -n 128

For a single large request (batch_size=1), the model is compute-bound for the first prefill token but then becomes memory-bandwidth-bound during generation (KV cache reads dominate). Tensor parallelism splits weight matrices across GPUs, effectively distributing memory bandwidth across all GPUs. N GPUs in tensor parallel mode reduces per-token generation time by ~N×, at the cost of all-reduce latency per layer. This makes tensor parallelism essential for latency-sensitive serving of large models (70B+) even for single requests.

Standard attention materializes the full N×N attention matrix in GPU HBM memory: memory = O(N^2). Flash Attention tiles Q, K, V into SRAM blocks and computes attention incrementally without ever writing the full matrix to HBM. Memory = O(N) for attention (only storing the final O output). For N=128K tokens with 8 heads: standard attention requires 128K^2 × 8 × 2 bytes ≈ 256 GB; Flash Attention fits in SRAM per block. This enables long-context inference that would OOM with standard attention.

TTFT (Time To First Token) is the latency from request arrival to the first generated token, dominated by the prefill step which processes the full input prompt in parallel. TPOT (Time Per Output Token) is the per-token generation latency during autoregressive decoding, dominated by memory bandwidth for reading model weights and KV cache. TTFT matters for user-perceived responsiveness (streaming). TPOT determines throughput. Prefill is compute-bound; generation is memory-bandwidth-bound — they have different optimization strategies.

Chunked prefill splits long prompt prefills into multiple smaller chunks processed over several steps, interleaving them with generation steps from active sequences. Without chunking, a long prefill stalls all ongoing generation requests, spiking their TPOT. By chunking, ongoing generation continues (reducing TPOT for active requests) while the long prefill progresses incrementally (slightly increasing its own TTFT). This is implemented in vLLM and TensorRT-LLM to balance latency fairness across concurrent requests.

Prefix caching (also called radix caching in vLLM) reuses KV cache from previous requests that share a common prefix (e.g., a long system prompt). When a new request starts with the same prefix, those KV values are retrieved from cache instead of recomputed, reducing TTFT to near-zero for the cached prefix. vLLM uses a radix tree to manage prefix cache; Anthropic's API and OpenAI's prompt caching features offer similar functionality. Cache hit rates are high for API deployments with shared system prompts.

Operator fusion merges multiple consecutive operations into a single GPU kernel, eliminating intermediate memory round-trips to HBM. For example, fusing LayerNorm + linear + GELU into one kernel avoids writing intermediate results to and reading from HBM between operations. CUDA kernel fusion is applied by TensorRT, torch.compile (TorchInductor), and XLA. For attention, Flash Attention is the canonical example of fusing QK^T × softmax × V into one kernel. Fusion can provide 2-5x speedup for memory-bandwidth-bound operations.

model = torch.compile(model, mode='max-autotune')  # enables fusion

QAT inserts fake quantization (quantize-then-dequantize) operations into the model during training, allowing the model to adapt its weights to minimize quantization error end-to-end. Unlike PTQ (which calibrates post-hoc), QAT minimizes task loss directly under quantization noise. QAT consistently outperforms PTQ at very low bit widths (INT4 and below) or for highly sensitive models, but requires full training compute. For INT8 with large models (7B+), PTQ (GPTQ, AWQ) is often preferred due to training cost.

Knowledge distillation trains a smaller student model to mimic a larger teacher model's output distributions (soft labels) rather than just hard class labels. The student learns not just correct answers but the teacher's confidence across all classes. Unlike quantization (which compresses an existing model), distillation produces a smaller, architecturally different model. It can achieve better quality-size tradeoffs for specific tasks. TinyBERT, DistilBERT, and specialized task distillation use this. LLM distillation often uses the teacher's logits or CoT traces as training targets.

# Distillation loss = CE(student_logits, labels) + alpha * KL(soft_teacher, soft_student)

Unstructured pruning zeros individual weight elements (produces sparse matrices) — high compression but requires sparse CUDA kernels for actual speedup. Structured pruning removes entire neurons, attention heads, or layers — the resulting model is a dense smaller model with no special sparse hardware required and immediate speedup on standard hardware. For example, removing 10 of 32 attention heads gives a dense model with 22 heads running ~30% faster. Structured pruning (SparseGPT, Wanda) is more practically impactful for inference acceleration.

The roofline model characterizes whether a workload is compute-bound or memory-bandwidth-bound by plotting achieved FLOPs/s vs arithmetic intensity (FLOPs/byte). Operations below the roofline's knee (low arithmetic intensity) are memory-bandwidth-bound; above it are compute-bound. LLM generation (batch=1): arithmetic intensity ≈ 1 FLOP/byte (reading weights once per token) — deeply memory-bound. Prefill (large batch): arithmetic intensity ≈ N (sequence length) — compute-bound. This determines whether to optimize for bandwidth (e.g., weight quantization) or compute (e.g., Flash Attention, FP8).

TensorRT-LLM is NVIDIA's production inference library for LLMs that applies: (1) weight-only INT4/INT8 quantization (GPTQ, AWQ support); (2) FP8 precision for H100 Tensor Cores; (3) custom fused attention kernels (xQA, Flash Attention variants); (4) continuous batching with inflight batching; (5) KV cache quantization (INT8); (6) multi-GPU tensor and pipeline parallelism; (7) speculative decoding with draft models. It is the primary backend for NVIDIA Triton Inference Server and targets maximum throughput on NVIDIA hardware.

Medusa (Cai et al., 2024) adds multiple draft heads (Medusa heads) directly to the LLM's final hidden state, each predicting tokens K steps ahead. During generation, all heads predict simultaneously in one forward pass. The main model verifies the full draft tree using a tree attention mask. Unlike speculative decoding with a separate draft model, Medusa requires no external model and adds only a small overhead (few MLP layers). Accepted paths can be multiple tokens; the best path is selected by the verification pass. Speedups of 2-3x are typical.

EAGLE (Li et al., 2024) trains a lightweight auto-regressive draft model that takes the target LLM's hidden states as input (rather than token embeddings) for drafting. Because it operates in the hidden state space, it requires far fewer parameters and training data than a standalone draft model. EAGLE-2 adds context-aware tree drafting where draft paths are expanded based on acceptance probability estimates. EAGLE achieves higher acceptance rates than Medusa-style heads and outperforms standard speculative decoding with small draft models on most tasks.

Static batching collects a fixed number of requests (batch size B) before processing — simple but wastes compute if requests arrive slowly. Dynamic batching groups requests within a time window (max_wait_time) to form batches that vary in size, improving GPU utilization when request rates fluctuate. Continuous/iteration-level batching (vLLM) goes further — it doesn't batch at request level but at token-generation steps, allowing different sequence lengths and new requests mid-generation. Dynamic batching is essential for production serving with variable load.

KV cache quantization compresses stored keys and values from FP16 to INT8 or even INT4, halving or quartering KV cache memory. Since KV values are read-only during decoding (no backward pass), quantization error only affects attention score computation, not weight updates. INT8 KV quantization (per-token or per-head scale factors) maintains acceptable quality for most tasks. This directly allows 2x more concurrent sequences or 2x longer contexts on the same GPU memory. TensorRT-LLM and vLLM support INT8/INT4 KV cache quantization.

Prefill processes the full input prompt in parallel (all tokens simultaneously) — it is compute-bound and benefits from large batch sizes, Flash Attention, and FP8/BF16 Tensor Core utilization. Decode processes one new token per step across the batch — it is memory-bandwidth-bound (limited by weight read bandwidth) and benefits from weight quantization (fewer bytes to read), tensor parallelism (more GPUs share bandwidth), and larger batch sizes (amortize weight reads). Different hardware configurations (more VRAM vs more compute) optimize for one phase vs the other.

Per-tensor quantization uses one scale factor for the entire weight matrix — simplest but worst quality (large range forces low precision for most values). Per-channel uses a scale per output channel — much better quality, still hardware-friendly (used by PyTorch PTQ for linear layers). Per-group (group quantization) uses a scale per small group of K elements (typically K=128) within each channel — best quality-compression tradeoff, used by GPTQ/AWQ for INT4. Smaller groups → better quality but more scale factor overhead in memory.

SmoothQuant (Xiao et al., 2022) migrates quantization difficulty from activations to weights by multiplying activation channels by a per-channel smoothing factor s (chosen based on calibration statistics) and dividing the corresponding weight channels by the same s. This is mathematically equivalent (linear layers absorb the scaling): Y = (Xs) × (W/s). After smoothing, both activations and weights have similar magnitudes and are easier to quantize to INT8. SmoothQuant enables INT8-INT8 GEMM (both activations and weights in INT8), maximizing throughput.

MLA (Chen et al., 2024, DeepSeek-V2) compresses KV cache by projecting keys and values into a low-dimensional latent space before caching. Instead of caching H×d_head K/V vectors per layer, only a small latent vector (dimension d_c << H×d_head) is cached. At attention time, K/V are reconstructed via up-projection. Combined with RoPE decoupling (caching RoPE-decoupled key separately), MLA achieves ~93% KV cache reduction vs standard MHA while maintaining near-MHA quality. DeepSeek-V3 uses MLA across all attention layers.

ONNX Runtime (ORT) is a cross-platform inference engine from Microsoft that executes models in ONNX format with optimized execution providers (CUDA, TensorRT, DirectML, CoreML, CPU MLAS). For LLMs, HuggingFace Optimum exports Transformers models to ONNX with operator-fused graphs. ORT applies layout optimization, fusion (e.g., multi-head attention fusion), and hardware-specific quantization. Optimum-NVIDIA targets maximum GPU throughput. ORT is particularly useful for edge deployment (mobile, Windows, embedded) where custom CUDA kernels are not available.

from optimum.onnxruntime import ORTModelForCausalLM
model = ORTModelForCausalLM.from_pretrained(name, export=True)

Self-speculation (e.g., Lookahead Decoding, Draft & Verify with the same model) avoids needing a separate draft model by reusing the target model itself for drafting. Methods include: (1) layer-skipping — run early layers only for draft; (2) Jacobi decoding — simultaneously solve a fixed-point equation for multiple tokens; (3) Lookahead — maintain a look-ahead window of speculative tokens from n-gram caches. These require only one model in memory and achieve 1.5-2x speedup with no quality degradation.

FP8 (E4M3 for weights/activations, E5M2 for gradient-like values) on NVIDIA H100 achieves the same 2000 TFLOPS throughput as INT8 GEMM but with wider dynamic range. FP8 is preferred over INT8 when: (1) activation distributions have wide ranges that INT8 cannot represent well; (2) model quality is sensitive to quantization noise; (3) per-token or per-tensor scaling suffices without per-channel calibration. INT8 is preferred when hardware support is wider (A100, A10) or when model is already well-optimized for INT8. Many production LLMs now deploy in FP8 on H100 clusters.

For a linear layer (weight matrix W: d_out × d_in), forward compute = 2 × d_out × d_in FLOPs; memory read = d_out × d_in × bytes (weights) + d_in × bytes (activation). Arithmetic intensity = FLOPs/bytes ≈ 2 × d_in / (d_in × bytes_per_weight) = 2/bytes_per_element per input token. For batch_size=1 and FP16: intensity ≈ 1 FLOP/byte. H100 bandwidth: 3.35 TB/s; compute: 2000 TFLOPS → crossover at ~600 tokens per batch. Below this batch size, inference is memory-bound — improving bandwidth (quantization) helps more than more compute.

Disaggregated serving (Zhong et al., 2024 — Splitwise, Mooncake) dedicates separate GPU pools for prefill and decode phases. Prefill GPUs (compute-optimized) handle prompt processing efficiently; decode GPUs (memory-bandwidth-optimized, e.g., with large VRAM) handle generation. KV cache is transferred between pools via PCIe/NVLink. This allows each pool to be sized and configured optimally, improving overall cluster utilization. It avoids resource contention where long prefills starve generation throughput on the same GPU.

Sparse attention reduces the O(N^2) attention cost by attending only to a structured subset of tokens. Main variants: (1) Sliding Window Attention (Mistral, Longformer) — each token attends to a local W-token window; (2) Dilated sparse attention — strided local windows at multiple scales; (3) Global+Local (BigBird, LongT5) — some tokens attend globally, rest locally; (4) Linear attention (Mamba-like) — reformulates attention to avoid O(N^2). Sliding window (GQA+SWA) is used in Mistral 7B to extend context while maintaining constant KV memory.

vLLM uses a priority queue with FCFS (first-come-first-served) as default, but supports priority-based scheduling. Under memory pressure, vLLM evicts (preempts) the KV cache of lower-priority sequences — either swapping to CPU RAM or simply re-prefilling when rescheduled. Block allocation uses the PagedAttention allocator; when fragmentation is high, compaction or preemption occurs. Anti-starvation mechanisms ensure low-priority requests are eventually served. This dynamic memory management enables higher concurrent request counts than static allocation systems.

For models too large for a single GPU, pipeline parallelism splits layers across GPUs: GPU 0 runs layers 0-11, GPU 1 runs 12-23, etc. At inference with a single request, only one GPU is active at a time (pipeline bubble = 1). For batched inference, micro-batching helps but the pipeline is less efficient than for training (fewer micro-batches). DeepSpeed-Inference and TensorRT-LLM support pipeline parallel inference. For latency-critical single-request serving, tensor parallelism (all GPUs active simultaneously) is preferred over pipeline parallelism.

max_new_tokens caps the number of tokens the model generates per request. Without it, slow requests that generate very long outputs monopolize GPU resources. Setting a reasonable max_new_tokens bound improves fairness and throughput by limiting worst-case generation time. Dynamic token budgeting (allocating tokens based on request priority or SLA) is used in production systems. Streaming (returning tokens as they're generated) allows users to read partial responses without waiting for max_new_tokens to complete, improving perceived latency.

outputs = model.generate(input_ids,
                         max_new_tokens=512,
                         do_sample=True)

LLM.int8() (Dettmers et al., 2022, bitsandbytes) uses mixed-precision matrix multiplication for INT8 inference: it identifies activation dimensions with very large magnitudes (outliers, typically >6σ from mean) present consistently across tokens. These outlier dimensions are extracted and computed in FP16; the remaining ~99.9% of dimensions are computed in INT8. The results are summed. This decomposition avoids the quality loss of naive INT8 quantization where outlier dimensions dominate the scale factor, compressing all other values into poor precision.

model = AutoModelForCausalLM.from_pretrained(name, load_in_8bit=True)

Standard KV cache pre-allocates a contiguous memory block of max_seq_len per request, wasting memory for shorter sequences (internal fragmentation) and limiting batch size. Paged KV Cache divides KV cache into fixed-size blocks (e.g., 16 tokens each). A request's KV cache grows by allocating new blocks on demand from a free-block pool — non-contiguous allocation is fine because block addresses are looked up in a block table. This eliminates internal fragmentation and allows efficient sharing of common prefix blocks across requests.

Mamba (Gu & Dao, 2023) is a selective state space model architecture that replaces attention with a recurrent SSM having input-dependent (selective) state transitions. At inference, Mamba runs in recurrent mode: one step per token with O(1) memory (only the SSM state, not a full KV cache) and O(1) compute per token. Unlike attention which has O(N) memory and O(N) compute per step, Mamba's constant-state inference is highly efficient for long sequences. Mamba-2 improved parallelism during training while maintaining efficient inference.

Throughput (tokens/s or requests/s per GPU) measures how much total work is done per unit time — maximized by large batches. Latency (TTFT + TPOT × output_length) measures user-perceived response time — minimized by small or no batching. These are in direct tension: batching improves throughput but increases individual request latency (requests wait for other batch members). Production serving uses dynamic batching with SLA constraints — maximize batch size while keeping 99th-percentile latency below a threshold, often configured via vLLM's --max-num-seqs and timeout parameters.

Activation offloading during inference moves unused KV cache blocks to CPU RAM (or NVMe) when GPU VRAM is full. When a sequence needs its KV cache for continued generation, the block is swapped back. This allows serving more concurrent sequences than GPU VRAM alone permits, at the cost of PCIe transfer latency (32 GB/s). It is beneficial when: requests have long input contexts, many concurrent users exist, and latency SLAs are relaxed enough to tolerate swap latency. vLLM supports KV swap via its CPU swap space configuration.

LoRA (Hu et al., 2021) freezes pretrained weights W and adds trainable low-rank decompositions: ΔW = BA, where B ∈ R^{d×r}, A ∈ R^{r×k}, r << min(d,k). Only A and B are updated during fine-tuning, reducing trainable parameters by ~100-1000x. At inference, LoRA weights can be merged back into W = W + αBA (lossless), adding zero inference overhead. QLoRA applies LoRA to 4-bit quantized base models, enabling fine-tuning of 70B+ models on a single GPU. Multiple LoRA adapters can be served and swapped dynamically (SLoRA).

from peft import get_peft_model, LoraConfig
config = LoraConfig(r=16, lora_alpha=32, target_modules=['q_proj','v_proj'])
model = get_peft_model(model, config)

Top-p sampling (Holtzman et al., 2020) samples from the smallest set of tokens whose cumulative probability exceeds p (e.g., p=0.9). It dynamically adjusts the number of candidate tokens per step — narrow distributions (sharp peaks) consider fewer tokens; broad distributions consider more. Top-p adds a sort-and-cumsum operation to the inference pipeline, which can be a CPU bottleneck in high-throughput systems. Inference engines like vLLM and TGI implement top-p as a batch CUDA kernel over logits. Top-p=1.0 degenerates to full vocabulary sampling.

output = model.generate(input_ids, do_sample=True,
                        top_p=0.9, temperature=0.7)

SLoRA (Sheng et al., 2023) enables serving thousands of LoRA fine-tuned adapters on a single GPU cluster without keeping all adapter weights in VRAM simultaneously. It stores LoRA A and B matrices in CPU memory, fetches them on-demand when requests for a specific adapter arrive, and batches together requests that use the same adapter. Custom CUDA kernels handle batched LoRA computation without needing to merge adapters into base weights. This enables fine-tuned model serving at scale without multiplicating GPU memory costs.

Eager inference runs each operation individually as Python calls, incurring Python overhead and lacking cross-operator optimization. Compiled inference via torch.compile captures the full computation graph and generates optimized fused CUDA kernels via TorchInductor. For LLM inference, compilation provides 15-40% speedup by fusing operations (attention projections, LayerNorm, activation), reducing kernel launch overhead, and enabling layout optimizations. The first call has compilation overhead; subsequent calls use the cached compiled kernel. torch.compile with mode='max-autotune' benchmarks kernel implementations.

model = torch.compile(model, mode='max-autotune', dynamic=True)

Token healing (Guidance, 2023) addresses a tokenization artifact: when a prompt ends mid-token boundary, the model generates next tokens starting from the split point, producing suboptimal completions (e.g., '192.168.' might be split, biasing generation). Token healing rolls back the last partial token, regenerates it fully, and then continues generation — ensuring the completion starts from a complete token. This is especially important for structured generation (code, JSON) where tokenization boundaries significantly affect output quality.

During prefill, the model processes all tokens of an input prompt in parallel using full causal attention — it is a dense matrix multiplication (batch_size × seq_len) that is compute-bound. During decode, only one new token is processed per step while reading the full KV cache — it is effectively a matrix-vector multiplication that is memory-bandwidth-bound. Batch prefill benefits from large batches and tensor core utilization (compute). Batch decode benefits from larger batch sizes to amortize weight reads across more output tokens, improving throughput.

RoPE encodes positions as rotation angles; out-of-distribution positions (beyond training length) cause attention degradation. Main scaling techniques: (1) Linear interpolation (Position Interpolation, Chen et al.) — divides position indices by a scale factor s, cramming longer sequences into the original range; (2) YaRN — applies different scaling factors to different RoPE frequency bands and adds a temperature correction; (3) LongRoPE — identifies non-uniform scaling factors per head via evolutionary search. All require continued training on long-context data after scaling for attention patterns to fully adapt.

Calibration determines per-layer quantization scale factors by running a representative set of inputs through the model and observing activation statistics (min, max, percentile distributions). For weight-only quantization (GPTQ, AWQ), calibration data affects scale quality — typically 128-512 samples from the target domain (e.g., C4, wikitext-2) are sufficient. For activation quantization (SmoothQuant, INT8 PTQ), calibration data must cover the expected activation range distribution. Poor calibration (out-of-domain data) leads to higher quantization error on target tasks.

During token-by-token decode, the model reads all weight tensors from GPU HBM for every token generated. For a 7B parameter FP16 model: 14 GB of weights read per token. H100 HBM bandwidth: 3.35 TB/s → maximum ~240 tokens/s single-batch. INT4 quantization (4 bits/weight) reduces weight size to ~3.5 GB → ~960 tokens/s theoretical maximum. In practice, quantization improves single-request decode throughput by 2-4x, making it the highest-impact optimization for latency-sensitive single-user inference on consumer hardware.