A summary of the video transcript with Tri Dao
Tri Dao, the inventor of Flash Attention and Flash Attention 2, shares his insights on the development of these techniques, the current state of AI research, and his perspectives on the future of the field. Let's break down the key points from the discussion. **Insights and Ideas behind Flash Attention** * Dao and his team focused on making attention more memory-efficient without approximations, as most practitioners preferred exact methods over approximate ones. * They realised that the majority of time in attention computation was spent on memory reading and writing, **not floating-point operations (FLOPs)**. * By combining ideas from the systems side (kernel fusion, tiling) and the machine learning side (the online softmax trick), they developed Flash Attention, which ***reduces memory usage from quadratic to linear in sequence length***. * Flash Attention provides a 2-4x speedup compared to standard attention implementation, allowing for training with longer sequences at the same cost. **Development of Flash Attention 2** * Dao started working on Flash Attention 2 as an excuse to play with **NVIDIA's CUTLASS library**, which provides primitives for efficient GPU operations. * He realised there were still inefficiencies in Flash Attention and made tweaks to reduce non-matmul FLOPs, as modern GPUs have specialised units (e.g., Tensor Cores) that make matmul much faster than non-matmul operations. * Flash Attention 2 also parallelises the computation across different thread blocks and partitions the work between warps to reduce shared memory access. * These optimisations result in a 2x speedup compared to Flash Attention, approaching the efficiency of highly-optimized matrix multiplication routines. **Explanation of GPU Performance Characteristics** * GPUs have a memory hierarchy consisting of **high-bandwidth memory (HBM)** and **on-chip SRAM (shared memory)**. * HBM is larger but slower, while SRAM is much smaller but faster. * The goal is to design algorithms that take advantage of this asymmetric memory hierarchy, loading data from HBM to SRAM, performing computations, and then writing results back to HBM. * As sequence lengths grow, the memory required for attention matrices increases quadratically, while the number of attention heads remains relatively small, leading to a memory inefficiency that Flash Attention aims to address. **Role of Academia and Industry in AI Research** * Academia and industry play complementary roles in AI research. * Industry has an advantage in scaling models and providing infrastructure, while academia focuses on developing new ideas, understanding and explaining models, and taking riskier bets on novel architectures. * Dao believes it's important for researchers to maintain a balance between exploring fundamentals and keeping up with the latest developments in the field. **Transformer Alternatives and Future Directions** * Dao is interested in ***exploring transformer alternatives***, such as recurrent neural networks (RNNs) and **state space methods**, which could be more efficient for long sequences and high-throughput generation. * He believes these alternatives could have advantages in specific use cases, such as processing high-resolution images, audio, and video. * However, he also acknowledges the "hardware lottery" and the challenges of developing new architectures that can compete with the heavily-optimised transformer on current hardware and software frameworks.Refresher: GPU Performance Characteristics
**Compute elements** GPUs consist of compute elements, such as floating-point arithmetic units, which perform the actual computations. These units are designed to handle parallel processing tasks efficiently. **Specialised matrix multiply units** Modern GPUs, like Nvidia's Tensor Cores, have specialised units that accelerate matrix multiplication operations in low-precision formats, such as FP16 (half-precision floating-point) or BF16 (bfloat16). These units are optimised for fast and efficient matrix multiplication, which is a common operation in deep learning workloads. **Memory hierarchy** GPUs have a memory hierarchy that consists of two main components: **High Bandwidth Memory (HBM):** HBM is a type of high-performance memory that offers high bandwidth and capacity. For example, the A100 GPU has 40-80GB of HBM with a bandwidth of 1.5-2.0TB/s. HBM is used to store large datasets and intermediate results during computation. **On-chip SRAM (Shared Memory):** GPUs also have on-chip SRAM, known as shared memory, which is a fast but limited-capacity memory. In the A100 GPU, each of the 108 streaming multiprocessors (SMs) has 192KB of shared memory with an estimated bandwidth of around 19TB/s. Shared memory is used for fast communication and data sharing between threads within a thread block. **L2 cache:** GPUs also have an L2 cache, which is a larger but slower cache compared to shared memory. However, the L2 cache is not directly controllable by the programmer and is managed by the GPU's cache hierarchy. **Execution Model** **Threads:** GPUs employ a massive number of threads to execute operations in parallel. These threads are lightweight and can be quickly switched to hide memory latency and maximise utilisation of the GPU's compute resources. **Thread blocks:** Threads are organized into thread blocks, which are scheduled to run on the GPU's streaming multiprocessors (SMs). Each thread block contains a set of threads that can cooperate and share resources, such as shared memory. **Warps:** Within each thread block, threads are further grouped into warps, typically consisting of 32 threads. Threads within a warp execute in lockstep, meaning they perform the same instruction on different data elements simultaneously (SIMD - Single Instruction, Multiple Data). Warps are the basic unit of execution on the GPU. **Communication and synchronization** **Shuffle instructions:** Threads within a warp can communicate using fast shuffle instructions, which allow them to exchange data directly without going through shared memory. This is useful for efficient data sharing and reduction operations within a warp. **Shared memory:** Warps within a thread block can communicate by reading from and writing to shared memory. Shared memory acts as a fast, software-managed cache that enables efficient data sharing and coordination between warps. **Kernel execution:** A GPU kernel is a function that is executed in parallel by multiple threads on the GPU. Each kernel loads inputs from HBM into registers and shared memory, performs computations, and then writes the outputs back to HBM. The kernel is launched with a specified grid and block size, which determines the number of thread blocks and threads per block. Understanding these GPU performance characteristics and the execution model optimises GPU code and achieves high performance. By leveraging the specialised matrix multiply units, using the memory hierarchy effectively, and organising threads and thread blocks appropriately, developers can maximise the utilisation of GPU resources and minimise memory latency, ultimately leading to **f*****aster execution of parallel workloads.***
From the above image you can see that SRAM is faster than the HBM, but SRAM has very much limited memory as its fast cache memory. In the above image you can compare the speed of HBM and SRAM. HBM is used to store Tensors, while SRAM is used to perform compute operations on those Tensors.
The problem with the sequence length being larger than the head dimension
The significance of **N** being much larger than **d** lies in the memory complexity of the standard attention implementation. Since the **attention scores matrix**** ****S** and the **attention probabilities matrix** **P** have dimensions **N×N**, the memory required to store these matrices grows quadratically with the sequence length **N**. For example, if **N** is 1,000 (1k), the **matrices** would have 1,000,000 (1 million) elements, and if **N** is 8,000 (8k), the **matrices** would have 64,000,000 (64 million) elements! In contrast, the **head dimension** **d** is typically **much smaller** than **N**. Even with a head dimension of 128, which is on the higher end of the typical range, it is still significantly smaller than the sequence lengths encountered in practice. This difference in scale between **N** and **d** highlights the memory inefficiency of the standard attention implementation. As the **sequence length** **N** grows, the ***memory required to store the attention matrices increases quadratically***, while the **head dimension** **d** remains relatively small. This leads to a bottleneck in terms of memory consumption, especially when dealing with long sequences, making it challenging to process such sequences efficiently using the standard attention implementation.
You can see the concept here - tiling means breaking or chopping the the matrix into multiple block and then performing the computation in SRAM instead of passing a complete N*N matrix from HBM to SRAM. This is because SRAM is very limited in capacity - the entire matrix cannot be loaded into SRAM.
From the paper: In the forward pass (left), we parallelise the workers (thread blocks) where each worker takes care of a block of rows of the attention matrix. In the backward pass (right), each worker takes care of a block of columns of the attention matrix
Work partitioning between different warps in the forward pass