LASP: Linear Attention Sequence Parallelism
Authors: Weigao Sun
Zhen Qin
Dong Li
Xuyang Shen
Yu Qiao
Yiran Zhong
Venue: Accepted by TMLR ()

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 8$\times$ longer than existing SP methods. Code is available at: \url{https://github.com/OpenNLPLab/LASP}.

Published in Transactions on Machine Learning Research (05/2025)
LASP: Linear Attention Sequence Parallelism
Weigao Sun' Shanghai Al Laboratory
sunureignoo@oulllo..cco
Zhen Qin Tap
shenging0501122@@mall.cco
Dong Li Shanghai Al Laboratory
liddalidd@gmail.cc
Xuyang Shen Shanghai Al Laboratory
Yu Qiao
viaoou@@yjla..org.com
Shanghai AI Laboratory
Yiran Zhong Shanghai Al Laboratory
shongyiran@gmail. com
Reviewed OpenReview: https: httpe://openrevieeu net/ forum? id= QUUs
openreview.
Abstract
Sequence parallelism (SP) serves as a prevalent strategy to handle long sequences that exceed the memory limit of single device. However, for linear sequence modeling methods like linear attention, existing SP approaches do not take advantage of their right-product-fisst feature, resulting in sub-optimal communication efficiency and usability. In this paper, we introduce Linear Attention Sequence Parallelism (LASP), efficient SP approach designed for linear attention-based transformer models. Specifically, we design efficient point-to- point ring-style communication mechanism to leverage the right-produet kernel trick of linear attention, which sharply decreases the communication overhead, comparing with existing SP methods. We enhance the computation efficiency LASP by performing kernel fusion and intermediate state caching, making the implementation of LASP hardware-friende on meticulously ensure the compatibility of sequence-lveel LASP with
GPUs. Furthermore, all types of batch-level data parallel methods, which is vital for distributed training 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 8x longer than existing SP methods. Code is available at: http://ggthaa.cccooo/pppaaala//a
1 Introduction
Linear sequence modeling methods (Katharopoulos et al.) 2020 Choromanski et al. 2022 Sun et al. 2025a
inchuding lincar attention (Qin et al. 2024d). state space models Dao Gu 2024) and linear RNN iQin et al, 2024et, are becoming increasingly popular due to their faster training and inference speed and comparable ble modeling performance vanilla Softmax attention- transformer models (Vaswani et al. 2017 Zeng et al,
2022 Touvron et. al. 2023a
Team 2023). The hybrid architecture, which interleaves Softmax attention and
"Equal Contribution Coorrspponding Author

---

Published in Transactions on Machine Learning Research (05/2025)
linear attention Transformer layers, has proven to be effective balance between their respective strengths. This approach has been successfully implemented in large-scale commercial models such as Minimax-01 (Li et al. 2025) and Tencent Hunyuan Turbo-S (Tencent 2025). as well as in smaller-scale hybrid models like
Samba (Ren et al. 2024). Jamba (Lieber et al. 2024).
As the size of large language models (LLMs) increases and sequence lengths extend, the capacity limitations of single GPU's memory become a significant challenge, constraining the maximum sequence length manageable by a large model. To address this, Sequence Parallelism (SP) techniques (Li et al., 2022 Korthikanti et al., 2022) are employed, which partition long sequence into multiple sub-sequences to be processed on separate devices. However, current implementations of SP methods do not fully exploit the right-product advantages of linear-complexity attention mechanisms Qin et al. (2024b). This results in less than optimal parallelism efficiency and reduced usability on linear sequence modeling methods.
In this paper, we present Lincar Attention Sequence Parallelism (LASP) approach for efficient SP on models with linear sequence modeling. Our approach takes linear allention (Katharopoulos et al. 2020) as an instance to design sophisticated point-to-point (P2P) ring-style communication mechanism during both forward and backward among devices within a node across multiple nodes. This design maximizes the utilization of right-product kernel tricks in linear attention, by only exchanging one single intermediate state instead of both of key and value states in other counterparts. Notably, our approach is independent of attention heads partitioning, allowing it to be applied to models with varying numbers or styles of attention heads, such as multi-head, multi-query, and grouped-query attentions. This flexibility exceeds the capabilities of existing SP methods in Megatron-LAt -LM (Shoeybi et al. 2019 Korthikanti et al. 2022) or DeepSpeed (Jacobs al. 2023
Our implementation of LASP incorporates system engineering optimizations such kernel fusion and KV State caching, resulting in significantly enhanced execution efficiency. Furthermore, we have taken great care in ensuring compatibility of LASP with various (sharded) distributed data-parallel (DDP) (Li et al., 2020) training methods during the implementation, which we refer to as the data-sequence hybrid parallelism. Through extensive experiments with linear transformer models of different parameter numbers, cluster sizes, and sequence lengths, we demonstrate the performance and efficiency of LASP when used with different. DDP instances. Specifically, LASP can extend sequence length up 4096K on 128 GPUs, which is 8x longer than existing SP methods.
Our primary contributions can be summarized as follows:
new SP approach called LASP that is designed for linear sequence modeling methods. LASP is able to perform soquence-level distributed training 8x longer sequences than existing SP methods while being significantly faster.
Sequence inngle--iddeennentt communication overhead. Our proposed P2P ring-style communication strategy leverages right-produs kernel trick of linear attention to ensure that the exchanging of linear attention intermediate states is sequence lengtt-independdnt.
GPU friendly implementation. We optimize the execution of LASP on GPU hardware through meticulous system engineering, including kernel fusion and KV State caching.
Data Pata-parallel compatibility. LASP is compatible with all batch-level DDP methods, including PyTorch/Legacy DDP, FSDP, and ZeRO-series optimizers.
2 Related Work
Linear Attention. Linear Transformer models bypass the use of Softmax attention by adopting various
approximation methods (Katharopoulos et al., 2020 Peng et al. 2021 Qin et al. 2022a Shen et al. 2024 instead. The central concept involves using the "kernel trick" to speed up the calculation of the atlention matrix, specifically by multiplying keys and values before tackling the computationally intensive n X n matrix multiplication (Sun et al.] 2025b). For instance, Katharopoulos et al. (2020) use + clu activation function, Qin al. (2022b) utilizes the cosine function imiiate Softmax characteristics. TransNormerLLM (Qin et al. 2024a) proposes Lightning Attention to accelerate linear attention via optimized TO operations, while
2

---

Published in Transactions on Machine Learning Research (05/2025)
Lightning Attention-2 (Qin al. 2024c]
enhances efficiency by separately processing inter- and intra-
block computations. RetNet (Sun et al. 2023)
integrates retention with attention for parallel training and
introduces data-independent gating and a hardware-efficient training algorithm. DeltaNet (Schlag et al. 2021 and its parallelized version (Yang et al. 2024b) improve long-context performance using a della-rule updale. GSA (Zhang et al. 2024), inspired by GLA, incorporales bounded-memory slot control to enhance recall-heavy tasks.
linear-time inference. GLA (Yang et al.) 2023
Despite significant research progress, the adoption of linear attention and similar lincar sequence modeling techniques in commercial large-scale models remains limited. However, some companies have started exploring
their use. For example, Minimax-01 incorporates Lightning Attention (Li et al. 2025). variant of lincar attention; Tencent's Hunyuan Turbo-S employs Mamba2 (Tencent, 2025), & variant of state-space model; ; and Together.Al AI integrates StripedHyena (Poli et al.
2023 a long convolution model.
Memory-Efficient Attention. Rabe & Staats
(021 first employs the online Softmar technique linear memory for attention,
efficiently compute numerically stable attention scores sequentially, resulting
yet still needs quadratic time complexity. While FlashAttention (Dao et al. 2022 Dao
2023 employs tiling
to minimize the number of memory reads/write /writes between GPU's high bandwidth memory (HBM) and on-chip
SRAM to reduce time and memory in the training process, PagedAttention (Kwon et al. 2023 optimizes the utilizat ion of the KV cache memory by reducing waste and enabling adaptable sharing among batched requests during inference. Ring Attention (Liu et al. 2023 reduces memory requirements for Transformer models when handling long sequences by distributing sequences across multiple devices and overlapping the communication of key-value blocks with blockwise attention computation.
Sequence Parallclism. SP widcly used method to train long sequences has been integrated into many large model training frameworks, including Megatron-LM, DeepSpeed, and Colossal-AL. Megatron- LM (Shoeybi et al.) 2019 implements SP along with model (tensor) parallelism (MP) to perform large matrix multiplications on GPUs. However, MP partitions the attention heads, which limits the maximum parallelism degree to be less than the number attention heads. DeepSpeed-Ulyseee Ulysses (Jacobs et al. 2023 uses an all-to-all communication primitive to reduce communication volume, but also partitions atlention heads and faces similar issues as Megatron-LM.
3 Method
3.1
Preliminary
Softmax Attention. Consider the standard attention masking in the transformer architecture, formulated as:
Vaswani et al. 2017 computation with causal
0 Softmax QK /VdMMVV
(1)
where d denotes the hidden dimension. The matrices Q, K, V
E Nxd represent query, key, and value matrices, respectively. These matrices are linear projections of the input X E Rxd i.e., Q = XWQ K XWK, V XWv. The output matrix is denoted as E R. and M E RNNN represents the causal mask matrix. The Softmax(-) operation introduces quadratic time complexity relative the input sequence length N, limiting the scalability of vanilla transformers extended input sequences.
Linear Attention. Linear attention is originally proposed in (Katharopoulos et al. 2020 with the climination of Softmax operation (Vaswani et al. 2017). Qin et al. (2022a) 2024a) propose to replace the Softmax operation with a normalization operation Norm(-), which turns to
O Norm (QK O M)V).
(2)
When considering bidirectional tasks, the above formulation can be simplified as O = Norm (QK V). Then by performing the associativity property of matrix products, it can be mathematically equivalently transformed into a right-product version:
O Norm(Q(K* V))- (V).
(3)
3

---

Published in Transactions on Machine Learning Research (05/2025)
Deston
Desize
Device
Deaion
Sead Acc
Seed Racy
Sens deee
2end Roc
IV
KV
KV
VI
0/+1
K2
X2 Xiet
Figure 1: Visualization of LASP. Left: The chunk-level linear attention computation with a causal mask can be segmented into two categories: intra-chunk and inter-chunk computations. Intra-chunk com- putations, corresponding to the diagonal elements (in diagonal orange boxes) of the mask matrix, utilize traditional left-product. methods. While inter-chunk computations, corresponding to the lower triangular boxes, employ efficient right-product methods for computation. Right: This panel illustrates the P2P communication mechanism employed by LASP. The input sequence X is divided into multiple sub-sequence chunks {.... ,X, ",,,, each processed by different model instances across distinet devices. For each device i, Qú Kü, and Vú are computed from its respective input chunk X,, Notably, the communication operations between devices are designed to complementary in the forward and backward passes. Specifically, in the forward pass, KV matrices are sent from device i to device + 1), and in the backward pass, dKV matrices sent back from device (i + 1) to device i.
This linear attention formulation facilitates recurrent prediction with a computational complexity O(Nd). And the recurrent update of KTV without needing compute the entire attention matrix makes its inference efficient.
While linear complexity offers significant advantages in terms of computational efficiency and memory optimization for linear attention, it still incurs proportional increase in computation and memory utilization on a single GPU as the sequence length N grows. This can lead to memory constraints on a single GPU, such as the 80GB limit in NVIDIA A100, for exceptionally long sequences. The challenge of achieving zero-redundancy (on sequence level) training for such long sequences using linear irrention-based LLMs across GPU clusters remains an open problem. Furthermore, the complexity of addressing this issue in a casual setting further intensifies the challenge. To address this, we propose LASP solution for parallelizing linear attention training at the sequence level, even in casual setting.
3.2 LASP
LASP tiles sequences over the cluster. Follow the oh LASP partitions the input sequences into multiple sub-sequence chunks, distributing these chunks individually across different GPUs. For linear attention in a casual setting, order to fully exploit the advantage of right-product in linear attention, we categorize the attention computation for chunks into two distinct types: intra-chunks and inter-chunks. Intra- chunks involve conventional attention computation, while inter-chunks leverage the kernel tricks associated with linear attention's 's right-product. Further details regarding the intricate mechanisms of LASP in data distribution, forward pass, and backward pass are expounded upon below. A visualization of LASP is presented Fig.
Data Distribution. LASP is designed for training long sequences on lincar transformers in distributed enviromment, achieved by partilioning the input data along ils sequence dimension. In this situation, each GPU within the distributed environment undertakes the training of a subset of sub-sequences, which serves diminish the large memory footprint associated with activation during the training of long sequences. Communication operations are introduced between GPUs to transmit intermediate states. The final trained model assimilates the knowledge derived from the entirety of the long sequences.
4

---

Published in Transactions on Machine Learning Research (05/2025)
Algorithm 1 LASP Data Distribution
1: Input: An input sequence in embedding space X R with sequence length N and hidden dimension d, distributed world size W and sequence parallel size T'.
Seq1
Sub-seg0
Sub-seg1
Sub-segg
Sub-seg
2: Obtain number of sequence parallel groups G = W/T.
Sed0
Sub-seab
Sub-seq!
Sub-sia2
Sub-ing3
3: Obtain sub-sequence length (or chunk size) C = N/T.
4: Get global rank list R = get__lobbl rank().
5: Obtain sequence parallel source rank list Rere [R/T] T.
GPU
6: Along sequence dimension, split. X into T chunks
SP-GroupO
SP-Group1
X2, of size C x d for each.
7: Transfer copics of data chunks X2,", Xx) to GPUs
with rank indices in Rave-
8: Scatter X2,,X2,". .r} from Rare to all ranks in respec- tive sequence parallel groups.
Figure 2: LASP Data Distribution. Left: An example of data distribution with two input sequences and eight GPUs. Right: Complete data distribution algorithm.
For an input sequence of length N, we establish its embedding space representation denoted as X E Nxd with a feature dimension of d. In the LASP framework, X is evenly partitioned into T chunks, where is called the sequence parallel size, which must be divisible by the distributed world size W. These segmented data chunks are subsequently assigned to the respective GPUs. It is essential to note that different sequence parallel groups receive dissimilar data batches. However, within the same group, all data chunks originate from an identical batch of data. A comprehensive depiction of the data distribution process in LASP is provided in Algorithm !
Additionally, an illustrative example of data distribution LASP is presented in Fig. 2 where the distributed world size is characterized by W 8, the sequence parallel size by T 4, the number of sequence parallel groups by G 2, and the sequence parallel source rank list by Rgre [0, 4]. For the first batch SEQO, the input sequence X undergoes partitioning into T chunks {X,,X2,-. xx....,xr along the sequence dimension, subsequently transmitted to the first rank in SP-GROUPO, which corresponds to global rank 0. The data chunks on global rank O are then scattered to global ranks {0, 1, 2, 3) within SP-GROUPO, where each rank only retains a single chunk. The subsequent batch SEQ1 follows a similar manner, being assigned to global ranks {4, 5, 6, 7} within SP-GROUP1.
Forward Pass. To streamline derivations, the Norm(.) operator in Eq. 2 is temporarily omitted. Addi- tionally, we consider normal case where W = T, indicating G = W/T= In this scenario, GPU with rank 0 consolidates all split sub-sequences in batch, subsequently distributing them to all GPUs across the entire distributed world. It is noteworthy that the scenario where the sequence parallel size is not equal to world size is discussed in B.5
We first define kv and KV as the intermediate memory state vector and matrix, respectively. Without loss of generality, we add À as the decay rate in linear attention with casual masking, choosing À yields the ordinary linear attention QQin et al. 2024a Sun al. (002).. In the forward pass of linear attention computation with casual masking, the s-th output can be calculated as
0 = 9. Σ x k,V
(4)
15.
Rewrite in recurrence form, we have
kvo =0 d kv.
o_ (kv.),
(5)
where
kv, = Eiss k.vV
(6)
is the activation memory state in the forward pass with s-th input.
5

---

Published in Transactions on Machine Learning Research (05/2025)
In SP, given data chunk X rank the query, key and value corresponding to X, is Q = X: K, = X,W K. V. = W. Note that we assume W here, their indices are thus equivalent, e.e., i. The output within the t-th chunk can be calculated as
Algorithm 2 LASP Forward Pass
1: Input: input sequence in embedding space X E NRxx with sequence length N and hidden dimension d, distributed world size W, sequence parallel size T = W, decay rate À E Rt.
2: Distribute input sequence X according to Algorithm |
3: Obtain sub-sequence length (or chunk size) C'= -N/T.
4: Initialize mask M RCxC, where Mi = = ffi2j. else
Or,intra [Q,K,) O M.
(7)
My = 0.
5: Initialize A = diag
(2, X² ..... A
The intra-chunk computation has no de- pendencies with other chunks on other GPUs, so it can be calculated parallelized on all ranks in the distributed world. How- ever, this result does not consider the im- pact of the previous 1 ~~ (t-1) chunks the 1-th chunk, which is called an inter- chunk. To calculate inter- churk, let us rearrange Eq. (4) as
6: Initialize activation stale KV = 0 Rddd
7: for chunk E {1,-.., at rank E do
in parallel
8:
Calculate Qr X.W K, = X,WKs, K, V, X,Wv according to its data chunk, of size C x d for each.
9: Compute O Or,intra = [[Q:K,] + M
10: end for
11:
for chunk LE (1,...,,)) at rank i E (1...,, W) do
12: Recv activation KV from rank (i-1).
13: Save KV,-1 as KV; for backward computation.
Σ
Compute O Oriiner AQ, (-1.
0 -9+C
Compute O, O iiiitraa O Oiittrr-
iSa+C
16: Update KV, = XVVV + (A -K)
C+
4s+C
Σ
X
(8)
17: Send activation KV, to rank (i 1).
18: end for
i=C+1
19: return O = [O,], with t E {1, (I....T) ,T).
+ X49
k.V.
ISC
The resulted first part in Eq. corresponds to the computation on previous chunks, and the second part corresponds to the computation on the current chunk. In SP, Eq. (8) can be rewritten in the chunk form
O t,inter AQ,KV
(9)
where KV, = kv,g- Note that the calculation of the inter-chunk for the f-th chunk depends on the activation state of previous (t 1) chunk, i..., KV--1, which calculated on rank (i - 1). Thus a P2P communccation operation Recv should be performed to pull KV,-1 from rank (i (i-1) 1) to rank Then the activation state KV, should be updated for subsequent inter-chunk attention computation at (t + 1)-th chunk. The update rule of KV, at t-th chunk is
IC
KV. = Σ 10- 100--ovv = = Σ 1 k.v +
Σ
1 KV
(10)
SSCC
.5(1-1)C
==(1-1)c-1
KV + (diagfA C-1
I/K,) -V
V
In correspondence the preceding Recv operation, another P2P communictiion operation Send is executed to transmit the acquired KV, in Eq. (1) to the subsequent rank (i + 1) for its inter-chunk computation.
It is noteworthy that in the backward pass, the t-th chunk necessitates KV,-1 as activation to calculate gradients. To minimize communcction operations, we cache KV on High-Bandwidth Memory (HBM) to accelerate computation. Integrating both the intra and inter parts, the final forward output as follows:
O Oiiutra Orjiter
(11)
We present the complete forward pass of TASPP with W T in Algorithm 2

---

Published in Transactions on Machine Learning Research (05/2025)
Backward Pass. For the backward pass, given do,, have [Katharopoulos et al.
2020
dq = do, E gixd dk. T = v, dkv. E IR lxd
dv = k. dkv, E I lxd dkv, = Σ x q,do ddd E Rdxd
(12)
12
By writing dkv, in a recursive form, we have
dkv =0 E Rdxd dkv,-1 = λdkv, + 9c-1do
(13)
In SP, have (Qr, Kt, Vt, Ot, dO,) which corresponds to the t-th sub-sequence chunk rank where E (1. . and i E {1, W). Same with the forward pass, the following derivations assume t = i, T= W.
We first calculate dQ with respective to the 1-th data chunk, which yields:
dQ intta intra [do, V, MJK,
(14)
Since the computation of intra is independent, its calculation can be parallelized on all GPUs. While the calculation of dQt,iteer reflects the inter-dependence of chunks to 1 on chunk In order compute the inter part, we transform Eq. 12 as
C+s Σ
dg do
Σ is++C
'v =o
X + xo s+C
v
(15)
i=+
ISC
The first part in Eq- 15 corresponds to the intra-chunk, while the second part corresponds to the inter-chunk. In SP, we can calculate dQ,,iutte as
dQ, tiinter = AdO,KV
(16)
Note that KV, has already been computed and cached during the forward pass, no communication is required here to obtain KV,.. Benefit from the KV state caching, the calculation of dQ, inter can also be exccuted parallel.
Next, dK within the 1-th chunk can be calculated parallel as
dK iintra = [dOO, ( M
(17)
Then we transform Eq.
as
C
dk. = v Σ
do = J Σ x do + +
= t=s
Σ Xi +2-+
(18)
2
where the term before plus sign corresponds to the intra-chunk, and the term after plus sign corresponds to the inter-chunk. The above equation can be rewrilien in terms of chunks as follow:
dk, intter = - A VVK.V 1+1
(19)
Here a Recv operation is required here to pull dK 1+ from the (t + 1)-th chunk. Then in order to compute dKV for the 1)-th chunk, dKV should be updated as:
(L+1)C
dKv, Σ 1- go = Σ (--(++1) a Ldo,
>> A
(20)
+>(++1)
s=10+1
= =XCaKV + (AQi) dO,

---

Published in Transactions on Machine Learning Research (05/2025)
Then Send operation is performed to push dKV, to rank (i - 1). Finally, for dV, its intra part can be calculated as dV i,intra [(Q.K)) @ "M dO,, Again we transform Eq. 12 as:
C d = k. Σ q = k, - q + +
Σ Xi q.
q.do
(21)
22.
120+1
The first and second terms corresponds to the computation of the intra- and inter-chunks, respectively. In SP, dV, iitte can be calculated as:
dV tinter =
4+1
(22)
Combine the intra and inter part, we obtain the final results of dQ2, dK, and dV,;
dQ, = dQq. inter dK, = iitra + Lntter, dV, = dV Lintra + dV, finter-
(23)
We provide the complete backward pass of LASP in Algorithm in Appendix A.1
3.3 Comparison
In LASP, is important to note that the forward pass requires communication for the KV Rdd state in each lincar attention module layer. The communication volume determined by Bd² /h, where B is the batch size and h is the number of heads. In comparison, Ring Attention also adopts P2P ring-style communication on states K, V E TR V xd which results a communication volume BNd/h. SP in Megatron-LM utilizes all-gather operations Iwice after two layer normalization layers within each transformer layer, and a reduce- scatter operation after the attention and Feedforward Neural Network (FFN) layers. This results in communication volume of Bd + d/Nd/T. DeepSpeed uses all-to-all collective communication (Thakur et al. 2005) for input K, V, and output O of each attention module layer, resulting in communication volume of 4BNd/T.
Table 1: Communication Volume Comparison. Sim-
plified Formulation: we eliminate the common factors Bd for case of comparison.
Full Formulation
Simplified Formulation
Method
LASP
B& /k 2BNd/h
d/h
Ring Attention
2N/h 4N/T
DeepSpeed-Ulysee Ulyssss
4BNd/T
Megatron-SP
2BNd + ABNd/T 2N + 4N/T
Table 1 displays a comparison communication volumes across three frameworks. d/h is the head dimension which is set at 128 as usual (L.n et al. 2020). In practical applications where N/T 2 32, LASP is able achieve the lowest theoretical communication volume. Furthermore, the communication volume of LASP is not impacted by changes in sequence length N or sub-sequence length C, which is huge advantage for SP with very-long sequences across large clusters.
It is worth to note that, although Ring Attention and LASP both use P2P ring-style communication, they have differences lie in both communication and computation sides. Communication: In both forward and backward, Ring Attention involves communicating two states K, V E Vxd, In contrast, LASP only communicates one single state KV E R which does not depend on the sequence length. LASP has a lower ical communication complexity. This makes LASP more efficient, especially in environments with slower interconnects where the commuicctionn-conptaaaiie Attention is specifically designed for standard attention, utilizing left-product manner, i.e., ((QK On the other hand, LASP is specifically tailored for linear attention-like sequence modeling methods, which leverages the right-product. kernel trick (Q(K" (Q) to achieve linear-time complexity.
overlap may not be optimal. Computation: Ring
3.4 System Engineering Optimization
Kernel Fusion. To improve the efficiency of LASP on GPU, we perform kernel fusion both the intra- chunk and inter-chunk computations, and also fused the updates of KV and dKV into the intra-chunk and inter-chunk computations.
8

---

Published in Transactions on Machine Learning Research (05/2025)
KV State Caching. To avoid recomputing activation KV during the backward pass, we choose to store it in the HBM of the GPU right after computing it in the forward pass. During the subsequent backward pass, LASP directly accesses KV for use. It is important to note that the size of the KV activation cached in HBM is d x d, which is not affected by the sequence length N. When the input sequence length N is exceptionally large, the memory usage of KV becomes negligible.
3.5 Hybrid Parallelism
Data-Sequence Hybrid Parallelism. As illustrated in Fig. 2 LASP allows for the specification of a smaller sequence parallel size that is divisible by the distributed world size. This configuration results in the input data being split, along both the batch and sequence dimensions, which is type of hybrid parallelism called data-sequence hybrid parallelism. The ZeROOseries optimizers (Rajbhandari et al. 2020) in DeepSpeed and FSDP (Zhao et al. 2023) in Py PyTorch propose to distribute model states, which include optimizer states, gradients, and model parameters, across all GPUs within the distributed environment. As variants of data parallclism, these techniques seamlessly align with LASP. Furthermore, their focus on minimizing the memory of model states complements LASP's objective of reducing activation memory on each GPU. By integrating these techniques, the training of large models handling long sequence lengths is rendered more practical.
Compatibility with Tensor Parallelism and Pipeline Parallelism. LASP supports both tensor parallelism (TP) and pipeline parallelism (PP). In PP, as exemplified by the GPipe (Kim et al., 2020) scheduling method, the model is initially partitioned across multiple devices, with each device holding a segment of the model. Data within a mini-batch is then divided into micro-batches, which are sequentially fed into the device containing the first segment. Each device processes its micro-batch and forwards the output to the next device the sequence, simultaneously preparing to receive and process the subsequent. micro-batch from the precoding device. This method of pipelining inputs effectively minimizes device idle times. When LASP is integrated with PP, micro-batches are substituted with sub-sequences derived from a mini-batch. Unlike standard PP, each device retains the intermediate states (KV in the forward and dkV in the backward) locally, rather than transmitting them to the next device as typically done in LASP alone. For TP, the integration with LASP is fluid. Linear attention layers utilize TP to segment matrix operations across both intra-chunk and inter-chunk computations.
Hybrid SP on Inter-layer Hybrid Models. The hybrid SP approach for inter-layer hybrid models applics the established Ring Attention SP to softmax attention Transformer layers while using LASP for lincar attention Transformer layers simultaneously. Since these two strategies operate independently within their respective layer types, they do not interfere with each other. This straightforward approach primarily serves as practical application of LASP to hybrid models. We conduct experiments to demonstrate the fcasibility of Hybrid SP in Appendix A.5.1
4 Experiments
We evaluate LASP on two representative lincar attention-based models: TransNormerLLM (TNL) QQnn el al.] 2023b 2024a) and Linear Transformer (Katharopoulos et al. 2020). TNL is the latest large language model purely built upon linear attention, while Linear Transformer is a classical linear transformer model recognized in the community. Our assessment focuses on three key areas: 1) the ability of LASP to scale up sequence length on scaling-out GPUs, 2) the convergence when using LASP, and 3) speed evaluation when using TAAP and its comparison with other SP methods. No activation checkpointing (AC) (Korthikanti le al. 2022 techniques are used in following experiments to reduce activation memory, except experiments in Section 4.3.2 This is because although the adoption of AC will further enables longer sequence lengths, it will cover up the ability of our sequence parallel method LASP. All experiments are conducted on a GPU cluster equipped with 128x A100 80G GPUs Our implementation is built on Metaseq (Zhang et al.] 2022),
a PyTorch-based sequence modeling framework with FairScale (FairScale authors 2021 b integrated. For
more details of hardware and software, and experimental setup, see Appendix A.2 & A.3 Note that when implement other SP methods (e.g., Ring Attentoin, DeepSpeed- instances for the purpose of comparison, we do not use the right-product kernel trick. We maintain the use of each method's original communícation primitives and computational manners as they originally proposed for softmax attention.
and Megatron-SP) on linear attention
9

---

Published in Transactions on Machine Learning Research (05/2025)
TML-1B,
16
DDP
15 GPUs
BPUs
GPUs
GPUs
Throuspuu
64 GPUs
GPUs
128 GPUs
128 GPUs
32K
64K
128K
200
512K
192M4 2048K 40AEK
24
37K
128K
2064
106K 204K
Sequence Length
Bequance Langth
TNL-1B, Backand
(68)
8 1GUUs
8 GPUs GPUs
80
32 CPUs
GPUs
50
GPUs
128 CPUs
128 GPUs
Merrory
30
24
46
8K
18K
12K
64K
120K
259K.
512K T024K 2045K 4250K #192X
10
2K
4K
56
16k
525
120K
230K
512K 1024K 204CK 4005K
Sequencce Length
Secuence Length
Figure 3: Scalability Evaluation LASP on Throughput (tokens/sec) and Memory Usage. Left: Integration of LASP with FSDP backend; Right: Integration of LASP with DDP backend. The TNL-1B model used, with a batch size of across up to 128x A100 80GB GPUs. The sign "x" with a dotted line represents occurring Out of Memory (0OM).
TNL-78
LASP
LASP
Throughout
Ring
Thooughas
DeepSpee Ulysesss
DeepSpeed Ulysses
Megatron-SP
Megatron.SS
32K
64K
250C 512K 1024K 2048K 4090X
2K
32K
120K
512K
1024K
Sequence Length
Sequence
Length
Figure 4: Speed Comparison (tokens/sec) of LASP Against Ring Attention, DeepSpeed- Sanappeed-Ulssses and Megatron-SP. The sign "x' with dotted line represents occurring an Out of Memory (OOM). The evaluation utilizes the TNL-1B and 7B models with a batch size of on 64x A100 80GB GPUs. The parallelism size for these three methods is configured to 64.
4.1 Scalability and Speed Comparison
The scalability results regarding throughput and memory usage with varying sequence lengths and number of GPUs are illustrated in Fig. 3 By using LASP, successfully scale the sequence length up to 4096K using the FSDP backend and 2048K with the DDP backend on a TNL model with 1B parameters, on 128 GPUs. We keep using a fixed batch size
to thoroughly assess the performance of LASP across a range of sequence lengths, from a typical 2K to an exceptionally long 4096K. By keeping the batch size constant at 1, we ensure that the experiment results are directly comparable, with only the sequence length varying.
Importantly, the implementation of LASP allows for linear increase in the maximum sequence length capacity, directly proportional (linear) to the number of GPUs used. For instance, sequence length of 512K can be trained using 16 GPUs, while 64 GPUs (4x) has is able to train 2048K (4x) sequence length. Enabling LASP maintains high throughput level even with more GPUs used. Furthermore, LASP demonstrates consistent scalability performance under both the FSDP and DDP backends. For more quantitat scalability results LASP, see Table I0 Appendix A.5
We furthermore conducted comparisons on TNL 1B and 7B models against existing SP methods: Ring Attention (Liu et al. 2023), DeepSpeed-Ulyssee (Jacobs al., 2023) and Megatron-SP (Korthikanti al. 2022). All results presented in Fig. are obtained 64 GPUs. LASP demonstrates notable enhancement in throughput, primarily due to its efficient communication design that facilitates the exchange of lincar attention intermediate states.
4.2 Convergence
Table 2 presents the convergence results of two linear attention based models: TNL Qin et al. 2024a and Linear Transformer (Katharopoulos et al.) 2020), and one transformer model (LLaMA (Touvron al. 2023a b)) with Softmax attention, evaluated on an byoch-by-epoch basis. The experiments were conducted using the same training corpus: the Pile (Gao et al. 2020 Both linear models has 0.4B parameters,
10

---

Published in Transactions on Machine Learning Research (05/2025)
Table 2: Convergence Performance of LASP. All experiments use 8x A100 80G GPUs, sequence length of 16K, and batch size of 1. The results cover various DDP backends conjunction with LASP. We explore the performance of two linear attention modcls: TransNormrLLM (TNL) and Lincar Transformer, and one transformer model (LLaMA) with Soflmax attention, all with 0.4B parameters, across 50K updates.
Model
Parameters
Method
Loss
Method
Loss
Translormer
0.4B
DDP
3.727
DDP
3.719
LASP DDP
3.715
Legacy DDP
3.709
LASP + Legacy DDP
3.705
TNL Qin et al.
FSDP
3.717
LASP FSDP
3.714
0.4B
2024a
ZeRO-1
3.653
LASP ZeRO-1
3.653
ZeRO-2
3.655
LASP ZeRO-2
3.649
ZeRO-3
3.656
LASP ZeRO-3
3.649
DDP
5.419
LASP DDP
5.408
Legacy DDP
5.425
LASP Legacy DDP
5.413
Linear
FSDP
5.428
LASP + FSDP
5.441
Transformer
0.4B
ZeRO-1
5.114
LASP ZeRO-1
5.118
Katharopoulos et al. 2020
ZeRO-2
5.105
LASP + ZeRO-2
5,120
ZeRO-3
5.110
LASP ZeRO-3 5.123
demonstrated consistent loss values when training with and without LASP. All experiments undergoes 50K steps. The uniform loss convergence across various DDP backends demonstrates that LASP does not negatively affect model convergence.
4.3
Ablation Study
4.3.1
Ablation on System Engineering Optimization
The system engineering optimization techniques, Kernel Fusion and KV State Caching, aimed at improving the practical execution efficiency of LASP. To better understand their effects, we perform ablation studies, and the results are presented in Table 3 We assess the training throughput and memory usage of a 1B TNL model with a batch size of 2 and a sequence length 8K, using A100 GPUs. The findings show that both Kernel Fusion and KV State Caching significantly enhance training throughput, with only a minimal effect memory usage.
Table 3: Ablation on System Engineering Optimizations Techniques Kernel Fusion and KV State Caching. Experiments are conducted on TNL-1B model with a batch size and sequence length of 8K, utilizing 2x A100 GPUs.
Kernel Fusion
KV State Cache
Throughput (tokens/s) /s)
Memory Usage Per GPU (GB)
No
No
37684.4
49.5
Yes
No
44691.0
49.5
No
Yes
41179.6
49.7
Yes
Yes
45915.2
49.6
4.3.2 Ablation on Activation Reducing Methods
LASP effectively reduces activation memory consumption during training on per-GPU basis, and this advantage becomes sequence dimension. Another widely used technique, activation checkpointing (AC), follows fundamentally different strategy but also contributes siguificantly to activation memory reduction. To further analyze their impact, we conduct ablation experiments to evaluate AC, LASP, and their combined effect. The results are summarized in Table
pronounced in larger clusters due to the distributed partitioning along the
The experimental results indicate that when using only DDP and FSDP, the maximum sequence lengths that be trained single node with 8 GPUs 12K and 16K, respectively. Both activation checkpointing
11

---

Published in Transactions on Machine Learning Research (05/2025)
(AC) and LASP substantially extend the maximum sequence length by significantly reducing activation memory consumption per GPU, although with a minor decrease in throughput. A key distinction between the two is that LASP exhibits scalability directly proportional to the number of GPUs, whereas AC does not. By combining LASP with AC, we achiove remarkable maximum sequence lengths of 496K and 768K on a single node using DDP and FSDP backends, respectively. This made possible by the complementary benefits of three techniques: linear attention, AC, and LASP, all of which contribute to efficient training with extremely long imput sequences.
Table 4: Ablation on Activation Reducing Methods. Both DDP and FSDP backends are tested. A single node equipped with 8x A100 80G GPUs is used to train a TNL-1B model.
Maximum Sequence Length
Maximum Sequence Length
Throughput (tokens/sec)
Throughput (tokens/sed
Method
Method
DDP
12K
131286.0
FSDP
16K
145303.6
DDP+AC
64K
117429.5
FSDP+AC
96K
114464.0
DDP+LASP
96K
126829.4
FSDP+LASP
120K
138598.8
DDP+AC+ LASP
496K
100837.8
ASDS+A++LAP
768K
106578.3
5 Discussion
Linear-complexity sequence modeling methods are emerging as important alternatives to traditional trans- formers (using Softmax attention) for next- next-generation foundational models due to their significantly faster training and inference times, coupled with performance that rivals conventional approaches. Recently, the AT community has seen a rapid development of novel lincar-complexity models, gaining considcrable interest. Examples include linear allention models such LransNormerLLM,, state space models (SSM) like Mamba and Jamba, and linear RNN models including RWKV, HGRN, and Griffin. We contend that the LASP design can be scamlessly integrated into most linear-complexity models. To underscore LASP's generalization, we use a generalized form of linear atlention in Appendix A.4 QQn et al. 2024b), demonstrating that other linear-compleiity models can also be accommodated within this the LASP framework.
Moreover, it is important to explore the compatibility of linear attention and LASP with the widely adopted softmax attention. Each of these mechanisms has distinct advantages in different scenarios. Softmax attention is highly effective for modeling short sequences with strong performance but suffers from quadratic complexity concerning sequence length, which limits its scalability to long-context tasks. On the other hand, linear attention and similar linear sequence modeling methods provide significantly better efficiency for long sequences but tend to be less effective in capturing complex dependencies. A practical solution to harness the benefits of both is a hybrid architecture that alternates softmax attention and lincar attention layers within Transformer models. This strategy has already been implemented in large-scale commercial models such as Minimax-01 (Li et al. 2025) and Tencent Hunyuan Turbo-S [Tencent] 2025), as well as in smaller hybrid models like Samba (Ren et al. 2024 Jamba (Lieber et al. 2024 and GetedDeltaNet (Yang et al. 2024a)
6 Conclusion
We presented LASP to effectively address the limitations of cxisting SP methods on lincar-complexity sequence modeling methods by leveraging their right-product features, which significantly euhanced communication and parallelism efficiency. Through the design of an efficient P2P ring-style communication mechanism and elaborated engineering optimizations including kernel fusion and KV state caching, LASP achieved notable reduction in communicatio traffic and improved hardware utilization on GPU clusters. Compatibility with all types of batch-level DDP methods ensured the practicability of LASP for large-scale distributed training with very-long sequences. Our experiments highlighted the advantages of LASP on scalability, speed, memory usage and convergence performance. In specific experimental setup, LASP achieves significant faster sequence-level distributed training speed on a maximum 8x longer sequence length than the out-of-the-bon SP methods.
12

---

Published in Transactions on Machine Learning Research (05/2025)
Broader Impact Statement
This work represents notable advancement in artificial intelligence and machine learning, particularly in improving the efficiency and scalability of linear allention-baset models. LASP enables the processing of much longer sequences compared to existing methods while significantly accelerating computation, making it highly beneficial for tasks like natural language understanding, genomic sequence analysis, and time-scrics lorecasting. However, the enhanced capabilities and efficiency introduced by LASP also raise ethical and societal considerations, such as the potential for misuse in generating persuasive but misleading content or in surveillance applications. Nevertheless, the contributions of LASP reducing computational overhead and energy consumption in training large models may also bring positive environmental impacts.
Acknowledgments
This work is supported by the Shanghai Al Laboratory.
References
Krzysztof Choromanski, Valerii Likhosherstov, David Dohan, Xingyou Song, Andreea Gane, Tamas Sarlos, Peter Hawkins, Jared Davis, Afroz Mohiuddin, Lukasz Kaiser, David Belanger, Lucy Colwell, and Adrian Weller. Rethinking attention with performers, 2022.
Tri Dao. Flashattention-22: Faster attention with better parallelism and work partitioning, arXiv preprint ar arXiv:2307.08691, 2023.
Tri Dao and Albert Gu. Transformers are ssms: Generalized models and efficient algorithms through structured state space duality. arXiv preprint arXiv:2405.2100 21600, 2024.
Tri Dao, Dan Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. Flashattention: Fast and memory-efficiet exact attention with io-awareness. Advances in Neural Information Processing Systems, 35:16344-16359, 2022.
FairScale authors. Fairscale: A general purpose modular pytorch library for high performance and large scale 2021.
Leo Gao, Stella Biderman, Sid Black, Laurence Golding, Travis Hoppe, Charles Foster, Jason Phang, Horace He, Anish Thite, Noa Nabeshima, Shawn Presser, and Connor Leaby. The pile: An 800gb dataset of diverse text for language modeling, 2020.
Albert Gu and Tri Dao. Mamba: Linear-time sequence modeling with selective state spaces. arXiv preprint arXiu:2312.00775 2023.
Albert Gu, Karan Goel, and Christopher Ré. Efficiently modeling long sequences with structured state spaces. The International Conference on Learning Representations (ICLR), 2022.
Ankit Gupla, Albert Gu, and Jonathan Berant. Diagonal state spaces are as ellective as structured state spaces, 2022.
Sam Ade Jacobs, Masahiro Tanaka, Chengming Zhang, Minjia Zhang, Shuaiwen Loon Song, Samyam Rajbhandari, and Yuxiong He. Deepspeed ulysses: System optimizations for enabling training of extreme long sequence transformer models, 2023.
Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, and François Fleuret. Transformers are rons: Fast. autoregressive transformers with lincar attention. In International Conference on Machine Learning, pp- 5156 5165. PMLR, 2020.
Tobias Katsch. Gateloop: Fully data-controlled linear recurrence for sequence modeling. ar preprint arXiv:2311.01197, 2023.
13

---

Published in Transactions on Machine Learning Research (05/2025)
Chiheon Kim, Heungsub Lee, Myungryong Jeong, Woonhyuk Back, Boogeon Yoon, Ildoo Kim, Sungbin Lim, and Sungwoong Kim. torchgpipe: On-the-fly pipeline parallelism for training giant models. 2020.
Vijay Korthikanti, Jared Casper, Sangkug Lym, Lawrence McAfee, Michael Andersch, Mohammad Shocybi, and Bryan Catanzaro. Reducing activation recomputation in large transformer models, 2022.
Woosuk Kwon, Zhuohan Li, Siyuan Zhuang, Ying Sheng, Lianmin Zheng, Cody Hao Yu, Joseph E. Gon- zalez, Hao Zhang, and Ion Stoica. Elficient memory management. for large language model serving with pagodattention, 2023.
Zhenzhong Lan, Mingda Chen, Sebastian Goodman, Kevin Gimpel, Piyush Sharma, and Radu Soricut. Albert: A lite bert for self-supervieed learning of language representations, 2020.
Aonian Li, Bangwei Gong, Bo Yang, Boji Shan, Chang Liu, Cheng Zhu, Chunhao Zhang, Congchao Guo, Da Chen, Dong Li, et al. Minimax-01: Scaling foundation models with lightning attention. arXiv preprint ar arXiv:2501.088 2025.
Shen Li, Yanli Zhao, Rohan Varma, Omkar Salpekar, Pieter Noordhuis, Teng Li, Adam Paszke, Jeff Smith, Brian Vanghan, Pritam Damania, and Soumith Chintala. Pytorch distributed: Experiences on accelerating data parallel training, 2020.
Shenggui Li, Fuzhao Xue, Chaitanya Baranwal, Yongbin Li, and Yang You. Sequence parallelism: Long sequence training from system perspective, 2022.
Opher Lieber, Barak Lenz, Hofit Bata, Gal Cohen, Jhonathan Osin, Itay Dalmedigos, Erez Safahi, Shaked Meirom, Yonatan Belinkov, Shai Shalev-Shwartz, et al. Jamba: A hybrid transformer-mambe language model. arXiv preprint arx:v22403. 19887, 2024.
Hao Liu, Matei Zaharia, and Pieter Abbeel. Ring attention with blockwise transformers for near-infinite context, 2023.
Huanru Henry Mao. Fine-tuning pre-trained transformers into decaying fast weights. In Proceedings of the 2022 Conference on Empirical Methods in Natural Language Processing, pp. 10236-10222 10242, Abu Dhabi, United Arab Emirates, December 2022. Association for Computational Linguistics. URL https //aclanthology. oog.aoooooooooooooooooooooo 2022. mumlp-main.co
Eric Martin and Chris Cundy. Parallelizing lincar recurrent neural nets over sequence length. In International
Conference on Learning Representations, 2018. URL. http://pppereeeeeeeeeeeeorrrrree
Bo Peng, Eric Alcaide, Quentin Anthony, Alon Albalak, Samuel Arcadinho, Stella Biderman, Huanqi Cao, Xin Cheng, Michael Chung, Leon Derczynski, Xingjian Du, Matteo Grella, Kranthi Gv, Xuzheng He, Haowen Hou, Przemyslaw Kazienko, Jan Kocon, Jiaming Kong, Bartlomiej Koptyra, Hayden Lau, Jiaju Lin, Krishna Sri Ipsit Mantri, Ferdinand Mom, Atsushi Saito, Guangyu Song, Xiangru Tang, Johan Wind, Stanislaw Wožniak, Zhenyuan Zhang, Qinghua Zhou, Jian Zhu, and Rui-Jie Zhu. RWKV: Reinventing RNNs for the transformer era. In Houda Bouamor, Juan Pino, and Kalika Bali (eds.), Pindings of the Association for Computational Linguistics: EMNLP 2023, pp. 14048 14077, Singapore, December 2023. Association for Computational Linguistics. doi: 10.18653/v1/22 /t2.//...n................ ceatingg-emmll.936.. URL
findings-emnlp.... 936
Hao Peng, Nikolaos Pappas, Dani Yogatama, Roy Schwartz, Noah A. Smith, and Lingpeng Kong. Random feature attention. In 9th International Conference on Learning Representations, ICLR 2021, Virtual Event, Austria, May 9-7, 2021. OpenReview.nce, 2021. URL
Michael Poli, Jue Wang, Stefano Massaroli, Jeffrey Quesnelle, Ryan Carlow, Eric Nguyen, and Armin Thomas. StripedHyena: Moving Beyond Transformers with Hybrid Signal Processing Models, 2023.
URL https://gittub..cm ott.///ttr/ggggttoo.oooo/erarraa http://hhh...moooeecceee/eeeeeep
Zhen Qin and Yiran Zhong. Accolerating toeplitz neural network with constant-time inference complexity. In Proceedings of the 2023 Conference on Empirical Methods in Natural Language Processing. Association for Computational Linguistics, December 2023.
14

---

Published in Transactions on Machine Learning Research (05/2025)
Zhen Qin, Xiaodong Han, Weixuan Sun, Dongxu Li, Lingpeng Kong, Nick Barnes, and Yiran Zhong. The devil in linear transformer. In Proceedings of the 2022 Conference on Empirical Methods in Natural Language Processing, pp. 7025 7041, Abu Dhabi, United Arab Emirates, December 2022a. Association for Computational Linguistics. URL https: http:::/accaantholog org/2022. 2022. maalp-mmiin..
Zhen Qin, Weixuan Sun, Hui Deng, Dongxu Li, Yunshen Wei, Baohong Lv, Junjic Yan, Lingpeng Kong, and Yiran Zhong. cosformer: Rethinking softmax in attention. In International Conference on Learning Representations, 2022b. URL
Zhen Qin, Xiaodong Han, Weixuan Sun, Bowen He, Dong Li, Dongxu Li, Yuchao Dai, Lingpeng Kong, and Yiran Zhong. Toeplitz neural network for sequence modeling. In The Eleventh International Conference on Learning Representations, 2023a. URL https://opparrevien wttp:///oooumiiii//////aaaaoaaaa
Zhen Qin, Dong Ti, Weigao Sun, Weixnan Sun, Xuyang Shen, Xiaodong Han, Yunshen Wei, Baohong Lv,, Fei Yuan, Xiao Luo, et al. Scaling transnormer to 175 billion paramcters. arXiv preprint ar Xic:2307.14996, 2023b.
Zhen Qin, Weixuan Sun, Kaiyue Lu, Hui Deng, Dongxu Li, Xiaodong Han, Yuchao Dai, Lingpeng Kong, and Yiran Zhong. Lincarized relative positional encoding. Transactions on Machine Learning Research, 2023c.
Zhen Qin, Dong Li, Weigao Sun, Weixuan Sun, Xuyang Shen, Xiaodong Han, Yunshen Wei, Baohong Lv, Xiao Luo, Yu Qiao, and Yiran Zhong. Transnormerllm: A faster and better large language model with improved transnormer, 2024a.
Zhen Qin, Xuyang Shen, Weigao Sun, Dong Li, Stan Birchfield, Richard Hartley, and Yiran Zhong. Unlocking the secrets of linear complexity sequence model from a unified perspective. ar preprint arXiv:2405.1735. 2024b.
Zhen Qin, Weigao Sun, Dong Li, Xuyang Shen, Weixuan Sun, and Yiran Zhong. Lightning attention-2: A free lunch for handling unlimited sequence lengths in large language models. arXiv preprint arXiv:2401.04658, 2024c.
Zhen Qin, Weigao Sun, Dong Li, Xuyang Shen, Weixnan Sun, and Yiran Zhong. Various lengths, constant speed: Efficient language modeling with lightning attention. arXiv preprinl Xiv: 17381, 2024d.
Zhen Qin, Songlin Yang, Weixuan Sun, Xuyang Shen, Dong Li, Weigao Sun, and Yiran Zhong. Hgrn2: Gated linear rnns with state expansion. arXiv preprint arXiv:2404.0700 07904, 2024c.
Zhen Qin, Songlin Yang, and Yiran Zhong. Hierarchically gated recurrent neural network for sequence modeling. Advances in Neural Information Processing Systems, 36, 2024f.
Markus N. Rabe and Charles Staats. Self-attention does not need o(n²) memory. CoRR, abs 2212201000 2021. URL
Samyam Rajbhandari, Jeff Rasley, Olatunji Ruwase, and Yuxiong He. Zero: Memory optimizations toward training trillion parameter models, 2020.
Liliang Ren, Yang Liu, Yadong Lu, Yelong Shen, Chen Liang, and Weizhu Chen. Samba: Simple hybrid state space models for efficient unlimited context language modeling. arXiv preprint arXiv:2406.0755 2024.
Imanol Schlag and Jürgen Schmidhuber. Gated fast weights for associative retrieval. 2018.
Imanol Schlag, Kazuki Irie, and Jürgen Schmidluber. Linear transformers are secretly fast weight programmers. In International Conference on Machine Learning, 2021.
Xuyang Shen, Dong Li, Ruitao Leng, Zhen Qin, Weigao Sun, and Yiran Zhong. Scaling laws for lincar complexity language models. ar preprint xrXiv:2/06116690 2024.
Mohammad Shocybi, Mostofa Patwary, Raul Puri, Patrick LeGresley, Jared Casper, and Bryan Catanzaro. Megatron-Im: Training multi-billion parameter language models using model parallelism. arXiv preprint 2019.
15

---

Published in Transactions on Machine Learning Research (05/2025)
Jimmy T. H. Smith, Andrew Warrington, and Scott W. Linderman. Simplified state space layers for sequence modeling. CoRR, abs/ 220/222208033333 2022. doi:
URL httpt::tt//aat /tttps://doo.oor/10
48550/arXiv. 2208 04933
Weigao Sun, Zhen Qin, Weixuan Sun, Shidi Lı, Dong Li, Xuyang Shen, Yu Qiao, and Yiran Zhong. Co2: Effi- cient distributed training with full communicaiion- overlap, arXiv preprint 2rX1::440.11265, 2024.
Weigao Sun, Disen Lan, Tong Zhn, Xiaoye Qu, and Yu Cheng. Lincar-moc: Linear sequence modeling meets mixture-of-experts. arxiv preprint or urXic:2503.0544 2025a.
Weigao Sun, Yongtuo Liu, Xiaqiang Tang, and Xiaoyu Mo. Sequence accumulation and beyond: Infinite context length on single gpu and large clusters. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 39, pp. 217333,,,3, 2025b.
Yutao Sun, Li Dong, Shaohan Huang, Shuming Ma, Yuqing Xia, Jilong Xue, Jianyong Wang, and Furu Wei. Retentive network: successor to transformer for large language models. orXiv preprint wxXiv:2307006222 2023.
I Team. Internlm: A multilingual language model with progressively enhanced capabilities, 2023.
Tencent. Tencent hunyuan turbo-s.
hww.aatttiimmmmmm.a...amammm turbo-s turbo-s 2025.
Rajeev Thakur, Rolf Rabenseifner, and William Gropp. Optimization of collective communication operations in mpich. The International Journal of High Performance Computing Applications, 19(1):49-60 66, 2005.
Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie- Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, Aurelien Rodriguez, Armand Joulin, Edouard Grave, and Guillaume Lample. Llama: Open and efficient foundation language models. arXiv preprint 2302 13971, 2023a.
Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Nikolay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, Dan Bikel, Lukas Blecher, Cristian Canton Perrer, Moya Chen, Guillem Cucurull, David Esiobu, Jude Fernandes, Jeremy Fu, Wenyin Fu, Brian Fuller, Cynthia Gao, Vedanuj Goswami, Naman Goyal, Anthony Hartshorn, Saghar Hosseini, Rui Hou, Hakan Inan, Marcin Kardas, Viktor Kerkez, Madian Khabsa, Isabel Kloumann, Artem Korenev, Punit Singh Koura, Maric-Anne Lachaux, Thibaut Lavril, Jenya Lee, Diana Liskovich, Yinghai Lu, Yuning Mao, Xavier Martinet, Todor Mihaylov, Pushkar Mishra, Igor Molybog, Yixin Nie, Andrew Poulton, Jeremy Reizenstein, Rashi Rungta, Kalyan Saladi, Alan Schelten, Ruan Silva, Eric Michael Smith, Ranjan Subramanian, Xiaoqing Ellen Tan, Binh Tang, Ross Taylor, Adina Williams, Jian Xiang Kuan, Puxin Xu, Zheng Yan, Ilyann Zarov, Yuchen Zhang, Angela Fan, Melanie Kambadur, Sharan Narang, Aurelien Rodriguez, Robert Stojnic, Sergey Edunov, and Thomas Scialom. Llama 2: Open foundation and fine-tuned chat models, 2023b.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neurol information processing systems, 30, 2017.
Songlin Yang, Bailin Wang, Yikang Shen, Rameswar Panda, and Yoon Kim. Gated linear attention transformers with hardware-effcient training. arXiv preprint wxXio:2312.06633 2023.
Songlin Yang, Jan Kautz, and Ali Hatamizadch. Gated delta networks: Improving mamba2 with delta rulc. arXiv preprint arXi::442.06440 2024a.
Songlin Yang, Bailin Wang. Yu Zhang, Yikang Shen, and Yoon Kim. Parallelizing lincar transformers with the delta rule over sequence length. Xix preprint arXiv:2406. 06400.064444444 2024b.
Aohan Zeng, Xiao Liu, Zhengxiao Du, Zihan Wang, Hanyu Lai, Ming Ding, Zhuoyi Yang, Yifan Xu, Wendi Zheng, Xiao Xia, et al. Glm-130b: An open bilingual pre-trained model. arXiv preprint arXiv:2210.0241 2022.
16

---

Published in Transactions on Machine Learning Research (05/2025)
Susan Zhang, Stephen Roller, Naman Goyal, Mikel Artetxe, Moya Chen, Shuohui Chen, Christopher Dewan, Mona Diab, Xian Li, Xi Victoria Lin, Todor Mihaylov, Myle Ott, Sam Shleifer, Kurt Shuster, Daniel Simig, Pumit Singh Koura, Anjali Sridhar, Tianlu Wang, and Luke Zettlemoyer. Opt: Open pre-trained transformer language models, 2022.
Yu Zhang, Songlin Yang, Ruijic Zhu, Yue Zhang, Leyang Cui, Yiqiao Wang, Bolun Wang, Wei Bi Freda Shi, Bailin Wang, Peng Zhou, and Guohong Fu. Gated slot altention for elficient linear-time sequence modeling. arXiv preprint arXiv:2409.071 2024.
Yanli Zhao, Andrew Gu, Rohan Varma, Liang Luo, Chien-Chin Huang, Min Xu, Less Wright, Hamid Shojanazeri, Myle Ott, Sam Shleifer, et al. Pytorch fsdp: experiences scaling fully sharded data parallel. arXiv preprint ar Xic:2304.11277, 2023.
Beitong Zhou, Jun Liu, Weigao Sun, Ruijuan Chen, Claire Tomlin, and Ye Yuan. pbsgd: Powered stochastic gradient descent methods for accelerated non-convex optimization. In IJCAI, pp. 3258 3266, 2020.
17

---

Published in Transactions on Machine Learning Research (05/2025)
A Appendix
A.1 Backward Pass Algorithm
Algorithm LASP Backward Pass
Input:
Sequence Length N, Distributed world size W, sequence parallel size T, decay rate X E R++, Qu K,, V,, O, dO, RCxd for & E 21.,,.... .T).
Obtain sub-sequence length (or chunk size) C=
if Initialize mask M RCXC, where My =x
else My 0.
4: Initialize A diag X?. ... ; E
5: Initialize dKV = 0 Rdxd
6: for E 21,2. ".,,, at rank e f1,2, W] in parallel do
7: Compute dQttintra = [(dO, MjK,
Compute dQ inter AdO,KV 1-1-
Compute dK,,iinex (dO,V. O M
Q,
10:
Compute dV t,ittra [QQ,, @ M] dO,-
11: end for
12:
for E (T..... 2, at rank i E ( ... 22,1} do
13:
Recv activation dKV++1 from rank 1).
14:
Compute ,in,r = (X A VV 4+1
15:
Compute dV tiiter = (A A KK) dKV++1 c+1.
16:
Load KV, as KV, on rank i. Combine intra- and inter-chunks of dQ,: dV,
17:
dQ: intra + dQ, inters
dk, dK, rintra + dK, tinttrss
dV, dV, intra + dV, iiinte-
18: Compute dKV, = XcdV + (AQ:) do..
19: Send activation dKV, to rank i.
20: end for
21: return dQ [dQ,), dK [dK,), dV [dV,), with & E (1,2,...,T)
A.2 Hardware and Software
Hardware. Our experimental configuration involves maximum of 16x DGX-A100 servers, each equipped with 8x A100 GPUs, these GPUs are interconnected through NVSwitch, ensuring inter-GPU bandwidth employ RoCE (RDMA over Converged Ethernet) technology,
of 600GBps. For inter-node communication, utilizing 8 RoCE RDMA adapters in each server. This setup facilitates efficient inter-server communication with a bandwidth capacity of 800Gbps.
Software. Experiments are implemented in PyTorc 2.1.1 and Triton 2.0.0 with CUDA 11.7, cuDNN 8.0, and NCCL 2.14.3. Our algorithm is developed upon Metaseq and DeepSpeed.
A.3 Experimental Setup
The training configuration is set with specific hyperparameters: a learning rate of 0.0005 to control the optimization step size, a cap of 50,000 updates to define the training duration, and a 2,000-update warmup period to stabilize early training by gradually adjusting the learning rate. Additionally, weight decay rate of 0.01 is used for regularization to avoid over-fitting (Sun al.) 2024). The Adam optimizer, with beta values of 0.9 and 0.999, is chosen for managing the momentum and scaling of gradients, aiding in effective and stable training convergence (Zhou et al. 2020). Different DDP backends, including Torch DDP (abbr.
18

---

Published in Transactions on Machine Learning Research (05/2025)
DDP), Legacy DDP, FSDP, ZeRO-serics, selected in experiments for cross-validation of compatibility with LASP.
A.4 Generalization of LASP
While LASP initially inspired by linear attention mechanisms, we aim to show its broader applicability to various linear sequence modeling approaches. This section investigates the generalization of LASP through both theoretical analysis and empirical validation.
In the theoretical aspect, we first define the following terms: Memory State m, E Rg Input State E R Expand State C₁ E R Oscillation State O, E R Shrink State E R& and write a general form of recurrent memory as (Qin et al. 2024b
m. 0 t.i.
(24)
which is general form of the recurrence form of Linear Attention in Eq. 5 with specified o, and e,:
kv. = Xkv,1 + k.v.
(25)
The design of LASP can be seamlessly applied to models which is able to be generally expressed by Eq. (24] These models include: $4 (Gu et al. 20222) S5 (Smith et al. 2002), DSS (Gupta et al. 2022) TNN (Qin et al.
2023a Lincar Attention Katharopoulos et al.; 2020). TNL iQin et al., 2004a). Ret.et Sun et al. 2023
Mamba (Gu & Dao 2023). RWKV-4 (Peng el al. 2023 Cosformer (Qin et al. 2022b), Lrpe (Qin et al.
2023ch, GLA Yang et al.
2023), GateLoop (Katsch 2023), DUR (Mao 2022 GFW (Schlag & Schmidhuber
2018), HGRN (Qin al. 2024f cb, and LRN (Martin & Cundy
2018 We list all these models and their
corresponding elements in Table
Table 5: Checklist for Typical Sehear-Complexity Sequence Modeling Methods within the Defined General Form. For each method, the following states are outlined: Input State, Expand State, Oscillation Stale, Shrink State, and Memory Stale. If the slate directly linked to the inpul sequence, the subscript is emphasized. Note that use 1 (k) R* where (k) = for j 1, l,...,k, and (kdd 1 14) B
Method
Input i, Expand e, Oscillation o. Shrink st
Memory & d
X;
A A
S5
B
C
k d
DSS
X;
B
a1;
C
k x d
TNN
B
A
C
k d
Linear Attention
B:
(kd)
C₁
k d
TNL/RetNet
X.
B.
(a)
C
k X d k d
Mamba
B.
A,
C
RWKV4
X:
exp(k:)
exp(-w)
C
1 X 1 k
Cosformer
X;
B.
exp(t0) ((d)
C
exp(i0) (a)
LRPE
X:
B.
C.
k X d k
GLA/ Gateeoop
X-
B.
g
C.
DUR/GFW
B,
gig.
C,
k d
HGRN/LRN
X; 1-At
A:
C.
1 1
We also give the complete explanation for each modeling method as below.
S4. In S4, we obtain 4, e Rd through linear projection from input x, and A E ki*k B,C R through SSM parameterization. The calculation as follows:
m, Amp-1 + BuT, m C.
Note that the original definition of S4 is defined channel-wise mappings fú, 1,..., . of Rax1 Rex1
19

---

Published in Transactions on Machine Learning Research (05/2025)
S5. The recurrence equation of S5 is the same as S4, with the only difference being the direct definition of the mapping Roxd IRaxd and B, C E IR
DSS. The recurrence equation of DSS is same as S4/S5, with the only difference being the direct definition of the mapping Raxd + Raxd and B, C Raxd A = Diaga E Blxk
2023 TNN can be losslessly converted SSM, where C -(kd) e R B R Diag)1..... Ak Rkxk get un from It through linear projection, and it can be
TNN. According QQin & Zhong
expressed recursive formula:
m Am Buy yt m, C.
Linear Attention. In Lincar Attention, obtain qucry a R$ key be R value 2/ E Rd from the input E Rd through linear projection, and recursively calculation is as follows:
kv = kv + ke = kv 9t-
TNL/RetNet. TNL/RetNet is a form of Linear Attention with exponential decay and the method for getting Qt, kz, Ut the same those in Linear Attention, and lambda is a predefined parameter that camnot be learned. Its recursive calculation is:
kv Akv-1 + kev, V1 kc, qt
Mamba. Mamba be seen data-dependent S4. It uses the similar method to get u, A, C, the Ac, Bt, Ct are computed throuth and A, B, C. Its recurrence equation is defined as:
m = A, @ m4-1 Buu Yt =m Cr. C.
RWKV-4. In RWKV-4, we get rı, ki ot through linear projection from input æt and learnable weight. Ignoring the denominator of RWKV-4, the recurrence equation can be simplified as:
exp(-w)m;-- + exp(k:)v Us = =m,
Similar to S4, RWKV4 uses channel-wise mapping Si, = 1,...,d d of mRax1 Raxi
Cosformer. In Cosformer, we obtain query E IR key Kt e RR, value E IRd from the input æ; E IRd and a prefined (not learnable). Then recursively calculate as follows:
kv; = exp((())kkk + hout Yt = Rel(ku]] [g.
Lrpe. In Lrpe, we obtain query qe E R& key hi E RA, value os E Rdd from the input de E Rd, 0 as a learnable weight and recursively calculate as follows:
kv = Akv;-1 + kbd A =
exp(idx)), d1 = Rel[ku]] qr.
GLA/ CIA/GateLoop. In GLA/GateLoop, we obtain query de R key ke E IR*, value is Rd decay gr e R from the input 3, Rdd and recursively calculate follows:
kv, Diag(g; kur kev, We kv, gc-
DUR/GFW In DUR/GFW, we obtain query dt E R key be I& value us Rd, decay ge E Ri g: gd
from the input 24 Rd and recursively calculate follows:
kv, (grg. T) ④ 1 kev, B1 [ku]] 9t-
20

---

Published in Transactions on Machine Learning Research (05/2025)
HGRN/LRN In HGRN/LRN, obtain output gate or RR, forget gate E R imput state in E R from the input E RR, and recursively calculate as follows:
h1 = ft o 1 + (1 (1 fe) f 1/t h4 Or.
Similar S4, HGRN/LRN use channel-wise mapping Je, = 1, ,,...,d d of R Rax1
To empirically validate the generalization of LASP, we adopt the experimental setup from Table 2 and apply LASP to three additional linear sequence modeling methods listed in Table [ namely Cosformer, RetNet, and Mamba. The convergence results, presented in Table 0 demonstrate that LASP does not negatively affect convergence and achieves performance on par with the basclines.
Table 6: Convergence Results LASP on Cosformer, RetNet and Mamba. TNL with 0.4B
parameters are lested with a batch size of 2 and sequence length 16K.
Model
Parameters
Method
Loss
Method
Loss
Cosformer
0.4B
DDP
4.001
LASP+DDP
4.005
Cosformer
0.4B
ZeRO-1
4.013
LASP+ZeRO4
3.969
RetNet
0.4B
DDP
4.312
LASP+DDP
4.306
RetNet
0.4B
ZeRO-1
4.312
LASP+ZeRO-1
4.309
Mamba
0.4B
DDP
4.116
LASP+DDP
4.122
Mamba
0.4B
ZeRO-1
4.108
LASP+ZeRO-1
4.110
A.5
Additional Experiment Results
A.5.1 Hybrid SP Results on Hybrid Models
We perform a small-scalc experiment (8x A100, 1B parameters, DDP backend) evaluate the fcasibility of the hybrid SP approach for hybrid models. this setup, a "1/4 hybrid" model denotes a configuration where one out of every four layers is softmax attention Transformer layer. Using "S" to represent softmax attention and "L" for linear attention, a 16-layer "1/4 hybrid" model follows the pattern "LLISLLISLIISLAA results in Table 7 demonstrate that hybrid SP effectively extends the maximum trainable sequence length for both TNL and Linear Transformer, while incurring only a slight reduction in training speed.
The
Table Hybrid SP Results on Inter-layer Hybrid Models. "1/4 hybrid" refers a model where one out of every four layers is a softmax attention Transformer layer. Maximum Sequence Length and Throughput (tokens/scc) are reported.
Model
Method
Maximum Sequence Length
Throughput
1/4 Hybrid TNL
No SP
12K
128684
1/4 Hybrid TNL
Hybrid SP Solution No SP
90K
125397
1/4 Hybrid Linear Transformer
12K
129253
1/4 Hybrid Linear Transformer
Hybrid SP Solution
90K
125883
A.5.2 Evaluation Results on Downstream Tasks
We conduct an experiment with extended training duration 300K steps (which consumes 40B tokens) to assess the performance of LASP, and its evaluation results on downstream tasks. Both TNL and Lincar Transformer with 0.4B parameters are investigated. We evaluate the performance of the trained models on multiple downstream benchmarks, including PIQA, HellaSwag (HS), WinoGrande (WG), ARC-E, ARC-C, OBQA, and CSR-AVG. The results are presented in the Tables and LASP does not negatively affect downstream task performance.
A.5.3 Quantitative Scalability Results
See Table 10 in next page.
21

---

Published in Transactions on Machine Learning Research (05/2025)
Table 8: Convergence Results of LASP with Extended 300K Steps. Both TNL and Linear Transformer with 0.4B parameters are tested with a batch size and sequence length of 16K.
Model
Parameters
Steps
Method
Loss
PPL
Method
Loss
PPL
TNL
04BB
300K
DDP
3.218
9.318
LASP+DDP
3.218
9.321
Linear
0.4B
300K
DDP
4.164
17.972
LASP+DDP
4.145
17.730
Transformer
Table 9: Evaluation Results on Downstream Tasks. HS: HellaSwag, WG: WinoGrande. A higher score
indicates better performance.
Model Method
Tokens
PIQA
HS
WG
ARC-E ARC-C OBQA
CSR-AVG
TNL
DDP
40B
55.71
28.21
51.30
28.87
23.72
26.00
35.64
TNL
DDSP+DDP
40B
54.30
28.17
51.54
31.27
24.06
29.60
36.49
Linear
DDP
40B
52.18
25.68
49.80
26.81
25.60
26.40
34.93
Transformer
Linear
LASP+DDP
40B
52.18
26.07
49.25
26.22
26.71
27.00
35.44
Transformer
22

---

Published in Transactions on Machine Learning Research (05/2025)
Table 10: Quantitative Scalability Results of LASP on Throughput (tokens/sed /sc)) and Memory Usage Per GPU (GB). Experiments are performed on TNL-1B, scaling sequence length from 2K to 4096K with a batch size of Both DDP and FSDP backends are tested.
LASP DDP
LASP + FSDP
Sequence Length
GPUs
Throughput
Memory
Throughput
Memory
16
1893.3
22.5
1780.5
6.9
32
1645.4
22.5
1671.2
6.6
2K
64
1639.7
22.5
1589.8
6.4
128
1610.9
22.5
1566.2
6.2
16
3686.9
22.5
3519.9
6.9
32
3458.4
22.5
3304.4
6.6
4K
64
3245.3
22.5
3152.2
6.4
128
3211.5
22.5
3075.7
6.2
16
7076.9
22.5
6924.8
6.9
32
7319.3
22.5
6472.9
6.6
8K
64
6869.1
22.5
6459.4
6.4
128
6793.6
22.5
6398.4
6.2
16
14036.8
22.5
13513.7
6.9
32
14671.7
22.5
12978.9
6.6
16K
64
13828.6
22.5
12569.4 1
6.4
128
13484.5
22.5
12184.5
6.2
16
28354.6
24.4
25727.2
6.9
32
27863.6
22.3
26646.4
6.6
32K
64
25275.9
22.5
25201.4
6.4
128
24523.8
22.5
25638.9
6.2
52993.1
28.3
48542.8
11
32
53393.2
24.1
19648.6
6.6
64K
64
52024.2
22.5
49780.5
6.4
128
51983.3
22.5
49833.3
6.2
16
107682
36.1
84901.9
19
32
93371.5
28.3
92718.8
10.6
128K
64
100046
24.4
96771.6
6.4
128
95828.5
22.5
98975.9
6.2
16
202057
51.7
136765
35.2
32
190675
36.1
159326
18.7
256K
64
193341
28.3
170996
10.4
128
187317.7
24.4
178628.4
6.3
16
OOM
OOM
201791
67.5
32
323596
51.7
250663
34.8
512K
64
304366
36.1
284803
18.5
128
295128.5
28.3
298755
10.1
16
OOM
0OM
0OM
0OM
32
OOM
0OM
358478
67.1
1024K
64
523119
51.7
437728
34.6
128
508383
36.1
459794
18.2
16
OOM
0OM
0OM
0OM
32
OOM
0OM
0OM
0OM
2048K
64
OOM
0OM
585326
66.9
128
658432
51.7
597953
33.8
16
OOM
OOM
OOM
0OM
32
OOM
0OM
0OM
OOM
4096K
64
OOM
0OM
0OM
00M
128
OOM
OOM
792705
66.2
23


Reviews:
[{"review_id": "1zByM0T3af", "reviewer": "Reviewer_DPgT", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths\n- Provides improved efficiency for linear attention-based models by leveraging the right-product kernel trick to reduce communication overhead while incorporating hardware optimizations such as kernel fusion and state caching.\n- Scales on 8x longer sequences than Megatron-SP and about 4x than Deepspeed-Ulysses. When using 128 GPUs it can fit sequences of about 4M length, which is remarkable.\n- Follows an implementation that is independent of attention heads partitioning, hardware-friendly for GPUs and compatible with different types of linear attention architectures (multi-head, multi-query or grouped-query).\n- Evaluates the convergence in terms of training loss but also downstream performance that showcases that the approach is effective and sound.\n\nWeaknesses\n- Even though linear attention-based models can scale efficiently on long sequences, they tend to underperform full attention methods in practical applications. Consequently, the scope of this work is limited, as the proposed approach is specifically designed for a particular type of linear attention methods (kernel-based). \n- The hardware-friendly implementation has increased complexity as it involves several optimizations and it may not be easily reproducible, limiting the impact and adoption of this work.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper introduces a sequence parallelism approach specifically tailored for linear attention-based transformer models. The key characteristic of the approach is that its communication overhead is not dependent on the sequence length by leveraging peer-to-peer ring-style communication and the right-product kernel trick of linear attention. The implementation is hardware-friendly for GPUs (by using fused kernels, state caching) and supports all types of batch-level data parallel methods. Experiments covering two linear attention models shows that that the proposed approach can converge to similar loss and scale up to 8x longer sequences than other sequence parallelism approaches such as DeepSpeed-Ulysses and Megatron-SP.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "Here are my recommendations for the recommendation of acceptance:\n1. Enhance the accessibility of the work by providing additional documentation on the implementation or sharing an open-source reference implementation.\n2. Improve the motivation of the work by emphasizing kernel-based attention mechanisms that closely resemble softmax attention and exhibit comparable performance (for instance, Performer, RFA), and hybrid attention mechanisms that are closer to state-of-the-art.\n3. Provide a proof-of-concept demonstrating the adaptation of the proposed approach to other kinds of attention-based transformer models beyond those that use kernel-based attention.\n4. Consider extending the work to incorporate hybrid transformers that utilize both linear and softmax attention, which are competitive with state-of-the-art models.\n\nThe first three suggestions are crucial, and the fourth suggestion would further strengthen the work in my opinion.", "broader_impact_concerns": "There is a broader impact statement that covers high-level concerns and implications.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "aQx58S0JQ7", "reviewer": "Reviewer_3Q8p", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "I find this work practically valuable, particularly given its high scalability. The efficient parallel computation structure makes it well-suited for handling the growing size of input tokens in large-scale models.\n\nThe P2P ring-style communication mechanism is a simple yet highly effective approach that significantly reduces computational complexity while maintaining performance.\n\nThe experimental results are promising, and the method appears to be compatible with recent parallelization techniques.\n\nLimitations and Concerns\n\nHowever, I also observed some limitations in this work:\n1. Dependence on Linear Attention: While linear attention is computationally efficient, state-of-the-art (SOTA) models primarily rely on Softmax-based nonlinear attention due to its superior performance in real-world applications. This raises concerns about the applicability of LASP to mainstream transformer architectures.\n\n2. Discussion on Compatibility with Softmax Attention: To address this limitation, the authors should discuss whether and how LASP could be adapted for Softmax-based transformers. Alternatively, it would be beneficial to highlight specific domains where linear attention models can be more advantageous compared to Softmax-based architectures.\n\n3. Need for More Real-World Task Evaluations: Additional experiments on practical downstream tasks would provide a clearer understanding of the real-world impact of LASP. Evaluations on tasks such as language modeling, machine translation, or large-scale retrieval could further strengthen the claims of efficiency and scalability.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper presents an efficient sequence parallelism approach for linear attention-based transformer models, such as : TransNormerLLM (TNL) and Linear Transformer. The proposed method, Linear Attention Sequence Parallelism (LASP), introduces a ring-style communication mechanism that significantly reduces communication overhead in distributed computations. The implementation is compatible with various parallelization strategies and demonstrates strong scalability.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "It would be great if the authors could provide additional evidence, particularly through further experiments, to address the concerns outlined above.", "broader_impact_concerns": "None", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "9jPnlE7cL2", "reviewer": "Reviewer_1E6g", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n\n- The authors have done a full-featured implementation, including several optimizations such as kernel fusion and KV state caching to ensure the approach is tested in a maximally efficient scenario.\n- Experimental results look good.\n- The lengthy derivation through pages 5 and 6 is thorough, and indicates that the combination of linear attention and sequence parallelism presented here is non-trivial\n\nWeaknesses:\n\n- Results are presented without activation checkpointing, and it is implied that the two methods may mask one another. The discussion of relevant experiments in the appendix (A.5.3) is of notably lower writing quality than the rest of the paper, and is quite difficult to follow. It feels like we may not be getting an apples to apples throughput comparison in the table either - the numbers didn’t make much sense to me.\n- I feel like the case for linear attention is somewhat overstated in this paper. It feels like the scalability of sequence parallelism is the crucial factor, and throughput is a nice-to-have that may or may not be worth the sacrifices that come with more complex code or changes to the loss surface. It would be nice to see some acknowledgement of the lack of (known) use of linear attention in commercial models.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper describes a method for doing both linear attention and sequence parallelism, in order to enable high throughput processing of very long sequences in a manner that allows sequence length to scale with more processing units. In particular, they propose a novel P2P ring communication method that leverages the right-product kernel trick of linear attention. They provide experiments with different model architectures (TNL, Linear Transformer), different model sizes (1B, 7B) and different back-ends (FSDP, DDP) to show that their LASP method retains the scaling properties of Ring Attention but at higher throughputs.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "It would be nice to have the section on activation checkpointing cleaned up and moved into the main paper. If it warrants a two sentence explanation for why it’s not included and the experiments are done, they might as well be in the main text.", "broader_impact_concerns": "None", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: There were initial concerns raised by reviewers, but they have reached unanimous positive ratings after the rebuttal. All reviewers have agreed that this work has clear merit, despite of clear limitations as well. I believe this work provides a useful contribution to the large scale sequence modeling community, and would be a good fit to TMLR.
Decision: Accept (poster)