Making gpt-oss-120b TPU-native with architecture surgery

Intro

We rearranged gpt-oss-120b's attention architecture to fit TPU silicon, uptrained it on 3B tokens, and got 2.9x input (prefill) throughput at 128k context - matching H100s on a chip that previously ran at a third of that speed. Here's why open-weight models are slow on TPUs, and how we fixed it.

TPUs are suddenly everywhere: Google's accelerators have become the most credible challenge to NVIDIA's dominance, with Anthropic, Meta, and others making major TPU commitments. Yet startups and smaller companies barely touch them.

The TPU ecosystem

Google is building a highly optimized open-source stack to serve popular open-weight models on TPU. For the first time, we can run most open-source models on TPU with minimal coding, thanks to the new vLLM TPU backend.

But running models is not the same as running them fast. The open-weight ecosystem moves much faster than the TPU ecosystem, so performance lags far behind GPUs - limiting what we can readily serve to customers. To close this gap, Google introduced TorchTPU, making PyTorch and many of its optimizations run on TPU natively.

TorchTPU gets us roughly 80% of the way to GPU parity. But the most fundamental problem for inference on TPUs is still unsolved: today's open-weight models are optimized to run efficiently only on GPUs.

To understand why the models themselves are the bottleneck, you have to look at the silicon.

GPUs vs TPUs

GPUs are general-purpose parallel processors. They run thousands of small, flexible cores, and they handle irregular work well: odd matrix sizes, scattered memory reads, dynamic shapes. Custom kernels can smooth over almost any model quirk.

TPUs are built for one thing: AI math. Each chip centers on a huge matrix unit that multiplies large, fixed-size blocks of numbers. When work arrives as big, aligned, predictable matrix multiplies, TPUs are extremely fast and cost-efficient.

Historically, TPUs were designed for two workloads: training, and inference on the embedding-heavy ranking models behind Search and Ads. Both are FLOPs problems (enormous, regular matmuls), so the chips skewed heavy on compute and light on memory. That's exactly backwards for LLM serving, where weights and KV cache strain memory capacity and decode is bound by bandwidth, not math.

TPU v6e (fka Trillium) is the turning point. Doubled HBM (32 GB per chip) and higher bandwidth shifted the optimal inference workload from embedding models to long-prefill, short-decode LLMs: the matrix unit, now 256×256 wide, chews through long inputs, and an 8-chip host finally fits a mid-size model and its KV cache. The sweet spot: long inputs, fast and cheap.

But the bigger matrix unit cuts both ways. A 256-wide array is only fast if your matmuls are 256-wide; feed it anything smaller and most of the silicon idles.

The catch: open-weight models are built and tuned for GPUs. Design choices like small attention heads are free on a GPU, but leave most of a TPU's matrix unit multiplying zeros.

gpt-oss-120b turned out to be a textbook case.

Our gpt-oss-120b, and our solution

We serve models on Trillium, which is optimized for long-input inference, particularly for small and medium LLMs. This is why we picked gpt-oss-120b as the first model we serve at scale.

On long-input benchmarks, we found that it underperformed due to an architecture largely incompatible with TPU silicon. gpt-oss uses 64-dimensional attention heads, while Trillium's MXU is a 256x256 systolic array, so the core attention matmuls (QK^T and attention-weighted V) could never occupy more than 25% of the MXU's lanes. The other 75% multiplied padded zeros.

To fix this, we performed targeted "architecture surgery" on the model. The original gpt-oss-120b has 64 attention heads, each 64-wide. We modified the model to use 16 256-wide attention heads, then uptrained it on a 3B-token synthetic dataset to calibrate the model.

On long-input benchmarks, this doubled-to-tripled prefill throughput in many cases:

With the original model weights, we found that Trillium performed at parity at short contexts below 32k tokens. The surgery gains concentrate at long context, where attention dominates prefill compute, exactly the workload we serve. In absolute terms, that's enough to match H100s at 32–64k and pull ahead at 128k:

Speed is worthless if the model got worse, so we checked. To verify the surgery preserved model quality, we evaluated the checkpoint on math, science, health, general-knowledge, and long-context benchmarks. It achieves quality close to the baseline model (scores averaged over 5 runs):

Quality is largely preserved across the board, including on tool-use (AIME with tools) and long-context (AA-LCR) tasks, with small deltas that are comparable to run-to-run variance on most benchmarks.

The big picture

Today's open-weight models are GPU-native by accident, not by design. With vLLM on TPU and TorchTPU, the models are now the bottleneck. Our gpt-oss-120b results show that this last gap can be closed with targeted architecture surgery and a small uptraining budget: H100-level prefill throughput per chip, on a chip that costs a fraction of an H100 per hour, and draws less than half the power.

The surgery also pays beyond serving on Trillium. 256-wide heads feed the matrix unit just as well during training. And they map exactly onto TPU7x's 256×256 MXU, so the model is native to the next generation of silicon, not just the current one.

We're applying the same pipeline to more models, and in the future we'll offer managed deployments for enterprises with custom models.

Longer term, we hope this work makes itself obsolete: the best outcome is a model ecosystem that isn't quietly designed for a single vendor's silicon.