TorchTPU makes TPU a native PyTorch device; Transformers, Diffusers and Kernels run by changing the device name
No more torch_xla or learning JAX: Hugging Face models run with to("tpu").
// Key points
- TorchTPU is a native PyTorch backend for TPU, with TPU-tuned Pallas kernels, distributed training support and a torch.compile integration.
- In Transformers, passing device_map="tpu" loads a model directly onto the TPU; eager mode is the default, with torch.compile for more speed.
- Diffusers treats TPU as a first-class device, and models too big for one chip can be sharded across chips at load time with a TensorParallelConfig.
- With native support, Optimum TPU is no longer required to run Transformers and Diffusers.
Builder's takeRead with yesterday's TorchNeuron, PyTorch has taken both Trainium and TPU this week. For product teams, inference pricing is no longer GPU quotes alone; my advice is still to move one running inference service over unchanged and compare unit cost, rather than rewrite code for new hardware.
// Background · from #Open models
Full timeline →- Oct 21 PyTorch 2.15 on Mac: MoE models decode 5–6.5x faster, Qwen3.6-35B-A3B rising from 9.4 to 61.3 tokens/s on an M3 Ultra
- Oct 20 AWS contributes bi-directional KV transfer to vLLM, cutting time to first token in multi-turn agent sessions by up to 3.3x
- Oct 20 Hugging Face and AWS launch TorchNeuron: Transformers and friends run directly on Trainium, with 89% of the top 100 text-generation architectures working out of the box, up from 46%