What the Chip Is Actually Doing During Decode

Part 2 established that exactly two of eight stages run on the accelerator. This post is about those two — and mostly about the second one, because decode is where the request spends nearly all its life and where the hardware is least happy.

---

Decode is awkward for a systolic array specifically

Post 2’s section 13 makes the general argument: decode collapses arithmetic intensity, because producing one token means reading every weight from HBM to use each one once. That argument applies to any accelerator.

There is a second penalty that is specific to systolic arrays, and it compounds with the first.

The same array, two batch sizes

A systolic array holds weights stationary in the grid and streams activations through it. Data reaches the cell at row r, column c only after r + c steps — so there is a fill phase before the array is fully working and a drain phase after the last row enters. That pipeline cost is fixed. It depends on the size of the array, not on how much work you brought it.

With a large batch, many rows stream through back to back, the pipeline stays full behind the leading wavefront, and the fixed cost is amortized across all of them. With a batch of one, you pay the entire fill and drain to push a single row through.

So decode on a TPU is penalized twice, independently:

1. Memory-bound — every weight read from HBM, used once. (Post 2, §13.)
2. Pipeline-bound — the array’s fixed fill and drain cost charged to a single row.

Both are fixed by the same thing, which is why batching dominates every other lever in this series. It is not one optimization among several. It is the mechanism that makes the hardware make sense at all.

---

Pallas, and what it is actually for

XLA compiles your model graph and does a good job of it. Pallas is the escape hatch for when a good job is not enough: a kernel language for writing TPU (and GPU) kernels directly, operating on tiles of memory rather than whole tensors, when you need control XLA’s fusion heuristics won’t give you.

For serving, the motivating case is paged attention, and it’s worth understanding why it needs a custom kernel rather than being expressible as ordinary operations.

A textbook attention kernel assumes a sequence’s keys and values are contiguous in memory. Paged attention deliberately breaks that assumption — that’s the entire point, since it’s what lets the KV pool be allocated in fixed-size blocks and avoid the fragmentation that would otherwise waste most of Part 0’s ~3.8 GiB. So the kernel has to:

• Walk a block table — an indirection layer mapping each sequence’s logical positions to scattered physical blocks
• Gather from non-contiguous HBM without materializing a contiguous copy, which would defeat the purpose
• Handle ragged batches, where every sequence in the batch is at a different length
• Keep the whole thing fused, so intermediate attention scores never round-trip through HBM

Express that as ordinary tensor operations and you get a gather, a materialization, and a fusion boundary in exactly the wrong place. Hence a kernel.

---

Where idle time comes from

“The chip is idle” is not a diagnosis. On one machine there are three distinct causes, they have different fixes, and they are distinguishable in a trace.

Host starvation. The TPU finishes a step and waits for the CPU to prepare the next one — tokenization, sampling, detokenization, scheduler bookkeeping. This is Part 2’s argument, and on a box where six of eight stages are CPU work it is the default suspect. In a trace it looks like device gaps that line up with host activity.

Dispatch overhead at small shapes. Every step has a fixed cost to launch, independent of how much work it contains. Decode steps are small, so that fixed cost is a large fraction of each one. This does not look like waiting on anything — it looks like the step itself being disproportionately expensive relative to the arithmetic it performs.

Recompilation from shape churn. XLA compiles per shape. Serving generates a lot of shapes — varying prompt lengths, varying batch composition, sequences growing a token at a time. Without bucketing, you get compilations mid-benchmark: enormous, obvious spikes that poison percentiles. This is also the one most likely to be silently fixed by a warmup you didn’t realize you were relying on, which is why Part 1 insisted on discarding warmup explicitly.

The reason to separate these is that they look identical in the metric and completely different in the trace. MFU says “low” for all three. Only a profile says which.

That profile — the actual walkthrough, with a real trace — is the thing both this post and Part 1 are deferring until there is hardware to capture one from.

---

Tensor parallelism: a prediction, and why it might be wrong

This is the part that needs the v6e-4, and it is worth writing down what I expect before measuring, so the measurement can embarrass it.

Sharding a model across 4 chips with tensor parallelism does two things. Each chip holds a quarter of the weights, so per token each chip reads a quarter as much from HBM — decode latency should improve substantially. And there is four times the total HBM, which grows the KV pool by considerably more than 4×, since the weights on each chip shrank too.

The cost is a collective: the chips must combine partial results, which means an all-reduce over ICI, per layer, per token. With 48 layers that is 48 collectives for every single token generated.

Here is the tension, and the reason I expect this to be less impressive than it sounds. At batch size 1, the tensor being all-reduced is tiny — one activation vector. A collective that small is not bandwidth-bound, it is latency-bound: dominated by the fixed per-hop cost, which ICI’s bandwidth advantage does nothing to help. So the win on weight reads and the loss on collective latency are both real, and which dominates is not something I can reason my way to from a spec sheet.

---

What is drafted and what is booked

Part 0 committed to drafting Parts 3 and 4 before booking any v6e-4 time, on the grounds that a few hours of TPU is not enough to also be working out what to measure. This post is that draft. It names three things the hardware has to settle:

• The split between bandwidth and pipeline occupancy in decode’s inefficiency
• Whether a hand-written Pallas kernel beats what XLA emits, and by how much
• Whether tensor parallelism on four chips is worth its collective cost at this model size

Each is a question with a falsifiable answer, which is the point of writing them down first.

---

Part 4 is next: precision and quantization — the lever Part 0’s memory table already showed to be the largest available on this machine, taking the KV pool from ~3.8 GiB to ~16 GiB, and the one where the cost is paid in accuracy rather than latency.