[RFC]: Support Context Pipeline Parallelism(CPP).

[pisceskkk]

Motivation.

Currently, vLLM is supporting DCP and will support PCP to enhance long-sequence inference capabilities and optimize inference efficiency. However, in scenarios with varying sequence lengths, Context Parallelism performs poorly and may even lead to overall performance degradation. Additionally, to support and optimize ultra-long sequence inputs (e.g., reaching 1M tokens or more), introducing a new partitioning dimension is necessary. For ultra-long and variable-length sequence scenarios, combining Context Pipeline Parallelism(CPP, or named Chunked Pipeline Parallelism, Sequence Pipeline Parallelism) with an SLO-based scheduling strategy can achieve better overall inference performance. For reference, see the paper: 2409.17264.
tldr: CPP is a kind of fine-grained PP with computational-load-based chunked prefill strategy.

Here, we'd like to introduce CPP first and then gradually optimize the SLO-based priority scheduling strategy.

Vanilla PP

Currently, the pipeline parallelism implemented in vLLM operates as shown in the figure below. Each PP rank device has an independent scheduler that processes and selects requests to form an InputBatch. After processing the current InputBatch, each PP rank device passes the IntermediateTensors to the subsequent PP rank device. Then, the current PP device will wait the next scheduled InputBatch.
Since the current schedulers typically operate in an FCFS mode without considering other factors, the computational load of adjacent InputBatches often varies, leading to excessive bubble overhead.

As illustrated in the figure, there are currently four known types of bubbles:

• Bubble0: caused by blocking synchronous communication.
• Bubble1: caused by synchronous schedulers (asynchronous schedulers are currently under development).
• Bubble2/3: resulting from differing computational loads between adjacent InputBatches.

CPP Phase 1

We plan to implement a chunked prefill strategy based on computational load and request-grained pipeline parallelism. This aims to reduce the proportion of Bubble2/3 illustrated above by balancing the computational load across adjacent chunks.
Specifically, the computational load depends on the current chunk size and context length. For details, refer to the PR #20808.

CPP Phase 2

We intend to package multiple chunks of the same request into a single InputBatch to enhance the flexibility of the priority scheduler. This adjustment may impact certain dictionary variables keyed by req_id.

Tests

In preliminary tests, we observed that request-level PP exhibits higher TTFT compared to vanilla PP and chunk-level PP. This suggests that while reducing parallelism granularity can improve performance, it remains sensitive to bubbles present in each iteration. Therefore, an asynchronous scheduler and asynchronous communication are critical prerequisites for achieving performance gains with CPP.

Proposed Change.

We anticipate that the modifications for this feature will primarily reside in the gpu_worker and scheduler components, along with minor adjustments such as configuration options.

The changes to the gpu_worker will mainly focus on the Worker.execute_model function, where the Pipeline Parallelism granularity will be split and results merged. For reference, please see the following diagram.

The modifications to the scheduler will include FLOPs estimation (referencing the DCPP PR #20808), request sequence splitting, and SLO-based scheduling. And we will use computional load as budget instead of token number.

Support for asynchronous isend and irecv communication, potential issues (pytorch/pytorch#122597 (comment)).

vllm/vllm/v1/worker/gpu_worker.py

Lines 555 to 560 in 552cac9

	intermediate_tensors = IntermediateTensors(
	get_pp_group().recv_tensor_dict(
	all_gather_group=get_tp_group(),
	all_gather_tensors=all_gather_tensors,
	)
	)

Feedback Period.

This solution is still in the preliminary validation phase. Although some usable code already exists, there are still significant optimizations(async_scheduler, async_send/recv_tensor_dict) required to ensure the expected benefits are achieved. Therefore, the feedback period for the RFC is expected to be relatively long.

CC List.

@WoosukKwon @robertgshaw2-redhat @njhill @ywang96 @alexm-redhat @heheda12345 @ApostaC @LucasWilkinson @youkaichao

Any Other Things.

Some related issues or PR:
#20359 #23545

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.