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

Search

Close Search

Blog

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

By Corbin Robeck, Puyan Lotfi, Ian Barber, Shane Nay, Alexey Loginov, Oleksandr Stashuk, Wenyuan ChiJuly 15, 2026No Comments

Featured projects

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:

Custom transformation passes: single passes that can be inserted at arbitrary points in the pipeline without an associated dialect.

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.

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<br>import os<br>import sysconfig

dist_packages = sysconfig.get_paths()["purelib"]<br>libutlx_path = os.path.join(dist_packages, "utlx_plugin", "libutlx.so")<br>os.environ["TRITON_PLUGIN_PATHS"] = libutlx_path<br>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...