Complete LLM Inference Optimization Guide 2025: vLLM, TensorRT-LLM, KV Cache, Speculative Decoding

┌──────────────────────────────────────────────────────┐
│                 LLM Inference Pipeline                │
├──────────────────┬───────────────────────────────────┤
│   Prefill Phase  │         Decode Phase              │
│  (Prompt Proc.)  │      (Token Generation)           │
├──────────────────┼───────────────────────────────────┤
│ - Input tokens   │ - Generates 1 token at a time     │
│   in parallel    │ - Memory-Bound                    │
│ - Compute-Bound  │ - Low GPU utilization (5-15%)     │
│ - High GPU util  │ - Repeats for output length       │
│ - Runs once      │ - KV Cache read + append          │
│ - Creates KV $   │                                   │
└──────────────────┴───────────────────────────────────┘

# Without KV Cache (full recomputation each step)
# Total compute for n tokens: O(n^2 * d)

# With KV Cache (reuse previous results)
# Total compute for n tokens: O(n * d)
# But KV Cache memory: O(n * d) additional

KV Cache Size = 2 * num_layers * num_kv_heads * head_dim * seq_len * batch_size * dtype_size

Example: Llama-2 70B, seq_len=4096, batch_size=1, FP16
= 2 * 80 * 8 * 128 * 4096 * 1 * 2 bytes
= 1.34 GB (for a single sequence!)

With batch_size=32: 1.34 * 32 = 42.9 GB

┌─────────────────────────────────────────────┐
│    Traditional KV Cache Allocation           │
│                                             │
│  Request 1: [████████░░░░░░░░░░░░]  40% used│
│  Request 2: [████████████░░░░░░░░]  60% used│
│  Request 3: [██░░░░░░░░░░░░░░░░░░]  10% used│
│              ^^^^^^^^ wasted memory          │
└─────────────────────────────────────────────┘

┌─────────────────────────────────────────────┐
│       PagedAttention KV Cache Allocation     │
│                                             │
│  Physical Blocks: [B0][B1][B2][B3][B4][B5]  │
│                                             │
│  Request 1 -> Page Table: [B0, B3, B5]      │
│  Request 2 -> Page Table: [B1, B4, B6, B7]  │
│  Request 3 -> Page Table: [B2]              │
│                                             │
│  Near-zero internal fragmentation           │
│  Non-contiguous memory blocks utilized      │
│  Copy-on-Write for shared prompts           │
└─────────────────────────────────────────────┘

# Enable Prefix Caching in vLLM
from vllm import LLM

llm = LLM(
    model="meta-llama/Llama-3.1-70B-Instruct",
    enable_prefix_caching=True,
    max_model_len=8192,
)

# Requests using the same system prompt
# share the KV Cache for the system prompt portion

┌──────────────────────────────────────────────┐
│          FlashAttention Core Idea             │
│                                              │
│  GPU Memory Hierarchy:                       │
│  ┌─────────┐  19 TB/s   ┌─────────────────┐ │
│  │  SRAM   │<---------->│  Compute Units   │ │
│  │ (20 MB) │            └─────────────────┘ │
│  └────┬────┘                                 │
│       | 2-4.8 TB/s                           │
│  ┌────v────────────────┐                     │
│  │    HBM (80-141 GB)  │                     │
│  └─────────────────────┘                     │
│                                              │
│  Strategy: Split Q,K,V into tiles (blocks),  │
│  perform all computation in SRAM,            │
│  write only final results to HBM             │
└──────────────────────────────────────────────┘

┌─────────────────────────────────────────────────────┐
│    Multi-Head Attention (MHA)                       │
│    Q heads: [H1][H2][H3][H4][H5][H6][H7][H8]      │
│    K heads: [H1][H2][H3][H4][H5][H6][H7][H8]      │
│    V heads: [H1][H2][H3][H4][H5][H6][H7][H8]      │
│    KV Cache: 8x                                     │
├─────────────────────────────────────────────────────┤
│    Multi-Query Attention (MQA)                      │
│    Q heads: [H1][H2][H3][H4][H5][H6][H7][H8]      │
│    K heads: [        H_shared         ]             │
│    V heads: [        H_shared         ]             │
│    KV Cache: 1x (8x reduction)                      │
├─────────────────────────────────────────────────────┤
│    Grouped-Query Attention (GQA, 2 groups)          │
│    Q heads: [H1][H2][H3][H4] | [H5][H6][H7][H8]   │
│    K heads: [   K_group1   ] | [   K_group2   ]    │
│    V heads: [   V_group1   ] | [   V_group2   ]    │
│    KV Cache: 2x (4x reduction)                      │
└─────────────────────────────────────────────────────┘

Static Batching (traditional):
Time ──────────────────────────────────>

Req 1: [████████████████████████████████]  (long response)
Req 2: [████████░░░░░░░░░░░░░░░░░░░░░░]  (short response)
Req 3: [██████████████░░░░░░░░░░░░░░░░]  (medium response)
Req 4: [WAIT WAIT WAIT WAIT WAIT WAIT ]  (waiting)

░ = GPU idle (padding), WAIT = waiting for batch to complete
Entire batch must finish before next batch starts -> very low throughput

Continuous Batching:
Time ──────────────────────────────────>

Req 1: [████████████████████████████████]
Req 2: [████████]
Req 3:          [██████████████]
Req 4:                  [████████████████]
Req 5:                          [████████]

Completed requests immediately removed -> new requests immediately added
GPU idle time minimized -> 10-20x throughput improvement

# vLLM chunked prefill configuration
from vllm import LLM, SamplingParams

llm = LLM(
    model="meta-llama/Llama-3.1-8B-Instruct",
    enable_chunked_prefill=True,
    max_num_batched_tokens=2048,  # max tokens per iteration
)

# A long prompt (e.g., 32K tokens) is split into 2048-token chunks
# Decode requests can also be processed between chunks
# Slightly increases TTFT but improves overall system throughput and ITL

┌────────────────────────────────────────────────────┐
│           Speculative Decoding Flow                │
│                                                    │
│  Step 1: Draft Model (small, fast)                 │
│  "The capital of France is" -> [Paris][,][a][city] │
│  4 tokens predicted very quickly (4ms)             │
│                                                    │
│  Step 2: Target Model (large, accurate)            │
│  Single forward pass verifies all 4 tokens         │
│  [Paris OK] [, OK] [a FAIL->"known"] [city FAIL]  │
│                                                    │
│  Result: "Paris, known" (2 accepted + 1 corrected) │
│  Before: 3 forward passes needed -> now 1          │
│  Speedup: ~2-3x                                    │
└────────────────────────────────────────────────────┘

# 1. Separate Draft Model
from vllm import LLM, SamplingParams

llm = LLM(
    model="meta-llama/Llama-3.1-70B-Instruct",
    speculative_model="meta-llama/Llama-3.1-8B-Instruct",
    num_speculative_tokens=5,
    use_v2_block_manager=True,
)

# 2. Medusa Heads (additional MLP heads predict multiple positions)
# No draft model needed - adds lightweight heads to target model itself
# Requires training but minimal memory overhead

# 3. EAGLE (Extrapolation Algorithm for Greater Language-model Efficiency)
# Draft model reuses target model's hidden states
# Higher acceptance rate than separate draft models

Token position:    1        2        3
                 +-- Paris --+-- is -- ...
The -------------+           +-- was -- ...
                 +-- Lyon --- is -- ...
                 +-- capital - of -- ...

All tree paths verified in a single forward pass
-> Maximizes acceptance rate, improves throughput

┌───────────────────────────────────────────────────────┐
│            Quantization Technique Classification      │
├─────────────────────┬─────────────────────────────────┤
│  Post-Training      │  Training-Aware                 │
│  Quantization(PTQ)  │  Quantization                   │
├─────────────────────┼─────────────────────────────────┤
│  - GPTQ (INT4)      │  - QLoRA + Merge                │
│  - AWQ (INT4)       │  - QAT (Quantization-Aware      │
│  - GGUF (various)   │    Training)                    │
│  - bitsandbytes     │                                 │
│  - SmoothQuant      │                                 │
│  - FP8 Dynamic      │                                 │
└─────────────────────┴─────────────────────────────────┘

# GPTQ: Layer-wise optimal quantization (OBQ-based)
# Pros: Good quality even at INT4, optimized for GPU inference
# Cons: Requires calibration data, slow quantization

from transformers import AutoModelForCausalLM, GPTQConfig

gptq_config = GPTQConfig(
    bits=4,
    group_size=128,
    dataset="c4",
    desc_act=True,
)
model = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Llama-3.1-70B-Instruct",
    quantization_config=gptq_config,
    device_map="auto",
)

# AWQ: Activation-aware Weight Quantization
# Key: Finds and protects important weight channels (based on activation magnitude)
# Faster quantization than GPTQ, similar quality

from awq import AutoAWQForCausalLM

model = AutoAWQForCausalLM.from_pretrained(
    "meta-llama/Llama-3.1-70B-Instruct"
)
quant_config = {
    "zero_point": True,
    "q_group_size": 128,
    "w_bit": 4,
    "version": "GEMM"
}
model.quantize(tokenizer, quant_config=quant_config)

# bitsandbytes: Simple INT8/NF4 quantization
from transformers import BitsAndBytesConfig

bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype="bfloat16",
    bnb_4bit_use_double_quant=True,
)
model = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Llama-3.1-8B-Instruct",
    quantization_config=bnb_config,
)

┌──────────────────────────────────────────┐
│              vLLM Architecture            │
│                                          │
│  ┌─────────┐     ┌──────────────────┐   │
│  │ FastAPI  │---->│   LLM Engine     │   │
│  │ Server   │     │                  │   │
│  └─────────┘     │  ┌────────────┐  │   │
│                  │  │ Scheduler   │  │   │
│  ┌─────────┐    │  │ (Batching)  │  │   │
│  │ OpenAI  │---->│  └─────┬──────┘  │   │
│  │ compat  │     │        |         │   │
│  └─────────┘     │  ┌─────v──────┐  │   │
│                  │  │ Block Mgr   │  │   │
│                  │  │ (PagedAttn) │  │   │
│                  │  └─────┬──────┘  │   │
│                  │        |         │   │
│                  │  ┌─────v──────┐  │   │
│                  │  │  Worker(s)  │  │   │
│                  │  │ (GPU exec)  │  │   │
│                  │  └────────────┘  │   │
│                  └──────────────────┘   │
└──────────────────────────────────────────┘

# Start vLLM server (OpenAI API compatible)
# vllm serve meta-llama/Llama-3.1-70B-Instruct \
#     --tensor-parallel-size 4 \
#     --max-model-len 32768 \
#     --gpu-memory-utilization 0.90 \
#     --enable-prefix-caching \
#     --enable-chunked-prefill \
#     --max-num-batched-tokens 4096 \
#     --port 8000

# API call via Python
from openai import OpenAI

client = OpenAI(
    base_url="http://localhost:8000/v1",
    api_key="token-abc123",
)

response = client.chat.completions.create(
    model="meta-llama/Llama-3.1-70B-Instruct",
    messages=[
        {"role": "system", "content": "You are a helpful assistant."},
        {"role": "user", "content": "Explain quantum computing"},
    ],
    max_tokens=512,
    temperature=0.7,
)

# vllm serve meta-llama/Llama-3.1-8B-Instruct \
#     --enable-lora \
#     --lora-modules \
#         sql-lora=./adapters/sql-lora \
#         code-lora=./adapters/code-lora \
#         chat-lora=./adapters/chat-lora \
#     --max-loras 3 \
#     --max-lora-rank 64

# Select LoRA adapter by model name in API call
response = client.chat.completions.create(
    model="sql-lora",  # LoRA adapter name
    messages=[{"role": "user", "content": "SELECT ..."}],
)

# Multimodal model serving
# vllm serve Qwen/Qwen2-VL-7B-Instruct \
#     --max-model-len 8192 \
#     --limit-mm-per-prompt image=4

from openai import OpenAI
import base64

client = OpenAI(base_url="http://localhost:8000/v1", api_key="key")

response = client.chat.completions.create(
    model="Qwen/Qwen2-VL-7B-Instruct",
    messages=[{
        "role": "user",
        "content": [
            {"type": "text", "text": "What is in this image?"},
            {"type": "image_url", "image_url": {"url": "data:image/png;base64,..."}}
        ]
    }],
)

┌────────┐     ┌───────────┐     ┌──────────┐     ┌──────────┐
│HF Model│---->│ Convert   │---->│ TRT-LLM  │---->│ Triton   │
│(source) │     │ Checkpoint│     │ Engine   │     │ Serving  │
└────────┘     └───────────┘     └──────────┘     └──────────┘
                 Apply quant      Compile optim    API server

# Step 1: Checkpoint conversion + FP8 quantization
python convert_checkpoint.py \
    --model_dir meta-llama/Llama-3.1-70B-Instruct \
    --output_dir ./checkpoint_fp8 \
    --dtype bfloat16 \
    --tp_size 4 \
    --pp_size 1 \
    --use_fp8

# Step 2: Build TensorRT engine
trtllm-build \
    --checkpoint_dir ./checkpoint_fp8 \
    --output_dir ./engine_fp8 \
    --gemm_plugin auto \
    --max_batch_size 64 \
    --max_input_len 4096 \
    --max_seq_len 8192 \
    --paged_kv_cache enable \
    --use_paged_context_fmha enable \
    --workers 4

# Triton Inference Server + TensorRT-LLM backend
# model_config.pbtxt configuration
"""
backend: "tensorrtllm"
max_batch_size: 64

model_transaction_policy {
  decoupled: True    # Streaming response support
}

parameters: {
  key: "batching_type"
  value: {string_value: "inflight"}  # Enable Inflight Batching
}

parameters: {
  key: "max_tokens_in_paged_kv_cache"
  value: {string_value: "131072"}   # Limit KV Cache token count
}
"""

Tensor Parallelism (TP=4):

         Layer N weight matrix W
    ┌──────┬──────┬──────┬──────┐
    │ W_1  │ W_2  │ W_3  │ W_4  │
    │GPU 0 │GPU 1 │GPU 2 │GPU 3 │
    └──┬───┴──┬───┴──┬───┴──┬───┘
       |      |      |      |
       v      v      v      v
    [part1] [part2] [part3] [part4]
       |      |      |      |
       └──────┴──────┴──────┘
              All-Reduce
              (aggregate results)

Pros: Reduces latency (all GPUs compute simultaneously)
Cons: Requires inter-GPU communication (NVLink recommended)
Best for: GPUs within same node (low latency needed)

Pipeline Parallelism (PP=4, 80 layers):

GPU 0: [Layer 0-19]  -> GPU 1: [Layer 20-39]
                              -> GPU 2: [Layer 40-59]
                                       -> GPU 3: [Layer 60-79]

Pros: Minimal inter-GPU communication (one direction)
Cons: Pipeline bubbles (GPU idle time)
Best for: Cross-node distribution (higher latency tolerable)

Expert Parallelism (Mixtral 8x7B, EP=4):

GPU 0: Expert 0, 1 + Shared Layers
GPU 1: Expert 2, 3 + Shared Layers
GPU 2: Expert 4, 5 + Shared Layers
GPU 3: Expert 6, 7 + Shared Layers

Token routing: Each token sent to Top-2 Experts
-> Requires All-to-All communication between GPUs

# Serving Llama-3.1 405B (8x H100 80GB)
# Model size: ~810 GB (FP16) -> FP8 ~405 GB

# Option 1: TP=8 (all layers split across all GPUs)
vllm serve meta-llama/Llama-3.1-405B-Instruct-FP8 \
    --tensor-parallel-size 8 \
    --max-model-len 16384

# Option 2: TP=4, PP=2 (4 GPUs per pipeline stage, 2 stages)
vllm serve meta-llama/Llama-3.1-405B-Instruct-FP8 \
    --tensor-parallel-size 4 \
    --pipeline-parallel-size 2 \
    --max-model-len 16384

# KV Cache FP8 quantization in vLLM
# vllm serve meta-llama/Llama-3.1-70B-Instruct \
#     --tensor-parallel-size 4 \
#     --kv-cache-dtype fp8 \
#     --quantization fp8

# KV Cache memory savings:
# FP16 KV Cache: 1.34 GB / sequence (Llama-2 70B, 4K)
# FP8 KV Cache:  0.67 GB / sequence (50% savings)
# 2x more concurrent requests on same GPU

GPU Memory Distribution (A100 80GB, Llama-3.1 70B FP16):

┌─────────────────────────────────┐
│  Model Weights: ~35 GB (TP=2)   │  43.75%
├─────────────────────────────────┤
│  KV Cache: ~35 GB               │  43.75%
│  (gpu_memory_utilization=0.90)  │
├─────────────────────────────────┤
│  Activation Memory: ~2 GB       │  2.5%
├─────────────────────────────────┤
│  System Reserved: ~8 GB         │  10%
└─────────────────────────────────┘

KV Cache determines the maximum number of concurrent requests.

Monthly token usage cost comparison (Llama-3.1 70B class):

┌─────────────────────────────────────────────────┐
│ Cost                                            │
│ ($)                                             │
│ 5000|                                    /API   │
│     |                                  /        │
│ 3000|                    ---------- Self-Hosted  │
│     |              /  (H100x4 monthly fixed)    │
│ 1000|      / API                                │
│     |   /                                       │
│    0+--+----+----+----+----+----+------>        │
│     0  2B   5B  10B  20B  50B 100B  tokens/mo  │
│                                                 │
│  Break-even: ~10B tokens/month                  │
└─────────────────────────────────────────────────┘

# vLLM built-in benchmark (recommended)
# After running: python -m vllm.entrypoints.openai.api_server

# python benchmarks/benchmark_serving.py \
#     --backend vllm \
#     --model meta-llama/Llama-3.1-8B-Instruct \
#     --dataset-name sharegpt \
#     --dataset-path ShareGPT_V3_unfiltered.json \
#     --num-prompts 1000 \
#     --request-rate 10 \
#     --endpoint /v1/completions

# Example results:
# Successful requests:                     1000
# Benchmark duration (s):                  105.23
# Total input tokens:                      215000
# Total generated tokens:                  180000
# Request throughput (req/s):              9.50
# Output token throughput (tok/s):         1710.5
# Mean TTFT (ms):                          48.2
# Median TTFT (ms):                        42.1
# P99 TTFT (ms):                           125.3
# Mean ITL (ms):                           11.8
# Median ITL (ms):                         10.2
# P99 ITL (ms):                            35.7

Throughput vs Latency relationship (as concurrency increases):

Throughput                          Latency
(tok/s)                            (ms)
  |        ┌──────────              |              /
  |       /                         |            /
  |      /                          |          /
  |     /                           |        /
  |    /                            |      /
  |   /                             |    /
  |  /                              |  /
  | /                               |/
  +──────────────> concurrency      +──────────────> concurrency

  Optimal operating point: Just before throughput saturates (Knee point)
  Usually GPU utilization of 70-80%

┌──────────────────────────────────────────────────┐
│                Production Architecture            │
│                                                  │
│  Client -> Load Balancer -> API Gateway          │
│                              |                   │
│                    ┌─────────┼─────────┐         │
│                    v         v         v         │
│              ┌─────────┐┌────────┐┌────────┐    │
│              │ vLLM    ││ vLLM   ││ vLLM   │    │
│              │ Pod 1   ││ Pod 2  ││ Pod 3  │    │
│              │ (4xH100)││(4xH100)││(4xH100)│    │
│              └────┬────┘└───┬────┘└───┬────┘    │
│                   |         |         |          │
│              ┌────v─────────v─────────v────┐    │
│              │     Prometheus + Grafana     │    │
│              │    (Metrics collection)      │    │
│              └─────────────────────────────┘    │
│                                                  │
│  Autoscaling: Based on queue length / GPU util   │
│  Health Check: /health endpoint                  │
│  Graceful Shutdown: Complete in-flight requests  │
└──────────────────────────────────────────────────┘

apiVersion: apps/v1
kind: Deployment
metadata:
  name: vllm-llama3-70b
spec:
  replicas: 3
  selector:
    matchLabels:
      app: vllm-llama3
  template:
    metadata:
      labels:
        app: vllm-llama3
    spec:
      containers:
      - name: vllm
        image: vllm/vllm-openai:latest
        command: ["python", "-m", "vllm.entrypoints.openai.api_server"]
        args:
          - "--model=meta-llama/Llama-3.1-70B-Instruct"
          - "--tensor-parallel-size=4"
          - "--max-model-len=16384"
          - "--gpu-memory-utilization=0.90"
          - "--enable-prefix-caching"
          - "--enable-chunked-prefill"
        resources:
          limits:
            nvidia.com/gpu: "4"
            memory: "64Gi"
          requests:
            nvidia.com/gpu: "4"
            memory: "32Gi"
        ports:
        - containerPort: 8000
        readinessProbe:
          httpGet:
            path: /health
            port: 8000
          initialDelaySeconds: 120
          periodSeconds: 10
        livenessProbe:
          httpGet:
            path: /health
            port: 8000
          initialDelaySeconds: 180
          periodSeconds: 30
      nodeSelector:
        gpu-type: h100
      tolerations:
      - key: nvidia.com/gpu
        operator: Exists
        effect: NoSchedule