A Case Study in CUDA Kernel Fusion: Implementing FlashAttention-2 on NVIDIA Hopper Architecture using the CUTLASS Library
(2312.11918)Abstract
We provide an optimized implementation of the forward pass of FlashAttention-2, a popular memory-aware scaled dot-product attention algorithm, as a custom fused CUDA kernel targeting NVIDIA Hopper architecture and written using the open-source CUTLASS library. In doing so, we explain the challenges and techniques involved in fusing online-softmax with back-to-back GEMM kernels, utilizing the Hopper-specific Tensor Memory Accelerator (TMA) and Warpgroup Matrix-Multiply-Accumulate (WGMMA) instructions, defining and transforming CUTLASS Layouts and Tensors, overlapping copy and GEMM operations, and choosing optimal tile sizes for the Q, K and V attention matrices while balancing the register pressure and shared memory utilization. In head-to-head benchmarks on a single H100 PCIe GPU for some common choices of hyperparameters, we observe 20-50% higher FLOPs/s over a version of FlashAttention-2 optimized for last-generation NVIDIA Ampere architecture.
Overview
-
Presents a case study on optimizing the FlashAttention-2 algorithm using CUDA kernel fusion on NVIDIA Hopper architecture.
-
Describes how kernel fusion reduces data transfers between GPU memory and processors, beneficial for LLM training and inference.
-
Details implementation using the CUTLASS library, improving computational efficiency by 20-50% over past architectures.
-
Explains the selection of optimal tile sizes, register pressure management, and shared memory utilization.
-
Identifies potential optimization strategies for future work, including sophisticated warpgroup use and improved pipelining.
Introduction to Kernel Fusion
In the area of computer programming, particularly for GPU computing, kernel fusion is an advanced technique with significant benefits. It combines multiple computational kernels into a single kernel, thereby reducing the amount of data transferred between the GPU memory and its processors. This is particularly advantageous for applications where memory bandwidth is a limiting factor, which is often the case with modern GPU hardware where computational power has outpaced memory bandwidth improvements. Amongst the workloads that benefit from kernel fusion is the training and inference of LLMs, which are foundational to major breakthroughs in AI and natural language processing.
Implementing Advanced Attention Algorithms
The focus of the paper is on optimizing the forward pass of FlashAttention-2, an attention mechanism algorithm that is essential for transformer models like GPT-3. Transforming FlashAttention-2 into a fused CUDA kernel leverages the capabilities of the NVIDIA Hopper architecture and is accomplished using the CUTLASS library, which simplifies GPU kernel development through the use of various abstractions. The paper demonstrates how to fuse operations, specifically an online-softmax function with two General Matrix Multiply (GEMM) operations, using Hopper's specialized Tensor Memory Accelerator (TMA) and Warpgroup Matrix-Multiply-Accumulate (WGMMA) instructions. It reports a 20-50% improvement in computational efficiency when compared to a previous generation architecture.
Performance Enhancements and Coding Abstractions
The implementation carefully chooses and explains the intricacies of optimal tile sizes for matrices and handles trade-offs between register pressure and shared memory utilization. The fusion of operations involves dealing with layout transformations and orchestrating data copies and computations to increase parallelism and reduce overheads. The custom kernel developed exhibits notable performance gains over previous versions of the algorithm tailored for last-generation hardware.
Directions for Future Work
Although promising results were encountered, the study identifies potential areas for optimization that could be further explored. These include more sophisticated warpgroup use, enhanced pipelining to better overlap memory operations with computations, and utilizing new shared memory features provided by upcoming GPU architectures. The paper anticipates that future improvements to GPU hardware and attention-based algorithm implementations will continue to push the boundaries of what can be achieved in the processing efficiency and performance of LLMs.
Create an account to read this summary for free: