ByteDance-Seed/Triton-distributed

Distributed Compiler and Optimized Parallel Kernels

Triton-distributed – 計算と通信をオーバーラップさせる分散コンパイラ

概要 – Triton-distributedは、OpenAI Triton言語の上に構築されたオープンソースコンパイラです。重い計算(GEMM、MoE、flash-decodeなど)と通信プリミティブ(AllGather、AllReduce、AllToAllなど)を組み合わせたGPUカーネルを記述し、マルチGPUシステム上で両者を同時に実行できるようにします。このプロジェクトはNVIDIA (SM80/SM89/SM90a) およびAMD CDNA3 GPUをターゲットとしており、他のハードウェアにも拡張可能です。

重要性 – 現代の大規模モデル(LLM、混合エキスパートモデル、テンソル並列推論)は、GPU間のテンソル転送にかかる時間によって制限されています。Triton-distributedは、低レベルの「オーバーラップ」プリミティブを公開することで、通信レイテンシを隠蔽するカーネルの記述を可能にし、NVIDIA CUTLASSやDeepEPといった手動チューニングされたライブラリに匹敵または凌駕する高速化を実現します。

主な機能(READMEより)

  • 低レベルプリミティブ:分散メモリ操作用(put/get、シグナルベースの同期、NVSHMEM、RDMA、PCIe、IB、NVLink)。
  • 高性能カーネル:分散GEMM + AllReduce、MoEテンソル並列カーネル、flash-decodeカーネル、およびH800 GPUでDeepEPを上回るAllToAllカーネルが既に提供されています。
  • クロスハードウェアサポート:NVIDIA GPU (SM80/89/90a) およびAMD CDNA3。Ascendサポートは別ブランチで提供。
  • プラグインアーキテクチャ:アップストリームのTriton 3.7ソースに対するアウトオブツリープラグインとして構築されており、古いハードウェア向けにレガシーなTriton 3.4ブランチも用意されています。
  • パフォーマンスレポート:GEMM、AllGather、ReduceScatter、flash-decodeのノード内およびノード間スケーリングプロットが含まれています。
  • ロードマップ:高レベルプリミティブ、さらなるチュートリアル、追加のクロスノードカーネル、より広範なバックエンドサポートが計画されています。

はじめに

  1. ソースから:docs/build.mdに従い、Triton 3.7に対してプラグインをコンパイルします。
  2. ビルド済みwheel:PyTorchコンテナ内で必要なCUDA/NVSHMEM wheelをインストールし、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. カーネルの記述:関数を@triton_dist.jitでデコレートし、提供されているlibshmem_device API(例:putmem_nbi_block、signal_wait_until)を呼び出します。READMEには完全なAllToAllの例が含まれています。

典型的なユースケース

  • 大規模推論:テンソル並列MoEやflash-decodeに必要なall-to-all通信とattention/FFN計算をオーバーラップさせることで、8〜32 GPUでのトークン単位のデコードを高速化します。
  • トレーニングカーネル:カスタムGEMM + AllReduceやMoE固有の削減を実装し、勾配同期のコストを隠蔽します。

成熟度 – プロジェクトは活発にメンテナンスされています(2026年9月までのニュースエントリ)。安定した低レベルAPI、ビルド済みバイナリ、およびベンチマーク済みのカーネルが提供されています。高レベルの抽象化は現在進行中であり、一部のバックエンド(Triton 3.7上のAscendなど)は保留中です。

ライセンス – ほとんどのコードはMIT。一部(flash_decode.py)はApache-2.0です。


すべての記述はリポジトリのREADMEから直接引用したものであり、外部の推測は含まれていません。

関連

  • Dispatch
  • プロジェクト
  • プロジェクト
  • プロジェクト
  • プロジェクト