Inference Optimization & Serving¶

Overview¶

Serving is a latency-and-throughput engineering problem, not a modeling one: batch efficiently, cache aggressively, pin the GPU busy, and keep Python out of the hot loop. This file covers the pure-PyTorch patterns; frameworks (TorchServe, Triton) automate most of it but the principles are the same.

Batching amortizes kernel launch & memory cost; dynamic batching collects requests.
Caching (LRU on embeddings, KV-cache for LLMs) skips recomputation.
Async I/O: overlap request parsing + pre/post-processing with GPU compute.
Precision: fp16/int8/fp8 at serving time (Ch 06) — but each path needs its own accuracy check.
Keep the model eval + no_grad, warm up once, and avoid recompiles.

💡 The #1 serving perf bug: re-creating graphs or recompiling per request (dynamic shapes, fresh torch.no_grad() contexts). Warm once, serve steady-state.

Static Batching Pattern¶

import torch, torch.nn as nn, time

model = nn.Sequential(nn.Linear(128, 128), nn.GELU(), nn.Linear(128, 10)).cuda().eval()
model = torch.compile(model)                  # warm + fused kernels

WARMUP = torch.randn(64, 128, device='cuda')
with torch.no_grad():
    for _ in range(5):
        model(WARMUP)
torch.cuda.synchronize()

@torch.no_grad()
def run_batch(batch):
    return model(batch)

# measure batching gain
def bench(b, n=100):
    t0 = time.perf_counter()
    for _ in range(n): run_batch(b)
    torch.cuda.synchronize()
    return (time.perf_counter() - t0) / n * 1e3

print("batch 1: %.3f ms" % bench(torch.randn(1, 128, device='cuda')))
print("batch 32: %.3f ms" % bench(torch.randn(32, 128, device='cuda')))

Dynamic Batching — the Little Server¶

import queue, threading, time, torch, torch.nn as nn

class DynamicBatcher:
    def __init__(self, model, max_batch=32, timeout=0.05):
        self.model = model.eval().cuda()
        self.q = queue.Queue()
        self.max_batch, self.timeout = max_batch, timeout
        threading.Thread(target=self._worker, daemon=True).start()

    def _worker(self):
        while True:
            batch, max_wait = [], 0.0  # collect for up to timeout or max_batch
            while len(batch) < self.max_batch:
                try:
                    item = self.q.get(timeout=0.01)
                    batch.append(item)
                except queue.Empty:
                    if max_wait >= self.timeout: break
                    max_wait += 0.01
            if not batch: continue
            xs = torch.cat([t[0] for t in batch])
            with torch.no_grad():
                ys = self.model(xs)
            for t in batch:
                t[1].append(ys[t[2]].cpu())   # slice back per request

    def predict(self, x):
        out = []
        self.q.put((x, out, len(out) - len(out)))
        return out

⚠️ Slicing results back per outgoing request means the batch's rows map to requests — keep request→row bookkeeping exact.

Caching Patterns¶

Embedding-level LRU¶

from functools import lru_cache
import torch

# cache the *computed* outputs for frequent inputs (pure recompute is fine)
@lru_cache(maxsize=1024)
def cached_forward(key_bytes):
    # recompute from reconstructed tensor
    pass

LLM KV-Cache pattern¶

# keep past_key_values across calls, don't recompute prefix
past = None
for token in tokens:
    logits, past = model(token.view(1, 1), past_key_values=past)

Serving Frameworks (the big two)¶

Framework	Batching	Backends	Notes
TorchServe (PyTorch)	dynamic batching	TorchScript/export/ONNX	native integration
Triton (NVIDIA)	dynamic batching	ONNX/TensorRT/C++	best GPU perf, more ops

# TorchServe
torchserve --start --ncs --model-store model_store --models my=model.mar
# Triton
tritonserver --model-repository ./models

Both expect exported models (Ch 07-01) — serve the compiled artifact, not the eager module.

Latency vs Throughput Trade-Off¶

Goal	Config
Min latency (p99)	small batches, graphs, no queuing
Max throughput	large batches, deep queues, batching window
Balanced	dynamic batching window ~RTT/4

The measurement mantra¶

# Always report: tail latency (p99), not mean; and stable throughput (req/s)
# Benchmark under *your* concurrency profile, not a single request loop.

💡 Mean latency is a lie; p99 + throughput at target p99 decides capacity.

Serving Checklist¶

Export once (Ch 07-01); verify parity vs eager.
Warm up with representative shapes; avoid dynamic recompiles.
Apply inference-time precision (fp16/int8) only after accuracy check.
Batch dynamically; use a bounded queue with backpressure.
Cache where it's cheap (embeddings, prefixes); monitor hit rate.
Profile under load: GPU util %, p99, queue depth.

Key Takeaways¶

Batch amortizes launch/memory; dynamic batching collects requests over a window.
Serve exported/compiled artifacts; keep the model eval, warm, steady-state.
Cache embeddings/prefixes when recompute costs more than memory.
TorchServe (native) vs Triton (perf): pick by your stack.
Measure p99 + throughput under concurrency, not single-run means.