# DISTFLASHATTN: Distributed Memory-efficient Attention for Long-context LLMs Training

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

## Key facts

- Cited by 29 papers in the map (data through 14 Oct 2026).
- Builds on 19 of its 25 references that are in the map.
- First posted on arXiv 5 Oct 2023.

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

- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [LLaMA: Open and Efficient Foundation Language Models](https://mlsanity.com/paper/2302.13971/) (2023)
- [Judging LLM-as-a-Judge with MT-Bench and Chatbot Arena](https://mlsanity.com/paper/2306.05685/) (2023)
- [FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness](https://mlsanity.com/paper/2205.14135/) (2022)
- [Longformer: The Long-Document Transformer](https://mlsanity.com/paper/2004.05150/) (2020)
- [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)
- [Big Bird: Transformers for Longer Sequences](https://mlsanity.com/paper/2007.14062/) (2020)
- [GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints](https://mlsanity.com/paper/2305.13245/) (2023)
- [Training Deep Nets with Sublinear Memory Cost](https://mlsanity.com/paper/1604.06174/) (2016)
- [PyTorch FSDP: Experiences on Scaling Fully Sharded Data Parallel](https://mlsanity.com/paper/2304.11277/) (2023)
- [Ring Attention with Blockwise Transformers for Near-Infinite Context](https://mlsanity.com/paper/2310.01889/) (2023)
- [Reducing Activation Recomputation in Large Transformer Models](https://mlsanity.com/paper/2205.05198/) (2022)
- [DeepSpeed Ulysses: System Optimizations for Enabling Training of Extreme Long Sequence Transformer Models](https://mlsanity.com/paper/2309.14509/) (2023)
- [Online normalizer calculation for softmax](https://mlsanity.com/paper/1805.02867/) (2018)
- [Self-attention Does Not Need O(n2) Memory](https://mlsanity.com/paper/2112.05682/) (2021)
- [A Length-Extrapolatable Transformer](https://mlsanity.com/paper/2212.10554/) (2022)
- [Checkmate: Breaking the Memory Wall with Optimal Tensor Rematerialization](https://mlsanity.com/paper/1910.02653/) (2019)
- [Multi-head or Single-head? An Empirical Comparison for Transformer Training](https://mlsanity.com/paper/2106.09650/) (2021)

## Built on by (most cited first, 29 of 29)

- [Fast Video Generation with Sliding Tile Attention](https://mlsanity.com/paper/2502.04507/) (2025)
- [Scaling Laws of RoPE-based Extrapolation](https://mlsanity.com/paper/2310.05209/) (2023)
- [Reinforcement Learning Optimization for Large-Scale Learning: An Efficient and User-Friendly Scaling Library](https://mlsanity.com/paper/2506.06122/) (2025)
- [Why Does the Effective Context Length of LLMs Fall Short?](https://mlsanity.com/paper/2410.18745/) (2024)
- [MegaScale-MoE: Large-Scale Communication-Efficient Training of Mixture-of-Experts Models in Production](https://mlsanity.com/paper/2505.11432/) (2025)
- [Context Parallelism for Scalable Million-Token Inference](https://mlsanity.com/paper/2411.01783/) (2024)
- [Efficient-vDiT: Efficient Video Diffusion Transformers With Attention Tile](https://mlsanity.com/paper/2502.06155/) (2025)
- [FlexSP: Accelerating Large Language Model Training via Flexible Sequence Parallelism](https://mlsanity.com/paper/2412.01523/) (2024)
- [Efficient Parallelization Layouts for Large-Scale Distributed Model Training](https://mlsanity.com/paper/2311.05610/) (2023)
- [Helix Parallelism: Rethinking Sharding Strategies for Interactive Multi-Million-Token LLM Decoding](https://mlsanity.com/paper/2507.07120/) (2025)
- [Out of the Memory Barrier: A Highly Memory Efficient Training System for LLMs with Million-Token Contexts](https://mlsanity.com/paper/2602.02108/) (2026)
- [Hydraulis: Balancing Large Transformer Model Training via Co-designing Parallel Strategies and Data Assignment](https://mlsanity.com/paper/2412.07894/) (2024)
- [SALE : Low-bit Estimation for Efficient Sparse Attention in Long-context LLM Prefilling](https://mlsanity.com/paper/2505.24179/) (2025)
- [Arctic Long Sequence Training: Scalable And Efficient Training For Multi-Million Token Sequences](https://mlsanity.com/paper/2506.13996/) (2025)
- [LongStraw: Long-Context RL Beyond 2M Tokens under a Fixed GPU Budget](https://mlsanity.com/paper/2607.14952/) (2026)
- [HelixPipe: Efficient Distributed Training of Long Sequence Transformers with Attention Parallel Pipeline Parallelism](https://mlsanity.com/paper/2507.00394/) (2025)
- [Pixel-Resolved Long-Context Learning for Turbulence at Exascale: Resolving Small-scale Eddies Toward the Viscous Limit](https://mlsanity.com/paper/2507.16697/) (2025)
- [Why Attention Fails: A Taxonomy of Faults in Attention-Based Neural Networks](https://mlsanity.com/paper/2508.04925/) (2025)
- [A Scalable Distributed Framework for Multimodal GigaVoxel Image Registration](https://mlsanity.com/paper/2509.25044/) (2025)
- [Efficient Scaling of LLM Training with Flexible Context Parallelism](https://mlsanity.com/paper/2602.21788/) (2026)
- [LV-XAttn: Distributed Cross-Attention for Long Visual Inputs in Multimodal Large Language Models](https://mlsanity.com/paper/2502.02406/) (2025)
- [Scaling Generative Recommendations with Context Parallelism on Hierarchical Sequential Transducers](https://mlsanity.com/paper/2508.04711/) (2025)
- [TASP: Topology-aware Sequence Parallelism](https://mlsanity.com/paper/2509.26541/) (2025)
- [Data-Aware Random Feature Kernel for Transformers](https://mlsanity.com/paper/2603.04127/) (2026)
- [Stream-CQSA: Exact Out-of-Memory Recovery for Attention](https://mlsanity.com/paper/2604.20819/) (2026)
- [ELSA: Exact Linear-Scan Attention for Fast and Memory-Light Vision Transformers](https://mlsanity.com/paper/2604.23798/) (2026)
- [The Life of a Token: from Words to Bits on the Wire](https://mlsanity.com/paper/2609.19924/) (2026)
- [FlashBoB: I/O-Efficient Exact Backward-over-Backward for Softmax Attention](https://mlsanity.com/paper/2609.24089/) (2026)
- [FlashMask-3: Efficient and Expressive Mask-Aware Distributed Attention](https://mlsanity.com/paper/2610.17901/) (2026)

## Abstract (arXiv)

FlashAttention (Dao, 2023) effectively reduces the quadratic peak memory usage to linear in training transformer-based large language models (LLMs) on a single GPU. In this paper, we introduce DISTFLASHATTN, a distributed memory-efficient attention mechanism optimized for long-context LLMs training. We propose three key techniques: token-level workload balancing, overlapping key-value communication, and a rematerialization-aware gradient checkpointing algorithm. We evaluate DISTFLASHATTN on Llama-7B and variants with sequence lengths from 32K to 512K. DISTFLASHATTN achieves 8x longer sequences, 4.45 - 5.64x speedup compared to Ring Self-Attention, 2 - 8x longer sequences, 1.24 - 2.01x speedup compared to Megatron-LM with FlashAttention. It achieves 1.67x and 1.26 - 1.88x speedup compared to recent Ring Attention and DeepSpeed-Ulysses. Code is available at https://github.com/RulinShao/LightSeq.