arXiv:2504.16214cs.LGcs.AI2025-04被引 1

Hexcute自动优化GPU程序数据布局,提升性能并大幅减少代码量。

Hexcute: A Compiler Framework for Automating Layout Synthesis in GPU Programs

  • 将布局生成建模为约束求解问题,用类型推断算法自动搜索最优方案。
  • 在GEMM、Attention等算子上性能媲美cuBLAS,代码量减少1.27到7.94倍。
  • 适合需要高性能且希望减少手动调优的深度学习开发者使用。

高效的GPU编程对深度学习应用的高性能至关重要。GPU程序的性能取决于数据在线程间的并行方式以及在内存子系统中的布局方式。描述张量在GPU上映射的函数称为张量布局。低层编程框架如CUTLASS和Hidet提供灵活的布局抽象,但通常需要大量手工编写以指定最优布局。高层GPU编程语言如Triton依赖编译器启发式方法生成数据流、布局和流水线策略,但这些启发式方法难以推广到复杂算子。为平衡表达力与可编程性,我们提出Hexcute,一个自动化布局合成的编译框架,同时保留对数据流和流水线的显式控制。Hexcute将布局合成形式化为约束规划问题,并通过基于类型推断的算法求解。该方法支持系统性探索最优布局与指令序列。评估显示,Hexcute在GEMM、Attention及其变体上的性能达到cuBLAS和FlashAttention水平,代码量比CUTLASS减少1.27×–7.94×;对于混合类型专家模型(MoE)算子,平均提速6.46×;在vLLM端到端评估中,对DeepSeek-R1-AWQ模型最高提速2.60×,对基于Mamba的模型提速2.04×。

原文摘要 · Abstract (English)

Efficient GPU programming is crucial for achieving high performance in deep learning (DL) applications. The performance of GPU programs depends on how data is parallelized across threads and arranged within memory subsystems. The mapping functions describing tensors on GPUs are known as \emph{tensor layouts}. Low-level programming frameworks, such as CUTLASS and Hidet, provide expressive layout abstractions but often require \emph{considerable programming effort} to manually specify optimal layouts. High-level GPU programming languages, such as Triton, rely on compiler heuristics to generate dataflow, layouts, and pipelining strategies in GPU programs. However, the heuristics for dataflow and pipelining strategies are not generalizable to complex operators. To balance expressiveness and programmability, we propose Hexcute, a compiler framework that automates layout synthesis while providing explicit control over dataflow and pipelining. Hexcute formalizes layout synthesis as a constraint programming problem and solves it with a type-inference-based algorithm. This approach enables systematic exploration of optimal layouts and instructions. Our evaluation shows that Hexcute matches the performance of libraries like cuBLAS and FlashAttention on GEMM, Attention, and their variants, while reducing the amount of code by 1.27$\times$-7.94$\times$ compared to CUTLASS. For mixed-type mixture-of-experts (MoE) operators, Hexcute achieves an average speedup of 6.46$\times$ over Triton. In the end-to-end evaluations of vLLM, Hexcute delivers up to 2.60$\times$ speedup on DeepSeek-R1-AWQ and 2.04$\times$ on a Mamba-based model.

GPU优化编译器张量布局自动化

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。