Transformer accounting on TPU

Tick each objective you can deliver out loud without notes. Then write the one-paragraph explanation you would give a room of GPU engineers. Saved attempts stay below, so you can compare them over the week.

• I can count the FLOPs of any einsum, and derive 6 × parameters × tokens for training from 2× forward and 4× backward.
• I can count the parameters of a dense GQA transformer and of an MoE, split into MLP, attention and vocabulary.
• I can compute KV-cache bytes per token and per sequence, and say how many sequences fit on a chip after the weights.
• I can derive when attention FLOPs dominate (T > 8D) and why decode attention is never compute-bound.
• I can compute the compute-bound batch for an MoE layer and explain why it is E/k times larger than for a dense layer.
• I can estimate training memory (10 bytes per parameter plus activations) and show that large clusters are needed for FLOPs, not memory.
• I can name the TPU-specific consequences of model shape — padding to 256, head dimension, KV heads — and where they show up later.

Prompts to answer out loud

1. Estimate on a whiteboard, in two minutes, the FLOPs and the Ironwood-pod-days to train a 70B dense model on 15T tokens.
2. Why does an MoE make decode harder to make compute-bound, and by what factor?
3. A 70B model has 320 KB of bf16 KV cache per token. How many 32k-token conversations fit beside the weights on 8 Ironwood chips?
4. Why is decode attention memory-bound at every batch size, while the MLP becomes compute-bound at large enough batch?