Skip to content

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
# Benchmark under *your* concurrency profile, not a single request loop.

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

-

Serving Checklist

  1. Export once (Ch 07-01); verify parity vs eager.
  2. Warm up with representative shapes; avoid dynamic recompiles.
  3. Apply inference-time precision (fp16/int8) only after accuracy check.
  4. Batch dynamically; use a bounded queue with backpressure.
  5. Cache where it's cheap (embeddings, prefixes); monitor hit rate.
  6. 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.

-