Striped Attention: Faster Ring Attention for Causal Transformers

2023/11/15 by William T. Brandon, Brandon, William, Aniruddha Nrusimha +11 · 19 citations
Computer Science · Engineering · #Advanced Neural Network Applications #Computation and Language (cs.CL) #FOS: Computer and information sciences #Ferroelectric and Negative Capacitance Devices #Machine Learning (cs.LG) #Parallel Computing and Optimization Techniques

paper · pdf · doi:10.48550/arxiv.2311.09431

openalex publication_date 2023/11/15 · openalex created_date 2025/10/10 · openalex updated_date 2026/07/28

Abstract

To help address the growing demand for ever-longer sequence lengths in transformer models, Liu et al. recently proposed Ring Attention, an exact attention algorithm capable of overcoming per-device memory bottle- necks by distributing self-attention across multiple devices. In this paper, we study the performance characteristics of Ring Attention in the important special case of causal transformer models, and identify a key workload imbal- ance due to triangular structure of causal attention computations. We propose a simple extension to Ring Attention, which we call Striped Attention to fix this imbalance. Instead of devices having contiguous subsequences, each device has a subset of tokens distributed uniformly throughout the sequence, which we demonstrate leads to more even workloads. In experiments running Striped Attention on A100 GPUs and TPUv4s, we are able to achieve up to 1.45x end-to-end throughput improvements over the original Ring Attention algorithm on causal transformer training at a sequence length of 256k. Furthermore, on 16 TPUv4 chips, we were able to achieve 1.65x speedups at sequence lengths of 786k. We release the code for our experiments as open source

Cited by

• Libra: Taming Attention Workload Skew in Long-Context LLM Training with Bounded Sequence Pool
• Mesh-Attention: A New Communication-Efficient Distributed Attention with Improved Data Locality
• TEMP: A Memory Efficient Physical-aware Tensor Partition-Mapping Framework on Wafer-scale Chips
• Optimizing Long-context LLM Serving via Fine-grained Sequence Parallelism
• MTraining: Distributed Dynamic Sparse Attention for Efficient Ultra-Long Context Training
• TASP: Topology-aware Sequence Parallelism
• SlimPack: Fine-Grained Asymmetric Packing for Balanced and Efficient Variable-Length LLM Training
• Zeppelin: Balancing Variable-length Workloads in Data Parallel Large Model Training
• InfiniPipe: Elastic Pipeline Parallelism for Efficient Variable-Length Long-Context LLM Training
• BurstEngine: an Efficient Distributed Framework for Training Transformers on Extremely Long Sequences of over 1M Tokens
• TokenLake: A Unified Segment-level Prefix Cache Pool for Fine-grained Elastic Long-Context LLM Serving
• BlindSight: Harnessing Sparsity for Efficient Vision-Language Models
• ZeCO: Zero Communication Overhead Sequence Parallelism for Linear Attention
• Arctic Long Sequence Training: Scalable And Efficient Training For Multi-Million Token Sequences
• DDiT: Dynamic Resource Allocation for Diffusion Transformer Model Serving
• Hetu v2: A General and Scalable Deep Learning System with Hierarchical and Heterogeneous Single Program Multiple Data Annotations
• Taming the Titans: A Survey of Efficient LLM Inference Serving
• Efficient Pretraining Length Scaling
• SlimPipe: Memory-Thrifty and Efficient Pipeline Parallelism for Long-Context LLM Training

Related

• FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness
• Transformer Based Linear Attention with Optimized GPU Kernel Implementation
• FuseMax: Leveraging Extended Einsums to Optimize Attention Accelerator Design
• Ring Attention with Blockwise Transformers for Near-Infinite Context
• InAttention: Linear Context Scaling for Transformers