Keywords: Sequence Parallelism; Communication-Computation Overlap; Distributed Attention; CUDA Streams; Long Sequence Processing

1 Introduction↩︎

Large-scale Transformer models [1] have achieved breakthrough advances in natural language processing [2], [3], computer vision [4], and multimodal generation. Video generation represents one of the most prominent application scenarios [5]–[8]. Taking models such as Sora [9] and Wan2.2 as examples, generating 2-4 second videos at 720P resolution corresponds to sequence lengths of approximately 31K-56K tokens in latent space. This medium-long sequence processing requirement makes sequence parallelism a critical technique—distributing sequences across multiple GPUs to accelerate computation.

However, all-to-all based sequence parallelism schemes, represented by DeepSpeed Ulysses [10], exhibit significant efficiency issues: communication and computation execute strictly in serial, causing hardware resource underutilization. We observe that under typical configurations of 4-8 GPUs with PCIe interconnect and 31K-100K token sequence lengths, the communication ratio \(\rho\) typically ranges from 15% to 40%. This characteristic provides substantial room for communication optimization—through communication-computation overlap, a theoretical speedup upper bound of \(1/(1-\rho)\) can be achieved.

Existing sequence parallelism methods each have their limitations. Ring Attention [11] employs a ring communication pattern where overlap depends on block-level pipelined execution; when the number of blocks is small, overlap effectiveness is limited, and \(P-1\) rounds of serial communication are required. Although DeepSpeed Ulysses requires fewer communication rounds, its original implementation executes communication and computation completely in serial, failing to exploit overlap optimization.

A fundamental property of multi-head attention is that computations across different heads are mutually independent. This independence implies that the serial dependency of “communication \(\rightarrow\) computation” can be relaxed at the head granularity—rather than waiting for all heads’ input data to complete communication before starting computation, computation for any head can begin immediately once its data is ready. This property provides the theoretical foundation for achieving fine-grained communication-computation overlap at the head dimension.

Based on this insight, we propose the Head-Chunked Multi-Stream Pipeline (HCMS) method, which partitions attention heads into multiple independent chunks and achieves fine-grained pipelining through dual CUDA streams. HCMS supports uneven chunk partitioning with performance variance less than 1%, and is orthogonally compatible with existing attention optimizations such as FlashAttention and SDPA without requiring kernel modifications. Experiments demonstrate that HCMS achieves significant speedups across four GPU platforms at 2-8 GPU scales: for typical video generation sequence lengths of 31K-56K tokens, 4-GPU configurations achieve 17.5% speedup and 8-GPU configurations achieve 16.4% speedup, outperforming Ring Attention by 5%-14.5%. In end-to-end validation on the Wan2.2 video generation model, attention layers achieve 18.3% speedup and the complete pipeline achieves 6.8% speedup.

The main contributions of this paper are as follows:

• Head Chunking Partitioning Strategy: We leverage the computational independence of multi-head attention to partition attention heads into \(C\) chunks, thereby decoupling the originally serial communication-computation dependency into \(C\) independently schedulable subtasks, achieving fine-grained parallelism without altering computational semantics. This method supports uneven partitioning (when \(H \bmod C \neq 0\)), and experiments show that performance variance compared to even partitioning is less than \(1\%\).

• Theoretical Modeling and Applicability Analysis: We establish a performance model for HCMS and derive the optimal chunk count \(C^* = \sqrt{T_{comm}/\beta}\), clearly defining its applicability: significant benefits can be achieved when communication ratio \(\rho > 20\%\), while benefits are limited when \(\rho < 10\%\).

• Cross-Platform Experimental Validation: Across four GPU platforms with \(2\)–\(8\) GPU configurations, for sequence lengths of \(31\text{K}\)–\(56\text{K}\) tokens, HCMS achieves \(5\%\)–\(14.5\%\) speedup over Ring Attention and \(10\%\)–\(17.5\%\) speedup over the Ulysses baseline. Furthermore, the method is fully compatible with PyTorch autograd and can be directly used for distributed training, achieving \(3.5\%\) training speedup while maintaining numerical equivalence.

2 Related Work↩︎

2.1 Sequence Parallelism Methods↩︎

Sequence parallelism distributes the sequence dimension across multiple devices [12], [13], primarily including ring communication-based methods and all-to-all based methods. The computations of different heads in multi-head attention are mutually independent, and this computational independence forms the theoretical foundation of our method.

Ring Attention [11] employs a ring communication pattern, partitioning sequences into \(P\) blocks and achieving distributed attention computation through \(P-1\) rounds of KV block passing. Its advantage lies in high memory efficiency, enabling processing of arbitrarily long sequences. However, Ring Attention’s communication-computation overlap is limited by block granularity: overlap is only effective when computation time exceeds communication time. Furthermore, \(P-1\) rounds of serial communication make total latency proportional to \(P\). Subsequent work such as Striped Attention [14] optimized for causal attention, but the video generation scenarios we focus on use bidirectional attention, making these optimizations not directly applicable.

DeepSpeed Ulysses [10] employs an all-to-all communication pattern: the input all-to-all rearranges data from “sequence-sharded, heads-complete” to “sequence-complete, heads-sharded”, and after attention computation, the output all-to-all performs the inverse transformation. Compared to Ring Attention’s \(P-1\) communication rounds, Ulysses requires only 2 rounds of all-to-all, with fewer communication rounds and better compatibility with FlashAttention [15], [16]. However, in Ulysses’ original implementation, communication and computation execute completely in serial, which is precisely the optimization target of this paper.

Megatron Context Parallelism [17], [18] provides both ring communication and all-gather modes. USP [19] unifies Ring and Ulysses into a single framework and proposes tile-based communication overlap techniques; LoongTrain [20] employs 2D parallelism across sequence and head dimensions to support million-scale token training; DistFlashAttn [21] extends FlashAttention to distributed scenarios; Colossal-AI [22] and LightSeq [13] provide unified distributed training frameworks. These works focus on the combination and scheduling of parallelism strategies, while HCMS focuses on the underlying communication-computation overlap implementation based on Ulysses; the two are orthogonal and composable.

In summary, existing sequence parallelism methods have the following limitations in communication-computation overlap: (1) Ring Attention’s overlap granularity is constrained by block size, and requires \(P-1\) communication rounds; (2) Ulysses achieves fewer communication rounds but lacks overlap mechanisms in its original implementation; (3) Tile-based overlap methods require modifications to underlying attention kernels, limiting their compatibility. HCMS addresses these gaps by introducing head-level chunking that enables fine-grained overlap while maintaining native compatibility with existing attention implementations.

2.2 Communication-Computation Overlap↩︎

Communication-computation overlap is a classical technique in distributed optimization [23], [24], widely applied in gradient synchronization [25], [26], ZeRO [27] parameter prefetching, and pipeline parallelism [28]–[30]. In the domain of distributed attention, tile-based methods [31], [32] achieve communication-computation overlap at the matrix level but require modifications to underlying attention kernels. In contrast, HCMS operates at a higher abstraction level—the attention head dimension—and is natively compatible with existing attention optimizations such as FlashAttention and SDPA, enabling orthogonal composition.

3 Method↩︎

3.1 Problem Formulation↩︎

Consider all-to-all based sequence parallel attention [10]. Input \(\boldsymbol{X} \in \mathbb{R}^{B \times (L/P) \times D}\) is distributed across \(P\) GPUs, with the forward pass consisting of: QKV projection \(\rightarrow\) RoPE [33] positional encoding \(\rightarrow\) input all-to-all \(\rightarrow\) attention computation \(\rightarrow\) output all-to-all \(\rightarrow\) output projection. The input all-to-all rearranges data from sequence-sharded to head-sharded, and the output all-to-all performs the inverse transformation. In the baseline implementation, communication and computation execute strictly in serial, with total execution time: \[T_{baseline} = T_{comm} + T_{attn} + T_{other}, \quad T_{comm} = T_{in} + T_{out}\] where \(T_{other}\) represents fixed overhead such as QKV projection and positional encoding. Defining the communication ratio \(\rho = T_{comm}/T_{baseline}\), under 4-8 GPU PCIe configurations \(\rho\) typically ranges from 15% to 40%, making communication-computation overlap a substantial optimization opportunity.

3.2 Method Overview↩︎

The core idea of HCMS is to partition attention heads into multiple chunks and achieve pipelined overlap between communication and computation through dual CUDA streams. Figure 1 illustrates the overall architecture.

As shown in the figure, HCMS comprises three key components: (1) Head Chunking partitions attention heads into \(C\) independent chunks; (2) the dual-stream pipeline uses communication stream \(S_{comm}\) and compute stream \(S_{comp}\) for overlapped execution; (3) CUDA Event synchronization ensures correctness of data dependencies. We detail each component below.

3.3 Head Chunking Strategy↩︎

Based on Property 1, we can partition \(H\) attention heads into \(C\) chunks, with each chunk independently completing communication and computation. Let the \(c\)-th chunk contain \(H_c\) heads: \[H = \sum_{c=0}^{C-1} H_c, \quad H_c = \left\lfloor \frac{H}{C} \right\rfloor + \mathbb{1}_{c < H \bmod C}\]

After partitioning, the QKV tensors are correspondingly decomposed: \[\boldsymbol{Q} = [\boldsymbol{Q}_0 \| \boldsymbol{Q}_1 \| \cdots \| \boldsymbol{Q}_{C-1}]\] where \(\|\) denotes concatenation along the head dimension, and \(\boldsymbol{Q}_c \in \mathbb{R}^{B \times (L/P) \times H_c \times D}\).

Lemma 1 indicates that communication and computation of different chunks can be scheduled independently, providing the theoretical foundation for pipelined overlap.

3.4 Dual-Stream Pipeline Architecture↩︎

Based on Head Chunking, we design a dual CUDA stream pipeline architecture:

• Communication stream \(S_{comm}\): dedicated to executing all-to-all collective communication

• Compute stream \(S_{comp}\): dedicated to executing attention computation

Pipeline execution consists of three phases, as shown in Figure 2:

Phase 1 - Input Communication: The communication stream sequentially initiates input all-to-all for each chunk, recording event \(E_c^{comm}\) upon completion of each chunk’s communication.

Phase 2 - Attention Computation: The compute stream monitors \(E_c^{comm}\) events; once chunk \(c\)’s input data is ready, computation begins immediately. Event \(E_c^{comp}\) is recorded upon completion.

Phase 3 - Output Communication: The communication stream monitors \(E_c^{comp}\) events; once chunk \(c\)’s computation completes, output all-to-all is initiated immediately.

3.5 CUDA Event Synchronization Mechanism↩︎

Pipeline correctness depends on precise dependency management. We define two types of CUDA Events:

• \(E_c^{comm}\): recorded when chunk \(c\)’s input all-to-all completes

• \(E_c^{comp}\): recorded when chunk \(c\)’s attention computation completes

Dependencies are expressed as: \[\begin{align} \text{Attn}_c &\leftarrow \text{wait}(E_c^{comm}) \\ A_c^{out} &\leftarrow \text{wait}(E_c^{comp}) \end{align}\]

Key observation: Due to Lemma 1, \(A_{c+1}^{in}\) does not need to wait for \(\text{Attn}_c\) or \(A_c^{out}\) to complete. This means: \[A_{c+1}^{in} \parallel \text{Attn}_c \parallel A_c^{out}\] All three can execute in parallel, subject to hardware resource constraints.

3.6 Algorithm Description↩︎

Algorithm 3 provides the complete pseudocode for HCMS.

3.7 Theoretical Analysis↩︎

We establish a theoretical performance model for HCMS, deriving the optimal chunk count and speedup upper bound; detailed derivations are provided in Appendix 7. Let communication time be \(T_{comm}\), per-chunk overhead be \(\beta\), and communication ratio \(\rho = T_{comm}/T_{total}\). The core conclusions are:

Applicability: HCMS benefits are positively correlated with communication ratio \(\rho\). When \(\rho > 20\%\), HCMS is recommended and can achieve 3%-24% speedup; when \(\rho < 10\%\), benefits are limited to less than 1%. Factors affecting \(\rho\) include GPU count, sequence length, and interconnect type, with PCIe interconnect scenarios typically having higher \(\rho\) than NVLink, making HCMS benefits more pronounced.

3.8 Comparison with Ring Attention and Ulysses↩︎

Table 1 compares HCMS with existing methods across multiple dimensions.

3pt

Comparison with Ring Attention: Ring Attention’s overlap occurs at the block level, communicating one complete KV block per round, while HCMS’s overlap occurs at the head level with finer granularity. Regarding communication rounds, Ring Attention requires \(P-1\) rounds of serial communication, while HCMS requires only 2 rounds. The two have different optimal scenarios: Ring Attention excels in very long sequence, memory-constrained scenarios, while HCMS performs better in medium-long sequence scenarios with high communication ratio.

Comparison with Ulysses: The two differ fundamentally in their optimization dimensions. Ulysses focuses on the choice between all-to-all and P2P communication patterns, treating communication and computation as atomic operations executed serially; HCMS introduces Head Chunking as a new partitioning abstraction, leveraging the computational independence of multi-head attention (Property 1) to decompose atomic operations into independently schedulable subtasks, thereby achieving fine-grained pipelining at the head dimension. This concept is general: any all-to-all based sequence parallelism scheme can introduce communication-computation overlap through Head Chunking, with HCMS’s application to Ulysses being just one instance.

4 Experiments↩︎

This section validates the effectiveness of HCMS. The core claims are:

• Performance improvement: HCMS achieves 10%-17.5% speedup over Ulysses and 5%-14.5% over Ring Attention for typical video generation sequence lengths of 31K-56K tokens, validated across four GPU platforms at 2-8 GPU scales.

• End-to-end benefits: In the Wan2.2 video generation model, attention layers achieve 18.3% speedup and the complete pipeline achieves 6.8% speedup.

• No accuracy loss: HCMS output is identical to baseline across all configurations, supports training scenarios with 3.5% speedup, and reduces peak memory by 8.7%.

Appendix 9 contains complete experimental data and supplementary analysis.

4.1 Experimental Setup↩︎

Hardware environment: To validate generalization, we conduct experiments on four GPU platforms with configurations shown in Table 2. Software environment is PyTorch [34] 2.9.1, CUDA 12.8, collective communication via NCCL [35], with BF16 [36] numerical precision.

3pt

Model configuration: We use typical parameters from video generation models [37], [38], shown in Table 3.

Evaluation metrics: We use end-to-end latency, throughput, and speedup ratio for performance evaluation, while verifying numerical consistency through maximum/mean error.

Baseline method: We use the standard implementation of DeepSpeed Ulysses [10] as the baseline. Ulysses is the mainstream sequence parallelism approach, with communication pattern consisting of input all-to-all, attention computation, and output all-to-all in three stages. In the original Ulysses implementation, communication and computation execute strictly in serial, which is precisely the optimization target of HCMS.

Implementation details: HCMS employs device-level stream caching to avoid frequent creation overhead, and KV Cache uses a chunk-organized list structure to support chunked computation. When \(H\) is not divisible by \(C\), a “more-first-less-later” uneven partitioning strategy is used; experiments show performance variance is less than 1%.

4.2 Main Results↩︎

Table 4 shows performance across different GPU configurations on each platform.

From the GPU scaling perspective, speedup correlates positively with communication ratio \(\rho\): the A10 platform with higher \(\rho\) achieves 12.6%-13.9% speedup; the L20 platform achieves 3.7%-5.4% speedup; the RTX 5090 8-GPU configuration achieves 4.4% speedup, validating HCMS effectiveness in large-scale parallel scenarios.

Table 5 shows performance comparison across different sequence lengths.

3pt

Experimental results show clear sequence length dependency: shorter sequences yield higher speedups. The 56K sequence achieves 17.5% speedup on 4090 4-GPU configuration and 16.4% speedup on 5090 8-GPU configuration. The 4090 platform achieves overall higher speedups than L20 because both platforms use PCIe 4.0 interconnect with similar communication time, while 4090 has shorter total execution time (182.8ms vs 234.5ms), resulting in higher communication ratio \(\rho\) and thus greater HCMS benefits. The A10 platform achieves 12.6%-13.9% speedup at 131K tokens, further validating the “higher communication ratio, greater benefits” principle. For 8-GPU scalability, the 5090 platform still achieves 16.4% speedup at 56K tokens, while benefits decrease to 0.9% at 281K tokens, consistent with theoretical expectations. All four platforms (L20, 4090, A10, 5090) at 2-8 GPU scales validate HCMS effectiveness, demonstrating cross-hardware, cross-scale generalization.

Reason analysis: Attention computation complexity is \(O(L^2)\), while communication complexity is \(O(L)\). As sequences grow longer, computation time increases much faster than communication time, causing communication ratio \(\rho\) to decrease and HCMS overlap benefits to diminish.

Figure 4 visualizes this trend.

Comparison with Ring Attention: Table 6 presents a three-way comparison of HCMS, Ring Attention, and Ulysses. For typical video generation sequence lengths of 31K-56K tokens, HCMS leads comprehensively: 14.5% faster than Ring at 31K tokens, 5.0% faster at 56K. For longer sequences above 131K, Ring is slightly better by 1.7%, with the crossover point at approximately 60-80K tokens.

4.3 Ablation Studies↩︎

Table 7 shows the effect of different chunk counts on performance, with experimental configuration of 4 GPUs and 131K tokens.

Results show that as \(C\) increases from 1 to 5, speedup continuously improves because more chunks enable finer-grained communication-computation overlap. However, performance begins to decline when \(C\) exceeds 5, due to accumulated Event synchronization overhead and efficiency loss from overly fine computation granularity. In this configuration, the optimal point is \(C=5\), but \(C=4\) is already near-optimal and simpler to implement.

Uneven partitioning robustness: When \(H\) is not divisible by \(C\), a “more-first-less-later” strategy is used. Table 8 shows partitioning schemes and performance for different \(C\) values, with configuration \(H=40\), 4 GPUs, \(H_{per\_rank}=10\).

4pt

Even partitioning achieves average speedup of 1.044\(\times\), while uneven partitioning achieves 1.048\(\times\), with only 0.4% difference. This indicates that load imbalance effects are masked by pipelined execution, and the optimal \(C\) value need not satisfy divisibility.

Correctness and training support: HCMS output is identical to baseline across all configurations, with maximum error of 0. HCMS is compatible with PyTorch autograd, achieving 3.5% speedup for forward+backward combined execution and 8.7% reduction in peak memory; see Appendix 9 for details.

4.4 End-to-End Video Generation Validation↩︎

To validate HCMS effectiveness in real applications, we integrate it into the Wan2.2 video generation model [39] for end-to-end experiments.

Experimental setup: We use the Wan2.2 TI2V-5B model (Text-Image-to-Video 5B parameters) with 24 attention heads, 128 head dimension, 30-layer DiT [40], resolution 1280\(\times\)​704, and 4 GPUs with Ulysses sequence parallelism.

Attention layer benchmark: Table 9 shows attention layer performance for the TI2V-5B configuration with 24 heads and 4 GPUs across different sequence lengths.

End-to-end performance: Table 10 shows end-to-end performance for the complete video generation pipeline.

For attention layer performance, the 17-frame short video configuration achieves 18.3% speedup, and the 49-frame medium-length video configuration achieves 14.8% speedup. For end-to-end performance, 49-frame video per-step inference time decreases from 2.34s to 2.18s, achieving measured speedup of 6.8%. In this configuration, the optimal chunk count is \(C=6\), i.e., 24 heads partitioned into 6 chunks with 4 heads per chunk, where more chunks provide better pipeline overlap.

End-to-end benefit analysis: Attention layers achieve 14.8% speedup, but end-to-end is only 6.8%, because the video generation pipeline includes other components such as T5 text encoding and VAE decoding. For pure DiT inference scenarios, HCMS end-to-end benefits will be closer to attention layer speedup.

5 Conclusion↩︎

This paper proposes Head-Chunked Multi-Stream Pipeline (HCMS), optimizing communication-computation overlap for sequence parallel attention in medium-long sequence scenarios of 31K-100K tokens such as video generation. By partitioning attention heads into independent chunks through the Head Chunking strategy and achieving fine-grained pipelining via dual CUDA streams, efficient parallelization of communication and computation is achieved without changing computational semantics. Theoretical analysis shows HCMS benefits are positively correlated with communication ratio \(\rho\), and its use is recommended when \(\rho > 20\%\).

Experiments validate effectiveness across four GPU platforms at 2-8 GPU scales: for typical video generation sequence lengths of 31K-56K tokens, HCMS achieves 10%-17.5% speedup over the Ulysses baseline and 5%-14.5% over Ring Attention; in Wan2.2 end-to-end validation, attention layers achieve 18.3% speedup and the complete pipeline achieves 6.8% speedup. HCMS supports uneven partitioning with performance variance less than 1%, is compatible with PyTorch autograd with 3.5% training speedup, and is orthogonally composable with frameworks such as USP and LoongTrain.

6 Discussion↩︎

Limitations: HCMS benefits are limited in compute-dominated scenarios where \(\rho < 10\%\), such as 281K token long sequences achieving only 3.0% speedup. Additionally, when chunk count \(C\) is too large, Event synchronization overhead accumulates noticeably, and implementation introduces multi-stream management complexity.

Future work: This paper validates up to 8-GPU scenarios; future work will extend to 16-GPU and cross-node scenarios. Other directions include: applying HCMS to backward pass gradient communication overlap, and implementing adaptive scheduling mechanisms that dynamically select \(C\) based on runtime profiling.

7 Detailed Theoretical Derivations↩︎

This section provides complete derivations for the HCMS theoretical performance model.

7.1 Notation↩︎

For convenience of analysis, we define the notation shown in Table 11:

3pt

7.2 HCMS Time Model↩︎

Baseline execution time. In serial execution mode, total time is the sum of all stages: \[T_{base} = T_0 + T_{comm} + T_{attn}\]

HCMS execution time. In pipeline mode, \(H\) attention heads are partitioned into \(C\) chunks. Due to no data dependency between chunks (Lemma 1), computation of chunk \(i\) can execute in parallel with communication of chunk \(i+1\). Let per-chunk overhead be \(\beta\), then: \[\label{eq:hcms95time} T_{hcms}(C) = T_0 + \frac{T_{comm}}{C} + T_{attn} + C \cdot \beta\tag{1}\]

7.3 Optimal Chunk Count Derivation↩︎

Taking the derivative of Equation (1 ) with respect to \(C\) and setting it to zero yields \(C^* = \sqrt{T_{comm}/\beta}\). The second-order condition confirms this is a minimum.

7.4 Speedup Analysis↩︎

Let communication ratio \(\rho = T_{comm}/T_{total}\), the speedup is: \[\label{eq:speedup} S(C) = \frac{1}{1 - \rho \cdot \frac{C-1}{C} + \frac{C \cdot \beta}{T_{total}}}\tag{2}\]

When \(\beta \to 0\) and \(C \to \infty\), the speedup upper bound is \(S_{max} = 1/(1-\rho)\).

7.5 Theoretical Predictions vs. Experimental Validation↩︎

Table 12 shows comparison of theoretical predictions with measurements on the RTX 5090 platform. Time model prediction errors are all less than 1.4%.

3pt

7.6 Applicability Guidelines↩︎

Table 13 summarizes applicability for different communication ratios \(\rho\).

3pt

8 System Architecture↩︎

Figure 5 shows the complete system architecture of HCMS.

9 Detailed Experimental Data↩︎

9.1 Correctness and Training Validation↩︎

HCMS output is identical to baseline across all configurations, as shown in Table 14. Forward+backward combined execution achieves 3.5% speedup (Table 15), with 8.7% reduction in peak memory (Table 16).

9.2 Complete Experimental Results↩︎

Table 17 summarizes complete results for all experimental configurations.

2.5pt

References↩︎