Inference Optimization | AI Engineering

Inference vs Training

Training is throughput-bound over many epochs. Inference is latency and cost per request. The dominant cost at serving time is autoregressive decoding: each generated token requires a forward pass through the full model, and long contexts amplify memory via the KV cache.

Optimization targets include time-to-first-token (TTFT), inter-token latency, tokens per second per dollar, and maximum concurrent users on fixed hardware.

KV Cache

During decoding, keys and values from prior tokens need not be recomputed. The KV cache stores them per layer and head, growing linearly with sequence length. Memory scales roughly as:

2 * num_layers * num_kv_heads * head_dim * seq_len * bytes_per_element

GQA and MQA reduce KV heads and shrink cache size. PagedAttention (vLLM) stores KV blocks in non-contiguous pages like OS virtual memory, reducing fragmentation and enabling larger batch concurrency on the same GPU.

# Launch vLLM with continuous batching and prefix caching
python -m vllm.entrypoints.openai.api_server \
  --model meta-llama/Llama-3.2-3B-Instruct \
  --dtype bfloat16 \
  --max-model-len 8192 \
  --enable-prefix-caching \
  --gpu-memory-utilization 0.90

Quantization

Lower precision weights and activations cut memory and increase throughput:

FP16 / BF16 - standard serving baselines
INT8 weight-only - good speedups with modest accuracy loss on many models
INT4 / GPTQ / AWQ / GGUF - aggressive compression for edge and cost-sensitive cloud
FP8 - emerging on H100-class hardware with kernel support

Quantization sensitivity varies by model size and task. Always eval after quantizing; math and JSON tasks are often more fragile than open chat.

Rule: Quantize weights for memory; consider KV cache quantization separately for very long contexts where cache dominates VRAM.

Batching and Scheduling

Continuous batching

Dynamic batching adds new requests to in-flight GPU work as others finish sequences, unlike static batching that waits for padding alignment. Essential for high GPU utilization in chat APIs.

Speculative decoding

A small draft model proposes multiple tokens; the large model verifies them in parallel. Accepted tokens advance faster than one-by-one decoding. Works best when draft and target distributions align.

Prefix caching

Shared system prompts and RAG contexts hash to reused KV blocks across requests, slashing TTFT for repeated prefixes.

Kernel and Framework Choices

FlashAttention-2/3, fused MLP kernels, and tensor parallel inference reduce memory bandwidth pressure. Serving frameworks (vLLM, TensorRT-LLM, TGI, llama.cpp) implement these optimizations out of the box compared to naive PyTorch generate loops.

Pick frameworks matching your hardware (CUDA vs ROCm vs Apple), need for OpenAI-compatible APIs, LoRA hot-swapping, and multi-modal inputs.

Application-Level Optimizations

Route easy queries to smaller models
Cache exact prompt completions where safe
Truncate, summarize, or compress context before model call
Stream tokens to improve perceived latency
Set max_tokens aggressively per use case
Batch offline jobs separately from interactive traffic

Profile end-to-end: embedding, retrieval, and reranking often dominate before the LLM is even called. Optimizing only the model server leaves money on the table.