Abstract
We present context parallelism for long-context large language model inference, which achieves near-linear scaling for long-context prefill latency with up to 128 H100 GPUs across 16 nodes. Particularly, our method achieves 1M context prefill with Llama3 405B model in 77s (93% parallelization efficiency, 63% FLOPS utilization) and 128K context prefill in 3.8s. We develop two lossless exact ring attention variants: pass-KV and pass-Q to cover a wide range of use cases with the state-of-the-art performance: full prefill, persistent KV prefill and decode. Benchmarks on H100 GPU hosts inter-connected with RDMA and TCP both show similar scalability for long-context prefill, demonstrating that our method scales well using common commercial data center with medium-to-low inter-host bandwidth.
Abstract (translated)
我们提出了长上下文大语言模型推理的上下文并行方法，该方法在16个节点上的最多128个H100 GPU中实现了接近线性的长时间上下文预填充延迟扩展。特别是，我们的方法使用Llama3 405B模型实现1M上下文预填充仅需77秒（并行化效率为93%，FLOPS利用率为63%），并且可以在3.8秒内完成128K上下文的预填充。我们开发了两种无损精确环注意力变体：pass-KV和pass-Q，以覆盖广泛的使用场景，并达到最先进的性能：完全预填充、持久KV预填充和解码。在通过RDMA和TCP互连的H100 GPU主机上的基准测试都显示了长时间上下文预填充具有相似的可扩展性，这表明我们的方法可以很好地应用于具有中低带宽的商用数据中心环境。
URL
https://arxiv.org/abs/2411.01783
