自动编译张量公式为高效GPU代码,无需人工优化。
Nautilus: An Auto-Scheduling Tensor Compiler for Efficient Tiled GPU Kernels

- 通过逐级降低设计,统一优化高阶表达式与分块策略。
- 在多种Transformer模型上,比现有编译器最高提升42%吞吐量。
- 首次实现从数学描述自动生成FlashAttention-3类内核,适合深度学习开发者。
我们提出Nautilus,一种新型张量编译器,致力于实现从数学描述到高效分块GPU内核的全自动优化。Nautilus将张量算子的高层代数描述编译为高效的分块GPU内核。其逐级降低的设计使高层优化、表达式重写和分块优化可在单一端到端系统中协同应用。Nautilus引入一种新颖的自动调度器,能发现一系列高层优化序列,同时保持分块优化器所需的规则程序结构。该自动调度器捕捉了高层优化中的复杂交互与权衡,包括激进的全局变换如高级归约融合。Nautilus是首个能从注意力的数学描述出发,自动发现类似FlashAttention-3的内核的端到端张量编译器,彻底将优化负担从程序员转移到编译器。在NVIDIA GH200和RTX 5090 GPU上的五种Transformer模型及150个评估配置下,其吞吐量相比最先进编译器在GH200上最高提升23%,在RTX 5090上最高提升42%,并在许多长序列配置中达到或超越手动编写cuDNN内核的性能。
原文摘要 · Abstract (English)
We present Nautilus, a novel tensor compiler that moves toward fully automated math-to-kernel optimization. Nautilus compiles a high-level algebraic specification of tensor operators into efficient tiled GPU kernels. Nautilus's successive lowering design allows high-level optimizations, expression rewrites, and tile optimizations to be jointly applied in a single end-to-end system. Nautilus presents a novel auto-scheduler that discovers sequences of high-level optimizations, while preserving the regular program structure needed by tile optimizers. Nautilus's auto-scheduler captures complex interactions and trade-offs in the high-level optimizations, including aggressive global transformations like advanced reduction fusion. Nautilus is the first end-to-end tensor compiler capable of starting from a math-like description of attention and automatically discovering FlashAttention-3-like kernels, offloading the entire burden of optimization from the programmer to the compiler. Across five transformer-based models and 150 evaluation configurations on NVIDIA GH200 and RTX 5090 GPUs, Nautilus achieves up to 23% higher throughput than state-of-the-art compilers on GH200 and up to 42% on RTX 5090, while matching or exceeding manually written cuDNN kernels on many long-sequence configurations.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。