Efficiency Frontiers › Hardware Co-Design

Algorithm design alone does not determine efficiency. The same algorithm can be ten times faster or slower depending on how well it matches the underlying hardware. Understanding what makes AI accelerators fast is essential for designing architectures that deliver their theoretical efficiency gains.

Modern AI training and inference happens primarily on GPUs (and increasingly on specialized AI chips like TPUs, H100s, and Trainium). Understanding their architecture explains why some operations are fast and others are not.

A GPU contains thousands of CUDA cores organized into streaming multiprocessors (SMs). All these cores share a memory hierarchy with dramatically different bandwidth at each level:

• Registers: fastest, local to each thread, limited size
• Shared memory / L1 cache: on-chip, ~10-50 MB total, ~100x slower than registers
• L2 cache: on-chip, ~40-80 MB, shared across SMs
• HBM (High Bandwidth Memory): off-chip, ~80 GB, bandwidth ~3 TB/s on H100

The critical bottleneck for most transformer operations is memory bandwidth, not compute. Moving data from HBM to the compute cores is slow relative to the arithmetic operations themselves. An operation is compute-bound if the GPU's arithmetic units are the bottleneck, and memory-bound if data movement is the bottleneck. On modern hardware, most attention operations are memory-bound.

This memory-bandwidth bottleneck explains why attention is slower in practice than its FLOP count suggests. The standard attention computation reads the $Q$, $K$, and $V$ matrices from HBM, writes the intermediate attention matrix back to HBM, then reads it again to multiply with $V$. Every round trip to HBM costs time.

FlashAttention is one of the most impactful efficiency innovations in the transformer era. It computes exact (not approximate) attention using a tiling algorithm that minimizes HBM reads and writes.

The key insight is that the attention computation:

$$O = \text{softmax}\!\left(\frac{QK^T}{\sqrt{d_k}}\right)V$$

does not need to materialize the full $n \times n$ attention matrix $S = QK^T / \sqrt{d_k}$ in HBM. Instead, FlashAttention tiles the computation into blocks that fit in SRAM (shared memory). It processes one tile at a time, computing a local softmax and updating the output matrix, keeping track of the normalization constants needed for the global softmax using an online algorithm.

The online softmax computation maintains two accumulators per query: the running maximum $m_i$ and the running sum of exponentials $l_i$. When a new tile of keys arrives with scores $s$, the global softmax is updated without ever materializing the full row:

$$\begin{aligned} m_i^{\text{new}} &= \max(m_i^{\text{old}}, \max(s)) \\ l_i^{\text{new}} &= e^{m_i^{\text{old}} - m_i^{\text{new}}} l_i^{\text{old}} + \sum_j e^{s_j - m_i^{\text{new}}} \\ O_i^{\text{new}} &= \frac{e^{m_i^{\text{old}} - m_i^{\text{new}}} l_i^{\text{old}} O_i^{\text{old}} + \sum_j e^{s_j - m_i^{\text{new}}} V_j}{l_i^{\text{new}}} \end{aligned}$$

where:

• $m_i$ is the running maximum of attention logits for query $i$
• $l_i$ is the running sum of exponentiated, shifted logits
• $O_i$ is the partial output for query $i$

By keeping all intermediate values in fast SRAM rather than writing to HBM, FlashAttention reduces HBM reads/writes from $O(n^2)$ to $O(n)$, making it 2-4x faster than standard attention in practice and enabling context lengths that would otherwise cause out-of-memory errors.

FlashAttention-2 extended this by improving work partitioning across thread blocks, reducing non-matrix-multiplication FLOPs, and achieving ~70% GPU utilization compared to ~35% for the original. FlashAttention-3 specifically targets the H100 architecture, exploiting its asynchronous pipeline and FP8 support.

Training very large models requires distributing computation across multiple GPUs. The two main strategies are tensor parallelism and pipeline parallelism.

Tensor parallelism splits individual weight matrices across GPUs. For a feed-forward layer $y = xW + b$ with $W \in \mathbb{R}^{d \times 4d}$, we split $W$ column-wise across $p$ GPUs:

$$W = [W_1 | W_2 | \ldots | W_p] \quad \text{where } W_i \in \mathbb{R}^{d \times (4d/p)}$$

Each GPU computes its portion $xW_i$, and the results are gathered with an all-reduce operation. This allows a single layer to use more GPUs than would fit in a single device's memory. The communication cost is an all-reduce per layer per forward and backward pass.

Pipeline parallelism splits the model's layers across GPUs. GPU 1 holds layers 1-8, GPU 2 holds layers 9-16, etc. The input flows through GPUs in a pipeline. Without careful scheduling, most GPUs sit idle waiting for their predecessor. The GPipe pipeline fills this bubble by processing multiple micro-batches simultaneously, so each GPU always has work to do.

ZeRO (Zero Redundancy Optimizer) is a memory optimization that partitions optimizer states, gradients, and parameters across data-parallel workers instead of replicating them. The basic version reduces optimizer state memory by a factor equal to the number of workers, which can be 64x for large training runs.

Hardware supports different numerical formats with different tradeoffs between precision and memory:

• FP32: 32-bit float, standard training format, 4 bytes per value
• BF16: 16-bit bfloat, same exponent range as FP32, 2 bytes per value, now the standard for training
• FP16: 16-bit float, smaller exponent range, prone to overflow, requires loss scaling
• FP8: 8-bit float, supported on H100, approximately 2x throughput over FP16
• INT8: 8-bit integer, well-supported for inference on many platforms
• INT4: 4-bit integer, increasingly common for inference, requires careful quantization

The NVIDIA H100 GPU can perform FP8 matrix multiplications at approximately twice the throughput of FP16, with 2x more values fitting in the same memory bandwidth. Exploiting these low-precision formats requires that the model's activations and weights stay within the representable range.

Quantization-aware training (QAT) simulates low-precision arithmetic during training by inserting fake quantization operations:

$$x_q = \text{round}\!\left(\frac{x}{\Delta}\right) \cdot \Delta \quad \text{where } \Delta = \frac{x_{\max} - x_{\min}}{2^b - 1}$$

where $b$ is the number of bits and $\Delta$ is the quantization step size. During the forward pass, values are quantized. During the backward pass, gradients flow through using the straight-through estimator (the gradient of the rounding operation is treated as 1). This allows the model to learn representations that retain accuracy under low-precision arithmetic.