--- title: 'LASP: Linear Attention Sequence Parallel' url: https://www.emergentmind.com/topics/linear-attention-sequence-parallel-lasp type: topic --- # LASP: Linear Attention Sequence Parallel Linear Attention Sequence Parallel (LASP) encompasses a family of algorithmic and systems techniques for efficiently parallelizing linear attention models across sequence length in distributed settings. Linear attention—the replacement of the quadratic softmax attention with a kernel-based, linear-complexity form—enables scalable sequence modeling, but practical deployment over ultra-long contexts and large clusters depends critically on tailored parallelization and communication strategies. LASP and its derivatives systematically address the bottlenecks of memory, compute, and inter-device communication in this context, with a range of innovations extending from ring-style and all-gather protocols to contemporary zero-overhead primitives. ## 1. Motivation and Problem Context LASP is motivated by the need to train large-scale models (notably, LLMs and long-context sequence models) whose sequence lengths far exceed the memory capacity of a single device. While traditional sequence parallelism (SP) techniques exist for softmax attention-based architectures, their communication patterns and memory footprints are not optimal for models with linear attention. The fundamental observation is that linear attention mechanisms admit a “right-product kernel trick” permitting recurrence-style update rules and greatly reduced intermediate state. Efficiently distributing such models over clusters, while minimizing per-device memory and keeping communication independent of sequence length, is the core challenge LASP addresses [2404.02882]. ## 2. Core Algorithmic Techniques ### 2.1 Right-Product Kernel Decomposition In linear attention, the output at token $s$ can be computed as: $$ o_s^\top = q_s^\top \cdot \mathrm{KV}_s, \qquad \mathrm{KV}_s = \sum_{i=1}^s \lambda^{s-i} k_i v_i^\top $$ or, with chunking and right-product association, $$ \mathrm{KV}_t = \lambda^{C} \cdot \mathrm{KV}_{t-1} + (\lambda^{C} \Lambda^{-1} K_t)^\top V_t $$ where $C$ is the chunk length and $\Lambda = \mathrm{diag}(\lambda, \lambda^2, ..., \lambda^{C})$. This structure allows sequence-parallel implementations to exchange only fixed-size state tensors (e.g., $d\times d$), regardless of sequence length [2404.02882, 2502.07563]. ### 2.2 Parallelization Patterns Two main approaches have evolved: - **Ring-style Point-to-Point (P2P):** Each GPU computes intra-chunk outputs, receives the prior chunk’s state from its left neighbor, performs an inter-chunk update, then sends the new state rightward [2404.02882]. - **AllGather-based (LASP-2):** Each GPU computes an intermediate memory state (e.g., $M_t = K_t V_t$ for chunk $t$), and all memory slices are gathered and reduced via one all-gather per iteration. This allows computation/communication overlap [2502.07563]. These mechanisms are designed to ensure the volume of communicated data per device is independent of sequence length, depending only on model width and number of layers, and to allow scalable parallel decompositions. ### 2.3 Kernel Fusion and State Caching To further improve efficiency, LASP systems often aggressively fuse intra-chunk and inter-chunk computations into single, custom GPU kernels and implement intermediate memory state caching (to avoid recomputation during backward passes). This results in hardware-optimized runtime and reduced global memory traffic [2404.02882, 2312.06635]. ## 3. Communication, Scalability, and Performance ### 3.1 Quantitative Communication Efficiency A central goal is for the communication overhead to remain constant (per layer: $O(Bd^2/h)$, with $B$ batch size, $d$ model dim, $h$ heads), independent of total sequence length. In contrast, traditional SP methods donate per-device workloads (and thus communication) scaling with sequence length. Empirical studies show that LASP’s communication cost—and thus throughput—scales linearly or superlinearly as GPUs are added, supporting sequence lengths to 4096K tokens across clusters of 128 GPUs (8$\times$ longer than prior methods) [2404.02882]. LASP-2 improves training speed by 15.2% over the earlier LASP design and 36.6% over Ring Attention with explicit all-gather [2502.07563]. ### 3.2 Hybrid and Heterogeneous Parallelism Proper alignment with batch-level data parallel strategies is critical. LASP is explicitly engineered to be compatible with PyTorch DDP, FSDP, and ZeRO-based optimizations, permitting hybrid parallelism decomposing both batch and sequence dimensions [2404.02882]. FlexSP generalizes LASP by adaptively grouping sequences with similar lengths—a linear programming approach to optimal device assignment—yielding up to 1.98$\times$ end-to-end speedups in realistic LLM corpora, which exhibit highly heterogeneous sequence lengths [2412.01523]. ### 3.3 Theoretical Lower Bound and Further Advances ZeCO demonstrates “zero communication overhead” for linear attention SP by introducing the All-Scan primitive: a blockwise, pipelined receive–scan–send pattern that provides each device exactly and only the initial state it needs. Theoretical analysis shows ZeCO meets the lower bound for data transmission per iteration, avoiding the scaling issues present in all-gather or chained ring protocols. On 256 GPUs and 8M tokens, ZeCO provides about a 60% speedup over previously best-in-class SP [2507.01004]. ## 4. Extensions, Hybrids, and Alternative Models ### 4.1 Hybrid Attention Architectures LASP is extensible to models that blend linear and standard softmax attention. LASP-2H generalizes the all-gather workflow to standard attention: the $K$/$V$ matrices of each chunk are all-gathered and the attention is computed locally on each device, enabling efficient SP across a spectrum of hybrid architectures [2502.07563]. ### 4.2 State Space Models and Linear-MoE The sequence parallelism ideas of LASP are also applicable to modern state space models (SSMs), matrix-valued RNNs, and MoE-based architectures, as described in the ParallelFlow framework [2504.00492] and Linear-MoE system [2503.05447]. These systems exploit parallel scan/flow discretizations, signature kernel techniques, and blockwise or tiled update schemes for efficient scaling across extremely long sequences. ### 4.3 Adaptive Memory and Expressivity Recent approaches, such as log-linear attention [2506.04761] and LoLA [2505.23666], focus on mitigating the limited expressivity and recall of fixed-size state linear attention. While not strictly communication or parallel scaling techniques, these improvements (e.g., hierarchical hidden state growth, sparse-cached KV memories) are synergistic with LASP infrastructure, as they preserve the low communication and memory footprint central to the approach. ## 5. Empirical Evaluation and Practical Impact | SP Method | Comms Scaling | Max Seq. Length (on 128 GPUs) | Throughput Speedup | |-----------------|-----------------|-------------------------------|-----------------------------| | Megatron-SP/DS | O(seq len) | Up to 512K | Baseline | | LASP-1 (Ring) | O(1) | Up to 512K–1M | 2–4× | | LASP-2 (AllGather) | O(1) | Up to 2M–4M | 15–36% over Ring, 4–8× over baseline | | ZeCO (All-Scan) | Optimal (O(1)) | Up to 8M | 60% speedup over LASP-2 | Benchmarks across TransNormerLLM, Linear-Llama3, and MoE/SSM-backed models demonstrate that LASP and its descendants retain near-constant memory and computation per device as sequence length and world size increase, with convergence and loss curves essentially matched to single-device, non-SP baselines [2404.02882, 2502.07563, 2507.01004]. ## 6. Generalization and Future Directions LASP’s foundational insight—a right-product factorization and parallel scan-based communication protocol—has broad applicability. It is suitable for arbitrary linearized sequence modeling methods including kernel-based attention, cosine or polynomial approximations, delta-networks, or SSMs. Emerging research emphasizes extending LASP to handle chunked, bidirectional, and hybrid recurrences (as in LION or log-linear attention) [2502.16249, 2506.04761], and to integrate further adaptive, resource-driven workload assignment as in FlexSP [2412.01523]. Anticipated future developments include deeper system-level optimizations (e.g., deeper kernel fusion, intra/inter-chunk tiling [2503.14376]), theoretical advances linking rough path theory to scan-based algorithms [2504.00492], and deployment in domains beyond NLP, such as large-scale vision and multimodal sequence tasks [2405.18425]. ## 7. Summary Linear Attention Sequence Parallelism (LASP) represents a suite of optimization strategies and system designs to train linear attention-based sequence models on ultra-long contexts over distributed environments. By leveraging the mathematical structure of linear attention for communication minimization, kernel fusion, and compatibility with both batch and hybrid parallelism, LASP achieves scalable, efficient training while retaining or improving model performance. Continuing evolution in this area, exemplified by ZeCO, FlexSP, and the generalization to hybrid and alternative architectures, positions LASP techniques as foundational tools in next-generation large-scale sequence modeling [2404.02882, 2502.07563, 2507.01004, 2412.01523]. Source: https://www.emergentmind.com/topics/linear-attention-sequence-parallel-lasp