GPUs, TPUs and why AI is slow where it is: compute, memory bandwidth, the KV cache and the software in between

When a chatbot pauses before its first word and then types the rest at a steady pace, you are watching two different hardware limits in action. The pause is mostly maths. The steady typing is mostly memory. Once I understood that, a lot of confusing advice about AI performance (bigger batches, caching, quantisation, special chips for inference) suddenly made sense.

This post is the hardware explainer I wish I had read first. I start with what sits inside an AI server, then explain the difference between being compute-bound and memory-bandwidth-bound with a kitchen and a few small sums. From there I look at why training, prefill and decode stress different parts of the chip, what the KV cache and prompt caching actually hold, how confidential computing protects data while a chip is working on it, and how the software stack (CUDA, XLA, Pallas, PyTorch, JAX, vLLM) sits on top. I keep the ideas vendor-neutral first and then map them to Google Cloud. Every product fact is dated November 2026, because this area changes monthly.

What do I need to know first?

Six terms carry the whole post. Each one gets a worked example later.

Accelerator: a chip built to do one kind of maths, multiplying large grids of numbers (matrices), much faster than a general-purpose processor. GPUs and TPUs are both accelerators.
GPU (graphics processing unit): originally built for drawing pixels, now the most common AI accelerator. NVIDIA’s data-centre GPUs (H100, B200 and newer) do the matrix maths in units called Tensor Cores.
TPU (Tensor Processing Unit): Google’s own accelerator, designed only for machine learning. Its matrix maths happens in a matrix multiply unit (MXU), which Google describes as a 128 by 128 systolic array: a grid of tiny multiply-and-add cells that pass partial results to their neighbours, so numbers are reused many times without going back to memory.
HBM (high-bandwidth memory): the memory stacked right next to the accelerator. It holds the model’s weights and working data. It is very fast but limited in size (tens to a few hundred gigabytes per chip).
Memory bandwidth: how many bytes per second can move between HBM and the math units, measured in gigabytes or terabytes per second.
FLOP/s (floating-point operations per second): how many multiplications and additions the math units can do per second. A modern accelerator does hundreds of trillions.

The rest of the post is about one question: for a given job, which runs out first, the maths (FLOP/s), the memory bandwidth, the memory capacity, or something outside the accelerator altogether?

What is actually inside an AI server?

An AI server is a normal computer (the host CPU) with several accelerators attached; the CPU feeds and coordinates, and the accelerators do the matrix maths out of their own HBM.

A simple way to picture it is a restaurant. The front of house (the host CPU) takes orders, reads the recipes and hands out tickets. The kitchen stations (the accelerators) do the cooking. Each station has its own fridge right beside it (HBM), and the stations pass dishes to each other through hatches (chip-to-chip links). The front of house never cooks, but if it is slow to hand out tickets, every station stands idle.

inside one AI server: the host CPU reads data, tokenises text and schedules work; it feeds the accelerators over PCIe; each accelerator has math units for matrix maths and HBM right next to them for weights and KV cache; chip-to-chip links (NVLink on NVIDIA, ICI on TPU) join the accelerators

Here is what each part does on a real machine:

Host CPU (an Intel Xeon, AMD EPYC or Arm processor such as Google’s Axion): runs the operating system and the Python program, reads training files or incoming requests, turns text into tokens (tokenisation), shuffles and batches data, launches work on the accelerators and handles networking. Google’s TPU documentation puts it plainly: the XLA compiler turns the maths into TPU machine code, and “the rest of the program runs on the TPU host machine”.
Host RAM (DRAM): large and cheap, but far from the accelerators. Data is staged here before it is copied across.
PCIe: the standard link between the CPU and the accelerators. On a TPU, Google’s architecture page describes the host streaming data into an infeed queue that the TPU loads into HBM.
Math units: Tensor Cores on a GPU, MXUs on a TPU. They do the multiply-and-add work.
HBM: holds the weights, the KV cache (explained below) and intermediate results.
Chip-to-chip links: NVLink on NVIDIA systems, inter-chip interconnect (ICI) on TPUs. They let many chips act as one big machine, which matters when a model or its training state does not fit on one chip.

Where people get tripped up: when engineers say a model is “compute-bound” or “memory-bound”, they almost always mean inside the accelerator, not the host CPU. The CPU never does the model’s matrix maths. It becomes the bottleneck only when it cannot prepare data fast enough. Google’s TPU performance guide calls these input-bound models: “TPUs perform calculations very fast”, so you must keep “a steady stream of data” flowing, for example by reading files in parallel.

A worked example. Suppose a training job needs 4,000 tokenised examples per second to keep its accelerators busy, and the data loader, running on the host CPU with one worker, prepares 1,500 per second. The accelerators are idle more than half the time, and no faster chip will help. Four loader workers (6,000 per second, if the CPU has the cores) fix it. These numbers are illustrative, but the pattern is the first thing a profiler shows you.

What does compute-bound versus memory-bandwidth-bound mean?

A job is compute-bound when the math units are the limit, and memory-bandwidth-bound when the math units sit waiting for numbers to arrive from memory. Which one you hit depends on how much maths you do per byte you fetch.

Picture a kitchen with eight very fast chefs at a counter. Every ingredient lives in a pantry behind a narrow door. The chefs are the math units (FLOP/s). The pantry is HBM, and its size is memory capacity. The width of the door is memory bandwidth.

compute versus memory bandwidth as a kitchen: the pantry is HBM, the narrow door is memory bandwidth, the chefs are the math units; a recipe with one chop per ingredient leaves chefs waiting at the door (memory-bandwidth-bound, like decode); a recipe with hundreds of chops per ingredient keeps them busy (compute-bound, like prefill and training)

Now give the kitchen two recipes:

Recipe A, one chop per ingredient. Each chef grabs an onion, chops once, and is done. Then they wait at the door for the next onion. Hiring more chefs changes nothing; only a wider door helps. This is memory-bandwidth-bound.
Recipe B, 500 chops per ingredient. Each onion keeps a chef busy for a long time, and the door easily keeps up. Now more or faster chefs help, and a wider door does nothing. This is compute-bound.

The real term for “chops per ingredient” is arithmetic intensity: the number of floating-point operations done per byte moved from memory. Every chip has a balance point, its peak FLOP/s divided by its memory bandwidth. Below that intensity you are waiting on memory; above it, on maths. This picture is called a roofline.

The sum with real numbers. Google DeepMind’s How to Scale Your Model works it through for a TPU v5e: about 1.97 × 10¹⁴ FLOP/s divided by 8.2 × 10¹¹ bytes per second of HBM bandwidth gives a balance point of about 240 FLOPs per byte. It quotes a similar figure, about 280, for an NVIDIA H100. A matrix multiplication that reuses each weight for 240 or more tokens at once can keep the chip busy; one that uses each weight for a single token cannot.

The misconception to drop: “a faster chip makes everything faster”. Peak FLOP/s, the number on the marketing slide, only helps compute-bound work. For memory-bound work, the number that matters is HBM bandwidth in terabytes per second.

Why do training, prefill and decode stress different things?

Training and prefill push many tokens through each weight at once, so they are compute-bound; decode produces one token at a time per user and must re-read every weight for each token, so it is memory-bandwidth-bound unless you batch many users together.

Think of a teacher marking exams. Marking the whole class’s papers for question 1, then question 2, means the answer sheet for each question is fetched once and used thirty times. That is prefill and training: high reuse. A teacher who has to fetch every answer sheet again for each single word a pupil writes is decode: almost no reuse.

Training: compute-bound, and hungry for memory capacity

During training the model reads thousands of whole documents at once, runs them forward, measures the error, runs a backward pass to work out gradients, and updates every weight. Every weight is reused across the whole batch, so the maths dominates. Two other limits appear: HBM capacity, because training keeps much more than the weights in memory (next section), and the chip-to-chip network, because chips must share gradients every step.

Prefill: compute-bound

Prefill is the first phase of answering a request: the model reads the entire prompt in one pass. All 2,000 tokens of your question and instructions go through each layer together, so each weight is reused 2,000 times. How to Scale Your Model’s takeaway is direct: “During prefill, all matrix multiplications are basically always compute-bound.”

Decode: memory-bandwidth-bound

Decode is the second phase: the model writes the answer one token at a time, and each new token needs a full pass through the model. For a single user, every weight is fetched from HBM to produce one token, then fetched again for the next. The scaling book’s takeaway for attention during generation: “we are always memory bandwidth-bound”.

prefill versus decode on a toy timeline: prefill processes all prompt tokens in one compute-bound pass and builds the KV cache; decode then makes one token per step, and every step re-reads the weights and KV cache from memory, so it is memory-bandwidth-bound; time to first token covers prefill and the first step, tokens per second is set by the step time

The worked numbers (illustrative, using Google’s published peak figures for a Trillium (v6e) chip: 918 TFLOP/s at 16-bit precision and 1,638 GB/s of HBM bandwidth). Take a toy model with 8 billion parameters stored at 2 bytes each, so 16 GB of weights. Producing one token takes roughly 2 floating-point operations per parameter (one multiply, one add):

Decode, one user, one token
  bytes to read   16 GB   / 1,638 GB/s     = about 9.8 ms
  maths to do     16 GFLOP / 918 TFLOP/s   = about 0.017 ms
  -> the chip waits on memory about 500 times longer than it computes
  -> at best about 100 tokens per second for this one user

Prefill, one 2,000-token prompt
  bytes to read   16 GB (weights read once) = about 9.8 ms
  maths to do     2,000 x 16 GFLOP = 32 TFLOP / 918 TFLOP/s = about 35 ms
  -> the maths takes longer than the reading, so prefill is compute-bound

Real chips never reach peak, and attention adds work these sums ignore, but the gap is so large that the conclusion holds. This is also why the two user-facing speed numbers behave differently. The time to first token is dominated by queueing and prefill, so it grows with prompt length and responds to more compute. Tokens per second while the answer streams is set by the decode step time, so it responds to memory bandwidth.

Why batching wins for decode. If 32 users are decoding at the same moment, the server can read the 16 GB of weights once and produce 32 tokens, one per user. Reading time stays about 9.8 ms; maths grows to about 0.6 ms; throughput rises roughly 32 times. That is why serving engines work so hard to batch requests, and why How to Scale Your Model notes that “batching many concurrent requests together without affecting latency is hard”. It is also where memory capacity bites again, because each extra user needs room for their own KV cache.

Why this fails, and why it wins. An engineer who sees slow streaming and buys chips with more FLOP/s gets little: decode was waiting on memory. An engineer who serves the same model with a 4-bit quantised copy of the weights halves or quarters the bytes read per token, and streaming speeds up because the bottleneck shrank.

One nuance: some modern training, such as reinforcement learning, generates answers inside the training loop, and that generation phase behaves like decode.

What does training keep in memory that inference does not?

Serving a model needs roughly 2 bytes per parameter for 16-bit weights plus a KV cache; training with the common mixed-precision Adam recipe needs about 16 bytes per parameter before activations.

The ZeRO paper (Rajbhandari et al.) gives the breakdown for mixed-precision training with the Adam optimiser and Ψ parameters:

16-bit weights used in the forward and backward pass: 2 bytes per parameter.
16-bit gradients: 2 bytes per parameter.
Optimiser states: a 32-bit master copy of the weights, plus Adam’s momentum (a running average of recent gradients) and variance (a running average of their squares), 4 bytes each, so 12 bytes per parameter.

That is 2 + 2 + 12 = 16 bytes per parameter. The paper’s own example: a 1.5 billion parameter GPT-2 needs 3 GB for its 16-bit weights, yet its training state needs “at least 24 GB”. Activations, the intermediate results saved during the forward pass for use in the backward pass, come on top and grow with batch size and sequence length; for the same GPT-2 at a sequence length of 1K and a batch size of 32, the paper estimates about 60 GB.

what sits in accelerator memory for a toy 8-billion-parameter model: training with mixed-precision Adam keeps 16-bit weights and gradients plus 32-bit weights, momentum and variance, 16 bytes per parameter or 128 GB before activations; serving keeps 2 bytes per parameter, 16 GB, plus a KV cache that grows with users and context

For the toy 8 billion parameter model: training state is 8 × 16 = 128 GB before activations; serving weights are 8 × 2 = 16 GB. A Trillium chip has 32 GiB of HBM, so the serving copy fits on one chip with room for the KV cache, while the training state must be split across at least four chips. ZeRO, and the similar fully sharded data parallel (FSDP) approach, do exactly that: each chip keeps only a slice of the optimiser states, gradients and weights.

Correcting a common shortcut: you will read that optimiser states are “about four times” the inference weights. With mixed-precision Adam they are six times (12 bytes against 2), and the whole training state is eight times.

What is the KV cache, and where does it live?

The KV cache stores the attention keys and values the model has already computed for every earlier token, at every layer, in accelerator memory, so that decode does not recompute them for each new token.

The simple example is a reading group discussing a long book. Without notes, every time someone asks about the next page, the group re-reads the whole book from the start. With notes, they keep a short summary of every page they have read and only read the new page. The notes are the KV cache. They save huge amounts of work, but they take up table space, and the table (HBM) is shared by everyone in the room.

On the real thing: inside each Transformer layer, the attention step turns each token into a query, a key and a value vector. A new token’s query is compared against the keys of all earlier tokens, and the matching values are blended. The keys and values of earlier tokens never change, so the server keeps them. They are tensors of numbers held in HBM, not the attention function itself.

How big is it? For a model with 32 layers, 8 key-value heads of 128 numbers each, stored at 2 bytes (a shape similar to many 8B open models; the figures are illustrative):

per token  = 2 (key and value) x 32 layers x 8 heads x 128 numbers x 2 bytes
           = 131,072 bytes = 128 KiB
8,000-token conversation  = 8,000 x 128 KiB = about 1 GB
32 users at 8,000 tokens  = about 32 GB, twice the 16 GB of weights

So the KV cache, not the weights, is usually what limits how many users one chip can serve, and it is read in full at every decode step, which adds to the memory-bandwidth pressure.

PagedAttention and continuous batching. The vLLM paper (Kwon et al., SOSP 2023) observed that KV cache memory “is huge and grows and shrinks dynamically”, and when managed badly it is “significantly wasted by fragmentation and redundant duplication, limiting the batch size”. Its fix, PagedAttention, stores the cache in fixed-size blocks, the way an operating system pages memory, so no space is reserved for tokens that never arrive. Continuous batching (also called iteration-level scheduling) lets a new request join the batch as soon as another finishes, instead of waiting for the whole batch to end. Both keep more users in the batch, which is exactly what memory-bound decode needs.

What is prompt caching, and is it the same thing?

Prompt (or context) caching keeps the processed form of a repeated prompt opening between separate requests, so the next request can skip that part of prefill; the KV cache lives only for the length of one response.

Back to the reading group: the KV cache is the notes taken during tonight’s meeting. Prompt caching is keeping those notes for chapter one in a drawer, so next week’s group can start at chapter two.

In a self-hosted engine you can see the mechanism. vLLM’s documentation for automatic prefix caching says it caches “the kv-cache blocks of processed requests” and reuses them “when a new request comes in with the same prefix”. Hosted model APIs describe caching repeated prompt prefixes and bill cached tokens at a discount, but they do not publish exactly what they store internally, so I avoid claiming it is “the same KV matrices kept in RAM”.

The lifetime differs by provider and setting, which is why “five minutes” is not a universal rule:

Anthropic’s prompt caching documentation gives a cache lifetime of 5 minutes by default, with a 1-hour option at extra cost.
Google Cloud (Gemini on Agent Platform), per the context caching overview as of November 2026: implicit caching is on by default in every project, with a 90 per cent discount on cached tokens; it keeps prefixes in memory and in a durable tier, and an entry expires no later than 24 hours after creation (most are evicted earlier). Explicit caching lets you create a cache yourself with a default time to live (TTL) of 60 minutes that you can extend, a 90 per cent discount on Gemini 2.5 and later models, and a storage charge. The minimum cacheable prompt is 4,096 tokens for most Gemini 3 and Gemini 4 models.

A worked example. A support assistant sends the same 6,000-token block of instructions and policy text with every question, followed by a 200-token question. With the instructions first and the question last, requests after the first can reuse the 6,000-token prefix: prefill shrinks to about 200 new tokens and the cached part is billed at the discounted rate. Put a timestamp or the user’s name at the top and the prefix differs from the first token, so nothing is reused. The Forward deployed engineering performance post shows this ordering rule with a diagram.

How does confidential computing protect a model while it runs?

Encryption at rest and in transit protect data on disks and wires; confidential computing protects data in use, while it sits in memory being processed, by running the work inside a hardware-isolated trusted execution environment whose setup you can verify by attestation.

A simple example: a bank vault protects cash at rest, and an armoured van protects it in transit. But the cash still has to be counted somewhere. Confidential computing is a sealed counting room with one-way glass: the counters work inside, nobody outside (including the building’s owner) can reach in, and before you send the cash you check a tamper-proof certificate that says exactly who and what is inside the room.

confidential computing: data is protected at rest by storage encryption and in transit by TLS, but in use it sits in RAM or HBM while the chip computes; a trusted execution environment spanning a confidential VM (AMD SEV-SNP or Intel TDX) and an accelerator in confidential mode keeps it isolated from the hypervisor and operators, and attestation proves the setup before any keys are released

The real terms, step by step:

Encryption at rest: data on disk or object storage is encrypted. On Google Cloud this is on by default, and CMEK (customer-managed encryption keys) lets you hold the keys.
Encryption in transit: data moving over a network is wrapped in TLS.
Encryption in use: the gap. To compute, a chip needs the numbers in its memory. A trusted execution environment (TEE) is hardware that isolates that memory from everything outside it, including the hypervisor and the people who run the machine.
On the CPU side, AMD SEV-SNP and Intel TDX create confidential virtual machines whose memory is encrypted with keys held in the processor. Google’s Confidential VM overview describes SEV-SNP as adding hardware-based security to “help prevent malicious hypervisor-based attacks like data replay and memory remapping”, and TDX as creating “an isolated trust domain (TD) within a VM”.
On the accelerator side, NVIDIA introduced GPU confidential computing with the H100 (Hopper), which NVIDIA calls “the first ever GPU to introduce support for confidential computing”. The GPU joins the CPU’s TEE: data crossing PCIe goes through encrypted bounce buffers, and command buffers and kernels are encrypted and signed before they cross the bus.
Attestation: the hardware produces a signed report of what firmware and software are running, and a verifier checks it before any secrets (model keys, customer data) are released into the TEE.

A precise correction about HBM. A common claim is that the H100 encrypts everything it writes to HBM. NVIDIA’s own technical blog says the opposite for Hopper: “The on-package HBM memory is considered secure against everyday physical attack tools and is not encrypted.” That is why NVIDIA reports GPU raw compute and HBM bandwidth “at par with non-confidential mode”, while CPU to GPU transfers are limited by encryption to “roughly 4 GBytes/sec” in its measurements.

What changed with Blackwell and later. NVIDIA’s Blackwell page says Blackwell “is the first TEE-I/O capable GPU in the industry”, with “inline protection over NVIDIA NVLink”, and “nearly identical throughput performance compared to unencrypted modes”. NVIDIA’s confidential computing page lists support across Hopper, Blackwell and Rubin, and describes Vera Rubin NVL72 as a rack-scale security domain across NVLink. I do not quote a percentage overhead because NVIDIA does not publish one on these pages; measure your own workload.

Why it wins, and why it can fail. It wins for a model provider who wants to run proprietary weights in someone else’s data centre, or a hospital that must process records it cannot expose to the cloud operator. It fails if you skip attestation: an encrypted VM that nobody verifies proves nothing. It also does not stop a bad prompt or a leaky application inside the TEE.

How does the software stack fit together?

At the bottom, a compiler or kernel layer turns maths into chip instructions (CUDA on NVIDIA GPUs, XLA and Pallas on TPUs); in the middle, a framework such as PyTorch or JAX is where the model is written; at the top, a serving engine such as vLLM batches users and manages the KV cache, and the same engine can have different hardware backends.

A simple example is a printed book. The author writes in English (the framework). A typesetter turns the text into plates for one specific press (the compiler and kernels). The publisher decides how many copies to print and in what order to ship orders (the serving engine). Change the press and you need a different typesetter, but the author’s words and the publisher’s process can stay the same.

the AI software stack on NVIDIA GPUs and Google TPUs in November 2026, top to bottom: serving engine (vLLM, SGLang, TensorRT-LLM on GPU; vLLM TorchTPU, vLLM JAX and SGLang on TPU), framework (PyTorch and JAX on both), compiler and kernels (CUDA on GPU; XLA and Pallas on TPU), and the hardware itself

Layer by layer, as of November 2026:

Compiler and kernels. A kernel is a small program that runs on the accelerator, such as one matrix multiplication or one attention step. CUDA is NVIDIA’s “parallel computing platform and programming model” for its GPUs, with libraries of hand-tuned kernels; it does not run on TPUs. XLA (Accelerated Linear Algebra), developed in the open as OpenXLA, is a compiler that takes models from frameworks such as PyTorch, TensorFlow and JAX and optimises them for GPUs, CPUs and other hardware; Google’s TPU docs say code that runs on TPUs “must be compiled by” XLA. Pallas is “an extension to JAX for writing custom kernels for GPU and TPU”, for cases where XLA’s automatic optimisations fall short (attention kernels in serving engines are a typical use).
Frameworks. PyTorch is widely used on GPUs. JAX is Google’s numerical library, native on TPUs and also used on GPUs through XLA. For PyTorch on TPUs, the picture changed in October 2026: Google open-sourced TorchTPU, a PyTorch-native TPU backend co-developed with Meta’s PyTorch team, and the PyTorch on Cloud TPU page now says “TorchTPU replaces PyTorch/XLA”. It uses eager execution by default, and torch.compile compiles the graph to TPU machine code through OpenXLA and StableHLO, so XLA is still underneath.
Serving engines. vLLM started on NVIDIA GPUs and now lists NVIDIA and AMD GPUs and CPUs as built-in targets, with “hardware plugins such as Google TPUs”. On TPUs there are two vLLM plugins: vLLM TorchTPU (package vllm-torchtpu), the PyTorch-native backend whose README states the goal of keeping the same vLLM command line, OpenAI-compatible server and model definitions with no model code changes, and vLLM JAX (formerly tpu-inference), focused on serving JAX models. SGLang also lists Google TPU support. TensorRT-LLM is NVIDIA’s own serving library for its GPUs. JetStream, Google’s earlier TPU serving engine, was archived on 1 February 2026; its README points users to vLLM on TPU.

The misconception to drop: “vLLM is a TPU thing” or “CUDA and vLLM are alternatives”. They sit at different layers. vLLM is a serving engine that calls CUDA kernels on NVIDIA GPUs and TorchTPU or JAX with Pallas kernels on TPUs. And “the exact same command on both” is a goal, not a guarantee: check the plugin’s supported-models list before you assume a model runs unchanged.

How does this map to Google Cloud?

Concepts first, then products, all as of November 2026:

TPUs. TPU7x (Ironwood) is described in the docs as “the latest TPU available on Google Cloud”, with 192 GiB of HBM and 7,380 GBps of HBM bandwidth per chip, against 32 GiB and 1,638 GBps for Trillium (v6e). In April 2026 Google announced the eighth generation as two chips, TPU 8t for training and TPU 8i for low-latency inference, which it said would be generally available later in 2026; TPU 8i pairs 288 GB of HBM with 384 MB of on-chip SRAM sized “for the KV cache footprint of reasoning models”, and both run on Google’s Arm-based Axion host CPUs. On 9 November 2026 Google added TPU 8tx, “optimized for secure and confidential inference”, and TPU 8i Max, two TPU 8i chips in one package with double the compute and HBM bandwidth for very long contexts. The announcement does not give availability dates for these two.
Confidential AI on TPUs. Google says TPU 8tx is built on the same silicon as TPU 8t, “paired with 5th Generation AMD EPYC CPU hosts”, and keeps “sensitive data and model weights protected and isolated even while in active use”, with “a near-zero performance penalty”.
Confidential AI on GPUs. Confidential VM supports the A3 High machine type (a3-highgpu-1g) with an NVIDIA H100 on Intel TDX, and Confidential GKE Nodes can run H100 GPU workloads (guide). The Confidential VM overview notes that NVIDIA HGX B200 and GB300 NVL72 access is available to specific eligible customers.
Software on TPUs. JAX, TorchTPU, vLLM TorchTPU, vLLM JAX and SGLang, on TPU VMs or GKE. Google’s TPU Sync library (open source, November 2026) moves KV cache blocks between HBM, host RAM and storage, which is the same capacity problem described above, solved by tiering.
Managed models. If you call Gemini on Agent Platform, Google runs the hardware. Your levers are prompt length, caching (implicit and explicit), batch versus online, and Provisioned Throughput, not chip choice.

My view: most teams building agents never choose a chip. They choose a managed model and pay per token. The hardware still shapes their bill and latency through prompt length, caching and output length, so the mental model in this post is worth having even if you never rent a TPU.

Which bottleneck am I hitting?

A litmus test, cheapest check first:

Is the accelerator idle? Check utilisation in your profiler or monitoring. If it is low, look at the host first: data loading, tokenisation, Python overhead. More CPU workers or prefetching is the cheapest fix.
Is the slow part the first token or the stream? Long time to first token with a long prompt means prefill: shorten the prompt or cache its stable prefix. Slow streaming means decode: go to step 3.
Is decode slow for one user, or is throughput low for many? For one user, it is memory bandwidth: a smaller or quantised model, or a chip with more HBM bandwidth. For many users, check the batch size the engine actually reaches: PagedAttention-style engines and continuous batching raise it.
Are you running out of memory? Out-of-memory errors or a capped batch size usually mean KV cache (serving) or optimiser states and activations (training): shard across more chips, use ZeRO or FSDP, or offload the KV cache to host memory.
Are chips waiting on each other? In multi-chip jobs, time spent in communication means the chip-to-chip network is the limit: change the sharding strategy.
Only then buy more or faster chips, and buy for the limit you measured: FLOP/s for compute-bound work, HBM bandwidth and capacity for decode.

How do I lock it in?

Term                   What it is                                   Everyday picture
Host CPU               feeds data, tokenises, schedules             front of house
Accelerator            GPU or TPU doing matrix maths                the chefs
HBM                    fast memory beside the chip                  the pantry
Memory bandwidth       bytes per second from HBM to the maths       the pantry door
Compute-bound          maths is the limit (training, prefill)       chefs flat out
Memory-bound           waiting for bytes (decode)                   chefs at the door
Arithmetic intensity   FLOPs per byte fetched                       chops per ingredient
KV cache               stored keys and values per token per layer   tonight's notes
Prompt caching         processed prefix kept between requests       notes kept in a drawer
TEE plus attestation   isolated memory you can verify               sealed counting room
CUDA / XLA, Pallas     turn maths into chip instructions            the typesetter
vLLM, SGLang           batch users, manage KV cache                 the publisher

One honest gap: most of the performance figures in this post are peak numbers or toy models, and real utilisation is lower and workload-specific. Vendors also do not publish like-for-like confidential computing overheads. The only reliable answer for your model is a benchmark on your own prompts.

What can go wrong?

Buying FLOP/s for a decode problem. Streaming speed is set by memory bandwidth, so a chip with more peak maths may not help.
Starving the accelerators. A single-threaded data loader or tokeniser on the host leaves expensive chips idle. Profile before you scale.
Forgetting the KV cache in capacity plans. The weights fit, then 30 long conversations do not. Plan memory for weights plus KV cache at your real context length and concurrency.
Breaking the cache with prompt order. A timestamp at the top of the prompt defeats prefix caching on every request.
Sharing a prefix cache across tenants. vLLM’s docs warn that shared prefix caches can enable timing attacks that infer cached content, and offer a per-request cache_salt to isolate them.
Assuming confidential means verified. Without checking attestation, a confidential VM is just a VM. And HBM on Hopper is not encrypted, so know which threats your setup covers.
Assuming one command runs everywhere. vLLM’s TPU backends are separate plugins with their own supported-model lists and releases.

What did we learn?

An AI server is a host CPU that feeds and coordinates plus accelerators that do the matrix maths out of their own HBM; the CPU becomes the bottleneck only when it cannot prepare data fast enough.
Whether a job is compute-bound or memory-bandwidth-bound depends on arithmetic intensity, the maths done per byte fetched, compared with the chip’s balance point (about 240 FLOPs per byte on a TPU v5e).
Training and prefill reuse each weight across many tokens, so they are compute-bound; decode reads every weight for each token, so it is memory-bandwidth-bound unless many users are batched together.
Mixed-precision Adam training needs about 16 bytes per parameter before activations, against about 2 bytes per parameter for 16-bit serving, which is why training needs many more chips.
The KV cache keeps each earlier token’s keys and values in HBM during a response, and prompt caching keeps a processed prefix between requests; the KV cache usually limits how many users one chip serves.
Confidential computing adds protection in use through a hardware TEE and attestation; the accelerator stack runs from CUDA or XLA and Pallas, through PyTorch or JAX, up to serving engines such as vLLM that run on both GPUs and TPUs.

What comes next?

If tokens, context windows and how a model writes one token at a time are still new, start with Prerequisite 3: How a model works, which this post builds on. To turn these hardware limits into model choices for an agent (size, thinking, open-weight versus managed), read Part 13: Choosing and running models for agents, and for measuring time to first token, tokens per second and caching in practice, see Forward deployed engineering, Part 9.

Sources

Google Cloud: Introduction to Cloud TPU (MXU as a 128 by 128 systolic array, XLA compilation, the rest of the program on the host)
Google Cloud: TPU architecture (infeed queue, HBM, MXU, ICI, slices)
Google Cloud: TPU7x (Ironwood) (latest available TPU; HBM capacity and bandwidth, peak FLOP/s for v5p, v6e and TPU7x)
Google Cloud: Cloud TPU performance guide (input-bound models)
Google: Our eighth generation TPUs: two chips for the agentic era (TPU 8t and 8i, 288 GB HBM, 384 MB SRAM, Axion hosts)
Google Cloud Blog: Introducing TPU 8tx and TPU 8i Max (9 November 2026; confidential inference, AMD EPYC hosts, near-zero penalty)
Google Cloud Blog: Introducing TPU Sync (KV cache tiering across HBM, host RAM and storage)
Google Cloud Blog: TorchTPU is now open source (20 October 2026)
Google Cloud: PyTorch on Cloud TPU (TorchTPU replaces PyTorch/XLA; torch.compile through OpenXLA and StableHLO)
How to Scale Your Model: All about Transformer inference (TPU v5e balance point of 240, H100 about 280; prefill compute-bound, generation memory bandwidth-bound)
Rajbhandari et al.: ZeRO: Memory Optimizations Toward Training Trillion Parameter Models (16 bytes per parameter for mixed-precision Adam)
Kwon et al.: Efficient Memory Management for Large Language Model Serving with PagedAttention
Anyscale: How continuous batching enables 23x throughput in LLM inference (iteration-level scheduling)
vLLM documentation (supported hardware, PagedAttention, continuous batching, prefix caching)
vLLM: Automatic prefix caching (KV cache block reuse, cache_salt)
vLLM TorchTPU on GitHub
vLLM JAX (tpu-inference) on GitHub
JetStream on GitHub (archival notice, 1 February 2026)
SGLang documentation (hardware support including Google TPU)
TensorRT LLM documentation
NVIDIA: CUDA Programming Guide
OpenXLA: XLA
JAX: Pallas, a JAX kernel language
Google Cloud: Context caching overview (implicit and explicit caching, 60-minute default TTL, 90 per cent discount, minimum tokens)
Anthropic: Prompt caching (5-minute default lifetime)
Google Cloud: Confidential VM overview (AMD SEV-SNP, Intel TDX, NVIDIA Confidential Computing)
Google Cloud: Confidential VM supported configurations (A3 High with H100 on Intel TDX)
Google Cloud: Encrypt GPU workload data in use with Confidential GKE Nodes
NVIDIA Technical Blog: Confidential Computing on NVIDIA H100 GPUs (first confidential computing GPU; HBM not encrypted; bounce buffers)
NVIDIA: Blackwell architecture (first TEE-I/O capable GPU, inline NVLink protection)
NVIDIA: Confidential computing (Hopper, Blackwell and Rubin support)

Opinions are my own. Every fact in this post comes from the public documentation linked above.