The Decode Block Size Heuristic in TPU Ragged Paged Attention Reduces LLM Inference Throughput by 28 to 69 Percent

The efficiency of large language model serving depends on kernel level software choices whose individual effects are rarely measured in isolation. This paper measures one such choice: the block size that the Ragged Paged Attention v3 (RPA) kernel uses to read the key and value (KV) cache during token generation on Google TPUs. The kernel selects this size automatically, at compilation time, through a heuristic inside the vLLM TPU backend. We replaced the installed kernel with an instrumented, patched copy at import time and verified that the substitution itself adds no measurable overhead: runs with and without it agree within 0.17 percent on one model and 0.02 percent on the other. We then swept the block size on a Cloud TPU v6e-4 slice (four chips in a 2x2 mesh, tensor parallelism TP=4) under a fixed batch workload. On Gemma 2 27B-it, overriding the heuristic with a block size of 256 increased generation throughput by 27.7 percent (7,621 to 9,732 tokens per second). On Gemma 4 31B the effect was larger and formed a plateau: every block size of 128 or less produced a gain of 67.6 to 68.7 percent over the heuristic (4,776 to 8,057 tokens per second at block size 32), with the response flat within 0.7 percent from block size 16 to 128. The mechanism follows from arithmetic intensity: decode attention is memory bound at every batch size and grows linearly with batch, so it dominates the generation step exactly in the throughput regime, and an oversized block spends the kernel's memory bandwidth loading padding that is masked and discarded. An earlier study on a single TPU v5e chip (DOI: 10.5281/zenodo.21212010) measured the same parameter at 24 to 27 percent without sharding; the present results show that the effect persists, and increases, when the model is sharded across four chips with collective communication in every layer, because that communication accounts for roughly 7 percent of the step and per chip memory traffic accounts for the rest. All results derive from timestamped logs and per run JSON exports. The record contains the article (PDF) and the complete LaTeX source with figures.