Skip to content
Residual Stream
Go back

Beyond the GPU: TPUs, LPUs and What Specialization Buys

Updated:
Tan Le
ENVI

How a GPU Runs Deep Learning showed what a GPU spends its silicon on: thousands of small cores, registers and caches to keep them fed, warp schedulers that swap threads to hide memory latency, and tensor cores bolted on for matrix multiplication. Most of that machinery exists because a GPU must run any parallel program well. A chip built for one workload can drop it. This post looks at two that did, from opposite directions: Google’s TPU, which builds the matrix multiplication itself into the hardware, and Groq’s LPU, which keeps the whole model in on-chip memory and plans every cycle in advance.

GPU thousands of cores caches, shared memory warp schedulers tensor cores HBM runs any parallel program TPU systolic matrix units vector unit SparseCores HBM torus interconnect dense linear algebra, planned by XLA LPU (Groq) 230 MB SRAM, 80 TB/s functional slices no caches, no arbitration no HBM every cycle compile-time deterministic, one stream fast

Each design removes something the GPU keeps for generality.

At a glance
  • Specialized chips win by removing general-purpose machinery: caches, schedulers, speculation. The first TPU had none of them and ran Google’s inference 15 to 30 times faster than the CPUs and GPUs of its time.
  • A systolic array wires data reuse into hardware: an n × n array uses each value it reads n times. At batch-1 decoding it still waits on memory.
  • TPUs pair systolic matrix units with HBM, torus pods of up to 9,216 chips and a compiler that plans everything.
  • Groq’s LPU keeps the model in on-chip SRAM and schedules every cycle at compile time: 609 chips to hold a 16-bit 70B model, and the fastest single stream.

Table of contents

Open Table of contents

The systolic array

Matrix multiplication does n3n^3 multiply-adds on 2n22n^2 input numbers, so each number is used nn times. A GPU gets that reuse by staging tiles of the inputs in shared memory and registers (FlashAttention did the same by hand). A systolic array gets it from wiring. It is a grid of multiply-add cells, each connected only to its neighbors. Values enter at the edges and pulse through the grid one cell per clock cycle, like blood through the heart (hence “systolic”), and every cell does one multiply-add with whatever arrives and passes its inputs on. No cell fetches anything from memory.

A 3×3 version, computing C=ABC = AB, simulated cycle by cycle. In this layout each cell owns one output: rows of AA flow left to right, columns of BB top to bottom, each skewed by one cycle per row so that AikA_{ik} and BkjB_{kj} meet in cell (i,j)(i, j) at cycle i+j+ki + j + k:

C11 C12 C13 C21 C22 C23 C31 C32 C33 a11 a12 a13 a21 a22 a23 a31 a32 a33 b11 b21 b31 b12 b22 b32 b13 b23 b33 A’s rows flow right, B’s columns flow down, each skewed one cycle per row; aik​ meets bkj​ in cell (i, j) at cycle i + j + k no cell reads memory output-stationary 3×3 array; the 6 highlighted cells work at cycle 2 (skew shown schematically)

Data reuse built into the wiring. Schematic; at cycle 2, the six highlighted cells are working.

# A 3x3 output-stationary systolic array computing C = A @ B, cycle by cycle.
# Cell (i, j) owns C[i, j]. Row i of A enters from the left edge and moves
# one cell right per cycle; column j of B enters from the top and moves one
# cell down per cycle. Each input is skewed by one cycle per row (column),
# so A[i, k] and B[k, j] meet in cell (i, j) at cycle i + j + k. Every cell
# multiplies what arrives, adds to its own total, and passes both inputs on:
# no cell ever reads memory.
n = 3
A = np.arange(1, 10).reshape(n, n)
B = np.arange(1, 10).reshape(n, n)[::-1]
C = np.zeros((n, n), dtype=int)
a_reg = np.zeros((n, n), dtype=int)       # the A value each cell holds this cycle
b_reg = np.zeros((n, n), dtype=int)
macs = 0
for cycle in range(3 * n - 2):
    # shift: A moves right, B moves down; the edges take the next skewed input
    a_reg[:, 1:] = a_reg[:, :-1].copy(); b_reg[1:, :] = b_reg[:-1, :].copy()
    for i in range(n):
        k = cycle - i
        a_reg[i, 0] = A[i, k] if 0 <= k < n else 0
    for j in range(n):
        k = cycle - j
        b_reg[0, j] = B[k, j] if 0 <= k < n else 0
    C += a_reg * b_reg
    active = int(((a_reg != 0) & (b_reg != 0)).sum())
    macs += active
    print(f"cycle {cycle}: {active} of {n * n} cells multiplying")
print(f"\nresult equals A @ B: {np.array_equal(C, A @ B)}   ({macs} multiply-adds in {3 * n - 2} cycles)")
print(C)

# Memory traffic: each of A's and B's n^2 values is read from memory once
# and then reused n times as it passes through the array.
for n in (3, 128, 256):
    print(f"n = {n:3d}: {n ** 3:>10,d} multiply-adds from {2 * n * n:>7,d} values read: "
          f"each value used {2 * n ** 3 // (2 * n * n):3d} times")

Full script: systolic.py

cycle 0: 1 of 9 cells multiplying
cycle 1: 3 of 9 cells multiplying
cycle 2: 6 of 9 cells multiplying
cycle 3: 7 of 9 cells multiplying
cycle 4: 6 of 9 cells multiplying
cycle 5: 3 of 9 cells multiplying
cycle 6: 1 of 9 cells multiplying

result equals A @ B: True   (27 multiply-adds in 7 cycles)
[[ 18  24  30]
 [ 54  69  84]
 [ 90 114 138]]
n =   3:         27 multiply-adds from      18 values read: each value used   3 times
n = 128:  2,097,152 multiply-adds from  32,768 values read: each value used 128 times
n = 256: 16,777,216 multiply-adds from 131,072 values read: each value used 256 times

Twenty-seven multiply-adds in seven cycles, on nine cells, from 18 values read once each. The wavefront of active cells rises and falls (1, 3, 6, 7, 6, 3, 1) as the skewed inputs fill the array and drain from it. Each edge cell takes in one new value per cycle, from on-chip buffers, and the n2n^2 cells inside do all the multiply-adds; an n×nn \times n array uses each value it takes in nn times.

1/9 cycle 0 3/9 cycle 1 6/9 cycle 2 7/9 cycle 3 6/9 cycle 4 3/9 cycle 5 1/9 cycle 6 cells multiplying per cycle: 27 multiply-adds in 7 cycles, from 18 values read an n × n array uses each value it takes in n times: 128 times at n = 128

The wavefront fills the array and drains from it.

Data: the output above

That reuse is inside the array, between it and its on-chip buffers. It doesn’t by itself make a workload efficient against off-chip memory: arithmetic intensity against HBM (How a GPU Runs Deep Learning’s roofline) depends on the problem’s dimensions, not the array’s. Multiplying one 128×128 BF16 tile by another is at most 64 FLOPs per byte (2 × 128³ FLOPs over two 32 KiB inputs), well left of a TPU v6e’s ridge point of about 560. The array pays off when a weight tile, once loaded, meets many rows of activations: a large batch, or a long sequence in prefill. At batch-1 decoding, each weight still meets one row, and the array waits on memory like any other chip. (The original TPU paper’s own roofline showed most of Google’s inference workloads were memory-bound.)

The simulation is output-stationary: each cell keeps one output and the two inputs flow past. TPUs use the other common layout, weight-stationary: a tile of weights is loaded into the cells once and stays, activations flow in from one side, and partial sums flow out the other. That suits neural networks, where the same weights meet a long stream of inputs, and it doesn’t have to drain outputs out of the cells between tiles. Either way the price is inflexibility. The array does matrix multiplication and nothing else; everything else (activations, normalization, softmax, top-kk) runs on a much smaller vector unit beside it, which sits on the critical path and has grown in every recent generation.

The TPU

Google’s first TPU (Jouppi et al., 2017) was built in about 15 months for inference in its data centers: one 256×256 systolic array of 8-bit multiply-adders (65,536 of them), 92 trillion operations per second, 28 MiB of on-chip buffers, and weights streamed from ordinary DDR3 memory. It had no caches, no branch prediction, no out-of-order execution and no multithreading: a host CPU sent it instructions over PCIe, and the chip just ran the matrix math. On Google’s inference workloads it was 15 to 30 times faster than the CPUs and GPUs of the time, at far better performance per watt.

Later generations kept the core idea and added what training and large models need:

  • Matrix units (MXUs), systolic arrays of 128×128 up to v5, and 256×256 in v6e (Trillium) and TPU7x (Ironwood), in BF16 and lower precisions. Since v4, chips also carry SparseCores, small processors for the scattered memory lookups of embedding tables (and, increasingly, routing in mixture-of-experts models).
  • HBM, the same stacked memory GPUs use. A v6e chip has 32 GB at 1.6 TB/s; an Ironwood chip has 192 GiB at 7.4 TB/s.
  • The inter-chip interconnect (ICI): chips wired directly to their neighbors in a 2-D torus (v5e and v6e, 256 chips per pod) or a 3-D torus (v4 with 4,096 chips per pod, v5p, and Ironwood with up to 9,216), so a pod behaves like one very large machine for the collective operations (all-reduce, all-gather) that splitting a model across chips requires. Since v4, the 3-D pods use optical circuit switches to connect cubes of chips into slices of any shape and to route around failed ones.
  • A compiler that plans everything. TPU programs go through XLA, which sees the whole computation graph with fixed shapes and lays out data movement ahead of time. That’s why TPUs reward static shapes and punish code that changes shape every step.

The TPU’s bet is that at Google’s scale, a chip that only does dense linear algebra well, wired into very large pods, beats a more general one on cost. It trains Gemini; outside Google it is rented through Google Cloud.

The LPU: everything in SRAM, nothing left to chance

Groq’s chip (Abts et al., 2020, called the Tensor Streaming Processor; Groq now calls it the LPU, for language processing unit) removes a different thing: memory hierarchy and runtime uncertainty. It has no HBM at all. Each chip holds 230 MB of SRAM, read at 80 TB/s, and computes 188 TFLOPS at FP16 (750 TOPS at INT8). Fast on-chip memory isn’t new (an H100’s caches and shared memory are also far faster than its HBM); what’s new is making it the only memory, big enough to hold the weights, so the slow tier disappears.

The other half of the design is determinism. The chip is organized as functional slices (memory, vector, matrix, data movement) that data streams across in lockstep, and it has no caches, no arbitration and no reordering: nothing whose timing depends on the data. So the compiler knows exactly which cycle every operation will run and every value will arrive, and schedules the whole program in advance. A later paper (Abts et al., 2022) extended the same static scheduling to the links between chips. A multi-chip system runs like one big statically scheduled machine. There is no tail latency from contention, because there is no contention.

The SRAM capacity problem

230 MB is a lot of SRAM and very little for a language model. Holding the weights alone, against the HBM chips:

# Holding a 70B-parameter model, and the fastest one stream of tokens could
# go, on four kinds of chip (vendor specs). Decoding one token at batch 1
# reads every weight once, so a bound on its time is weight bytes divided by
# the memory bandwidth the weights sit behind, or FLOPs divided by compute,
# whichever is larger. Bounds, not measurements: they ignore the links
# between chips, the KV cache and any software overhead.
P = 70e9
chips = {   # name: (memory per chip in bytes, bandwidth to it in bytes/s, dense FLOP/s at 16-bit)
    "GroqChip (SRAM)":  (230e6,          80e12,   188e12),
    "H100 SXM (HBM)":   (80e9,           3.35e12, 989e12),
    "TPU v6e (HBM)":    (32e9,           1.638e12, 918e12),
    "TPU7x (HBM)":      (192 * 2 ** 30,  7.38e12, 2307e12),
}
for bytes_per_param, label in ((2, "16-bit weights"), (1, "8-bit weights")):
    W = P * bytes_per_param
    print(f"70B model, {label}: {W / 1e9:.0f} GB")
    print(f"  {'chip':18s} {'chips to hold it':>16s} {'tokens/s bound':>15s} {'on 8 chips':>11s}")
    for name, (mem, bw, flops) in chips.items():
        n = math.ceil(W / mem)
        bound = lambda k: 1 / max(W / (k * bw), 2 * P / (k * flops))
        if name.startswith("Groq"):
            # Layers are spread across chips and run one after another for a
            # given token: each chip reads only its own share, but in turn, so
            # one stream's time is all the bytes over one chip's bandwidth,
            # however many chips there are.
            print(f"  {name:18s} {n:16d} {1 / max(W / bw, 2 * P / flops):15.0f} {'-':>11s}")
        else:
            # Tensor parallelism: k chips each read their share at the same time.
            print(f"  {name:18s} {n:16d} {bound(n):15.0f} {bound(max(n, 8)):11.0f}")

Full script: sram.py

70B model, 16-bit weights: 140 GB
  chip               chips to hold it  tokens/s bound  on 8 chips
  GroqChip (SRAM)                 609             571           -
  H100 SXM (HBM)                    2              48         191
  TPU v6e (HBM)                     5              58          94
  TPU7x (HBM)                       1              53         422

70B model, 8-bit weights: 70 GB
  chip               chips to hold it  tokens/s bound  on 8 chips
  GroqChip (SRAM)                 305            1143           -
  H100 SXM (HBM)                    1              48         383
  TPU v6e (HBM)                     3              70         187
  TPU7x (HBM)                       1             105         843

A 70B model in 16-bit weights needs 609 LPUs just to hold it. If the layers are laid out across the chips as a pipeline, each chip reads only its own 230 MB, in turn, so one token’s pass takes the whole model’s bytes over one chip’s 80 TB/s: 1.75 ms, a bound of 571 tokens per second. Spreading the model over more chips doesn’t make that faster; it’s the fast SRAM that does, and the spreading is what it costs to have enough of it. Groq’s deployments use fewer chips than this (its Llama 2 70B demo linked 576 GroqChips, as The Next Platform reported in November 2023), which is only possible with weights stored in less than 16 bits (Groq uses mixed precision), so the 8-bit row is closer to what it runs.

70B model in 16-bit weights (140 GB) chips just to hold the weights (log scale) GroqChip (SRAM) 609 H100 SXM (HBM) 2 TPU v6e (HBM) 5 TPU7x Ironwood (HBM) 1 tokens per second for one stream, ideal bound 609 LPUs, pipelined 571 8 × H100, tensor-parallel 191 8 × TPU v6e 94 8 × TPU7x 422

SRAM buys speed per stream at the price of hundreds of chips.

Data: the output above

Against that, the HBM chips with eight chips each, using tensor parallelism (each layer’s weight matrices split across the chips, so all eight read their shares of every layer at once): H100s are bounded at 191 tokens per second, and Ironwood, with 7.4 TB/s per chip, at 422.

These are idealized bounds for one stream at batch 1, and the idealization isn’t even. Tensor parallelism needs two all-reduces per layer, 160 per token for an 80-layer model, and at batch 1 each is latency-bound, so the HBM bounds are further from reachable than the pipeline’s, which only passes small activations from stage to stage. Groq also uses tensor parallelism across chips, and speculative decoding (Speculative Decoding, Derived) gets more than one token per pass over the weights: Groq’s speculative endpoint for a 70B Llama was reported at about 1,665 tokens per second (The Register, December 2024, citing Groq), far above the 571 that one token per pass allows.

Throughput is a different question from one stream’s speed. A 609-stage pipeline can have hundreds of streams in flight, one per stage, so the LPU’s aggregate throughput is far above 571 tokens per second; what limits it is that every stream’s KV cache must also fit in SRAM next to the weights, which caps the number of sequences and their length. And prefill is compute-bound: at 188 TFLOPS per chip, an LPU has a fifth of an H100’s dense compute, so long prompts are where it is weakest. Cost per token depends on all of this, and on what 609 chips cost against 8, which this calculation leaves out.

When each wins

  • GPUs are the default: one chip family for training and inference, every framework, every model shape, and the largest software ecosystem. Their generality costs some efficiency on any single workload and buys the ability to run whatever comes next.
  • TPUs win on large, regular, dense workloads at scale, where the pod interconnect and compiler planning pay off, and for whoever can commit to Google’s stack (JAX or PyTorch through XLA).
  • LPUs win on latency for a single stream: the fastest tokens per second per user, and predictable timing. They lose on capacity per chip, which means many chips per model and little room for KV caches, and on compute for long prompts. In December 2025 Groq announced a non-exclusive license of its inference technology to Nvidia, with its founder joining Nvidia, a sign the low-latency, SRAM-heavy design is being taken seriously.

Other designs sit between these points. Cerebras puts a whole wafer’s worth of cores and 44 GB of SRAM on one wafer-scale chip (the WSE-3, by Cerebras’s announcement); AWS’s Trainium and Inferentia are built around systolic tensor engines in the TPU’s lineage. Every one of them is a different answer to the question this series keeps returning to: where do the bytes live, and how fast can the arithmetic get at them?

Where this leads

  • Where Models Run asks the same question one level up: which model should run where, on a server or on the device, and how to route between them.

Takeaways

  • Specialized chips gain by removing general-purpose machinery: caches, schedulers, speculation.
  • A systolic array wires data reuse into the hardware: an n×nn \times n array uses each value it takes in nn times. Simulated on 3×3, it does 27 multiply-adds in 7 cycles. Efficiency against HBM still depends on the problem: at batch-1 decoding, the array waits on memory.
  • TPUs pair systolic matrix units with HBM, torus interconnects of up to 9,216 chips, and a compiler that plans all data movement.
  • Groq’s LPU keeps the model in on-chip SRAM (230 MB per chip at 80 TB/s) and schedules every cycle at compile time, so its timing is deterministic.
  • Holding a 16-bit 70B model takes 609 LPUs; a pipeline over them is bounded at 571 tokens/s for one stream (more with speculative decoding), against 422 for 8 Ironwood chips in an idealized tensor-parallel bound. SRAM buys speed per stream at the price of many chips and little room for KV caches.

References

Start with the original TPU paper and the Scaling Book’s chapter on TPUs, then Groq’s ISCA 2020 paper on the deterministic design; Kung’s 1982 paper is the classic on systolic arrays.


Written by

Tan Le

AI engineer. I write these notes to learn LLM internals and AI engineering in depth, and to keep what I learn where I can find it again.