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

PyTorch-Triton 3.7 introduces the Triton Plugin Extensions system, enabling dynamic loading of custom compiler passes and DSL extensions without forking or recompiling. Meta's Triton Language Extensions (TLX) are now supported out of the box, delivering persistent GEMM kernels that match or exceed vendor library performance on NVIDIA H100 and AMD MI350 GPUs.

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 https://github.com/facebookexperimental/triton 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: python @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: python @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 https://stormy-sailor-96a.notion.site/GPU-MODE-Mini-Competition-3-AlphaFold-s-Triangle-Multiplicative-Update-207221cc2ffa8034b3eddff1d898dc14 – 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: python 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 https://colab.research.google.com/drive/1zANW7SP8dG9I7SXvWkH 9pYFRZYPAu y?usp=sharing - H100 Standalone Script: GitHub Gist https://gist.github.com/CRobeck/daec7724bd3fd1ef2b38af7024032bc4 - AMD MI350 Pipelined GEMM: GitHub Gist https://gist.github.com/CRobeck/6cfe9a32da9cdb6f446c8214a53d2293 - Plugin Documentation: triton/lib/Plugins/README.md https://github.com/triton-lang/triton/blob/main/examples/plugins/README.md - Extension Repository: triton-lang/triton-ext https://github.com/triton-lang/triton-ext - utlx PyPI Project: https://pypi.org/project/triton-utlx/ 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 https://github.com/triton-lang/triton-ext and the plugin documentation https://github.com/triton-lang/triton/blob/main/examples/plugins/README.md .