LLM Inference Optimization: Faster, Cheaper, Further

LLM development and LLM operations are two entirely different problems.

Running a GPT-4-class model requires hundreds of gigabytes of GPU memory, and generating a single token takes tens of milliseconds. One concurrent user is manageable, but when simultaneous requests scale to hundreds or thousands, the situation changes entirely. GPU costs explode, and response latency tanks.

Inference optimization is the collective term for techniques that close this gap. The goal is to handle more requests on the same hardware, respond faster, and reduce power consumption in the process. Research shows that combining quantization, batching, caching, and speculative decoding can reduce energy consumption by up to 73% compared to an unoptimized baseline.

This post gives a broad overview of the core inference optimization techniques, organized into three major layers.

---

1. Understanding the Bottleneck: Prefill vs. Decode

Before discussing optimization, we need to understand where LLM inference actually slows down.

Token generation in an LLM happens in two phases:

During the decode phase, every single step requires reading the entire model's weights from GPU memory — just to produce one token. Loading hundreds of gigabytes of weights to generate a single token results in extremely low hardware utilization. Most inference optimization techniques ultimately target this inefficiency.

---

2. Model-Level Optimization: Making the Model Lighter

2.1 Quantization

The most widely used technique. It converts model weights from 32-bit floating point (FP32) to a lower precision format.

FP32 (4 bytes) → FP16 (2 bytes) → INT8 (1 byte) → INT4 (0.5 bytes)

Why does it work? Neural network weights are typically concentrated around small values, so representing them at lower bit widths doesn't significantly shift the model's output distribution.

Key methods:

• GPTQ: Compresses weights layer-by-layer to INT4, using Hessian information to correct for quantization error. Has become the standard for post-training quantization (PTQ).
• AWQ (Activation-aware Weight Quantization): Keeps channels with large activations (i.e., salient weights) at higher precision while quantizing the rest, minimizing quality loss.
• W4A4KV4: An aggressive strategy that quantizes weights, activations, and the KV cache all to INT4. Results show negligible quality degradation compared to INT8, even at Llama 3.1 70B/405B scale.

Impact: 2–8× reduction in model memory usage, with corresponding throughput gains.

2.2 Pruning

Removes neurons, layers, or attention heads that contribute little to model output.

• Unstructured Pruning: Sets individual weights to zero. Increases sparsity, but achieving real-world speedups requires specialized sparse operation kernels.
• Structured Pruning: Removes entire attention heads or layers. Because the model architecture itself changes, speedups are immediately available on standard hardware.

Pruning is most effective when combined with knowledge distillation.

2.3 Knowledge Distillation

Trains a small student model to mimic the output distribution and internal representations of a large teacher model.

Teacher (70B) ──soft labels──▶ Student (7B)
                    └──hidden states──▶

A student trained on the teacher's confidence distribution — rather than hard labels alone — produces significantly higher quality. The cases where small models distilled from GPT-3.5 approach GPT-4 performance on certain benchmarks illustrate this well.

Compounding effect: Applying pruning → quantization → knowledge distillation in sequence can yield efficiency gains of 10× or more over the baseline.

---

3. Attention-Level Optimization: Efficient Core Operations

3.1 FlashAttention

Standard attention requires O(N²) memory. As sequence length grows, the data movement between high-bandwidth GPU memory (HBM) and fast on-chip SRAM becomes the bottleneck.

FlashAttention solves this with tiling. It splits the Q, K, and V matrices into small blocks, processes them in SRAM, and writes results back to HBM — dramatically reducing the number of HBM accesses.

Standard attention: HBM → SRAM → HBM round trips O(N²) times
FlashAttention:     tile-based processing, HBM accesses reduced to O(N)

This brings memory usage down to O(N) as well, while producing numerically identical results at significantly higher speed. All major serving frameworks — vLLM, TensorRT-LLM, and others — now use FlashAttention by default.

3.2 KV Cache

During decode, recomputing the Key and Value tensors for every previous token at every step is enormously wasteful. The KV cache stores these values in memory and reuses them.

The problem is that the KV cache consumes a lot of memory. As sequences get longer and concurrent requests increase, KV cache size grows explosively.

Approaches to managing this:

---

4. Serving-Level Optimization: Making the Whole System Efficient

4.1 Continuous Batching

With traditional static batching, all GPUs must wait until the longest sequence in the batch finishes. Requests with short responses complete early, but the GPU sits idle waiting for the rest.

Continuous batching fixes this. As soon as a request completes, a new one is inserted into the batch immediately, keeping the GPU busy at all times.

Static batching:     [req1════════] [req2══] [req3══════]  ← idle gaps
                                             ↑ GPU waits for req1 even after req2,3 finish

Continuous batching: [req1][req2][req3][req4][req5][req6]  ← GPU always busy

According to Anyscale benchmarks, continuous batching alone has produced up to 23× throughput improvements in some cases.

4.2 PagedAttention

Introduced by vLLM, PagedAttention applies the concept of virtual memory from operating systems to KV cache management.

The traditional approach pre-allocates a contiguous memory block sized to the maximum sequence length for each request. Short sequences waste most of that space (internal fragmentation), and when no sufficiently large contiguous block is available, longer sequences can't be served.

PagedAttention manages the KV cache in fixed-size, non-contiguous blocks (pages). A logical block table maps to physical blocks, which nearly eliminates fragmentation and substantially improves GPU memory utilization.

Logical KV blocks:    [block0] → [block1] → [block2]
                         ↓           ↓           ↓
Physical GPU memory: [pageA]    [pageC]    [pageB]   ← non-contiguous is fine

vLLM's combination of PagedAttention and continuous batching delivers 2–4× throughput improvement over the previous FasterTransformer baseline.

4.3 Speculative Decoding

The biggest waste in the decode phase is running a large model once per step to produce a single token. Speculative decoding inverts this structure.

1. A small draft model quickly generates several tokens in sequence.
2. A large target model verifies all draft tokens in parallel, in a single forward pass.
3. Tokens that pass verification are accepted; at the first failure, the target model generates the correct token and the draft is discarded from that point.

Draft model:   [tokenA] [tokenB] [tokenC] [tokenD]   (4 tokens generated quickly)
Target model:    ✓        ✓        ✗       (skip)    (parallel verification, fail at C)
Final output:  [tokenA] [tokenB] [correctedC]         (3 tokens from 1 target forward pass)

Critically, the output distribution is mathematically identical to running the large model alone — quality is preserved, only speed improves. Speedups of 1.5–3.5× are commonly reported.

4.4 Distributed Inference

When a model doesn't fit on a single GPU, it must be distributed across multiple GPUs.

• Tensor Parallelism: Splits the weight matrices of each layer across GPUs, computing them simultaneously. The most effective approach for reducing latency, but requires fast interconnects between GPUs (NVLink, etc.).
• Pipeline Parallelism: Groups layers into stages and assigns each stage to a different GPU. Useful for handling memory constraints, but inter-GPU dependencies can cause idle bubbles.
• Data Parallelism: Replicates the model across GPUs and distributes requests across replicas. The simplest and most effective approach for scaling throughput.

How to choose: If the bottleneck is request volume, use data parallelism. If it's GPU memory, use pipeline parallelism. If it's latency or compute, use tensor parallelism.

---

5. Architecture Level: Mixture-of-Experts (MoE)

MoE is more of an architectural choice than an optimization technique, but it has a decisive impact on inference efficiency.

While a dense model activates all parameters for every token, MoE selects and activates only a subset of expert FFNs — out of dozens or hundreds — per token. GPT-4, Mixtral, and DeepSeek-V3, among others, use this architecture.

Dense:  input → [all parameters activated] → output   (all 70B params used)
MoE:    input → [router] → [2–8 experts activated] → output  (far fewer active params)

This produces tokens at the same quality level with significantly less computation. The downside: when experts are distributed across multiple GPUs, routing introduces communication overhead that becomes a new bottleneck. Research like XShare attempts to mitigate this by sharing experts across requests within a batch.

---

6. Combining Techniques

Each technique is meaningful on its own, but the synergy from combining them is where the real gains come from.

---

7. A Production Serving Stack: vLLM

Rather than implementing these techniques from scratch, vLLM integrates PagedAttention, continuous batching, FlashAttention, speculative decoding, distributed inference, and prefix caching in a single framework. It exposes an OpenAI-compatible API out of the box, making production adoption straightforward.

# Example: launching a vLLM server
vllm serve meta-llama/Llama-3.1-8B-Instruct \
  --quantization awq \
  --tensor-parallel-size 2 \
  --enable-prefix-caching

Within the NVIDIA ecosystem, TensorRT-LLM pushes performance even further at the hardware-optimized kernel level.

---

8. Trends Heading into 2026

• Test-time compute scaling: Rather than scaling model size, this approach invests more compute at inference time to improve output quality. Reasoning models like o1, o3, and DeepSeek-R1 are leading this shift. As a result, research into efficiently handling long chain-of-thought sequences is rapidly gaining momentum.
• Ultra-long context: Contexts exceeding 1M tokens are becoming routine, making KV cache compression increasingly critical.
• Edge inference: Research into extreme model compression for running LLMs on smartphones and PCs is very active.
• Energy efficiency: As power constraints in data centers intensify, inference optimization is no longer just a cost problem — it's increasingly framed as a sustainability problem.

---

Closing Thoughts

Inference optimization can be approached layer by layer:

Lighten the model (quantization, pruning, distillation) → Efficient core operations (FlashAttention, KV cache) → Keep the whole system busy (continuous batching, PagedAttention, distributed inference) → Generate smarter (speculative decoding)

No single technique is a universal solution. The right combination depends on your service's characteristics (latency-first vs. throughput-first), hardware environment, and model size.

In a future post, I'll pick one of these techniques and go deeper with actual implementation and benchmarks.