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

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:

| Operation | Description |

|---|---|

tlx.local_alloc(shape, dtype, num_buffers) |

Allocate shared memory buffers for software pipelining. |

tlx.local_view(buffers, index) |

View a specific buffer within an allocation. |

tlx.async_load(src, dst, mask) |

Initiate an asynchronous load from global to shared memory. |

tlx.async_load_commit_group(tokens) |

Commit a group of async loads. |

tlx.async_load_wait_group(n) |

Wait for async load groups to complete. |

tlx.async_dot(a, b, acc) |

Asynchronous matrix multiply-accumulate. |

tlx.async_dot_wait(n, acc) |

Wait for async dot operations to complete. |

tlx.local_store(dst, src) |

Store data to shared memory. |

tlx.local_load(src) |

Load data from shared memory to registers. |

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.