[Feature] Support Decode Context Parallel (DCP) for MLA

[youzhedian]

This PR adds Decode Context Parallel (DCP) support for MLA inference, fully compatible with chunked prefill and APC.

You can enable DCP with --decode-context-parallel-size/-dcp xxx (only support flashmla backend now), and tp_size needs to be divisible by dcp_size, because the world size does not change by dcp, it simply reuse the GPUs of TP group, and split one TP group into tp_size//dcp_size DCP groups. e.g.

with -tp 8 -dcp 8 , we use 8 GPUs
with -tp 8 -dcp 4 , we use 8 GPUs
with -tp 4 -dcp 4 -pp 2 , we use 8 GPUs

and kvcache token budget always increased by `dcp` times.

This DCP implement store kvcache with an interleave style, the kvcache for the token whose token_idx is i is always stored on the GPU whose dcp_rank equals to i % dcp_world_size:

e.g. DCP2, req with prompt_len=5, generation_len=4:
kvcache store in dcp_rank0: 0, 2, 4, 6, 8 
kvcache store in dcp_rank1: 1, 3, 5, 7,

deepseek-ai/DeepSeek-V2-Lite-Chat gsm8k eval:

- TP4PP2DCP4
|Tasks|Version|     Filter     |n-shot|  Metric   |   |Value |   |Stderr|
|-----|------:|----------------|-----:|-----------|---|-----:|---|-----:|
|gsm8k|      3|flexible-extract|     5|exact_match|↑  |0.6687|±  | 0.013|
|     |       |strict-match    |     5|exact_match|↑  |0.6657|±  | 0.013|

-  TP4PP2DCP1
|Tasks|Version|     Filter     |n-shot|  Metric   |   |Value |   |Stderr|
|-----|------:|----------------|-----:|-----------|---|-----:|---|-----:|
|gsm8k|      3|flexible-extract|     5|exact_match|↑  |0.6657|±  | 0.013|
|     |       |strict-match    |     5|exact_match|↑  |0.6611|±  | 0.013|

more info pls ref introduce Doc

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

• DCP support MLA fullgraph
• Extend cp_gather_cache to handle scaled KV-cache and supersede gather_and_maybe_dequant_cache
• DCP support triton_mla/cutlass_mla, DCP only support flashmla backend now
• KV-cache deduplication via DCP in GQA (e.g., GQA8-TP16 with DCP2)
• DCP support for MTP
• Ring-CP style prefill context parallel (PCP) implement to optimize TTFT

[mergify]

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

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 Context Parallelism (CP) support for MLA inference, which is a significant feature enhancement. The changes are extensive, touching configuration, parallel state management, scheduling, KV cache, and attention backends. The implementation seems well-thought-out, with new CUDA kernels for CP-specific operations and corresponding Python wrappers and tests. The end-to-end tests comparing CP with TP are a good validation strategy.

My review found one critical bug fix in the cuda_communicator.py file, where a reduce_scatter operation was using a potentially non-contiguous tensor, which could lead to incorrect results. The provided patch correctly fixes this issue. The rest of the implementation for context parallelism appears solid.

[youkaichao]

thanks for the great work!

as discussed, there can be two types of cp, cp for prefill (where the world size is enlarged by cp) and cp for decode (where the world size does not change by cp). if possible, let's denote the current pr as decode-context-parallel-size and dcp_size to leave room for prefill cp in the future.

[youkaichao]

@youzhedian to accelerate the review and merge (especially ci testing), maybe we can split the kernel side changes to a separate PR and get it merged first. then follow-up PRs can use pre-compiled wheels from that PR, with much faster ci testing.

[hmellor]

I've just come across this PR adding cp to non-MLA attention #23703

[youzhedian]

@youzhedian to accelerate the review and merge (especially ci testing), maybe we can split the kernel side changes to a separate PR and get it merged first. then follow-up PRs can use pre-compiled wheels from that PR, with much faster ci testing.

#23791 as suggested

[LucasWilkinson]

Cool thanks for taking this on! I think this can be done without any GPU model runner changes; I was working on a prototype but got unfortunately it got backburned for few months 😞 anyways just sharing here for an alternative solution that doesn't require as much more of the core code but potentially more susceptible to imbalance (its not fully functional yet)

#22789

[facebook-github-bot]

@wushidonguc has imported this pull request. If you are a Meta employee, you can view this in D81728831.

[gary-wjc]

Is this PR compatible with #22668 ? @youzhedian