AI Inference Hardware

AI inference hardware refers to the processors and accelerators purpose-built or optimized for running trained AI models at production scale. While training a model requires powerful hardware over days or weeks, inference runs continuously — handling millions of requests per day. Hardware choice has a direct impact on cost, latency, and throughput.

Why Dedicated Hardware Matters

Neural networks are dominated by one operation: matrix multiplication. A standard forward pass through a transformer layer involves:

$$Y = XW^\top$$

Where $X \in \mathbb{R}^{B \times d}$ (batch × hidden dim) and $W \in \mathbb{R}^{d_{out} \times d}$. This is a dense GEMM (General Matrix Multiply) operation that must be performed thousands of times per inference call.

General-purpose CPUs are designed for low-latency, sequential, branchy code — not for the massively parallel matrix arithmetic that AI demands. Dedicated accelerators achieve 10–1000× better performance per watt for AI workloads.

GPUs: Graphics Processing Units

GPUs were originally designed for real-time 3D graphics — which also requires massively parallel matrix math. NVIDIA recognized this overlap early and built the CUDA programming model, making GPUs the dominant AI hardware platform.

GPU Architecture

Tensor Cores are the key innovation: they perform a $4 \times 4$ matrix multiply-accumulate in a single clock cycle in various precisions (FP32, BF16, FP16, INT8, FP8).

Key NVIDIA GPU Generations

Memory bandwidth is often the binding constraint for LLM inference — moving model weights from HBM to compute units is the bottleneck, not the computation itself.

Consumer vs. Data Center GPUs

Consumer GPUs (RTX 4090, RTX 5090) have less memory bandwidth, no HBM, and lack enterprise reliability features but are cost-effective for development and small-scale inference. The RTX 4090 (24 GB GDDR6X) can run 7B–13B parameter models locally.

TPUs: Tensor Processing Units

TPUs are Google’s custom ASIC (Application-Specific Integrated Circuit) designed specifically for neural network workloads. They are available exclusively through Google Cloud and are used internally for training and serving Google’s AI products.

TPU Architecture

TPUs are built around a Systolic Array — a grid of multiply-accumulate units that passes data through in a wave, eliminating the need to load/store intermediate results to memory.

The systolic array computes:

$$C = A \cdot B$$

by streaming $A$ row-wise and $B$ column-wise through the grid simultaneously, accumulating partial sums as they flow. This is extremely efficient for GEMM but less flexible than GPU compute.

TPU v5p and v5e (2024)

TPU pods interconnect thousands of chips into a single logical accelerator, enabling training and inference at scales impossible with GPU clusters of equivalent size.

NPUs: Neural Processing Units

NPUs are low-power AI accelerators integrated into consumer devices — smartphones, laptops, and edge devices. Their goal is on-device inference: running AI models locally without cloud connectivity.

NPU Examples

NPUs are optimized for INT8/INT4 quantized inference and excel at models that fit in device memory (typically 1–8 GB for mobile NPUs).

TOPS (Tera Operations Per Second) measures INT8 throughput, which is the primary metric for NPU comparisons.

Purpose-Built AI Chips

Several companies have designed chips specifically for LLM inference:

Groq LPU (Language Processing Unit)

Groq’s LPU uses a deterministic, compiler-scheduled architecture with no caches or branch predictors. All memory accesses are statically scheduled at compile time, eliminating the memory latency variability that degrades GPU LLM throughput.

Result: extremely high token throughput — Groq systems have demonstrated 500+ tokens/second on Llama models vs. 50–100 tokens/second on equivalent GPU setups.

AWS Trainium and Inferentia

• AWS Trainium 2: Amazon’s custom training chip — up to 3.5× better price-performance than GPU instances for training.
• AWS Inferentia 2: Inference-optimized; up to 4× lower latency and 10× higher throughput vs. comparable GPU instances for deployed models.

Cerebras Wafer-Scale Engine

Cerebras builds chips at wafer scale — a single chip the size of an entire silicon wafer (46,225 mm², vs ~800 mm² for an H100). This provides:

• 900,000 cores.
• 44 GB on-chip SRAM (no off-chip memory bottleneck).
• 20 PB/s memory bandwidth.

The elimination of off-chip memory is particularly valuable for LLM inference, where weight loading is the primary bottleneck.

Memory: The Critical Bottleneck

For LLM inference, memory capacity and bandwidth matter more than raw FLOPS:

• A 70B parameter model in FP16 requires ~140 GB of memory just to load weights.
• Each forward pass requires reading all weights from HBM to compute units.
• At 2 TB/s bandwidth, reading 140 GB takes ~70ms — limiting throughput even if compute is idle.

Strategies to address the memory wall:

• Quantization (INT8, INT4): Halves or quarters weight size.
• Continuous batching: Batch multiple requests to amortize weight reads.
• Speculative decoding: Use a small draft model to propose tokens, reducing large model calls.
• KV cache management: Efficiently manage the KV cache to maximize GPU memory utilization.

Choosing Hardware for AI Workloads

Further Reading

• NVIDIA H100 Architecture Whitepaper
• Google TPU v5p Cloud Documentation
• Groq LPU Architecture Overview
• Flash Attention: Fast and Memory-Efficient Attention — Dao et al., 2022