HCMS: Head-Chunked Multi-Stream Pipeline for Communication-Computation Overlap in Long-Sequence Parallel Attention

2026-07-02

AI Summary

This work addresses the low hardware utilization in full-all-to-all sequence parallelism for medium- to long-length sequences, where strict serialization of communication and computation limits efficiency. To overcome this, the authors propose a head-blocked multi-stream pipelining mechanism that partitions multi-head attention into multiple (possibly uneven) blocks and leverages dual CUDA streams to enable fine-grained overlap between communication and computation. The approach is orthogonal to existing optimizations such as FlashAttention, requires no modification to underlying operators, and preserves numerical equivalence. Evaluated on sequences of 31K–56K tokens, the method achieves 10%–17.5% speedup over Ulysses and outperforms Ring Attention by 5%–14.5%, delivering an end-to-end acceleration of 6.8% on the Wan2.2 model.

Abstract

All-to-all based sequence parallelism methods execute communication and computation strictly in serial when processing medium-long sequences, resulting in hardware resource underutilization. This paper proposes Head-Chunked Multi-Stream Pipeline (HCMS), which exploits the computational independence of multi-head attention by partitioning attention heads into multiple chunks and achieving fine-grained communication-computation overlap through dual CUDA streams. HCMS is orthogonally compatible with existing optimizations such as FlashAttention and SDPA, requires no modification to underlying kernels, supports uneven partitioning while maintaining numerical equivalence. Experiments validate the 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\% speedup over Ring Attention; end-to-end acceleration of 6.8\% is achieved on the Wan2.2 model. Theoretical analysis shows that HCMS benefits are positively correlated with communication ratio $ρ$, and its use is recommended when $ρ>20\%$.

Problem

• sequence parallelism
• communication-computation overlap
• multi-head attention
• long-sequence processing
• hardware underutilization

Innovation

• Head-Chunked Multi-Stream
• Communication-Computation Overlap
• Sequence Parallelism
• Multi-Head Attention
• CUDA Streams