[RFC]: Support Prefill Context Parallel (PCP)

[pisceskkk]

Motivation.

A follow-up to add prefill context parallel(PCP) after DCP (PR #23734) to optimize TTFT.

In DCP, the KV Cache is partitioned along the sequence dimension into several segments (interleave style), achieving context parallelism during the decode phase.

Similarly, PCP will split the entire request along the sequence dimension during the prefill phase and further partitions the KVCache along the sequence dimension.

Proposed Change.

Aside from the attention module, almost all other modules do not involve contextual dependencies and can naturally adapt to context parallelism.

In the attention module, we adopt the following strategy to implement PCP (chunked prefill is not considered for now), marked as red.
First, for the KV, we perform an AllGather op along the sequence dimension within the PCP group to obtain the complete KV values. Then, the kvcache is stored according to the slot_mapping.
For attention computation, since we have obtained the complete KV, we only need to carefully design the custom mask and perform normal attention.

During the decode phase, modules other than attention involve redundant computations (since num new tokens = 1) across PCP group.
In the attention module, we first execute the original DCP computation logic within the respective PCP group.
Then, before updating the attention output using lse, we perform an AllGather within the PCP group to obtain complete sequence information. And the subsequent steps remain unaffected.

Summary

• The creation of Prefill Context Parallel communication domains. Unlike DCP communication domains, which are subdivisions within the TP domain, PCP communication domains stand alongside DP, PP, and TP, and affect the total device count allocation.
• Modification of slot_mapping calculation.​​ Building upon the ​​virtual block​​ concept introduced by DCP, PCP will continue to utilize virtual blocks within each PCP group. Each PCP group will be responsible for storing the portion of the KV Cache that corresponds to its assigned segment of the sequence.
• Modification of parameter calculations in the ModelRunner.​​ The calculation of parameters such as sequence length and token count must be updated, as these values are now influenced and scaled by the ​​PCP size​​.
• Modification of the attention backend computation logic. The core attention calculation requires changes to handle the distributed KV Cache across PCP groups. The specific algorithmic updates are detailed in the red markings within the provided schematic diagram.

Feedback Period.

We will complete the above modifications and submit a PR within 1-2 weeks

CC List.

@youkaichao @njhill @WoosukKwon @ruisearch42 @youzhedian

Any Other Things.

Here are some related issues and PRs:
DCP: PR #23734, PR #24864, PR #25438
PCP: RFC #22693, RFC #7519

Roadmap

• Basic feature; PR [Feature] Support Prefill Context Parallel (PCP) for GQA flashinfer #26864
• Make PCP of flashinfer compatible with DCP after PR [DCP] Support Decode Context Parallel (DCP) for GQA with Flashinfer #25438 is merged;
• Make block-level interleaved KV cache storage compatible after PR [DCP] Support dcp kv_cache interleave size > 1 #26696 is merged;
• PCP support for PIECEWISE CUDAGraph; 1598b45
• PCP support for chunked-prefill and prefix caching features;

Feature works (These items will be tackled in follow-up PRs; community contributions are warmly welcomed.):

• PCP support for MLA and other backends;
• PCP support for MTP;
• PCP support for CUDAFullGraph;
• Ring-CP style attention backend algorithm, ref RFC [RFC]: Support Context Parallelism with Fully Sharded KV Cache and Ring Attention #26133.

[LucasWilkinson]

Thanks for the RFC! I appreciate the detailed diagrams!

I couple of initial reactions:

1. I think from an interface perspective the fact that tokens are sharded across (dcp*pcp) GPUs a bit confusing, I think ideally the tokens would always be shared across dcp GPUs (less confusing for the users); One option could be if we just did -tp 6 -dcp 6 -pcp 2 (as opposed to -tp 6 -dcp 3 -pcp 2) then we could in vLLM now to break down ranks like:

   2 dcp groups: [0, 1, 2, 6, 7, 8] and [3, 4, 5, 9, 10, 11]
   2 tp/pcp groups: [0, 1, 2, 3, 4, 5] and [6, 7, 8, 9, 10, 11]

using this approach I think we could just do an all-to-all within this wider DCP group instead of having to do an all-to-all (DCP) followed by an all-gather (please correct me if this is incorrect)

2. How we would we handle a multi-request batch? Or do you think we'd have to restrict this to one prefill at a time (feels like a pretty harsh restriction). If we wanted to support a multi-request batch, I imagine we would almost treat this as a form of token level data-parallelism? i.e. if we have batch with [req1: (0, 1024), req2: (1024, 7168), req3: (7168, 8192)], we'd split the batch into [req1: (0, 1024), req2_1: (1024, 4096)] and [req2_2: (4096, 7168), req3: (7168, 8192)] and then attention would have to know how to handle context parallelism for req2 (this is basically identical to how dual batch overlap, DBO, would split the batch)

[pisceskkk]

Thanks for your comments! @LucasWilkinson
For point 1, I think this approach effectively distinguishes the concepts of PCP/DCP and avoids user confusion. We will improve this part based on your suggestion.

For point 2, we will split each request independently. As in your example, requests would be split into [req1_1 (0, 512), req2_1 (512, 3584), req3_1 (3584, 4096)]; [req1_2 (4096, 4608), req2_2 (4608, 7608), req3_2 (7608, 8192)]. I will reference the more detailed WIP PR at a later time.

[pisceskkk]

For point 2, we will split each request independently. As in your example, requests would be split into [req1_1 (0, 512), req2_1 (512, 3584), req3_1 (3584, 4096)]; [req1_2 (4096, 4608), req2_2 (4608, 7608), req3_2 (7608, 8192)]. I will reference the more detailed WIP PR at a later time.

Specific details can be reviewed in the WIP PR #25852

[LucasWilkinson]

Cross posting here (#26133 (comment)); @pisceskkk it would be nice to know your opinion on the relationship between this RFC and #25749

Also cross posting #22693 (comment) which has similar ideas to #25749 when it comes to splitting up the request accounting for the causal mask

[LucasWilkinson]

For point 2, we will split each request independently. As in your example, requests would be split into [req1_1 (0, 512), req2_1 (512, 3584), req3_1 (3584, 4096)]; [req1_2 (4096, 4608), req2_2 (4608, 7608), req3_2 (7608, 8192)]. I will reference the more detailed WIP PR at a later time.

Got it; thanks! This makes sense. I do think it would be interesting to experiment with only splitting the "middle" request to save on communication but I think this could be left for future experimentation.

[pisceskkk]

Also cross posting #22693 (comment) which has similar ideas to #25749 when it comes to splitting up the request accounting for the causal mask

Actually, the authors of this RFC, @zhenwenqi2024 @FENP, and I are from the same group, meaning that this RFC can be understood as an improved version: it modifies the cp/sp concepts to avoid confusion, maintains compatibility with DCP, etc. I overlooked this relationship, so I omitted some implementation details in this RFC; this was my mistake XD

Based on the current situation, our solution has many similarities with RFC #26133, and ring-attn would be a better improvement to our solution.

[zhenwenqi2024]

For point 2, we will split each request independently. As in your example, requests would be split into [req1_1 (0, 512), req2_1 (512, 3584), req3_1 (3584, 4096)]; [req1_2 (4096, 4608), req2_2 (4608, 7608), req3_2 (7608, 8192)]. I will reference the more detailed WIP PR at a later time.

Got it; thanks! This makes sense. I do think it would be interesting to exp

For point 2, we will split each request independently. As in your example, requests would be split into [req1_1 (0, 512), req2_1 (512, 3584), req3_1 (3584, 4096)]; [req1_2 (4096, 4608), req2_2 (4608, 7608), req3_2 (7608, 8192)]. I will reference the more detailed WIP PR at a later time.

Got it; thanks! This makes sense. I do think it would be interesting to experiment with only splitting the "middle" request to save on communication but I think this could be left for future experimentation.

Hello,Do you have any suggestions for the current PCP program? And can the VLLM community accept it?

[aarondou]

@zhenwenqi2024 Thanks for driving the PCP feature. Wonder what is the current status - has the design been finalized and implementation started? do we have a rough timeline on when it can be supported in vllm?

[aarondou]

Hi @pisceskkk , sorry I think I tagged the wrong alias above.
Thanks for driving the PCP feature. Wonder what is the current status - has the design been finalized and implementation started? do we have a rough timeline on when it can be supported in vllm?

[pisceskkk]

Wonder what is the current status - has the design been finalized and implementation started? do we have a rough timeline on when it can be supported in vllm?

@aarondou
We have already begun implementing the feature and completed the adaptation for the flashinfer backend. For specific details, you can check our draft pull request. We plan to submit an official pull request today and promote its merge into vLLM.
Adaptations for other backends are still in the planning stage. Due to time and resource constraints from other priorities, we cannot guarantee a timeline for these additional adaptations. We welcome support from the community for the adaptation of this feature.