AI chips: GPUs, TPUs, NPUs, and the real bottleneck of memory bandwidth

Why the chip story became a memory story

When people hear 'AI chips', they usually think of how fast a chip can compute. That is why FLOPS (floating-point operations per second) sits at the top of chip spec sheets. But when you actually serve a large language model (LLM), the compute units sit idle for longer than you might expect, because the data they need to compute on has not yet arrived from memory.

The reason lies in how LLMs generate text. An LLM is an autoregressive model: it produces one token at a time, then appends that token to the input to produce the next one. Each new token requires reading most of the model's weights from memory again. Yet the computation done with those weights, once read, amounts to multiplying a matrix by a single vector. The in-DRAM parallel computing study paper points out that on edge devices, most autoregressive decoding time is spent on memory-bound matrix-vector multiplication (GeMV).

The NPU speculative decoding study paper also calls the problem faced when running LLM inference on NPU hardware the 'Memory Wall'. No matter how fast the compute units get, if memory bandwidth (the amount of data that can be moved per second) cannot keep up, overall speed is set by memory. That is why competition among AI chips today is often fought over memory and data movement rather than compute units.

Key idea: a kitchen where moving ingredients is slower than the cook

Picture a restaurant kitchen. The cook (the compute unit) works very fast, but ingredients come from a distant storeroom (memory), a fixed amount per trip. If the cook makes one customer's dish one spoonful at a time, and every spoonful requires fetching all the ingredients from the storeroom again, then no matter how many cooks you add, food cannot come out faster than the ingredients are delivered.

Seen through this analogy, chip designers can do three main things. First, widen the corridor. HBM (high-bandwidth memory), which stacks memory right next to the chip to increase bandwidth, takes this approach. Second, make several dishes from one delivery of ingredients. Batching, which processes multiple requests together, and speculative decoding, which verifies several tokens at once, fall into this group. Third, cook inside the storeroom. That is the idea behind PIM (Processing-In-Memory), which performs computation inside memory.

The same analogy explains the differences among GPUs, TPUs and NPUs. All three are kitchens with many 'fast cooks', but they differ in which dishes they were built to make and how well-equipped they are with kitchen tools (software).

How do GPUs, TPUs and NPUs differ?

The GPU (graphics processing unit) was originally built to compute the many pixels on a screen at the same time. Its design for massively parallel execution of the same computation fits the matrix multiplications of neural networks well, so it came to be widely used for AI training and inference. The GPU's biggest strength is less the hardware itself than the programming ecosystem called CUDA. Most developers, libraries and inference engines are built on it first.

The TPU (tensor processing unit) is a dedicated chip Google designed for the matrix operations of neural networks. It gives up general-purpose flexibility to focus on streaming matrix multiplications efficiently. NPU (neural processing unit) is a broad term for chips specialized for neural network computation, ranging from small chips inside smartphones and edge devices to inference chips for data centers. Statically scheduled chips such as Groq's TSP, which drop caches and hardware schedulers entirely and have the compiler decide the timing of every operation in advance, are one extreme of this family.

The more specialized a chip is, the better its power and cost efficiency can be for particular operations. But if the software does not keep up, those advantages go unused. The problem that the NPU speculative decoding study paper started from was that domestic (Chinese) infrastructure did not natively support mainstream speculative decoding algorithms. Optimizations that work out of the box on GPUs become work that has to be rebuilt by hand on a new chip. This is the core trade-off between general-purpose and specialized chips.

How it works: how a single inference uses memory

LLM inference has two stages. The first, prefill, processes the user's entire prompt at once. Because many tokens are computed together, weights read once can be reused many times, and the compute units are easy to keep fully busy. During this stage, intermediate results for each token are stored in the KV cache so the same computation does not have to be repeated in the next stage.

The second stage, decoding, produces tokens one by one. To produce each token, the model weights and the KV cache accumulated so far are read from memory, the computation is done, and the resulting token is fed back in as input. Because little computation is done per piece of data read, this stage is bound by memory bandwidth. The longer the conversation, the larger the KV cache grows. That is why LoopServe paper notes that as conversation history grows, both the compute and memory burden on the LLM increase.

In short, 'memory capacity' determines whether the model weights and KV cache can fit on the chip, and 'memory bandwidth' determines decoding speed. In terms of what users notice, prefill mainly affects the time to first token (TTFT), and decoding mainly affects how fast text streams out after that.

How the field developed: three attempts to get past the memory wall

The first is algorithms that read memory less. In speculative decoding, a small model guesses the next few tokens ahead of time, and the large model verifies those candidates in one pass. Since the large model's weights are read once to confirm several tokens, round trips to memory are reduced. The NPU speculative decoding study paper implemented this technique to run on NPUs as an end-to-end speculative inference acceleration system for the OpenPangu-7B model. In the same line of work, the hierarchical verification tree study paper reported reducing inference time and energy by verifying the most likely candidates first and pruning unpromising candidates early.

The second is shrinking the KV cache intelligently. LoopServe paper computes only the important parts of the attention matrix for each input during prefill, and during decoding progressively compresses the KV cache based on the most recently generated tokens. It grew out of the observation that cutting by fixed rules based only on position cannot follow the actual flow of a conversation. The study also released a benchmark of 11 multi-turn datasets that reflect real question positions and conversational dependencies.

The third is changing the hardware architecture itself: PIM, which computes inside memory. PRIMAL paper combined LoRA (low-rank adaptation) with a PIM-based LLM inference accelerator. It connects several types of PIM processing units in a 2D mesh network, and overlaps LoRA swapping with computation by reprogramming SRAM and powering off idle parts. In a setting with Llama-13B and LoRA rank 8, it reported 1.5x the throughput and 25x the energy efficiency of an NVIDIA H100. It is an example of how cooking inside the storeroom, instead of moving ingredients, can greatly improve power efficiency.

Where things are heading now

Recent research is digging deeper into in-memory computation. The in-DRAM parallel computing study paper handles decoding's GeMV with PUD (Processing-using-DRAM), which computes directly inside DRAM. DRAM has enough parallelism to compute across 65,536 columns at once. But the weights of a single LLM layer are too small to fill that parallelism; for example, the QKV matrix of Llama3 8B has only 6,144 rows. The researchers designed both an algorithm that packs the next layer's computation together with the current layer's and a data layout tailored to the DRAM structure. As a result, they reported decoding up to 1.41x faster than a PUD baseline and 2.45x faster than processor-based execution.

The second trend is measuring the efficiency of 'a group of chips' rather than 'a single chip'. The cluster efficiency study paper argues that cheap accelerators have little memory per chip, so many must be combined, and chip-to-chip communication then decides success or failure. On Tenstorrent Galaxy, which links Blackhole chips, optimizations such as overlapping communication with computation yielded 1.07–2.24x the tokens per dollar of a server with eight RTX PRO 6000 Blackwell GPUs, based on three years of purchase and electricity costs. Conversely, when communication on the same chips was slowed to emulate TCP-level speeds, throughput dropped 23.5–67.4%, and the cost advantage disappeared in every configuration evaluated. This means the interconnect linking chips matters as much as memory bandwidth.

The third trend is trying to use AI to make up for new chips' weak point, software. According to the multi-agent kernel optimization study paper, current LLMs write single-operation CUDA kernels correctly most of the time, but their accuracy collapses when faced with chips and languages they have rarely seen. The researchers built an agent loop targeting the FuriosaAI RNGD NPU, which is programmed in a poorly documented Rust-based language. The agents wrote and verified their own documentation summaries, measured chip performance, and recorded compiler rules. Over 12 days, they optimized three Gemma-4-12B decoding kernels, achieving a geometric-mean speedup of 9.293x over the official baseline. Each kernel came within 1.25x of the cycles needed just to move the weights. In the end, this confirms again that the ceiling on optimization is also the amount of memory movement.

Limitations, open questions, and practical decision criteria

The most common mistake is comparing chips on FLOPS numbers alone. Services that generate long answers, such as chatbots and coding agents, spend a large share of time in decoding, so memory bandwidth and capacity determine perceived speed. If you handle long contexts, first check whether there is enough memory capacity to hold the KV cache. If you plan to combine multiple chips, the cluster efficiency study paper shows that a slow interconnect can wipe out the cost advantage entirely.

The second mistake is generalizing numbers from papers. PRIMAL's paper 25x energy efficiency was measured in a specific setting of Llama-13B with LoRA rank 8. The Galaxy results also assume a specific model and a three-year cost model. Keep in mind, too, that many PIM studies are at the design and simulation stage, not on chips in mass production. Always check what the baseline is (an H100, or a baseline from the same family). The 1.41x and 2.45x in the in-DRAM computing study paper differ because their baselines differ.

The third is underestimating software maturity. A dedicated NPU may have good hardware efficiency yet not natively support optimizations such as speculative decoding paper. Writing kernels also takes considerable effort. Even in the multi-agent kernel study paper, it took 12 days because of limited hardware time and noisy measurements. When deciding on adoption, it is safer to first check whether the models and inference engines you currently use run on the chip out of the box.

As practical metrics, it is a good habit to look not only at 'tokens per second' but also at 'tokens per dollar' and 'tokens per unit of energy (joule)'. Even if a single chip is fast, total cost changes if you need to combine many of them or they draw a lot of power.

Further study

To get a feel for the memory wall problem itself, start with the problem description in the in-DRAM parallel computing study paper. It uses concrete numbers to show why decoding is memory-bound and why hardware parallelism and model size do not match.

On the algorithm side, reading the NPU speculative decoding study paper and LoopServe paper in turn lets you compare the two branches of the 'read less' strategy (verifying tokens in batches, and shrinking the KV cache). On the hardware side, PRIMAL paper is a good entry point into PIM accelerator design.

For readers considering actual adoption, the cluster efficiency study paper is recommended. It shows how chip prices, electricity costs and interconnects combine into a single cost metric. If you are interested in the software problems of new NPUs, the multi-agent kernel optimization study paper candidly documents the difficulties on the ground.