Chasing Speed of Light on TPU v6e

In a blog post, the team at an AI infrastructure company detailed their optimization of Gemma 4 31B on Google's TPU v6e, boosting prefill efficiency from ~32% to ~63% MFU. They highlighted the chip's trade-offs, including lower memory bandwidth and slower inter-chip interconnect compared to NVIDIA H100, but noted its cost-effectiveness for certain workloads.

We recently released HTDYM /blog/htdym , our performance modeling infrastructure which helps us assess which chips are most cost effective to deploy models on. If HTDYM always spat out that serving on NVIDIA or AMD was best, our jobs would be easy. These platforms have strong software ecosystems and highly optimized off the shelf codepaths in existing inference engines we can take advantage of. Unfortunately, this is often not the case. Depending on the model, more niche/specialized accelerators often offer performance at a substantial discount relative to the big players. We found this to be the case with Gemma 4 31B, which perhaps unsurprisingly being a Google model , appeared on paper to be a good fit for Google’s TPU v6e, and decided to give it a shot. In this post, we’ll break down our journey optimizing Gemma 4 on TPU v6e, taking prefill from an unimpressive ~32% MFU to a substantial ~63%. To understand what optimizations were necessary and why, let’s first look at what this chip is good at, and where it’s likely to struggle. Meet the v6e TPU v6e is a pretty weirdly spec’d accelerator by modern standards. Compared to an H100, it matches it in BF16 FLOPs, but has 2.5x less HBM, and less than half the memory bandwidth. 1 fn-1 Its inter-chip interconnect ICI is also constrained compared to H100. On a standard 8x H100 node, every GPU is connected to every other GPU via NVLink. On TPU v6e, each chip is connected only to its neighbors. In a 2x2 topology, the 4 chips form a ring. The ICI bandwidth numbers Google and NVIDIA quote are perhaps intentionally difficult to compare directly. In practice, an of a matrix sharded across all chips on a single host reaches an effective bandwidth of ~180 GB/s on TPU v6e vs. ~400 GB/s on H100. 2 fn-2 Low memory capacity, low memory bandwidth, and slow ICI... why would we ever want to use a chip like this? Memory is expensive, some estimates find HBM alone is north of 63% of the cost of producing an AI accelerator today.<sup 3 fn-3 </sup Not all models and workloads require a ton of memory/bandwidth: a small to medium sized model used for prefill-heavy tasks summarization, classification, etc. can function well in this regime.<sup 4 fn-4 </sup By skimping out on HBM but still equipping the TPU v6e with two giant 256x256 systolic arrays https://jax-ml.github.io/scaling-book/tpus/ appendix-b-how-does-a-systolic-array-work for matrix multiplication, Google created a chip with an incredible FLOPs/$ value. But even with a small-ish model, realizing that value in the real world is tricky. 32 GB of memory is barely enough to hold a 16B parameter model, let alone its KV cache. Thus serving a medium-sized model like Gemma 4 31B requires sharding it across multiple chips, which introduces substantial inter-chip communication requirements throughout the forward pass, making it easy for TPU v6e’s slow ICI to become a bottleneck. Let’s spin up an off-the-shelf implementation of Gemma 4 on TPU v6e and see if our fears come true here. Baseline After spinning up SGLang-JAX on a TPU v6e 2x2 4 total VM, we can follow SGLang-JAX’s guide https://github.com/sgl-project/sglang-jax/blob/main/docs/developer guide/benchmark and profiling.md to profiling, feed in a prompt of 8192 tokens, and open the .trace.json.gz file JAX generates on ui.perfetto.dev https://ui.perfetto.dev . This trace is correct but pretty minimal. There’s more detail in the .xplane.pb files JAX generates alongside them, which we can visualize using Google’s XProf https://openxla.org/xprof trace viewer. However, XProf’s trace viewer is kinda crummy to use for a number of reasons. A script like this https://gist.github.com/transmissions11/0de193188187a1a590a669fbcde26240 will let you convert .xplane.pb files into .trace.json.gz files, which you can then visualize with Perfetto. Much better. Digging into this trace, we can see that each group of 6 layers 5 swa : 1 global attention takes about 42.8 ms. Is this good or bad? We can use Model FLOPs Utilization MFU as a benchmark. MFU asks a simple question: if the only thing we did were the model’s required floating point operations, executed at the hardware’s peak FLOPs rate, how long would its forward pass take relative to what we actually measured? Concretely: For a simple dense transformer, ignoring dot-product-attention FLOPs, we can approximate that compute-only floor as: For Gemma 4 31B on 4 v6es, that gives us a floor of about 138 ms, while the baseline takes about 428 ms: So roughly two-thirds of the wall-clock time is overhead beyond the model’s bare FLOPs. Some of that is unavoidable memory movement and communication, but the question is how much. Let’s dig into the trace to see where there’s fat to cut. Attention One of the first things that jumps out in the trace is how much time is spent on attention. Each sliding window layer spends ~1.539 ms in the Ragged Paged Attention kernel. This is almost half as much time as we’re spending in our entire MLP block , which should make us suspicious. During prefill at relatively short sequence lengths we’re measuring at , time spent in MLP should dominate time spent in attention. With a little roofline math<sup 5 fn-5 </sup we find that across our 4 chips, even assuming no overlap of compute and memory , this kernel should run 10x faster : Digging further, we find the kernel isn’t slow because it’s poorly implemented, but rather just because it’s poorly tuned, causing the same keys and values to be read over and over. The RPA kernel doesn’t just read the KV cache once. Instead, it works through the prompt in blocks of queries, and for each block it loads in every key that block attends to while the queries stay parked in VMEM, so each key gets pulled out of HBM once per query block that wants it. In Gemma’s sliding-window layers a key is wanted by the 1024 queries following it. By default SGLang-JAX was using query blocks of 32, so those queries get spread over 32 separate blocks and thus each key gets fetched from HBM 32+ times. Tuning block size to 512 takes the sliding layers from 1.539 ms to 261 μs, within ~2x of the serial roofline computed above. 6 fn-6 Collectives Continuing to look over the trace, we can see there are two large s per layer. The presence of these is not surprising. We are sharding the model using tensor parallelism across 4 TPUs, so each TPU computes only a portion of the attention and MLP matmuls. After these sharded matmuls, the chips need to combine their partial results before proceeding, which is achieved via . Unfortunately, these s take up a significant amount of time, approximately ~788 μs each, or 1.576 ms per layer. That’s almost 27% of the 5.88 ms we’re spending in total per sliding attention layer This is not because the s are implemented poorly. In fact they achieve ~93% of the theoretical limit for an of this size on this topology<sup 7 fn-7 </sup : This is exactly what we were worried about when we were evaluating the v6e earlier in this post: the ICI is just not that fast, and thus can become a bottleneck. Luckily, we don’t just have to accept slow collectives in our critical path. Collective Matmuls Just because we can’t make the s faster they’re already running close to the physical limit , doesn’t mean we have to wait for them. What we’d like to do is overlap our collectives with our compute, hiding the ICI time behind matmuls we were going to run anyway. To make this possible, the first thing we want to do is split each into two halves. An is just a followed by an , which both take equal time, so our 788 μs becomes two ~370 μs collectives. On its own, splitting buys us nothing, we still have to move the same number of bytes. But unlike a monolithic , the and don’t have to happen back-to-back. We can defer the , leaving the residual sharded until the beginning of the next block, right before the matmuls that actually need the full set of tokens again. Now we just need to overlap each of these smaller collectives with their preceding and succeeding matmuls respectively... but how? Consider the for the MLP’s up/gate projections. Each chip starts holding a quarter of the tokens and a quarter of the weights. We need to ultimately feed all tokens through each chip’s shard of the weights, but we don’t need to feed them all through at once . Each chip can send its current slice of tokens to its neighbor while matmuling the slice it already has. When the next slice arrives, it repeats the process: matmul that slice while also forwarding it to its neighbor. After 3 hops, every chip has multiplied every token by its local weight shard. As long as each slice’s matmul takes at least as long as an ICI network hop, the adds zero actual wall-clock time. The pattern is basically the mirror image. For it we’ll produce the matmul output in slices and send each finished slice around the ring while computing the next one. These fused matmul-collective pairs are called collective matmuls, and Wang et al. https://dl.acm.org/doi/10.1145/3567955.3567959 describe the general construction. Do we actually have enough compute to hide behind? In the MLP block, yes, plenty: The attention block is tighter. The qkv slices come in just under the hop time, and the tiny o proj has only ~200 μs of compute which we cannot hide 370 μs of behind. It’s okay if we can’t get perfect overlap: 300 μs or so of exposed ICI per layer is still a lot better than 1.58 ms. Asking XLA Nicely In theory, XLA the ML compiler stack used for the TPU already knows this trick, and should do it for us automatically. In practice, it required a bit of fiddling to trigger for us, but eventually we see: The trace clearly shows that the is gone, and each half of the MLP is executed via a collective fusion. That’s what we want to see Unfortunately, the implementation XLA has generated is just not very good, only shaving off ~0.45 ms per-layer of the ~1.3 ms we should ideally be able to hide. XLA is awesome when it works, but when it doesn’t, we’re sort of just left scratching our head and trying out various poorly documented compiler flags in hopes of coaxing XLA into generating better code. We tried a lot to get XLA to generate collective matmuls that got the performance we expected, but were not able to push them much beyond the disappointing performance we got out of the box. Fortunately, as of 2023, we have an escape hatch we can use in cases like these, when we know where XLA is falling short and want to take matters into our own hands: Pallas. Pallas Pallas is JAX’s kernel language for TPU. With Pallas, we get much finer control over memory movement, pipelining, and each chip’s communication with its neighbors over ICI. The docs https://docs.jax.dev/en/latest/pallas/tpu/distributed.html have worked examples of how to implement and in Pallas, exactly the collectives we need, and we used these as the starting point for our collective matmul kernels. They are teaching examples, though, and the gap between correct and fast in practice turned out to be a fair amount of work. Most of this was straightforward but nonetheless time consuming kernel engineering, but we also ran into quite a few surprising potholes with JAX and Mosaic the Pallas compiler that we didn’t expect, which we’ll lay out below. Mesh device ordering When we were measuring our first kernels, they were mysteriously about twice as slow as they should have been. The cause was how we were initializing our mesh: What’s the problem? Well, Mesh is sensitive to the ordering of the devices array you pass it, and on a 2x2 topology jax.devices returns the order 0, 1, 2, 3 . Makes sense right? Well, look at what this device 0 → 1 → 2 → 3 path looks like on the physical mesh: This naive ordering results in a path that would require diagonal links between devices, which we do not have, silently downgrading our performance by almost a factor of ~2x. This took us a while to notice because we were comparing against XLA’s , which is agnostic to device order, so nothing hinted that it was the mesh not Pallas that was the problem. Despite having some presence in the docs, this way of initializing a Mesh is not recommended. The jax.make mesh function exists specifically to solve this problem by “automatically comput ing a good mapping from a set of logical axes to a physical mesh,” and on a 2x2 v6e machine, it does. However on a 2x4 machine make mesh actually gets this topology wrong, emitting an ordering with three diagonal hops in it It turns out, this is because JAX optimizes the ordering with TPU v5e in mind here, which has a 4x2 topology rather than a 2x4 . Google is aware of the issue and we hope to see it fixed soon. 8 fn-8 Not so free transpose Our kernel needed to grab one 672-wide column block of its weight per ring step, but Pallas TPU BlockSpec s require the final block dimension to be a multiple of 128 or equal the full array dimension .<sup 9 fn-9 </sup To work around this we wanted to repack the weight so each column block sits on its own leading axis, which requires a transpose. We were only doing this once at load time, so it should have been costless during the rest of the forward pass. In practice, this ended up adding a ton of overhead, because XLA doesn’t actually move the bytes of the tensor around at load time if you perform a transpose, it just silently relabels the array as column-major. That’s normally a great trick, but Pallas kernels can’t work with a relabeled array, they need the bytes physically in the order the shape claims. Thus before every invocation of these kernels XLA was forced to wire a real copying transpose into the forward pass, adding ~173 μs/layer of overhead. To fix this, we use device put with a Format argument at load time to force each weight into plain row-major order. Since the bytes aren’t in that order yet, JAX has to actually move them, which is what we want pay for the transpose once, at load, instead of on every call . 10 fn-10 Mysterious SyncWait s At various points while optimizing, seemingly minor changes to block size or other components of our kernels would sometimes result in unexpected slowdowns, with no clear explanation of where the time went on the trace. After fiddling with some XLA flags<sup 11 fn-11 </sup , we were able to elicit some more information that helped us figure out where the missing time was going: “ SyncWait sequencer overlay ” events. What do these mean? The name turned out to be fairly literal once you decode the TPU jargon. Each TPU TensorCore has a core sequencer that fetches instructions from a small dedicated instruction memory IMEM and dispatches them to the core’s compute units.<sup 12 fn-12 </sup Each TPU generation’s IMEM size isn’t explicitly documented anywhere, but it appears to be around ~4 MiB for v6e according to reverse engineering efforts.<sup 13 fn-13 </sup Programs that outgrow it are split into instruction overlays and streamed in from HBM during execution.<sup 14 fn-14 </sup Those “ SyncWait sequencer overlay ” events were instruction-fetch stalls : the sequencer waiting for overlays to stream in from HBM. Small changes like larger block sizes can result in dramatic code size increases, as Mosaic seems to prefer unrolling jnp.dot rather than emitting a compact loop e.g. a 2048, 5632 × 5632, 1536 matmul via jnp.dot compiles to ~4.45 MB . Once we exceed the 4 MiB IMEM limit, we’re at the mercy of the opaque TPU runtime, which can’t always perfectly overlap overlay fetches from HBM with compute, and leaves us with nearly invisible stalls in the critical path. By using smaller block sizes, we were able to keep our kernels under the IMEM limit and avoid sequencer issues going forward. The Long Tail After solving those issues, our fused kernels close most of the gap to the roofline<sup 15 fn-15 </sup : | Callsite | Pattern | XLA | Our Pallas | Roofline | |---|---|---|---|---| | MLP up proj | | 1528 μs | 1178 μs | 1032 μs | | Attention qkv proj | | 845 μs | 575 μs | 465 μs | | MLP down proj | | 1049 μs | 690 μs | 516 μs | | Attention o proj | | 830 μs <sup 16 fn-16 </sup | 532 μs | 416 μs | We’re now at ~52% MFU, pretty good However, there’s still plenty of room to improve, an extra copy here and unfused op there, which adds up quickly. Chasing down these death-by-a-thousand-cuts slowdowns is what brought us all the way to ~63% MFU. Here’s two of the most impactful: - Gate mul/gelu between the MLP up proj / down proj The MLP gate proj / up proj kernel writes its output 176 MB at to HBM, but then right after, an XLA fusion implementing reads all of it back out and writes 88 MB back. Interleaving the gate/up columns inside the packed weights makes it easy for us to fold this gate computation into the kernel’s output stage while they’re still sitting in VMEM, eliminating this redundant HBM read. - Norms, transposes, and copies after qkv proj Between each qkv proj kernel and attention kernel invocation, we have around 500 μs of copies, dtype conversions, slices, transposes caused by XLA struggling to optimize across the Pallas boundary. Folding the q/k/v RMS norms and RoPE computation into the Pallas kernel gets rid of the slow conversions, and emitting q/k/v as three separate arrays instead of a packed array gets rid of the extra copies. Adding It All Up In the end, we achieve the following performance for an 8192-token prefill on a 4-chip 2x2 v6e node: | | Stock SGLang-JAX | Ours | Speed of Light | |---|---|---|---| | Throughput | 18,228 tok/s | 36,669 tok/s | 58,431 tok/s | | Forward pass time | 449.4 ms | 223.4 ms | 140.2 ms | | MFU | 32% | 63% | 100% | We doubled the throughput of stock SGLang-JAX, so every prefill token now costs half as much as it did when we started which was already pretty good thanks to v6e’s low cost . We’re even slightly outperforming the MFU SGLang gets out of the box for Gemma on H100, on a chip with a fraction of the memory/ICI bandwidth and much thinner software ecosystem. Our forward pass is now almost entirely spent in Pallas kernels, with some small norms in-between: Besides unavoidable real world chip overheads, the ~83 ms gap to 100% MFU can be attributed to: - The RPA kernel running at ~48% of its serial floor and ~30% of its overlap floor - The ~200 μs per layer of unhideable ICI communication after the o proj - The ~200 μs per layer of norms, residual adds, and required padding For a sense of where the ceiling is: as far as we can tell, the highest prefill MFU anyone has publicly reported is Google’s own 76%, for PaLM 540B on TPU v4.<sup 17 fn-17 </sup And that’s with an architecture heavily co-designed for high throughput serving on that hardware e.g. parallel attention/MLP , so we’re pretty happy to be within striking distance. 18 fn-18 Takeaways/Opinion Zooming out, nothing in this post actually required deep, preexisting TPU knowledge. All these optimizations came from the same core loop: 1. Compute what an operation should cost from first principles 2. Measure what it actually costs 3. Investigate the gap as a bug until proven otherwise Knowing how to compute rooflines is essential for figuring out where to press further and when to move on. Otherwise, an attention kernel taking 1.539 ms and a 788 μs both just look like “how long that operation takes on this chip.” It’s only when you put them next to a roofline that it becomes clear which is a poorly tuned kernel versus a physical limit you have to design around. When it comes to TPUs specifically, I have a couple complaints about JAX/XLA. To be clear: XLA is a very strong compiler, it’s capable of sophisticated fusion and generally achieves impressive performance out of the box. But one aspect of its programming model I found frustrating is how often performance failures are silent . When JAX/XLA notices something is wrong weights stored in the wrong orientation, mesh ordering is bad, etc. , it just finds a workaround e.g. insert a big copy and moves on. If your JAX can be compiled to something correct and functional, it will be, no matter the performance impact. This is great when you just want to scale your research code up quickly, but when we’re trying to achieve high MFU, this is frustrating I’d much rather JAX/XLA stop in their tracks and warn me that my weights have an inconvenient layout, or my mesh’s device order is unfamiliar, and allow me to opt-in to a workaround, rather than silently generating terrible code. The same frustrations show up with collective matmuls too. If overlapping an with a matmul is key for getting peak performance on this hardware, I don’t want to have to write code that looks like an followed by a matmul and hope XLA recognizes that it can fuse them. As Horace He has opined about before https://www.youtube.com/watch?v=139UPjoq7Kw , fiddle with your code until the compiler recognizes it can do this very specific optimization is a terrible programming model I believe these optimizations should be something the programming model lets you express explicitly , with predictable semantics and performance characteristics. shard map https://docs.jax.dev/en/latest/201/shard-map.html is a good step in this direction, and I’d like to see more “strict/explicit mode”-esque JAX features in the future. Despite these frustrations, overall I quite enjoyed working with TPUs, and found the software stack to be well thought out and approachable, though more examples of real, performant Pallas kernels would go a long way. The level of fusion XLA is able to achieve out of the box is really quite impressive, and JAX profiles are super readable and well attributed. I think we’ll be seeing a lot more of TPUs going forward, especially now that they’re increasingly being deployed outside of Google’s walled garden https://www.blackstone.com/news/press/blackstone-announces-joint-venture-with-google-to-create-new-tpu-cloud/ . P.S. If chasing speed of light on unfamiliar hardware /blog/improving-decode-throughput-gaudi-3 sounds like your kind of thing, Sail is hiring /careers . Footnotes 1. With H100 SXM as the reference point. H100 PCIe uses HBM2e memory. ↩ fnref-1 2. See the networking details appendix of the GPUs chapter of the Scaling Book https://jax-ml.github.io/scaling-book/gpus/ appendix-b-more-networking-details . ↩ fnref-2 3. epoch.ai/data-insights/ai-chip-component-cost-shares https://epoch.ai/data-insights/ai-chip-component-cost-shares ↩ fnref-3 4. Of course, thanks to prefill-decode disaggregation the composition of our workload doesn’t really matter, because we can just do all our prefill on one chip and all our decode on another, but colocated serving is much simpler to reason about, and our workloads do skew prefill-heavy anyway. ↩ fnref-4 5. Each query attends to max keys, and thus the kernel computes QK pairs. Each pair costs FLOPs per query head: two matmuls and , both -length dot products. In terms of memory, a flash attention kernel loads , , , and through HBM exactly once each, and the score matrix never leaves VMEM. Thus: ↩ fnref-5 6. PR 913 https://github.com/sgl-project/sglang-jax/pull/913 integrated RPA v3 from vLLM’s TPU backend into SGLang-JAX. vLLM-TPU’s RPA v3 implementation scales the query block with the head ratio, which lands on 512 for this shape good . But for some reason, the SGLang-JAX team replaced it with a flat size of 32 , “to match RPA v2 precision characteristics.” As far as we can tell, v2’s accuracy came from keeping its softmax accumulators in f32, not from small blocks, and the error is identical to four significant figures even after changing the block size from 32 to 512 . ↩ fnref-6 7. jax-ml.github.io/scaling-book/sharding/ what-have-we-learned https://jax-ml.github.io/scaling-book/sharding/ what-have-we-learned ↩ fnref-7 8. We reported this to Google, who identified the mesh ordering issue and gave us a workaround. ↩ fnref-8 9. docs.jax.dev/en/latest/pallas/tpu/details.html https://docs.jax.dev/en/latest/pallas/tpu/details.html ↩ fnref-9 10. Note: HLO prints layouts minor-to-major {1,0} is row-major , and jax.experimental.layout.Layout takes them major-to-minor 0,1 is row-major , which can mess you up and indeed messed us up if you’re not careful ↩ fnref-10 11. The “Tensor Core Sync Flag” track only appears if you profile with the TRACE COMPUTE AND SYNC trace mode, and the waits are only labeled if you set --xla xprof register llo debug info=true in LIBTPU INIT ARGS . ↩ fnref-11 12. Norrie et al., A Domain-Specific Supercomputer for Training Deep Neural Networks https://dl.acm.org/doi/10.1145/3360307 . ↩ fnref-12 13. Via Grigory Evko’s libtpu notes: 65536 IMEM bundles https://gh.evko.io/crucible-notes/libtpu/appendix/per-gen-comparison-matrix.html × 64 B/bundle https://gh.evko.io/crucible-notes/libtpu/isa/bundle-model-overview.html = 4 MiB of IMEM. ↩ fnref-13 14. This is from reading in-between the lines of Google’s public papers https://doi.org/10.1109/MM.2021.3058217 “While an instruction cache backed by HBM would have been nice, a DMA target for software-managed instruction overlays was easier” combined with Grigory Evko’s libtpu reverse engineering https://gh.evko.io/crucible-notes/libtpu/barrier/special-purpose-sync-flags.html work. ↩ fnref-14 15. Here, we define the roofline as the pipelined floor for each collective matmul. With the work split into four slices around the ring, one slice’s matmul is always exposed and the other three overlap with the three ICI hops, so the floor is . We use 918 TFLOP/s of peak compute, a 122 μs hop 22 MB of activations at 180 GB/s , and Gemma 4 31B’s shapes at tokens at . A pure bound that ignores the exposed slice would be 393 μs for qkv proj and 367 μs for o proj , and leave the compute-bound callsites up proj , down proj unchanged. ↩ fnref-15 16. Slightly inflated by ~110 μs/layer of weight relayout we could fix just like we did for our Pallas kernels. ↩ fnref-16 17. Pope et al., Efficiently Scaling Transformer Inference https://arxiv.org/abs/2211.05102 . ↩ fnref-17 18. We’ve focused here on prefill; pursuing further optimizations for decode is a natural direction for future work. ↩ fnref-18