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

Contribution

• DISTFLASHATTN extends FlashAttention to distributed sequence parallelism while targeting high GPU utilization and low communication overhead.It distributes token chunks across devices while retaining memory-efficient attention benefits.Source: §1 · pp. 1–2

Token-level workload balancing

• Causal attention creates quadratic prefix-dependent work, leaving earlier-token workers idle; load-balanced scheduling routes later-token work to idle workers and doubles throughput versus the unbalanced version.The proposed schedule reduces the asymptotic idle fraction toward zero as worker count grows.Source: §1 · p. 2

Communication overlap

• Communication and computation overlap hides key-value communication inside computation, producing a 1.32× end-to-end speedup over a non-overlapping version.The method pre-fetches tensors using the attention dependencies.Source: §1 · p. 2, Figure 2 · p. 3

Gradient checkpointing

• Rematerialization-aware checkpointing avoids unnecessary FlashAttention recomputation and yields a 1.31× speedup without numerical differences.This addresses checkpointing overhead that becomes important when attention dominates long-sequence computation.Source: §1 · p. 2

Evaluation

• 8× longer sequences are supported, with 1.24–2.01× speedup over Megatron-LM with FlashAttention and 1.67× or 1.26–1.88× speedup over Ring Attention and DeepSpeed-Ulysses.The evaluation uses LLaMA models and compares against four distributed systems.Source: §1 · pp. 2–3

Memory-efficient attention

• Memory-efficient attention reduces single-device peak memory, but naive tensor or pipeline parallel combinations incur excessive communication and do not scale effectively with sequence length.These methods therefore remain limited by the capacity and communication costs of distributed training.Source: §2 · p. 3

Sequence parallelism and ring attention

• Ring Self-Attention and Ring Attention distribute sequence activations but lack either memory-efficient attention support or causal-language-modeling scheduling, while Ring Attention incurs 2× extra computation.DISTFLASHATTN is designed to optimize both concerns.Source: §2 · p. 3
• DeepSpeed Ulysses uses hybrid tensor and sequence parallelism to address distributed attention challenges, providing a direct comparison baseline for DISTFLASHATTN.The paper reports head-to-head comparison in Table 4.Source: §2 · p. 3

Gradient checkpointing

• Existing checkpointing designs do not account for memory-efficient kernels that internally recompute intermediates, whereas DISTFLASHATTN changes checkpoint positions to avoid this recomputation without numerical differences.The approach targets redundant kernel-level recomputation created by conventional layer-boundary checkpoints.Source: §2 · p. 3

3.1 DISTFLASHATTN: distributed memory-efficient attention via sequence parallelism

• DISTFLASHATTN distributes one sequence across workers and retains FlashAttention’s IO-aware memory efficiency for long-context training.Each worker stores activations for only a subsequence, while attention is computed over fetched key-value blocks rather than gathering the full sequence.Source: §3.1 · p. 4

3.2 Load balanced scheduling with communication and computation overlap

• Each worker processes only its causal prefix, creating workload imbalance because workers with larger positions compute more attention blocks.The baseline schedule leaves roughly half of workers idle asymptotically.Source: §3.2 · p. 4
• Load-balanced scheduling assigns completed workers to help heavier workers and reduces idle time asymptotically toward zero for even worker counts.Helpers communicate partial outputs and softmax statistics back to the original worker for correct rescaling.Source: §3.2 · p. 5, Figure 3 · p. 5
• Peer-to-peer key-value communication is prefetched in parallel with attention computation, so the next block is available without blocking the next iteration.Separate GPU streams enable communication and computation to run concurrently.Source: §3.2 · p. 5

3.3 Rematerialization-aware checkpointing strategy

• Conventional layer-boundary checkpointing redundantly recomputes FlashAttention because the kernel already rematerializes intermediates during backward.This redundancy arises from combining outer-system recomputation with FlashAttention’s internal recomputation.Source: §3.3 · pp. 5–6
• Checkpointing FlashAttention outputs lets subsequent-module recomputation and preceding-attention backward share each materialized result, so FlashAttention’s forward is computed once, saving about 7 seconds for a 64K Llama-7B example.Source: §3.3 · p. 6

4.1 Comparison with Megatron-LM on MHA and GQA models

• 1.24× and 1.44× speedup over Megatron-LM is achieved on LLaMA-7B in single-node and cross-node settings.The result combines overlapping communication with rematerialization-aware checkpointing.Source: §4.1 · p. 7, Table 1 · p. 7
• DISTFLASHATTN gains higher speedup on GQA because its key-value communication volume decreases while Megatron-LM’s remains unchanged.Source: §4.1 · p. 7

4.2 Comparison with Megatron-LM on models with irregular or less number of heads

• 1.50× and 2.01× speedups are observed on irregular-head models, while Megatron-LM wastes 45.5% computation padding dummy heads.Source: §4.2 · p. 7, Table 1 · p. 7
• 2× and 8× longer sequences are supported than Megatron-LM alternatives on 16 A100 GPUs, where data and pipeline parallelism constrain memory or utilization.Pipeline stages consume 18–32GB for a 64K sequence in the LLaMA-2H experiment.Source: Table 2 · p. 7, §4.2 · p. 8

4.3 Comparison with Ring Self-Attention (RSA) and Ring Attention

• DISTFLASHATTN supports at least 8× longer sequences and runs 4.45×–5.64× faster than Ring Self-Attention.The speedup includes causal workload balancing, overlapping, and distributed memory-efficient attention.Source: §4.3 · p. 8, Table 3 · p. 8
• 1.67× speedup over Ring Attention and 1.26×–1.88× speedup over DeepSpeed-Ulysses are reported in the evaluated settings.The Ring Attention comparison uses a PyTorch ablation implementation, while DeepSpeed-Ulysses results cover regular- and irregular-head models.Source: §4.3 · p. 8, §4.4 · p. 8, Table 4 · p. 8

4.5 Ablation Study

• Balanced scheduling raises measured saturation from 4.5× to 7.5× over single-GPU FlashAttention as sequence length scales from 4K to 256K.Theoretical expected speedup rises from 4.5× unbalanced to 7.2× balanced.Source: §4.5 · p. 9, Figure 4 · p. 9
• At 128K global sequence length, overlapping reduces communication overhead from 105% to 44%, with only 8% and 1% overhead when communication is fully overlappable.Source: §4.5 · p. 9, Figure 4 · p. 9

• DISTFLASHATTN combines load balancing, overlapped communication, and rematerialization-aware checkpointing for distributed long-context attention.Source: §5 · pp. 9–10

• The eight-worker analysis compares workload and communication across time steps before and after applying causal-modeling load balancing.Figure 5 shows the pre-balancing workload; Figure 6 shows the post-balancing workload.Source: B Load-balancing Algorithm for Causal… · p. 13

• Megatron-LM’s memory consumption is uneven across pipeline stages when training with tensor and pipeline parallelism.Source: C Memory Consumption for Pipeline… · p. 13

• Its communication schedule fetches remote keys and values and exchanges partial results between workers during distributed attention.The ring-style pseudocode includes remote key-value fetches, remote result fetches, and returned partial attention state.Source: D Communication and memory analysis · p. 14
• Pipeline-parallel Megatron-LM shows highly uneven memory consumption across stages at 128K sequence length on 16 A100 40GB GPUs.Source: Table 6 · p. 14
• DISTFLASHATTN reduces communication volume to 3Nd versus Megatron-LM’s 14Nd with gradient checkpointing, a 4.7x reduction.Its forward volume is Nd and backward volume is 2Nd; the rematerialization-aware strategy avoids additional forward communication.Source: D Communication and memory analysis · p. 14
• The communication analysis notes that FSDP communication scales with model-weight size rather than long sequence length and is orthogonal to DISTFLASHATTN.The passage cautions that this theoretical analysis may differ from wall-clock behavior because of implementation factors.Source: D Communication and memory analysis · p. 14
• DISTFLASHATTN distributes sequence tokens across workers and processes local and remotely fetched key-value chunks through causal attention.The forward pseudocode initializes local state, receives remote key-value blocks, updates attention state, and returns output and logsumexp values.Source: D Communication and memory analysis · pp. 13–15

E Implementation Details

• DISTFLASHATTN modifies the FlashAttention2 Triton kernel and implements load balancing, overlapping scheduling, and checkpointing across Python, NCCL PyTorch bindings, and PyTorch.The kernel uses block size 128 and one stage for best performance in the authors’ cluster.Source: E Implementation Details · p. 16
• Evaluations use FSDP so DISTFLASHATTN consumes memory similar to Megatron-LM, while comparisons use the same attention backend.Megatron-LM support avoids materializing causal masks and adds head padding when attention heads cannot be evenly divided across devices.Source: E Implementation Details · pp. 16–17
• Timing results use Adam, 10 warm-up iterations, and averages over 10 additional iterations, with figures showing load balancing across 8 workers and attention-versus-module time breakdowns.The time breakdown is measured with Flash-Attention on a single 40GB A100 GPU in milliseconds.Source: E Implementation Details · p. 17, Figure 6 · p. 17, Figure 7 · p. 17

F Discussion on sparse attention

• The paper focuses on exact attention and presents sparse-attention solutions as possible approaches intended to inspire future work.The discussion specifically mentions load balancing for local sliding windows and global attention.Source: F Discussion on sparse attention · p. 17

Local sliding windows

• For local sliding windows, workload is naturally near-balanced, so workers can ignore non-local attention and fetch only the required neighboring key-value blocks.Implementation changes only the loop’s worker range, such as restricting worker 7 to workers 6–7.Source: F Discussion on sparse attention · p. 18

Global attention

• For global attention, workers can replicate all global tokens or split them evenly and all-gather during computation to reduce memory requirements.Replication avoids all-gathering global tokens at every time step, whereas splitting further reduces memory use.Source: F Discussion on sparse attention · p. 18