您可以通过以下渠道进一步了解我们👇

Triton-distributed

原版 Triton README | 中文版 README

Triton-distributed 是一款基于 OpenAI Triton 的分布式编译器，专为计算与通信重叠设计。

借助 Triton-distributed，程序员可以开发出媲美高度优化库的高效内核（包括 Distributed-GEMM 和 FLUX）。目前，Triton-distributed 主要面向 Nvidia GPU 和 AMD GPU，未来也可移植到其他硬件平台。如果您希望在自己的硬件上使用 Triton-distributed，请随时联系我们。

新闻

• 2025年12月22日 ✨✨✨：更新了 EP 函数，支持低延迟模式、令牌节省和 Mega-EP。
• 2025年10月21日 🔥🔥🔥：Triton-distributed 在 2025 年 Triton 开发者大会 上进行了展示，详情请参阅 演讲视频。
• 2025年3月9日 ✨✨✨：推出了内核级性能分析器，详情请参阅 文档。
• 2025年8月24日 ⚡⚡⚡：支持对 ByteDance-Seed/Seed-OSS-36B-Instruct 进行推理加速，速度提升至 1.33 倍。
• 2025年8月13日 ✨✨✨：引入了 MegaTritonKernel，并在 H20/H800 上提供了 Qwen3 TP 示例，详情请参阅 文档。
• 2025年8月6日 ✨✨✨：支持在 H800 上进行 GEMM+AllReduce 操作，并在 L20 上支持 MoE 算子，详细信息请参见 GEMM+AR 测试 和 MOE 测试。
• 2025年7月24日 🤖🤖🤖：推出了端到端推理加速演示，统一支持 NVIDIA 和 AMD GPU，详情请参阅 文档。
• 2025年7月11日 ✨✨✨：利用 Triton-distributed 实现了快速 AllReduce，详情请参见 AllReduce 测试。
• 2025年7月11日 ✨✨✨：改进了张量并行下的 MoE 算子，详情请参见 AG+MoE 测试 和 MoE+RS 测试。
• 2025年7月11日 ✨✨✨：支持 Triton 3.4，并集成 NVSHMEM4py（MR）。无需修改 NVSHMEM 代码即可通过 pip install 安装。
• 2025年5月12日 🚀🚀🚀：我们的论文《TileLink：利用以块为中心的原语生成高效的计算-通信重叠内核》已被 MLSys 2025 接受。

入门指南

安装 Triton-distributed

方法一：从源码编译

请参阅 从源码构建。

方法二：使用 pip 安装

准备 PyTorch 容器

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

安装依赖项

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

然后，使用 pip 安装 Triton-distributed。

# 移除随 torch 安装的 Triton
pip uninstall triton
pip uninstall triton_dist # 移除之前的 Triton-distributed
rm -rf /usr/local/lib/python3.12/dist-packages/triton
# 安装 Triton-distributed
VERSION=v0.0.2 # 使用最新版本
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

如何使用 Triton-distributed

Triton-distributed 提供了一组易于使用的原语，以支持分布式计算与通信重叠内核的开发。这些原语分为低层原语和高层原语。目前，我们已发布了低层原语，并计划在未来发布高层原语。

Triton-distributed 原语

借助这些原语，用户可以轻松地编写通信内核。例如，下面展示了一个低延迟的 AllToAll 操作（其延迟优于用于推理的 DeepEP）。 在 32 张 H800 GPU 上，该示例的性能为 137 微秒（每张卡 128 个标记，topk=8，隐藏层大小为 7168，数据类型为 fp8），而 DeepEP 的性能则为 182 微秒（注意：DeepEP 在推理时不使用 NVLink）。

@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,  # 现在我们使用 `int32` 来表示分割信息
        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,
        )

此外，用户还可以将通信部分与计算部分结合，设计出重叠内核。我们在 python/triton_dist/kernels 中提供了示例实现。

性能

Triton-distributed 的性能可以达到与手工优化库相当甚至更好的水平。

单节点 H800x8 上的 AllGather GEMM

单节点 H800x8 上的 GEMM ReduceScatter

两节点 H800x8 上的 AllGather GEMM

两节点 H800x8 上的 GEMM ReduceScatter

分布式 Flash-Decode 从 1 张 GPU 扩展到 32 张 GPU 的性能

解码时的批次大小为 1（单个查询）。

其他平台上的性能

AMD GPU

路线图

功能

• 发布低层原语
• 发布高层原语
• 教程
• 预编译二进制文件

内核

• 发布单节点 GEMM TP 重叠内核
• 发布单节点 MoE TP 重叠内核
• 发布单节点分布式 Flash-Decoding 内核
• 发布单节点 MoE EP 重叠内核
• 发布跨节点 GEMM TP 重叠内核
• 发布跨节点 MoE TP 重叠内核
• 发布跨节点分布式 Flash-Decoding 内核
• 发布跨节点 EP 全对全通信内核（类似于 DeepEP）
• 提供内核实现教程

后端

计算

• 支持 Nvidia SM90a
• 支持 Nvidia SM80
• 支持 Nvidia SM89
• 支持 AMD CDNA3

通信

• NVLink
• IB
• PCIe

性能

• 性能报告

许可证

Triton-distributed 项目采用 MIT 许可证。我们的部分代码采用 Apache-2.0 许可证：

• python/triton_dist/kernels/flash_decode.py

引用

如果您在科学出版物中使用 Triton-distributed，我们鼓励您在相关论文中添加以下参考文献：

@misc{zheng2025tritondistributed,
      title={Triton-distributed: 使用 Triton 编译器在分布式 AI 系统上编程重叠内核},
      author={Size Zheng、Wenlei Bao、Qi Hou、Xuegui Zheng、Jin Fang、Chenhui Huang、Tianqi Li、Haojie Duanmu、Renze Chen、Ruifan Xu、Yifan Guo、Ningxin Zheng、Ziheng Jiang、Xinyi Di、Dongyang Wang、Jianxi Ye、Haibin Lin、Li-Wen Chang、Liqiang Lu、Yun Liang、Jidong Zhai、Xin Liu},
      year={2025},
      eprint={2504.19442},
      archivePrefix={arXiv},
      primaryClass={cs.DC},
      url={https://arxiv.org/abs/2504.19442},
}

@article{zheng2025tilelink,
  title={Tilelink: 使用以 Tile 为中心的原语生成高效的计算-通信重叠内核},
  author={Zheng, Size、Fang, Jin、Zheng, Xuegui、Hou, Qi、Bao, Wenlei、Zheng, Ningxin、Jiang, Ziheng、Wang, Dongyang、Ye, Jianxi、Lin, Haibin 等},
  journal={arXiv 预印本 arXiv:2503.20313},
  year={2025}
}

关于 字节跳动 Seed 团队

字节跳动 Seed 团队成立于 2023 年，致力于打造业界最先进的 AI 基础模型。该团队立志成为世界一流的研究团队，为科学和社会的进步做出重大贡献。

讨论与贡献

请使用议题或拉取请求进行讨论和贡献（参见 CONTRIBUTING.md）。

Triton-distributed 快速上手指南

Triton-distributed 是由字节跳动 Seed 团队开发的分布式编译器，基于 OpenAI Triton 构建。它专为计算与通信重叠（Compute-Communication Overlapping）设计，支持开发者编写高效的内核，性能可媲美高度优化的手工调优库（如 Distributed-GEMM、FLUX）。目前主要支持 NVIDIA GPU 和 AMD GPU。

环境准备

系统要求

• 操作系统: Linux (推荐 Ubuntu)
• GPU:
  • NVIDIA: 支持 SM80, SM89, SM90a 架构 (如 A100, H800, H100 等)
  • AMD: 支持 CDNA3 架构
• 容器环境: 推荐使用官方 PyTorch Docker 镜像

前置依赖

在开始安装前，请确保已安装以下基础软件：

• Docker (用于容器化部署)
• NVIDIA Driver 及 CUDA Toolkit (若宿主机直接安装)
• Python 3.12 (推荐在容器中运行)

安装步骤

推荐使用 Docker 容器 + pip 的方式进行安装，以避免环境冲突。

1. 启动 PyTorch 容器

拉取并运行包含必要驱动环境的 PyTorch 容器：

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

2. 安装依赖库

在容器内安装 cuda.core, nvshmem 等关键依赖：

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

3. 安装 Triton-distributed

首先移除可能随 PyTorch 预装的旧版 Triton，然后安装最新版本的 Triton-distributed：

# 卸载冲突的旧版本
pip uninstall triton
pip uninstall triton_dist
rm -rf /usr/local/lib/python3.12/dist-packages/triton

# 安装 Triton-distributed (当前版本 v0.0.2)
VERSION=v0.0.2
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

注意: 如果网络访问 GitHub 较慢，可尝试配置国内 pip 镜像源（如 pip install -i https://pypi.tuna.tsinghua.edu.cn/simple ...），但 wheel 文件链接需保持原样或下载后本地安装。

基本使用

Triton-distributed 提供了一套底层原语（Primitives），用于开发分布式的计算通信重叠内核。以下是一个低延迟 AllToAll 通信内核的最小化示例，展示了如何使用 @triton_dist.jit 装饰器和 libshmem_device 进行显存直接通信。

import triton
import triton.language as tl
import triton_dist

@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)
    # 获取线程索引 (假设 tid 函数已在上下文中定义或通过 triton 原生方式获取)
    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

    # 执行非阻塞内存复制 (Put)
    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,
        pid,
    )
    
    # 可选：处理 Scale 数据并发送信号
    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,
        )

使用说明：

1. 将上述代码保存为 .py 文件。
2. 在多卡环境下（需配置好 NCCL/NVSHMEM 环境），通过多进程启动该脚本。
3. 更多高级用法（如结合 GEMM 计算的重叠内核）请参考项目目录下的 python/triton_dist/kernels 示例。