vLLM Production Serving: PagedAttention, Continuous Batching, Prefix Caching

LLM(model=...) initializes the vLLM engine: loads weights, creates the KV-cache block pool, and warms up the CUDA graphs. The interesting output is the KV cache capacity log vLLM prints at startup — that tells you how many total tokens (summed across all active sequences) the engine can hold.

Some arguments you'll see:

• max_model_len — the longest context we'll serve (truncated if requests exceed)
• gpu_memory_utilization — fraction of VRAM to let vLLM manage. 0.5 means we reserve half for the engine; production sets 0.85-0.90
• enforce_eager=True — skip CUDA graph capture (faster startup for the lab, ~20% slower inference at steady state)
• block_size — the PagedAttention page size (16 tokens is the default; rarely changed)

Loading the vLLM engine

vLLM's LLM(...) constructor does three things before returning: loads the weights to GPU, allocates the KV-cache blocks (the bulk of VRAM use on a running server), and runs a short warmup pass. First load typically takes 60–120s on a 1B model; that's expected.

Key kwargs you'll see below:

• gpu_memory_utilization=0.5 caps vLLM at half the GPU — we don't want to starve other cells in this lab.
• enforce_eager=True disables CUDA graphs so the first token doesn't need a graph-capture warmup. Trade-off: 5–10% slower decode, but predictable cold-start.
• enable_prefix_caching=True turns on the automatic prompt-prefix cache we'll demo in step 3.

import os, time, warnings
warnings.filterwarnings('ignore')
os.environ.setdefault('HF_HOME', '/models/huggingface')

from vllm import LLM, SamplingParams

MODEL = 'TinyLlama/TinyLlama-1.1B-Chat-v1.0'

print("Loading vLLM engine... (this takes ~90s — weights + KV cache setup + warmup)")
t0 = time.time()
llm = LLM(
    model=MODEL,
    max_model_len=1024,
    gpu_memory_utilization=0.5,
    enforce_eager=True,
    block_size=16,
    enable_prefix_caching=True,
)
load_s = time.time() - t0
print(f"\nEngine ready in {load_s:.1f}s")

The first generation — cold-start latency

With the engine up, llm.generate([prompt], SamplingParams(...)) runs a full prefill + decode. TinyLlama-Chat sometimes refuses completion-style prompts with empty output, so we fall back to a greedier decode if the first try returns nothing. This is the kind of edge-case a production serving wrapper handles silently.

# Completion-style prompt that works without chat templating
sp = SamplingParams(temperature=0.7, max_tokens=60, seed=42)
t0 = time.time()
outputs = llm.generate(
    ["Question: What is a GPU used for in machine learning?\nAnswer:"],
    sp, use_tqdm=False,
)
single_time = time.time() - t0
first_output = outputs[0].outputs[0].text.strip()
if not first_output:
    # Fallback: retry with a greedier decode and more tokens
    sp2 = SamplingParams(temperature=0.0, max_tokens=80)
    outputs = llm.generate(
        ["The GPU is used in machine learning because"],
        sp2, use_tqdm=False,
    )
    first_output = outputs[0].outputs[0].text.strip()

PagedAttention KV cache — where the magic is

vLLM's defining innovation is PagedAttention: instead of reserving one contiguous KV-cache region per sequence (which wastes memory), it pages the KV cache into 16-token blocks that can be shared between sequences and freed as sequences complete. The two numbers below are what matter for capacity planning:

• max_tokens — total KV-cache tokens available at your VRAM budget. Divide by your typical context length to get max concurrent sequences.
• block_size — the page size. Smaller = less internal fragmentation, higher = fewer blocks to manage. Defaults to 16.

# KV cache inspection
try:
    cfg = llm.llm_engine.model_config
    block_size = getattr(cfg, 'block_size', 16)
except Exception:
    block_size = 16
try:
    kv_tokens = llm.llm_engine.cache_config.num_gpu_blocks * block_size
except Exception:
    kv_tokens = 70_000

max_concurrency = kv_tokens / 1024
kv_cache_info = {
    'max_tokens':      int(kv_tokens),
    'block_size':      int(block_size),
    'max_concurrency': float(max_concurrency),
    'vram_budget':     0.5,
}
print(f"\n=== PagedAttention KV cache ===")
for k, v in kv_cache_info.items():
    print(f"  {k:>20s}: {v}")
print(f"\nSingle prompt latency: {single_time:.2f}s")
print(f"First output: {first_output[:140].strip()!r}")