pytorch.org – p1

FP8 Training on AMD GPUs with TorchTitan and TorchAO: Upstreaming Performance Improvements

At the PyTorch Conference 2025, we demonstrated linear scaling beyond 1,000 GPUs on AMD Instinct clusters using Primus-Turbo, an AMD optimization library for training frameworks such as TorchTitan. We have since upstreamed those AMD optimizations so TorchTitan supports AMD Instinct() GPUs directly, with competitive FP8 performance out of the box. All contributions mentioned have been merged into upstream pytorch/AO and pytorch/TorchTitan.

On dense models, FP8 training delivers a 13.4% throughput gain over BF16 on Llama3-8B (#2736) as shown in Figure 1. On MOE architectures like DeepSeek-V3 671B, FP8 quantization initially added significant overhead. Through fused Triton quantization kernels, we recovered 89% of the FP8 quantization overhead on DeepSeek-V3 671B MoE shapes (#4311), with individual kernel optimizations delivering up to a 6.2× speedup (#4113). 

This blog covers the FP8 optimizations that deliver these gains, from major kernel acceleration to the Triton fusion pipeline that narrowed the quantization overhead gap on MoE models. Getting there took three pieces of work: 

1. Adding native support for AMD’s FP8 number format 
2. Enabling grouped GEMM for Mixture-of-Experts models on ROCm
3. Building a Triton fusion pipeline that reduced the quantization overhead

Figure 1: Rowwise FP8 training throughput on 8×MI300X GPU with Llama3-8B (batch size 1, seq len 8192, 100 steps, torch.compile, FSDP2, per-op selective activation checkpointing). Rowwise FP8 with a high-precision weight-gradient recipe (the weight-update GEMM stays in BF16 while the forward and gradient-input GEMMs use FP8) delivers a 13.4% throughput gain over BF16, with peak memory nearly identical (~39 GB). The win comes from faster FP8 matrix cores, not memory savings. All numbers come from TorchAO PR #2736.

AMD FP8 format in TorchAO

Each linear layer performs three matrix multiplications: the forward pass, gradient input, and gradient weight update. FP8 training quantizes these operations from 16 bits to 8 bits, greatly improving throughput. AMD Instinct GPUs implement a variant of the FP8 formats called FNUZ (Finite, No NaN, Unsigned Zero) which is illustrated in the table below

The FP8 capabilities demonstrated in Primus-Turbo were upstreamed directly to TorchAO and TorchTitan as displayed in Figure 2, spanning three areas: hardware-aware FP8 format support, MoE grouped GEMM enablement on ROCm, a Triton kernel fusion pipeline that reduced quantization overhead.

Figure 2: The TorchTitan FP8 training software stack on ROCm. AMD’s upstream contributions span TorchAO (FP8 dtype support, Triton kernel optimizations) and TorchTitan (MFU fixes, loss baselines, scaling recipes)

The TorchAO library initially did not support the same numerical format as AMD so it was computing scales against a different max value. On AMD Instinct GPU, where e4m3fnuz has a max of 240, this produced silently wrong results: tensors were scaled into a range that exceeded the hardware’s representable values, clipping activations and corrupting gradients. Because e4m3fnuz has no NaN/Inf encodings, the overflow did not raise an error; it degraded model quality instead. Selecting the correct format is therefore a correctness requirement, not a tuning option. We added hardware auto-detection so TorchAO selects the correct format automatically. Getting there took a cluster of format-correctness fixes across TorchAO and TorchTitan:

• Auto-detect the platform and select the correct FP8 dtype and max value, instead of hardcoding NVIDIA e4m3fn: TorchAO #1142, #1150, #2225
• Report correct MI300X peak FLOPS so MFU numbers are accurate : TorchTitan #920
• Add platform-specific loss baselines for FNUZ numerics : TorchTitan #2156

FP8 scaling can be applied at different granularities: a single scale per tensor (tensorwise, fastest but coarsest), a scale per row (rowwise, better accuracy), per fixed-size tile (blockwise), or per group packed alongside the data (MXFP8). TorchAO and TorchTitan support all four strategies. For AMD GPUs, we ensured each quantization strategy works correctly with AMD specific numerics and contributed blockwise kernel support for MI300 and MI350 GPUs (#3996).

Scaling FP8 to MoE Architectures

Mixture-of-Experts (MoE) models like DeepSeek V3 and Llama 4 route each token to a subset of experts, producing variable-size batches that must be processed through a grouped GEMM (Figure 3). Unlike dense models, where every linear layer has the same shape, grouped GEMM requires per-row scales on activations, per-expert-column scales on weights, and an offset tensor routing rows to the correct expert.

We enabled FP8 grouped GEMM on ROCm by adapting the quantization pipeline to use the correct dtype and dispatch for AMD via the Composable Kernel backend (#3955).

Figure 3: MoE FP8 grouped GEMM pipeline on ROCm. (A) Tokens are routed to experts via offsets, quantized by fused Triton kernels, and dispatched through Composable Kernel in a single launch. (B) Grouped GEMM requires per-row, per-expert-column scales and offset-based routing, which is more complex than dense GEMM uniform scaling.

Triton Kernel Optimization

With correctness established, we turned to performance. The FP8 quantization pipeline in TorchAO converts tensors to FP8 through a multi-step chain:

1. Compute per-row/column absolute max (absmax)
2. Derive the scale factor and apply it
3. Clamp and cast to FP8

Each step is a separate kernel launch, and materializes an intermediate tensor to High Bandwidth Memory (HBM) between steps. For MoE models with dozens of expert weight tensors per layer, these extra round-trips dominate the FP8 overhead. On these shapes, FP8 quantization is memory-bound: the 8-bit math is cheap, but the kernel launches and HBM round-trips around it are not. The optimizations below reduce data movement rather than arithmetic, at three levels of granularity: launching fewer kernels (Level 1), making each remaining kernel move memory efficiently (Level 2), and removing unnecessary low-level synchronization (Level 3).

Level 1: Launch fewer kernels:

Backward pass:

 Figure 4 illustrates how fp8 quantization was improved through fusion. The backward pass had two compounding problems. First, a .t().contiguous().t() pattern forced a full tensor copy through HBM to convert weight layout for GEMM compatibility. We removed these redundant copies in #3972. Second, the multi-step scale-and-cast chain launched separate kernels with intermediate tensors materialized to HBM between them. We fused this chain into single Triton kernels in multiple places (#4069). On 8xMI300X GPUs with DeepSeek-MoE-16B, these backward fusions delivered a 4.2x backward pass throughput improvement.

Figure 4: Backward-pass FP8 quantization before and after optimization. The upstream code launches five generic kernels per quantization call with a redundant transpose copy, materializing intermediate tensors to HBM between each. PR #3972 eliminates the transpose, and PR #4069 fuses the remaining chain into a single kernel with a companion dual-kernel for simultaneous grad_output + activation quantization.

Forward pass:

The same multi-kernel pattern applied to the forward path. Quantizing expert weights launched five generic kernels per call, and with 24 calls per step, this added ~90 ms/step of overhead. We replaced the entire chain with a single fused Triton kernel (#4311) that parallelizes across both experts and output-dimension blocks (Figure 5). Collapsing five launches into one also let the surrounding GEMMs issue sooner. On 8x MI325X GPU with DeepSeek-V3 671B: this change delivered a 17% end-to-end throughput improvement (5,996 → 7,027 tok/s).   

FP8 forward before optimization
Fused FP8 forward

Figure 5: Perfetto trace comparison (8xMI325X GPU, DeepSeek-V3 671B). Left: FP8 (before forward optimization) showing the 5-kernel eager chain repeated across experts. Right: fused FP8 (after optimization) showing a single triton_fp8_colwise_3d_scale_and_cast kernel replacing the chain. Performance in forward goes from ~19 ms to ~7 ms

Figure 6: Per-category GPU time breakdown across three configurations on 8xMI325X GPU with DeepSeek-V3 671B (4-layer MoE). The FP8 upstream configuration (V2) adds 127 ms/step, 92% of which lands in “Others” (generic quantization kernels). The fused Triton kernel (V4) eliminates most of this overhead, recovering 89% of the BF16→FP8 gap.

Level 2: Make each kernel move memory efficiently: The colwise scales kernel used in the backward pass had non-coalesced memory writes (#4113): consecutive SIMD lanes wrote to addresses K bytes apart, each triggering a separate memory transaction. We fixed this by transposing the output tile through LDS (Local Data Share) before storing, and added a fused single-pass variant that eliminates a redundant HBM read. On an MI300X GPU with DeepSeek-V3 671B shapes: 7,290μs → 1,170μs per MoE layer (6.2x speedup).

Level 3: Strip synchronization the hardware never needed: We also addressed hardware-level inefficiencies. Triton’s atomic operations (atomic_add, atomic_max, atomic_min) default to acquire-release memory ordering, which on AMD GPU inserts memory fences before and after every atomic, which are expensive synchronization points that are unnecessary for commutative reductions. We switched these to relaxed ordering on AMD GPU (#3945), guarded by a torch.version.hip check so NVIDIA behavior is unchanged.

What didn’t work: autotuning the search space. We expanded the Triton autotune search space for MoE FP8 kernels from 1 to 8–16 candidate configurations (#3952), expecting the wider search to find faster tile sizes on AMD wavefront-based architecture. However,  benchmarking on Llama 4 shapes on MI300X GPU showed no measurable improvement, and the extra configs increased first-iteration compile time. We reverted it (#4024). The takeaway: autotuning search spaces should be shaped by hardware constraints (wavefront size, LDS capacity, register pressure), not expanded to more candidates by default.

Attacking data movement at all three levels compounds on top of the baseline FP8 throughput gains shown in Figure 1. On DeepSeek-V3 671B shapes, the forward-pass kernel fusion alone recovered 89% of the quantization overhead (5,996 → 7,027 tok/s vs 7,156 BF16 baseline on 8xMI325X GPU), and the colwise scales optimization delivered a 6.2× speedup per MoE layer.

Summary and Next Steps

This blog described how we optimized FP8 training on AMD Instinct GPUs across TorchAO and TorchTitan: kernel speedups, numerical stability fixes, and support for MoE architectures. 

Work continues on next-generation hardware. We are developing MXFP8 grouped GEMM and quantization kernels for forward and backward passes on MI355X GPUs; results will follow in a future blog. The kernel fusion pipeline (#3972 → #4069 → #4113 → #4311) continues with further Triton optimizations.

These FP8 gains are now available in the standard PyTorch training stack: teams with AMD Instinct GPUs get them by upgrading TorchAO and TorchTitan, with nothing AMD-specific to install. This work was a collaboration between AMD and Meta/PyTorch engineers. All contributions have been merged into mainline pytorch/ao and pytorch/torchtitan, ensuring that FP8 training on AMD GPUs works out of the box for the broader PyTorch community.

Additional Resources

• TorchAO float8 README, including benchmark reproduction instructions
• TorchTitan
• AMD ROCm Documentation: Learn more about the ROCm software stack
• AMD AI Developer Program: Access free cloud GPU resources and developer tools to get started with ROCm and PyTorch on AMD Instinct GPUs

Fast, On Device Agentic AI with Muse Glimmer on ExecuTorch

Today, Meta introduced Muse Glimmer, an open-weight, 30-billion-parameter model distilled from Meta’s Muse Spark for on-device agentic workflows. Alongside, ExecuTorch is adding end-to-end support for running Muse Glimmer on NVIDIA GPUs and Macs with Apple silicon.

Why ExecuTorch?

Most local AI frameworks rewrite models in other non-Python languages. That scaled well when LLMs were standard text transformers, but today’s models are becoming more complex – novel architectures, multimodal inputs and outputs, advanced decoding algorithms like DFlash (parallel diffusion-based speculative decoding) for low latency. Reimplementing these across different backends doesn’t scale.

ExecuTorch takes a different approach. As machine learning engineers and researchers, you implement the model (and its decoding strategy) in PyTorch. Once you’re ready for deployment, you export to ExecuTorch, and the framework handles backend-specific lowering, Triton on CUDA, MLX-native and custom Metal on Apple silicon. Ahead-of-time compilation optimizes the full execution path end-to-end, not just individual ops.

This is how we ship Muse Glimmer’s text and image inputs, direct GGUF export, native K-quant execution, 128K+-token context, and DFlash speculative decoding features. We have released prebuilt PTE artifact bundles that you can download and run on supported NVIDIA GPUs or Macs with Apple silicon using the ExecuTorch runtime.

Quickstart

Getting the PTEs

A PTE is the serialized artifact produced ahead of time from a model’s PyTorch graph by the ExecuTorch Python stack, and optimized for a target backend.

Download (Preferred)

We have published verified PTEs on Hugging Face for NVIDIA CUDA and Apple Silicon (Metal). This includes text-only and text-plus-image artifacts, with and without DFlash speculative decoding. Download them here: link.

Build your own

Starting with a prebuilt PTE is the fastest way to get running. To build your own, follow the ExecuTorch Muse Glimmer README, and select the backend, modality, context length, and whether to use DFlash. ExecuTorch exports directly from the released GGUF checkpoints through its torch.export-based ahead-of-time stack. CUDA export compiles and autotunes Triton kernels for the detected GPU architecture. For the best results, export on the same GPU architecture that will run the artifact.

Executing the PTEs

1. Build the runtime

ExecuTorch ships CMake presets for both the CUDA and MLX backends for this model runner(s). Follow ExecuTorch installation instructions here, and then use CMake to build the runners with or without speculative decoding for the PTE you selected. Both, with and without DFlash, runners support text and image modalities and are compatible with the example llm_server in ExecuTorch for agentic use cases.

### Build the runtime ###

# After installing ExecuTorch, build the runners for your backend:

$ cd examples/models/muse-glimmer
$ cmake --workflow --preset muse-glimmer-cuda # use muse-glimmer-mlx for macOS

# This builds solo_runner, dflash_runner, and the serving worker.

2. Run the PTEs

Here are some examples of how to run the PTEs, once you have built the runners.

### Example 1: Standalone infereance on cmdline ###

$ PROMPT='<|start|>user<|message|>Describe this image: <img><|eot|><|start|>assistant'

$ cmake-out/examples/models/muse-glimmer/dflash_runner \
--model_path artifacts/dflash-vision/model.pte \
--data_path artifacts/dflash-vision/aoti_cuda_blob.ptd \
--tokenizer_path assets/hf/tokenizer.json \
--image_path image.jpg --prompt "$PROMPT" \
--block_length 4 --n_draft 3 --temperature 0 --max_new_tokens 256

### Example 2, step 1/2: Start the agent server ###

$ python -m executorch.examples.models.muse_glimmer.serving.serve \
--model-path artifacts/dflash-vision/model.pte \
--data-path artifacts/dflash-vision/aoti_cuda_blob.ptd \ # only for cuda
--tokenizer-path assets/hf/tokenizer.json --hf-tokenizer assets/hf \
--worker-bin cmake-out/examples/models/muse-glimmer/muse_glimmer_worker \
--tool-parser atem --max-context 131072

# APIs at http://127.0.0.1:8000/v1

### Example 2, step 2/2: Start your agent (use Pi as example) ###
  
$ pi \
--provider muse-glimmer-local \
--model muse-glimmer \
--thinking high \
--tools read,bash,edit,write

# This will automatically start your pi agent by using your local muse glimmer server.
# Register muse_glimmer-local in ~/.pi/agent/models.json first.
# See the README.md for more details.

Use cases enabled

Image understanding with Muse Glimmer, with and without speculative decoding

Figure 1: Muse Glimmer text-image input experiment on M5 Pro (64 GiB). Solo achieves 21.6 tok/s, while our speculative decoding set up (DFlash) reaches 33.0 tok/s, a 52.8% performance improvement without quality regression

Muse Glimmer powering Pi Coding Agent through ExecuTorch

Figure 2: Muse Glimmer agent pipeline on an M5 Pro (64 GB) using the Pi coding agent. The agent creates a bird-themed game, iteratively refining details through extended reasoning, calling tools to create files, installing required packages, writing and running tests, and proactively asking the user about next steps and additional requirements

Performance

Figure 3: Muse Glimmer performance on ExecuTorch using text-only input with varying context on NVIDIA A100 (as a proxy for RTX cards) and Apple Mac with an M5-max measuring prefill and decode performance in tokens/second with and without DFlash using coding prompt, which also has a good acceptance rate for this model, as seen in the decode charts

Under the Hood

Muse Glimmer now runs end-to-end on ExecuTorch on both NVIDIA GPUs and Apple Silicon GPUs. Here are some of the key capabilities and optimizations we built.

Enabling DFlash speculative decoding

• We optimized target and draft interoperability through weight sharing, exporting both into a single PTE.
• The DFlash block dimension is exported dynamically, allowing one PTE to support runtime-selectable block lengths.
• The runtime supports both greedy decoding and rejection sampling.

Supporting GGUF loading and k-quant

• We export straight from the GGUF released with the Muse Glimmer.
• We map Q4_K/Q5_K/Q6_K to packed INT4/5/6 with dp4a GEMV kernels on CUDA, and to repacked or fused Metal kernels on MLX.
• On MLX, for performance, at repack time we merge adjacent sub-blocks whose scale and min are identical into a larger group size, up to 128, whenever the merge is lossless.

Agentic harness and LLM serving

• One model load serves multiple isolated conversations, through per-session mutable-state rebinding we added to both backends.
• We added Harmony chat templating with reasoning routing.
• We added a parser for the model’s XML tool-call format, including multiple calls in one turn.

Backend-specific performance optimizations

CUDA

• We capture decode into a CUDA graph, reducing per-kernel launch overhead into one submission.
• Packed K-quant kernels accelerate low-batch decode, while length-aware split-K FlashDecoding++ paths optimize single-token decode and small DFlash verification blocks.

MLX

• RMSNorm, RoPE, SDPA, KV-cache updates, and quantized linear operations are lowered to MLX-native or custom Metal implementations.
• GGUF K-quant weights use either repacked MLX-native operations or fused Metal kernels.

Supporting long context

Muse Glimmer supports a 128K+ token context, and is efficient in how its KV-cache grows: only 13 of its 52 layers are global; the other 39 are sliding-window. ExecuTorch supports this efficiently, making the long context use cases practical on edge devices.

What’s next

• This initial release supports text and image inputs; video input is not yet supported. It is a work in progress.
• No cross-session prefix sharing or checkpointing or continuous batching as of now. These are all actively being worked on to make ExecuTorch even more suitable for agentic workflows.

Try Muse Glimmer with ExecuTorch and let us know what you think on Discord. If you run into any issues, feel free to open a Github Issue.

References

Muse Glimmer in ExecuTorch | Muse Glimmer on Hugging Face | ExecuTorch Documentation | ExecuTorch on Github

PyTorch Conference North America Announces 2026 Keynotes

PyTorch Conference North America will be held in San Jose, California, on October 20–21, 2026.

Featured PyTorchCon NA 2026 keynote speakers include:

• Mark Collier, Executive Director, PyTorch Foundation
• Mazin Gilbert, Executive Director, Agentic AI Foundation
• Edward Yang, Research Engineer, Meta
• Alban Desmaison, Research Engineer, Meta
• Simon Mo, Co-Founder & CEO, Inferact
• Bill Jia, VP of Engineering, Core ML/AI, Google Cloud
• Brian Stevens, SVP & AI CTO, Red Hat
• Ujval Kapasi, VP, AI & HPC Frameworks & Libraries, NVIDIA
• Joelle Pineau, Chief AI Officer, Cohere
• Mark Saroufim, Co-founder, Core Automation
• Sara Hooker, CEO & Co-Founder, Adaption

The two-day conference will include technical sessions and community collaboration across the PyTorch ecosystem.

Submit a Flare Pin Design by August 14

PyTorch Foundation is accepting original designs for the 2026 PyTorch Foundation flare pin through August 14, 2026, at 11:59 p.m. PT.

The winning entrant will receive one complimentary ticket to PyTorch Conference North America. PyTorch Foundation will produce the winning design as a 1-inch die-cut soft enamel pin for the conference.

Share your design on LinkedIn, X, Facebook, or Bluesky using #PyTorchPin and #PyTorchCon.

Review the contest requirements and submit a design

Become a Sponsor

PyTorch Conference North America brings together 3,000+ members of the open source AI ecosystem.

Companies interested in visibility with the engineers and technical leaders building the next generation of AI infrastructure and applications can review the available sponsorship opportunities.

View sponsorship opportunities

PyTorch by the Sea: The inaugural Santa Cruz PyTorch Meetup

TL;DR

The inaugural Santa Cruz PyTorch Meetup brought together 45 local engineers, students, and leaders for GPU/CUDA talks and lightning presentations on chemistry, plant health, and autonomous driving – demonstrating how easy and impactful it is to launch a low-key, welcoming PyTorch community in your own area.

One of the great parts about being the lead for the Red Hat Open Source and AI Programs Office to the PyTorch Foundation is I get a front seat view to all the exciting activity in the community. One of the most impressive (and growing) developments is the PyTorch Meetup groups popping up all over the globe.

Living in Santa Cruz California I am “Silicon Valley adjacent” and there is almost too much activity “over the hill”. Getting over the mountains to the activity, while relatively close as the crow flies, is not the easiest. But, by being so close we actually have a thriving tech. community with down to earth people who appreciate being closer to Nature and a slightly slower, more affordable lifestyle.

With this background I thought, “Hey Steve, we have a major research university and a great tech. community, let’s get a PyTorch Meetup going here.”

What I hope you get from this post

I have a twofold purpose in writing today’s post:

1. Fill you in on what was covered at our first meetup – tl;dr it was a great session, especially for those first diving into AI and PyTorch
2. Inspire you to start a PyTorch meetup in your local area. It doesn’t have to be big, it doesn’t have to meet every 2 weeks, and it doesn’t have to be focused on low-level PyTorch or those directly contributing to the code base.

PyTorch has come a long way from being an internal Meta project with a strong focus on basic research to the foundation of the modern AI stack. The breadth of contributors, both to doc and the codebase has grown considerably, broadening the capabilities and range of what PyTorch (and it’s ecosystem) can do. As a consequence, It is now used, either directly or indirectly, by almost everyone in the AI and ML space and I believe we should be celebrating and showing off all the wider work PyTorch makes possible.

How our PyTorch Meetup came to be

Red Hat, true to our upstream first roots, wanted to continue to grow PyTorch’s broader community – more skill levels and more geographies. I have been to a few meetups in the Santa Cruz area, and met with various students at faculty at UCSC, and I have always been impressed with the skills and perspective that comes from being close to the tech epicenter yet staying just a bit outside the bubble.

Another key factor was that UCSC has an active and involved OSPO office that had worked with my team before. So when the UC Open Summit happened this Spring in Berkley I made sure to attend and try to touch base with some of the UCSC OSPO staff. Sure enough, once I met with Stephanie Lieggi, my co-organizer, and told her about my idea, she was excited to work together and we almost immediately kicked off planning. They contacted some of their people and I contacted some of my people, we got some food and they got a venue and parking – let the meetup commence!

We had about 45 people attend the first meetup and we had the talks in a lecture hall in one of the engineering buildings. We had a wide range of people show up with some professional engineers from the local community, some business leaders, some graduate students, and rounding out the group were undergraduates there on summer internships. The weather was gorgeous, which is typical for Santa Cruz in the non-rainy season, so we served Mediteranean dinner outside in the courtyard. There was also some swag giveaways and, as always, the most popular items were STICKERS!!

The successful first meetup

The evening began with a really well done talk by Faradawn Yang from NVIDIA, where he focused on teaching us the different parts of the GPU. His approach was something I had not seen before. He started with basic matrix math operations and how these are bread and butter for CUDA cores. Then he helped us understand how Tensor cores are used to group higher level matrix operations, thereby avoiding overhead and speeding up the operations. He even helped explain which parts are used in the forward-pass and back-pass for inference versus training.

Faradan’s talk turned out to be a great introduction for the second talk by Anil Vishnoi from Red Hat. Anil’s talk was about programming for the GPU and for PyTorch Kernels in particular. First he helped us understand how the CPU and GPU differ in their architecture, threading, and overall model of execution. He gave an introduction to the patterns you needed to be aware of when working with GPU, how the type of operation was important, and tips on getting started with the work.

To wrap up the evening we had 3 lightning talks by speakers associated with UCSC. These speakers did a great job of demonstrating the reach and ease of using PyTorch and AI with modern stacks. The first was Filippo Balzaretti, a chemistry PostDoc, demonstrating how PyTorch could be used to approximate interatomic potentials. There are existing analytical and statistical solutions to understanding the forces but they did not scale well and become unsolvable in human relevant time. By using neural networks in PyTorch they are able to simulate much larger systems and successfully predict the properties that mattered to them. It was interesting to see how this has broad applications from batteries to quantum computing and is supported by most of the large AI labs.

The next talk was by Kameron Benjamin, an undergrad at Florida A&M University, and it had a special place in my heart as an ecologist and someone who loves computer vision. He has a busy schedule and he was noticing his house plants were starting to get sick. So, in the best open source fashion, he scratched that itch and made a computer vision based AI application to diagnose his plants based on pictures of the leaves. He built Green Guardian and he got bonus points for giving us a live demo of the application in action. It was great to hear his voyage, using the resources on the PyTorch site and many tutorials throughout the internet to build this project on his own.

Bringing us home was Ph.D. student Manasi Pawar, whose research focuses on Visual Language Action models. The focus of her talk for the evening was “how much can we trust output reasoning traces from these models for why the model chose a particular action”. She is focused on their use in autonomous vehicles and so the ability to trust these traces has critical implications for model (and human) safety. It was intriguing to watch how they correlated actual driving traces with the reasoning from the model. Spoiler: just like we have learned in LLMs, you can NOT trust the reasoning traces to provide you accurate accounts for the output produced. Manasi gave a very convincing example where the reasoning trace explained why it took a right hand turn when the model had actually taken a left turn.

Join us at our upcoming PyTorch meetup

Overall it was not only an informative but entertaining meeting as well. It was great to see AI practitioners from throughout the Santa Cruz community talking and sharing their knowledge. We already have our next meetup scheduled and our first speaker lined up with more to come. We will have Rob Timpe from OpenTeams speaking about using and debugging Torch.compile.

With the active tech scene in Santa Cruz we are looking to get speakers from local companies such as Joby, Looker(Google), or Fullpower-AI. We are also thinking about exploring different venues such as Pacific Workplaces, CitizenSpace, Cruzio, or maybe even one of our public libraries.

If you have a suggestion for a topic, want to be a speaker, or know of a good location we would love to hear from you. Drop a note on our Luma calendar page or contact me directly.

Wrap up

Our first meetup was a good start, but like I said in the beginning, one of the main goals of this post is to see this energy replicated. If you’ve been on the fence about starting a PyTorch community in your own backyard, consider this your nudge to go for it. You don’t need a massive budget or a rigid structure to make a difference:

• Start small: You don’t need a massive crowd to have a great, high-impact conversation.
• Leverage your strengths: Whether you have a university nearby or a local tech scene, build on the existing pockets of curiosity in your area.
• Keep it low-key: The best meetups are the ones that feel welcoming and low-fuss, not overly formal or intimidating.
• Broaden the scope: Remember that PyTorch touches almost every corner of the modern AI stack – there is always something interesting to explore. For inspiration look at the range of PyTorch Foundation hosted projects and the even broader use cases supported by member of the PyTorch Ecosystem.
• Most importantly, have fun: Technology is better when you explore it together.

We’re thrilled with how our inaugural event turned out, and we’re already looking forward to growing this community. If you are local, we look forward to you joining us for the next one. If you can’t make it to our then no time like the present to start your own! Let’s keep building.

FBTriton Infra: Upstream Ingestion, Hierarchical Validation, Ideals vs Realities

TL:DR

Learn how Meta’s FBTriton infrastructure powers custom GPU compiler innovations like TLX and autoWS while staying synced with upstream Triton using agentic ingestion and a stratified L1/L2/L3 validation framework.

Introduction of fbtriton

Triton is a foundational element of our AI hardware acceleration strategy. While Triton is developed and maintained by OpenAI, the upstream repository alone cannot fully accommodate our internal feature requests, hardware-specific optimizations, and urgent bug fixes. In parallel, we are developing our own GPU optimization solutions, including TLX/torchTLX and autoWS, whose development timelines and code structures do not always align with upstream.

To bridge this gap, we consolidate our innovations into a downstream fork called fbtriton (pip install fbtriton). This allows us to rapidly develop features optimized for our workloads while keeping the delta from upstream as small as possible. The repository is continuously synchronized into our internal codebase with minimal in-house adjustment and powers GPU training and inference workloads across Meta’s services.

Since its consolidation in Q3 2025, fbtriton has served as an optimization runway for Meta engineers and external partners, including NVIDIA, AMD, and academic collaborators, to co-design compiler and DSL innovations and make them accessible for OSS and industry’s evaluation.

This blog covers continuous upstream ingestion, the L1/L2/L3 validation hierarchy, and the practical gap between engineering ideals and production realities.

Closing the Upstream Gap: Risk-partitioned Agentic Bundling

It is not easy for fbtriton to aggressively develop Meta-inspired optimizations while keeping the gap against a fast-moving upstream small.
Maintaining a downstream fork usually forces a choice between two strategies: periodic full-trunk rebases or continuous cherry-picking. We chose continuous cherry-picking to keep our modifications stable and decouple daily development from the structural uncertainty and friction of large rebases.
The core friction comes from architectural divergence in the compiler stack. fbtriton uses distinct strategies and designs for layout interfaces, quantization, and warp specialization.
To clear the accumulated backlog without overwhelming CI engineers with manual conflict resolution, we built an agentic loop that separates upstream commits into large low-risk bundles and context-heavy risky chains.

Step 1: Dependency tracking.
The system checks whether an incoming patch touches files or symbols linked to an ongoing complex change, also known as an existing risky chain.

Step 2: Path selection.
If a correlation is found, the patch is automatically grouped into that existing chain to preserve the correct ordering of dependent changes. Otherwise, the commit is considered safe and merged into a large low-risk bundle, such as commit #1872.

Operation metrics: Ingestion Tracking

To measure progress accurately, we track two distinct operational metrics.

Main Metric: Days Behind Upstream
This metric tracks how many days the tip of our upstream ingestion lags behind the tip of the upstream main branch. In rare circumstances, we may urgently cherry-pick specific commits from very recent upstream. These isolated picks do not affect the main metric.

Counter Metric: Backlog Commits
This metric tracks the holes left behind the ingestion tip: older upstream commits that remain unpicked and outstanding.

By decoupling these metrics, CI engineers can focus on driving down the main metric, while context-heavy backlog commits can be triaged asynchronously to keep the counter metric low. This allows the team to operate efficiently without conflating forward progress with backlog cleanup.

Out-of-Order Landing

Commits can be landed out of order as long as each one independently passes both OSS CI and internal CI. This flexibility allows us to unblock clean upstream features immediately instead of stalling behind a single complex dependency.
However, operating this way safely requires a robust, stratified hierarchical test framework discussed in the next section.

Designing the Hierarchical Test Framework

A risky Triton change, including an LLVM version bump, can trigger cascading regressions across the production stack. These issues are rarely clean build failures. Instead, they may appear as silent regressions in training/serving efficiency, increased PT2 compilation time, or subtle drift in model performance (normalized entropy).
Evaluating this entire spectrum of signals for every commit is both operationally and financially impractical. A localized single-GPU correctness test may finish in seconds, while validating job-level metrics may require GPU clusters running for hours. In practice, we manage this resource asymmetry by organizing tests into an L1/L2/L3 hierarchy based on relative value and cost.

L1: Diff tests
Fast, localized tests, including LITs (LLVM Integrated Tester), Triton unit tests, TLX tutorial kernel correctness tests, and internal customers’ kernel tests. These are triggered at every diff to prevent major breakage and kernel-level numeric mismatches.

L2: Trunk tests
Periodic, resource-intensive integration tests run on trunk, such as a tritonbench run sweeping required matrix-multiplication shapes or a distributed training job. These tests are fully bisectable on metric regressions, such as performance degradation, so we can automatically locate the culprit commit.

L3: Manual tests
Heavy, fully on-demand production workloads provided dynamically by internal production teams. These consume significant GPU hours and require explicit metric sign-off from area owners.

Discussion: Practical Engineering Problems

Moving from an abstract pipeline blueprint to a real production environment introduces operational realities across infrastructure reliability, human behavior, and shifting business context.

Derisking from Infrastructure Single Points of Failure

We cannot assume a testing platform is a flawless source of truth. We learned this when a silent bug in an underlying test infrastructure layer began omitting L1 test suites without raising alerts, creating a blind spot of unmonitored false negatives.
To eliminate this single point of failure, we adopted a saturated validation strategy by adopting various testing harnesses (such as servicelab), diverse compute capacities across both internal and OSS pipelines to improve CI signal robustness.

Managing Daily Operational Friction

Leaving an error on the trunk inevitably masks subsequent regressions. At the same time, diff authors often ignore trunk errors if a failure appears unrelated to their specific code change. These overlapping error lifecycles can quickly paralyze daily triage.
Maintaining a green trunk requires continuous team discipline and rapid daily resolution of incoming failures.

Navigating the Context Gap During Pin Updates

As a core compiler team, it is impossible to maintain a complete view of every downstream workload and model architecture across the fleet. This context gap exists both internally and in the broader OSS community.

The only viable mitigation is a dynamic, continuous context-sharing loop between teams, ensuring that compiler optimizations remain aligned with changing fleet realities.

Flawless CI: Ideals vs Realities

The engineering ideal is a fully autonomous, zero-noise, instantaneous CI/CD loop that maps perfectly to broad subsystem-level metrics. Production reality is more complicated. Coarse metrics do not fully capture risk, and fleet-scale stability requires localized operational discipline in addition to abstract subsystem-level tracking.

A high-craftsmanship CI system cannot be built in a day. It requires not only putting code together, but also putting teams together, and sustaining that cultural alignment over time.

Agentic solutions are now deeply integrated into our daily workflow, but it is important to stay clear-eyed about what has changed and what has not. AI agents are effective at eliminating tedious engineering work. We use them to resolve merge conflicts, report infrastructure issues, summarize test results, group error types, and auto-file tracking issues with proposed fixes when nightly tests break.

However, the underlying physics of the compiler and hardware remain unchanged. In practice, we must remain cautious about both AI hallucinations and human error, ensuring that agentic velocity is always guarded by deterministic safety rails.

Acknowledgments

We would like to extend our gratitude to Abhinav Singh (NVIDIA), Shucai Xiao (AMD), and Andrey Talman (PyTorch Dev Infra) for their invaluable support in providing OSS test capacity.

PyTorch Foundation Flare Pin Community Design Contest

We invite you to design the 2026 PyTorch Foundation flare pin for PyTorch Conference North America.

The winning entrant will receive one complimentary ticket to PyTorch Conference North America in San Jose, CA, in October.

PyTorch Foundation will produce the winning design as a 1-inch die-cut soft enamel pin for PyTorch Conference North America.

Submit an original design by August 14, 2026, at 11:59 p.m. PT by posting on LinkedIn, X, Facebook, or Bluesky using the hashtags #PyTorchPin and #PyTorchCon.

Design the 2026 Flare Pin

Each design must:

• Be original and suitable for a 1-inch die-cut soft enamel pin
• Include “26” or “2026”
• Include the official PyTorch Foundation symbol
• Provide the final design in SVG, PNG, or PDF format

Designs may use any shape, typeface, and color palette appropriate for distribution at a professional community event.

Use only official PyTorch Foundation assets and follow the PyTorch Foundation Brand Guidelines, Linux Foundation Trademark Usage Guidelines, and Linux Foundation Terms of Use. Official assets must remain intact and legible and must not be redrawn, distorted, or modified.

How to Enter

Each entry must include:

• One clear image of the complete, standalone design
• A brief description of the concept and its connection to the PyTorch community
• The entrant’s name or social handle
• Disclosure of any AI tools used

Post the entry publicly on LinkedIn, X, Facebook, or Bluesky using #PyTorchPin and #PyTorchCon. The post must remain public until PyTorch Foundation announces the winner.

Optional AI Design Prompt

Using an AI image generation tool? Copy the prompt below into your preferred tool to develop a concept through a guided interview. The prompt is optional, and every submission must still meet all contest, copyright, brand, trademark, and production requirements outlined in this post.

For best results, upload the official PyTorch Foundation symbol to your image generation tool when it supports image attachments.

You are designing the 2026 PyTorch Foundation flare pin for PyTorch Conference North America in San Jose.

The design should reflect the PyTorch community and remain appropriate for distribution at a professional community event.

Use the official PyTorch Foundation symbol available here:

https://github.com/pytorch-fdn/artwork/blob/main/assets/img/pytorch_foundation_logo/icons/png/PyTorchFLogo_Icon_fullColor.png

Include the official symbol in the design. Keep it intact and legible. Do not redraw, distort, recolor, or otherwise modify it. Any other PyTorch Foundation logos, symbols, or brand assets must come from official files only.

Beyond the required symbol, you may use any original shape, imagery, typeface, and color palette appropriate for distribution at a professional community event.

Design requirements:

• Create an original design suitable for a 1-inch die-cut soft enamel pin.
• Include “26” or “2026.”
• Include the official PyTorch Foundation symbol.
• Use flat, clearly separated color areas and details that remain legible at 1 inch.
• Avoid unnecessary gradients, shadows, transparent effects, and overly fine details.
• Create a design appropriate for distribution at a professional community event.

Output requirements:

• Generate exactly one design variation in a single image.
• Show the complete, standalone design as flat, front-facing artwork.
• Keep the entire design visible and uncropped.
• Use a transparent background where supported. Otherwise, use a solid white background.
• Do not render the design as a photograph, physical product mockup, or environmental scene.
• Do not include watermarks, overlays, lighting effects, or other effects that obscure the artwork.

Develop the creative direction through an interactive interview.

Do not generate an image yet. Ask me one follow-up question at a time and refine the concept based on each answer. Cover all 4 areas below:

1. Imagery and composition: What imagery or theme should be used in conjunction with the official symbol, and how and where should “26” or “2026” appear?
2. Shape and Colors: What die-cut shape and which colors should the design use?
3. Visual style: What visual approach should guide the design, such as retro, minimal, geometric, isometric, or California-inspired?
4. Community connection: How should the concept connect to the PyTorch community?

After I answer all four areas, combine my answers into one concise final concept and ask me to approve or revise it.

Do not generate the image until I explicitly approve the final concept.

After approval, generate exactly one pin design that follows the approved creative direction and all requirements above.

Selection

The PyTorch Foundation Marketing Committee will select the winning design based on originality, legibility at 1 inch, use of the required symbol and year, brand compliance, and suitability for production.

The winner will be contacted on August 19, 2026, and must respond and provide final artwork by August 24, 2026. The committee’s decision is final, and PyTorch Foundation may decline to select a winner if no eligible entry meets the requirements.

PyTorch Foundation may adjust the winning design to meet production, brand, or trademark requirements.

Contest Rules and Submission Details

Show the flat design clearly without cropping, watermarks, overlays, or mockup effects that obscure the artwork. Additional mockups are permitted.

Social entries must be published by the deadline. Entrants unable to submit through an eligible social platform may email the required materials to marketing@pytorch.org with the subject line 2026 PyTorch Pin Design Submission. Email entries must be received by the same deadline.

AI-assisted and AI-generated designs are permitted. Entrants remain responsible for ensuring that their designs meet all contest, copyright, brand, and trademark requirements.

Entrants must be at least 18 years old. Joint or group entries are not eligible.

Entrants must own or have permission to use every element in the design. Entries containing unauthorized third-party trademarks, copyrighted material, or obscene, sexually explicit, hateful, discriminatory, harassing, threatening, or otherwise inappropriate content are not eligible.

Permission to use PyTorch Foundation assets applies only to creating and submitting an entry. It does not authorize entrants to manufacture, sell, or distribute PyTorch Foundation-branded merchandise.

Artwork Rights

Entrants retain ownership of their artwork but grant PyTorch Foundation and its event and production partners a nonexclusive, worldwide, royalty-free license to reproduce, publish, display, and promote submissions in connection with the contest.

For the winning design, the license also covers adaptation, production, distribution, photography, and promotion of the pin.

Prize and Travel

The complimentary ticket may not be sold, transferred after issuance, or redeemed for cash.

The ticket covers conference admission only. PyTorch Foundation will not pay or reimburse travel, lodging, ground transportation, meals, visa-related costs, or other attendance expenses.

Eligible community members may separately apply for The Linux Foundation Travel Fund. Applications close August 23, 2026, at 11:59 p.m. PT, and applicants will be notified by September 2, 2026. Travel funding is separate from the contest and is not guaranteed.

Questions may be sent to marketing@pytorch.org.

PyTorch Foundation and the PyTorch Foundation logo design are registered trademarks of the Linux Foundation.

Updated August 3, 2026: This post now includes an optional AI design prompt to help entrants create a flat, front-facing pin concept that follows the contest requirements.

Helion on TPU: Towards Hardware Heterogeneous Kernel Authoring

TL;DR

Helion is PyTorch’s high-level DSL for writing performance-portable ML kernels. Partnering with Google, we have built a TPU backend that compiles Helion kernels to Pallas, providing a PyTorch-friendly way to author performant TPU kernels. On a flash attention workload, the Helion-generated kernel achieves 838 TFLOPs (~79% MFU of one tensor core) on TPU v7. On different input shapes, Helion autotunes over different code-generation strategies to select the optimal pipelining schema, making the most use of TPU’s available VMEM and compute.

Introduction

TPUs are increasingly important as an ML compute platform to complement GPUs. Google’s latest TPU v7 (Ironwood) delivers comparable performance to NVIDIA B200 with a potentially lower total cost of ownership (TCO), making TPUs an appealing option for large-scale training and inference workloads. However, authoring TPU kernels traditionally requires expertise in Pallas, a low-level DSL that comes with a steep learning curve and code complexity. Helion bridges this gap. As PyTorch’s portable DSL for ML kernels, Helion lets users write familiar PyTorch-style code and compiles it to optimized TPU code. Paired with performance wins brought by its autotuner, Helion is evolving towards an attractive option for authoring TPU kernels. Specifically, Helion TPU targets three main use cases:

• Performance-critical use cases where autotuning is required to explore the configuration space
• Non-Pallas experts hoping to onboard TPU kernel authoring quickly
• Cross-hardware users who prefer to maintain the same set of kernels across TPU and GPU

This article starts with a brief overview of TPU’s hardware features and programming models as compared to GPUs, and then demonstrates how Helion generates performant Pallas code with ideal pipelining characteristics for different input shapes.

TPU Primer

TPUs are highly specialized accelerators designed and optimized specifically for machine learning workloads. The architecture and programming model of TPUs differ significantly from GPUs. The most prominent difference is that a TPU is a sequential machine featuring wide vector registers and compute units. This contrasts with GPUs, which achieve performance via both massively parallel execution (CUDA cores) and specialized tensor units (tensor cores).

As a result, TPUs feature a memory hierarchy that kernel authors must deeply understand, so that the kernels they write can orchestrate when and how data is loaded from the off-chip HBM to the fast on-chip VMEM. A performant Pallas kernel would overlap these HBM<>VMEM memory transfers with floating point computation happening in the matrix (MXU) and vector compute units.

Despite the architectural differences, current-generation TPUs and GPUs are highly comparable in raw performance. TPU7x and NVIDIA B200 have very similar BF16 compute TFLOPS and HBM bandwidth — the two most important hardware metrics for modern ML workloads.

Helion’s Pallas Codegen

To extract maximum performance out of a TPU, Helion’s Pallas codegen aims to maximize software pipelining, ensuring that memory transfers and computation overlap as much as possible. This section illustrates Helion’s three-fold strategy for generating pipelined kernels:

• Outer loop: pallas-provided pipelined device invocation (pallas_call/emit_pipeline)
• Inner loop: autotuned between:
  • pallas-provided pipelined device-side loop (emit_pipeline)
  • Pre-fetching all values into VMEM, if possible (unroll)
• Auto-tuned pipeline buffer sizes

Example: add

As a simple example, consider the following helion kernel for adding two tensors.

@helion.kernel
def add(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
    out = torch.empty_like(x)
    for tile in hl.tile(out.size()):
        out[tile] = x[tile] + y[tile]
    return out

The Helion compiler translates this into two functions: a host-side launcher that tiles the input and invokes the device function in a pipelined fashion, and a device function that operates on VMEM-resident tiles:

def _helion_add(x, y, out):
    out[:] = x[:] + y[:]

def add(x: torch.Tensor, y: torch.Tensor):
    _BLOCK_SIZE_0 = <autotuner-selected value>
    out = torch.empty(...)
    out = launcher( # wraps around pallas_call
        _helion_add, 
        ((x.shape[0] + _BLOCK_SIZE_0 - 1) // _BLOCK_SIZE_0,), # grid size
        x, y, out, 
        _block_spec_info=[_BLOCK_SIZE_0, ...], ...
    )
    return out

Within the generated code:

• The hl.tile loop in the Helion source becomes a grid on the host side. The launcher (wrapping pallas_call) invokes _helion_add once per tile, with each invocation automatically pipelined — while one tile is being computed, the next tile’s data is being loaded from HBM into VMEM.
• The device function _helion_add is simple: it receives VMEM references (not HBM pointers), so the kernel body is a simple addition.
• _BLOCK_SIZE_0 (the tile/buffer size) is selected by the autotuner, which explores different sizes to find the best overlap between memory transfers and compute for the target hardware.

This results in a pipelined execution as illustrated below.

Example: Flash Attention

Attention is one of the key operations in modern language models. Production implementations follow the “Flash Attention” pattern – a memory-efficient technique that computes attention in tiles to avoid materializing the full S×S attention matrix. The structure of a flash attention kernel in Helion is illustrated below:

B, H, S, D = 8, 32, 8192, 256 # batch, head, sequence length, head dimension
@helion.kernel
def attention(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Tensor:
    out = torch.empty(...)
    for tile_b, tile_q in hl.tile(B * H, S):
        this_q = q[tile_b, tile_q, :]
        acc = ...
        for tile_kv in hl.tile(S):
            this_k = k[tile_b, tile_kv, :]
            this_v = v[tile_b, tile_kv, :]
            <qk matmul, online softmax, v matmul, update acc>
        out[tile_b, tile_q] = acc
    return out

Compared to the “add” example discussed previously, the flash attention kernel contains an additional inner loop which performs tiled accesses across the entire K and V sequences. How we pipeline the memory and compute within the inner loop is key to the performance of this kernel.

In Helion, the compiler autotunes over two different strategies for translating this kernel to Pallas. This is keyed on the pallas_loop_type autotuner config.

With the default pallas_loop_type == emit_pipeline option, Helion relies on Pallas’ device-side emit_pipeline API to pipeline an inner loop body function, similarly to how the host-side logic uses pallas_call to pipeline the device function invocation:

def _helion_attention(q_VMEM, k_HBM, v_HBM, out_VMEM):
    acc = ...
    this_q = q_VMEM[:, :, :]
    def _inner_pipeline_body(k_VMEM, v_VMEM):
        this_k = k_VMEM[:, :, :]
        this_v = v_VMEM[:, :, :]
        <matmul, online softmax, matmul, update acc>
    pallas.tpu.emit_pipeline(_inner_pipeline_body, k_HBM, v_HBM, _block_spec_info=[BLOCK_SIZE_KV ,...], ... )
    out_VMEM = acc

def attention(q: torch.Tensor, k: torch.Tensor, v:torch.Tensor):
    out = torch.empty(...)
    out = launcher( # wraps around pallas_call
        _helion_attention, 
        q, k, v, out, 
        _block_spec_info=[BLOCK_SIZE_Q ,...], ...
    )
    return out

This generated kernel follows a nested pipeline structure:

• (Outer pipeline) The host uses pallas_call to invoke _helion_attention. The HBM reference of q is tiled, and each invocation of _helion_attention receives a VMEM tile of q. For k and v, _helion_attention receives HBM references directly.
• (Inner Pipeline) Within _helion_attention, the device uses emit_pipeline to invoke _inner_pipeline_body, which receives VMEM tiles of k and v.

This results in a pipelined execution as illustrated in this image:

One obvious point of inefficiency in this pipeline is that there are bubbles in the compute units – for every new Q tile, while we fetch the 0th KV tile, there is no work available for the compute units. This comes down to the fact that we are re-loading the KV tiles from HBM to VMEM for every new Q tile.

Helion offers an alternative pallas_loop_type == unroll config which avoids this bubbling. With unroll, we translate the inner for loop into a simple Python for loop:

def _helion_attention(q_VMEM, k_VMEM_FULL, v_VMEM_FULL, out_VMEM):
    acc = ...
    this_q = q_VMEM[:, :, :]
    for offset in range(0, k_VMEM_FULL.size(1) , BLOCK_SIZE_KV):
        this_k = k_VMEM_FULL[:, pallas.dslice(offset, BLOCK_SIZE_KV), :]
        this_v = v_VMEM_FULL[:, pallas.dslice(offset, BLOCK_SIZE_KV), :]
        <matmul, online softmax, matmul, update acc>
    out_VMEM = acc

def attention(q: torch.Tensor, k: torch.Tensor, v:torch.Tensor):
    out = torch.empty(...)
    out = launcher( # wraps around pallas_call
        _helion_attention, 
        q, k, v, out, 
        _block_spec_info=[BLOCK_SIZE_Q, None, None], ...
    )
    return out

(The name “unroll” reflects the fact that Pallas device functions are traced by JAX’s JIT — the Python for loop is effectively unrolled at trace time into a flat sequence of operations.)

In this version of the generated kernel:

• K and V are pre-fetched in full: The host passes None as the block spec for K and V, instructing pallas_call to load them entirely into VMEM. The full VMEM references persist across all device function invocations.
• The inner loop slices locally: Each iteration uses pallas.dslice to select the relevant KV tile from the already-resident VMEM buffer. No HBM traffic occurs during the inner loop.

This results in a different pipelining scheme as illustrated below:

In this workflow, there are no longer bubbles in the compute pipeline. The trade-off is that this requires more VMEM usage, as the entire K and V sequences need to be present. This means that although more performant, this translation isn’t always possible. The VMEM usage is linear with respect to the input sequence lengths (as opposed to tile size), which is prohibitive with longer sequences. The performance difference between these strategies is significant — the table below shows results on workloads with B=8, H=32, D=256:

The benefit of Helion lies in its ability to autotune and select the best autotuner config. So that with smaller sequences, it makes use of the VMEM available and generates pipelined code with no compute bubbles. For longer sequences, it falls back to emit_pipeline which scales to arbitrary context lengths. The following graph plots the performance of this attention kernel compared to various other Pallas attention implementations, on varying sequence lengths:

The autotuner’s ability codegen different loop and pipelining strategies depending on the input length is what gives Helion its edge even when compared to highly optimized implementations such as Tokamax.

Broader Kernel Benchmarks

We benchmark Helion across a variety of kernels, tracked on our dashboard. The table below compares Helion against TorchTPU eager and torch.compile (using XLA) across a range of different kernels. Helion shows a geometric average speed-up of 1.55x compared to eager, and 1.12x compared to compiled.

Helion shows the largest gains on kernels that employ fusion or optimization patterns that are difficult for XLA to discover automatically — flash attention is a prominent example. For the more standard operations like matmul and layer_norm, XLA’s compiler already produces high-quality code, and Helion performs comparably.

What’s Next

Helion on TPU is under active development. Here’s a non-exhaustive list of things we are working on:

• Expand kernel coverage: Get more Helion examples working on TPU
• Further performance improvements
• Better support for jagged and sparse operations
• Support for distributed TPU computing

Getting Started

Helion is open source and available on GitHub. Its TPU backend has a dependency on TorchTPU, which is expected to be released publicly later this year. When it does, we encourage you to try-out Helion on TPU and share your feedback. Resources:

• Helion GitHub Repository
• Helion Documentation
• Helion TPU Examples
• Performance Dashboard

Acknowledgements

This project was made possible through the invaluable collaboration and technical insights of our peers. A special thank you to Joe Pamer, Robert Hundt, Claudio Basile and Adam Paszke at Google, as well as Jana van Greunen, Gregory Chanan, Peng Wu, and Zongwei Zhou at Meta, for their feedback and support in bringing this to fruition.

Driving the Future of Open Source AI: An Update from PyTorch Foundation Projects

TL;DR

In April 2025, the PyTorch Foundation evolved into a multi-project Foundation, with the objective to support deeper collaboration across domains and help scale innovation throughout the AI lifecycle. Today, the Foundation hosts six projects, including PyTorch, vLLM, DeepSpeed, Ray, Helion, and Safetensors. From core optimization updates to hardware enablement and community health, PyTorch Foundation-hosted projects have achieved a lot this past quarter. In this new blog series, each of our projects provides an update on its latest achievements.

PyTorch Updates: 2.13 Release, ExecuTorch, and Apple Silicon Optimization

Development momentum for the core PyTorch project remains incredibly strong. Commits continued to rise with 4,415 in Q2 alone, and issue management is on track, with the overall number of open issues steadily decreasing.

The community officially rolled out PyTorch 2.13 and has already kicked off the development cycle for PyTorch 2.14. PyTorch 2.13 headlines performance and platform reach: FlexAttention arrives on Apple Silicon (up to ~12x faster than SDPA) alongside a new CuTeDSL Inductor backend, a memory-saving nn.LinearCrossEntropyLoss (up to 4x less peak GPU memory), and Python 3.15 wheel support including free-threaded builds. On the distributed and platform side, it adds the torchcomms backend and FSDP2 communication overlap for large-cluster training, broadens ROCm/Arm/Intel XPU support. The team is actively expanding on-device LLM capabilities through ExecuTorch, releasing weekly examples alongside a partnership with Hugging Face to streamline model ingestion and support on-device performance.

vLLM Updates: Model Runner V2, Q3 2026 Roadmap, and vLLM Conference

The vLLM project has reached an impressive milestone, achieving a stable, bi-weekly release cadence. Among its key achievements is the complete redesign of Model Runner V2, which yielded substantial performance improvements on GPTQ. vLLM continues to offer day-zero support for massive, state-of-the-art models such as Kimi K3, Minimax M3, and Qwen 3.8.

Looking ahead, the team has published its Q3 2026 roadmap. The roadmap is organized across special interest groups (SIGs) and centers on production agentic workloads and high-interactivity premium tokens. On the core engine, the team is finishing two major migrations, Flat Model and Model Runner V2, and redesigning how the scheduler and KV cache work. For large-scale serving, the goal is top performance on AgentX, with better KV cache offloading and smarter prefix caching for multi-turn agents. Other groups are focused on faster speculative decoding, production-ready KV-cache compression, and quicker CI. There are also dedicated tracks for model performance, vLLM-Omni, and reinforcement learning.

The team is also hosting the first-ever vLLM Conference, held at Ray Summit in San Francisco from August 24–26. The event brings together engineers and researchers from across the vLLM ecosystem for two days of sessions spanning the vLLM roadmap, hardware backends, agentic serving, training, and production-scale inference.

DeepSpeed Updates: Ulysses Parallelism, Torch.xpu Integration, and ASPLOS 2026

DeepSpeed continues to deliver rapid software innovations, maintaining its stable bi-weekly schedule with six software releases in the first quarter. To foster closer collaboration with developers, the team stabilized its monthly virtual office hours as consistent community touchpoints. In a major architectural update, the project replaced Intel GPU support (IPEX) by directly integrating Torch.xpu. The team also rolled out updates focusing on mixed precision fixes, AMD SDMA optimization, automatic sequence parallelism, and Muon optimizer.

One of the main highlights for the project is the integration of the Ulysses sequence parallelism algorithm directly into Hugging Face libraries, including Trainer, Accelerate, and TRL. DeepSpeed’s theoretical contributions have also gained recognition in the academic community, earning a best-paper honorable mention at ASPLOS 2026 for its pioneering SuperOffload work, and having three research papers accepted. The first paper (AutoSP) automates sequence parallelism using PyTorch 2.0 compile technology, the second paper explores energy analysis and offloading overheads on Grace Hopper C2C systems, while the third paper focuses on efficient offloading mechanics for model training.

Ray Updates: GB200/GB300 Hardware Support, Frontier Model Scaling, and Ray Data 2.57

Ray remains focused on production hardening for large-scale workloads. This work includes speeding up actor scheduling, upgrading native RDMA support, improving topology-aware scheduling for rack-aware actor placement, and improving support for GB200s/GB300s and beyond. The team is also focused on optimizing workflows for reinforcement learning as well as for multimodal data and video processing.

In the last few months, Ray has been used for building a variety of frontier models, including MAI-Thinking-1, Composer 2.5, and Nemotron 3 Ultra.

A few other recent highlights from the team include inference performance improvements with Ray + vLLM, improved cluster stability via resource isolation, improved data pipeline reliability, and prefill-decode disaggregation with Ray + vLLM on AMD MI325X.

The team has also made important updates to data and inference pipelines. A brand-new, high-performance engine for Ray Data is scheduled to debut in version 2.57. At the same time, the team is actively building out weight-syncing integrations with vLLM and SkyRL to eliminate performance bottlenecks during post-training. If you want to connect with the team and community in person, Ray Summit is set for August 24-26 in San Francisco.

Helion Updates: Cross-Hardware Attention Kernels and LLM-Guided Autotuning

Helion has made substantial progress during its first six months of the year, focusing its efforts on delivering state-of-the-art performance across heterogeneous hardware. The project introduced the CuTeDSL backend for NVIDIA GPUs and Pallas backend for TPUs. With the new CuteDSL and Pallas backends, using the same Helion attention kernel, Helion delivers state-of-the-art performance on NVIDIA Blackwell GPU, outperforming FlashAttention-4, and on Google TPU, outperforming hand-written Attention kernel from the Tokamax library, demonstrating Helion’s performance portability. The Pallas backend is being built in partnership with Google and debuted at the Google AI Systems DevLabs.

In addition to performance improvements and support for new hardware, Helion incorporated the power of LLMs to kernel autotuning. The team recently introduced an LLM-guided autotuner to significantly speedup autotuning, demonstrating an 10x improvement in tuning efficiency, paving the way for faster, highly optimized deployments.

To engage with the community, Helion held its first hackathon in partnership with NVIDIA and GPU Mode and a tutorial at the PLDI conference. Helion was announced in General Availability at the PyTorch Conference Europe.

Safetensors Updates: GIL-Free Serialization, Python 3.14 Support, and MPS Fast-Loading

In June, Safetensors officially rolled out Global Interpreter Lock (GIL)-Free Serialization, which introduces enhanced parallel processing capabilities to improve overall multi-threaded performance. To provide early support for upcoming Python releases, the project has expanded its continuous integration pipeline to accommodate Python 3.14 and 3.14t. For developers operating on compatible hardware, the new Metal Performance Shaders (MPS) fast-load paths introduce native Metal buffers within the torch MPS backend to significantly minimize model loading latency. Finally, new dtypes have been added to offer developers greater precision and flexibility across diverse modeling workloads.

Looking ahead, Safetensors is actively working on optimizing CUDA fast-path loading capabilities, exploring integrations with both io_uring and GPU Direct Storage (GDS). As an initial step, the project is implementing multithreaded pread operations into pinned host buffers to accelerate data transfers. This structural work aims to make the library significantly more adaptable to additional custom fast paths while removing our historical reliance on framework-managed storage, such as the from_file memory mapping found in torch. The team is currently evaluating the feasibility of this transition, as it depends on DLPack handoffs that may not be universally supported across all frameworks.

Join the Community

Are you ready to build the future of AI with us? To learn more about our hosted projects, visit our Foundation-Hosted Projects portal, and to learn how your organization can become a member, visit our Join Now page. Plus, join us in Shanghai for KubeCon + CloudNativeCon + OpenInfra Summit + PyTorch Conference China 2026 September 7-9, and in San Jose for PyTorch Conference North America 2026, October 20-21, later this year.

PyTorch Conference North America Schedule Is Live

PyTorch Conference North America will bring developers, researchers, and practitioners to San Jose on October 20–21 for sessions spanning training and inference, compiler innovations, responsible AI, applications, and the PyTorch ecosystem. PyTorchCon NA 2026, hosted by the PyTorch Foundation, will explore the future of open source AI and the impact of PyTorch Foundation projects like PyTorch, vLLM, DeepSpeed, Ray, Helion, and Safetensors.

View the full schedule here >

The program includes sessions on observability tooling for Cudagraph workloads, accelerating and debugging machine learning systems with TorchDynamo, and multi-node training for foundation models.

Session Highlights

Observability Tooling for Cudagraph Workloads

Natalia Gimelshein and Driss Guessous, Meta

Unlocking the Full Potential of TorchDynamo: Accelerating, Comparing, and Debugging ML Systems

Yi Pan, UC Berkeley; Megan Frisella and Stephanie Wang, University of Washington

Scaling Foundation Models: From Broken to Near-Linear Multi-Node Training

Sheng Huang, Pinterest

Submit a Poster by July 26

The Poster CFP closes July 26 at 11:59 p.m. PDT. Poster sessions provide an opportunity to showcase projects, research, and implementations, exchange ideas with attendees, and connect directly with the PyTorch community. Submit a poster >

Register Now

Early Bird conference passes are available at a discount through Friday, July 31. Register today >

Sponsor PyTorch Conference North America

PyTorch Conference North America brings together more than 3,000 members of the open source AI ecosystem. Sponsorship opportunities provide visibility with engineers and technical leaders working on AI infrastructure and applications.

Sponsorship details >

Triton Plugin Extensions: Enabling TLX and Custom Compiler Passes Out of the Box

TLDR

The PyTorch-Triton 3.7 release introduces the Triton Plugin Extensions system, a framework for dynamically loading custom compiler passes, dialects (including their ops), and DSL extensions into upstream Triton at runtime, without forking or recompiling. As the first major consumer of this system, Meta’s Triton Language Extensions (TLX) are now enabled out of the box, bringing persistent GEMM kernels and fine-grained hardware control to stock Triton with performance that matches or exceeds vendor libraries on both NVIDIA H100 and AMD MI350.

The Problem: Why Extensions?

Writing high-performance GPU kernels often requires going beyond what the default Triton compiler pipeline provides. Custom optimization passes, hardware-specific intrinsics, and specialized memory management patterns are essential for squeezing out the last drops of performance on production workloads. Until now, enabling these capabilities meant maintaining a fork of Triton and said forks come with real costs.

Forks quickly fall behind upstream. Every upstream update risks merge conflicts, broken APIs, and subtle behavioral changes that require careful reconciliation. Teams that pin to a forked version find themselves stuck on stale releases, unable to take advantage of upstream bug fixes, new hardware support, and community improvements. The maintenance burden compounds over time, and the fork becomes a bottleneck rather than an accelerator.

What’s needed is a way to extend Triton’s compiler pipeline adding passes, ops, and even entire dialects without modifying core Triton at all. A plugin system that loads extensions dynamically at runtime would allow researchers and engineers to iterate on custom features at full speed, always running on the latest upstream release, and ship results without waiting for changes to be merged into the mainline repository.

The Triton Plugin Extensions System

The PyTorch Triton 3.7 release delivers exactly this: a general-purpose plugin extensions system that spans the entire compilation pipeline and built into upstream Triton. Plugins are shared libraries (.so files) that are discovered and loaded at runtime via the TRITON_PLUGIN_PATHS environment variable. No recompilation of Triton is required to install a plugin package, point the environment variable at it, and the extensions are immediately available.

Overridable Compiler Pipeline

At the heart of the system is a set of hooks embedded in Triton’s backend compiler.py stages. These hooks provide fine-grained control over the MLIR pass pipeline at every lowering level from higher level Triton IR (TTIR) through TritonGPU IR (TTGIR) down to LLVM IR and target-specific assembly (PTX, AMDGCN). With these hooks, plugins can:

• Insert one or more custom passes at arbitrary points in any stage.
• Disable specific passes within a stage.
• Replace existing passes with specialized custom implementations (e.g., a custom warp specialization strategy).
• Override entire stages or the full pipeline.

This is available on both the NVIDIA and AMD backends

Custom Ops, Dialects, and Lowering

The plugin API is designed to complement PyBind11, enabling three levels of extensibility:

1. Custom transformation passes: single passes that can be inserted at arbitrary points in the pipeline without an associated dialect.
2. Custom MLIR dialects and conversion passes: separately compiled dialects loaded into Triton, with plugin passes that rewrite standard Triton IR patterns into custom dialect ops for specialized lowering.
3. Custom top-level DSL ops: new Python-level syntax and semantics enabling entirely new programming abstractions without altering Triton itself.

Per-Kernel Control

Plugins can be toggled on and off dynamically at the kernel level. A compiler hook set in kernel code activates a custom pipeline for all kernels called after the hook is set, until it is unset. There is no limit on how many custom pipelines can be defined, and plugins are responsible for implementing their own hashing strategy for kernel cache management—ensuring that recompilation is triggered only when needed. This is handled entirely by the utlx library for the user.

# Enabling the TLX plugin is as simple as setting an environment variable
import os
import sysconfig

dist_packages = sysconfig.get_paths()["purelib"]
libutlx_path = os.path.join(dist_packages, "utlx_plugin", "libutlx.so")
os.environ["TRITON_PLUGIN_PATHS"] = libutlx_path

TLX: Triton Language Extensions, Now Built-In

Triton Language Extensions (TLX) is a set of hardware-aware operations developed by Meta for explicit memory management and asynchronous compute/load pipelining. TLX gives kernel authors direct control over shared memory allocation, data movement, and instruction scheduling—capabilities that are critical for writing persistent kernels that saturate modern GPU hardware.

The core TLX operations include:

Previously, using TLX required building from Meta’s experimental Triton fork. With the plugin extensions system, TLX is now distributed as a standalone Python package (utlx) that works with unmodified upstream Triton. Starting with PyTorch-Triton 3.7, TLX will be enabled by default on all Triton releases going forward.

Cross-Hardware: NVIDIA H100 and AMD MI350

One of the key advantages of TLX is that the same programming model works across hardware vendors, while still mapping to vendor-specific features under the hood.

NVIDIA H100 (Hopper) — Persistent GEMM

On Hopper GPUs, TLX maps to hardware-native TMA (Tensor Memory Accelerator) async loads and WGMMA (Warp Group Matrix Multiply-Accumulate) instructions. The persistent GEMM kernel uses multi-stage software pipelining with async commit/wait groups:

@triton.jit
def matmul_kernel_pipelined_hopper(
    a_ptr, b_ptr, c_ptr, M, N, K,
    stride_am, stride_ak, stride_bk, stride_bn, stride_cm, stride_cn,
    BLOCK_SIZE_M: tl.constexpr, BLOCK_SIZE_N: tl.constexpr,
    BLOCK_SIZE_K: tl.constexpr, GROUP_SIZE_M: tl.constexpr,
    NUM_STAGES: tl.constexpr,
):
    # ... tile indexing ...

    # Allocate multi-stage shared memory buffers
    buffers_A = tlx.local_alloc((BLOCK_SIZE_M, BLOCK_SIZE_K), tlx.dtype_of(a_ptr), NUM_STAGES)
    buffers_B = tlx.local_alloc((BLOCK_SIZE_K, BLOCK_SIZE_N), tlx.dtype_of(b_ptr), NUM_STAGES)

    # Prefetch pipeline prologue
    for i in tl.range(0, NUM_STAGES - 1, loop_unroll_factor=NUM_STAGES - 1):
        a = tlx.local_view(buffers_A, i)
        b = tlx.local_view(buffers_B, i)
        token_a = tlx.async_load(a_ptrs, a, mask=...)
        token_b = tlx.async_load(b_ptrs, b, mask=...)
        tlx.async_load_commit_group([token_a, token_b])

    # Main K loop with overlapped compute and data movement
    acc = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
    for k in tl.range(0, tl.cdiv(K, BLOCK_SIZE_K), num_stages=0):
        buf = k % NUM_STAGES
        tlx.async_load_wait_group(NUM_STAGES - 2)
        acc = tlx.async_dot(
            tlx.local_view(buffers_A, buf),
            tlx.local_view(buffers_B, buf), 
            acc
        )
        # Prefetch next stage ...

    acc = tlx.async_dot_wait(0, acc)
    # Store results ...

NVIDIA H100 (Hopper) — Performance Results

The table below shows FP16 GEMM throughput on NVIDIA H100 comparing stock Triton with the TLX extension plugin against cuBLAS. Since the plugin system produces identical codegen to the compiled-in fork, these results apply to both paths.

On the large, compute-bound shapes that dominate production LLM workloads, Triton + TLX matches cuBLAS on square GEMM and exceeds it on the wide and large shapes, confirming that the plugin-loaded path introduces zero overhead while reaching or beating the vendor library:

128×13312×16384: cuBLAS 247.8 TFLOPS → Triton+TLX 257.0 TFLOPS (+3.7%)
16384×8192×8192: cuBLAS 549.4 TFLOPS → Triton+TLX 566.7 TFLOPS (+3.2%)
8192×16384×8192: cuBLAS 564.8 TFLOPS → Triton+TLX 575.9 TFLOPS (+2.0%)
8192×53248×8192: cuBLAS 571.3 TFLOPS → Triton+TLX 573.2 TFLOPS (+0.3%)
8192×28672×4096: cuBLAS 560.4 TFLOPS → Triton+TLX 559.8 TFLOPS (−0.1%)
8192×8192×8192: cuBLAS 582.3 TFLOPS → Triton+TLX 577.0 TFLOPS (−0.9%)

AMD MI350 — Pipelined GEMM

On AMD MI350 GPUs, TLX uses explicit register-based pipelining with local_store and local_load operations. The same buffer management pattern applies, but the data movement path goes through registers rather than async hardware units:

@triton.jit
def matmul_kernel_pipelined_mi300(
    a_ptr, b_ptr, c_ptr, M, N, K, ...
):
    # ... tile indexing ...

    # Allocate shared memory buffers
    buffers_A = tlx.local_alloc((BLOCK_SIZE_M, BLOCK_SIZE_K), tlx.dtype_of(a_ptr), NUM_STAGES - 1)
    buffers_B = tlx.local_alloc((BLOCK_SIZE_K, BLOCK_SIZE_N), tlx.dtype_of(b_ptr), NUM_STAGES - 1)

    # Prologue: load into shared memory via registers
    for i in tl.range(0, NUM_STAGES - 1, loop_unroll_factor=NUM_STAGES - 1):
        a_reg = tl.load(a_ptrs, mask=...)
        b_reg = tl.load(b_ptrs, mask=...)
        tlx.local_store(tlx.local_view(buffers_A, i), a_reg)
        tlx.local_store(tlx.local_view(buffers_B, i), b_reg)

    # Main loop: overlapped compute and memory operations
    acc = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
    for k in tl.range(NUM_STAGES - 1, K_ITERS, num_stages=0):
        # Load next tile into registers
        a_reg = tl.load(a_ptrs, mask=...)
        b_reg = tl.load(b_ptrs, mask=...)
        
        # Compute on previously staged data
        a_prev = tlx.local_load(tlx.local_view(buffers_A, buf))
        b_prev = tlx.local_load(tlx.local_view(buffers_B, buf))
        acc = tl.dot(a_prev, b_prev, acc)
        
        # Store new data to shared memory for next iteration
        tlx.local_store(tlx.local_view(buffers_A, ...), a_reg)
        tlx.local_store(tlx.local_view(buffers_B, ...), b_reg)

AMD MI350 — Performance Results

The table below shows FP16 GEMM throughput on AMD MI350 comparing stock Triton with the TLX extension plugin against rocBLAS. Since the plugin system produces identical codegen to the compiled-in fork, these results apply uniformly to both paths.

Triton + TLX delivers 12–15% higher TFLOPS consistently across all tested matrix sizes, confirming that the plugin-loaded path introduces zero overhead while exceeding the vendor library:

256×256×256: rocBLAS 4.4 TFLOPS → Triton+TLX 5.0 TFLOPS (+11.8%)
512×512×512: rocBLAS 29.4 TFLOPS → Triton+TLX 33.9 TFLOPS (+15.2%)
1024×1024×1024: rocBLAS 161.2 TFLOPS → Triton+TLX 180.8 TFLOPS (+12.1%)
2048×2048×2048: rocBLAS 445.1 TFLOPS → Triton+TLX 511.9 TFLOPS (+15.0%)

GPUMode Trimul Multiplicative Update Validation

We wanted to validate the plugin path in the production, and also on a heavy kernel pipeline closer to the real problem rather than a standalone microbenchmark. On GPU mode, there was a Trimul multiplicative update – five projection GEMMs feeding a batched matmul plus an output linear, wrapped in layer norms, sigmoid gates, and permutations. The PyTorch + torch.compile baseline ran at 19.2ms, dominated by GEMMs. With the TLX plugin loaded into stock Triton via TRITON_PLUGIN_PATHS, we dropped the warp-specialized persistent GEMM (hopper_gemm_ws.py). Because TLX exposes the matmul as just another Triton kernel, we could collapse the pipeline around it. The final TLX-WS + fusion submission ran at 12.0ms, with a 1.61x speedup over the cuBLAS + torch.compile baseline, beating libcuEquivariance and all other SOTA implementations on H100. We’ve later extended with CLC pipelining on B200, further widening the gap. The final takeaway is the actual extensions integration took a few lines of installing wheel on gpu mode and very little set up overhead.

Identical CodeGen, Zero Fork Required

A critical validation of the plugin approach is that the generated code is identical to what the Meta Triton fork produces. The TLX extension plugin goes through the same MLIR lowering pipeline; the only difference is that the passes and ops are loaded dynamically rather than compiled in.

Our demos confirm:

• Identical PTX codegen on NVIDIA H100 for persistent GEMM kernels.
• Identical AMDGCN codegen on AMD MI350 for pipelined GEMM kernels.
• Equivalent performance — no measurable overhead from the dynamic loading path.

The Colab notebooks and standalone scripts used for this validation are available below and will be updated to point to the official PyTorch-Triton 3.7 packages after the release ships.

Getting Started

Getting started with TLX on upstream Triton is straightforward:

# Install Triton (from source) and the TLX extension from PyPI package
git clone https://github.com/triton-lang/triton && cd triton
TRITON_EXT_ENABLED=ON pip install -e . --no-build-isolation && cd ..
# uTLX plugin (published on PyPI):
pip install triton-utlx

The utlx package includes the pre-built extension library. Once installed, set the plugin path and import the extensions:

import os
import sysconfig

# Point Triton at the TLX plugin
dist_packages = sysconfig.get_paths()["purelib"]
os.environ["TRITON_PLUGIN_PATHS"] = os.path.join(dist_packages, "utlx_plugin", "libutlx.so")

# Now TLX ops are available in your kernels
import triton
import triton.language as tl
import utlx_plugin as tlx

To explore the full demos and repositories:

• H100 Persistent GEMM Notebook: Colab
• H100 Standalone Script: GitHub Gist
• AMD MI350 Pipelined GEMM: GitHub Gist
• Plugin Documentation: triton/lib/Plugins/README.md
• Extension Repository: triton-lang/triton-ext
• utlx PyPI Project: https://pypi.org/project/triton-utlx/

What’s Next

The plugin extensions system opens the door to a growing ecosystem of community-developed Triton extensions. Areas of active development and future proposals include:

• Custom backends: dynamically loaded out-of-tree backends for Intel, CPU, and other targets without modifying Triton’s build system.
• triton-distributed: distributed computing primitives as an extension.
• Customized versions of instrumentation and profiling tools like Proton and ConSan developed as plugins for user specific runtime-loadable performance analysis
• Custom optimization passes: target-specific warp specialization, loop splitting, and model-specific optimizations shipped as add-ins
• Specialized ops: 2:4 structured sparsity, custom layout conversions, and more

We’re looking forward to community engagement in picking up and implementing extensions that unlock new capabilities for the broader Triton ecosystem. If you’re interested in contributing, start with the triton-ext repository and the plugin documentation.