{"slug": "triton-plugin-extensions-enabling-tlx-and-custom-compiler-passes-out-of-the-box", "title": "Triton Plugin Extensions: Enabling TLX and Custom Compiler Passes Out of the Box", "summary": "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.", "body_md": "### Featured projects\n\n**TLDR**\n\nThe 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.\n\n## The Problem: Why Extensions?\n\nWriting 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.\n\nForks 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.\n\nWhat’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.\n\n## The Triton Plugin Extensions System\n\nThe 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.\n\n### Overridable Compiler Pipeline\n\nAt 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:\n\n- Insert one or more custom passes at arbitrary points in any stage.\n- Disable specific passes within a stage.\n- Replace existing passes with specialized custom implementations (e.g., a custom warp specialization strategy).\n- Override entire stages or the full pipeline.\n\nThis is available on both the NVIDIA and AMD backends\n\n### Custom Ops, Dialects, and Lowering\n\nThe plugin API is designed to complement PyBind11, enabling three levels of extensibility:\n\n- Custom transformation passes: single passes that can be inserted at arbitrary points in the pipeline without an associated dialect.\n- 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.\n- Custom top-level DSL ops: new Python-level syntax and semantics enabling entirely new programming abstractions without altering Triton itself.\n\n### Per-Kernel Control\n\nPlugins 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.\n\n```\n# Enabling the TLX plugin is as simple as setting an environment variable\nimport os\nimport sysconfig\n\ndist_packages = sysconfig.get_paths()[\"purelib\"]\nlibutlx_path = os.path.join(dist_packages, \"utlx_plugin\", \"libutlx.so\")\nos.environ[\"TRITON_PLUGIN_PATHS\"] = libutlx_path\n```\n\n## TLX: Triton Language Extensions, Now Built-In\n\n[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.\n\nThe core TLX operations include:\n\n| Operation | Description |\n|---|---|\n`tlx.local_alloc(shape, dtype, num_buffers)` |\nAllocate shared memory buffers for software pipelining. |\n`tlx.local_view(buffers, index)` |\nView a specific buffer within an allocation. |\n`tlx.async_load(src, dst, mask)` |\nInitiate an asynchronous load from global to shared memory. |\n`tlx.async_load_commit_group(tokens)` |\nCommit a group of async loads. |\n`tlx.async_load_wait_group(n)` |\nWait for async load groups to complete. |\n`tlx.async_dot(a, b, acc)` |\nAsynchronous matrix multiply-accumulate. |\n`tlx.async_dot_wait(n, acc)` |\nWait for async dot operations to complete. |\n`tlx.local_store(dst, src)` |\nStore data to shared memory. |\n`tlx.local_load(src)` |\nLoad data from shared memory to registers. |\n\nPreviously, 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`\n\n) that works with unmodified upstream Triton. Starting with PyTorch-Triton 3.7, TLX will be enabled by default on all Triton releases going forward.\n\n## Cross-Hardware: NVIDIA H100 and AMD MI350\n\nOne 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.\n\n### NVIDIA H100 (Hopper) — Persistent GEMM\n\nOn 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:\n\n``` python\n@triton.jit\ndef matmul_kernel_pipelined_hopper(\n    a_ptr, b_ptr, c_ptr, M, N, K,\n    stride_am, stride_ak, stride_bk, stride_bn, stride_cm, stride_cn,\n    BLOCK_SIZE_M: tl.constexpr, BLOCK_SIZE_N: tl.constexpr,\n    BLOCK_SIZE_K: tl.constexpr, GROUP_SIZE_M: tl.constexpr,\n    NUM_STAGES: tl.constexpr,\n):\n    # ... tile indexing ...\n\n    # Allocate multi-stage shared memory buffers\n    buffers_A = tlx.local_alloc((BLOCK_SIZE_M, BLOCK_SIZE_K), tlx.dtype_of(a_ptr), NUM_STAGES)\n    buffers_B = tlx.local_alloc((BLOCK_SIZE_K, BLOCK_SIZE_N), tlx.dtype_of(b_ptr), NUM_STAGES)\n\n    # Prefetch pipeline prologue\n    for i in tl.range(0, NUM_STAGES - 1, loop_unroll_factor=NUM_STAGES - 1):\n        a = tlx.local_view(buffers_A, i)\n        b = tlx.local_view(buffers_B, i)\n        token_a = tlx.async_load(a_ptrs, a, mask=...)\n        token_b = tlx.async_load(b_ptrs, b, mask=...)\n        tlx.async_load_commit_group([token_a, token_b])\n\n    # Main K loop with overlapped compute and data movement\n    acc = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)\n    for k in tl.range(0, tl.cdiv(K, BLOCK_SIZE_K), num_stages=0):\n        buf = k % NUM_STAGES\n        tlx.async_load_wait_group(NUM_STAGES - 2)\n        acc = tlx.async_dot(\n            tlx.local_view(buffers_A, buf),\n            tlx.local_view(buffers_B, buf), \n            acc\n        )\n        # Prefetch next stage ...\n\n    acc = tlx.async_dot_wait(0, acc)\n    # Store results ...\n```\n\n### NVIDIA H100 (Hopper) — Performance Results\n\nThe 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.\n\nOn 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:\n\n128×13312×16384: cuBLAS 247.8 TFLOPS → Triton+TLX 257.0 TFLOPS (+3.7%)\n\n16384×8192×8192: cuBLAS 549.4 TFLOPS → Triton+TLX 566.7 TFLOPS (+3.2%)\n\n8192×16384×8192: cuBLAS 564.8 TFLOPS → Triton+TLX 575.9 TFLOPS (+2.0%)\n\n8192×53248×8192: cuBLAS 571.3 TFLOPS → Triton+TLX 573.2 TFLOPS (+0.3%)\n\n8192×28672×4096: cuBLAS 560.4 TFLOPS → Triton+TLX 559.8 TFLOPS (−0.1%)\n\n8192×8192×8192: cuBLAS 582.3 TFLOPS → Triton+TLX 577.0 TFLOPS (−0.9%)\n\n### AMD MI350 — Pipelined GEMM\n\nOn 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:\n\n``` python\n@triton.jit\ndef matmul_kernel_pipelined_mi300(\n    a_ptr, b_ptr, c_ptr, M, N, K, ...\n):\n    # ... tile indexing ...\n\n    # Allocate shared memory buffers\n    buffers_A = tlx.local_alloc((BLOCK_SIZE_M, BLOCK_SIZE_K), tlx.dtype_of(a_ptr), NUM_STAGES - 1)\n    buffers_B = tlx.local_alloc((BLOCK_SIZE_K, BLOCK_SIZE_N), tlx.dtype_of(b_ptr), NUM_STAGES - 1)\n\n    # Prologue: load into shared memory via registers\n    for i in tl.range(0, NUM_STAGES - 1, loop_unroll_factor=NUM_STAGES - 1):\n        a_reg = tl.load(a_ptrs, mask=...)\n        b_reg = tl.load(b_ptrs, mask=...)\n        tlx.local_store(tlx.local_view(buffers_A, i), a_reg)\n        tlx.local_store(tlx.local_view(buffers_B, i), b_reg)\n\n    # Main loop: overlapped compute and memory operations\n    acc = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)\n    for k in tl.range(NUM_STAGES - 1, K_ITERS, num_stages=0):\n        # Load next tile into registers\n        a_reg = tl.load(a_ptrs, mask=...)\n        b_reg = tl.load(b_ptrs, mask=...)\n        \n        # Compute on previously staged data\n        a_prev = tlx.local_load(tlx.local_view(buffers_A, buf))\n        b_prev = tlx.local_load(tlx.local_view(buffers_B, buf))\n        acc = tl.dot(a_prev, b_prev, acc)\n        \n        # Store new data to shared memory for next iteration\n        tlx.local_store(tlx.local_view(buffers_A, ...), a_reg)\n        tlx.local_store(tlx.local_view(buffers_B, ...), b_reg)\n```\n\n### AMD MI350 — Performance Results\n\nThe 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.\n\nTriton + 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:\n\n256×256×256: rocBLAS 4.4 TFLOPS → Triton+TLX 5.0 TFLOPS (+11.8%)\n\n512×512×512: rocBLAS 29.4 TFLOPS → Triton+TLX 33.9 TFLOPS (+15.2%)\n\n1024×1024×1024: rocBLAS 161.2 TFLOPS → Triton+TLX 180.8 TFLOPS (+12.1%)\n\n2048×2048×2048: rocBLAS 445.1 TFLOPS → Triton+TLX 511.9 TFLOPS (+15.0%)\n\n### GPUMode Trimul Multiplicative Update Validation\n\nWe 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.\n\n## Identical CodeGen, Zero Fork Required\n\nA 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.\n\nOur demos confirm:\n\n- Identical PTX codegen on NVIDIA H100 for persistent GEMM kernels.\n- Identical AMDGCN codegen on AMD MI350 for pipelined GEMM kernels.\n- Equivalent performance — no measurable overhead from the dynamic loading path.\n\nThe 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.\n\n## Getting Started\n\nGetting started with TLX on upstream Triton is straightforward:\n\n```\n# Install Triton (from source) and the TLX extension from PyPI package\ngit clone https://github.com/triton-lang/triton && cd triton\nTRITON_EXT_ENABLED=ON pip install -e . --no-build-isolation && cd ..\n# uTLX plugin (published on PyPI):\npip install triton-utlx\n```\n\nThe utlx package includes the pre-built extension library. Once installed, set the plugin path and import the extensions:\n\n``` python\nimport os\nimport sysconfig\n\n# Point Triton at the TLX plugin\ndist_packages = sysconfig.get_paths()[\"purelib\"]\nos.environ[\"TRITON_PLUGIN_PATHS\"] = os.path.join(dist_packages, \"utlx_plugin\", \"libutlx.so\")\n\n# Now TLX ops are available in your kernels\nimport triton\nimport triton.language as tl\nimport utlx_plugin as tlx\n```\n\nTo explore the full demos and repositories:\n\n- H100 Persistent GEMM Notebook:\n[Colab](https://colab.research.google.com/drive/1zANW7SP8dG9I7SXvWkH_9pYFRZYPAu_y?usp=sharing) - H100 Standalone Script:\n[GitHub Gist](https://gist.github.com/CRobeck/daec7724bd3fd1ef2b38af7024032bc4) - AMD MI350 Pipelined GEMM:\n[GitHub Gist](https://gist.github.com/CRobeck/6cfe9a32da9cdb6f446c8214a53d2293) - Plugin Documentation:\n[triton/lib/Plugins/README.md](https://github.com/triton-lang/triton/blob/main/examples/plugins/README.md) - Extension Repository:\n[triton-lang/triton-ext](https://github.com/triton-lang/triton-ext) - utlx PyPI Project:\n[https://pypi.org/project/triton-utlx/](https://pypi.org/project/triton-utlx/)\n\n## What’s Next\n\nThe plugin extensions system opens the door to a growing ecosystem of community-developed Triton extensions. Areas of active development and future proposals include:\n\n- Custom backends: dynamically loaded out-of-tree backends for Intel, CPU, and other targets without modifying Triton’s build system.\n- triton-distributed: distributed computing primitives as an extension.\n- Customized versions of instrumentation and profiling tools like Proton and ConSan developed as plugins for user specific runtime-loadable performance analysis\n- Custom optimization passes: target-specific warp specialization, loop splitting, and model-specific optimizations shipped as add-ins\n- Specialized ops: 2:4 structured sparsity, custom layout conversions, and more\n\nWe’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).", "url": "https://wpnews.pro/news/triton-plugin-extensions-enabling-tlx-and-custom-compiler-passes-out-of-the-box", "canonical_source": "https://pytorch.org/blog/triton-plugin-extensions-enabling-tlx-and-custom-compiler-passes-out-of-the-box/", "published_at": "2026-07-15 15:09:15+00:00", "updated_at": "2026-07-15 15:35:51.615485+00:00", "lang": "en", "topics": ["developer-tools", "machine-learning", "ai-infrastructure"], "entities": ["PyTorch", "Triton", "Meta", "TLX", "NVIDIA", "AMD", "H100", "MI350"], "alternates": {"html": "https://wpnews.pro/news/triton-plugin-extensions-enabling-tlx-and-custom-compiler-passes-out-of-the-box", "markdown": "https://wpnews.pro/news/triton-plugin-extensions-enabling-tlx-and-custom-compiler-passes-out-of-the-box.md", "text": "https://wpnews.pro/news/triton-plugin-extensions-enabling-tlx-and-custom-compiler-passes-out-of-the-box.txt", "jsonld": "https://wpnews.pro/news/triton-plugin-extensions-enabling-tlx-and-custom-compiler-passes-out-of-the-box.jsonld"}}