Context Parallelism for Scalable Million-Token Inference

2024/11/04 by Amy Yang, H. J. Yang, Yang, Amy +15 · 8 citations
Computer Science · #Advanced Data Compression Techniques #Algorithms and Data Compression #Artificial Intelligence (cs.AI) #Distributed #FOS: Computer and information sciences #Machine Learning (cs.LG) #Machine Learning in Healthcare #Parallel #and Cluster Computing (cs.DC)

paper · pdf · doi:10.48550/arxiv.2411.01783

openalex publication_date 2024/11/04 · openalex created_date 2025/10/10 · openalex updated_date 2026/07/28

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.

Cited by

• Mesh-Attention: A New Communication-Efficient Distributed Attention with Improved Data Locality
• TEMP: A Memory Efficient Physical-aware Tensor Partition-Mapping Framework on Wafer-scale Chips
• Design Space Exploration of DMA based Finer-Grain Compute Communication Overlap
• Chopper: A Multi-Level GPU Characterization Tool & Derived Insights Into LLM Training Inefficiency
• Inferix: A Block-Diffusion based Next-Generation Inference Engine for World Simulation
• Optimizing Long-context LLM Serving via Fine-grained Sequence Parallelism
• PRISM: Probabilistic Runtime Insights and Scalable Performance Modeling for Large-Scale Distributed Training
• TokenLake: A Unified Segment-level Prefix Cache Pool for Fine-grained Elastic Long-Context LLM Serving

Related

• APB: Accelerating Distributed Long-Context Inference by Passing Compressed Context Blocks across GPUs
• InfiniteHiP: Extending Language Model Context Up to 3 Million Tokens on a Single GPU
• ParallelComp: Parallel Long-Context Compressor for Length Extrapolation
• HGCA: Hybrid GPU-CPU Attention for Long Context LLM Inference
• MOM: Memory-Efficient Offloaded Mini-Sequence Inference for Long Context Language Models