---
title: "PyTorch"
description: "6 posts about PyTorch, summarised, each linking to the original."
---

# PyTorch
> 6 posts about PyTorch, summarised, each linking to the original.

## Articles

### [Profiling TorchNeuron (Part 1): System-level profiling on AWS Trainium.](https://yomu.fyi/post/profiling-torchneuron-part-1-system-level-profiling-on-aws-trainium.md)
- Company: [Hugging Face](https://yomu.fyi/company/hugging-face.md)
- Author: Josh Longenecker
- Published: Oct 21, 2026
- Type: Tutorial

The post introduces system-level profiling of inference performance on AWS Trainium using TorchNeuron. It explains how to configure torch.profiler with the Neuron backend by passing a NeuronConfig and using NeuronProfiler to capture traces. The walkthrough profiles three implementations of RMSNorm: naive PyTorch operations, the fused torch.rms\_norm, and naive operations compiled with torch.compile. The key\_averages() table shows that the naive implementation spends 2,355 µs on the Neuron device across six operator dispatches, while both fused and compiled versions reduce device time to about 400 µs with a single dispatch. Neuron Explorer visually confirms the lifecycle of each operation from framework dispatch through runtime scheduling to hardware execution, revealing the gaps that fusion eliminates.


### [Announcing New vLLM and SGLang Inference Engines on TPU](https://yomu.fyi/post/announcing-new-vllm-and-sglang-inference-engines-on-tpu.md)
- Company: [Google](https://yomu.fyi/company/google.md)
- Author: Qi Zhou, Chris Chan
- Published: Oct 20, 2026
- Type: Announcement

Google announced new PyTorch-native backends for vLLM and SGLang inference engines on TPU, co-developed with Inferact and RadixArk. Built on TorchTPU, the engines make TPU a native PyTorch device with standard tensor management, eager execution, and torch.compile via XLA, using a platform-plugin design that reuses each framework’s worker model and scheduler. They support production workloads including agentic multi-turn, interactive low-latency, high-throughput batch, long context, disaggregated serving, and multimodal models, with custom Pallas kernels called as PyTorch custom operators. Features include automatic prefix caching, speculative decoding with MTP, EAGLE-3, DFlash, DSpark and n-gram drafters, KV cache offloading to host memory and Managed Lustre, and TPU Sync for efficient cache transfer between hosts. Both engines are running in production at Meta and preview customers, with public benchmarks showing industry-leading cost per million tokens at 2.01 cents on Ironwood TPU, and a community contribution TPU capacity pool is available for developers.


### [Integrating TorchTPU into TorchTitan and Optimizing Models for TPU](https://yomu.fyi/post/integrating-torchtpu-into-torchtitan-and-optimizing-models-for-tpu.md)
- Company: [Google](https://yomu.fyi/company/google.md)
- Author: Aleksey Vlasenko, Rob Torres, Will Constable
- Published: Oct 20, 2026
- Type: Announcement

TorchTitan TPU is a fork of Meta’s TorchTitan that enables PyTorch model training on Google Cloud TPUs using the TorchTPU native backend with minimal device initialization changes. SimpleFSDP replaces FSDP2 hooks with DTensor parametrization, making all-gathers and reduce-scatters ordinary nodes in the traced graph for compiler scheduling. The Graph Trainer captures the full forward, loss, and backward step as one XLA program via make\_fx, achieving 51-52% MFU on Qwen3-8B with functional AdamW. On Llama 3 405B, TorchTitan TPU reaches about 57% MFU on up to 1,024 chips, and a weak-scaling test of Qwen3 235B MoE retained 92% of per-chip throughput from 64 to 8,192 chips. The repository includes an optimization playbook, a coding-agent model-onboarding skill, and supports over 500 Hugging Face models through a Transformers backend.


### [Hugging Face ❤️ TorchTPU: TPUs, Natively](https://yomu.fyi/post/hugging-face-heart-torchtpu-tpus-natively.md)
- Company: [Hugging Face](https://yomu.fyi/company/hugging-face.md)
- Author: Alvaro Moran, Jingya Huang, Simon Pagezy, Jeff Boudier, Régis Pierrard
- Published: Oct 20, 2026
- Type: Announcement

TorchTPU adds a native PyTorch backend for Google TPUs, eliminating the need for torch\_xla or JAX. Installing it on a TPU VM requires a single pip command with a CPU-only torch index. Users can load any Hugging Face model onto a TPU with device\_map=“tpu” and run it in eager mode without code changes. Diffusers pipelines also support TPU natively, requiring only pipe.to(“tpu”) for eager execution or an optional compile call for deployment. The Kernels Hub provides ready-to-use TPU kernels written in Pallas, which can accelerate MoE models by over 20× at large batch sizes, as demonstrated with Qwen3-30B-A3B. The Hugging Face team is continuously testing TPU support, with 96.5% of tests passing and upstream fixes for discovered issues, including bugs in StaticCache and Qwen3 tensor-parallel plans.


### [MPS in PyTorch 2.15](https://yomu.fyi/post/mps-in-pytorch-2-15.md)
- Company: [Hugging Face](https://yomu.fyi/company/hugging-face.md)
- Author: Irakli Salia, Khushi Agrawal
- Published: Oct 20, 2026
- Type: Announcement

PyTorch 2.15 brings FP8 support to MPS, storing float8\_e4m3fn values in half the space of BF16 and enabling scaled matrix multiplication directly on the GPU with tensorwise or rowwise scaling. Language-model decoding accelerates significantly: mixture-of-experts models decode 5–6.5× faster thanks to a native grouped matmul with a dedicated decoding path—on an M3 Ultra with torch.compile, Qwen3.6-35B-A3B reaches 61.3 tokens/s, up from 9.4. Dense models benefit from specialized matrix-vector kernels for single-row linear layers, and softmax over the vocabulary is about 8× faster. Video generation memory usage drops dramatically: Wan2.2’s VAE decode now peaks at 8 GiB instead of 270 GiB. The Hugging Face team runs real models to find and fix MPS shortcomings. The release also adds complex LU and triangular solves, faster small-matrix inverses and QR decomposition, public MPS streams, and requires macOS 15 or later, with many operations moved from MPSGraph to dedicated Metal kernels.


### [TorchNeuron: running HF models on AWS Trainium is easier than ever](https://yomu.fyi/post/torchneuron-running-hf-models-on-aws-trainium-is-easier-than-ever.md)
- Company: [Hugging Face](https://yomu.fyi/company/hugging-face.md)
- Author: Régis Pierrard, Simon Pagezy, Michael Benayoun, Jingya Huang, Alvaro Moran, David Corvoysier, Jeff Boudier
- Published: Oct 19, 2026
- Type: Announcement

TorchNeuron is a new PyTorch plugin that adds a native Neuron backend for AWS Trainium, supporting eager mode, torch.compile, and distributed APIs like FSDP and TP. It will be available as a public beta starting October 20, 2026, and can be installed with pip install torch-neuronx. Hugging Face libraries such as Transformers, Diffusers, Accelerate, PEFT, and TRL now run directly on Trainium without needing the Optimum Neuron wrapper, using the same classes and scripts as on any other accelerator. For text generation, 89% of the top 100 most-used model architectures are supported out of the box, compared to 46% with Optimum Neuron v0.4.5. Training uses standard TRL scripts with only a few Neuron-specific settings like static shapes and torch.compile, and the resulting checkpoints are hardware-agnostic.