We worked on this project at the end of 2023 and in early 2024. NVIDIA’s Hopper GPUs had just been on the market, but popular AI frameworks offered little to no support for 8-bit inference at the time. Several other optimizations were also waiting for our attention. In this case study we’ll give some insight on how we reached 1 PFLOPS Flash Attention on a single H100 GPU.
About the customer
HippoML was an innovative startup that changed how AI compute works. They saw an opportunity to bring generative AI to lower GPU costs, and we worked closely together on several projects. This case study describes the first.
NVIDIA has since acquired HippoML, so we can now share a few details about these projects.
The challenge
The Attention algorithm sits at the core of many modern transformer-based AI models: large language models, vision models, and other generative AI. It computes how a transformer weighs the relevance of every other token when processing one. Because Attention often dominates the computational cost of training and inference – especially for long sequences – any performance improvement here is very valuable.
Flash Attention performs the entire Attention operation in a single kernel. By eliminating intermediary memory operations, it delivers a significant speedup over classical Attention. Most AI frameworks adopted the primary implementation, by Tri Dao and contributors, as their main Attention method.
The client asked us to optimize this algorithm for NVIDIA’s Hopper and Lovelace architectures. At the time, the primary Flash Attention implementation was not ready for the quantized-AI era: it supported fp16 and bf16, but no quantized types. The client had novel ideas for improving the algorithm, but needed help implementing them at good performance.
Our approach
We modified the original Flash Attention to support int8 and fp8. We also added Bias, Paged Attention, and Varlen Attention, and we optimized the existing fp16 and bf16 paths.
Making the most of Hopper
Lovelace was a fairly conventional upgrade over the previous generation, but NVIDIA designed Hopper specifically for AI. It introduced specialized features never seen before on GPUs. To reach the best performance, we made extensive use of the new TMA engine and WarpGroup-GEMM instructions.
A guessing game on Lovelace fp8
Supporting Lovelace was not straightforward either. When we first discussed the implementation with the client, NVIDIA’s marketing pages claimed that this generation of GPUs supported accelerated fp8 matrix instructions. Yet the most recent CUDA toolkit at the time, 12.3, did not include this support.
We studied the input and output layouts of other accelerated matrix instructions. From those, we made an educated guess about the layout a future fp8 matrix instruction would use. To keep development moving, we built a temporary fallback with conventional bf16 matrix instructions.
CUDA 12.4 then added accelerated fp8 matrix instructions on Lovelace; and the layout matched our expectations. That let us ship day-1 support for fp8 Flash Attention on Lovelace.
Fixing a transpose bottleneck
One bottleneck in the Flash Attention kernel was the transposition of one of the input matrices. Pre-transposing this matrix improved end-to-end model inference performance. The results below do not include this optimization.
Work on Cutlass underneath
Flash Attention builds on NVIDIA’s Cutlass library, which is notorious for its advanced and difficult-to-understand mathematical basis. We also worked on Cutlass directly, adding 8- and 16-bit integer convolution with dynamic quantization, plus fp8 GEMV. The client excelled at the high-level design of how to string these algorithms together; they needed help building performant implementations. Given the end result, we succeeded.
The result
Our implementation was the world’s first to measure over 1 PFLOPS on a single GPU. The HippoML team announced it on 7 February 2024:
After further optimization, we are thrilled to announce that HippoAttention has achieved 1 PFLOPS on a single H100 SXM GPU.
— HippoML on their blog
Another lab reached the same milestone some weeks later. We were first.
Key takeaways
- New silicon demands new kernels to reach peak performance. That means purpose-built kernels per architecture – not portable code.
- Together with HippoML, we showed that low-bit precision is the path to affordable AI.
- Innovators work with us to get winning implementations first. Together, we built the world’s first 1 PFLOPS/s Flash Attention.
Did you enjoy reading about how we made the AI faster which you potentially use to generate images? You can read about more projects we worked on by checking out our other case studies.