TorchTPU · Gemma 2B greedy decoding on TPU v5e-1: Boosting inference to 21,645 tok/s

TorchTPU is a PyTorch backend that runs PyTorch programs on Google TPUs through a nativetpu device. This report measures greedy text generation with Gemma 2B on one TPU v5e chip usingthe public release torch_tpu 0.2.0. No serving engine exists for TorchTPU, so the model, the keyand value cache, the decoding loop, and a decode attention kernel in Pallas were written for thisstudy. The workload repeats a previous vLLM experiment by the author on the same chip class:512 sequences and 128 generated tokens per sequence. Four experiments add one change at a time.Compiling the decode step with torch.compile raises throughput at batch 128 by a factor of 5.1,from 2,095 to 10,624 tokens per second. At batch 512, the Pallas kernel registered as a PyTorchoperation adds 5.6 percent over compiler generated attention, and running eight decode steps insideone compiled program with the scan operator adds 21.1 percent. With the real weights the finalconfiguration reaches 21,645.8 tokens per second, with a batch time to first token of 222 milliseconds,an inter token latency of 22.1 milliseconds, and identical tokens across repeated runs. The vLLMexperiment had reached 19,511 tokens per second; the two systems differ in scheduling and cachedesign, which the report lists. An estimate of the arithmetic places the final configuration at about55 percent of the chip’s peak bfloat16 compute. The report explains each component, how each resultwas verified, and the practical lessons for engineers who run inference on TorchTPU.