ByteDance-Seed/Triton-distributed

Triton‑distributed – a distributed compiler for overlapping compute & communication

What it is – Triton‑distributed is an open‑source compiler built on top of the OpenAI Triton language. It lets you write GPU kernels that combine heavy computation (e.g. GEMM, MoE, flash‑decode) with communication primitives (AllGather, AllReduce, AllToAll, etc.) so the two can run simultaneously on multi‑GPU systems. The project targets NVIDIA (SM80/SM89/SM90a) and AMD CDNA3 GPUs and can be extended to other hardware.

Why it matters – Modern large‑scale models (LLMs, mixture‑of‑experts, tensor‑parallel inference) are limited by the time spent moving tensors between GPUs. By exposing low‑level “overlap” primitives, Triton‑distributed lets developers write kernels that hide communication latency, achieving speed‑ups that rival or surpass hand‑tuned libraries such as NVIDIA CUTLASS or DeepEP.

Key features (as of the README)

• Low‑level primitives for distributed memory operations (put/get, signal‑based sync, NVSHMEM, RDMA, PCIe, IB, NVLink).
• High‑performance kernels already shipped: distributed GEMM + AllReduce, MoE tensor‑parallel kernels, flash‑decode kernels, and an AllToAll kernel that beats DeepEP on H800 GPUs.
• Cross‑hardware support – NVIDIA GPUs (SM80/89/90a) and AMD CDNA3; Ascend support via a separate branch.
• Plugin architecture – built as an out‑of‑tree plugin against the upstream Triton 3.7 source, with a legacy Triton 3.4 branch for older hardware.
• Performance reporting – includes intra‑node and inter‑node scaling plots for GEMM, AllGather, ReduceScatter, and flash‑decode.
• Roadmap – high‑level primitives, more tutorials, additional cross‑node kernels, and broader backend support are planned.

Getting started

1. From source – follow docs/build.md to compile the plugin against Triton 3.7.
2. Pre‑built wheel – in a PyTorch container, install required CUDA/NVSHMEM wheels, then pip install https://github.com/ByteDance-Seed/Triton-distributed/releases/download/v0.0.2/triton_dist-3.4.0‑cp312‑cp312‑manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl.
3. Write a kernel – decorate a function with @triton_dist.jit and call the provided libshmem_device APIs (e.g., putmem_nbi_block, signal_wait_until). The README includes a full AllToAll example.

Typical use‑case

• Large‑scale inference – accelerate token‑wise decoding on 8‑32 GPUs by overlapping the attention/FFN compute with the all‑to‑all communication required for tensor‑parallel MoE or flash‑decode.
• Training kernels – implement custom GEMM + AllReduce or MoE‑specific reductions that hide the cost of gradient synchronization.

Maturity – The project is actively maintained (news entries up to Sep 2026). It ships a stable low‑level API, pre‑built binaries, and several benchmarked kernels. High‑level abstractions are still in progress, and some backends (e.g., Ascend on Triton 3.7) are pending.

License – MIT for most code; a subset (flash_decode.py) is Apache‑2.0.

---

All statements are taken directly from the repository’s README; no external speculation is added.