Inference Optimization Patterns

Overview

Inference (serving predictions) has different requirements than training:

• Latency matters more than throughput
• Memory is often the bottleneck (not compute)
• Batching combines multiple requests
• Caching reduces recomputation
• Precision can be reduced (fp32 → fp16 → int8)

Understanding inference patterns is critical for production ML systems.

-

Single-Request Inference

Basic Inference

import torch
import torch.nn as nn

# Load model
model = nn.Linear(10, 5)
model.load_state_dict(torch.load('model.pth'))
model.eval()

# Single inference request
input_data = torch.randn(1, 10)

# Forward pass
with torch.no_grad(): # Skip gradient tracking
 output = model(input_data)
 prediction = output.argmax(dim=-1)

print(f"Prediction: {prediction}")

# Characteristics:
# - Small input (1 sample)
# - No gradient computation
# - Latency critical

Device Optimization

# Move model to GPU for faster inference
device = 'cuda' if torch.cuda.is_available() else 'cpu'
model = model.to(device)

# Move input to same device
input_data = input_data.to(device)

# Inference
with torch.no_grad():
 output = model(input_data)

Inference Time Patterns

import time

def benchmark_inference(model, input_shape, device='cuda', num_runs=100):
 """Benchmark inference latency."""
 model = model.to(device).eval()

 # Warmup
 for _ in range(10):
 with torch.no_grad():
 _ = model(torch.randn(*input_shape, device=device))

 # Measure
 torch.cuda.synchronize() # Wait for GPU
 start = time.time()

 for _ in range(num_runs):
 with torch.no_grad():
 _ = model(torch.randn(*input_shape, device=device))

 torch.cuda.synchronize() # Wait for completion
 elapsed = time.time() - start

 latency_ms = (elapsed / num_runs) * 1000
 print(f"Latency: {latency_ms:.2f}ms per request")
 print(f"Throughput: {num_runs / elapsed:.0f} requests/sec")

 return latency_ms

# Benchmark
model = nn.Linear(10, 5)
latency = benchmark_inference(model, (1, 10))

---

Batch Inference

Processing Multiple Requests

# Multiple inference requests (batch)
batch_inputs = torch.randn(32, 10) # 32 samples

model.eval()
with torch.no_grad():
 batch_outputs = model(batch_inputs)
 batch_predictions = batch_outputs.argmax(dim=-1)

# Characteristics:
# - Larger batch size
# - Better GPU utilization
# - Higher throughput, higher latency
# - Use for offline inference (e.g., batch scoring)

Asynchronous Batching

from queue import Queue
import threading
import time

class AsyncBatcher:
 """Collect requests and batch inference."""

 def __init__(self, model, batch_size=32, timeout=0.1):
 self.model = model.eval()
 self.batch_size = batch_size
 self.timeout = timeout
 self.request_queue = Queue()
 self.response_queue = {}

 # Start batch inference thread
 self.thread = threading.Thread(target=self._batch_worker, daemon=True)
 self.thread.start()

 def _batch_worker(self):
 """Worker thread that batches and processes requests."""
 while True:
 # Collect up to batch_size requests
 batch = []
 request_ids = []

 start_time = time.time()
 while len(batch) < self.batch_size:
 try:
 elapsed = time.time() - start_time
 remaining = max(0, self.timeout - elapsed)

 request_id, input_data = self.request_queue.get(timeout=remaining)
 batch.append(input_data)
 request_ids.append(request_id)
 except:
 break

 # Process batch
 if batch:
 batch_tensor = torch.stack(batch)

 with torch.no_grad():
 outputs = self.model(batch_tensor)

 # Return results
 for req_id, output in zip(request_ids, outputs):
 self.response_queue[req_id] = output

 def infer(self, input_data):
 """Submit inference request (async)."""
 request_id = id(input_data)
 self.request_queue.put((request_id, input_data))

 # Wait for response
 while request_id not in self.response_queue:
 time.sleep(0.001)

 return self.response_queue.pop(request_id)

# Usage
batcher = AsyncBatcher(model, batch_size=32)

# Submit requests (will be batched together)
results = []
for i in range(100):
 input_data = torch.randn(10)
 output = batcher.infer(input_data)
 results.append(output)

-

KV Cache: Reducing Computation

KV Cache in Transformers

# Problem
# 
# Without cache:
# Token 1
# Token 2
# Token 3
# → O(n²) computation
#
# With cache:
# Token 1
# Token 2
# Token 3
# → O(n) computation

import torch
import torch.nn as nn

class CachedAttention(nn.Module):
 def __init__(self, hidden_size, num_heads):
 super().__init__()
 self.hidden_size = hidden_size
 self.num_heads = num_heads
 self.head_dim = hidden_size // num_heads

 self.query = nn.Linear(hidden_size, hidden_size)
 self.key = nn.Linear(hidden_size, hidden_size)
 self.value = nn.Linear(hidden_size, hidden_size)
 self.out = nn.Linear(hidden_size, hidden_size)

 def forward(self, x, cache=None):
 """
 x: (batch_size, 1, hidden_size) - single token
 cache: (batch_size, seq_len, hidden_size) or None
 """
 batch_size = x.size(0)

 # Compute Q, K, V for current token only
 q = self.query(x) # (batch_size, 1, hidden_size)
 k = self.key(x) # (batch_size, 1, hidden_size)
 v = self.value(x) # (batch_size, 1, hidden_size)

 # Concatenate with cached K, V from previous tokens
 if cache is not None:
 k_cache, v_cache = cache
 k = torch.cat([k_cache, k], dim=1) # (batch_size, seq_len, hidden_size)
 v = torch.cat([v_cache, v], dim=1)

 # Attention with full K, V but only current Q
 attn = (q @ k.transpose(-2, -1)) / (self.head_dim ** 0.5)
 attn = torch.softmax(attn, dim=-1)
 output = attn @ v
 output = self.out(output)

 # Return output and updated cache for next iteration
 return output, (k, v)

# Inference with cache
model = CachedAttention(hidden_size=768, num_heads=12)
model.eval()

cache = None
generated_tokens = []

for step in range(100):
 # Current token
 if step == 0:
 x = torch.randn(1, 1, 768) # First token
 else:
 x = embedding(generated_tokens[-1]) # Previous token embedding

 # Forward with cache
 with torch.no_grad():
 output, cache = model(x, cache)

 # Sample next token
 token = output.argmax(dim=-1)
 generated_tokens.append(token)

 # Cache reused for next iteration

Memory Savings with KV Cache

# Calculate memory savings
# 
# Batch size
# Sequence length
# Hidden size
# Data type

seq_len = 512
batch_size = 32
hidden_size = 768
dtype_bytes = 4

# Without cache
attention_matrices = seq_len * seq_len * batch_size * hidden_size * dtype_bytes
print(f"Without cache: {attention_matrices / 1e9:.2f}GB") # ~3.1GB

# With cache
cache_memory = seq_len * batch_size * hidden_size * dtype_bytes * 2 # K + V
print(f"Cache memory: {cache_memory / 1e6:.2f}MB") # ~402MB

# Savings

-

Mixed Precision Inference

Float16 Inference

# Float32
# Float16
# Use float16 when accuracy loss is acceptable

model = nn.Linear(10, 5)
model = model.half() # Convert to float16

# Input also needs to be float16
input_data = torch.randn(1, 10, dtype=torch.float16)

with torch.no_grad():
 output = model(input_data)

# Memory savings
# Speed improvement
# Accuracy impact

Automatic Mixed Precision (AMP)

from torch.cuda.amp import autocast

model = nn.Linear(10, 5)
model.eval()

# Use autocast to select precision automatically
# Low-precision ops (matrix multiply, conv) → float16
# High-precision ops (reductions) → float32

with torch.no_grad():
 with autocast(device_type='cuda'):
 output = model(torch.randn(32, 10))

# Automatic selection gives best accuracy/speed trade-off

Quantization for Inference

import torch.quantization as quantization

# Dynamic quantization (simple, no calibration needed)
model = nn.Linear(10, 5)
model = quantization.quantize_dynamic(
 model,
 qconfig_spec={nn.Linear}, # Quantize Linear layers
 dtype=torch.qint8
)

# Static quantization (better, needs calibration)
model.qconfig = quantization.get_default_qconfig('fbgemm')
quantization.prepare(model, inplace=True)

# Calibration
for batch in calibration_loader:
 with torch.no_grad():
 model(batch)

quantization.convert(model, inplace=True)

# Benefits:
# - 4x memory reduction (float32 → int8)
# - 4x faster (CPU) or 2x (GPU)
# - <1% accuracy loss typically

-

Batching Strategies

Dynamic Batching

from collections import defaultdict
import heapq

class DynamicBatcher:
 """Batch requests with similar sizes together."""

 def __init__(self, model, max_batch_size=32, timeout=0.1):
 self.model = model.eval()
 self.max_batch_size = max_batch_size
 self.timeout = timeout
 self.batches_by_size = defaultdict(list)

 def infer(self, input_data):
 """Batch inference by input size."""
 input_size = input_data.shape[1] # Sequence length

 # Add to appropriate size bucket
 self.batches_by_size[input_size].append(input_data)

 # Batch when full or timeout
 batch = self.batches_by_size[input_size]
 if len(batch) >= self.max_batch_size:
 outputs = self._process_batch(batch)
 self.batches_by_size[input_size] = []
 return outputs[-1] # Return last result

 return None

 def _process_batch(self, batch):
 """Process batch of same-sized inputs."""
 stacked = torch.stack(batch)
 with torch.no_grad():
 outputs = self.model(stacked)
 return list(outputs)

# Usage
batcher = DynamicBatcher(model, max_batch_size=32)

# Requests with different sequence lengths batched separately
outputs = []
for seq_len in [10, 20, 10, 20, 10]:
 input_data = torch.randn(1, seq_len)
 output = batcher.infer(input_data)
 if output is not None:
 outputs.append(output)

Bucket Batching

# Group similar sequence lengths into buckets
# Reduces padding waste

def create_buckets(lengths, bucket_size=50):
 """Create length buckets."""
 buckets = defaultdict(list)
 for i, length in enumerate(lengths):
 bucket_idx = (length - 1) // bucket_size * bucket_size
 buckets[bucket_idx].append(i)
 return buckets

# Lengths
buckets = create_buckets([10, 15, 20, 25, 30, 35, 40, 45, 50, 100], bucket_size=50)

# Result:
# Bucket 0
# Bucket 50
# → Reduced padding, better performance

-

Practical: Production Inference Server

import torch
import torch.nn as nn
from flask import Flask, request, jsonify
import logging

class InferenceServer:
 def __init__(self, model_path, device='cuda'):
 self.device = device
 self.model = nn.Linear(10, 5)
 self.model.load_state_dict(torch.load(model_path))
 self.model = self.model.to(device).eval()

 # Optimization settings
 if device == 'cuda':
 self.model = self.model.half() # Mixed precision

 self.logger = logging.getLogger(__name__)

 def preprocess(self, data):
 """Convert input to tensor."""
 import numpy as np
 tensor = torch.from_numpy(np.array(data)).float()
 return tensor.to(self.device)

 def infer(self, input_data):
 """Run inference."""
 with torch.no_grad():
 output = self.model(input_data)

 return output

 def postprocess(self, output):
 """Convert output to JSON."""
 result = output.cpu().numpy().tolist()
 return result

# Flask app
app = Flask(__name__)
server = InferenceServer('model.pth', device='cuda')

@app.route('/predict', methods=['POST'])
def predict():
 try:
 # Parse input
 data = request.json['data']

 # Preprocess
 tensor = server.preprocess(data)

 # Infer
 output = server.infer(tensor)

 # Postprocess
 result = server.postprocess(output)

 return jsonify({'prediction': result})

 except Exception as e:
 logging.error(f"Error: {e}")
 return jsonify({'error': str(e)}), 500

# Run
# python -m flask run
# curl -X POST http://localhost:5000/predict -d '{"data": [1,2,3,4,5,6,7,8,9,10]}' -H "Content-Type: application/json"

---

Performance Checklist

# Inference optimization checklist

def optimize_inference(model, device='cuda'):
 """Apply standard optimizations."""

 # 1. Move to GPU
 model = model.to(device)

 # 2. Set eval mode (disables dropout, batchnorm)
 model.eval()

 # 3. Convert to mixed precision
 if device == 'cuda':
 model = model.half()

 # 4. Fuse operations (BatchNorm + Conv, etc.)
 torch.nn.utils.fusion.fuse_conv_bn_eval(model)

 # 5. Quantize if needed
 # model = torch.quantization.quantize_dynamic(model)

 # 6. TorchScript compilation
 example_input = torch.randn(1, 10).to(device)
 try:
 model = torch.jit.trace(model, example_input)
 except:
 pass # Not all models can be traced

 return model

# Benchmark improvements
model = optimize_inference(model)

# Expected improvements:
# - Mixed precision
# - TorchScript
# - Quantization
# - Combined

---

Summary: Inference vs Training

Aspect	Training	Inference
Throughput	Maximize	Lower priority
Latency	Lower priority	Critical
Memory	Per-batch	Per-request
Precision	High (fp32)	Can reduce (fp16/int8)
Caching	Not used	Essential (KV cache)
Batching	Fixed	Dynamic
Optimization	Modest	Extreme

-

Related Topics

• [04 Context Managers & Resource Management](/05-py3/01-fundamentals/(04-context-managers-resource-management/) - GPU resource management
• 00 Readme - Memory-efficient computation
• [04 Profiling & Performance Analysis](/05-py3/09-bytecode-and-execution/(04-profiling-performance-analysis/) - Profiling inference
• 05 Custom Operators - Optimized kernels