Hardware Optimization for ML

A neural network is mostly one operation done billions of times: multiply a matrix of inputs by a matrix of weights. Which chip you run that on changes your speed by orders of magnitude, because the chips are built for different jobs.

• CPU — a handful of very fast, flexible cores. Brilliant at branchy control code and small models, but it does those huge matrix multiplies one small chunk at a time. Use it for tiny models, preprocessing, or when no GPU is available.
• GPU — thousands of simpler cores that crunch the same matrix multiply in parallel. This is the default for both training and inference. NVIDIA GPUs add Tensor Cores, dedicated units for fast FP16/INT8 matrix maths.
• TPU — Google's chip built almost entirely around a giant matrix-multiply array plus very high-bandwidth memory. Superb on large transformers and CNNs and very power-efficient, but less flexible than a GPU and tied to Google's stack.
• Accelerators / NPUs — purpose-built inference chips (AWS Inferentia, Apple Neural Engine, edge NPUs). They sacrifice generality for efficiency on a fixed set of operations, ideal on phones and at the edge.

Beginners assume a model is slow because the chip can't do enough maths (not enough FLOPs, floating-point operations per second). For large-model inference the opposite is usually true: the chip is starving, waiting for weights to arrive from memory. This is the memory-bandwidth bottleneck.

Here's why. To generate one token, an LLM has to read every weight in the model from memory, but only does a tiny bit of arithmetic with each one. So the limit isn't compute — it's how many bytes per second you can stream out of memory (the bandwidth). A useful back-of-the-envelope estimate for single-stream speed is:

# Single-stream LLM speed is bounded by memory bandwidth.
tokens_per_sec  ≈  memory_bandwidth  /  model_size_in_bytes

# Example: a 7B model in FP16 is ~14 GB. On a GPU with 2 TB/s:
#   2,000 GB/s  /  14 GB  ≈  143 tokens/sec

Two big consequences follow. First, shrinking the model (lower precision, Section 4) directly raises speed because there are fewer bytes to read. Second, batching (Section 3) helps because you read each weight once and reuse it for the whole batch — turning a bandwidth-bound job into a compute-bound one that keeps the cores busy.

Running one input at a time wastes the hardware: you pay the cost of reading the weights but only process a single item. Batching stacks many inputs and runs them together, so the expensive weight read is shared across the whole batch. Two metrics matter, and they pull in opposite directions:

• Throughput — items processed per second. Batching raises it sharply, because the per-batch overhead is amortized over more items.
• Latency — time for one request to finish. Batching increases it, because each request waits for the batch to fill and complete.

The model is simple: a batch costs a fixed overhead (kernel launch, reading weights) plus a small per-item compute cost. Throughput is batch / time_per_batch. It climbs as the batch grows — until the hardware saturates, after which throughput flattens while latency keeps rising. The job is to pick the batch size that maximizes throughput within your latency budget.

Every weight is stored in a number format, and the format's bit width decides how many bytes you move and how fast the hardware runs. Cutting precision is the single easiest large speedup in inference.

Because speed scales with bytes read, going FP32 → FP16 roughly doubles tokens/sec and FP16 → INT8 roughly doubles it again. The catch is accuracy: lower precision rounds more, so always benchmark the quantized model on your task before shipping. You'll feel this effect directly in Your Turn #2.

A model is a graph of small operations — a matrix multiply, then add a bias, then a ReLU activation. Run naively, each operation launches its own GPU kernel (a unit of GPU work) and writes its result back to slow memory before the next one reads it. Those round-trips dominate the time.

Operator fusion merges a chain of ops into a single kernel that keeps the intermediate values in fast registers — e.g. matmul + bias + ReLU becomes one fused kernel, eliminating the memory round-trips. You rarely write fusion by hand; graph compilers do it for you:

• torch.compile — one line in PyTorch. Traces, fuses, and picks fast kernels. ~1.3-2x for free.
• ONNX Runtime — export to ONNX (the universal model format), then run with cross-platform graph optimization on CPU or GPU.
• TensorRT — NVIDIA's inference compiler. Aggressive fusion, FP16/INT8 kernels, and kernel auto-tuning for ~3-5x on NVIDIA GPUs.
• OpenVINO / CoreML — the equivalents for Intel CPUs and Apple Silicon.

The usual pipeline: PyTorch/TF model → export to ONNX → compile with the right backend for your hardware → deploy. The two worked examples below show torch.compile and a TensorRT FP16 export.

The most expensive mistake in this whole lesson is optimizing by guessing. The bottleneck is rarely where you assume — a slow model is often waiting on data loading or Python overhead, not compute. Profile first.

# PyTorch profiler — see where the time actually goes.
import torch
from torch.profiler import profile, ProfilerActivity

model = torch.nn.Linear(1024, 1024).eval()
x = torch.randn(64, 1024)

with profile(activities=[ProfilerActivity.CPU]) as prof:
    with torch.no_grad():
        for _ in range(10):
            model(x)

# Print the ops that took the most total time.
print(prof.key_averages().table(sort_by="cpu_time_total", row_limit=5))

# The table shows each op's total time — read it to find the real
# hotspot before you change a single line of model code.

For GPU work, NVIDIA Nsight Systems shows whether the GPU is busy or idle and whether transfers overlap compute. Low GPU utilization with slow inference almost always means the GPU is waiting: tiny batches, data-loading stalls, or per-call Python overhead. Fix what the profiler points at — then measure again.

One line wraps your model and gives you fused kernels with no accuracy change. The first call compiles (slow); every call after runs the optimized graph.

# Graph compilation with torch.compile — one line, no accuracy change.
# It traces the model, fuses ops into bigger kernels, and picks the
# best kernel for your GPU. Typical inference speedup: 1.3-2x.

import torch
import torch.nn as nn

model = nn.Sequential(
    nn.Linear(1024, 1024), nn.ReLU(),
    nn.Linear(1024, 1024), nn.ReLU(),
    nn.Linear(1024, 10),
).eval()

x = torch.randn(64, 1024)

# Compile the graph. The FIRST call is slow (it compiles); every call
# after that runs the optimized, fused kernels.
fast_model = torch.compile(model)

with torch.no_grad():
    out = fast_model(x)        # warm-up: triggers compilation
    out = fast_model(x)        # now running the fused kernels

print("output shape:", tuple(out.shape))

# Expected output:
#   output shape: (64, 10)

The production pipeline in code: export to ONNX, then build a TensorRT engine with FP16 enabled. The FP16 flag is the speed switch that turns on half-precision kernels.

# TensorRT: NVIDIA's inference compiler. Export to ONNX, then let
# TensorRT fuse ops, pick fast kernels, and run in FP16 for ~3-5x.

import torch
import torch.nn as nn

model = nn.Sequential(nn.Linear(512, 512), nn.ReLU(), nn.Linear(512, 10)).eval()
dummy = torch.randn(1, 512)

# 1) Export the PyTorch model to ONNX (the universal interchange format).
torch.onnx.export(model, dummy, "model.onnx", input_names=["x"],
                  output_names=["y"], dynamic_axes={"x": {0: "batch"}})

# 2) Build a TensorRT engine from the ONNX file with FP16 enabled.
#    (Pseudocode — TensorRT's Python API in one place for clarity.)
import tensorrt as trt
logger  = trt.Logger(trt.Logger.WARNING)
builder = trt.Builder(logger)
config  = builder.create_builder_config()
config.set_flag(trt.BuilderFlag.FP16)   # <- the speed switch

print("ONNX exported, TensorRT FP16 engine configured")

# Expected output:
#   ONNX exported, TensorRT FP16 engine configured

These five mistakes sink most first attempts at optimizing inference: