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

arXiv: https://arxiv.org/abs/2410.07531 · HTML: https://mlsanity.com/paper/2410.07531/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 23 of its 47 references that are in the map.
- First posted on arXiv 10 Oct 2024.

## Builds on (most cited first, 23 of 23)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [The Llama 3 Herd of Models](https://mlsanity.com/paper/2407.21783/) (2024)
- [Llama 2: Open Foundation and Fine-Tuned Chat Models](https://mlsanity.com/paper/2307.09288/) (2023)
- [FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness](https://mlsanity.com/paper/2205.14135/) (2022)
- [FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning](https://mlsanity.com/paper/2307.08691/) (2023)
- [Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism](https://mlsanity.com/paper/1909.08053/) (2019)
- [Splitwise: Efficient generative LLM inference using phase splitting](https://mlsanity.com/paper/2311.18677/) (2023)
- [FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision](https://mlsanity.com/paper/2407.08608/) (2024)
- [DeepSpeed-MoE: Advancing Mixture-of-Experts Inference and Training to Power Next-Generation AI Scale](https://mlsanity.com/paper/2201.05596/) (2022)
- [Reducing Activation Recomputation in Large Transformer Models](https://mlsanity.com/paper/2205.05198/) (2022)
- [LLM Inference Unveiled: Survey and Roofline Model Insights](https://mlsanity.com/paper/2402.16363/) (2024)
- [Sequence Parallelism: Long Sequence Training from System Perspective](https://mlsanity.com/paper/2105.13120/) (2021)
- [Vidur: A Large-Scale Simulation Framework For LLM Inference](https://mlsanity.com/paper/2405.05465/) (2024)
- [To Repeat or Not To Repeat: Insights from Scaling LLM under Token-Crisis](https://mlsanity.com/paper/2305.13230/) (2023)
- [FLUX: Fast Software-based Communication Overlap On GPUs Through Kernel Fusion](https://mlsanity.com/paper/2406.06858/) (2024)
- [MuxServe: Flexible Spatial-Temporal Multiplexing for Multiple LLM Serving](https://mlsanity.com/paper/2404.02015/) (2024)
- [LLM Inference Serving: Survey of Recent Advances and Opportunities](https://mlsanity.com/paper/2407.12391/) (2024)
- [Dropout Reduces Underfitting](https://mlsanity.com/paper/2303.01500/) (2023)
- [DropAttention: A Regularization Method for Fully-Connected Self-Attention Networks](https://mlsanity.com/paper/1907.11065/) (2019)
- [T3: Transparent Tracking & Triggering for Fine-grained Overlap of Compute & Collectives](https://mlsanity.com/paper/2401.16677/) (2024)
- [Optimizing Distributed ML Communication with Fused Computation-Collective Operations](https://mlsanity.com/paper/2305.06942/) (2023)
- [Etalon: Holistic Performance Evaluation Framework for LLM Inference Systems](https://mlsanity.com/paper/2407.07000/) (2024)

## Abstract (arXiv)

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.