LiquidCP: Decoupling Attention Tasks from Communication and Kernel Boundaries in Context Parallelism
Hongyu Chen, Yikang Ruan, Yuhai Long, Ziyu Huang, Yiyu Luo, Liukun Yu, Yangjie Zhou, Zihan Liu 等 17 位作者
原文摘要
Context parallelism (CP) partitions long sequences across GPUs for LLM training and inference. For variable-length training and full prefill, attention-balanced placement can leave token loads uneven, while ring execution repeats overhead at each computation and communication step. In stage-serialized decode CP (DCP), long-running tasks delay dependent work and leave substantial GPU hardware idle. We propose LiquidCP, which unifies training and inference through a common model of Transfer, Compute, and Reduction tasks with tile-level dependencies, decoupling attention tasks from communication and kernel boundaries. For training and full prefill, MegaRing incorporates BuddyRing placement to balance attention and token loads, and coalesces ready KV regions during forward attention to amortize overhead. For decoding and mixed batches, MegaDCP selects history-KV splits by jointly modeling the completion of attention and its dependent reductions. A shared persistent GPU-resident engine executes both designs. Compared with the fastest baseline per configuration on eight H100 GPUs, MegaRing achieves geometric-mean core-attention speedups of $1.23\times$ forward and $1.18\times$ backward, and accelerates a Transformer layer by $1.08$--$1.16\times$. MegaDCP achieves geometric-mean core-attention speedups of $1.20\times$ for decode-only and $1.35\times$ for mixed batches. In vLLM serving with KV prefix cache, it achieves an average time-to-first-token speedup of $1.22\times$.
cs.DC