From 1,540 to 19,511 Tokens per Second on a Single TPU v5e Chip: A Measurement Study of Large Language Model Inference Optimization

This paper reports a measurement study of large language model inference on a single Google TPU v5e chip. The subject model was Gemma 2B, served through the vLLM inference engine on the JAX-based TPU software stack. The study asked one question: how much of the chip's potential does a default serving configuration leave unused, and how far can disciplined measurement recover it without changing the model or hardware. The answer was a factor of 12.7. The starting configuration produced 1,540 tokens per second; the final configurations produced 16,700 at the recommended operating point and 19,511 at the throughput maximum. The gain came from three interventions, each identified by profiling the chip rather than by guesswork: a change to how output tokens are selected, worth 2.3x on its own; a systematic increase in the number of requests processed together, worth a further 4.3x; and a retuning of the attention kernel's block size to match the workload's sequence lengths, delivered through a custom Pallas kernel harness, worth a further 24 percent. No model weights were changed, no accuracy was sacrificed, and no additional hardware was used. Because the chip and its power envelope stayed constant, every efficiency metric improved by the same factor: the energy bound per token fell from roughly 128 millijoules to 10, and at published on-demand pricing the cost of generation fell from roughly 4.6 million tokens per dollar to 58.5 million. The study is equally defined by what did not work. Speculative decoding reduced throughput by a factor of six in this workload's operating regime; weight and cache quantization turned out to be unavailable on this software stack for this model; and a precision optimization that looked certain on paper reduced throughput by 14 percent when applied, and was reverted. Part I reports the results; Part II discloses the methods, code paths, patches, and verification procedures, strategy by strategy. Google ML Developer Programs and the Google Developers Program supported this work by providing Google Cloud credits.