ByteDance-Seed/Triton-distributed

Triton-distributed – 연산과 통신을 오버랩하기 위한 분산 컴파일러

개요 – Triton-distributed는 OpenAI Triton 언어 기반의 오픈 소스 컴파일러입니다. 무거운 연산(예: GEMM, MoE, flash-decode)과 통신 프리미티브(AllGather, AllReduce, AllToAll 등)를 결합한 GPU 커널을 작성하여 멀티 GPU 시스템에서 두 작업이 동시에 실행되도록 합니다. 이 프로젝트는 NVIDIA (SM80/SM89/SM90a) 및 AMD CDNA3 GPU를 타겟으로 하며, 다른 하드웨어로 확장 가능합니다.

중요성 – 현대의 대규모 모델(LLM, Mixture-of-Experts, 텐서 병렬 추론)은 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 소스에 대한 아웃오브트리(out-of-tree) 플러그인으로 빌드되며, 구형 하드웨어를 위한 레거시 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에서 직접 가져온 것이며, 외부의 추측은 추가되지 않았습니다.