Google's TPUs Power Meta's AI Services as Amazon Releases PyTorch Code for Trainium Chips

Major technology companies are releasing code designed to run AI models on their own chips, reducing reliance on Nvidia.

Google's custom AI accelerator chips are now being used for AI workloads in Meta's data centers. This follows the release of a PyTorch-native version of the SGLang inference engine. A corresponding version of the vLLM inference engine was also made available as open source.

Previously, both engines required Google's own JAX AI framework to operate on TPUs. The new software versions enable the hardware to function as a native option for PyTorch, mostly without rewriting code. "SGLang-TorchTPU is running in production at Meta today," said Meta's PyTorch-TPU lead Zongwei Zhou, following a collaboration between Meta, Google, and the SGLang team. That team works at RadixArk, a startup founded by SGLang's creators, which built the engine with Google and Meta. The engine, he said, provides Meta "a remarkably mature, PyTorch-native serving stack."

Jana van Greunen, director of PyTorch engineering at Meta, stated in a Google blog post that Meta "cannot be constrained by a single hardware architecture" and that the new software "enables us to deploy our PyTorch workloads to Google TPUs." Google indicated it developed the software jointly with Meta's PyTorch team. This is not Meta's first attempt to utilize TPUs. Google identifies PyTorch/XLA, the older bridge from PyTorch to TPUs, as a "predecessor" to the new software, and SemiAnalysis reported last year that Meta research teams used TPUs through this bridge from 2020 to 2023 before Meta canceled those contracts. Google's PyTorch-on-TPU push was aimed mainly at Meta, which did not want to migrate to JAX.

It remains uncertain where Meta's production TPUs are located. Discussions last year reportedly involved Google Cloud TPU rentals in 2026 with purchases in 2027, and Meta later reportedly signed a multi-year rental agreement.

This announcement came one day after Amazon said it was open-sourcing TorchNeuron, its own native PyTorch backend for its Trainium2 and Trainium3 chips, at the same PyTorch Conference. Both actions appear aimed at loosening Nvidia's software lock-in.

Regarding performance, a tuned vLLM setup serving Moonshot AI's open-weight, 2.8-trillion-parameter Kimi K3 model reportedly exceeds 700 tokens per second for a single user on the seventh-generation Ironwood TPU. This speed figure comes from Inferact, the startup founded by vLLM's creators that helped Google build vLLM-TorchTPU. Its test ran 16 Ironwood chips with speculative decoding, a method of speeding up large language model inference with a small draft model, against vLLM's published setup on 16 Nvidia GB200s, which reached 452 tokens per second. Its kernels for Kimi K3 and Qwen 3.8 27B deliver roughly 1.4 to 2 times the GB200 baseline's throughput at small batch sizes without speculative decoding.

SGLang-TorchTPU will keep adding models, starting with DeepSeek V4.1 Flash and Reflection Beam. It's also working toward Day 0 TPU support for new open models, meaning support on release day, by the first quarter of 2027, according to RadixArk. Google, Meta and RadixArk plan to share more at SGLang Summit 2026, Nov. 12 to 13 in San Francisco.