# Triton-distributed: Programming Overlapping Kernels on Distributed AI Systems with the Triton Compiler

arXiv: https://arxiv.org/abs/2504.19442 · HTML: https://mlsanity.com/paper/2504.19442/

## Key facts

- Cited by 25 papers in the map (data through 14 Oct 2026).
- First posted on arXiv 28 Apr 2025.

## Built on by (most cited first, 25 of 25)

- [Robix: A Unified Model for Robot Interaction, Reasoning and Planning](https://mlsanity.com/paper/2509.01106/) (2025)
- [MPK: A Compiler and Runtime for Mega-Kernelizing Tensor Programs](https://mlsanity.com/paper/2512.22219/) (2025)
- [TokenWeave: Efficient Compute-Communication Overlap for Distributed LLM Inference](https://mlsanity.com/paper/2505.11329/) (2025)
- [Least-Loaded Expert Parallelism: Load Balancing An Imbalanced Mixture-of-Experts](https://mlsanity.com/paper/2601.17111/) (2026)
- [UniEP: Unified Expert-Parallel MoE MegaKernel for LLM Training](https://mlsanity.com/paper/2604.19241/) (2026)
- [Event Tensor: A Unified Abstraction for Compiling Dynamic Megakernel](https://mlsanity.com/paper/2604.13327/) (2026)
- [Kareus: Joint Reduction of Dynamic and Static Energy in Large Model Training](https://mlsanity.com/paper/2601.17654/) (2026)
- [Axe: A Simple Unified Layout Abstraction for Machine Learning Compilers](https://mlsanity.com/paper/2601.19092/) (2026)
- [ZeCO: Zero Communication Overhead Sequence Parallelism for Linear Attention](https://mlsanity.com/paper/2507.01004/) (2025)
- [Eliminating Multi-GPU Performance Taxes: A Systems Approach to Efficient Distributed LLMs](https://mlsanity.com/paper/2511.02168/) (2025)
- [Revisiting Parameter Server in LLM Post-Training](https://mlsanity.com/paper/2601.19362/) (2026)
- [Optimizing Long-context LLM Serving via Fine-grained Sequence Parallelism](https://mlsanity.com/paper/2511.06247/) (2025)
- [Design Space Exploration of DMA based Finer-Grain Compute Communication Overlap](https://mlsanity.com/paper/2512.10236/) (2025)
- [Stage-Aware Communication Scheduling for Disaggregated LLM Serving](https://mlsanity.com/paper/2603.17456/) (2026)
- [FEPLB: Exploiting Copy Engines for Nearly Free MoE Load Balancing in Distributed Training](https://mlsanity.com/paper/2604.19654/) (2026)
- [Searching for Plans You Can Actually Build: A Realizability-Aware Full-Space Optimizer for MoE Training and Serving](https://mlsanity.com/paper/2607.18631/) (2026)
- [Iris: First-Class Multi-GPU Programming Experience in Triton](https://mlsanity.com/paper/2511.12500/) (2025)
- [Throughput Optimization as a Strategic Lever in Large-Scale AI Systems: Evidence from Dataloader and Memory Profiling Innovations](https://mlsanity.com/paper/2603.26823/) (2026)
- [NanoCP: Request-Level Dynamic Context Parallelism for Data-Expert Parallel Decoding](https://mlsanity.com/paper/2605.21100/) (2026)
- [DynaFlow: Transparent and Flexible Intra-Device Parallelism via Programmable Operator Scheduling](https://mlsanity.com/paper/2605.21603/) (2026)
- [HyperParallel-MoE: Multi-Core Interleaved Scheduling for Fast MoE Training on Ascend NPUs](https://mlsanity.com/paper/2605.23764/) (2026)
- [HetCCL: Enabling Collective Communication For Mixed-Vendor Heterogeneous Clusters](https://mlsanity.com/paper/2605.31000/) (2026)
- [mKernel: Fast Multi-GPU, Multi-Node Fused Kernels](https://mlsanity.com/paper/2609.13585/) (2026)
- [Weave: Fine-Grained Dynamic SM Scheduling in an MoE Megakernel for Compute-Communication Overlap](https://mlsanity.com/paper/2609.21483/) (2026)
- [Mixture-of-Kittens: MoE Megakernel for NVL72s](https://mlsanity.com/paper/2609.36070/) (2026)

## Abstract (arXiv)

In this report, we propose Triton-distributed, an extension of existing Triton compiler, to overcome the programming challenges in distributed AI systems. Triton-distributed is the first compiler that supports native overlapping optimizations for distributed AI workloads, providing a good coverage of existing optimizations from different frameworks. First, we integrate communication primitives compliant with the OpenSHMEM standard into the compiler. This enables programmers to utilize these primitives with a higher-level Python programming model. Second, we illustrate how to achieve complex joint optimization of computation, memory access, and communication with the assistance of the compiler. In particular, we show how to use overlapping techniques to hide latency and present our compiler-based programming methods in both single-node and multi-node scenarios. Finally, we showcase the performance of the code generated by our compiler. In a test environment with up to 64 devices, our compiler can fully utilize heterogeneous communication and computation resources to provide effective overlapping and high performance. In many cases, the performance of the generated code can even outperform hand-optimized code. Moreover, the development difficulty and the time cost for development using our compiler are far less than those of low-level programming such as CUDA/C++, which clearly demonstrates significant productivity advantages.