Context Parallelism for Scalable Million-Token Inference

Context Parallelism for Scalable Million-Token Inference

Speakers: Amy Yang, Jingyi Yang, Aya Ibrahim, Xinfeng Xie, Jongsoo Park, Jianyu Huang

Conference: MLSys 2025

YouTube: https://www.youtube.com/watch?v=None

Overview

This article delves into a pivotal presentation from MLSys 2025 titled "Context Parallelism for Scalable Million-Token Inference," delivered by a team of researchers from Meta. The talk introduces Context Parallelism (CP) as an innovative and highly effective technique designed to mitigate the substantial latency associated with processing extremely long contexts—up to 1 million tokens—in large language models (LLMs) during inference. This capability is becoming increasingly critical for emerging AI applications, such as the ingestion and analysis of extensive multimodal data (e.g., images and videos), comprehensive code analysis for co-pilot systems, and processing medium to large codebases. The core challenge addressed is the prohibitive inference latency incurred by long contexts, which severely degrades user experience without advanced performance optimizations.

The primary objective of this research is to enable fast and scalable LLM inference for these demanding long-context use cases without necessitating any major architectural modifications to existing models, thereby avoiding the costly and time-consuming processes of retraining or fine-tuning. The Meta team's work, specifically benchmarked on Llama 3 405B models, showcases remarkable efficiency gains. A key contribution is the development of two distinct ring attention algorithm variants, Pass KV and Pass Q, which are tailored to optimize latency during different stages of multi-turn chat interactions involving long contexts. The efficacy of their approach is demonstrated through impressive benchmarks, where the Llama 3 405B model, processing up to 1 million tokens, achieved prefill completion in under 77 seconds with a remarkable 93% parallelization efficiency on H100-based systems. This work represents a significant step forward in making ultra-long context LLM inference practical and performant.

Background

▶ Watch: Long context demand and inference latency problem (0:00)

The demand for processing increasingly long contexts in LLMs, extending to a million tokens and beyond, stems from rapidly evolving AI applications. These include sophisticated multimodal data processing, where vast amounts of visual or auditory information need to be contextualized; advanced code analysis tools that operate on entire software repositories; and intelligent assistants that maintain extensive conversational histories. However, the inherent architectural characteristics of transformer-based LLMs present significant challenges when scaling to such immense context lengths, primarily manifesting as prohibitive latency during inference.

The latency challenge for long-context inference can be systematically broken down into three core components: compute, memory, and communication. On the compute front, the dominant factor during the prefill phase (processing the initial prompt) is the attention mechanism. Using conventional dense attention, typically implemented with techniques like Flash Attention that leverage lower-triangle block diagonal computation, the computational cost scales roughly quadratically with the growth of the context length. Beyond a certain threshold, attention compute rapidly becomes the primary bottleneck, dominating the end-to-end prefill latency. Memory consumption is another critical concern, as the KV cache (storing key and value vectors for previously processed tokens) grows linearly with context length, quickly exhausting the memory capacity of individual accelerators. Finally, communication overhead escalates with the degree of parallelization across multiple hosts, becoming a significant bottleneck, particularly in distributed inference scenarios.

To illustrate the severity of this problem, consider the empirical data presented for a Llama 3 405B model, quantized to FP8 precision, served from a single NVIDIA H100 host. At an 8K context length, the Time to First Token (TTFT)—a key metric reflecting prefill latency—is approximately 1.6 seconds. However, when the context length is scaled up 16-fold to 128K tokens, the TTFT drastically increases to 42 seconds, representing a nearly 26-fold increase in latency. In stark contrast, the Time to Incremental Token (TTIT), which measures decode latency (the time to generate subsequent tokens), shows a much less perceptible increase. For the same Llama 3 405B model, TTIT grows from roughly 30 milliseconds at 8K context to only 100 milliseconds at 128K context—a mere 3X increase. This disparity highlights that prefill latency is the dominant and most critical challenge for user experience in long-context scenarios, making it the central focus of this research.

Existing parallelization schemes, such as Tensor Parallel (TP), are effective at reducing serving latency but often suffer from poor inter-node scaling due to high all-reduce communication overhead. The message size in TP typically scales linearly with both the token length and the model dimension (number of heads multiplied by head dimension). This inherent limitation motivates the search for alternative parallelization strategies that can more effectively manage communication costs while enabling distributed computation for extremely long sequences. The need for a solution that avoids costly model architectural changes, retraining, or fine-tuning underscores the practicality and immediate applicability of the proposed Context Parallelism approach.

Key Findings

▶ Watch: Technical challenges: compute, memory, and communication (2:00)

The research introduces Context Parallelism (CP) as a highly effective and novel technique for addressing the prohibitive latency of LLM inference in long-context scenarios. The core contributions and findings can be summarized as follows:

1. Introduction of Context Parallelism (CP): CP is presented as an effective method to tackle the latency problem for LLM inference at long context lengths, specifically targeting the prefill phase. Unlike Tensor Parallelism (TP), which distributes model weights, CP distributes input tokens across multiple computational ranks, enabling parallel processing of the sequence dimension.
2. Novel Ring Attention Algorithms: Two distinct ring attention algorithm variants, Pass KV and Pass Q, were developed. These algorithms facilitate low-latency prefill and decode operations, particularly in multi-turn chat scenarios involving long contexts. The choice between Pass KV and Pass Q is dynamic, optimized based on the specific stage (full prefill, partial prefill, decode) and characteristics like KV cache hit rate.

• Pass KV is designed for full prefill and partial prefill with low to medium KV cache hit rates. It operates by keeping query (Q) tensors stationary on each CP rank while circulating key (K) and value (V) tensors (KV cache) among ranks. This strategy is efficient when KV tensors are significantly smaller than Q tensors, leading to lower inter-host communication.
• Pass Q is more effective for partial prefill with high KV cache hit rates or for decode. In this variant, KV tensors remain stationary, and Q tensors are passed around the CP group. This is advantageous when Q tensors are smaller than KV tensors, or when the communication pattern for Q tensors is more favorable for overlap.

1. Exceptional Prefill Performance: Benchmarks using the Llama 3 405B model, quantized to FP8, demonstrated remarkable prefill performance. For a 1 million token inference prefill, the system achieved completion in under 77 seconds with an impressive 93% parallelization efficiency. This benchmark was conducted on NVIDIA H100 based systems, showcasing near-linear scaling efficiency. For a 128K token prefill, the system completed it in just 3.8 seconds on 16 H100 nodes.
2. Communication-Computation Overlap: A critical finding is that for the Pass KV variant, the communication of KV tensors can be almost entirely overlapped with computation, contributing significantly to the high scaling efficiency observed during prefill. This overlap is crucial for minimizing the impact of distributed communication on overall latency.
3. Decode Latency Regression: While highly effective for prefill, the current CP algorithm introduces a regression in decode latency (TTIT). The TTIT was observed to increase with the number of CP ranks, primarily because decode operations are memory bandwidth-bound and inherently fast, making it difficult to fully overlap the communication overhead with computation.
4. Recommendation for System Decoupling: Given the decode regression, the researchers recommend decoupling prefill and decode systems. CP is best utilized for the prefill stage, while a different parallelization scheme might be more suitable for the decode stage to maintain low TTIT.
5. Architectural Agnosticism: A significant advantage of CP is its ability to be combined seamlessly with existing retrieval-based approximate methods and, crucially, it does not require any changes to the underlying model architecture, retraining, or fine-tuning. This makes it a highly practical and deployable solution for existing LLMs.

These findings collectively establish Context Parallelism as a powerful paradigm for scaling LLM inference to unprecedented context lengths, addressing a critical bottleneck in the deployment of advanced AI applications.

Technical Deep Dive

▶ Watch: Prefill (TTFT) latency growth with context length (2:50)

The core innovation of Context Parallelism lies in its strategic approach to distributing the computational workload of long sequences across multiple nodes, fundamentally differing from traditional parallelization schemes like Tensor Parallelism (TP). The challenges, as previously outlined, stem from the quadratic scaling of attention compute, linear growth of KV cache memory, and increasing communication overhead with parallelization.

Context Parallelism (CP) vs. Tensor Parallelism (TP)

Tensor Parallelism (TP) is a widely used technique where model weights are sharded across devices. While effective for reducing the memory footprint and computation for large models, TP suffers from poor inter-node scaling efficiency, particularly due to the high overhead of all-reduce operations. The message size exchanged between hosts in TP roughly scales linearly with the token_length and the model_dimension (number of heads multiplied by the dimension of each head). This makes it less ideal for scenarios where the context length becomes the dominant factor.

Context Parallelism (CP), on the other hand, distributes the input tokens (i.e., the sequence dimension) across different computational ranks or hosts. This allows for parallel processing of different segments of the input sequence. A critical design choice in CP is whether to pass Query (Q) tensors or Key/Value (KV) tensors between ranks. The inter-host message size in CP can be significantly smaller than in TP. For instance, if CP chooses to pass KV tensors, the message size could scale by number_of_KV_heads * head_dimension. This is particularly advantageous for models like Llama 3 405B, which employs Grouped-Query Attention (GQA) or Multi-Query Attention (MQA), having significantly fewer KV heads (e.g., 8 KV heads) compared to Q heads (e.g., 128 Q heads). In this specific case, the number of KV heads is 16 times smaller than the number of Q heads, leading to a potentially much larger reduction in inter-host message size by passing KV tensors.

Attention Compute Patterns in Multi-Turn Chat

To optimize CP for real-world use cases, the researchers characterized attention compute patterns across three stages of a multi-turn chat:

1. Full Prefill: This occurs when a user submits an initial prompt. The model must process the entire prompt, computing attention for all tokens, and store the resulting KV cache. The attention compute pattern is a classic lower triangle (or block diagonal) where each token attends to itself and all preceding tokens in the sequence. This leads to the quadratic computational complexity with respect to sequence length.
2. Decode: This is the auto-regressive generation phase, where the model generates one token at a time. Each newly generated token attends to all previous tokens in the sequence (both prompt and generated). This is typically a single-row attention computation.
3. Partial Prefill: This stage occurs when a user provides a follow-up prompt. The new prompt tokens must attend to both themselves (forming a lower triangle pattern for the new prompt) and all previously cached tokens from prior turns (initial prompt and machine-generated responses).

Load Balancing in CP

Achieving high efficiency in any parallelization scheme, especially CP, hinges on effective load balancing across hosts, both in terms of compute and memory. Imbalances can lead to idle resources waiting for the slowest node, eroding overall efficiency. The proposed strategy for load-balanced sharding on the sequence dimension is as follows:

For an input sequence with N tokens to be parallelized across R CP ranks, the sequence is divided into 2 * R chunks. Each CP rank then takes two chunks. For example, with 4 input tokens and 2 CP ranks (CP0, CP1), the sequence is broken into 4 pieces (token 0, token 1, token 2, token 3). CP0 might take the first and last chunks (token 0 and token 3), while CP1 takes the second and third chunks (token 1 and token 2). This ensures that each CP rank processes an equal number of tokens, thereby achieving perfect memory load balancing. Furthermore, this specific sharding pattern also ensures perfectly balanced compute blocks across ranks (e.g., 5 compute blocks for each rank in the 4-token, 2-rank example for full prefill), preventing compute bottlenecks. The same sharding logic applies to partial prefill, ensuring balanced memory and compute even with previously cached tokens.

Ring Attention Algorithms

The core of the Context Parallelism implementation lies in two novel ring attention algorithms:

1. Pass KV Ring Attention Algorithm:

• Applicability: Most effective for full prefill and partial prefill with low to medium KV cache hit rates. This is because KV tensors are often much smaller than Q tensors, making their communication less costly.
• Mechanism: All Q tensors are kept stationary on their respective CP ranks. KV tensors are passed around in a circular fashion (a ring) between the CP ranks within a group.
• Computation & Communication Overlap: As KV tensors rotate, each CP rank performs partial attention computation between its stationary Q tensors and the incoming rotating KV tensors. A key advantage here is that the communication of KV tensors can be completely overlapped with the ongoing computation, significantly boosting scaling efficiency.
• Final Output: Once all KV tensors have circulated through the group, each rank possesses partial attention outputs. These are then combined using a merge attention operation to reconstruct the final, full attention output.

1. Pass Q Ring Attention Algorithm:

• Applicability: More effective for partial prefill with high KV cache hit rates or for decode operations. In these scenarios, passing Q tensors might result in a smaller message size or a more advantageous communication pattern.
• Mechanism: KV tensors are kept stationary on each CP rank. Instead, Q tensors are passed around the CP group in a ring.
• Computation & Communication: Each CP rank computes partial attention between its stationary KV tensors and the rotating Q tensors.
• Final Output: After Q tensors complete their rotation, each rank has partial attention outputs. To reconstruct the final output, an extra permutation step is required on each CP rank to restore the original order of partial attention outputs. Subsequently, an all-to-all communication operation is performed to gather the partial attention outputs for different Q tensors back to their original ranks. Finally, a merge attention operation combines these outputs into the full attention result.
• Decode Regression Explanation: The Pass Q algorithm's reliance on an all-to-all operation for output recovery is identified as a primary contributor to the observed regression in decode latency. Decode iterations are typically very fast and memory bandwidth-bound. The communication overhead, especially the all-to-all on the critical path, cannot be fully overlapped with the rapid computation, thus increasing the TTIT. The overall performance for decode becomes bottlenecked by this communication rather than computation.

By strategically choosing between Pass KV and Pass Q based on the specific phase of inference and the characteristics of the attention tensors, Context Parallelism provides a flexible and highly optimized solution for scaling LLM inference to unprecedented context lengths.

Experimental Setup & Results

▶ Watch: Attention compute patterns: full prefill, decode, partial prefill (5:20)

The effectiveness of Context Parallelism was rigorously evaluated using a comprehensive experimental setup, yielding significant performance improvements for long-context prefill.

Model and Quantization

The primary model used for benchmarking was the Llama 3 405B model. To optimize for memory and throughput, the model was completely quantized to FP8 precision. A crucial architectural detail of the Llama 3 405B model, which heavily influenced the design of the CP algorithms, is its attention head configuration: it features 8 KV heads and 128 Q heads. This 16-fold disparity means that KV tensors are significantly smaller than Q tensors, a characteristic that the Pass KV algorithm leverages for reduced communication overhead.

Hardware Platforms

The benchmarks were conducted on two distinct hardware platforms at Meta to assess performance under varying network conditions:

1. GTT (Grand Training Platform): This platform is equipped with a high-bandwidth, low-latency RDMA network, providing 400 gigabits per second (Gbps) per GPU. This represents an ideal, high-performance distributed computing environment.
2. GTI (Grand Inference Platform): This platform, designed for inference, uses a front-end TCP/IP interconnect with 100 Gbps per GPU. This network bandwidth is 4X lower than the GTT platform, providing a more realistic and challenging environment for evaluating communication efficiency.

Parallelization Strategy

The overall parallelization strategy combined Tensor Parallelism (TP) and Context Parallelism (CP):

• Intra-host parallelization: An 8-way Tensor Parallel (TP8) scheme was used to partition the model across GPUs within a single host.
• Inter-host parallelization: Context Parallelism was applied across multiple hosts to distribute the long input sequences.

Prefill Performance Results

The results for prefill performance demonstrated near-linear scaling efficiency with CP, particularly due to the ability to fully overlap communication with computation for the Pass KV variant.

• 128K Token Prefill: Using 16 H100 nodes, the system successfully completed a 128K token prefill in a rapid 3.8 seconds.
• 1 Million Token Prefill: Scaling up to a 1 million token prefill on H100-based systems, the inference completed in under 77 seconds. This was achieved with an impressive 93% parallelization efficiency, underscoring the effectiveness of the CP approach for extremely long contexts.

Pass KV vs. Pass Q Performance

The researchers also investigated the crossover point between the Pass KV and Pass Q algorithms based on the KV cache miss rate. It was found that Pass Q begins to outperform Pass KV when the KV cache miss rate exceeds approximately 10%. This indicates the dynamic nature of the optimal algorithm choice, depending on the characteristics of the inference workload (e.g., how frequently new KV data needs to be processed versus reusing cached data).

Decode Performance (TTIT)

Despite the significant gains in prefill latency, the current CP algorithm introduced a regression in decode latency (TTIT). The TTIT was observed to increase with the number of CP ranks. The primary reasons for this regression were identified as:

• Memory Bandwidth Bound: Decode operations are typically memory bandwidth bound, meaning each iteration is inherently very fast.
• Communication Overhead: The communication introduced by CP, particularly the all-to-all operation in the Pass Q setup, cannot be fully overlapped with these rapid decode computations. This communication becomes part of the critical path, directly adding to the latency.
• Recommendation: Based on these findings, the researchers strongly recommend decoupling prefill and decode systems. CP should be primarily utilized for the prefill stage, while alternative parallelization schemes (e.g., optimized Tensor Parallel or other memory-centric approaches) might be more suitable for the decode phase to maintain low TTIT. This ensures that the benefits of CP for prefill do not negatively impact the interactive experience during token generation.

Overall, the experimental results convincingly validate Context Parallelism as a powerful solution for scaling LLM prefill to unprecedented context lengths, while also highlighting the nuanced challenges and trade-offs involved in optimizing for both prefill and decode performance in distributed environments.

Practical Implications

▶ Watch: Context Parallelism load balancing and sharding strategy (7:00)

The introduction of Context Parallelism (CP) for scalable million-token inference carries profound practical implications for a wide range of stakeholders involved in the development, deployment, and utilization of large language models.

For practitioners and model builders, CP offers a direct and efficient pathway to leverage LLMs with extremely long contexts, unlocking new application possibilities. Use cases such as comprehensive multimodal data understanding, where vast amounts of visual or auditory information need to be processed alongside text, or sophisticated code analysis tools that can reason over entire codebases, become genuinely feasible. Crucially, CP achieves this without requiring any modifications to the underlying model architecture, meaning existing, pre-trained LLMs like Llama 3 405B can immediately benefit from these optimizations without the prohibitive costs and time associated with retraining or fine-tuning. This preserves the model's performance characteristics and avoids introducing new complexities into the model development lifecycle.

Infrastructure teams and deployers will find CP a valuable tool for optimizing their LLM serving stacks. The ability to complete a 1 million token prefill in under 77 seconds with 93% efficiency on H100 systems represents a significant leap in throughput and responsiveness for long-context requests. This enables more ambitious service level agreements (SLAs) for applications that are heavily reliant on processing large inputs. However, the findings also introduce a critical design consideration: the observed regression in decode latency (TTIT) when CP is applied to token generation. This necessitates a strategic architectural decision to decouple prefill and decode systems. Infrastructure teams might opt to deploy dedicated CP-optimized clusters for handling high-volume, long-context prefill requests, while routing decode operations to systems optimized for low-latency, memory-bandwidth-bound token generation (perhaps using a different parallelization scheme like optimized Tensor Parallel). This nuanced approach ensures that the benefits of CP for prefill do not compromise the interactive experience of auto-regressive generation.

The tradeoffs and limitations of CP are important to consider. While highly effective for prefill, its impact on decode latency means it's not a silver bullet for all inference stages. The choice between Pass KV and Pass Q ring attention algorithms also introduces a dynamic element; systems will need to intelligently switch between these based on the specific workload characteristics, such as the KV cache hit rate (e.g., Pass Q being more beneficial above ~10% miss rate). This requires sophisticated runtime management and scheduling. Furthermore, the efficiency of CP, particularly the ability to overlap communication and computation, benefits significantly from high-bandwidth, low-latency inter-node networks, as demonstrated by the superior performance on RDMA-equipped platforms compared to those with lower-bandwidth TCP/IP. Deployments on less performant networks might see reduced efficiency gains.

In summary, Context Parallelism provides a robust and practical solution for tackling one of the most pressing challenges in LLM inference: scaling to extremely long contexts. It empowers practitioners to build more capable AI applications, while guiding infrastructure teams towards optimized, decoupled system architectures that balance the demands of initial prompt processing with real-time token generation, all without altering the core LLM.

Key Takeaways

• Context Parallelism (CP) revolutionizes long-context LLM prefill: It significantly reduces latency for processing sequences up to 1 million tokens, making previously impractical applications feasible.
• Two dynamic ring attention algorithms for optimization: Pass KV (for low-medium KV cache miss rates, full/partial prefill) and Pass Q (for high KV cache hit rates, partial prefill/decode) allow for adaptive performance tuning.
• Exceptional prefill performance demonstrated: The Llama 3 405B model achieved a 1 million token prefill in under 77 seconds with 93% parallelization efficiency on H100 systems.
• Communication-computation overlap is key: For Pass KV, communication of KV tensors is largely hidden by computation, driving high efficiency.
• Decode latency regression necessitates system decoupling: CP introduces overhead for decode operations; thus, it is recommended to use CP primarily for prefill and separate systems for decode to maintain low Time to Incremental Token (TTIT).
• Architecturally agnostic and highly practical: CP integrates seamlessly with existing LLMs without requiring model retraining, fine-tuning, or architectural changes, simplifying its adoption.

About the Speaker(s)

The research presented on Context Parallelism for scalable million-token inference was a collaborative effort by Amy Yang, Jingyi Yang, Aya Ibrahim, Xinfeng Xie, Jongsoo Park, and Jianyu Huang. This team comprises colleagues from Meta, where this work was conducted, leveraging their expertise in large language models and distributed systems to address critical challenges in LLM inference at scale.