Hexcute: A Tile-based Programming Language with Automatic Layout and Task-Mapping Synthesis

2025/04/22 by Zhang, Xiao, Ding, Yaoyao, Hu, Yang +1 · 1 citation
#Artificial Intelligence (cs.AI) #FOS: Computer and information sciences #Machine Learning (cs.LG) #Programming Languages (cs.PL)

paper · doi:10.48550/arxiv.2504.16214

Abstract

Deep learning (DL) workloads mainly run on accelerators like GPUs. Recent DL quantization techniques demand a new matrix multiplication operator with mixed input data types, further complicating GPU optimization. Prior high-level compilers like Triton lack the expressiveness to implement key optimizations like fine-grained data pipelines and hardware-friendly memory layouts for these operators, while low-level programming models, such as Hidet, Graphene, and CUTLASS, require significant programming efforts. To balance expressiveness with engineering effort, we propose Hexcute, a tile-based programming language that exposes shared memory and register abstractions to enable fine-grained optimization for these operators. Additionally, Hexcute leverages task mapping to schedule the GPU program, and to reduce programming efforts, it automates layout and task mapping synthesis with a novel type-inference-based algorithm. Our evaluation shows that Hexcute generalizes to a wide range of DL operators, achieves 1.7-11.28× speedup over existing DL compilers for mixed-type operators, and brings up to 2.91× speedup in the end-to-end evaluation.

Citations

• Linear Layouts: Robust Code Generation of Efficient Tensor Computation Using \mathbbF₂
• Task-Based Tensor Computations on Modern GPUs
• FlashInfer: Efficient and Customizable Attention Engine for LLM Inference Serving
• MARLIN: Mixed-Precision Auto-Regressive Parallel Inference on Large Language Models
• Ascend-CC: Confidential Computing on Heterogeneous NPU for Emerging Generative AI Workloads
• FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision
• Towards a high-performance AI compiler with upstream MLIR
• Mamba: Linear-Time Sequence Modeling with Selective State Spaces
• Efficient Memory Management for Large Language Model Serving with PagedAttention
• Recent Advancements in End-to-End Autonomous Driving using Deep Learning: A Survey
• AWQ: Activation-aware Weight Quantization for LLM Compression and Acceleration
• TPU v4: An Optically Reconfigurable Supercomputer for Machine Learning with Hardware Support for Embeddings
• LLaMA: Open and Efficient Foundation Language Models
• GPTQ: Accurate Post-Training Quantization for Generative Pre-trained Transformers
• SparseTIR: Composable Abstractions for Sparse Compilation in Deep Learning
• TensorIR: An Abstraction for Automatic Tensorized Program Optimization
• FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness
• High-Resolution Image Synthesis with Latent Diffusion Models
• High-Resolution Image Synthesis with Latent Diffusion Models
• Multi-Task Learning in Natural Language Processing: An Overview
• Ansor: Generating High-Performance Tensor Programs for Deep Learning
• Language Models are Few-Shot Learners
• A Survey of Autonomous Driving: Common Practices and Emerging Technologies
• Optimal DNN Primitive Selection with Partitioned Boolean Quadratic Programming
• Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer
• Deep Residual Learning for Image Recognition
• Fast R-CNN

Cited by

• Ada-MK: Adaptive MegaKernel Optimization via Automated DAG-based Search for LLM Inference

Related

• TVM: An Automated End-to-End Optimizing Compiler for Deep Learning
• Tiramisu: A Polyhedral Compiler for Expressing Fast and Portable Code
• ALCOP: Automatic Load-Compute Pipelining in Deep Learning Compiler for AI-GPUs
• Neptune: Advanced ML Operator Fusion for Locality and Parallelism on GPUs
• Salus: Fine-Grained GPU Sharing Primitives for Deep Learning Applications