Sail Research's engineering team published a breakdown of prefill optimization for the Gemma 4 31B model on a node of four TPU v6e chips in the SGLang-JAX stack: chip utilization increased from ~32% to ~63% MFU, throughput for an 8192-token prompt increased from 18,228 to 36,669 tokens per second, and time to first token (TTFT) decreased from 449.4 to 223.4 ms. The recipe is built on enlarging query blocks in Ragged Paged Attention, collective matrix multiplications instead of monolithic AllReduce, and custom Pallas kernels.

image
image
image

What happened

On September 9, 2026, Sail Research, authored by Ben Mayer, released the material "Chasing Speed of Light on TPU v6e" — a 16-minute engineering breakdown of how to squeeze the most out of the prefill phase of Gemma 4 31B on a node of four TPU v6e chips connected in a 2x2 topology, on top of the SGLang-JAX stack. The first technique is increasing the query block size in Ragged Paged Attention from 32 to 512, so that attention is processed in large chunks and fewer resources are spent on the overhead of small operations. The second is replacing monolithic AllReduce with collective matrix multiplications (collective matmuls) overlaid directly on computations over the ring inter-chip ICI network. The third is custom Pallas kernels, written because the automatic fusion of collective operations in the XLA compiler proved insufficiently efficient. Every claim in the material is accompanied by profiles in Perfetto and XProf, and the breakdown stems from the recent open release of the HTDYM (How To Deploy Your Model) tool.

Context

Prefill is the inference phase in which the model processes the entire input prompt before generating the first token; with long contexts and batch processing, it is precisely this phase that determines response speed, so RAG systems, long-context agents, and batch processing are prefill-heavy workloads. MFU shows what fraction of the chip's peak performance is actually spent on computations, so the initial values meant that more than two-thirds of the chips' capabilities did not reach useful work. TPU v6e is an accelerator whose BF16 FLOPs are on par with H100, but whose HBM memory volume is 2.5 times smaller, and whose inter-chip ICI network throughput is around 180 GB/s compared to approximately 400 GB/s for AllGather on H100: the chip is strong in computation but limited by memory and inter-chip exchange. The material continues the open release of HTDYM (How To Deploy Your Model) from August 27, 2026 — a simulator that allows modeling inference performance and choosing a chip before purchasing hardware.

Why this matters for the industry

For the industry, the material changes the logic of choosing hardware for inference of small and medium models: the question "which chip is faster" gives way to the question "which chip is cheaper for a specific workload." According to Sail Research's estimate, HBM memory accounts for more than 63% of the production cost of an accelerator, so a chip with comparable FLOPs but smaller memory volume wins on price precisely for prefill loads. Against this backdrop, TPU v6e becomes an economically attractive option for RAG systems, long-context agents, and batch processing, and the presented optimization path is a reproducible recipe for increasing MFU on non-standard architectures. Teams already working on TPU v6e and JAX receive a proven sequence of steps to reduce TTFT and increase throughput without replacing hardware, and HTDYM allows justifying the choice of a chip through modeling before purchase. If the results are confirmed on other models and topologies, the niche of accelerators that are strong in FLOPs but weak in memory may systematically establish itself in the inference market.

Why this matters for users

For the reader, the material is useful as a practical case, not an announcement: it contains exact figures, profiles in Perfetto and XProf, and an explanation of when the XLA compiler stops coping and one has to drop down to the level of custom Pallas kernels. The main practical lesson is that in inference, the bottleneck is often not FLOPs, but communications between chips and the "tail" of small operations, so profiling and tuning pay off faster than changing hardware. Two artifacts are available immediately: the open HTDYM simulator, which can be used to estimate performance and cost before buying chips, and the optimization recipes themselves for the SGLang-JAX stack. For developers on JAX, this is a ready-made sequence of actions; for others, it is public benchmarks for MFU, throughput, and TTFT that can be relied upon when choosing hardware for their own prefill tasks.

What is still unknown / limitations

The material is an engineering blog case, not a peer-reviewed article with an ablation breakdown: there is no breakdown of the MFU increase by individual techniques, so the contribution of each of the three optimizations individually is unknown. The publication lacks repeated runs and variance assessment, as well as comparison with baseline configurations and other inference stacks. The figures were obtained for one model, one 2x2 topology of TPU v6e, and an 8192-token prompt, so transferring the results to other models, chips, and context lengths remains unproven: reproducibility here should be understood as principled, not as a guaranteed outcome. HTDYM is described only as a recent open release, without data on its accuracy and validation.

Sources

Author

Look at AI, editorial team