[Feature] Support Prefill Context Parallel (PCP) for GQA flashinfer

[pisceskkk]

Purpose

This PR, splited from full PR #26864, adds the supports for the Prefill Context Parallelism (PCP) with GQA flashinfer, following PR #28718. For specific implementation details, please refer to the RFC #25749.
TL;DR: PCP enhances long-sequence inference capabilities by partitioning the sequence dimension during the prefill stage.

The current implementation primarily includes the following changes:

• Modified ModelRunner.py for PCP partitioning logic for tokens;
• Modified flashinfer.py to adapt the FlashInfer backend for GQA to PCP.
• Added PrefillContextParallelMetadata shared across attention backends;
• Renamed variables and functions shared by both PCP and DCP.

Test Plan

Qwen/Qwen2.5-3B

export VLLM_ATTENTION_BACKEND='FLASHINFER'
vllm serve Qwen/Qwen2.5-3B --tensor-parallel-size 4 --decode-context-parallel-size 2 --prefill-context-parallel-size 2 --dcp-kv-cache-interleave-size 8

Test Result

gsm8k eval

tp4 17c540a

dataset	version	metric	mode	vllm-api-general-stream
gsm8kdataset	-	avg@5	gen	72.78

tp4 dcp2 interleave 8

dataset	version	metric	mode	vllm-api-general-stream
gsm8kdataset	-	avg@5	gen	72.43

tp4 pcp2 interleave 8

dataset	version	metric	mode	vllm-api-general-stream
gsm8kdataset	-	avg@5	gen	72.51

tp4 dcp2 pcp2 interleave 8

dataset	version	metric	mode	vllm-api-general-stream
gsm8kdataset	-	avg@5	gen	72.98

CC @LookAround0301 @FENP @gjc0824 @LucasWilkinson

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @pisceskkk.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[gemini-code-assist]

Code Review

This pull request introduces support for Prefill Context Parallelism (PCP) for GQA with flashinfer, which is a significant feature for enhancing long-sequence inference. The changes are extensive, touching configuration, parallel state management, attention backends, and the model runner. Overall, the implementation looks solid, but I've identified a few critical issues that need to be addressed. These include a duplicated command-line argument, a syntax error, a typo in a variable name, and incorrect tensor indexing, all of which could lead to runtime errors or prevent the code from running.

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

ℹ️ About Codex in GitHub

Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @pisceskkk.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[LucasWilkinson]

Left some comments on #28988 which I think similarly apply here

[LucasWilkinson]

Thanks for the contribution! A few more comments

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @pisceskkk.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[github-actions]

This pull request has been automatically marked as stale because it has not had any activity within 90 days. It will be automatically closed if no further activity occurs within 30 days. Leave a comment if you feel this pull request should remain open. Thank you!

[github-actions]

This pull request has been automatically closed due to inactivity. Please feel free to reopen if you intend to continue working on it. Thank you!