Search papers, labs, and topics across Lattice.
Nautilus, a novel tensor compiler, automates the generation of efficient tiled GPU kernels from high-level algebraic specifications by jointly optimizing high-level transformations and tile optimizations. Its auto-scheduler discovers optimization sequences, including advanced reduction fusion, while preserving program structure. Experiments on transformer models show Nautilus achieves up to 23% and 42% higher throughput than state-of-the-art compilers on NVIDIA GH200 and RTX 5090 GPUs, respectively, and matches or exceeds manually written cuDNN kernels in many cases.
Nautilus can automatically discover FlashAttention-3-like kernels from a math-like description of attention, completely automating the optimization process.
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.
