A Case Study in CUDA Kernel Fusion: Implementing FlashAttention-2 on NVIDIA Hopper Architecture using the CUTLASS Library (2312.11918v1)

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.

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 paper 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.