Reducing the Cost of Dropout in Flash-Attention by Hiding RNG with GEMM

Dropout, a network operator, when enabled is likely to dramatically impact the performance of Flash-Attention, which in turn increases the end-to-end training time of Large-Language-Models (LLMs). The main contributor to such performance degradation is the Random Number Generation (RNG) phase. The state-of-the-art optimization is to fuse RNG into the Flash-Attention kernel. However, while RNG and Attention do not compete on compute or memory resources, they are bounded on the same lower-level architecture bottlenecks. Fusion can hardly hide RNG latency within the Attention kernel. We propose overlapping RNG with previous GEMM layers in the network to hide RNG latency and improve end-to-end performance. RNG and GEMM have distinct resource requirements and hardware bottlenecks, so they can run together without compromising each other's performance. We propose a fine-grained analytical performance model that analyzes low-level architecture resource utilization to evaluate RNG-GEMM overlapping performance benefits. This model, cross-validated by silicon results, shows 1.26x speedup for overlapping RNG and GEMM layers over a sequential implementation on one Transformer Block (one LLM layer including multi-head attention and feed-forward layers), and 1.22x over state-of-the-art fusion implementation, for Llama3 on GH100 GPUs with FP8 precision. Because the kernel patterns are regular, the findings of the shared bottlenecks, as well as the achievable performance benefits, can be generalized to different model architectures, software implementations and hardware configurations.