LLM Inference Optimization: KV Cache, Batching, and Speculative Decoding

Why Inference Is Slow by Default

Text generation is autoregressive: each new token depends on all previous tokens. The model generates one token, appends it to the sequence, then generates the next. This creates a fundamental bottleneck. Modern GPUs are designed for massively parallel computation, but autoregressive decoding uses only a fraction of available compute on each step because the model processes one token at a time (per request).

The result is that GPU utilization during text generation is typically 10% to 30% without optimization. The GPU spends most of its time waiting for memory transfers rather than computing. Every inference optimization technique addresses some aspect of this utilization gap, either by eliminating redundant memory operations, sharing GPU compute across multiple requests, or parallelizing the generation process itself.

KV Cache Management

The KV (key-value) cache is the single most important concept in LLM inference optimization. Understanding it explains why inference frameworks make the architectural choices they do.

During the attention mechanism, the model computes key and value vectors for every token in the context. On the first token, the model processes the entire prompt and computes KV vectors for all prompt tokens. On subsequent tokens, it only needs to compute the KV vector for the newest token, but it must attend to all previous KV vectors. Without caching, the model would recompute KV vectors for all previous tokens on every generation step, making generation time quadratic with sequence length.

The KV cache stores previously computed key and value vectors in GPU memory so they do not need to be recomputed. This turns generation from quadratic to linear with sequence length. For a 70B model with a 32K context, the KV cache consumes approximately 4GB of GPU memory. For 128K context, approximately 16GB. This is why larger context windows require significantly more VRAM even though the model weights stay the same size.

PagedAttention (used by vLLM) manages KV cache memory like virtual memory in an operating system. Instead of pre-allocating a fixed block for each request's maximum possible context length, it allocates small pages on demand. Memory is only consumed as tokens are generated, and pages from completed requests are immediately recycled. This eliminates the 50% to 80% memory waste typical of fixed-allocation systems and allows 2x to 4x more concurrent requests on the same hardware.

Prefix caching reuses KV cache across requests that share the same prompt prefix. If 100 requests share the same system prompt (common in production), the system prompt's KV vectors are computed once and shared. This is especially valuable for applications with long, consistent system prompts or RAG pipelines where retrieved context is often repeated.

Continuous Batching

Traditional batching waits for a fixed number of requests to accumulate, processes them together, and returns all results. This is wasteful because shorter requests finish before longer ones, but their GPU slots remain occupied until the entire batch completes.

Continuous batching (also called in-flight batching) adds new requests to the batch as soon as GPU resources become available from completed requests. When one request finishes generating, its KV cache memory is freed and a waiting request immediately takes its slot. The GPU never idles while requests are queued.

The throughput improvement from continuous batching is dramatic: 3x to 8x higher tokens per second compared to sequential processing, depending on the ratio of short to long requests. vLLM, SGLang, and TensorRT-LLM all implement continuous batching. Ollama does not, which is a primary reason it is not suitable for multi-user production serving.

The key configuration parameter is the maximum batch size (maximum number of concurrent sequences). Higher values increase throughput but also increase per-request latency because each generation step must process more sequences. For interactive applications where time-to-first-token matters, a batch size of 32 to 64 balances throughput and latency. For batch processing where total throughput is the priority, batch sizes of 128 to 256 maximize GPU utilization.

Speculative Decoding

Speculative decoding uses a small "draft" model to predict multiple future tokens, then verifies them in a single forward pass of the larger "target" model. The target model can verify N predicted tokens in approximately the same time it takes to generate a single token, because verification processes all tokens in parallel.

When the draft model's predictions are correct (which happens 60% to 80% of the time for a well-chosen draft model), the system effectively generates multiple tokens per step. When predictions are wrong, the system falls back to the first correct token and continues normally. The net effect is a 1.5x to 3x speedup in tokens per second with zero quality degradation, because the final output is identical to what the target model would produce without speculation.

The ideal draft model is much smaller than the target (typically 10x to 20x fewer parameters) but trained on similar data. For a Qwen 3 72B target, the Qwen 3 0.6B or 1.7B models serve as effective drafts. The draft model must run fast enough that its generation time is negligible compared to the target model's verification time. On the same GPU, a 1B draft model generates tokens 50x to 100x faster than a 72B target model, making the overhead minimal.

vLLM supports speculative decoding with the --speculative-model flag. The configuration requires both the target and draft models to fit in GPU memory simultaneously, so account for the draft model's memory (typically 1 to 2GB at Q4) when planning VRAM usage.

Prefill Optimization

Inference has two phases: prefill and decoding. Prefill processes the entire input prompt and populates the KV cache. Decoding generates tokens one at a time. These phases have different computational characteristics. Prefill is compute-bound (processing many tokens in parallel), while decoding is memory-bandwidth-bound (reading the full KV cache to generate each token).

Chunked prefill splits long prompts into chunks that can interleave with ongoing decoding requests. Without chunking, a long prompt (say 50K tokens) would monopolize the GPU for seconds, stalling all other requests in the batch. Chunked prefill processes the long prompt in segments, yielding the GPU to decoding requests between chunks. This dramatically improves time-to-first-token for queued requests at the cost of slightly longer prefill for the long-context request.

Prompt caching (different from KV prefix caching) stores the fully processed KV state for common prompt patterns. When a new request matches a cached prompt, the system skips prefill entirely and starts decoding immediately. Prompt caching is most effective for applications where the system prompt or context template is consistent across requests.

Quantization for Speed

Beyond reducing memory usage, quantization also increases inference speed. Lower-precision operations are faster on GPU hardware that supports them natively. INT4 operations on recent NVIDIA GPUs are approximately 2x faster than FP16 operations per element. FP8 operations on Hopper GPUs (H100, H200) are similarly accelerated with dedicated tensor cores.

The speed benefit compounds with the memory benefit: quantized models are smaller (more can fit in VRAM), and the reduced memory footprint means fewer memory transfers per operation (reducing the primary bottleneck). A Q4 model typically generates tokens 30% to 50% faster than the same model at FP16 on the same hardware, in addition to using 75% less memory.

FlashAttention

FlashAttention is a foundational optimization that restructures how the attention mechanism accesses GPU memory. Standard attention computes the full attention matrix, writes it to GPU high-bandwidth memory (HBM), then reads it back for the softmax operation. FlashAttention fuses these operations and keeps intermediate results in the GPU's much faster SRAM (on-chip cache), avoiding the expensive round-trip to HBM.

The performance impact is substantial. FlashAttention 2 delivers 2x to 4x speedup on the attention computation compared to standard attention, and reduces memory usage from quadratic to linear with sequence length. FlashAttention 3, optimized for Hopper GPUs, adds asynchronous execution that overlaps computation with memory transfers for an additional 1.5x speedup on H100 hardware.

In practice, you rarely configure FlashAttention directly. vLLM, SGLang, and modern versions of llama.cpp use FlashAttention automatically when available. On NVIDIA GPUs with compute capability 8.0+ (A100, RTX 3090, and newer), FlashAttention activates by default. On older GPUs, the framework falls back to standard attention. Verify FlashAttention is active by checking vLLM's startup logs for "Using FlashAttention backend" or similar messages.

Disaggregated Serving

Disaggregated serving splits the two phases of inference, prefill and decode, onto separate hardware. This emerging architecture recognizes that prefill is compute-bound (processing the full prompt in parallel) while decode is memory-bandwidth-bound (reading the KV cache for each new token). Running both phases on the same GPU means the hardware is optimized for neither.

In a disaggregated setup, high-compute GPUs handle prefill (processing incoming prompts and populating the KV cache), then transfer the populated KV cache to high-bandwidth GPUs that handle decode (generating tokens one by one). The prefill GPUs can be oversubscribed because prefill completes quickly, while the decode GPUs run at high utilization because they receive a steady stream of ready-to-decode requests.

The practical benefit is 30% to 50% cost reduction for large-scale serving compared to homogeneous GPU pools. Prefill machines can be older, compute-heavy GPUs (V100, A100) while decode machines benefit from newer, bandwidth-optimized hardware. Frameworks like Mooncake, DistServe, and recent vLLM experimental builds support disaggregated serving. For most deployments under 8 GPUs, the operational complexity outweighs the savings, but at data center scale this technique becomes increasingly standard.

Memory Bandwidth Optimization

During the decoding phase, the GPU reads the entire KV cache to compute attention for each new token. The bottleneck is memory bandwidth (how fast data can be read from GPU memory), not compute (how fast the GPU can process data). This is why consumer GPUs with high bandwidth (RTX 4090 at 1TB/s) sometimes generate tokens faster than older data center GPUs with more compute but lower bandwidth.

Techniques that reduce bandwidth pressure include: KV cache quantization (storing KV vectors at INT8 instead of FP16, halving cache memory reads), grouped-query attention (GQA, where multiple query heads share KV heads, reducing cache size proportionally), and multi-query attention (MQA, an extreme version of GQA with all query heads sharing a single KV head).

Most modern open source models use GQA by default (Llama 4, Qwen 3, Mistral), which provides a 4x to 8x reduction in KV cache size compared to multi-head attention. This design choice was made specifically to improve inference speed and reduce memory requirements.

Concrete numbers illustrate the impact: a Llama 4 70B model with GQA generates approximately 45 tokens per second on an A100 80GB at Q4. The same architecture with traditional multi-head attention would generate roughly 12 tokens per second at the same context length, because the KV cache would be 8x larger, consuming most of the available memory bandwidth. GQA is the reason 70B models are practical to serve on single GPUs at interactive speeds.

Practical Optimization Checklist

For single-user local inference (Ollama): Use Q4_K_M quantization for the best speed-quality balance. Set context length to only what you need (not the model's maximum). On NVIDIA GPUs, ensure CUDA drivers are up to date. On Apple Silicon, use Ollama 0.19+ for MLX acceleration.

For multi-user production (vLLM): Enable continuous batching (on by default). Use FP8 on Hopper GPUs or AWQ on older GPUs. Set max_model_len to your actual maximum context needs. Enable chunked prefill for workloads with long prompts. Consider speculative decoding if latency is critical. Add a response cache for repetitive queries.

For agent workloads (SGLang): Use SGLang's RadixAttention for prefix sharing across agent tool calls. Enable structured output optimization for JSON and schema-constrained generation. Cache common tool schemas in the prefix tree.

Benchmarking before and after: When applying optimizations, measure with realistic workloads, not synthetic benchmarks. Generate a test set of 100 to 500 prompts representative of your production traffic, with varying lengths and complexity. Run this test set through each configuration change and record tokens per second (throughput), time-to-first-token (TTFT), and P99 latency. Typical improvements from a naive baseline:

Adding Q4 quantization: 30% to 50% faster generation, 75% less VRAM. Adding continuous batching (vLLM vs sequential): 3x to 8x throughput increase under concurrent load. Adding speculative decoding: 1.5x to 3x faster per-request generation. Adding prefix caching with shared system prompts: 40% to 60% reduction in TTFT for requests after the first. Adding FlashAttention (if not already active): 2x faster attention computation. Adding response caching for repetitive queries: 30% to 60% reduction in GPU compute load. Combined, these take a single-request naive server and turn it into a system serving 10x to 20x more users on the same hardware.

Monitor throughput (tokens/second), latency (time-to-first-token and total generation time), GPU utilization, and KV cache occupancy. Optimization is iterative: measure, change one parameter, measure again. The biggest gains typically come from the first few optimizations (quantization, continuous batching), with diminishing returns on finer tuning.