ArXiv: 2507.01004

🎯 Pitch

ZeCO reimagines sequence parallelism for linear attention to slash communication costs to their theoretical minimum, enabling training on 1M-token sequences across 64 GPUs in roughly the same time as a 16K sequence on one GPU. It introduces All-Scan, a new pipelined collective that delivers over 60% throughput gains at 8M sequence length by completely eliminating per-device state aggregation overhead.


1. Executive Summary

This paper introduces ZeCO (Zero Communication Overhead) sequence parallelism, a new distributed training method for linear attention models that eliminates the communication bottleneck plaguing existing approaches. Evaluated on 1B-parameter Gated Linear Attention (GLA) models, ZeCO reformulates sequence parallelism around a novel collective communication primitive called All-Scan, which employs a pipelined receive-scan-send pattern to provide each device with precisely the initial operator state it requires while maintaining the theoretically minimum communication volume of just one state tensor per device. Empirically, All-Scan achieves up to 3.9× faster communication than existing methods, enabling ZeCO to deliver over 60% throughput improvement on 256 GPUs with 8M-token sequences and near-linear scaling efficiency from 8 to 256 devices, establishing that sequence parallelism for linear attention can approach the ideal data-parallel throughput upper bound when communication is reduced to its theoretical minimum.

2. Context and Motivation

The Core Problem: Sequence Parallelism Becomes the Bottleneck It's Supposed to Solve

The fundamental tension this paper addresses is a paradox in distributed training of linear attention models: sequence parallelism (SP), which is essential for distributing ultra-long sequences across devices, has itself become the primary bottleneck preventing efficient scaling. This is not merely a theoretical concern — it is the practical reality that limits whether linear attention models can fulfill their architectural promise for long-context training.

To understand why this matters, we need to appreciate the computational landscape of long-context LLM training. Standard Transformer self-attention has quadratic complexity O(L2)O(L^2) in sequence length LL. When scaling from, say, a 4K context to a 128K context, the attention FLOPs explode by over 1000× (Section 1). This makes full pretraining on long sequences prohibitively expensive, forcing practitioners into workarounds: they pre-train on short sequences, then perform a specialized "mid-training" adaptation phase with longer contexts (Abdin et al., 2024), rather than training with long sequences from scratch. The consequence is that models don't fully internalize long-range dependencies during their foundational learning phase, potentially limiting their capabilities for document-level reasoning, multimodal understanding over extended contexts, and retrieval-augmented generation with large knowledge bases.

Linear attention mechanisms offer an algorithmic escape hatch. By replacing the softmax operation with a kernel trick and rearranging the computation order, they achieve O(Ld2)O(Ld^2) complexity where dd is the hidden dimension — linear in sequence length rather than quadratic (Katharopoulos et al., 2020). The key mechanism is compressing the Key-Value cache (which in standard attention grows with sequence length) into a fixed-size hidden state representation, denoted StS_t in the paper's notation, that is updated recurrently as new tokens arrive. This is what enables processing sequences of 1M tokens or more with manageable compute.

But here is where the paradox emerges. To actually train such models on sequences long enough to leverage this linear complexity, we need to distribute the workload across multiple devices — and that distribution strategy is sequence parallelism. The paper identifies a stark reality:

"While linear attention provides these algorithmic advantages, Sequence Parallelism (SP), essential for distributing such computationally intensive workloads, paradoxically becomes a bottleneck that impedes efficient scaling across multiple devices." (Section 1)

In other words, the very mechanism we need to make long-sequence training practical is itself so inefficient that it cancels out much of the benefit of using linear attention in the first place. This is the gap the paper targets: current SP methods suffer from either serial execution bottlenecks or communication overhead that grows with the number of devices, making them fundamentally unscalable for ultra-long sequences.

Why This Problem Is Important: Both Practical and Theoretical Stakes

Practical impact. The inability to efficiently train on long sequences has direct consequences for model quality and deployment economics. When long-context pretraining is relegated to a brief adaptation phase rather than being integrated throughout training, the model's ability to reason over extended contexts may be fundamentally limited. Moreover, the computational inefficiency of existing SP methods means that even when organizations invest in large GPU clusters, they cannot achieve linear throughput scaling — adding more GPUs yields diminishing returns. The paper's Figure 4 demonstrates this concretely: with LASP-1 and LASP-2, per-GPU throughput degrades as more devices are added (from ~38K tokens/sec at 8 GPUs down to ~15K tokens/sec at 256 GPUs for LASP-2 on 16K sequences), whereas ZeCO maintains near-constant per-GPU throughput (~44K down to ~34K), approaching the ideal data-parallel baseline. This isn't just an academic benchmark — it represents real dollar costs and wall-clock time for training runs.

The paper sets an ambitious target: "training a model with a 1M sequence length across 64 devices using ZeCO takes roughly the same time as training with a 16K sequence on a single device" (Abstract). This is the kind of scaling that would make long-context pretraining economically viable at scale.

Theoretical significance. Beyond the practical implications, the problem touches on fundamental systems questions: what is the minimum communication necessary for sequence parallelism in linear attention? Can we design algorithms that achieve that lower bound? The paper positions its All-Scan primitive as not merely an engineering improvement but a theoretically optimal solution — one that achieves the minimum possible communication volume (dk×dvd_k \times d_v per device, independent of device count PP) while also overlapping communication with computation to hide latency. This transforms sequence parallelism from an ad-hoc engineering problem into one with known theoretical limits that can be provably achieved.

Where Existing Approaches Fall Short

The paper identifies three categories of prior work, each with fundamental limitations that ZeCO addresses.

LASP-1: Serial Execution as a Dealbreaker

LASP-1 (Weigao Sun, Qin, et al., 2025) distributes the input sequence into contiguous chunks across devices and has each device compute its local attention state sequentially. The communication pattern is P2P: device pp waits for device p1p-1 to finish computing its final state Sp1S_{p-1}, receives that state, updates it with its own local computation to produce SpS_p, and passes SpS_p to device p+1p+1.

The problem is immediately visible in the paper's Equation (15). The total time for LASP-1 scales as:

TLASPP(PL)=P×(Tideal-SP1(L)+τ(dk×dv))T_{\text{LASP}}^P(PL) = P \times (T_{\text{ideal-SP}}^1(L) + \tau(d_k \times d_v))

where Tideal-SP1(L)T_{\text{ideal-SP}}^1(L) is the time to process a single device's local sequence, and τ(dk×dv)\tau(d_k \times d_v) is the communication time to send one state tensor. The total time grows linearly with PP because devices must execute in strict serial order. At any given moment, only one device is doing useful computation while the others idle. This is the most basic form of inefficiency: adding more devices increases capacity but doesn't increase parallelism, because the dependency chain forces sequential execution.

The consequences are stark in the experimental results. In Table 3, with 16K sequence length and 128 GPUs, LASP-1 takes 113.71ms for a forward-backward pass of the attention operator — nearly 18× slower than the single-device baseline of 6.55ms. This is worse than simply running on one device and accepting the memory limitations. The throughput data in Table 5 confirms this: at 8 GPUs with 16K sequences, LASP-1 achieves 26,428 tokens/sec/GPU, but at 64 GPUs this drops to 12,596 — you're getting less total throughput from 64 GPUs than you'd want from proportional scaling, and far less per-GPU efficiency.

LASP-2: Parallelism at the Cost of Communication Explosion

LASP-2 (Weigao Sun, Lan, et al., 2025; A. Li et al., 2025) fixes the serial bottleneck by replacing P2P communication with an All-Gather operation. Every device collects the local states from all other devices, then independently performs the same scan operation to compute all global states. This enables full parallelism — all devices compute simultaneously — but introduces a new problem: the communication volume scales with PP.

Specifically, each device must receive state tensors from all P1P-1 other devices, resulting in a communication volume of (P1)×dk×dv(P-1) \times d_k \times d_v per device. As the paper notes in Equation (16), the total time becomes:

TLASP-2P(PL)=Tideal-SP1(L)+P×τ(dk×dv)T_{\text{LASP-2}}^P(PL) = T_{\text{ideal-SP}}^1(L) + P \times \tau(d_k \times d_v)

The computation is now parallel (the first term doesn't multiply by PP), but the communication cost grows linearly with device count. At large scale — the paper tests up to 256 GPUs — this All-Gather becomes the dominant cost. Figure 3 shows that at 256 GPUs, All-Gather communication takes 8.51ms versus 2.16ms for All-Scan — a 3.9× gap. And this is just the communication time for the state tensors; it doesn't include the additional computation LASP-2 performs for scan reduction and state updates (Appendix A.2, Table 1 shows extra log(P)De+NDe\log(P)De + NDe computation).

The practical failure mode is visible in Table 5: LASP-2 throughput degrades from 37,812 tokens/sec/GPU at 8 GPUs to 15,196 at 256 GPUs for 16K sequences — a 60% drop. The method scales better than LASP-1 but still falls far short of the ideal linear scaling curve that ZeCO approaches.

Full Attention SP Methods: Irrelevant but Illustrative

The paper also discusses full-attention SP methods like Megatron-LM's context parallelism (Shoeybi et al., 2020) and Ring Attention (Liu, Zaharia, and Abbeel, 2023; Brandon et al., 2023), as well as Ulysses (Jacobs et al., 2023). These are included primarily to show that the communication problem in linear attention SP is fundamentally different — and simpler — than in full attention, which creates an opportunity for an optimal solution.

In full attention, each device needs access to all Key-Value pairs from all other devices to compute self-attention, leading to communication volumes that scale with both sequence length and device count (e.g., 2PLD2PLD for Megatron CP, as shown in Table 1). The computation itself also scales with PP (L2DPL^2DP), making full attention fundamentally unscalable for ultra-long sequences regardless of the communication strategy. Linear attention, by compressing the KV cache into a fixed-size state, reduces the computation per device to LDeLDe — independent of PP — and only requires communicating one state tensor per device. This structural difference means linear attention SP can theoretically achieve near-linear scaling if the communication is handled correctly, which is exactly what ZeCO demonstrates.

A Deeper Issue: The State of the Art Fails to Exploit the Structural Simplicity of Linear Attention

The paper's key intellectual move is recognizing that existing methods don't take advantage of a fundamental property of linear attention: the global state can be decomposed as a linear update from an initial state, as proven in Equation (6) and Appendix A.1. Specifically, for any chunk nn, the global state S(p1)L+nCS_{(p-1)L+nC} can be expressed as:

S(p1)L+nC=(γ~[n]T1)S(p1)L+S[n]S_{(p-1)L+nC} = (\tilde{\gamma}_{[n]}^T \mathbf{1}) \odot S_{(p-1)L} + S_{[n]}

where S(p1)LS_{(p-1)L} is the incoming global state from the previous device, γ~[n]\tilde{\gamma}_{[n]} is the cumulative decay, and S[n]S_{[n]} is the local contribution computed independently. This means each device needs only the final state from its predecessor — not all intermediate states, not all states from all devices. LASP-1 recognizes this (it only sends one state), but serializes the computation. LASP-2 parallelizes the computation, but over-communicates by gathering all states when only one is needed per device.

ZeCO's insight is that you can have both: receive only the one state you need (like LASP-1), but overlap that communication with local computation so that devices don't sit idle waiting (unlike LASP-1). The pipelined All-Scan primitive achieves this by splitting the state tensor into KK blocks and beginning the scan update as soon as the first block arrives, rather than waiting for the full tensor. This is what the theoretical analysis in Equation (12) captures: the communication latency is τ(dk×dv)+(P1)τ(dk×dv)K\tau(d_k \times d_v) + \frac{(P-1)\tau(d_k \times d_v)}{K}, where the second term — representing the non-overlappable boundary overhead — vanishes as KK grows large.

How ZeCO Positions Itself

The paper doesn't position itself as yet another SP variant but as the theoretically optimal solution that previous methods approximated with different compromises. The claim is explicit: ZeCO satisfies two necessary and sufficient conditions for optimality (Section 3.3):

  1. Zero Communication Overhead: Each device communicates exactly one state tensor (dk×dvd_k \times d_v), which is the theoretical minimum — you cannot do sequence parallelism with less information transfer because each device genuinely needs the accumulated state from all previous tokens.
  2. Optimal Extra Cost: The communication is overlapped with computation (intra-chunk diagonal attention in the forward pass), and the auxiliary I/O and computation for SP-specific operations (maintaining cumulative decay vectors γ~[n]\tilde{\gamma}_{[n]}, updating local to global states) is provably negligible — less than 1% overhead for typical configurations.

This dual optimality is what distinguishes ZeCO from incremental improvements. It's not claiming to be 10% better than LASP-2; it's claiming to achieve the theoretical ceiling, with LASP-1 and LASP-2 representing points on the Pareto frontier that trade off serialization against communication volume, neither achieving both minima simultaneously.

The paper also positions All-Scan as a general-purpose collective communication primitive, not merely a technique for linear attention. The statement that it "provides a foundational innovation for advancing distributed computing in the linear model community" (Section 5) signals an ambition for All-Scan to become a standard primitive alongside All-Gather and All-Reduce, one that is specifically designed for scan-like operations where each element in a sequence needs the accumulated result from all preceding elements — a pattern that appears beyond just linear attention (e.g., in state-space models, prefix sums, and other recurrent computations).

3. Technical Approach

3.1 Reader Orientation

This paper presents a distributed training system for linear attention models—specifically, a sequence parallelism (SP) method called ZeCO that distributes a long input sequence across multiple GPUs and coordinates them to compute the attention output correctly while keeping communication costs to the theoretical minimum. The problem ZeCO solves is that existing SP methods either force GPUs to execute in serial order (LASP-1) or flood the network with redundant communication that grows with the number of devices (LASP-2); ZeCO's solution shape is a pipelined scan communication primitive (All-Scan) that sends exactly one state tensor per device boundary while overlapping that communication with local computation, achieving near-ideal throughput scaling.

3.2 Big-Picture Architecture (Diagram in Words)

The ZeCO system has four interconnected components that transform a distributed input sequence into correct attention outputs:

  1. Input Partitioning Layer: The full sequence (total length P×LP \times L across PP devices) is divided into contiguous chunks of size CC, with each device receiving L/C=NL/C = N chunks. Each device projects its local tokens into Query (Q[n]Q_{[n]}), Key (K[n]K_{[n]}), Value (V[n]V_{[n]}), and decay factor (G[n]G_{[n]}) matrices—this is standard linear attention preprocessing.

  2. Local State Computer: Each device independently computes a series of NN local chunk states S[n]S_{[n]} using the GLA recurrence formula, starting from a zero initial state. This produces intermediate states that represent the local contribution of each chunk, along with a cumulative decay vector γ~[n]\tilde{\gamma}_{[n]} that tracks how much historical information survives through decay. Crucially, this step requires no inter-device communication.

  3. All-Scan Communication Engine: This is the novel primitive. Each device splits its final local state S[N]S_{[N]} into KK blocks and sends them in a pipelined fashion to the next device. Upon receiving the first block of the predecessor's final global state S(p1)LS_{(p-1)L}, the device immediately begins the global update computation—it doesn't wait for all blocks. This receive-update-send pipeline propagates the correct global initial state across all devices with minimal latency.

  4. Global Output Computer: Using the received global initial state S(p1)LS_{(p-1)L} and the pre-computed local states S[n]S_{[n]}, each device computes the final attention outputs O[n]O_{[n]} for all its chunks. This involves adding the inter-chunk contribution (the decayed global state transformed by query) to the intra-chunk contribution (local masked attention within the chunk). The computation is parallelized across chunks and overlapped with the All-Scan communication.

Information flows forward as: raw tokens → projection into Q/K/V/G → local recurrence (producing S[n]S_{[n]} and γ~[n]\tilde{\gamma}_{[n]}) → All-Scan pipeline (propagating SpLS_{pL} between devices) → global output computation (combining S(p1)LS_{(p-1)L} with local states to produce O[n]O_{[n]}). The backward pass reverses this flow, propagating gradients through the same All-Scan pattern in the opposite direction.

3.3 Roadmap for the Deep Dive

The explanation proceeds through five layers, each building on the previous:

  • First, the Gated Linear Attention (GLA) formulation—the mathematical recurrence that ZeCO parallelizes. Understanding GLA is essential because ZeCO's correctness depends on the linear decomposition property of its state update.

  • Second, local state computation—how each device computes its independent chunk states and cumulative decay vectors without any communication. This establishes what information exists before communication begins.

  • Third, the global state update formula—the mathematical relationship that allows a device to transform its local states into globally correct states given only the final state from the previous device. This is what makes minimal communication possible.

  • Fourth, the All-Scan collective communication primitive—the pipelined receive-scan-send algorithm that delivers the necessary initial state to each device while overlapping communication with computation.

  • Fifth, the complete forward and backward algorithms—how all components integrate, including the overlapping of All-Scan with intra-chunk computation, the auxiliary I/O and computation costs, and the proof that these overheads are negligible.

This order follows the natural dependency chain: the mathematics (GLA) justifies the parallelization strategy (local + global decomposition), which determines the communication requirement (one state per device boundary), which motivates the All-Scan design (pipelined to hide latency), which integrates into the full algorithm.

3.4 Detailed, Sentence-Based Technical Breakdown

This is a systems paper whose core idea is that sequence parallelism for linear attention can achieve the theoretical minimum communication volume (dk×dvd_k \times d_v per device, independent of the number of devices) while simultaneously overlapping that communication with computation, yielding near-ideal throughput scaling. The key insight is recognizing that the GLA state update is a linear operation that decomposes into independent local contributions plus a decayed initial condition, meaning each device needs only one incoming state (not all states from all devices) to compute globally correct outputs.

Gated Linear Attention (GLA) Formulation — The Mathematics ZeCO Parallelizes

The paper uses Gated Linear Attention (GLA) as a representative linear attention operator. Understanding its recurrence is essential because ZeCO's correctness proof depends on a linear decomposition property of this recurrence.

The per-token recurrent update. For each token tt, the attention state StRdk×dvS_t \in \mathbb{R}^{d_k \times d_v} is updated as:

St=(αt1)St1+KtVtS_t = (\alpha_t^\top \mathbf{1}) \odot S_{t-1} + K_t^\top V_t

where αt(0,1)dk\alpha_t \in (0, 1)^{d_k} is a per-token decay factor (a vector of values between 0 and 1, one per key-dimension element, controlling how much of the previous state is retained), KtRdkK_t \in \mathbb{R}^{d_k} is the key vector for token tt, VtRdvV_t \in \mathbb{R}^{d_v} is the value vector for token tt, \odot denotes element-wise (Hadamard) product, and 1\mathbf{1} is a vector of ones used to broadcast the decay factor across columns.

What it computes: the new state StS_t is a gated combination of the previous state St1S_{t-1} (element-wise multiplied by the decay factor αt\alpha_t) plus the outer product of the current key and value vectors (KtVtK_t^\top V_t, which is a rank-1 matrix of size dk×dvd_k \times d_v). The decay factor αt\alpha_t acts as a forget gate: values near 1 retain most historical information, values near 0 aggressively forget. The output for token tt is then Ot=QtStO_t = Q_t S_t, where QtRdkQ_t \in \mathbb{R}^{d_k} is the query vector—this dot-product retrieves information from the accumulated state.

Why this form: this recurrence is what gives linear attention its O(L)O(L) complexity. Instead of computing pairwise interactions between all L2L^2 token pairs (as in softmax attention), the model maintains a fixed-size state matrix StS_t that summarizes all past information, updating it with each new token. The cost per token is O(dkdv)O(d_k d_v), and the total cost for LL tokens is O(Ldkdv)O(L d_k d_v)—linear in sequence length. The decay factor αt\alpha_t provides a mechanism for the model to learn which historical information to retain or discard, analogous to gates in LSTMs but applied element-wise to the key dimension.

Chunk-wise formulation for training parallelism. During training, processing tokens one-by-one is inefficient because it prevents parallel computation across the sequence dimension. GLA therefore partitions the sequence into NN chunks of length CC and reformulates the recurrence at chunk granularity. For chunk ii containing tokens iCiC through (i+1)C1(i+1)C-1:

The chunk-level state S[i]Rdk×dvS_{[i]} \in \mathbb{R}^{d_k \times d_v} (the state after processing all tokens in chunk ii) is computed as:

S[i]=(γ[i]1)S[i1]+(K[i]Γ[i])V[i]S_{[i]} = (\gamma_{[i]}^\top \mathbf{1}) \odot S_{[i-1]} + (K_{[i]} \odot \Gamma_{[i]})^\top V_{[i]}

where γ[i]=j=1CαiC+jRdk\gamma_{[i]} = \prod_{j=1}^C \alpha_{iC+j} \in \mathbb{R}^{d_k} is the cumulative decay across all tokens in chunk ii (element-wise product of the per-token decay factors, so γ[i],m=j=1CαiC+j,m\gamma_{[i],m} = \prod_{j=1}^C \alpha_{iC+j,m} for each dimension mm), K[i]RC×dkK_{[i]} \in \mathbb{R}^{C \times d_k} is the key matrix for all tokens in chunk ii, V[i]RC×dvV_{[i]} \in \mathbb{R}^{C \times d_v} is the value matrix, and Γ[i]RC×dk\Gamma_{[i]} \in \mathbb{R}^{C \times d_k} is a token-wise scaling factor defined as ΓiC+j=b(i+1)CbiC+j\Gamma_{iC+j} = \frac{b_{(i+1)C}}{b_{iC+j}} where bt=s=1tαsb_t = \prod_{s=1}^t \alpha_s is the cumulative decay from the beginning of the sequence to token tt.

What it computes: the chunk-level recurrence is the batch version of the per-token recurrence. The first term (γ[i]1)S[i1](\gamma_{[i]}^\top \mathbf{1}) \odot S_{[i-1]} applies the cumulative chunk decay to the previous state—this is exactly what you'd get by applying all CC per-token decays sequentially. The second term (K[i]Γ[i])V[i](K_{[i]} \odot \Gamma_{[i]})^\top V_{[i]} computes the contribution of all tokens in chunk ii, where the scaling Γ[i]\Gamma_{[i]} adjusts each token's contribution to account for its position within the chunk (tokens earlier in the chunk undergo more decay before reaching the chunk boundary).

Why this form: the chunk-wise formulation enables parallel computation within each chunk. The matrix multiplication (K[i]Γ[i])V[i](K_{[i]} \odot \Gamma_{[i]})^\top V_{[i]} is a single dk×Cd_k \times C times C×dvC \times d_v operation that processes all CC tokens simultaneously, exploiting GPU tensor cores. This is O(Cdkdv)O(C d_k d_v) rather than O(C2)O(C^2) for full attention. The tradeoff is that we lose token-level resolution within the chunk for the inter-chunk computation, but we can recover it in the intra-chunk computation.

Output computation. The output for chunk ii, denoted O[i]RC×dvO_{[i]} \in \mathbb{R}^{C \times d_v} (the attention output for all tokens in the chunk), is computed as:

O[i]=(Q[i]Λ[i])S[i1]O[i]inter+[((Q[i]Λ[i])(K[i]Γ[i]))M]V[i]O[i]intraO_{[i]} = \underbrace{(Q_{[i]} \odot \Lambda_{[i]}) \cdot S_{[i-1]}}_{O_{[i]}^{\text{inter}}} + \underbrace{\left[ \left( (Q_{[i]} \odot \Lambda_{[i]}) \cdot (K_{[i]} \odot \Gamma_{[i]})^\top \right) \odot M \right] \cdot V_{[i]}}_{O_{[i]}^{\text{intra}}}

where ΛiC+j=biC+jbiC\Lambda_{iC+j} = \frac{b_{iC+j}}{b_{iC}} is the token-wise scaling for queries (analogous to Γ\Gamma but from the chunk start rather than the chunk end), Q[i]RC×dkQ_{[i]} \in \mathbb{R}^{C \times d_k} is the query matrix, and MRC×CM \in \mathbb{R}^{C \times C} is a causal mask (lower triangular, ensuring token jj within the chunk can only attend to tokens 1..j1..j, not future tokens).

What it computes: the output has two independent components. The inter-chunk term O[i]interO_{[i]}^{\text{inter}} computes the contribution from all tokens before chunk ii—it multiplies the scaled queries (Q[i]Λ[i])(Q_{[i]} \odot \Lambda_{[i]}) by the state S[i1]S_{[i-1]} that summarizes the entire history up to chunk i1i-1. The intra-chunk term O[i]intraO_{[i]}^{\text{intra}} computes the contribution from tokens within chunk ii—it's a chunk-local attention computation with causal masking, producing CC output vectors. The sum gives the full attention output.

Why this form: this decomposition is the key enabler for ZeCO. The inter-chunk term depends only on S[i1]S_{[i-1]}, the state before the current chunk. If we know the correct (global) S[i1]S_{[i-1]}, we can compute O[i]O_{[i]} entirely from local data. The intra-chunk term depends only on tokens within the chunk and can be computed independently and in parallel across chunks. This separation of concerns is what makes minimal-communication SP possible.


Local State Computation — What Each Device Computes Independently

The first phase of ZeCO executes purely locally on each device, requiring zero communication. This phase transforms the input tokens into a set of intermediate states and decay vectors that will later be combined with the incoming global state to produce correct outputs.

Setup. Each device p{0,1,,P1}p \in \{0, 1, \dots, P-1\} receives a contiguous segment of the full sequence: tokens at positions pLpL through (p+1)L1(p+1)L-1 (using 0-indexing for the global sequence). This segment is further partitioned into N=L/CN = L/C chunks of size CC, with chunk indices n=1,,Nn = 1, \dots, N (local indexing within the device). The input XRL×dX \in \mathbb{R}^{L \times d} is projected to produce Q[n],K[n],G[n]RC×dkQ_{[n]}, K_{[n]}, G_{[n]} \in \mathbb{R}^{C \times d_k} and V[n]RC×dvV_{[n]} \in \mathbb{R}^{C \times d_v} for each chunk nn.

The local recurrence. Starting from an initial state S[0]=0Rdk×dvS_{[0]} = \mathbf{0} \in \mathbb{R}^{d_k \times d_v} (the zero matrix) and initial cumulative decay γ~[0]=1Rdk\tilde{\gamma}_{[0]} = \mathbf{1} \in \mathbb{R}^{d_k} (a vector of all ones, meaning no decay has been applied yet), the device iterates through its NN chunks sequentially:

S[n]=(γ[n]T1)S[n1]+K~[n]TV[n],for n=1,,NS_{[n]} = \left( \gamma_{[n]}^T \mathbf{1} \right) \odot S_{[n-1]} + \tilde{K}_{[n]}^T V_{[n]}, \quad \text{for } n = 1, \dots, N

where K~[n]=K[n]Γ[n]RC×dk\tilde{K}_{[n]} = K_{[n]} \odot \Gamma_{[n]} \in \mathbb{R}^{C \times d_k} is the scaled key matrix (the per-token Γ\Gamma scaling is applied element-wise to the key vectors, adjusting for each token's position within the chunk).

Simultaneously, the device maintains and stores the cumulative decay vector:

γ~[n]=i=1nγ[i],for n=1,,N\tilde{\gamma}_{[n]} = \prod_{i=1}^n \gamma_{[i]}, \quad \text{for } n = 1, \dots, N

where the product is element-wise, so γ~[n],m=i=1nγ[i],m\tilde{\gamma}_{[n],m} = \prod_{i=1}^n \gamma_{[i],m} for each dimension m{1,,dk}m \in \{1, \dots, d_k\}.

What it computes: S[n]S_{[n]} is the local contribution of chunks 1 through nn, computed as if the initial state were zero. It represents the accumulated key-value outer products from all tokens up to chunk nn, properly decayed within the local sequence—but because the initial state S[0]S_{[0]} is zero, it lacks the contribution from tokens before this device's segment. The vector γ~[n]\tilde{\gamma}_{[n]} records the total multiplicative decay from the device's first chunk through chunk nn, which will be needed to correctly incorporate the incoming global state from the previous device.

Why this form: zero-initialization is the crucial design choice. By starting from zero, the device can compute its local contributions entirely independently—no waiting for upstream devices. The resulting S[n]S_{[n]} is incomplete (it lacks history), but the linearity of the recurrence (proven in Appendix A.1) means we can later add the missing historical contribution using only the final state S(p1)LS_{(p-1)L} from the previous device and the decay vector γ~[n]\tilde{\gamma}_{[n]}. This is the mathematical property that enables minimal communication.

What gets stored. At the end of this local recurrence, the device has produced a list of chunk states {S[0],S[1],,S[N]}\{S_{[0]}, S_{[1]}, \dots, S_{[N]}\} (where S[0]=0S_{[0]} = \mathbf{0}) and cumulative decay vectors {γ~[0],γ~[1],,γ~[N]}\{\tilde{\gamma}_{[0]}, \tilde{\gamma}_{[1]}, \dots, \tilde{\gamma}_{[N]}\} (where γ~[0]=1\tilde{\gamma}_{[0]} = \mathbf{1}). These are written to HBM (high-bandwidth memory, the GPU's main memory) for later use in the output computation. The vectors γ~[n]\tilde{\gamma}_{[n]} are small—dkd_k elements each, compared to dk×dvd_k \times d_v for the state matrices—so storing them costs only 1/dv1/d_v of the state storage, a negligible overhead as proven in Section 3.3.


Global State Update Formula — The Linear Decomposition Property

The mathematical insight that makes ZeCO possible is that the global state at any position can be expressed as a linear combination of the incoming global state and the local contribution. This is proven in Appendix A.1 and stated as Equation (6) in the main text.

The update formula. Let S(p1)LS_{(p-1)L} be the final global state from the previous device p1p-1 (the state after processing all tokens up to position (p1)L(p-1)L). For device pp, the global state at the end of chunk nn (which corresponds to global position (p1)L+nC(p-1)L + nC) is:

S(p1)L+nC=(γ~[n]T1)S(p1)L+S[n]S_{(p-1)L+nC} = (\tilde{\gamma}_{[n]}^T \mathbf{1}) \odot S_{(p-1)L} + S_{[n]}

where γ~[n]\tilde{\gamma}_{[n]} is the cumulative decay from the device's first chunk through chunk nn (computed locally above), S(p1)LS_{(p-1)L} is the incoming state from the predecessor device, and S[n]S_{[n]} is the local chunk state computed from zero initialization.

What it computes: the global state is the sum of two independent contributions. The first term (γ~[n]T1)S(p1)L(\tilde{\gamma}_{[n]}^T \mathbf{1}) \odot S_{(p-1)L} applies the cumulative decay to the incoming historical state—it accounts for the fact that information from before device pp has undergone additional decay by the time we reach chunk nn. The second term S[n]S_{[n]} is the local contribution, computed exactly as if the initial state were zero. Because the recurrence is linear in the initial state (as proven by unfolding the recurrence in Equation 19 of the appendix), these two contributions are additive and independent.

Why this form: this decomposition is what enables minimal communication. It means device pp does not need to know the detailed history of tokens before its segment—it only needs the single summary state S(p1)LS_{(p-1)L}, which compresses all that history into a fixed-size matrix. Moreover, it does not need any intermediate states from device p1p-1 (unlike LASP-2's All-Gather approach), because the linearity lets it apply the correct decay and addition in one shot. This is a fundamentally different property from full attention, where computing the output for a token requires access to the individual key-value pairs of all previous tokens, not just a summary.

The proof sketch (from Appendix A.1). Starting from the chunk-wise recurrence with a non-zero initial state S(p1)LS_{(p-1)L}:

S(p1)L+nC=i=1n(j=i+1nγ[j]T1)(K~[i]TV[i])+(j=1nγ[j]T1)S(p1)LS_{(p-1)L+nC} = \sum_{i=1}^n \left( \prod_{j=i+1}^n \gamma_{[j]}^T \mathbf{1} \right) \odot (\tilde{K}_{[i]}^T V_{[i]}) + \left( \prod_{j=1}^n \gamma_{[j]}^T \mathbf{1} \right) \odot S_{(p-1)L}

The first summation is exactly S[n]S_{[n]} (the local contribution from zero initialization), and the product j=1nγ[j]T\prod_{j=1}^n \gamma_{[j]}^T is exactly γ~[n]T\tilde{\gamma}_{[n]}^T. This establishes the decomposition.

Practical consequence for communication. The only information that must flow between devices is the final global state of each device. Device pp needs S(p1)LS_{(p-1)L} from device p1p-1, uses it to compute all its global states via the formula above, and then sends its own final global state SpLS_{pL} to device p+1p+1. The communication volume per device is exactly S=dk×dv|S| = d_k \times d_v elements—the theoretical minimum.

Why one state is sufficient. A natural question: could we need more information? For full attention, yes—you need all individual key-value pairs. For linear attention, no—the state matrix SS is a sufficient statistic for the entire history. The recurrence St=(αtT1)St1+KtTVtS_t = (\alpha_t^T \mathbf{1}) \odot S_{t-1} + K_t^T V_t is a Markov process: the future depends on the past only through the current state. Therefore, knowing S(p1)LS_{(p-1)L} gives device pp all the information about the past that can possibly be relevant for computing future outputs. Any additional communication would be redundant.


All-Scan Collective Communication — The Pipelined Receive-Scan-Send Primitive

All-Scan is the novel communication primitive at the heart of ZeCO. Its job is to propagate the final global state from each device to the next device in the sequence, enabling each device to compute its global states, while minimizing latency through pipelining.

The basic communication requirement. The dependency chain is: device 0 computes its local states and sends its final global state SLS_{L} to device 1; device 1 receives SLS_{L}, uses it to update its local states to global states, computes its own final global state S2LS_{2L}, and sends it to device 2; and so on. A naive implementation would have device pp wait to receive the entire state S(p1)LS_{(p-1)L} before beginning its global update computation. For large states (e.g., dk=128d_k = 128, dv=128d_v = 128 gives a 128×128=16,384128 \times 128 = 16,384 element matrix) and many devices, this serial waiting would dominate runtime.

The pipelining strategy. All-Scan avoids this waiting by splitting the state tensor into KK contiguous blocks along the dkd_k dimension and processing them in a pipeline. Specifically, the state SRdk×dvS \in \mathbb{R}^{d_k \times d_v} is partitioned as:

S=[S(1),S(2),,S(K)],S(k)RdkK×dvS = \left[ S^{(1)}, S^{(2)}, \dots, S^{(K)} \right], \quad S^{(k)} \in \mathbb{R}^{\frac{d_k}{K} \times d_v}

and the decay factor is correspondingly split as γ~(k)R1×dkK\tilde{\gamma}^{(k)} \in \mathbb{R}^{1 \times \frac{d_k}{K}}. Each block represents a subset of the key-dimension's contribution to the state.

The per-block update. For block kk, device pp receives S(p1)L(k)S_{(p-1)L}^{(k)} from device p1p-1, immediately computes:

SpL(k)=(γ~[N](k)T1)S(p1)L(k)+S[N](k)S_{pL}^{(k)} = (\tilde{\gamma}_{[N]}^{(k)T} \mathbf{1}) \odot S_{(p-1)L}^{(k)} + S_{[N]}^{(k)}

where SpL(k)S_{pL}^{(k)} is block kk of device pp's final global state, and immediately sends SpL(k)S_{pL}^{(k)} to device p+1p+1.

What it computes: this per-block update is the same mathematical operation as the global state update formula (Equation 6), applied to the final chunk NN of the device rather than intermediate chunks. The result SpL(k)S_{pL}^{(k)} is block kk of the device's final global state, incorporating both the historical contribution (decayed incoming state) and the device's full local contribution (S[N](k)S_{[N]}^{(k)}). By computing and forwarding each block independently, device p+1p+1 can begin processing block kk of its own update as soon as it receives SpL(k)S_{pL}^{(k)}, without waiting for blocks k+1,,Kk+1, \dots, K.

Why this form: the pipelining exploits the fact that the global state update is an element-wise operation along the dkd_k dimension—each row of SS (corresponding to one key-dimension element) can be updated independently using its corresponding decay factor. There is no cross-row dependency. This means splitting along dkd_k introduces no sequential bottlenecks: block kk's update uses only γ~(k)\tilde{\gamma}^{(k)}, S(p1)L(k)S_{(p-1)L}^{(k)}, and S[N](k)S_{[N]}^{(k)}, with no need for information from other blocks. The element-wise nature of the decay operation is what makes this independent block processing valid.

The pipeline latency analysis. The total time for All-Scan, as derived in Equation (12), is:

TAll_Scan=τ(dk×dv)+(P1)τ(dk×dv)KT_{\text{All\_Scan}} = \tau(d_k \times d_v) + \frac{(P-1)\tau(d_k \times d_v)}{K}

where τ(dk×dv)\tau(d_k \times d_v) is the time to communicate a full state tensor of size dk×dvd_k \times d_v.

  • The first term τ(dk×dv)\tau(d_k \times d_v) represents the time to transmit the full state once—this is the minimum possible communication time, since at least one full state must cross each device boundary.
  • The second term (P1)τ(dk×dv)K\frac{(P-1)\tau(d_k \times d_v)}{K} represents the non-overlappable boundary overhead. When K=1K=1 (no pipelining), this is (P1)τ(dk×dv)(P-1)\tau(d_k \times d_v), which would mean each device waits for the full state before forwarding—terrible scaling. When KK is large, this term approaches zero because the pipeline fills quickly and steady-state throughput approaches the link bandwidth.

Why this achieves near-optimal communication time: as KK \to \infty, TAll_Scanτ(dk×dv)T_{\text{All\_Scan}} \to \tau(d_k \times d_v), which means the total communication time approaches the time to send just one state tensor, even though PP devices are involved. The intuition: the pipeline keeps all links busy simultaneously—device 0 sends block 1 to device 1, device 1 sends block 1 to device 2, etc.—so the bandwidth is fully utilized, and the only serial delay is the pipeline fill time, which becomes negligible as the number of blocks grows. In practice, KK is chosen large enough that the boundary overhead is dwarfed by the computation it's overlapped with.

Implementation details (Algorithm 2). All-Scan operates as a CUDA stream separate from the main computation stream, enabling hardware-level overlap. The algorithm handles three device roles:

  • Start device (rank 0): sends its local state S[N]S_{[N]} (which equals SLS_{L} since there's no incoming state to incorporate) block-by-block to device 1.
  • Middle devices (ranks 1 to P2P-2): for each block kk, receive S(p1)L(k)S_{(p-1)L}^{(k)} from device p1p-1, compute Ssendk=S[N](k)+(γ~[N](k)1)×Srecv(k)S_{\text{send}}^k = S_{[N]}^{(k)} + (\tilde{\gamma}_{[N]}^{(k)} \mathbf{1}) \times S_{\text{recv}}^{(k)}, and immediately send SsendkS_{\text{send}}^k to device p+1p+1.
  • End device (rank P1P-1): receives blocks from device P2P-2 and computes the update, but does not forward (there is no device PP).

Why a new primitive rather than using existing ones: One might ask why not use MPI_Scan or NCCL ReduceScatter. Standard scan operations compute prefix sums where each process contributes a value and receives the accumulated sum of all preceding processes' values—this matches the conceptual pattern (each device needs the accumulated state from all previous devices). However, existing implementations are designed for small-element operations (scalars or small vectors), not large matrix tensors with element-wise decay scaling. The decay operation (γ~T1)\odot (\tilde{\gamma}^T \mathbf{1}) is not a simple sum; it's a scaled addition. All-Scan is a custom scan that (a) handles the specific linear attention update formula, (b) pipelines at block granularity to overlap communication with computation, and (c) operates on a separate CUDA stream for hardware parallelism.

Relation to the backward pass. The backward pass uses the same All-Scan pattern but in the reverse direction (Algorithm 3). Gradients of the loss with respect to the states (dSdS) flow from later devices to earlier devices, following the chain rule. The backward All-Scan propagates dSpLdS_{pL} from device pp to device p1p-1 using the same pipelined block-wise pattern, with the decay factor γ~[0]\tilde{\gamma}_{[0]} (the cumulative decay from the beginning, computed in reverse). The mathematical justification is symmetric: just as the forward state depends linearly on the incoming state, the backward gradient depends linearly on the incoming gradient, enabling the same pipelined communication pattern.


The Complete Forward Algorithm — Integrating Computation and Communication

Algorithm 1 presents the full forward pass of ZeCO, which orchestrates local computation, All-Scan communication, and output computation to exploit maximal parallelism.

Phase 1: Local state computation (lines 4-12). This phase runs on the main CUDA stream without any communication. For each chunk n=1,,Nn = 1, \dots, N:

  • Load K[n],G[n],V[n]K_{[n]}, G_{[n]}, V_{[n]} from HBM to SRAM (on-chip memory).
  • Compute γ[n]\gamma_{[n]} (chunk-level decay) and Γ[n]\Gamma_{[n]} (token-wise key scaling) from the per-token decay factors G[n]G_{[n]}.
  • Compute K~[n]=K[n]Γ[n]\tilde{K}_{[n]} = K_{[n]} \odot \Gamma_{[n]} (scaled keys).
  • Update cumulative decay: γ~=γ~γ[n]\tilde{\gamma} = \tilde{\gamma} \odot \gamma_{[n]}.
  • Write γ~\tilde{\gamma} to HBM as γ~[n]\tilde{\gamma}_{[n]}.
  • Update state: S=(γ[n]T1)S+K~[n]TV[n]S = (\gamma_{[n]}^T \mathbf{1}) \odot S + \tilde{K}_{[n]}^T V_{[n]}.
  • Write SS to HBM as S[n]S_{[n]}.

The output is the set of local chunk states {S[0],,S[N]}\{S_{[0]}, \dots, S_{[N]}\} and cumulative decay vectors {γ~[0],,γ~[N]}\{\tilde{\gamma}_{[0]}, \dots, \tilde{\gamma}_{[N]}\}, all in HBM.

Phase 2: Overlapped communication and intra-chunk computation (lines 13-23). After local state computation, two independent streams execute in parallel:

Stream 1 (line 15): Executes All-Scan(S_{[N]}, \tilde{\gamma}_{[N]}), which sends device pp's final local state to device p+1p+1 and receives the predecessor's final global state S(p1)LS_{(p-1)L}, using the pipelined block-wise protocol. This stream handles all inter-device communication.

Stream 2 (lines 17-23): Computes the intra-chunk attention scores for all chunks in parallel. For each chunk nn:

  • Load Q[n],K[n],G[n]Q_{[n]}, K_{[n]}, G_{[n]} from HBM.
  • Construct causal mask MRC×CM \in \mathbb{R}^{C \times C} (lower triangular, ensuring causality within the chunk).
  • Compute Λ[n]\Lambda_{[n]} (token-wise query scaling) and Q~[n]=Q[n]Λ[n]\tilde{Q}_{[n]} = Q_{[n]} \odot \Lambda_{[n]}, K~[n]=K[n]/Λ[n]\tilde{K}_{[n]} = K_{[n]} / \Lambda_{[n]} (note: this uses bare keys, not Γ\Gamma-scaled keys—these are for the intra-chunk term).
  • Compute P=(Q~[n]K~[n]T)MRC×CP = (\tilde{Q}_{[n]} \tilde{K}_{[n]}^T) \odot M \in \mathbb{R}^{C \times C} (the intra-chunk attention matrix, masked for causality).
  • Write PP to HBM as P[i]P_{[i]}.

What happens during overlap: while the All-Scan pipeline is transmitting state blocks between devices, the GPU compute units are busy computing intra-chunk attention matrices. This is possible because the intra-chunk computation depends only on QQ, KK, and GG—all of which are fully available locally by the end of Phase 1. It does not depend on the incoming global state S(p1)LS_{(p-1)L}. By scheduling these on separate CUDA streams, the GPU's compute and communication resources (NVLink/InfiniBand) operate simultaneously, hiding communication latency.

Why this specific overlap: the intra-chunk computation is chosen for overlap because it's the most compute-intensive local operation that doesn't depend on S(p1)LS_{(p-1)L}. The alternative—overlapping with the inter-chunk output computation—wouldn't work because that computation requires S(p1)LS_{(p-1)L}, which is exactly what All-Scan is in the process of delivering. By the time All-Scan completes (stream 1 finishes), stream 2 has already computed all intra-chunk attention matrices, so the final output computation can proceed without waiting for this compute-heavy step.

Phase 3: Output computation (lines 24-32). After both streams complete (enforced by a stream barrier at line 24), the final outputs are computed for each chunk nn:

  • Load Q[n],G[n],V[n],S(p1)LQ_{[n]}, G_{[n]}, V_{[n]}, S_{(p-1)L} (the received global state), S[n]S_{[n]} (the local chunk state), γ~[n1]\tilde{\gamma}_{[n-1]}, and PP (the pre-computed intra-chunk attention matrix).
  • Compute Λ[n]\Lambda_{[n]} and Q~[n]=Q[n]Λ[n]\tilde{Q}_{[n]} = Q_{[n]} \odot \Lambda_{[n]}.
  • Compute O[n]inter=Q~[n](S[n1]+(γ~[n1]T1)S(p1)L)O_{[n]}^{\text{inter}} = \tilde{Q}_{[n]}(S_{[n-1]} + (\tilde{\gamma}_{[n-1]}^T \mathbf{1}) \odot S_{(p-1)L}) — this is the inter-chunk contribution, where S[n1]+(γ~[n1]T1)S(p1)LS_{[n-1]} + (\tilde{\gamma}_{[n-1]}^T \mathbf{1}) \odot S_{(p-1)L} is the globally correct state before chunk nn, decomposed as the local state plus the decayed global state.
  • Compute O[n]intra=PV[n]O_{[n]}^{\text{intra}} = P V_{[n]} — this uses the pre-computed attention matrix PP to weight the value vectors, producing the intra-chunk contribution.
  • Compute O[n]=O[n]inter+O[n]intraO_{[n]} = O_{[n]}^{\text{inter}} + O_{[n]}^{\text{intra}} — the full output for chunk nn.
  • Write O[n]O_{[n]} to HBM.

A subtle point about the output formula: note that the inter-chunk output uses S[n1]+(γ~[n1]T1)S(p1)LS_{[n-1]} + (\tilde{\gamma}_{[n-1]}^T \mathbf{1}) \odot S_{(p-1)L}, which is the global state before chunk nn. This differs slightly from the earlier presentation (Equation 6) which gave the state after chunk nn. The output computation needs the state before chunk nn because it computes how the queries at chunk nn attend to all previous tokens, which are summarized in the state prior to processing chunk nn's tokens. This is consistent with the GLA formulation: output Ot=QtSt1O_t = Q_t S_{t-1}, so chunk nn's output uses the state after chunk n1n-1.

Why the decomposition into S[n1]S_{[n-1]} and S(p1)LS_{(p-1)L}: this is another manifestation of the linearity property. The correct global state before chunk nn is the sum of (a) the local contribution from chunks 1 through n1n-1, which is S[n1]S_{[n-1]}, and (b) the contribution from before this device, which is the incoming global state decayed by γ~[n1]\tilde{\gamma}_{[n-1]} to account for the additional decay through chunks 1 to n1n-1 on this device. This avoids recomputing the full global state for each chunk separately—we just combine the pre-computed local state with the appropriately decayed global state.


The Backward Algorithm and Gradient Flow

The backward pass (Algorithm 3) computes gradients of the loss with respect to Q,K,V,GQ, K, V, G (the gate/decay parameters) by reversing the forward computation. The structure mirrors the forward pass but in reverse order, with gradients flowing through the same All-Scan pattern in the backward direction.

Phase 1: Local gradient accumulation. For each chunk nn in reverse order (from NN down to 1), the device accumulates the gradient with respect to the state:

dS=(γ[n]T1)dS+Q~[n]TdO[n]dS = (\gamma_{[n]}^T \mathbf{1}) \odot dS + \tilde{Q}_{[n]}^T dO_{[n]}

starting from dS=0dS = 0. This computes how the loss changes with respect to the state before chunk nn (i.e., S[n1]S_{[n-1]}). The gradient of the state update operation involves two terms: the decayed gradient from later chunks (γ[n]T1)dS(\gamma_{[n]}^T \mathbf{1}) \odot dS, and the contribution from the current chunk's output Q~[n]TdO[n]\tilde{Q}_{[n]}^T dO_{[n]} (since O[n]=Q~[n]S[n1]+O_{[n]} = \tilde{Q}_{[n]} S_{[n-1]} + \dots, the gradient w.r.t. S[n1]S_{[n-1]} is Q~[n]TdO[n]\tilde{Q}_{[n]}^T dO_{[n]}).

Phase 2: Overlapped backward All-Scan and intra-chunk gradient computation. Two streams run in parallel:

Stream 1: Executes backward All-Scan to propagate dS[0]dS_{[0]} (the gradient w.r.t. the state before this device's first chunk, which is S(p1)LS_{(p-1)L}) to the previous device, and receive dSpLdS_{pL} from the next device. The direction is reversed: gradients flow from device p+1p+1 to device pp (when DIR = BWD in Algorithm 2, send_rank = p - 1, recv_rank = p + 1).

Stream 2: Recomputed the local states S[n]S_{[n]} (needed for gradient computation of other parameters) and computes intra-chunk gradients:

  • Recompute S[n]S_{[n]} using the correct initial state S(p1)LS_{(p-1)L} (received during forward pass, stored in HBM).
  • For each chunk, compute dQ[n],dK[n]dQ_{[n]}, dK_{[n]} from the intra-chunk attention—this involves computing dP=(dO[n]V[n]T)MdP = (dO_{[n]} V_{[n]}^T) \odot M and backpropagating through the attention score computation.

Phase 3: Combining inter and intra-chunk gradients. The final gradient for each parameter combines contributions from both the inter-chunk path (through the state) and the intra-chunk path (through the attention matrix):

  • For K[n]K_{[n]}: gradient includes the intra-chunk contribution dK~[n]d\tilde{K}_{[n]} and the inter-chunk contribution V[n](dS[n1]T+(γ~[n1]T1)dSpLT)V_{[n]}(dS_{[n-1]}^T + (\tilde{\gamma}_{[n-1]}^T \mathbf{1}) \odot dS_{pL}^T), where the second term reflects how changing the keys affects the state update and thus all future outputs.
  • For V[n]V_{[n]}: similarly combines intra-chunk (P[n]TdO[n]P_{[n]}^T dO_{[n]}) and inter-chunk (K~[n](dS[n1]T+)\tilde{K}_{[n]}(dS_{[n-1]}^T + \dots)) contributions.
  • For Q[n]Q_{[n]}: greedily combines intra-chunk and inter-chunk contributions through the query-state product.

What gets communicated in the backward pass: only one gradient state tensor dSdS flows between adjacent devices, just as in the forward pass. The communication volume is again dk×dvd_k \times d_v per device, the theoretical minimum. The backward All-Scan uses the same block-wise pipelining as the forward All-Scan, achieving the same latency bound of τ(dk×dv)+(P1)τ(dk×dv)K\tau(d_k \times d_v) + \frac{(P-1)\tau(d_k \times d_v)}{K}.


Proof of Negligible Extra Cost — Why ZeCO is Optimal

Section 3.3 provides the theoretical analysis that ZeCO achieves the minimum possible overhead for sequence parallelism. This analysis decomposes the total runtime into components and shows each is either optimal or negligible.

The runtime decomposition. For processing a total sequence of length PLPL on PP devices, the total time is:

TZeCOP(PL)=Tideal-SP1(L)Toverlapped_comp+τ(dk×dv)+ϵT_{\text{ZeCO}}^P(PL) = T_{\text{ideal-SP}}^1(L) - T_{\text{overlapped\_comp}} + \tau(d_k \times d_v) + \epsilon

where Tideal-SP1(L)T_{\text{ideal-SP}}^1(L) is the time to process a sequence of length LL on a single device with ideal (zero-overhead) parallelism—this is the theoretical lower bound for any SP method, since each device must at minimum process its local LL tokens. Toverlapped_compT_{\text{overlapped\_comp}} is the computation time that runs concurrently with All-Scan (the intra-chunk attention computation). τ(dk×dv)\tau(d_k \times d_v) is the non-overlapped communication time—the minimum time to transmit one state tensor, which is unavoidable since at least this much information must cross each device boundary. ϵ\epsilon represents the extra computation and I/O overhead from ZeCO-specific operations (maintaining γ~\tilde{\gamma} vectors, the global state update), which is proven to be negligible.

Why the extra I/O is negligible. Algorithm 1 introduces two additional I/O operations compared to a non-SP GLA implementation: writing γ~[n]\tilde{\gamma}_{[n]} to HBM (line 9) and loading γ~[n1]\tilde{\gamma}_{[n-1]} from HBM (line 26). These are vectors of size dkd_k elements each. For comparison, the state tensor S[n]S_{[n]} is of size dk×dvd_k \times d_v. The I/O cost ratio is:

extra I/Ostate I/O=dkdk×dv=1dv\frac{\text{extra I/O}}{\text{state I/O}} = \frac{d_k}{d_k \times d_v} = \frac{1}{d_v}

For typical configurations (e.g., dv=128d_v = 128), this is less than 1%. Moreover, the γ~\tilde{\gamma} vectors are reused across all NN chunks on a device—they're written once per chunk but loaded NN times (once per chunk for output computation). Compared to the total I/O of loading Q,K,VQ, K, V for each chunk (which are C×dkC \times d_k or C×dvC \times d_v each), the γ~\tilde{\gamma} I/O is amortized to 1N\frac{1}{N} of the chunk I/O. For L=8192,C=64L = 8192, C = 64, we have N=128N = 128, making the overhead 11280.8%\frac{1}{128} \approx 0.8\%. The element-wise multiplications for state updates (using γ~\tilde{\gamma}) are scalar-vector operations, costing O(dk×dv)O(d_k \times d_v) flops per chunk, compared to O(Cdkdv)O(C d_k d_v) for the attention computation—a factor of CC less.

Why this makes ZeCO optimal. The paper identifies two necessary conditions for optimal SP:

  1. Communication volume must be the theoretical minimum (dk×dvd_k \times d_v per device). ZeCO achieves this (proven in Equation 9), while LASP-2 incurs (P1)×dk×dv(P-1) \times d_k \times d_v.
  2. Communication must be overlapped with computation to hide latency. ZeCO overlaps All-scan with intra-chunk attention, while LASP-1 serializes everything.

The paper also identifies a sufficient condition: extra computation and I/O introduced by SP must be negligible. ZeCO satisfies this with the ϵ\epsilon bound above.

Comparison with LASP-1 and LASP-2. The theoretical runtime formulas (Equations 15-16) quantify the suboptimality:

  • LASP-1: TLASPP(PL)=P×(Tideal-SP1(L)+τ(dk×dv))T_{\text{LASP}}^P(PL) = P \times (T_{\text{ideal-SP}}^1(L) + \tau(d_k \times d_v)) — both computation and communication scale with PP due to serialization. At P=128P=128, this is 128×128\times the single-device time.
  • LASP-2: TLASP-2P(PL)=Tideal-SP1(L)+P×τ(dk×dv)T_{\text{LASP-2}}^P(PL) = T_{\text{ideal-SP}}^1(L) + P \times \tau(d_k \times d_v) — computation is parallel (good) but communication scales with PP (bad). At P=128P=128, communication is 128×128\times the minimum.
  • ZeCO: TZeCOP(PL)Tideal-SP1(L)Toverlapped_comp+τ(dk×dv)T_{\text{ZeCO}}^P(PL) \approx T_{\text{ideal-SP}}^1(L) - T_{\text{overlapped\_comp}} + \tau(d_k \times d_v) — computation is parallel, communication is constant (independent of PP), and the overlapped computation further reduces the effective latency. This is the theoretical optimum.

The experimental results in Tables 3-4 validate this analysis. At 128 GPUs with 16K sequences, ZeCO takes 9.88ms for a forward-backward pass (only 3ms more than the 6.55ms single-device baseline), while LASP-1 takes 113.71ms (17×17\times slower) and LASP-2 takes 35.72ms (5.5×5.5\times slower). The near-constant per-GPU throughput in Tables 5-6 (ZeCO drops only from 44,497 to 34,400 tokens/sec/GPU at 8 to 256 GPUs for 16K, versus LASP-2's drop to 15,196) confirms that the theoretical optimality translates to practical near-linear scaling.

4. Key Insights and Innovations

Innovation 1: Reframing Sequence Parallelism from a Communication-Avoidance Problem to a Communication-Optimality Problem

The dominant framing in sequence parallelism research has been to treat communication as an unavoidable cost to be minimized through clever scheduling, but never eliminated. LASP-1 minimizes communication volume (one state per device) but pays with serial execution. LASP-2 achieves parallel execution but pays with communication volume that grows linearly with device count. Both approaches implicitly accept a tradeoff: you can have low communication or high parallelism, but not both.

ZeCO's distinctive intellectual move is to reframe the problem entirely. Rather than asking "how can we reduce communication overhead?" — which implies communication is a cost to be trimmed — ZeCO asks "what is the theoretically minimum communication required for correctness, and can we design a strategy that achieves exactly that minimum while also achieving full parallelism?" The answer it provides is yes: the minimum communication is one state tensor (dk×dvd_k \times d_v elements) per device boundary, and All-Scan achieves this while overlapping that communication with computation.

This reframing has a subtle but profound implication: it converts sequence parallelism for linear attention from an empirical engineering problem (where methods are compared by relative throughput improvements) into a theoretically bounded problem (where methods can be proven optimal or suboptimal against known lower bounds). The paper's Equation (9) — showing VZeCO(p)=dk×dvV_{\text{ZeCO}}^{(p)} = d_k \times d_v — is not just a claim about ZeCO's communication volume. It is a claim that no algorithm can do better, because each device genuinely needs the accumulated state from all previous tokens to compute correct outputs. This transforms the evaluation criterion from "is ZeCO better than LASP-2?" to "is ZeCO provably optimal?", and the answer to the second question is a stronger claim than any empirical comparison.

The significance of this reframing extends beyond ZeCO itself. It establishes a template for analyzing other distributed learning algorithms: identify the information-theoretic minimum that must be communicated for correctness, then design a scheduling strategy that achieves that minimum with maximal overlap. This is analogous to how the Chinchilla scaling laws (Hoffmann et al., 2022) reframed pretraining from "how do we make bigger models?" to "what is the optimal allocation of compute between model size and data?", establishing a theoretical framework that subsequent work builds on.

Innovation 2: The Linear Decomposition Property as an Exploitable Structural Insight, Not Just a Mathematical Convenience

The linear decomposition property of GLA states — that S(p1)L+nC=(γ~[n]T1)S(p1)L+S[n]S_{(p-1)L+nC} = (\tilde{\gamma}_{[n]}^T \mathbf{1}) \odot S_{(p-1)L} + S_{[n]} — is presented in the paper as a mathematical proof (Appendix A.1), but its significance is deeper than a correctness lemma. Prior work on linear attention SP implicitly understood that the recurrence is linear, but did not exploit the full implication: that the local computation can be completely decoupled from the global history, with the only connection being a single summary state that undergoes a simple element-wise decay operation.

LASP-1 recognized that only one state needs to cross each device boundary, but did not recognize that the local computation on each device could proceed entirely independently of when that state arrives. Instead, it serialized the full pipeline: compute local states, wait for incoming state, update, send. LASP-2 recognized that local computation could be parallelized, but did not recognize that each device needs only one incoming state — it gathered all states from all devices, essentially computing the same scan operation redundantly on every device.

ZeCO's insight is that the linear decomposition creates a structural separation between local and global computation that can be exploited for scheduling. Because S[n]S_{[n]} depends only on local tokens (computed from zero initialization) and γ~[n]\tilde{\gamma}_{[n]} depends only on local decay factors, the entire local computation phase (Phase 1 in Algorithm 1) can run without any communication whatsoever. The global update — combining S(p1)LS_{(p-1)L} with S[n]S_{[n]} — is a cheap element-wise operation that can begin as soon as the first block of S(p1)LS_{(p-1)L} arrives. And the intra-chunk attention computation (P[n]=(Q~[n]K~[n]T)MP_{[n]} = (\tilde{Q}_{[n]} \tilde{K}_{[n]}^T) \odot M) depends on neither S(p1)LS_{(p-1)L} nor S[n]S_{[n]}, so it can run concurrently with the communication pipeline.

This three-way decomposition — local state computation (no communication needed), global state update (minimal communication, element-wise operation), intra-chunk attention (compute-heavy, no communication needed) — is what makes the overlap possible. Prior work conflated these stages because they didn't recognize the full independence: LASP-1 serialized everything, LASP-2 redundantly computed the global update everywhere. ZeCO's contribution is not just that it overlaps communication with computation (many systems do this), but that it identified which specific computation can overlap without creating dependencies — the intra-chunk attention, which is both the most compute-intensive local operation and completely independent of the communication content.

This is a case study in how a theoretical property (linearity of the recurrence) can have direct system design consequences when fully exploited. The paper does not merely apply linearity as a convenience for proving correctness; it uses it as the architectural principle that determines which operations go on which CUDA streams, which tensors are stored versus recomputed, and where the pipeline boundaries are drawn.

Innovation 3: A New Communication Primitive (All-Scan) That Fills a Gap in the Collective Communication Landscape

The collective communication primitives available in standard distributed computing libraries — All-Gather, All-Reduce, Reduce-Scatter, Broadcast, Scatter, Gather — are designed for specific patterns of data movement: all-to-all, one-to-all, or many-to-one. None of them directly support a scan pattern, where process pp needs the accumulated result from processes 00 through p1p-1, and each process contributes to the accumulation.

This is not an accident of library design. Scan operations (also called prefix sums) have been well-studied in parallel computing (Blelloch, 1990), but their implementations in MPI (MPI_Scan, MPI_Exscan) are designed for scalar or small-vector operations, not for the large matrix tensors with element-wise decay scaling that linear attention requires. The gap is not that scan operations don't exist — it's that existing scan operations assume the accumulation operation is a simple associative binary operator (like addition), while linear attention requires a scaled addition where the incoming state is element-wise multiplied by a decay vector before being added to the local state.

ZeCO's All-Scan fills this gap by implementing a custom scan that (a) handles the specific Sout=(γT1)Sin+SlocalS_{\text{out}} = (\gamma^T \mathbf{1}) \odot S_{\text{in}} + S_{\text{local}} update formula, (b) pipelines the communication at block granularity to minimize latency, and (c) operates on a separate CUDA stream to enable hardware-level overlap with computation. The result is not just a faster implementation of an existing primitive, but a genuinely new primitive that provides a communication pattern previously unavailable.

The significance of this as a contribution extends beyond linear attention. The scan pattern — where each element in a sequence needs the accumulated result from all preceding elements — appears in many recurrent computations: state-space models (like Mamba), prefix sum computations in sorting and ranking algorithms, cumulative distribution functions in statistics, and autoregressive generation where each token depends on all previous tokens. The paper's claim that All-Scan "provides a foundational innovation for advancing distributed computing in the linear model community" (Section 5) is an understated way of saying: we've built a primitive that could become as standard as All-Reduce is for gradient synchronization in data-parallel training, but for a different — and increasingly important — class of computations.

The empirical evidence (Figure 3, Table 2) supports this claim: All-Scan at 256 GPUs takes 2.16ms versus 8.51ms for All-Gather — a 3.9× speedup — despite having the correct communication volume for its use case (one state per device rather than all states). This demonstrates that purpose-built communication primitives can dramatically outperform general-purpose ones when the communication pattern is known and structured. If the linear model community grows as some predict, All-Scan could become the default communication backbone for distributed training of these architectures.

Innovation 4: Empirical Proof That Near-Ideal Scaling Is Achievable for Linear Attention SP — A Line-in-the-Sand Result

Prior to ZeCO, it was unclear whether sequence parallelism for linear attention could approach the ideal scaling curve of data parallelism (where throughput scales linearly with device count). LASP-1's serial execution made it clearly suboptimal. LASP-2's communication overhead made it degrade at scale. The field lacked a positive existence proof — a demonstration that some method could achieve near-linear throughput scaling for long sequences.

ZeCO provides that existence proof. Figure 4 (bottom rows) and Tables 5-6 show that ZeCO maintains near-constant per-GPU throughput as device count increases from 8 to 256 GPUs. For 16K sequences, per-GPU throughput drops from 44,497 tokens/sec at 8 GPUs to 34,400 at 256 GPUs — a 23% reduction, compared to LASP-2's 60% reduction (37,812 to 15,196). For 32K sequences, ZeCO drops from 47,369 to 40,967 (14% reduction), while LASP-2 drops from 42,946 to 25,402 (41% reduction). The ideal data-parallel baseline drops from 47,594 to 42,838 (10% reduction) for 16K, meaning ZeCO's degradation is only marginally worse than the inherent overhead of distributing any workload.

This is not just an incremental improvement over LASP-2. It is a line-in-the-sand result that establishes what is achievable: sequence parallelism for linear attention can, in practice, approach the theoretical upper bound of data-parallel throughput scaling. Any future method that claims to improve on ZeCO must either approach data-parallel scaling even more closely (which is hard, since ZeCO is already within ~10-20% of the ideal) or achieve comparable scaling with lower implementation complexity.

The broader implication is for the viability of long-context pretraining. The paper's abstract claim — "training a model with a 1M sequence length across 64 devices using ZeCO takes roughly the same time as training with a 16K sequence on a single device" — quantifies what this scaling means in practical terms. If a 1M-token sequence can be processed in the same wall-clock time as a 16K-token sequence on one GPU, then long-context pretraining becomes economically feasible in a way that it wasn't with LASP-1 or LASP-2. The bottleneck shifts from "can we efficiently distribute long sequences?" to "do we have enough GPUs and memory?", which is a much more tractable problem given modern GPU clusters.

This result also implicitly answers a question the paper doesn't pose explicitly: is there a fundamental reason why sequence parallelism must be less efficient than data parallelism? ZeCO's answer is no — the inherent communication requirement for linear attention SP is small (one fixed-size state tensor per device), and with proper scheduling, that communication can be almost entirely hidden. The inefficiencies of prior methods were artifacts of their design choices, not fundamental limitations of the problem.

Innovation 5: A Unified Framework for Analyzing SP Methods via Communication Volume, Computation Cost, and Extra Overhead

While the paper doesn't explicitly label this as an innovation, Appendix A.2 (Table 1) provides a unified analysis framework that is itself a conceptual contribution. By decomposing sequence parallelism methods into three dimensions — communication volume, computation cost, and additional computation overhead — the paper creates a structured way to compare methods across different attention architectures (full vs. linear) and different design choices (serial P2P vs. All-Gather vs. pipelined scan).

This framework reveals what was previously obscured: for linear attention, the communication volume's dependence on PP is the primary bottleneck, not the computation cost. Table 1 shows that all methods (including ZeCO) have computation costs that are independent of PP (or nearly so, with just NDe+NdNDe + Nd extra for ZeCO versus log(P)De+NDe\log(P)De + NDe for LASP-2). The difference is in communication: LASP-1 has PDePDe communication volume multiplied by serialization, LASP-2 has PDePDe communication volume from All-Gather, and ZeCO has DeDe — the theoretical minimum. This decomposition makes clear why ZeCO outperforms: it's not that ZeCO is generally "better," it's that ZeCO fixes the specific dimension (communication volume) where prior methods were suboptimal, while maintaining comparable performance on the other dimensions.

The framework also clarifies why full attention SP methods (Ulysses, Megatron CP) cannot achieve similar scaling: their computation costs grow with PP (L2DPL^2DP), making them fundamentally constrained by the attention algorithm itself, not just by communication. This distinction — between methods that are communication-bound (linear attention SP) and those that are computation-bound (full attention SP) — provides a clear diagnostic for where effort should be invested. For linear attention, invest in better communication primitives. For full attention, no amount of communication optimization can overcome the quadratic compute in sequence length.

This framework is not a one-time analysis; it is a reusable tool for evaluating future SP methods. Any new method can be characterized by its communication volume, computation cost, and extra overhead, and compared against the theoretical lower bounds for each dimension. The paper's optimality proof for ZeCO — showing it achieves the minimum in all three dimensions — establishes the Pareto frontier against which future methods can be benchmarked.

5. Experimental Analysis

Evaluation Methodology

  • Dataset and model configuration. The experiments use a 1B-parameter Gated Linear Attention (GLA) model with 20 layers and 32 attention heads, based on the Flash Linear Attention implementation (Yang and Zhang, 2024). The hidden dimension dd is 2048, and experiments are run with per-device sequence lengths LL of 16K and 32K tokens, chunk size CC of 64 tokens. The focus is purely on systems performance (communication speed, operator runtime, training throughput); no downstream task accuracy is measured.

  • Hardware. All experiments run on a cluster of 256 H100 80GB GPUs connected via NVLink and InfiniBand. The training framework is Meta Lingua (Videau et al., 2024), a PyTorch-based distributed training library. The scale is deliberately large — up to 256 GPUs and effective sequence lengths up to 8M tokens (P×L=256×32KP \times L = 256 \times 32\text{K}) — to stress the communication subsystem and reveal scaling bottlenecks that would not appear at small device counts.

  • Metrics. Three metrics are reported, at different levels of granularity. At the communication level, the metric is communication runtime (milliseconds), measured by timing the collective communication operation in isolation (5 warm-up rounds, average of 50 runs). At the operator level, the metric is SP algorithm runtime (milliseconds), measuring the total time for one forward and backward pass of the GLA attention operator including all communication, computation, and I/O. At the model level, the metric is per-GPU throughput (tokens per second per GPU), measured during end-to-end training of the full 1B-parameter GLA model (5 warm-up steps, average of 100 training steps). All three metrics are compared against an ideal baseline.

  • Baselines. Three baselines are used, spanning the design space of existing SP methods for linear attention. LASP-1 (Sun, Qin, et al., 2025) uses serial P2P communication — each device passes its final state to the next device in sequence — achieving minimal communication volume but enforcing strict serial execution. LASP-2 (Sun, Lan, et al., 2025; Li et al., 2025) uses All-Gather communication — each device collects local states from all other devices and independently performs the scan operation — achieving parallel computation but with communication volume that grows linearly with device count. The GLA (baseline) is the single-device GLA attention operator running without any sequence parallelism, serving as the theoretical lower bound for operator runtime and the theoretical upper bound for per-GPU throughput. At the communication level, All-Gather (the communication primitive underlying LASP-2) and All-Reduce (a standard collective for comparison) are also benchmarked.

  • Compute accounting. Communication runtime comparisons use the workload that each SP method actually requires for correct training. For LASP-2, this is All-Gather of PP state tensors of size dk×dvd_k \times d_v each. For ZeCO, this is All-Scan with one state tensor per device boundary, with the tensor partitioned into blocks of size 16,384. Operator runtime comparisons measure the full forward-backward pass time, which includes both communication and computation — this is the fairest micro-level comparison because it accounts for both the communication volume and the degree of parallelism. Throughput comparisons measure tokens/sec/GPU during actual training, which captures all system-level effects including communication-computation overlap, memory pressure, and framework overhead.

  • Statistical protocol. For communication runtime, each kernel is warmed up for 5 rounds, then 50 rounds are averaged. For operator runtime, 5 warm-up rounds are followed by averaging over 50 rounds. For model throughput, 5 warm-up steps are followed by averaging over 100 training steps. No cross-validation or statistical significance testing is reported — the metrics are deterministic system measurements on fixed hardware, so variance is expected to be small. However, the paper does not report standard deviations or confidence intervals, which would help assess whether small differences (e.g., ZeCO's 9.88ms vs. baseline 6.55ms at 128 GPUs in Table 3) are reliable or within noise.


Main Quantitative Results

Communication Speed: All-Scan Outperforms All-Gather by up to 3.9×

Figure 3 (right panel) and Table 2 present the communication runtime of All-Scan versus All-Gather (the communication primitive underlying LASP-2) and All-Reduce (a reference collective). The experiment measures the time to execute the communication pattern required by each SP method, with H=32H=32, d=4096d=4096, and L=8192L=8192 per device, across device counts from 8 to 256 GPUs.

  • At 8 GPUs, All-Scan takes 0.226ms versus 0.375ms for All-Gather — a 1.7× speedup.
  • At 64 GPUs, All-Scan takes 0.739ms versus 2.543ms for All-Gather — a 3.4× speedup.
  • At 256 GPUs, All-Scan takes 2.165ms versus 8.514ms for All-Gather — a 3.9× speedup.

The gap widens with device count because All-Gather's communication volume grows linearly with PP (each device must receive P1P-1 state tensors), while All-Scan's communication volume is constant at one state tensor per device. The All-Gather operation at 128 and 256 GPUs encounters out-of-memory errors (noted in the paper's Figure 3 description), suggesting that even the memory footprint of collecting all states becomes prohibitive at scale — a practical failure mode beyond just speed.

All-Reduce, included as a reference point, is faster than All-Scan (0.604ms at 256 GPUs) because it performs element-wise reduction of same-sized tensors, which is a simpler operation than the scan with decay scaling. This is not an apples-to-apples comparison — All-Reduce cannot implement the scan semantics needed for SP — but it provides a lower bound on what communication hardware can achieve, showing that All-Scan is within roughly 3.6× of the fastest possible collective on this hardware.

The headline result — 3.9× faster communication than the existing SP communication method — is strong evidence that All-Scan is a genuine systems contribution, not merely a theoretical one. It demonstrates that the pipelined block-wise design translates to wall-clock speedups on real GPU hardware.

SP Algorithm Runtime: ZeCO Approaches the Single-Device Baseline

Tables 3 and 4 and Figure 4 (top two rows) show the operator-level runtime — the total time for one forward and backward pass of the GLA attention operator using each SP method. The experiment varies the number of GPUs from 8 to 128 (for 16K sequences) and from 8 to 128 (for 32K sequences), with H=16H=16, d=2048d=2048, and C=64C=64.

For 16K sequences (Table 3):

  • The single-device GLA baseline takes 6.39–6.55ms across all GPU counts (the minor variation is measurement noise, since the baseline is always single-device).
  • ZeCO takes 7.32ms at 8 GPUs and 9.88ms at 128 GPUs — only 0.93ms to 3.33ms slower than the baseline, and the increase with GPU count is modest (2.56ms from 8 to 128 GPUs).
  • LASP-2 takes 19.39ms at 8 GPUs and 35.72ms at 128 GPUs — 13–29ms slower than baseline, and degrading by 16.33ms from 8 to 128 GPUs.
  • LASP-1 takes 22.59ms at 8 GPUs and 113.71ms at 128 GPUs — catastrophically worse, with runtime growing from 3.5× to 17.4× the baseline as GPU count increases.

For 32K sequences (Table 4), the pattern is similar but with larger absolute gaps:

  • The baseline takes 11.74–11.79ms.
  • ZeCO takes 12.12ms at 8 GPUs and 15.06ms at 128 GPUs — 0.38ms to 3.32ms slower than baseline.
  • LASP-2 takes 27.57ms at 8 GPUs and 42.50ms at 128 GPUs — 15.83ms to 30.71ms slower.
  • LASP-1 takes 41.64ms at 8 GPUs and 217.20ms at 128 GPUs — 29.90ms to 205.46ms slower.

The key observation is that ZeCO's runtime increase with GPU count is substantially smaller than either baseline's. From 8 to 128 GPUs, ZeCO adds 2.56ms (16K) or 2.94ms (32K), while LASP-2 adds 16.33ms (16K) or 14.93ms (32K). This is the direct consequence of the theoretical analysis in Section 3.3: ZeCO's communication cost is τ(dk×dv)+(P1)τ(dk×dv)K\tau(d_k \times d_v) + \frac{(P-1)\tau(d_k \times d_v)}{K}, where the PP-dependent term is divided by KK (the number of pipeline blocks), making it nearly constant when KK is large. LASP-2's All-Gather cost grows as P×τ(dk×dv)P \times \tau(d_k \times d_v) with no such mitigation.

The practical implication: at 128 GPUs with 32K sequences (4M total tokens), ZeCO processes a forward-backward pass in 15.06ms versus 42.50ms for LASP-2 — a 2.8× speedup. This is the operator-level complement to the 3.9× communication speedup, showing that faster communication translates nearly directly to faster end-to-end operator execution.

Model Throughput: Near-Linear Scaling from 8 to 256 GPUs

Tables 5 and 6 and Figure 4 (bottom two rows) report per-GPU throughput (tokens/sec/GPU) for the full 1B GLA model during training, with H=32H=32, 20 layers, d=2048d=2048, and per-device sequence lengths of 16K and 32K. This is the most important metric because it measures what practitioners actually care about: how many tokens per second can be processed in a real training run, accounting for all system overheads, not just the attention operator in isolation.

For 16K sequences (Table 5):

  • The data-parallel (DP) baseline achieves 47,594 tokens/sec/GPU at 8 GPUs and 42,838 at 256 GPUs — a 10.0% reduction, representing the inherent overhead of distributing the training workload (gradient synchronization, etc.).
  • ZeCO achieves 44,497 tokens/sec/GPU at 8 GPUs and 34,400 at 256 GPUs — a 22.7% reduction from 8 to 256 GPUs, and consistently within 6.5% to 19.7% of the DP baseline.
  • LASP-2 achieves 37,812 tokens/sec/GPU at 8 GPUs and 15,196 at 256 GPUs — a 59.8% reduction, falling from 79.5% to 35.5% of the DP baseline.
  • LASP-1 achieves 26,428 tokens/sec/GPU at 8 GPUs and degrades so severely at higher counts (12,596 at 64 GPUs) that measurements at 128 and 256 GPUs are not reported (marked "-" in the table).

For 32K sequences (Table 6), the same pattern:

  • DP baseline: 49,633 at 8 GPUs, 46,588 at 256 GPUs (6.1% reduction).
  • ZeCO: 47,369 at 8 GPUs, 40,967 at 256 GPUs (13.5% reduction, 87.9–95.4% of DP).
  • LASP-2: 42,946 at 8 GPUs, 25,402 at 256 GPUs (40.9% reduction, 54.5–86.5% of DP).
  • LASP-1: 27,014 at 8 GPUs, degrading to 12,268 at 64 GPUs, with higher counts not reported.

The scaling behavior is what makes these results significant. ZeCO's total throughput (per-GPU × PP) grows from 356K tokens/sec at 8 GPUs to 8.81M tokens/sec at 256 GPUs for 16K sequences — a 24.7× throughput increase from a 32× increase in GPUs, representing 77.3% scaling efficiency. For LASP-2, total throughput grows from 302K to 3.89M tokens/sec — only 12.9× increase from 32× GPUs (40.2% efficiency). For 32K sequences, ZeCO achieves 90.0% scaling efficiency (28.8× throughput from 32× GPUs) versus 61.5% for LASP-2 (19.7× from 32×).

The headline figure from the abstract — 60% throughput improvement on 256 GPUs with 8M sequences — corresponds to ZeCO's 34,400 tokens/sec/GPU versus LASP-2's 15,196 at 256 GPUs with 16K sequences (the 8M figure comes from P×L=256×32K=8MP \times L = 256 \times 32\text{K} = 8\text{M} total sequence length using 32K per device, where ZeCO achieves 40,967 versus 25,402 — a 61.3% improvement). This is the paper's strongest empirical claim: not just that ZeCO is faster, but that it fundamentally changes the scaling behavior, making sequence parallelism viable at scales where prior methods collapse.

Comparison Against the Theoretical Upper Bound

The paper implicitly benchmarks against two theoretical bounds: the single-device operator runtime (lower bound for how fast one device's work can be done) and the DP throughput (upper bound for how efficiently many devices can be used). ZeCO's operator runtime is within 3.33ms of the single-device bound at 128 GPUs with 32K sequences (15.06ms vs. 11.74ms), meaning communication and SP overhead add only 28% to the ideal operator time. ZeCO's throughput is within 12.1% of the DP bound at 256 GPUs with 32K sequences (40,967 vs. 46,588 tokens/sec/GPU), meaning the SP communication overhead translates to only a 12% throughput penalty at scale.

This is the practical implication of the theoretical claim that ϵ\epsilon (the extra computation and I/O) is negligible and that All-Scan's communication latency approaches τ(dk×dv)\tau(d_k \times d_v) as KK grows. The empirical gap between ZeCO and the theoretical bound — roughly 12–28% depending on metric and scale — provides an estimate of how much headroom remains for further optimization.


Ablation Studies and Robustness Checks

The paper does not present traditional ablation studies in the machine learning sense — there is no variation in model architecture, training hyperparameters, or data. However, several experimental design choices serve as implicit ablations that test the robustness of the claims.

Scale ablations via device count scaling (8 to 256 GPUs): All three metrics (communication runtime, operator runtime, model throughput) are evaluated across device counts spanning a 32× range. This serves as an implicit ablation of the scaling dimension: the near-constant per-GPU operator runtime and the near-linear total throughput growth confirm that ZeCO's design is robust to device count, with no emergent bottlenecks appearing at the largest scales tested. The fact that LASP-1 and LASP-2 degrade markedly at higher device counts (while ZeCO does not) is itself a robustness check — it demonstrates that the scaling behavior is not an artifact of the particular GPU count chosen.

Sequence length ablation via two per-device sequence lengths (16K and 32K): Tables 3-4 (operator runtime) and Tables 5-6 (throughput) report results for both L=16KL=16\text{K} and L=32KL=32\text{K}. The consistent relative ordering of methods across both sequence lengths, and the similar scaling patterns, confirms that ZeCO's advantage is not specific to a particular sequence length. The 32K results show slightly smaller relative gaps for ZeCO versus the DP baseline (12.1% at 256 GPUs for 32K vs. 19.7% for 16K), which is expected because longer sequences increase computation proportionally more than communication, diluting the communication cost. This is consistent with the theoretical analysis: Tideal-SP1(L)T_{\text{ideal-SP}}^1(L) grows with LL while τ(dk×dv)\tau(d_k \times d_v) stays constant, so the communication fraction of total time decreases.

Backward pass with All-Scan (Algorithm 3, Appendix A.4): The paper includes the full backward pass algorithm but does not separately benchmark forward versus backward communication. The operator runtime metric (forward + backward combined) implicitly tests that the backward All-Scan works correctly and efficiently, since any backward bottleneck would inflate the total runtime. The fact that total forward-backward time remains close to the baseline confirms that the backward pass does not introduce hidden costs.

Communication workload fairness (Appendix A.2, Table 1): The paper explicitly states that communication experiments use the workload "sufficient for correct training" for each method. This means All-Gather for LASP-2 communicates PP state tensors while All-Scan for ZeCO communicates 1 state tensor — a seemingly unfair comparison. However, this is the correct comparison because it reflects what each SP method actually requires: LASP-2 genuinely needs all states to compute the scan, while ZeCO genuinely needs only one. To equalize the communication workload would be to force ZeCO to communicate unnecessary data, which would test an algorithm nobody would implement. The fairness lies in comparing the communication cost of achieving the same end result (correct global states for all devices), not in equalizing intermediate data movement.

Negative result: LASP-1 and LASP-2 scaling collapse. The experiments document clear negative results for the baselines. LASP-1's throughput at 64 GPUs (12,596 tokens/sec/GPU for 16K) is less than half of its 8-GPU throughput (26,428), with measurements at 128 and 256 GPUs not even reported — likely because they were too slow or encountered errors. LASP-2's throughput degrades by 60% from 8 to 256 GPUs for 16K sequences, and the All-Gather communication encounters memory errors at 128 and 256 GPUs. These negative results are important because they demonstrate that the scaling problems ZeCO addresses are real and severe, not merely theoretical inefficiencies that don't manifest in practice.


Critical Assessment

Claim: "ZeCO achieves near-linear throughput scaling from 8 to 256 devices"

This claim is well-supported for the specific configuration tested (1B GLA model, 16K-32K per-device sequences, H100 GPUs). The data in Tables 5-6 show ZeCO's total throughput growing from 356K to 8.81M tokens/sec (24.7× for 32× GPUs) for 16K sequences, and from 379K to 10.49M (27.7× for 32× GPUs) for 32K sequences — both representing 77-87% scaling efficiency.

However, three qualifications are warranted. First, the comparison benchmarks ZeCO against a data-parallel baseline of the same model, but the DP baseline processes short sequences (16K-32K) on each GPU — it cannot process the 4M-8M total sequence lengths that ZeCO handles. The throughput comparison is fair in machine throughput terms (tokens/sec), but it doesn't capture that ZeCO enables training on sequence lengths that DP simply cannot handle due to memory constraints. The "approaching DP throughput" framing understates ZeCO's value: ZeCO achieves near-DP throughput while solving a problem DP cannot solve at all.

Second, the scaling efficiency is measured against ZeCO's own 8-GPU throughput, not against an extrapolation of the 1-GPU throughput. The paper does not report single-GPU ZeCO throughput (which would be the true ideal scaling reference), so the 77-87% efficiency is relative to an 8-GPU starting point that already includes some overhead. The true efficiency from 1 to 256 GPUs might be lower if the 1-to-8 scaling already incurs non-trivial overhead.

Third, the paper does not report scaling results at intermediate model sizes (e.g., 7B, 13B) or alternative configurations (different chunk sizes, different numbers of attention heads). The claim of "near-linear scaling" is empirically validated only for one model architecture at one scale, running on one GPU type. It is plausible that the scaling behavior would hold for larger models (since the communication volume depends on dk×dvd_k \times d_v, which grows with model size), but larger models also change the computation-to-communication ratio in ways that could affect the degree of overlap achievable.

Claim: "All-Scan achieves up to 3.9× communication speedup"

This claim is directly supported by Figure 3 and Table 2 for the specific tensor sizes and device counts tested. The 3.9× figure comes from 256 GPUs, where All-Scan takes 2.165ms versus All-Gather's 8.514ms — a clear and substantial gap.

The measurement methodology is sound: dedicated communication benchmarks with warm-up and averaging over 50 runs eliminate cold-start effects and provide reliable timing. However, the paper does not report communication time for the P2P send/recv pattern used by LASP-1, which would be a useful additional baseline. LASP-1's P2P communication theoretically has the same volume as ZeCO (one state tensor per boundary) but without pipelining, so it should be faster than All-Gather but slower than All-Scan — the paper does not quantify this.

More importantly, the communication benchmark measures the collective operation in isolation, not how much communication time remains after overlap with computation in the full training setting. The theoretical analysis proves that the overlapped portion Toverlapped_compT_{\text{overlapped\_comp}} reduces effective communication latency, but the paper does not experimentally separate overlapped from non-overlapped communication time in the full model throughput measurements. It is possible that the full-training communication advantage is smaller than the isolated-communication advantage if the overlap is imperfect (e.g., if intra-chunk computation completes before All-Scan finishes, leaving some communication exposed). The operator runtime and throughput results suggest the overlap is effective — ZeCO is close to the DP baseline — but they don't directly measure what fraction of communication is successfully hidden.

Claim: "ZeCO introduces only negligible time and space overhead"

This claim has theoretical support (the ϵ\epsilon analysis in Section 3.3) and circumstantial empirical support (ZeCO's operator runtime is close to the baseline), but no direct measurement. The paper proves that the extra I/O for γ~\tilde{\gamma} vectors is 1/dv1/d_v of state I/O and that the extra computation is O(dkdv)O(d_k d_v) versus O(Cdkdv)O(C d_k d_v) for the main computation, making it formally negligible. However, formal negligibility and practical negligibility are different: memory pressure, cache effects, and kernel launch overheads could make the "negligible" extra operations more costly than the FLOP count suggests.

A direct ablation would measure operator runtime with and without the γ~\tilde{\gamma} maintenance (e.g., by comparing a version that recomputes γ~\tilde{\gamma} from scratch versus storing it). The paper does not report such an experiment. The circumstantial evidence is that ZeCO's operator runtime is within 3.33ms of the baseline at 128 GPUs with 32K sequences (15.06ms vs. 11.74ms, a 28% overhead), and this gap includes both the extra computation/I/O and the non-overlapped communication. Since the theoretical non-overlapped communication is τ(dk×dv)\tau(d_k \times d_v) — which, based on the All-Scan time of 2.165ms at 256 GPUs, is roughly 1–2ms at 128 GPUs — the extra computation and I/O overhead is plausibly in the 1–2ms range, or roughly 10–17% of total operator time. This is small but not so negligible that it could be ignored in a rigorous accounting. The paper's claim of "negligible" is reasonable but would be stronger with direct measurement.

Claim: "Training a 1M sequence across 64 devices takes roughly the same time as training a 16K sequence on a single device"

This claim from the abstract is an extrapolation, not a directly measured result. The paper measures ZeCO at up to 256 GPUs with per-device sequences of 32K (8M total) and 16K (4M total), but the specific configuration of 64 devices with 1M total sequence length (i.e., L16KL \approx 16\text{K} per device) is never explicitly benchmarked. The claim is a reasonable extrapolation from the measured scaling behavior: at 64 GPUs with 16K per device, ZeCO achieves 39,832 tokens/sec/GPU (Table 5), which is 83.7% of the 8-GPU throughput — close enough to "roughly the same time" given the near-constant per-GPU throughput that ZeCO maintains. However, the paper should be clearer that this is an inference from the scaling curve, not a directly measured datapoint.

Missing experiments that would strengthen the paper

Scaling to different model sizes. All experiments use a 1B model. The communication volume depends on dk×dv=(d/h)×(d/h)d_k \times d_v = (d/h) \times (d/h), where dd is hidden dimension and hh is number of heads. A 7B model with larger dd would have proportionally larger state tensors, changing the ratio of communication to computation. Testing at multiple model scales would demonstrate whether the near-linear scaling generalizes or is specific to the 1B configuration.

Varying chunk size CC. The chunk size affects N=L/CN = L/C, the number of chunks per device, which in turn affects how many times γ~\tilde{\gamma} is stored/loaded and how much computation is available for overlap. The paper uses C=64C=64 throughout but does not justify this choice or explore sensitivity. Larger chunks mean fewer state updates (less extra I/O) but less granular overlap opportunities — there may be a sweet spot that the paper does not characterize.

Comparison against a simpler pipelined P2P baseline. The paper compares against LASP-1 (serial P2P, no overlap) and LASP-2 (All-Gather, redundant communication), but not against a straightforward optimization of LASP-1: P2P communication with manual overlap (send the final state to the next device while computing intra-chunk attention, without block-wise pipelining). This would help separate the benefit of overlap alone from the benefit of block-wise pipelining. If simple overlap achieved similar results, All-Scan's block-wise design would be less novel; if simple overlap was significantly worse, the block-wise approach would be more strongly validated.

Memory footprint analysis. The paper does not report peak GPU memory usage for each method. All-Gather for LASP-2 encounters OOM errors at 128-256 GPUs (noted in Figure 3), which is a significant practical failure mode that the paper mentions but does not quantify. A memory comparison would show whether ZeCO's constant communication volume also translates to constant memory overhead, which is as important as speed for very long sequences.

Gradient synchronization overhead in model throughput. The throughput measurements include the full training loop, which involves gradient synchronization across devices. The paper does not separate how much of the gap between ZeCO and the DP baseline is due to SP-specific communication versus general distributed training overhead (gradient All-Reduce, optimizer steps). If a substantial fraction of the 12-20% gap is from gradient synchronization that all methods share, then ZeCO's SP-specific overhead is even smaller than the throughput numbers suggest.

Overall assessment

The experimental section is strong in its primary goal: demonstrating that ZeCO achieves substantially better scaling than existing SP methods and approaches the ideal DP throughput bound. The multi-level evaluation (communication, operator, model) provides convergent evidence that the theoretical optimality proven in Section 3.3 translates to practical performance gains. The scaling range (8 to 256 GPUs) is sufficient to reveal the asymptotic behavior that distinguishes ZeCO from LASP-2, whose communication-bound degradation becomes severe at larger scales.

The main weakness is the narrowness of the experimental scope — one model size, one model architecture (GLA), one GPU type, two sequence lengths, one chunk size — which limits the generalizability of the claims. The paper's theoretical analysis suggests that ZeCO's advantages should hold broadly for linear attention models since they all share the same structural property (fixed-size state, linear decomposition), but without experiments on at least one other linear attention variant (e.g., Mamba, RetNet, cosFormer), the empirical evidence is confined to GLA. The lack of memory measurements and the absence of a simple P2P-with-overlap baseline are omissions that would be straightforward to address and would strengthen the experimental case.

That said, the core empirical finding — that ZeCO maintains 77-90% scaling efficiency from 8 to 256 GPUs while LASP-2 degrades to 40-62% — is robust and well-documented. The 3.9× communication speedup and 60% throughput improvement at 256 GPUs are specific, measured quantities that support the paper's central claim: redesigning the communication primitive for linear attention SP, rather than adapting general-purpose collectives, can transform sequence parallelism from a bottleneck into a near-transparent layer.

6. Limitations and Trade-offs

Single Model Architecture and Scale

The assumption or constraint. Every experiment in the paper—communication benchmarks, operator runtime, and end-to-end model throughput—uses exactly one model: a 1B-parameter Gated Linear Attention (GLA) model with 20 layers, 32 attention heads, hidden dimension 2048, and chunk size 64. The paper does not evaluate ZeCO on any other linear attention variant (e.g., Mamba, RetNet, cosFormer, HGRN2), nor on any model scale beyond 1B parameters. The authors position GLA as "one of the generalization forms of linear models" (Section 2.1), implying representativeness, but provide no empirical evidence that the scaling behavior transfers to other architectures or scales. This is particularly significant because linear attention is a diverse family—different variants use different state update rules (diagonal decay in GLA, data-dependent gating in Mamba, exponential decay in RetNet), and ZeCO's correctness proof in Appendix A.1 depends specifically on the linear decomposition property of the GLA recurrence S(p1)L+nC=(γ~[n]T1)S(p1)L+S[n]S_{(p-1)L+nC} = (\tilde{\gamma}_{[n]}^T \mathbf{1}) \odot S_{(p-1)L} + S_{[n]}, which may not hold in identical form for all linear attention mechanisms.

The consequence. A practitioner considering ZeCO for training, say, a Mamba model or a 7B GLA model cannot know from this paper whether the near-linear scaling will hold. Larger models increase dk×dvd_k \times d_v (the state size), which changes the ratio of computation to communication—the All-Scan time τ(dk×dv)\tau(d_k \times d_v) grows quadratically with dd if dkd_k and dvd_v are scaled proportionally, while the local computation per token grows as dk×dvd_k \times d_v (for the attention) and d×dffd \times d_{\text{ff}} (for the feed-forward layers). If the feed-forward layers dominate at larger scales, the communication fraction may shrink, making ZeCO's advantage over LASP-2 smaller; if attention dominates, the advantage may grow. The paper's theoretical framework (Section 3.3) predicts ZeCO remains optimal regardless, but the magnitude of the throughput gap against LASP-2—the 60% figure that headlines the paper—could vary substantially. Moreover, if a different linear attention variant requires communicating something other than a single state matrix (e.g., Mamba's state is larger or differently structured), the lower-bound proof of VZeCO(p)=dk×dvV_{\text{ZeCO}}^{(p)} = d_k \times d_v would need re-derivation, and All-Scan's block partitioning along the dkd_k dimension might need modification.

What evidence exists in the paper. None. The paper contains zero experiments varying model architecture, model size, or linear attention variant. The 1B GLA configuration is held constant across all experiments in Sections 4.1 and 4.2. The theoretical analysis (Section 3.3, Appendix A.1) is specific to the GLA recurrence, though the authors suggest in Section 5 that extending to "various forms, including matrix transform structures" is future work. The communication benchmarks (Figure 3, Table 2) isolate the collective operation and are independent of the model architecture, but the operator and throughput benchmarks are inherently tied to the specific GLA implementation from the Flash Linear Attention repository (Yang and Zhang, 2024).

Mitigation status. The paper does not attempt to address this limitation empirically. Section 5 explicitly lists "generalize the sequence parallelism algorithm for linear attention beyond diagonal decay, extending it to support various forms, including matrix transform structures" as a future direction, which is an acknowledgment that the current validation is narrow. The theoretical framework is general enough that extension should be possible, but without at least one additional architecture as a validation point, the claim that ZeCO is broadly applicable to linear attention models remains a hypothesis, not a demonstrated fact.


Difficulty Estimation Cost Is Not Accounted for in System Performance

The assumption or constraint. ZeCO's design makes a fundamental assumption about the communication pattern: each device needs exactly one incoming state tensor (S(p1)LS_{(p-1)L}) from its predecessor, and the All-Scan pipeline delivers it with near-optimal latency. This is a structural assumption about how the sequence is partitioned—contiguous chunks, each device handling N=L/CN = L/C chunks in order—not a runtime estimation that can be amortized. However, there is a practical deployment consideration that interacts with this assumption: the decision of how many devices to use for a given total sequence length (i.e., choosing PP for a desired total P×LP \times L) requires estimating the memory and throughput characteristics of the model at that scale. The paper does not discuss how a practitioner would determine the optimal PP a priori, nor does it model the cost of profiling runs to determine whether a given model configuration fits in GPU memory under ZeCO partitioning. While this is a general distributed training concern, it becomes more acute for ZeCO because the method's primary value proposition is enabling sequence lengths (1M–8M tokens) that are infeasible with other methods, meaning that misconfiguration—choosing too few devices and running out of memory, or too many and wasting resources—has higher stakes.

The consequence. In practice, deploying ZeCO for a new model configuration requires either: (a) analytical memory modeling to predict the peak memory for a given PP, LL, dd, and CC, which the paper does not provide; or (b) empirical profiling runs that consume GPU hours before productive training begins. For configurations near the memory limit (e.g., training the largest possible sequence length on available hardware), this profiling is non-trivial—a configuration that fails with an out-of-memory error may waste significant time. The paper's headline claim that "training a model with a 1M sequence length across 64 devices using ZeCO takes roughly the same time as training with a 16K sequence on a single device" (Abstract) implicitly assumes that the 64-device configuration is correctly chosen—that the sequence length fits in memory, that the per-device batch size is viable, and that the interconnect bandwidth is sufficient for the All-Scan pipeline to maintain its near-optimal latency. If any of these assumptions is violated, the claim becomes invalid, but the paper provides no guidance for ensuring they hold.

What evidence exists in the paper. The paper reports throughput results for specific, working configurations (1B model, 16K/32K per device, 8–256 GPUs in Tables 5–6), but does not report any failed configurations, memory footprint measurements, or configuration guidelines. The only mention of memory limitations is the observation in Figure 3 that All-Gather for LASP-2 encounters memory errors at 128–256 GPUs—a failure mode that ZeCO avoids, but that says nothing about ZeCO's own memory behavior at extreme scales. The paper does not provide a peak memory formula for ZeCO versus LASP-2 or versus DP, which would allow practitioners to predict whether a desired sequence length fits on their hardware without trial-and-error profiling.

Mitigation status. Not addressed. This is a practical deployment concern rather than a fundamental algorithmic limitation—ZeCO's communication overhead is near-optimal regardless of configuration—but it matters for the specific claim about enabling previously intractable sequence lengths. The paper could partially address this by providing memory analysis (ZeCO stores local states S[n]S_{[n]} for all NN chunks, plus the cumulative decay vectors γ~[n]\tilde{\gamma}_{[n]}, which adds N×dk×dv+N×dkN \times d_k \times d_v + N \times d_k elements to the HBM footprint compared to a method that doesn't store intermediate states) and showing that this overhead is manageable at the reported scales. The absence of memory analysis makes it difficult to assess whether ZeCO's throughput advantage comes partly at the cost of higher memory usage, which would be a meaningful tradeoff for practitioners operating near hardware limits.


No Results on Downstream Task Quality or Training Convergence

The assumption or constraint. All reported metrics in the paper are systems performance metrics: communication runtime (Figure 3, Table 2), operator forward-backward time (Tables 3–4), and training throughput in tokens/sec/GPU (Tables 5–6). The paper does not train any model to completion, report any validation loss curves, or evaluate any downstream task (e.g., language modeling perplexity, long-context retrieval accuracy, synthetic reasoning benchmarks). The implicit assumption is that ZeCO's sequence parallelism is a numerically equivalent reordering of the GLA computation—that is, it produces exactly the same outputs and gradients as the non-parallel GLA operator, just faster. This assumption is supported by the proof in Appendix A.1 showing that the global state update formula recovers the correct states, but the proof does not address potential numerical issues: floating-point accumulation order differences between the chunk-wise parallel computation and the sequential reference, or precision effects from the pipelined block-wise state updates in All-Scan versus a single full-tensor update.

The consequence. If ZeCO introduces subtle numerical differences (e.g., due to different summation order in the All-Scan pipeline versus sequential computation, or due to block-wise decay application producing slightly different floating-point results than full-tensor decay), these could accumulate over many training steps and affect convergence behavior or final model quality. This is not a hypothetical concern—distributed training methods that change computation order (e.g., model parallelism with pipeline scheduling, gradient accumulation with different microbatch boundaries) are known to produce non-bitwise-identical results that can, in rare cases, affect training dynamics. For ZeCO, the All-Scan pipeline splits the state update SpL(k)=(γ~[N](k)T1)S(p1)L(k)+S[N](k)S_{pL}^{(k)} = (\tilde{\gamma}_{[N]}^{(k)T} \mathbf{1}) \odot S_{(p-1)L}^{(k)} + S_{[N]}^{(k)} into KK independent blocks that are computed and forwarded at different times. If the decay factors γ~(k)\tilde{\gamma}^{(k)} are computed or applied differently across blocks (e.g., due to different numerical ranges in different head dimensions), the global state could diverge from the reference implementation in ways that the correctness proof (which assumes exact real arithmetic) does not capture.

A practitioner deploying ZeCO for production training would need assurance that the model trained with ZeCO SP achieves the same final quality as a model trained with a reference implementation (or at least that any quality difference is negligible). Without convergence curves, perplexity numbers, or task evaluations, the paper provides no such assurance. The paper's exclusive focus on throughput implicitly claims that "faster throughput → same model quality in less time," but this is an untested assumption.

What evidence exists in the paper. None. The paper contains zero training curves, zero loss values, and zero task evaluations. The only evidence that ZeCO is numerically correct is the theoretical proof in Appendix A.1 and the fact that the operator runs without errors (the throughput measurements imply successful forward and backward passes, but provide no information about numerical accuracy). The paper does not even report whether ZeCO's output tensors are bitwise-identical to a reference GLA implementation's output—a simple experiment that would take negligible compute compared to the scaling benchmarks and would substantially strengthen confidence in correctness.

Mitigation status. Not addressed. This is perhaps the most significant omission for a systems paper claiming to enable training of "next-generation LLMs on previously intractable sequence lengths" (Section 5). The claim implies that ZeCO can be used for real training runs, not just for throughput benchmarking. While it is common for systems papers to focus on performance metrics and defer quality evaluation to subsequent work (or to assume numerical equivalence), the complete absence of any correctness validation—even a comparison of output tensors for a single forward pass—leaves a gap between the demonstrated speedup and the claimed capability of enabling practical long-context training. The paper would be significantly strengthened by at minimum showing that ZeCO's forward and backward passes produce outputs within floating-point epsilon of a reference sequential implementation, and ideally by showing a short training run (even 100 steps) where loss curves match.


Chunk Size CC Is Fixed and Its Impact Is Unexplored

The assumption or constraint. All experiments in the paper use a fixed chunk size C=64C = 64 tokens (stated in Appendix A.3: "the tensor size of each chunk of segmentation is 16384"). The chunk size is a critical hyperparameter in chunk-wise linear attention because it controls the tradeoff between parallelism and accuracy: smaller chunks mean more chunk boundaries (more frequent state updates, more I/O for storing/loading γ~[n]\tilde{\gamma}_{[n]}, but finer-grained parallelism opportunities; larger chunks mean fewer boundaries (less overhead) but coarser granularity for the intra-chunk attention computation that ZeCO overlaps with All-Scan. For ZeCO specifically, the chunk size determines N=L/CN = L/C, the number of local chunk states that must be stored (NN matrices of size dk×dvd_k \times d_v each) and the number of cumulative decay vectors γ~[n]\tilde{\gamma}_{[n]} that must be maintained (NN vectors of size dkd_k each). It also determines how much computation is available for the overlap in Phase 2 of Algorithm 1: the intra-chunk attention computation is O(Cdk2)O(C d_k^2) per chunk, and this is what runs concurrently with All-Scan. If CC is too small, the intra-chunk computation may complete before All-Scan finishes, leaving communication exposed.

The consequence. A practitioner deploying ZeCO cannot know whether C=64C=64 is optimal for their configuration, or whether a different chunk size would yield better throughput or scaling. The theoretical analysis in Section 3.3 proves that the extra I/O overhead from γ~\tilde{\gamma} is 1/dv1/d_v per state access and 1/N1/N amortized across chunks—but this is expressed in terms of NN without evaluating how the choice of CC (and thus NN) affects the practical overhead in milliseconds. For very long sequences (e.g., L=32KL=32\text{K}), C=64C=64 gives N=512N=512, meaning 512 local states and decay vectors must be stored and later accessed. This is a non-trivial memory and I/O burden that is provably small in relative terms but may be absolutely large in a memory-constrained deployment. Furthermore, the overlap between All-Scan and intra-chunk attention depends on the intra-chunk computation time exceeding the All-Scan communication time—if CC is small, the intra-chunk computation is fast, and the overlap is incomplete, meaning effective communication latency is higher than the theoretical τ(dk×dv)/K\tau(d_k \times d_v)/K bound.

What evidence exists in the paper. The paper uses C=64C=64 throughout all experiments (Section 4.1, 4.2, and Appendix A.3), but never justifies this choice, never varies it in an ablation, and never reports sensitivity to chunk size. The paper does not provide the chunk-level computation time or the intra-chunk attention time separately from the total operator runtime, making it impossible for a reader to assess whether Phase 2's intra-chunk computation is indeed long enough to fully overlap with All-Scan. The theoretical Formula 12 shows that boundary overhead decreases as KK (the number of pipeline blocks for All-Scan) increases, but CC controls NN, not KK—these are independent parameters, and the paper does not explore their interaction.

Mitigation status. Not addressed. The fixed chunk size is a missing ablation that would be straightforward to perform—running the operator runtime benchmark (Tables 3–4) with C{32,64,128,256}C \in \{32, 64, 128, 256\} would reveal whether ZeCO's advantage is robust to chunk size or tuned to a specific sweet spot. The paper's theoretical framework does not predict a failure mode from varying CC, but the practical performance impact could be non-trivial, especially for configurations where the intra-chunk computation time is close to the All-Scan time. A practitioner implementing ZeCO for a different model (with different dkd_k, dvd_v, or sequence length) would need to tune CC empirically, and the paper provides no guidance.


The assumption or constraint. All experiments run on a single homogeneous cluster: 256×H100256 \times \text{H100} 80GB GPUs with NVLink for intra-node communication and InfiniBand for inter-node communication (Section 4). This is a top-tier hardware configuration with high-bandwidth interconnects that can support the fine-grained pipelined communication pattern of All-Scan. The All-Scan primitive relies on the ability to split a state tensor into KK blocks and stream them between devices with low latency—this requires not just high bandwidth but also low per-message overhead, which NVLink and InfiniBand provide. On hardware with weaker interconnect (e.g., PCIe-based GPU communication, older GPU generations with less NVLink bandwidth, or cloud instances with variable network performance), the pipeline efficiency would degrade because the per-block transmission time τ(dk×dv/K)\tau(d_k \times d_v / K) would increase, and the boundary overhead (P1)τ(dk×dv)K\frac{(P-1)\tau(d_k \times d_v)}{K} in Equation 12 would become larger relative to the computation time available for overlap.

The consequence. The paper's headline throughput numbers (e.g., 34,400 tokens/sec/GPU at 256 GPUs with 16K sequences, Table 5) are hardware-specific and may not generalize to more modest clusters. For a practitioner using, say, A100 GPUs with slower interconnect, or cloud instances where inter-node bandwidth is shared and variable, the degradation could be substantial. The theoretical analysis in Section 3.3 provides a lower bound on communication time (τ(dk×dv)\tau(d_k \times d_v)) but does not model how τ\tau depends on hardware—it's treated as a black-box function. In practice, τ\tau for a tensor of size dk×dvd_k \times d_v could easily be 2–5× larger on PCIe-connected A100s than on NVLink-connected H100s, which would directly increase the non-overlappable communication cost and reduce the throughput advantage over LASP-2. Moreover, the All-Scan pipeline requires reliable, ordered message delivery between adjacent devices—on networks with packet loss or reordering (which InfiniBand avoids but Ethernet does not), additional protocol overhead or retransmission delays could break the pipeline's steady-state throughput assumption.

What evidence exists in the paper. None beyond the hardware specification. The paper does not benchmark ZeCO on any other GPU type, any other interconnect technology, or any cloud configuration. The communication runtime benchmarks (Table 2, Figure 3) are measured on the specific H100 cluster and are likely near the best-case values achievable for the given tensor sizes. The paper does not provide a roofline model or communication cost model parameterized by bandwidth and latency, which would allow practitioners to estimate performance on their hardware. The statement that All-Scan "could run independently with other CUDA stream" (Section 3.2) assumes hardware support for concurrent communication and computation—a capability available on H100s but less effective on older architectures where the copy engine has lower bandwidth or where streams are not fully independent.

Mitigation status. Not addressed. Hardware generality is a common limitation of systems papers—it is expensive and time-consuming to benchmark on multiple hardware configurations—but the paper's strong scaling claims ("near-linear scalability," "training a model with a 1M sequence length across 64 devices using ZeCO takes roughly the same time as training with a 16K sequence on a single device") implicitly assume a high-performance interconnect. The paper would benefit from at least a theoretical analysis of how the throughput scales with interconnect bandwidth and latency, parameterizing τ(dk×dv)\tau(d_k \times d_v) as α+β×(dk×dv)\alpha + \beta \times (d_k \times d_v) (where α\alpha is per-message latency and β\beta is inverse bandwidth) so that practitioners can substitute their hardware parameters. Alternatively, a single experiment on a lower-bandwidth configuration (e.g., disabling NVLink and forcing communication through the host PCIe bus) would provide a lower bound on the expected throughput gap.


No Combined Evaluation of ZeCO with Other Parallelism Strategies

The assumption or constraint. ZeCO is presented and evaluated as a standalone sequence parallelism method, but real-world LLM training typically combines multiple parallelism strategies: data parallelism (DP) for batch scaling, tensor parallelism (TP) for distributing individual layer computations, pipeline parallelism (PP) for distributing layers across devices, and sequence parallelism (SP) for distributing long sequences. The paper's throughput experiments compare ZeCO-SP against a DP baseline (single-device sequence, replicated across devices with gradient synchronization) and against LASP-based SP methods, but never evaluates ZeCO in combination with TP or PP. The model configuration used (1B parameters, 20 layers, 2048 hidden dimension) fits comfortably on a single H100 (80GB), so there is no necessity for TP or PP in the reported experiments—but this also means the experiments avoid the communication interference and scheduling complexity that arise when multiple parallelism strategies share the same interconnect.

The consequence. In a realistic training run for a larger model (e.g., 7B or 13B parameters) where TP or PP is also required, the All-Scan communication would share NVLink and InfiniBand bandwidth with TP's All-Reduce operations and PP's P2P activation transfers. The paper does not analyze whether All-Scan's pipelined block-wise communication pattern interacts poorly with other collectives—for instance, if TP's All-Reduce bursts saturate the interconnect and cause head-of-line blocking for All-Scan blocks, the pipeline would stall, and the effective communication latency would increase beyond the theoretical τ(dk×dv)/K\tau(d_k \times d_v)/K bound. Conversely, All-Scan's continuous stream of small block transfers could fragment the interconnect and reduce throughput for other collectives. A practitioner deploying ZeCO in a hybrid parallelism setting cannot know from the paper whether the near-linear scaling holds when communication resources are shared.

Furthermore, ZeCO's design assumes a specific device ordering (a linear chain from device 0 to device P1P-1) for the All-Scan pipeline. In a 2D or 3D parallelism topology where devices are organized into TP groups, PP stages, and DP replicas, the SP group might not map cleanly onto the physical topology with adjacent NVLink connections. All-Scan's latency depends on devices pp and p+1p+1 being connected with the lowest possible latency—if the SP group is spread across nodes connected only by InfiniBand (higher latency than NVLink), the per-block transmission time increases, and the pipeline efficiency degrades.

What evidence exists in the paper. None. The paper does not discuss hybrid parallelism, does not report any experiment combining ZeCO with TP or PP, and does not analyze the interaction between All-Scan and other communication patterns sharing the interconnect. The related work discussion of full-attention SP methods (Section 2.3) mentions Ulysses's incompatibility with TP ("easy to implement but incompatible with tensor parallelism (TP) and limited by the number of heads"), implying that compatibility with other parallelism strategies is an important practical consideration, but the paper does not analyze ZeCO's own compatibility.

Mitigation status. Not addressed. This is a forward-looking limitation—the paper achieves its claimed near-linear scaling in isolation, and hybrid parallelism is likely the next step for practical deployment. Section 5 mentions investigating "efficient parallel topologies for sequence parallelism in large-scale models" as a future direction, which is an implicit acknowledgment that the current evaluation is simplified relative to production training setups. However, for a practitioner currently training large models with hybrid parallelism who wants to add ZeCO-SP for long-context training, the absence of any analysis or guidance on interoperability is a significant practical gap. The paper could partially address this by characterizing All-Scan's communication pattern (message sizes, frequency, bandwidth usage) in a way that enables a practitioner to reason about interference with other collectives, even without running the full hybrid experiment.

7. Implications and Future Directions

How This Work Changes the Landscape

ZeCO causes a reframing of sequence parallelism from an empirical tradeoff problem to a theoretically bounded optimization problem, and in doing so establishes that near-ideal throughput scaling for linear attention SP is achievable in practice—a previously open question.

The significance of this reframing becomes clear when we trace the implicit assumption embedded in prior work. Both LASP-1 and LASP-2 approached sequence parallelism as an engineering compromise: LASP-1 traded parallelism for minimal communication, LASP-2 traded communication for parallelism. Neither claimed optimality, and the field implicitly accepted that SP must incur some irreducible overhead—that distributing long sequences across devices was inherently less efficient than training on shorter sequences with data parallelism. ZeCO's contribution is not merely a better point on the tradeoff curve; it is the demonstration that the tradeoff curve itself is an artifact of suboptimal design, not a fundamental constraint. By proving that the theoretical minimum communication volume is dk×dvd_k \times d_v per device (Equation 9) and that this minimum can be achieved while overlapping communication with computation (Equation 12), ZeCO shows that linear attention SP can, in principle, approach the efficiency of single-device training. LASP-1 and LASP-2 were not failed attempts at a hard problem—they were solutions that didn't recognize the problem's true structure.

This reframing has a methodological consequence that extends beyond ZeCO itself: it establishes a template for analyzing distributed learning algorithms by first deriving the information-theoretic lower bound on communication, then designing a scheduling strategy that achieves that bound with maximal overlap. The paper's three-part decomposition—identify what must be communicated for correctness (the state SS), minimize that communication (one tensor per boundary via linear decomposition), and overlap it with independent computation (intra-chunk attention)—is a reusable analytical pattern. Future work on distributed training of state-space models, linear RNNs, or any architecture where a fixed-size hidden state summarizes history can follow the same recipe: prove the linear decomposition property, derive the minimum communication, then design the pipeline. The paper's own Table 1 in Appendix A.2, which decomposes methods into communication volume, computation cost, and extra overhead, is the skeleton of this template.

The paper also resolves a tension that was visible but unarticulated in the literature: why do linear attention models, which are algorithmically simpler and more efficient than full attention, struggle with distributed training? The answer, as ZeCO reveals, is that the systems community was applying communication patterns designed for full attention (All-Gather in LASP-2, adapted from Megatron CP) to a problem that has fundamentally different communication requirements. Full attention SP must communicate all KV pairs because the attention computation is non-separable—token tt's output depends on individual key-value pairs from all previous tokens, not just on a summary. Linear attention's separability (the Markov property of the state recurrence) means only the summary is needed. By recognizing this structural difference and designing a custom primitive (All-Scan) rather than adapting an existing one (All-Gather), ZeCO achieves what was previously assumed impossible. This is a case study in how algorithmic properties of a model architecture should drive the design of its distributed training systems, rather than treating systems design as architecture-agnostic.

In terms of which research directions become more or less attractive:

  • More attractive: custom communication primitives for emerging architectures. ZeCO demonstrates a ~4× communication speedup from a purpose-built primitive over a general-purpose one. As architectures diversify beyond the standard Transformer (state-space models like Mamba, linear attention variants, gated RNNs), the return on investment for designing architecture-specific collectives—rather than forcing everything through All-Gather, All-Reduce, and P2P send/recv—has now been quantified. This paper makes a strong case that the systems community should invest in building a library of scan-like primitives tailored to different recurrence structures, rather than assuming existing MPI collectives are sufficient.

  • More attractive: long-context pretraining from scratch. ZeCO's demonstration that training a 1M-token sequence on 64 GPUs takes roughly the same time as training a 16K sequence on a single GPU (Abstract) removes the throughput argument against full-length pretraining. Prior to ZeCO, the inefficiency of SP meant that even if you could fit a long sequence in memory across many GPUs, the training time penalty made it impractical. With ZeCO's near-linear scaling, the economic calculus shifts: the cost of pretraining on 1M-token sequences is now proportional to the number of tokens (as in standard training), not super-linearly penalized by communication overhead. This could enable a new generation of models that learn long-range dependencies during their foundational pretraining phase rather than in a post-hoc adaptation stage.

  • Less attractive: research on incremental improvements to All-Gather-based SP for linear attention. The paper's theoretical analysis shows that All-Gather-based approaches (LASP-2) have communication volume that grows linearly with PP, which is provably suboptimal when the minimum is O(1)O(1). Empirical results confirm that this matters at scale: LASP-2 degrades by 60% from 8 to 256 GPUs (Table 5) while ZeCO degrades by only 23%. This suggests that further optimization of All-Gather-based SP for linear attention—e.g., compressing the state tensors before gathering, or using hierarchical All-Gather—is unlikely to close the gap with scan-based approaches, because the gap arises from communicating redundant information, not from inefficient implementation of All-Gather. Research effort is better directed at designing scan primitives for other linear attention variants.

  • Less attractive: full attention SP for ultra-long sequences. The paper's Table 1 makes explicit what was implicit in the literature: full attention SP methods (Megatron CP, Ring Attention) have computation costs that grow with PP (L2DPL^2DP) and communication volumes that grow with both LL and PP. Even with perfect communication (zero latency, infinite bandwidth), the quadratic computation in sequence length makes full attention SP fundamentally unsuitable for the 1M+ token regime. ZeCO, by enabling linear attention SP to achieve near-ideal scaling, strengthens the case for linear attention as the architectural path forward for long-context models—if the systems bottleneck is solved, the algorithmic advantages of linear attention can be fully realized.


Follow-Up Research This Work Enables

Extending All-Scan to state-space models (Mamba, Mamba-2) and measuring the throughput gap versus ZeCO on identical hardware and sequence lengths. The paper evaluates ZeCO exclusively on Gated Linear Attention, but Mamba (Gu and Dao, 2024) and its variants represent a large and growing fraction of the linear recurrent model landscape. Mamba's state update differs from GLA in two important ways: the state size is larger (state dimension NN times head dimension dd), and the state transition is input-dependent (the A, B, C matrices are functions of the input, not fixed decay factors). A direct extension of ZeCO to Mamba would need to handle (a) larger state tensors, which would increase τ(dk×dv)\tau(d_k \times d_v) and potentially change the optimal number of pipeline blocks KK, and (b) input-dependent transitions, which might require communicating additional information beyond just the final state (since the transition matrices differ per token). A strong follow-up would implement All-Scan for Mamba, benchmark on the same H100 cluster at 1B and 7B scales, and report both throughput scaling and any necessary modifications to the communication pattern. A negative result—e.g., Mamba's larger state makes All-Scan's pipeline fill time dominant, or input-dependent transitions require communicating transition matrices that increase volume beyond one state tensor—would refine our understanding of when scan-based SP is optimal versus when other approaches are needed.

Training a model from scratch with ZeCO on 1M-token sequences and measuring downstream long-context performance versus a mid-training adaptation baseline. ZeCO's primary value proposition is enabling efficient long-context pretraining, but the paper provides no evidence that such training actually improves model quality. A critical follow-up would train two versions of a 1B GLA model on a standard pretraining corpus: one trained from scratch on 1M-token sequences using ZeCO-SP across 64 GPUs, and one pretrained on 16K sequences then adapted to 1M tokens through a mid-training phase (the current standard practice). Both models would be evaluated on long-context benchmarks (e.g., Long Range Arena, SCROLLS, needle-in-a-haystack retrieval at 128K+ tokens). The key question is whether full-length pretraining with ZeCO yields qualitatively different long-range reasoning capabilities compared to the adaptation approach—for instance, better performance on tasks requiring integration of information separated by hundreds of thousands of tokens, or more robust length generalization. This experiment would transform ZeCO from a systems contribution into a demonstration that systems efficiency unlocks new capabilities, which is the impact trajectory that justifies the paper's strong opening claims about LLM long-context capabilities.

Characterizing the numerical precision of ZeCO versus a sequential reference implementation and measuring the effect on training loss trajectories over 10K+ steps. The paper's correctness proof assumes exact real arithmetic, but GPU floating-point operations are non-associative—different summation orders produce different results. ZeCO's pipeline splits the state update into KK blocks that are computed at different times and accumulated with different ordering than a sequential scan, and the global output computation O[n]inter=Q~[n](S[n1]+(γ~[n1]T1)S(p1)L)O_{[n]}^{\text{inter}} = \tilde{Q}_{[n]}(S_{[n-1]} + (\tilde{\gamma}_{[n-1]}^T \mathbf{1}) \odot S_{(p-1)L}) adds the local and global contributions in a specific order that may differ from the reference implementation. A thorough follow-up would: (a) compare ZeCO's forward pass output tensors and backward pass gradient tensors against a sequential chunk-wise GLA implementation, reporting maximum absolute error and relative error; (b) train a 1B model for 10K steps with both ZeCO and the reference implementation (on short enough sequences that the reference doesn't OOM), plotting training loss and gradient norm trajectories; (c) if discrepancies emerge, identify whether they come from the block-wise state update, the γ~\tilde{\gamma} accumulation, or the separation of inter and intra-chunk computation. A negative result—ZeCO produces bitwise-different outputs that cause training loss to diverge after many steps—would be practically important and would motivate research into numerically stable SP-specific reduction trees. A positive result—training loss matches to within floating-point noise—would substantially increase confidence in ZeCO for production use.

Developing a performance model that predicts ZeCO's throughput given hardware parameters (bandwidth, latency, FLOPs) and model parameters (dd, hh, LL, CC), validated across at least two GPU generations. The paper's theoretical analysis provides asymptotic bounds but no quantitative throughput model. A practitioner considering ZeCO for their hardware needs to know: given my specific GPU type and interconnect, how many GPUs do I need for a given sequence length, and what throughput can I expect? Building a roofline-style model for ZeCO would require benchmarking the All-Scan kernel at various tensor sizes to measure α\alpha (per-message latency) and β\beta (inverse bandwidth) on different interconnects, measuring the intra-chunk attention computation time as a function of CC and dkd_k, and modeling the degree of overlap as a function of the ratio of these two quantities. This model could then be validated against measured throughput on H100s (the paper's hardware) and at least one older generation (A100s with NVLink 3.0) or one cloud configuration (A100s with PCIe). Such a model would transform ZeCO from a single-hardware result into a practical tool that practitioners can use for capacity planning, and would identify the hardware characteristics that most constrain ZeCO's scaling—is it bandwidth, latency, or computation throughput?

Combining ZeCO with tensor parallelism and pipeline parallelism, measuring throughput interference between All-Scan and TP All-Reduce on shared interconnects. Real-world training of models larger than 1B parameters requires hybrid parallelism. A direct follow-up would implement ZeCO in a framework supporting 2D or 3D parallelism (e.g., Megatron-LM or FSDP with tensor parallelism), train a 7B GLA model on 256 GPUs with ZeCO-SP + TP, and measure: (a) total throughput as a function of the SP group size and TP group size, compared to a baseline using only TP (shorter sequences) or using LASP-2 SP + TP (same sequence length); (b) whether All-Scan's pipelined block transfers cause head-of-line blocking for TP's All-Reduce operations when they share NVLink bandwidth, or vice versa; (c) the impact of SP group topology—whether mapping the All-Scan pipeline to devices within a single node (NVLink-connected) versus across nodes (InfiniBand-connected) changes throughput. This experiment would directly address the deployment gap between ZeCO's standalone evaluation and production training scenarios, and would generate practical guidance on how to configure hybrid parallelism with ZeCO-SP.

Investigating the optimal CC (chunk size) for ZeCO as a function of sequence length and model dimension, and whether adaptive chunking provides benefits. The paper uses a fixed C=64C=64 for all experiments without justification. The chunk size determines N=L/CN = L/C (the number of local states and decay vectors stored), the intra-chunk computation time (which is O(Cdk2)O(C d_k^2) and must be long enough to overlap All-Scan), and the number of chunk boundaries (each of which incurs an I/O cost for loading/storing γ~\tilde{\gamma}). A systematic study would vary CC from 16 to 256 for different LL (8K, 32K, 128K) and dd (2048, 4096) configurations, measuring operator runtime and memory footprint. The goal would be to produce a heuristic—e.g., "set CC such that intra-chunk attention time exceeds All-Scan time by at least 2× to ensure full overlap"—that generalizes across configurations. A more ambitious extension would explore dynamic chunking: use larger chunks for devices in the middle of the All-Scan pipeline (where the pipeline is fully utilized and overlap is most critical) and smaller chunks at the ends (where pipeline fill/drain time dominates), potentially improving overall throughput.


Practical Applications and Downstream Use Cases

Long-context pretraining of linear attention LLMs from scratch, rather than mid-training adaptation on top of a short-context base model. Current practice for long-context models (e.g., Gemini 1.5 Pro with 1M context, GPT-4 with 128K context) is to pretrain on relatively short sequences (4K–8K tokens) and then run a specialized continued pretraining phase with longer sequences to extend the context window. This two-phase approach exists partly because pretraining directly on long sequences has been computationally prohibitive—the quadratic attention cost in full Transformers, and the communication bottlenecks in linear attention SP. ZeCO's demonstration that training on 1M sequences across 64 GPUs takes roughly the same time as 16K on one GPU (Abstract) removes the economic barrier to single-phase long-context pretraining for linear attention models. An organization training a 7B GLA model could allocate 256 GPUs for pretraining directly on 1M-token documents, achieving approximately 40,967 × 256 ≈ 10.5M tokens/sec total throughput (extrapolating from Table 6's 32K per-device throughput). This would process a 1T-token corpus in approximately 26 GPU-hours—feasible for a well-resourced team—and would produce a model that learns long-range dependencies throughout its entire training, not just in a brief adaptation phase.

Cost-efficient fine-tuning of long-context models on domain-specific document collections. Many enterprise applications—legal document analysis, scientific literature review, financial report processing—require models that can reason over very long documents (100K+ tokens) in specialized domains. Fine-tuning a pretrained linear attention model on a corpus of such documents requires distributing each document across multiple GPUs via SP. With LASP-2, a fine-tuning run on 128 GPUs with 100K-token documents would see per-GPU throughput drop to roughly 50–60% of the ideal DP throughput (based on Table 5 extrapolation), wasting substantial compute. ZeCO's near-linear scaling means the fine-tuning cost is essentially proportional to the number of documents—the SP overhead is only ~10–20% above ideal. For a research lab fine-tuning on 100K legal opinions or 100K scientific papers, ZeCO could reduce the GPU-hours required by ~40% compared to LASP-2 (based on the throughput gap at 128 GPUs in Table 5: 37,955 vs. 22,386 tokens/sec/GPU), making domain-adaptive long-context models more accessible to teams without massive compute budgets.

Enabling on-device or edge-deployment training scenarios where sequence length exceeds single-device memory. While ZeCO is evaluated on datacenter GPUs, its core principle—minimal communication of a fixed-size state—applies to any distributed setting where devices are connected with some bandwidth. Consider a scenario where a small linear attention model (e.g., 300M parameters) needs to process a very long sensor stream (e.g., 500K time steps from wearable health monitors) distributed across multiple edge devices (phones, embedded GPUs) connected via WiFi or Bluetooth. The communication volume per device boundary is dk×dvd_k \times d_v floats—for a model with d=1024d=1024, h=16h=16, this is 64×64=409664 \times 64 = 4096 floats ≈ 16KB—which is feasible even on low-bandwidth links. ZeCO's pipelined All-Scan would overlap this communication with local intra-chunk computation on each edge device, enabling near-linear throughput scaling even with consumer-grade interconnect. This is not a near-term application (edge training of LLMs is nascent), but ZeCO provides the communication template that would make it feasible.

Self-improvement and iterative training pipelines that alternate between generation and training on long sequences. When using LLMs to generate training data for themselves (STaR, ReSTEM^{EM}-style self-improvement), the model must process both the prompt and its own long generated outputs as context. If the prompt-plus-generation exceeds typical training sequence lengths, SP is necessary. ZeCO makes these self-improvement loops economically viable: the throughput penalty for training on long (prompt + generation) sequences is small, so the cost of each self-improvement iteration is dominated by the generation phase (which benefits from linear attention's O(L)O(L) complexity for autoregressive decoding) rather than the training phase. This could enable self-improvement on tasks requiring very long reasoning chains—mathematical proofs, code generation with long execution traces, multi-step planning—where the model's own outputs serve as training data for subsequent iterations, with each iteration's context window growing as the model learns to produce longer and more sophisticated outputs.


When to Prefer This Method

The paper articulates a clear tradeoff between ZeCO and existing SP methods based on the number of devices and the sequence length, grounded in the theoretical analysis of Equations 15–16 and the empirical scaling data in Tables 3–6. The following decision rule emerges:

  • Prefer ZeCO over LASP-1 when using more than a minimal number of GPUs. LASP-1's serial execution causes total time to grow as P×Tideal-SP1(L)P \times T_{\text{ideal-SP}}^1(L), making it fundamentally unscalable. The operator runtime data in Table 3 shows LASP-1 at 113.71ms versus ZeCO at 9.88ms on 128 GPUs with 16K sequences—a 11.5× gap. LASP-1 should only be considered for single-digit GPU counts where the serial dependency chain is short enough that the simplicity of P2P send/recv (no pipelining, no block management) outweighs the throughput penalty. The paper does not identify a crossover point where LASP-1 becomes competitive—the gap is already 3.1× at 8 GPUs (22.59ms vs. 7.32ms in Table 3)—suggesting LASP-1 is strictly dominated for any multi-GPU setting.

  • Prefer ZeCO over LASP-2 when scaling to device counts where communication volume dominates. The theoretical analysis in Equation 16 shows LASP-2's communication cost is P×τ(dk×dv)P \times \tau(d_k \times d_v) versus ZeCO's approximately τ(dk×dv)\tau(d_k \times d_v). At small PP (e.g., 4–8 GPUs), LASP-2's All-Gather overhead is modest—Table 3 shows LASP-2 at 19.39ms vs. ZeCO's 7.32ms at 8 GPUs, a 2.6× gap that may be acceptable if the implementation complexity of All-Scan is a barrier. At large PP (128+ GPUs), the gap widens to 3.6× (35.72ms vs. 9.88ms), and LASP-2 encounters out-of-memory errors for All-Gather at 128–256 GPUs (Figure 3), making it not just slower but non-viable. The crossover point—where LASP-2's simpler implementation might be preferred despite lower throughput—is not quantified in the paper, but based on Tables 3–4, the gap is already substantial at 16 GPUs (24.48ms vs. 7.45ms), suggesting LASP-2 is only competitive for very small deployments where communication is not yet a bottleneck.

  • Prefer ZeCO when training on sequences that exceed single-GPU memory. The paper does not provide peak memory measurements, but the structural design of ZeCO—communicating one state tensor rather than gathering PP state tensors—means its memory footprint for SP-specific tensors is O(dk×dv)O(d_k \times d_v) rather than O(P×dk×dv)O(P \times d_k \times d_v). LASP-2's All-Gather requires storing all PP local states simultaneously to perform the scan, which the paper notes causes OOM errors at 128–256 GPUs in the communication benchmarks (Figure 3). If the total sequence length pushes the limits of available GPU memory, ZeCO's constant memory overhead is a hard requirement, not just a throughput preference.

  • Prefer ZeCO when the computation-to-communication ratio is favorable for overlap. ZeCO's throughput advantage depends partly on successful overlap of All-Scan with intra-chunk attention computation. If the intra-chunk computation is too short relative to All-Scan's latency—for example, with very small chunk sizes CC, small hidden dimensions dkd_k, or very high-bandwidth interconnects where communication is extremely fast—the overlap may be incomplete, and ZeCO's advantage over a well-optimized P2P implementation may narrow. The paper does not characterize this boundary in the experimental results (chunk size is fixed at C=64C=64), so practitioners using configurations very different from the 1B GLA model on H100s should profile the degree of overlap before assuming the full throughput advantage.