Meta, Blackwell GPU용 하드웨어 네이티브 GPU 컴파일러 확장 TLX 공개

• Meta는 고수준 생산성과 저수준 하드웨어 제어 사이의 격차를 메우기 위해 Triton Low-Level Extensions(TLX)를 개발하고, 명시적인 워프 그룹 오케스트레이션을 위한 MIMW(Multiple Instructions, Multiple Warps) 추상화를 도입함
• NVIDIA Blackwell GPU에서 GEM 모델 학습 시 TLX는 O(10^24) 총 학습 flops를 달성하고 20-25%의 end-to-end Model FLOPs Utilization(MFU)를 기록했으며, 생산성과 성능의 균형 측면에서 CUTLASS와 CuTeDSL을 능가함
• TLX 기반 커널은 Generalized Dot Product Attention(GDPA)에서 최대 2배의 forward speedup을 제공하며, In-Kernel Broadcast Optimization(IKBO)을 통해 RecSys 추론에서 compute-intensive net latency를 최대 3분의 2까지 감소시킴
• 대규모 LLM 사전 학습을 위한 distributed grouped GEMM을 지원하며, Mixture-of-Experts(MoE) 모델에서 end-to-end latency를 줄이기 위해 계산과 통신을 통합함

Design and Architecture

• TLX는 Triton을 second execution-oriented layer로 확장하여 계산(Triton tensor ops)과 오케스트레이션(TLX warp-group roles: producer, consumer, epilogue)을 분리함
• Tensor Memory(TMEM), 비동기 메모리 전송(TMA), mbarrier를 통한 warp-group-level synchronization 등 하드웨어 원语에 대한 직접 접근을 제공함
• CUDA나 CUTLASS와 달리 명시적인 layout specification을 요구하지 않아 컴파일러가 기본값을 추론하고 fusion 비용을 줄일 수 있음

Production Deployment and Future Plans

• TLX는 NVIDIA GPU 기반의 RecSys 학습/추론 및 대규모 LLM 사전 학습을 위해 Meta의 ML 인프라 전체에 배포되었으며, NVIDIA Hopper, Blackwell 및 AMD GPU를 대상으로 함
• 향후 개발 계획에는 MTIA accelerator 지원 확대, 자동 워프 그룹 오케스트레이션 추론을 위한 "auto-TLX" 시스템, 그리고 torch.compile을 통한 PyTorch 2.x 통합 및 자동 커널 생성이 포함됨