Measuring What the TPU Scaling Playbook Predicts: Tensor Parallelism, Collectives, and a Per Layer Type Optimization in vLLM Serving

Abstract

Analytical cost models guide sharding decisions for large language model serving on TPUs: tensor parallelism adds two AllReduce collectives per transformer layer, decode steps are bound by streaming the weights from high bandwidth memory, and per chip memory is the weights divided by the shard count plus the key and value cache. This study tests these statements as preregistered quantitative predictions, at the level of individual compiled kernels, on a four chip TPU v6e slice serving Gemma 4 31B through vLLM. Seven hypotheses were recorded with their numbers before the campaign ran. The decode mechanism is confirmed: at matched batch and memory utilization, step time scales with the inverse of the shard count (measured ratio 1.99, predicted 2.0), and communication is 2 percent or less of the step. Two estimates are refined by measurement: collective times sit on a per operation overhead of 14 to 24 microseconds rather than on the modeled bandwidth and hop terms, and a two way sharded deployment has an initialization footprint of 58.3 GiB, larger than the four way total, so runtime residency rather than division sets memory feasibility. A documented per layer cost crossover between sliding window and global attention layers replicates across builds, and its cost ratio is found to vary with the sharding degree. Finally, a per layer type block size policy is deployed as a two line kernel change, verified in the trace, and measured end to end: decode attention per step falls 36.0 percent and throughput at 8,000 token histories rises 18.9 percent, within the 15 to 25 percent range estimated in prior work. The record includes the paper and the complete LaTeX source; the harness, kernel patches, and analysis tools are released