Expose safe low level gpu programming to improve tensor core performance
Exo-GPU: Safe, Imperative, User-schedulable Programming for Tensor Cores
Programming Languages
Summary
Using modern GPUs for fast math needs careful control of how work is split and coordinated. The authors designed Exo-GPU, a programming language that lets developers write low-level GPU code safely by marking parallelism as annotations rather than core parts of code flow. This design helps ensure parallel execution behaves the same as running steps one after another. They used Exo-GPU to build matrix multiplication programs that run very close to GPU speed limits and sometimes beat the official library.
What this means in practice
- •For gpu performance engineers: Write safer low-level GPU kernels by verifying parallel and sequential equivalence to achieve high-performance matrix operations on tensor cores.
- •For high-performance computing developers: Develop efficient GPU compute kernels that nearly reach peak hardware performance and outperform vendor libraries for large-scale linear algebra tasks.
Authors
David Zhao Akeley, Yuka Ikarashi, Jonathan Ragan-Kelley
Abstract
Modern GPUs require not only SIMT-style parallelism but also software-managed concurrency between compute and data movement to reach maximum performance. Performance engineers must reason about subdividing work into the hierarchy of computation resources (threads, warps, warpgroups, blocks, clusters), and, in many cases, also must use asynchronous tensor core and memcpy instructions on different levels of the memory hierarchy (registers, tensor core accumulators, shared memory, global memory). Unlike CPUs, where out-of-order execution is managed by hardware and hidden from programmers, GPUs expose explicit instruction reordering to software through these asynchronous instructions. Well-established GPU programming languages generally offer either direct low-level control without safety guarantees (e.g., CUDA C++ inline assembly or intrinsics) or easier-to-analyze, high-level abstractions (e.g., Triton's tile-based model) that hide asynchronous instructions in the compiler backend, which may prevent performance engineers from maximizing performance by tuning critical details. We propose Exo-GPU, an imperative, low-level language that creates minimal abstraction over CUDA. Our key idea is to treat parallelism and synchronization as mere annotations on sequential code rather than as fundamental control flow primitives, enabling verification that these constructs do not alter the program semantics. The benefit is twofold: programmers can reason about code without hidden control flow or mutation, while allowing the Exo-GPU compiler to verify sequential-parallel equivalence--guaranteeing that parallel execution is functionally equivalent to its sequential interpretation. We used Exo-GPU to author GEMM kernels for the H100 GPU, using wgmma, TMA, and split-k. Our kernels achieved over 80% of theoretical peak on large problem sizes, in some cases outperforming the vendor-provided CUBLAS library.