You can get to know us better through the following channels👇

Triton-distributed

Original Triton README | README in Chinese

Triton-distributed is a distributed compiler designed for computation-communication overlapping, which is based on OpenAI Triton.

Using Triton-distributed, programmers are able to develop efficient kernels comparable to highly-optimized libraries (including Distributed-GEMM and FLUX). Triton-distributed currently mainly targets Nvidia GPU and AMD GPU. It can also be ported to other hardware platforms. Feel free to contact us if you want to use Triton-distributed on your own hardware.

News

• 12/22/2025 ✨✨✨: Updated EP functions, support low-latency mode, token saving, and Mega-EP.
• 21/10/2025 🔥🔥🔥: Triton-distributed is presented at Triton Conference 2025, see the talk for details.
• 09/03/2025 ✨✨✨: Introduced Intra-Kernel Profiler, See the doc for details.
• 08/24/2025 ⚡⚡⚡: Support inference acceleration for ByteDance-Seed/Seed-OSS-36B-Instruct, achieving a 1.33x speedup.
• 08/13/2025 ✨✨✨: Introduced the MegaTritonKernel and provided a Qwen3 TP demo on H20/H800, See the doc for details.
• 08/06/2025 ✨✨✨: Support GEMM+AllReduce on H800 and support MoE operators on L20, see GEMM+AR Test and MOE Test for detail.
• 07/24/2025 🤖🤖🤖: Introduced end-to-end inference acceleration demo with unified support for both NVIDIA and AMD GPUs. See the doc for details.
• 07/11/2025 ✨✨✨: Fast AllReduce implemented with Triton-distributed, see AllReduce Test.
• 07/11/2025 ✨✨✨: Improved MoE operators for tensor parallel. See AG+MoE Test and MoE+RS Test.
• 07/11/2025 ✨✨✨: Triton 3.4 support with NVSHMEM4py (MR). pip install is also supported without any need to modify NVSHMEM code.
• 05/12/2025 🚀🚀🚀: Our paper TileLink: Generating Efficient Compute-Communication Overlapping Kernels using Tile-Centric Primitives accepted by MLSys 2025.

Getting started

Install Triton-distributed

Method 1. From source

See build from source.

Method 2. Using pip

Prepare PyTorch container

docker run --name triton-dist --ipc=host --network=host --privileged --cap-add=SYS_ADMIN --shm-size=10g --gpus=all -itd nvcr.io/nvidia/pytorch:25.04-py3 /bin/bash
docker exec -it triton-dist /bin/bash

Install Dependencies

pip3 install cuda.core==0.2.0 nvidia-nvshmem-cu12==3.3.9 Cython==0.29.24 nvshmem4py-cu12==0.1.2
pip3 install cuda-python==12.4 setuptools==69.0.0 wheel pybind11

Then, pip install triton-dist.

# Remove triton installed with torch
pip uninstall triton
pip uninstall triton_dist # remove previous triton-dist
rm -rf /usr/local/lib/python3.12/dist-packages/triton
# Install Triton-distributed
VERSION=v0.0.2 # use the latest version
pip install https://github.com/ByteDance-Seed/Triton-distributed/releases/download/${VERSION}/triton_dist-3.4.0-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl

How to use Triton-distributed

Triton-distributed provides a set of easy-to use primitives to support the development of distributed compute-communication overlapping kernels. The primitives are divided into low-level primitives and high-level primitives. Currently, we have released our low-level primitives, and we plan to release high-level primitives in future.

Triton-distributed Primitives

Using these primitives, users can program communication kernels easily. For example, a low-latency AllToAll (with better latency than DeepEP for inference) is shown below. The performance of this example on 32 H800 GPUs is 137us (128 tokens per rank, topk=8, hidden_size=7168, dtype=fp8), while DeepEP is 182 us (note: DeepEP doesn't use NVLink for inference).

@triton_dist.jit
def all_to_all_kernel(
    data_src,
    data_dst,
    splits_src,
    splits_dst,
    signal,
    splits_cumsum,
    scale_src,
    scale_dst,
    rank: int,
    call_count: int,
    WITH_SCALE: tl.constexpr,
    WORLD_SIZE: tl.constexpr,
    HIDDEN: tl.constexpr,
    MAX_M: tl.constexpr,
    EXPERTS_PER_RANK: tl.constexpr,
    NUM_TOT_EXPERTS: tl.constexpr,
    ELEMENT_SIZE: tl.constexpr = 2,
    SCALE_ELEMENT_SIZE: tl.constexpr = 4,
):
    pid = tl.program_id(0)
    threadidx = tid(axis=0)

    exp_st = pid * EXPERTS_PER_RANK
    exp_ed = exp_st + EXPERTS_PER_RANK

    m_st = tl.load(splits_cumsum + exp_st)
    m_ed = tl.load(splits_cumsum + exp_ed)
    num_rows_cur_block = m_ed - m_st

    src_off = m_st
    dst_off = rank * MAX_M

    split_src_ptr = splits_src + exp_st
    off0 = exp_st + tl.arange(0, EXPERTS_PER_RANK)
    off1 = exp_st + tl.arange(0, EXPERTS_PER_RANK) + 1
    cumsum_sts = tl.load(splits_cumsum + off0)
    cumsum_eds = tl.load(splits_cumsum + off1)
    tl.store(split_src_ptr + tl.arange(0, EXPERTS_PER_RANK), cumsum_eds - cumsum_sts)

    act_pos = call_count % 2
    data_dst_ptr = data_dst + act_pos * WORLD_SIZE * MAX_M * HIDDEN + dst_off * HIDDEN
    split_dst_ptr = splits_dst + act_pos * NUM_TOT_EXPERTS + rank * EXPERTS_PER_RANK
    signal_ptr = signal + act_pos * WORLD_SIZE + rank

    libshmem_device.putmem_nbi_block(
        data_dst_ptr,
        data_src + src_off * HIDDEN,
        num_rows_cur_block * HIDDEN * ELEMENT_SIZE,
        pid,
    )
    libshmem_device.putmem_nbi_block(
        split_dst_ptr,
        split_src_ptr,
        EXPERTS_PER_RANK * 4,  # now we use `int32` for splits
        pid,
    )
    if WITH_SCALE:
        scale_dst_ptr = scale_dst + act_pos * WORLD_SIZE * MAX_M + dst_off
        libshmem_device.putmem_signal_nbi_block(
            scale_dst_ptr,
            scale_src + src_off,
            num_rows_cur_block * SCALE_ELEMENT_SIZE,
            signal_ptr,
            call_count,
            libshmem_device.NVSHMEM_SIGNAL_SET,
            pid,
        )

    libshmem_device.fence()
    if threadidx == 0:
        if not WITH_SCALE:
            libshmem_device.signal_op(
                signal_ptr,
                call_count,
                libshmem_device.NVSHMEM_SIGNAL_SET,
                pid,
            )
        libshmem_device.signal_wait_until(
            signal + act_pos * WORLD_SIZE + pid,
            libshmem_device.NVSHMEM_CMP_EQ,
            call_count,
        )

Also, users can combine the communication part with computation part to design overlapping kernels. We have provided example implementations in python/triton_dist/kernels.

Performance

Triton-distributed can achieve comparable or better performance than hand-tuned libraries.

AllGather GEMM on single node of H800x8

GEMM ReduceScatter on single node of H800x8

AllGather GEMM on 2 nodes of H800x8

GEMM ReduceScatter on 2 nodes of H800x8

Scaling of Distributed Flash-Decode from 1 GPU to 32 GPUs

The batch size is 1 (one query) for decoding.

Performance on Other Platforms

AMD GPUs

Roadmaps

Functionalities

• Release low-level primitives
• Release high-level primitives
• Tutorials
• Pre-built binary

Kernels

• Release single-node GEMM TP overlapping kernels
• Release single-node MoE TP overlapping kernels
• Release single-node distributed Flash-Decoding kernels
• Release single-node MoE EP overlapping kernels
• Release cross-node GEMM TP overlapping kernels
• Release cross-node MoE TP overlapping kernels
• Release cross-node distributed Flash-Decoding kernels
• Release cross-node EP all-to-all kernels (similar to DeepEP)
• Provide tutorials for kernel implementation

Backends

Computation

• Nvidia SM90a support
• Nvidia SM80 support
• Nvidia SM89 support
• AMD CDNA3 support

Communication

• NVLink
• IB
• PCIe

Performance

• Performance report

License

The Triton-distributed project is under MIT license. Part of our code is under Apache-2.0 License:

• python/triton_dist/kernels/flash_decode.py

Citation

If you use Triton-distributed in a scientific publication, we encourage you to add the following reference to the related papers:

@misc{zheng2025tritondistributed,
      title={Triton-distributed: Programming Overlapping Kernels on Distributed AI Systems with the Triton Compiler},
      author={Size Zheng and Wenlei Bao and Qi Hou and Xuegui Zheng and Jin Fang and Chenhui Huang and Tianqi Li and Haojie Duanmu and Renze Chen and Ruifan Xu and Yifan Guo and Ningxin Zheng and Ziheng Jiang and Xinyi Di and Dongyang Wang and Jianxi Ye and Haibin Lin and Li-Wen Chang and Liqiang Lu and Yun Liang and Jidong Zhai and Xin Liu},
      year={2025},
      eprint={2504.19442},
      archivePrefix={arXiv},
      primaryClass={cs.DC},
      url={https://arxiv.org/abs/2504.19442},
}

@article{zheng2025tilelink,
  title={Tilelink: Generating efficient compute-communication overlapping kernels using tile-centric primitives},
  author={Zheng, Size and Fang, Jin and Zheng, Xuegui and Hou, Qi and Bao, Wenlei and Zheng, Ningxin and Jiang, Ziheng and Wang, Dongyang and Ye, Jianxi and Lin, Haibin and others},
  journal={arXiv preprint arXiv:2503.20313},
  year={2025}
}

About ByteDance Seed Team

Founded in 2023, ByteDance Seed Team is dedicated to crafting the industry's most advanced AI foundation models. The team aspires to become a world-class research team and make significant contributions to the advancement of science and society.

Discussion and Contribution

Please use issues or pull requests for discussion and contribution (see CONTRIBUTING.md).