Abstract
Sequence parallelism (SP) serves as a prevalent strategy to handle long sequences that exceed the memory limit of a single device. However, for linear sequence modeling methods like linear attention, existing SP approaches do not take advantage of their right-product-first feature, resulting in sub-optimal communication efficiency and usability. In this paper, we introduce Linear Attention Sequence Parallelism (LASP), an efficient SP approach designed for linear attention-based transformer models. Specifically, we design an efficient point-to-point ring-style communication mechanism to leverage the right-product kernel trick of linear attention, which sharply decreases the communication overhead, comparing with existing SP methods. We enhance the computation efficiency of LASP by performing kernel fusion and intermediate state caching, making the implementation of LASP hardware-friendly on GPUs. Furthermore, we meticulously ensure the compatibility of sequence-level LASP with all types of batch-level data parallel methods, which is vital for distributed training on large clusters with very-long sequences. We also discuss the generalization of LASP on other linear sequence modeling methods. Extensive experiments on linear attention-based models are conducted with varying sequence lengths from 2K to 4096K. LASP scales sequence length up to 4096K on 128 GPUs, which is 8x longer than existing SP methods.
Introduction
As Large Language Models (LLMs) strive for increasingly long context windows, they encounter a significant hardware barrier: the memory capacity of a single GPU. Sequence Parallelism (SP) addresses this by splitting a single sequence across multiple GPUs. However, traditional SP methods (e.g., Megatron-LM, DeepSpeed) were designed for standard Softmax attention. These methods often require communication that scales with the sequence length ($N$), failing to exploit the mathematical properties of linear attention models. This results in communication bottlenecks that hinder the efficiency these linear models were intended to provide.
The LASP Methodology
The core innovation of LASP is a point-to-point (P2P) ring-style communication mechanism tailored for the "right-product" property of linear attention.
The Right-Product Trick
Standard attention computes $(QK^\top)V$, creating a massive $N \times N$ matrix. In contrast, linear attention can be computed as $Q(K^\top V)$. The term $(K^\top V)$ results in a fixed-size $d \times d$ matrix (where $d$ is the head dimension), regardless of the sequence length.
P2P State Passing
LASP divides the sequence into chunks across GPUs. During the forward pass, each device computes its local "KV state" and passes this small $d \times d$ matrix to the next GPU in the ring. This allows each device to compute its portion of the sequence while maintaining global context without the overhead of transmitting large activation matrices.
System Optimizations
To ensure hardware efficiency, LASP incorporates several optimizations:
- Kernel Fusion: Intra-chunk and inter-chunk computations are fused into single GPU kernels to reduce memory access overhead.
- KV State Caching: Intermediate KV states are cached to accelerate gradient computation during the backward pass.
- Hybrid Compatibility: LASP is designed to be "plug-and-play" with existing data-parallel strategies like PyTorch FSDP and ZeRO-series optimizers.
Key Results and Scalability
LASP's primary advantage is that its communication overhead is independent of sequence length. While competing methods like Ring Attention or DeepSpeed-Ulysses have communication costs that scale with $N$, LASP's cost is tied only to the head dimension $d$.
In empirical evaluations, LASP successfully trained a 1B parameter model with a sequence length of 4096K (4.1 million tokens) on 128 A100 GPUs—an 8x improvement over existing SP methods. Convergence tests on The Pile and downstream benchmarks (PIQA, HellaSwag, ARC, etc.) confirmed that LASP achieves zero accuracy loss compared to non-parallel baselines.
Conclusion
LASP provides a highly scalable solution for training linear-complexity models with massive context windows. By making communication overhead independent of sequence length, it enables the processing of millions of tokens efficiently. While designed for linear attention, the authors demonstrate that LASP's logic generalizes to other modern architectures, including State Space Models (Mamba), RetNet, and Linear RNNs (RWKV).