How PyTorch Became the Socket Every AI Chip Plugs Into

On 19 October, AWS released TorchNeuron, which lets PyTorch run directly on its Trainium chips. A day later, Google and Meta released TorchTPU, which does the same for Google's TPUs. Both arrived at the PyTorch Conference, both are open source, and both connect to PyTorch through the same small hook.

That last part is the story. Two companies with no reason to coordinate reached the same conclusion: the road onto a working engineer's laptop runs through one extension point in PyTorch.

Write a + b on two tensors and PyTorch does not immediately add anything. It looks up a name, aten::add, in a table.

The part doing that lookup is called the dispatcher. It checks where the tensors live, reads a label off them — CPU, CUDA, and so on — and uses it to pick the function that actually does the addition on that hardware. Those hardware-specific functions are called kernels. A backend is a full set of kernels filed under one label.

PyTorch keeps one label deliberately empty: PrivateUse1. Any hardware vendor can claim it and register their kernels against it from their own repository, without needing their code accepted into PyTorch itself. TorchTPU registers there. TorchNeuron uses the same registration mechanism. A new chip becomes a new label, and model.to("tpu") is what that looks like from your code.

A backend is a set of kernels filed under one dispatch label. TorchTPU names PrivateUse1; AWS says only "PyTorch's built-in mechanism for registering new hardware". Sources: PyTorch blog; Amazon; InfoQ.

This replaces something clumsier. PyTorch on TPUs used to go through JAX, Google's other framework. The TorchTPU team is blunt: the old approach "shoehorn[ed] PyTorch into JAX's modus operandi." The new backend has no JAX in it at runtime, and it runs operations one at a time as you write them, which is how most people debug. The chip speaks PyTorch now, instead of being translated first.

The coverage looks real. Google tested 4,917 models from Hugging Face: 98% ran without crashing, and of those it checked closely, 97% produced the same numbers a GPU would. Amazon says it tests TorchNeuron against hundreds of thousands of models on every build.

Here is the TorchTPU team's own framing, caveat included:

"Porting a workload to TPU should be as simple as changing model.to("cuda") to model.to("tpu"). (One cannot achieve peak TPU performance with that simple change... however, the model should run out of the box.)"

Read the bit in brackets, because that is where the work actually is.

Illustrative: the bar shows the direction of each documented caveat, not a measured benchmark. Source: PyTorch blog, TorchTPU.

So the part that looked hard, getting the code to run, is nearly solved. The part nobody puts in a launch announcement is that your tuned kernels, your checked numbers and your network settings do not move with.to(). That is months of work per platform, and it is where lock-in actually lives now.

Most of TorchTPU's performance numbers come from Google, and most of its examples do not compare against a GPU on equal footing. The 64-chip test against JAX/MaxText (1,828 against 1,937 tokens per second) is one of the few straight comparisons, and it is against JAX, not a GPU. The project calls itself "in active development." AWS does not mention CUDA at all. Putting these two releases in one frame is my argument, not either company's claim.

The moat was never CUDA the language. It was the thousands of hand-tuned kernels underneath it, and PrivateUse1 is the agreed-on door for building a rival set.

If switching chips costs one line of code and three months of tuning, is that portable?