Attention, where every word in a prompt weighs every other word, dominates the cost of serving long inputs. On Google's Tensor Processing Unit (TPU), a chip built around matrix units that multiply grids of numbers, the arithmetic surrounding those multiplications leaves the units idle much of the time.
The team rearranged data so loaded operands get reused, replaced the repeated rescaling inside softmax (the step that turns raw scores into weights) with a cheaper checked scheme, and interleaved independent work to fill idle cycles.
Attention computation runs up to 4.14 times faster than the previous best TPU kernel, and the serving system vLLM answered long prompts up to 1.56 times faster with comparable accuracy.