about
A Case Study in CUDA Kernel Fusion (arxiv.org)
1 point by veryluckyxyz on May 25, 2024 | hide | past | pdf | discuss on HN

In plain words: A custom graphics-chip program packs FlashAttention-2's forward pass into one fused operation on NVIDIA's newest chips, merging the step-by-step softmax with matrix multiplies and using built-in hardware to overlap copying and computing. It runs 20-50% faster at raw math than the version tuned for the previous chip generation.

Abstract · A Case Study in CUDA Kernel Fusion: Implementing FlashAttention-2 on NVIDIA Hopper Architecture using the CUTLASS Library

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.

Ganesh Bikshandi, Jay Shah
arXiv:2312.11918 · cs.LG, cs.DC · submitted Dec 19, 2023
abstract · pdf · html · 13 pages, comments welcome

add comment on HN