[RFC]: Context Parallelism && Sequence Parallelism

[zhenwenqi2024]

Motivation.

As large language models (LLMs) support increasingly longer token contexts, vLLM currently employs the chunkprefill mechanism to handle long sequences: long sequences are split into individual chunks, and ring_attention is called to compute the attention results.

However, the implementation of chunkprefill has two issues:
（1）Each chunk is computed serially, so a single request fails to fully leverage the advantages of concurrency in a multi-GPU environment;
（2）Each GPU within the tensor parallelism (TP) domain stores identical KV caches (kvcache), leading to resource waste.

To address the above issues, we aim to resolve them by introducing the CP function and enhancing the SP capability.

As mentioned in the paper, https://arxiv.org/pdf/2411.01783,The advantages of Context parallel lie in:

（1）Compute parallelization: CP distributes computation across multiple GPUs in order to reduce latency, in contrast with pipeline parallelization (PP) that improves throughput but not latency.
（2）Communication message size reduction: Compared to tensor parallelism (TP), CP demands less communication bandwidth in multi-host environments, by maintaining a communication size that is orders of magnitude smaller than TP, especially for inter-node communication.
（3）KV cache distribution: Key and value (KV) embeddings grow linearly with context length. CP distributes the storage of KV embeddings across multiple GPUs, enabling larger batch sizes with the addition of more CP ranks.

for Context Parallel, the request is split into CP size, each CP group deals a part of the request, and store kvcache in its own cp group.

for Sequecne Parallel, vllm supprts it using compilation pass，what we do is to save kvcache in each sp group.

Proposed Change.

Suppose there are 4 cards, with CP=2, SP=2（This means TP=2 plus enabled sequence parallelism.）, batch size=1, seqlen=511, and block_size=128.
CP communication domains: [0, 2], [1, 3]
SP communication domains: [0, 1], [2, 3]
The splitting is performed along two dimensions: the KV cache management dimension and the computation dimension.
KV Cache Management Dimension Splitting:
This is primarily aimed at applying for the block table. The sequence is split based on SP and CP:
Card 0 is assigned 128 tokens;
Card 1 is assigned 128 tokens;
Card 2 is assigned 128 tokens;
Card 3 is assigned 127 tokens.
In the prefill phase, each card performs its own Attention computation and then invokes ring-attention to obtain the final result. After the computation is completed:
Card 0 stores the KV cache of the first 128 tokens;
Card 1 stores the KV cache of tokens 129–255;
And so on, with each card only saving the KV cache corresponding to its assigned segment of the entire sequence.
Computation Dimension Splitting:
This mainly focuses on load balancing. The sequence is evenly split into each CP group, and within each CP group, the corresponding computation is performed through SP.

for Sequence Parallel:

For Context Parallel:

The changes that need to be made are as follows:
（1）The creation of Context Parallel communication domains(It is not reflected in the figure).
（2）Sequence Parallel communication domains reuses the Tensor Parallel communication domains（It is not reflected in the figure）
（3）Modify Schedule to allocate block_tables
（4）Modify Prepare_inputs for ModelRunner
（5）Modify attention for cp/cp commution

Feedback Period.

No response

CC List.

No response

Any Other Things.

Someone submitted an RFC, but it was put on hold.:
#7519
The idea is roughly the same, and we want to complete it.

Before submitting a new issue...

• Make sure you already searched for relevant issues, and asked the chatbot living at the bottom right corner of the documentation page, which can answer lots of frequently asked questions.

[LucasWilkinson]

Do your think you can elaborate on:

The changes that need to be made are as follows:
（1）The creation of Context Parallel communication domains(It is not reflected in the figure).
（2）Sequence Parallel communication domains reuses the Tensor Parallel communication domains（It is not reflected in the figure）
（3）Modify Schedule to allocate block_tables
（4）Modify Prepare_inputs for ModelRunner
（5）Modify attention for cp/cp commution

these are pretty major systems in vLLM so it would be nice to have a more detailed description of the changes we can comment on

[zhenwenqi2024]

Do your think you can elaborate on:

The changes that need to be made are as follows:
（1）The creation of Context Parallel communication domains(It is not reflected in the figure).
（2）Sequence Parallel communication domains reuses the Tensor Parallel communication domains（It is not reflected in the figure）
（3）Modify Schedule to allocate block_tables
（4）Modify Prepare_inputs for ModelRunner
（5）Modify attention for cp/cp commution

these are pretty major systems in vLLM so it would be nice to have a more detailed description of the changes we can comment on

Thanks for your comment.
（1）SP/CP is built on the basis of TP. original tp_size = cp_size*sp_size, assume original tp_size=8, if cp_size=2, then sp_size=4, new tp_size=4。The CP communication domain needs to be created，The SP communication shares with Tp communication domain

（2）For scheduler, the core is to apply for the block_table of each request.As shown in the figure,We need to calculate the number of tokens allocated to each cp_sp group and then apply for the corresponding block_table.

（3）The request is distributed across multiple cards, and the input needs to be updated based on the number of tokens allocated to each card.eg positions, slot_maoping and logits_indices(because of padding for cp/sp) etc.

（4）CP/SP primarily affects the computation of attention.for prefill, The biggest change is that each card no longer needs to store the full amount of KV cache, but instead stores a portion of it.for decode,The biggest change is that the required KV cache is distributed across various cards, and a ring communication mechanism is needed to obtain the final result.

for prefill，after key and value is computed,do all_gather to get all kvcache,and then save kvcache in corresponding slots.

for decode, we have full query, and kvcache is stored in each cards, we should do alltoall to get full attention result。

[roywei]

cc @luccafong @frank-wei @wushidonguc

[vasushyam]

Please consider using our algorithm from https://arxiv.org/abs/2408.04093 for context parallel decoding: https://github.com/Zyphra/tree_attention/blob/main/tree_shard_test.py (jax) https://gist.github.com/lucidrains/fe5d6c2c896bf1f935acf63d45339916 (torch)

[youkaichao]

Note: the "Sequence Parallelism" mentioned in this RFC, is essentially the PR #23734 which we call decode context parallel.

[FENP]

Hi, everyone! As PR #23734 has introduced Decode Context Parallelism, we present a detailed design proposal for Prefill Context Parallelism here. Welcome to discuss.

co-authored-by @FENP @zhenwenqi2024

Motivation

We plan to support context parallelism for long-context LLM inference to reduce the prefill latency.

1. Blockwise Parallel Transformer for Large Context Models
2. Ring Attention with Blockwise Transformers for Near-Infinite Context
3. Context Parallelism for Scalable Million-Token Inference

Background

Context Parallelism (“CP”) is a parallelization scheme on the dimension of sequence length. By sharding the long input sequence across multiple devices, each device works on a different set of tokens and communicates with each other to compute the final attention score.
Ring Attention is a high-performance attention implementation in distributed manner to support CP. (Detailed workflow ref: https://coconut-mode.com/posts/ring-attention/).

Proposed Change

Args

--context-parallel-size (default value is 1)

CP Group

world_size_across_dp = TPxCPxPPxDP
world_size = TPxCPxPP
Assume there is a node with 8 GPUs, dp_size=1, pp_size=2, cp_size=2, tp_size=2. These 8 GPUs are grouped as shown below.
● 4 tensor model-parallel groups: [g0, g1], [g2, g3], [g4, g5], [g6, g7]
● 4 pipeline model-parallel groups: [g0, g4], [g1, g5], [g2, g6], [g3, g7]
● 4 context parallel groups: [g0, g2], [g1, g3], [g4, g6], [g5, g7]

Prefill

The key issue is how we split the sequence. To ensure load balancing for causal attention, we cannot simply split the sequence into each CP rank in token order.

Solution 1

The sequence partitioning method is different between the KV Cache storage stage and the forward process. This is a flexible solution, but it requires managing the block pool for each CP rank.
V1/core：Split Sequence and Alloc KV Cache Block for Each CP rank
To ensure the correctness of hash computation for prefix cache, we must guarantee that the tokens within a block are contiguous.

V1/Model Runner：prepare inputs & forward
For model forward, we can split the sequence using the DualChunkSwap strategy to ensure load balancing, while storing KV into the KV Cache according to the KV block allocation scheme.

Solution 2

The sequence partitioning method is exactly the same between the KV Cache storage stage and the forward process, using a striped partitioning strategy.
V1/core：Split Sequence and Alloc KV Cache Block for Each CP worker

Benefits of setting global_block_size to cp_block_size × cp_size:

1. Minimal modifications required for the scheduler and KV Cache Manager. The global block ID and CP block ID are mapped one-to-one, without maintaining additional per CP rank block pool.
2. The number of tokens in the prefix cache on each rank is guaranteed to be identical, avoiding any imbalance in the prefix cache distribution.

V1/Model Runner：prepare inputs & forward

V1/Attention Backend

Modify the attention backend to support prefill with CP. The attention computation needs to be completed through multiple iterative rounds, with exchanging the KV of each rank via the CP Group.
The calculation of cp is universal, whether it is for MLA, GQA, or MHA.
In the context of sequence processing, sequences are partitioned within Context Processing (CP) groups, and then each CP rank is assigned a segment of the sequence. However, a core challenge arises: Attention computation inherently requires full-sequence information (as dependencies between tokens are often global, not limited to segment boundaries), which conflicts with the segment-wise processing of CP ranks. To resolve this, we can draw inspiration from the design philosophy of FlashAttention (FA) 。
Step 1: Preprocess and Distribute Segments to CP Groups
Step 2: Tiled Attention Computation (FA-Style Local Processing)
Step 3: Cross-Segment Update and Final Output Fusion（FA_Update）

Additionally, the key to CP lies in achieving load balancing.To achieve load balancing. Take DualChunkSwap strategy as an example, we split the sequence into 2*cp_size parts instead of splitting it into cp_size parts. To split a sequence into 2×cp_size parts and achieve load balancing, the following steps can be implemented:

1. Determine Core Parameters and Objectives
   ● Parameter Definition: Let cp_size be the number of Context Parallel (CP) groups involved in computation, denoted as G; the total sequence length is L; the total number of sub-segments to be split is M=2×G, meaning each CP group will eventually be assigned 2 sub-segments.
   ● Load Balancing Objectives: Ensure that the deviation of the total number of tokens processed by each CP group does not exceed 5%; the deviation of the computational load (such as dependency calculation amount, memory access frequency, etc.) of each CP group does not exceed 10% ~ 15%.
2. Sequence Preprocessing
   ● Length Padding: If the total sequence length L is not divisible by M, calculate the number of tokens to be padded pad_len=M−(L%M) (if L%M=0, then pad_len=0). Select padding tokens matching the task (such as  for NLP tasks, 0 values or mean values for time - series tasks), and pad pad_len tokens at the end or evenly at the beginning and end of the sequence to make the total length L′=L+pad_len, satisfying L′%M=0.
3. Sub-Segment Splitting
   Based on the preprocessed sequence S′ (with length L′), split it into M=2×G equal - length sub-segments in the order of consecutive tokens, denoted as s0​,s1​,s2​,…,sM−1​. The splitting rule is: the i - th sub-segment si​ contains tokens in the range [i×l,(i+1)×l−1] (index starts from 0), where the basic length of the sub-segment l=L′/M.
4. Sub-Segment Pairing and CP Group Assignment
   Adopt a "symmetric pairing" strategy to pair "head sub-segments of the sequence" with "tail sub-segments of the sequence", so that each CP group undertakes a mixed load of "high - dependency + low - dependency". The specific rule is: for the k-th CP group (k∈[0,G−1]), assign a pair of symmetric sub-segments (sk​,sM−1−k​). In this way, using the characteristic that "head and tail sub-segments of the sequence have strong dependencies, and middle sub-segments have weak dependencies", the difference in calculation amount between groups is flattened.
5. Sub-Segment Merging and Local Preprocessing within CP Groups
   ● Sub-Segment Merging: Concatenate the two sub-segments sk​ and sM−1−k​ of CP group k in the original sequence order to generate a complete sequence Sk​ within the group.
   ● Local K/V Calculation: Each CP group k independently performs linear projection on Sk​ to generate a local Key matrix Kk​ and a local Value matrix Vk​ within the group. Using the parallelism of CP groups, each group calculates simultaneously without mutual dependence, further ensuring load balancing.
6. Load Balancing Verification and Dynamic Adjustment
   ● Verification Indicators: Calculate length balance (deviation of the number of tokens between the largest and smallest groups), dependency calculation amount balance (deviation of the total number of interactions between sub-segments), and memory access balance (deviation of the number of simulated read K/V tiles).
   ● Dynamic Adjustment: If the deviation of a certain indicator exceeds the standard, methods such as sub-segment splitting and re-pairing, and adjusting the sub-segment splitting boundary can be used to reduce the load deviation to within the acceptable threshold.
7. Coordination with Attention Calculation
   Coordinate the segmentation strategy with optimized attention calculation (such as FlashAttention), split the sequence of each CP group into tiles suitable for the on-chip memory of hardware, exchange K/V tiles between groups according to a ring topology to ensure bandwidth load balancing, and finally merge the attention results of each group in the original sequence order.

Auto Prefix Cache

Each CP worker needs to compute attention score of new tokens without causal mask on the prefix cache. It gathers the prefix KV from the KVCache of all workers in the CP group and then execute the attention computation.

RoadMap

• Basic Implementation
  • CP Distributed Group support
  • Enable CP in Scheduler and KV Cache Manager
  • Enable CP in Model Runner
  • Enable CP in Attention Backend
• Compatible with other features
  • Auto Prefix Cache support
  • Chunk Prefill support
  • ...