Engineers doubled Gemma 4 31B prefill speed on TPU v6e

Sail Research published a 16-minute engineering breakdown in which prefill optimization for the Gemma 4 31B model on a four-chip TPU v6e node (2x2 topology) raised MFU from ~32% to ~63%, increased throughput from 18,228 to 36,669 tok/s, and reduced TTFT from 449.4 to 223.4 ms.

Compiled by Sergey KostenchukPublished 2026-09-09Updated 2026-09-18

2026-09-09 Research

🤖 Engineers doubled Gemma 4 31B prefill speed on TPU v6e

Sail Research published a breakdown of prefill optimization on a 4-chip TPU v6e node with the SGLang-JAX stack: MFU rose from ~32% to ~63%, throughput at an 8192-token prompt went from 18,228 to 36,669 tok/s, and TTFT dropped from 449 to 223 ms. Techniques: query block 512 instead of 32 in Ragged Paged Attention, collective matmuls instead of monolithic AllReduce, and custom Pallas kernels — XLA failed to auto-fuse collectives.

🌍 HBM makes up more than 63% of accelerator cost, so TPU v6e with H100-like FLOPs but 2.5x less memory is cost-effective for prefill-heavy inference of small and medium models.

👤 The main inference overhead is not FLOPs but communications and the small “tail” of operations: in the case study, Perfetto/XProf profiling and a breakdown of when XLA is powerless and a custom kernel is needed.

Source 1: https://www.sailresearch.com/blog/tpu-v6e-gemma Source 2: https://www.sailresearch.com/blog/htdym

Sources