The design space

Where the weights live, who decides when data moves, and how big the engines are: a map of the families of AI chips, and the decoder layer every chapter runs.

Loading the animation…

Concept

Every AI chip does the same sums: multiply a vector of activations by matrices of weights, again and again, layer after layer. What separates the families is not the arithmetic but three choices around it.

Where the weights live. A GPU or a TPU keeps them in off-chip DRAM (HBM) and streams them in for every step. A wafer-scale engine or a Groq-style LPU keeps them in on-chip SRAM, spread across as many chips as it takes to hold them. Dataflow chips and meshes of small cores use both tiers, and in-memory designs compute where the weight is stored.

Who decides when data moves. A GPU's hardware schedules warps at run time and its L2 cache decides what stays close. A TPU or an LPU is fed by a compiler that fixed every transfer before the program started. In between, a wafer's cores fire when data arrives over routes the compiler set up.

How big the engines are. A few large matrix units (a TPU's arrays), a hundred or so cores (a GPU's SMs, a mesh's Tensix cores), or 900,000 tiny cores on one wafer.

The map places the eight families on the first two axes, with the third drawn as each marker's grain. The same running workload follows us through every chapter: a Llama decoder generating text. Its prefill (reading a whole prompt at once) is a large matrix multiply; its decode (one new token per step) reads every weight to do two flops with each. So the first question for any chip is the one under the map: where do 70,548,193,280 weights fit? In 8-bit form, with one 1,024-token KV cache, Llama 3.1 70B needs 70.7 GB: one H100 (80 GB of HBM), but 308 first-generation LPUs at 230 MB of SRAM each. In 16-bit form it needs 4 wafers, which is the number Cerebras published.

Family	Weights live in	Who schedules	Grain
GPUsNVIDIA H100, B200	off-chip DRAM[hopper-blog]	hardware, at run time[hopper-blog]	tens to hundreds of cores[hopper-blog]
Systolic designsGoogle TPU, AWS Trainium	off-chip DRAM[tpu-7x]	compiler, ahead of time[tpu-arch]	a few large engines[tpu-arch]
Wafer-scaleCerebras WSE-3	on-chip SRAM[cerebras-inference]	mixed[cerebras-wp]	hundreds of thousands of cores[cerebras-wp]
Compiler-scheduledGroq LPU	on-chip SRAM[groqchip-brief]	compiler, ahead of time[groq-isca22]	a few large engines[groq-isca22]
Reconfigurable dataflowSambaNova RDU, Graphcore IPU	tiered SRAM + DRAM[sn40l]	compiler, ahead of time[sn40l]	tens to hundreds of cores[ipu-dissect]
Meshes of small coresTenstorrent Blackhole	tiered SRAM + DRAM[tt-hardware]	mixed[tt-hardware]	tens to hundreds of cores[tt-hardware]
FPGAs and edge NPUsAMD Versal AI Engines, Apple Neural Engine	tiered SRAM + DRAM[amd-aie]	compiler, ahead of time[amd-aie]	tens to hundreds of cores[amd-aie]
In-memory and analoguephase-change crossbars (research)	in the memory array[pcm-imc]	compiler, ahead of time[pcm-imc]	tens to hundreds of cores[pcm-imc]

One decoder layer, the running workload

Every chapter runs the same decoder: Llama 3.1 8B and 70B, with the layer shapes Meta's Llama 3 paper gives (its Table 3), and Llama 2 70B where a published result used it. One layer of the 70B model has 855,638,016 weights in four matrices: the attention projections (queries, and the keys and values of 8 grouped heads) and the gated MLP. A decode step multiplies one activation vector per sequence by all of them, so a batch of one does about two flops per weight read: 2.03 flop per byte at 8 bits. A prefill of 2,048 tokens reuses each weight 2,048 times. The chapters follow those two steps on each family.

Maths

For a model of $P$ parameters at $b$ bytes each, serving $B$ sequences of $T$ tokens with $k_{\text{KV}}$ bytes of KV cache per token:

n_{\text{chips}} = \left\lceil \frac{P b + B T k_{\text{KV}}}{M_{\text{chip}}} \right\rceil, \qquad k_{\text{KV}} = 2 \, L \, h_{\text{KV}} \, d_{\text{head}} \cdot b_{\text{KV}}.

For Llama 3.1 70B, $L = 80$ layers, $h_{\text{KV}} = 8$ and $d_{\text{head}} = 128$ , so $k_{\text{KV}}$ is 163,840 elements per token. $M_{\text{chip}}$ is the chip's HBM when it has some, and its SRAM when it has none.

Code

Where the weights go, cut from src/lib/accel/model.ts:

const off = fval(pid, "off_bytes");
let where: "off-chip" | "sram";
let cap: number;
let bw: number;
if (off !== null && off > 0) {
  where = "off-chip";
  cap = off;
  bw = fval(pid, "off_bw") as number;
} else {
  where = "sram";
  cap = fval(pid, "sram_bytes") as number;
  bw = fval(pid, "sram_bw") as number;
}
const chips = Math.ceil(need / cap);