Google and AWS make PyTorch native on their chips and loosen CUDA’s grip
On October 22, 2026, at PyTorch Conference North America, Google open-sourced TorchTPU for its TPUs and AWS launched TorchNeuron in public beta for its Trainium chips, two native PyTorch backends that avoid JAX or CUDA dependencies. Getting started is easy, but peak performance still needs chip-specific tuning, with a 30 to 40% penalty for TorchTPU’s default approach.
October 22, 2026. At PyTorch Conference North America 2026, Google open-sources TorchTPU, a native PyTorch stack for its TPUs, and AWS launches TorchNeuron in public beta, its native PyTorch backend for Trainium. April 2026. TorchTPU had been in preview with select customers, after work that began in mid-2025. Both releases aim at the same target: loosening the grip of CUDA, the Nvidia software ecosystem that has dominated AI development for years.
For an AI infrastructure lead staring at a GPU bill, the announcement fits in one sentence: there are now two official paths to run PyTorch on non-Nvidia accelerators, but ease of entry does not yet translate into peak performance — and that is exactly what will decide your migration.
Two native backends to break out of CUDA
The reasoning of both hyperscalers is identical, expressed almost word for word. Google wants to make its TPU software ecosystem “more open and developer-friendly” for the PyTorch community; AWS extends the same approach to its own chips. The shared goal is to make TPUs and Trainium practical options for developers already using PyTorch, without forcing a rewrite.
The context shows how much has changed. For years, running PyTorch on a TPU was an “awkward dance”: workloads went through JAX or XLA paths that required adapting the code. On Trainium, you had to understand the Neuron SDK before hoping for a gain. The two new backends short-circuit that friction by putting PyTorch at the center, where most teams already live.
TorchTPU: PyTorch on TPUs without JAX
TorchTPU, developed in collaboration with Meta’s PyTorch team, runs existing PyTorch workloads with no JAX dependency at runtime. Rather than requiring extensive code changes, migration boils down, per PyTorch, to “trivial code changes.”
Two execution modes cover a model’s lifecycle. Eager execution is for debugging, running operations as they come, the way you would locally. torch.compile then optimizes TPU utilization and reduces memory overhead. The promise is a continuum: debug simply, then compile for production.
The release ships with a train of integrations announced at the conference — Hugging Face, vLLM, SGLang, llm-d, TorchTitan, Ray, and TorchRec — plus a contribution of LiteRT to the PyTorch ecosystem. Google frames them as pieces “designed to help you run PyTorch effectively on Google Cloud or in your own infrastructure,” a signal that the target goes beyond its own cloud.
TorchNeuron: PyTorch on Trainium in beta
AWS plays the same card with TorchNeuron, a native PyTorch backend for Trainium, its family of AI chips built for training and inference. The public beta supports Trainium2 and Trainium3 instances.
The advertised simplicity fits in one line of code: moving a model to Trainium can be as simple as switching the device to neuron via .to(), for eager execution. Compiled models must select the Neuron backend, either as the default or in a torch.compile call. TorchNeuron also integrates with the Neuron Kernel Interface (NKI), the bare-metal language and compiler for programming Trainium and Inferentia devices, to build custom operators and optimize performance.
AWS and PyTorch have published detailed guidance on the backend’s execution modes and their trade-offs — an admission that the “trivial” path and the “fast” path are not the same.
Ease of entry is not peak performance
This is the point neither announcement puts in one line. PyTorch acknowledges that a simple code change gets models running on TPUs, but peak performance still requires TPU-specific optimization.
The most telling number concerns distributed workloads: TorchTPU’s default Skinny Collectives approach can impose an end-to-end penalty of roughly 30 to 40% versus an optimized approach, because it prevents the XLA compiler from overlapping computation and communication. For those targeting peak performance, PyTorch recommends extra tuning and production-grade stacks — “Pallas kernels and stacks such as vLLM-TorchTPU, SGLang-TorchTPU, or TorchTitan-TorchTPU.”
AWS is more cautious in its wording: teams can try Trainium with a configuration change rather than a code or model rewrite “in most cases” — which leaves open what the remaining cases are. A case study shows why: to adapt i14’s Laplace 3B model to Trainium2, AWS’s Generative AI Innovation Center had to make three changes — use native PyTorch for the GLA attention kernel, redesign how Muon communicates, and compile the sequential backward pass.
Verdict
If your fleet runs heavily on CUDA and your GPU bill is exploding, open an evaluation track: TorchTPU and TorchNeuron lower the entry cost to alternative accelerators, and a first “trivial” test over a few days will tell you whether your models port without a rewrite.
If you are chasing peak performance on a critical model, do not take the “trivial” promise at face value: budget for per-chip optimization — Pallas or NKI kernels, collective tuning — and measure the 30 to 40% penalty of the default approach before comparing prices. The switching cost is no longer in the code but in the optimization; that is what you must price before signing.