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

arXiv:2504.16214v3 Announce Type: replace-cross
Abstract: 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 emphtensor layouts. Low-level programming frameworks, such as CUTLASS and Hidet, provide expressive layout abstractions but often require emphconsiderable 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.