ML-Triton通过多级编译提升GPU算子性能,支持细粒度控制。
ML-Triton, A Multi-Level Compilation and Language Extension to Triton GPU Programming
- 采用从工作组到瓦片、内建指令的逐级降低编译流程
- 在Intel GPU上实现超95%专家手写核函数性能,几何均值达标
- 新增编译提示与瓦片级编程,加速研究迭代无需等待编译器更新
在大模型时代,密集运算如GEMM和MHA是关键组件,适合基于分块的并行执行。传统GPU编程依赖CUDA或SYCL等底层接口,而Triton作为更高层级的领域特定语言(DSL)提供了更友好的编程体验和可移植性。当前Triton从工作组(workgroup)级别开始,直接降至线程级别,再通过一系列优化传递信息,我们认为此为过早降级。原因包括:1. 现代GPU具有物理与逻辑上的层次结构,支持以瓦片(warp)或瓦片组(warpgroup)为单位直接操作的SIMD单元,如分块加载和分块矩阵乘法(blocked MMA);2. 多级渐进式降级可使编译器解耦清晰,分离层间与层内考虑;3. 核心开发者常需精细控制以获得最新硬件性能,例如FlashAttention2通过显式瓦片间数据划分提升性能。为此,我们提出ML-Triton,具备多级编译流程与编程接口扩展。其从工作组层级逐步降低至瓦片与内建指令层级,契合GPU的层次结构。同时,扩展Triton语言以支持用户设定编译提示及瓦片级编程,使研究人员可在无需等待编译器更新的情况下获得良好开箱性能。实验表明,该方法在Intel GPU上达到专家手写核函数95%以上的性能,几何均值衡量。
原文摘要 · Abstract (English)
In the era of LLMs, dense operations such as GEMM and MHA are critical components. These operations are well-suited for parallel execution using a tilebased approach. While traditional GPU programming often relies on low level interfaces like CUDA or SYCL, Triton has emerged as a DSL that offers a more user-friendly and portable alternative by programming at a higher level. The current Triton starts at the workgroup (aka threadblock) level, and directly lowers to per-thread level. And then attempt to coalesce and amend through a series of passes, promoting information from low-level representation. We believe this is pre-mature lowering based on the below observations. 1. GPU has a hierarchical structure both physically and logically. Modern GPUs often feature SIMD units capable of directly operating on tiles on a warp or warpgroup basis, such as blocked load and blocked MMA. 2. Multi-level gradual lowering can make compiler decoupled and clean by separating considerations inter and intra a logical layer. 3. Kernel developers often need fine control to get good performance on the latest hardware. FlashAttention2 advocates explicit data partition between warps to make a performance boost. In this context, we propose ML-Triton which features multi-level compilation flow and programming interface. Our approach begins at the workgroup level and progressively lowers to the warp and intrinsic level, implementing a multilevel lowering align with the hierarchical nature of GPU. Additionally, we extend triton language to support user-set compiler hint and warp level programming, enabling researchers to get good out-of-the box performance without awaiting compiler updates. Experimental results demonstrate that our approach achieves performance above 95% of expert-written kernels on Intel GPU, as measured by the geometric mean.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。