Memory Layout and Data Movement on Accelerators

Two programs can run the same model and feel very different in speed and running cost. Part of the reason is the software they are written in. A large part, however, is the hardware underneath them. Some optimization techniques work on any machine, while others only pay off once you understand the machine you are running on.

This article starts from the hardware side. It looks at the chips that run AI workloads and asks what those chips are built to do. The question it returns to again and again is simple: when a computation has a particular shape, which chip design uses its resources well?

The word specific is the key to the definition. A general-purpose processor is built to handle a wide mix of tasks. An accelerator gives up some of that breadth in exchange for doing one kind of work well. Within the family of AI accelerators, the dominant type today is the GPU, the graphics processing unit.

GPUs are not special because they are fast at everything. They are special because their design fits the arithmetic that AI models repeat constantly. The source also points to NVIDIA as the biggest economic driver of the AI boom in the early 2020s, which reflects how central this chip type became.

The link between models and hardware is older than the current boom. A shortage of sufficiently powerful computers was one factor in the first AI winter in the 1970s. The return of interest in deep learning in 2012 was closely tied to compute. One commonly cited reason for the influence of the AlexNet paper is that it was the first to successfully use GPUs to train neural networks. Before that, training a model at AlexNet's scale generally meant running on thousands of CPUs. A couple of GPUs were far more affordable for PhD students and researchers than thousands of CPUs, and that accessibility helped set off the deep learning research boom.

Hardware is also shaped by the kinds of data a computation handles. The source identifies three shapes. A scalar is a single value. A vector is an ordered list of values. A tensor generalizes these to higher-dimensional arrays. Compute units are specialized for one or more of these shapes, and chips often mix them. GPUs include tensor cores, and TPUs are built around tensor operations.

The central difference between a CPU and a GPU is their purpose. A CPU is designed for general-purpose use, meaning it is meant to handle a wide mix of tasks. A GPU is designed for parallel processing, meaning it is meant to do many pieces of work at the same time. Most of the other differences between the two follow from that single choice.

A CPU has a few powerful cores. Many of them handle multi-threaded workloads well, but they are at their best when a task needs high single-thread performance. Running an operating system, managing input and output operations, and handling complex sequential processes all fall into this category.

A GPU takes the opposite approach. It has thousands of smaller, less powerful cores. Those cores are optimized for work that can be broken into many small, independent calculations, such as graphics rendering and machine learning. Each core is weaker than a CPU core, but the sheer number of them is what matters when the work splits cleanly.

GPUs are optimized for tasks that can be broken into many smaller, independent calculations. The operation that makes up most machine learning workloads is matrix multiplication, and it is highly parallelizable. The key property is independence. If one piece of the answer does not need another piece's result, the two pieces can be worked out at the same time on different cores.

The same logic scales up. A real model multiplies matrices far larger than two by two, which means far more independent entries. This is why a chip with thousands of cores can keep many of them busy on one multiplication.

Consider what happens as a batch of independent matrix entries arrives at each kind of chip, starting with the CPU. The CPU has a few powerful cores, so the batch is spread across a small number of strong workers. Each worker takes on a share of the entries. Because each core is powerful, it handles its share with high single-thread performance. The batch is finished by a handful of workers rather than a crowd.

Now run the same batch on a GPU. The same entries are spread across thousands of smaller cores. Each core is less powerful, but because the entries are independent, none of those cores has to wait for another. At any moment, many entries are in progress together. The state of the batch changes from a few entries being worked on slowly to a large set being worked on at once.

Notice what changed and what did not. The work itself is the same, and the answer is the same. What changed is how the work is spread across hardware. That is the core idea behind matching a chip to a workload: the shape of the computation determines which design makes good use of its resources.

Many chips can handle both training and inference, but one clear trend is the rise of chips specialized for inference. Inference means running a trained model to produce outputs. A survey by Desislavov and colleagues (2023) reports that inference can exceed the cost of training in commonly used systems, and that inference accounts for up to 90 percent of machine learning costs for deployed AI systems. For anyone deploying models, that makes inference hardware a central decision.

Inference hardware is often optimized for lower precision and faster memory access. The reason is that inference aims to minimize latency, meaning the delay before an output is produced. A chip that is ideal for building a model is not automatically ideal for serving it, because the priorities change once the model is in use.

A chip's specification sheet holds many details, and only some of them matter for a given use case. Across use cases, three characteristics stand out. They should be weighed together rather than read one at a time.

Computational capability describes how much work the chip can do. Memory size and bandwidth describe how much data it can hold and how quickly that data can move. Power consumption describes the energy it draws while doing that work. Because each one is a different kind of limit, improving one does not automatically improve the others. A chip that looks strong on computation may still be a poor fit if its memory cannot keep up with the work, and a chip that looks efficient may not have the capability a workload needs.

A CPU and a GPU differ in purpose before they differ in anything else. The CPU is built for a wide mix of tasks with a few powerful cores. The GPU is built for parallel work with thousands of smaller cores. Matrix multiplication, the core operation of most machine learning workloads, splits into independent entries, which is exactly the shape a GPU is built to exploit. Inference hardware is often tuned for lower precision and faster memory access, and the three characteristics of computation, memory, and power should be weighed together when choosing a chip.