Hardware

Google DeepMind’s Free Scaling Book Explains the Memory Wall Its New TPUs Are Sold On

Updated

  • Book Review
  • AI Chips
  • Google TPU
  • AI Inference
Google DeepMind’s Free Scaling Book Explains the Memory Wall Its New TPUs Are Sold On
Uploaded by the desk

How to Scale Your Model, a free online textbook by Google DeepMind engineers, shows with simple arithmetic why serving large language models is limited more by memory bandwidth than by raw compute. A day after Google launched a TPU it pitches at “the memory wall”, the book is the best tool available for checking claims like that. It is also a reminder of whose chips it was written to explain.

On 9 November 2026 Google announced the TPU 8tx and TPU 8i Max, its third and fourth eighth-generation TPU systems this year. The company said the 8i Max “is built for a different problem: the memory wall,” and that it delivers “double the compute, double the SRAM and double the HBM bandwidth of TPU 8i.” Chip launches now use this vocabulary. Buyers hear about bandwidth, context windows and price-performance per token, but they rarely get the means to test any of it.

How to Scale Your Model supplies those means. Jacob Austin, Sholto Douglas, Roy Frostig, Anselm Levskaya, Charlie Chen, Sharad Vikram, Federico Lebron, Peter Choy, Vinay Ramasesh, Albert Webson and Reiner Pope published it online on 4 February 2025. It is free to read, and its source is on GitHub. A chapter on Nvidia GPUs, credited to Austin, Swapnil Patil, Adam Paszke and Pope, followed in August 2025, and the TPU chapter was revised for the seventh-generation Ironwood chip (TPU7x) on 19 October 2026. For this review I read every chapter of the current web edition, about 75,000 words of text including the worked problems. It is the most useful thing I know for reading AI-hardware marketing critically. It is also openly a Google book about Google’s hardware, and that shapes it more than its tone suggests.

What the book argues

The thesis comes in the introduction: “Relatively simple principles apply everywhere — from dealing with a single accelerator to tens of thousands.” Everything rests on what the authors call roofline analysis. Any computation is bounded by three resources: how fast a chip can do arithmetic, how fast it can move bytes (from its own memory or across the network to other chips), and how much memory it has. Divide the arithmetic by the compute rate, divide the bytes by the bandwidth, and the larger number is roughly how long the operation must take.

The ratio between the two is arithmetic intensity, the number of operations performed per byte moved. Chapter 1 turns it into a rule a reader can remember: “for a bfloat16 matmul to be compute-bound on most TPUs, we need our batch size to be greater than 240 (tokens).” On an Nvidia H100, using Nvidia’s dense (non-sparse) figures, the same calculation gives about 295. Below that threshold, a chip that costs thousands of dollars a month spends most of its time waiting for data.

The authors argue that this is a research problem as much as an engineering one: “A 20% win on benchmarks is irrelevant if it comes at a 20% cost to roofline efficiency. Promising model architectures routinely fail either because they can’t run efficiently at scale or because no one puts in the work to make them do so.”

How it is built

The book has twelve chapters in four parts. Chapters 1 to 3 cover rooflines, the TPU as a chip and as a networked system, and how to multiply matrices that are split across devices. Chapter 4 does what the authors call “all the Transformer math you need to know”: parameter counts, FLOPs and the size of the key-value (KV) cache, the store of past tokens that a model keeps in memory while it generates. Chapters 5 and 7 are the core, on how to parallelise training and inference. Chapters 6 and 8 apply that to Meta’s LLaMA 3, and chapters 9 and 10 cover profiling and the JAX framework. The GPU chapter is a long appendix of about 14,000 words.

The format is a large part of the value. Every chapter ends with problems, with the answers hidden behind a click, and many of them use real spec sheets. One asks the reader to find the compute-bound batch size for an H100 from Nvidia’s published data, and the answer notes that the headline Tensor Core numbers are “twice the true value since they’re only achievable with structured sparsity.” It is a small lesson in reading vendor specifications, which most of the industry needs.

The authors are candid about the drafting. “This series is probably longer than it needs to be,” the introduction says, and “for the time being, this is a draft and will continue to be revised.” The repository’s history bears that out: since September 2026 alone there have been commits adding a from-scratch LLaMA inference exercise, revising the training text and updating the TPU chapter for Ironwood.

Where it is strongest: why serving costs what it does

Chapter 7, on inference, which the acknowledgements say Charlie Chen wrote, is the one to read if you read only one. It separates the two phases of running a model. Prefill processes the prompt in one large, parallel pass and is “basically always compute-bound.” Generation then produces one token at a time, and that is where the economics turn. For attention during generation, the chapter shows that arithmetic intensity collapses to roughly one operation per byte, whatever the batch size. In the book’s words, “during generation our arithmetic intensity is low and constant, so we are always memory bandwidth-bound.” Or, more bluntly: “We’re doing a tiny amount of FLOPs while loading a massive KV cache.”

Chapter 8 then puts numbers on it. LLaMA 3-70B stores about 160kB of KV cache per token even in int8. At a 32,000-token context that is 5.3GB per sequence, and a batch of 240 sequences needs about 1.3TB, against about 70GB for the model’s weights. The book calculates that this alone would take about 86 TPU v5e chips just to hold. Its conclusion for that model is that “we are strongly KV cache memory bandwidth-bound (and HBM-bound) in almost all of these configurations.” An interactive figure shows the trade-off that drives inference pricing: “At the cost of doubling per-token latency, we can achieve a roughly 100x reduction in per-token cost.”

That is the mechanism behind Google’s latest pitch. When Google says long-context agent workloads run into a memory wall, Chapter 7 shows why: each extra token of context adds KV cache that has to be read again for every generated token. Doubling HBM bandwidth per package, as Google says the 8i Max does, attacks the term that dominates the book’s own worked examples. A reader who has done the Chapter 8 problems can see why a chip vendor would build that product. They can also see what the announcement leaves out: how much memory capacity it has, what it costs per hour and what utilisation it achieves in practice. Those numbers decide whether the claim holds.

The book also explains disaggregated serving, which runs prefill and generation on separate pools of chips, and shows why the ratio between those pools depends on prompt and output lengths. In one worked example, prompts of 8,192 tokens and outputs of 512 call for about three prefill servers per generation server. Few public sources explain this so clearly.

Running the book’s numbers on its own price table

Chapter 8 opens with a heuristic: “what hardware should we serve on? The answer is basically, whichever is cheapest in FLOPs / dollar.” It qualifies this at once (“sometimes more HBM or ICI bandwidth is critical rather than FLOPs”) and then compares three chips using Google Cloud on-demand prices from February 2025. The H100 is listed at $10.80 an hour, the TPU v5p at $4.20 and the TPU v5e at $1.20. That gives the v5e about 1.8 times the H100’s bfloat16 FLOPs per dollar and the v5p about 1.2 times.

The book’s own argument, though, is that generation is bound by memory bandwidth. So I redid the comparison by memory bandwidth per dollar, using the book’s price table and its own bandwidth figures (3.35TB/s for the H100, 2.8e12 bytes/s for the v5p and 8.2e11 bytes/s for the v5e). An hour of H100 time works out to about 1.1 petabytes of HBM traffic per dollar. The v5p and v5e both come to about 2.4 petabytes per dollar. By the book’s own reasoning, the v5p’s case for bandwidth-bound decoding is much stronger than its FLOPs-per-dollar figure suggests: it roughly matches the cheaper v5e per dollar and has six times the memory per chip (96GiB against 16GiB).

This is VibeCodAI’s calculation, not the book’s, and it carries the book’s limits. These are list prices from one cloud, Google’s own, on one date. Negotiated and committed-use prices differ, and peak bandwidth is not achieved bandwidth. The point is about method. The tools in the book are good enough to question the book’s own shortcut, and that is the best thing one can say about a textbook.

What it gets wrong, or leaves out

It is a vendor’s view. The book never hides that it was written at Google DeepMind about TPUs, and its subtitle is “A Systems View of LLMs on TPUs.” But a few judgements read as house positions rather than findings. The introduction calls the TPU story “a resounding success.” Chapter 1 says that “H100s and B200s can usually only achieve around 80-85% of the claimed peak FLOPs, while TPUs can get closer to 95% in normal use,” and gives no measurement or citation for either figure. That is plausible, and it may well be true. But it is the kind of comparative claim this publication would not accept from Nvidia without evidence, and it should not accept it from Google either. The authors do flag elsewhere that a raw speed comparison “does not normalize to cost,” and the GPU chapter is commendably sceptical of Nvidia’s numbers, noting that while Nvidia claims about 450GB/s over H100 NVLink, “it is difficult in practice to exceed 370 GB/s.” The same scrutiny is not applied to the TPU figures.

The worked economics are fixed in early 2025. The October 2026 revision moved the TPU chapter’s specifications to TPU7x, with 192GiB of HBM, about 7.4e12 bytes/s of bandwidth and 2.3e15 bfloat16 FLOPs/s per chip. The cost chapters, however, still serve LLaMA 3-70B on the TPU v5e with February 2025 prices. A reader therefore learns the method on hardware several generations old, and nothing in the book covers the eighth-generation chips Google has launched since April. That is understandable for a living document, but it means the book cannot settle the current argument about which accelerator is cheapest. SemiAnalysis, an industry research firm, says benchmark results it first published in September showed Ironwood delivering up to 50% better performance per dollar than Nvidia’s B200 and B300 in an FP8 comparison on one model. The book shows how to reason about results like that. It does not provide them.

Dense models anchor the story. The two applied chapters use LLaMA 3, a dense model, although the frontier has largely moved to mixture-of-experts (MoE) designs, which activate only some of their parameters for each token. The book is not blind to this. One of the best problems in Chapter 4 shows that an MoE needs a much larger batch to keep its chips busy, working out a threshold of 120 × 256 ÷ 8, or 3,840 tokens, for a DeepSeek-style model with int8 weights on TPU v5e. The book calls that “a remarkably large batch size at generation time.” The GPU chapter discusses expert parallelism, and DeepSeek V3’s training setup, at length. But there is no full worked serving example for an MoE model. Chapter 7’s list of KV-cache reduction techniques (grouped-query attention, local attention, KV sharing across layers, quantisation and paged attention) also omits multi-head latent attention, the compression scheme associated with DeepSeek’s models.

It is narrow on hardware. “GPUs” in the book means Nvidia GPUs. AMD gets a passing mention and Amazon’s Trainium none. Readers who need to compare the three platforms that a lab such as Anthropic says it uses (TPUs, Trainium and Nvidia GPUs) will have to do the extension themselves. The roofline method transfers easily; the network topologies the book describes in detail do not.

The theory is cleaner than practice, and the book admits it. Chapter 7 lists the ways real systems miss their rooflines, including that perfect overlap of memory reads and computation “is not realistic, since our compiler (XLA) is fallible.” That candour helps. It also means the book’s numbers are best treated as limits and order-of-magnitude estimates, which is how the authors present them.

Who it is for

The stated audience is engineers and researchers who know the basics of Transformers and, ideally, some JAX. That is accurate. Most chapters need comfort with matrix shapes and some algebra. But the first chapter and the inference chapters are readable for a technically minded analyst, investor or journalist who is willing to work through a few problems. For that reader, the payoff is large. After Chapter 8 you can look at a claim about tokens per second, context length or cost per million tokens and ask which resource is the bottleneck, and whether the numbers add up.

It pairs naturally with Vijay Janapa Reddi’s open machine-learning systems textbook, which we reviewed this week and which deliberately stops at a single machine. How to Scale Your Model starts more or less where Reddi’s first volume ends: thousands of chips, linked by networks with finite bandwidth.

How it meets the moment

The contest over AI accelerators has become a contest over who can serve tokens most cheaply, and Google has moved from keeping TPUs mostly for itself to supplying large blocks of capacity to outside labs. Anthropic said in October 2025 that it would use up to one million TPUs and in April 2026 that it had signed for multiple gigawatts more from 2027. Claims about memory walls and price-performance will keep coming from every vendor, Google included.

That makes this book more important than when it appeared, and it makes its gaps matter more. It teaches the right questions about where the bytes go, why decoding is memory-bound and what the KV cache costs. The answers it works through are for the v5e, LLaMA 3 and February 2025 prices. Read it to learn the method, apply the same scepticism to its TPU claims that it applies to Nvidia’s, and keep checking whether the October revision is followed by updated cost chapters. Until then, the book tells you how to evaluate Google’s new chips but not what the answer is.