Search papers, labs, and topics across Lattice.
This paper introduces the Head-Chunked Multi-Stream Pipeline (HCMS), a novel approach that enhances sequence parallelism in multi-head attention by partitioning attention heads into chunks, allowing for concurrent communication and computation. By leveraging dual CUDA streams, HCMS achieves significant speedups in processing long sequences, demonstrating 10%-17.5% improvements over existing methods like Ulysses and Ring Attention across multiple GPU platforms. The findings suggest that HCMS is particularly beneficial when the communication ratio exceeds 20%, thereby optimizing resource utilization in large-scale AI models.
Achieving up to 17.5% faster processing of long sequences, HCMS redefines efficiency in multi-head attention by enabling true parallelism.
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 $\rho$, and its use is recommended when $\rho>20\%$.
