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: Transformer autoregressive generation recomputes keys/values
# 
# Without cache:
# Token 1: Compute K, V for token 1
# Token 2: Compute K, V for tokens 1-2 (redundant!)
# Token 3: Compute K, V for tokens 1-3 (very redundant!)
# → O(n²) computation
#
# With cache:
# Token 1: Compute K, V for token 1, save
# Token 2: Reuse cached K1,V1 + compute K2,V2
# Token 3: Reuse cached K1,V1,K2,V2 + compute K3,V3
# → 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: 32
# Sequence length: 512
# Hidden size: 768
# Data type: float32 (4 bytes)

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

# Without cache: recompute everything
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: save K, V
cache_memory = seq_len * batch_size * hidden_size * dtype_bytes * 2  # K + V
print(f"Cache memory: {cache_memory / 1e6:.2f}MB")  # ~402MB

# Savings: 3.1GB - 402MB = 2.7GB (87% reduction)

---

Mixed Precision Inference

Float16 Inference

# Float32: larger, slower, more accurate
# Float16: smaller, faster, less accurate
# 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: 2x (4 bytes → 2 bytes per value)
# Speed improvement: 1.5-2x (depending on GPU)
# Accuracy impact: Usually <0.1% top-1 accuracy loss

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: run on representative data
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: [10, 15, 20, 25, 30, 35, 40, 45, 50, 100]
buckets = create_buckets([10, 15, 20, 25, 30, 35, 40, 45, 50, 100], bucket_size=50)

# Result:
# Bucket 0: indices for lengths 10-50 (pad to 50)
# Bucket 50: indices for lengths 50-100 (pad to 100)
# → 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: 1.5-2x faster
# - TorchScript: 1.2-1.5x faster
# - Quantization: 4x faster (CPU) or 2x (GPU)
# - Combined: 3-8x total speedup

---

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 - GPU resource management
• 00 Readme - Memory-efficient computation
• 04 Profiling & Performance Analysis - Profiling inference
• 05 Custom Operators - Optimized kernels