LLM Inference and Serving

Every hosted LLM API response involves the same handful of engineering decisions underneath the HTTP layer: does this server reuse work across generation steps, does it process the prompt differently than it generates new tokens, and does it serve one request at a time or many together. All three are measurable, even on a laptop CPU with a small model.

Baseline Throughput and Latency

import time
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

tok = AutoTokenizer.from_pretrained("distilgpt2")
model = AutoModelForCausalLM.from_pretrained("distilgpt2")
model.eval()

prompt = "The history of artificial intelligence began"
ids = tok(prompt, return_tensors="pt")
N_TOKENS = 60

start = time.perf_counter()
with torch.no_grad():
    out = model.generate(**ids, max_new_tokens=N_TOKENS, do_sample=False, use_cache=True)
elapsed = time.perf_counter() - start

new_tokens = out.shape[1] - ids["input_ids"].shape[1]
print(f"Generated {new_tokens} tokens in {elapsed:.3f}s")
print(f"Throughput: {new_tokens / elapsed:.1f} tokens/sec")
print(f"Latency per token: {elapsed / new_tokens * 1000:.1f} ms")

$ python baseline.py
Generated 60 tokens in 1.054s
Throughput: 56.9 tokens/sec
Latency per token: 17.6 ms

~57 tokens/sec on a CPU with an 82M-parameter model — every number in this lesson is specific to this hardware and model size, but the relationships between them (cache vs. no cache, prefill vs. decode, batched vs. sequential) hold at any scale.

The KV Cache Is the Single Biggest Win

import time
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

tok = AutoTokenizer.from_pretrained("distilgpt2")
model = AutoModelForCausalLM.from_pretrained("distilgpt2")
model.eval()

prompt = "The history of artificial intelligence began"
ids = tok(prompt, return_tensors="pt")
N_TOKENS = 40

for use_cache in [True, False]:
    start = time.perf_counter()
    with torch.no_grad():
        out = model.generate(**ids, max_new_tokens=N_TOKENS, do_sample=False, use_cache=use_cache)
    elapsed = time.perf_counter() - start
    print(f"use_cache={use_cache}: {elapsed:.3f}s for {N_TOKENS} tokens ({N_TOKENS/elapsed:.1f} tok/s)")

$ python kv_cache.py
use_cache=True: 0.656s for 40 tokens (61.0 tok/s)
use_cache=False: 1.279s for 40 tokens (31.3 tok/s)

Turning the cache off makes generation roughly 2x slower on this short a sequence — and the gap widens as the sequence gets longer, because without caching, generating token N means recomputing attention over all N-1 previous tokens from scratch, every single step. This is not an optional optimization; every production serving stack has a KV cache on by default.

Prefill Is Parallel, Decode Is Sequential

import time
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

tok = AutoTokenizer.from_pretrained("distilgpt2")
model = AutoModelForCausalLM.from_pretrained("distilgpt2")
model.eval()

long_prompt = "Artificial intelligence has transformed many industries. " * 20
ids = tok(long_prompt, return_tensors="pt")
print(f"Prompt length: {ids['input_ids'].shape[1]} tokens")

# Prefill: process the whole prompt at once
start = time.perf_counter()
with torch.no_grad():
    out = model(**ids, use_cache=True)
prefill_time = time.perf_counter() - start
print(f"Prefill ({ids['input_ids'].shape[1]} prompt tokens at once): {prefill_time*1000:.1f} ms")

# Decode: one new token at a time, reusing the cache
past = out.past_key_values
next_id = out.logits[:, -1:].argmax(dim=-1)
decode_times = []
for _ in range(10):
    start = time.perf_counter()
    with torch.no_grad():
        step_out = model(input_ids=next_id, past_key_values=past, use_cache=True)
    decode_times.append(time.perf_counter() - start)
    past = step_out.past_key_values
    next_id = step_out.logits[:, -1:].argmax(dim=-1)

avg_decode = sum(decode_times) / len(decode_times) * 1000
print(f"Decode (1 new token, cache reused): {avg_decode:.1f} ms/token average over 10 steps")
print(f"Prefill per-token cost: {prefill_time*1000/ids['input_ids'].shape[1]:.2f} ms/token (batched)")

$ python prefill_decode.py
Prompt length: 142 tokens
Prefill (142 prompt tokens at once): 221.7 ms
Decode (1 new token, cache reused): 18.0 ms/token average over 10 steps
Prefill per-token cost: 1.56 ms/token (batched)

1.56ms per token during prefill, 18ms per token during decode — roughly 11x more expensive per token to decode than to prefill. Prefill processes 142 tokens in one parallel matrix multiplication; decode processes exactly one token per step, so it can’t use the same parallelism and instead pays the full cost of reading the (growing) KV cache from memory for every single new token. This asymmetry is why “time to first token” and “tokens per second after that” are reported as two separate numbers by every production API.

Batching Multiple Requests Together

import time
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

tok = AutoTokenizer.from_pretrained("distilgpt2")
tok.pad_token = tok.eos_token
tok.padding_side = "left"
model = AutoModelForCausalLM.from_pretrained("distilgpt2")
model.eval()

prompts = [
    "The best programming language is",
    "My favorite type of food is",
    "The weather this weekend will be",
    "The most important skill for engineers is",
]
N_TOKENS = 20

start = time.perf_counter()
for p in prompts:
    ids = tok(p, return_tensors="pt")
    with torch.no_grad():
        model.generate(**ids, max_new_tokens=N_TOKENS, do_sample=False)
sequential_time = time.perf_counter() - start

batch_ids = tok(prompts, return_tensors="pt", padding=True)
start = time.perf_counter()
with torch.no_grad():
    model.generate(**batch_ids, max_new_tokens=N_TOKENS, do_sample=False)
batched_time = time.perf_counter() - start

total_tokens = len(prompts) * N_TOKENS
print(f"Sequential: {sequential_time:.3f}s for {total_tokens} tokens ({total_tokens/sequential_time:.1f} tok/s)")
print(f"Batched:    {batched_time:.3f}s for {total_tokens} tokens ({total_tokens/batched_time:.1f} tok/s)")
print(f"Speedup: {sequential_time/batched_time:.2f}x")

$ python batching.py
Sequential: 1.473s for 80 tokens (54.3 tok/s)
Batched:    0.481s for 80 tokens (166.5 tok/s)
Speedup: 3.06x

Three requests processed together beat three requests processed one after another by roughly 3x, because the model’s weights only get loaded into fast memory once per batch step instead of once per request — most of the cost of a forward pass on modern hardware is moving weights into compute, not the compute itself, so serving more requests per pass is close to free until the batch gets memory-bound. This is exactly the throughput win vLLM and TGI chase — and why they use continuous batching instead of the static batching shown here: a static batch waits for its slowest member to finish before starting the next one, wasting capacity the whole time.

Practice

import time
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

tok = AutoTokenizer.from_pretrained("distilgpt2")
model = AutoModelForCausalLM.from_pretrained("distilgpt2")
model.eval()

for reps in [10, 20]:
    prompt = "Artificial intelligence has transformed many industries. " * reps
    ids = tok(prompt, return_tensors="pt")
    start = time.perf_counter()
    with torch.no_grad():
        model(**ids, use_cache=True)
    elapsed = time.perf_counter() - start
    print(f"{ids['input_ids'].shape[1]:4} tokens: {elapsed*1000:.1f} ms")

  72 tokens: 119.5 ms
 142 tokens: 185.1 ms

import time
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

tok = AutoTokenizer.from_pretrained("distilgpt2")
model = AutoModelForCausalLM.from_pretrained("distilgpt2")
model.eval()

ids = tok("The history of computing", return_tensors="pt")
with torch.no_grad():
    out = model(**ids, use_cache=True)
past = out.past_key_values
next_id = out.logits[:, -1:].argmax(dim=-1)

times = []
for step in range(120):
    start = time.perf_counter()
    with torch.no_grad():
        out = model(input_ids=next_id, past_key_values=past, use_cache=True)
    times.append(time.perf_counter() - start)
    past = out.past_key_values
    next_id = out.logits[:, -1:].argmax(dim=-1)

early = sum(times[:10]) / 10 * 1000
late = sum(times[-10:]) / 10 * 1000
print(f"Decode latency, steps 1-10:    {early:.2f} ms/token")
print(f"Decode latency, steps 110-120: {late:.2f} ms/token")

Decode latency, steps 1-10:    16.85 ms/token
Decode latency, steps 110-120: 18.77 ms/token

import time
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

tok = AutoTokenizer.from_pretrained("distilgpt2")
tok.pad_token = tok.eos_token
tok.padding_side = "left"
model = AutoModelForCausalLM.from_pretrained("distilgpt2")
model.eval()

base_prompts = [
    "The best programming language is", "My favorite type of food is",
    "The weather this weekend will be", "The most important skill for engineers is",
]

for batch_size in [2, 8]:
    prompts = (base_prompts * 2)[:batch_size]
    batch_ids = tok(prompts, return_tensors="pt", padding=True)
    start = time.perf_counter()
    with torch.no_grad():
        model.generate(**batch_ids, max_new_tokens=20, do_sample=False)
    elapsed = time.perf_counter() - start
    total = batch_size * 20
    print(f"batch_size={batch_size}: {elapsed:.3f}s, {total/elapsed:.1f} tok/s total")

batch_size=2: 0.466s, 85.9 tok/s total
batch_size=8: 0.592s, 270.2 tok/s total