Triton 3.7, 플러그인 확장으로 TLX를 업스트림에서 지원

도입

PyTorch Blog가 소개한 PyTorch-Triton 3.7의 핵심은 단일 커널 최적화가 아니라 Triton을 확장하는 방식의 변화다. 새 Triton Plugin Extensions는 개발자가 업스트림 Triton을 수정하거나 다시 빌드하지 않고도 커스텀 컴파일러 Pass, MLIR 방언, DSL 확장을 런타임에 로드할 수 있게 한다. 그동안 자체 Triton 포크를 유지해 온 팀에게는 개발 속도와 유지보수 측면에서 중요한 변화다.

핵심 내용

• 포크 유지 부담 완화: 고성능 GPU 커널은 하드웨어 특화 명령, 맞춤형 lowering, 세밀한 메모리 스케줄링을 요구하는 경우가 많다. 기존에는 이를 위해 Triton 포크를 관리해야 했고, 업스트림 변경을 따라가는 비용이 컸다.
• 런타임 플러그인 로딩: 플러그인은 .so 공유 라이브러리로 배포되며 TRITON_PLUGIN_PATHS 환경 변수로 경로를 지정하면 Triton이 이를 발견해 로드한다. Triton 자체를 다시 컴파일할 필요가 없다.
• 컴파일 파이프라인 제어: 백엔드 컴파일 단계에 추가된 훅을 통해 TTIR, TTGIR, LLVM IR, PTX 또는 AMDGCN 출력에 이르는 여러 단계에 개입할 수 있다. Pass 추가, 비활성화, 교체는 물론 특정 단계나 전체 파이프라인 재정의도 가능하다.
• 방언과 DSL 확장: 단일 변환 Pass뿐 아니라 별도로 컴파일된 MLIR 방언, Python 레벨에서 노출되는 새로운 DSL 연산까지 지원한다.
• 커널 단위 활성화: 플러그인은 커널 수준에서 켜고 끌 수 있으며, 필요한 경우에만 재컴파일되도록 캐시 해시 전략을 플러그인 측에서 관리한다.

TLX가 중요한 이유

Meta의 Triton Language Extensions, 즉 TLX는 이 시스템의 첫 주요 활용 사례다. TLX는 공유 메모리 명시적 할당, 전역 메모리에서 공유 메모리로의 비동기 로드, 비동기 행렬 곱 누산, 로컬 로드와 스토어, 소프트웨어 파이프라이닝을 위한 연산을 제공한다. 이는 최신 GPU를 충분히 활용하는 지속형 GEMM 커널을 작성할 때 필요한 저수준 제어다.

이전에는 TLX를 사용하려면 Meta의 실험적 Triton 포크가 필요했다. 이제는 utlx 독립 패키지로 배포되며 수정되지 않은 업스트림 Triton과 함께 사용할 수 있다. NVIDIA H100에서는 Hopper의 TMA 비동기 로드와 WGMMA 명령으로 매핑되고, AMD MI350에서도 같은 프로그래밍 모델을 해당 백엔드 기능에 연결한다. 블로그는 플러그인 경로가 기존 포크에 컴파일해 넣는 방식과 동일한 코드 생성 경로를 제공한다고 설명한다.

의미와 영향

이번 변화는 Triton 코어를 안정적으로 유지하면서도 연구자, 하드웨어 업체, 프로덕션 커널 팀이 새로운 최적화를 더 빠르게 실험할 수 있게 만든다. 새로운 Pass, 방언, DSL 추상화는 메인 저장소에 병합되기 전에도 배포하고 사용할 수 있으며, 오래된 포크에 갇힐 필요가 줄어든다. 결과적으로 Triton 생태계는 더 모듈화되고, 빠르게 변하는 GPU 하드웨어와 워크로드에 적응하기 쉬워진다.

출처: PyTorch Blog