Google TPUs now running PyTorch workloads in Meta data centers
Meta is using Google's custom AI chips via a new PyTorch-native SGLang inference engine, reducing reliance on JAX and Nvidia.
Summary
Google's Tensor Processing Units (TPUs) are now being used in production at Meta for AI workloads. This was made possible through a collaboration between Meta, Google, and RadixArk to create a PyTorch-native version of the SGLang inference engine. Additionally, an open-source version of vLLM for TPUs has been announced, allowing PyTorch users to run workloads on Google hardware without needing to migrate to the JAX framework.
Technical shift from JAX to PyTorch
- New SGLang-TorchTPU engine allows Meta to run PyTorch workloads natively on Google hardware.
- An open-source version of vLLM for TPUs has been released.
- Previously, TPU support for these engines required Google's JAX framework.
- Performance tests on seventh-generation Ironwood TPUs showed the vLLM setup reaching 700 tokens per second for the Kimi K3 model using speculative decoding.
- Without speculative decoding, kernels for certain models delivered 1.4 to 2 times the throughput of an Nvidia GB200 baseline at small batch sizes.
Why it matters
This move, alongside Amazon's open-sourcing of TorchNeuron, aims to reduce the software lock-in provided by Nvidia's CUDA ecosystem by making alternative AI accelerators more accessible to PyTorch developers.
What to do
- PyTorch users can use the new SGLang-TorchTPU engine to run workloads on Google TPUs without migrating to JAX.
- PyTorch users can access the newly released open-source version of vLLM for TPUs.
Original reference: www.tomshardware.com →