[PD] Enable all CP ranks for KVCache transfer by ShangmingCai · Pull Request #19765 · sgl-project/sglang

ShangmingCai · 2026-03-03T09:10:45Z

Motivation

Changes:

Add an env var to allow all CP ranks for KVCache transfer instead of just cp rank 0. The env var is set to False by default since it has little performance improvement for MLA models.
Remove is_mla_backend check, now we can support Qwen3 CP as well (This PR can unblock Support Qwen3 MoE context parallel #18233).

Modifications

Accuracy Tests

Benchmarking and Profiling

Checklist

Format your code according to the Format code with pre-commit.
Add unit tests according to the Run and add unit tests.
Update documentation according to Write documentations.
Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.
Follow the SGLang code style guidance.

Review Process

Ping Merge Oncalls to start the PR flow. See the PR Merge Process.
Get approvals from CODEOWNERS and other reviewers.
Trigger CI tests with comments or contact authorized users to do so.
- /tag-run-ci-label, /rerun-failed-ci, /tag-and-rerun-ci
After green CI and required approvals, ask Merge Oncalls to merge.

Signed-off-by: Shangming Cai <csmthu@gmail.com>

gemini-code-assist · 2026-03-03T09:11:10Z

Summary of Changes

Hello, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request enhances the KVCache disaggregation mechanism by allowing all Control Plane (CP) ranks to actively participate in KVCache data transfer. Previously, this responsibility was limited to CP rank 0. By introducing a new environment variable and modifying the KVCache sender logic across various disaggregation backends, the system can now distribute the KVCache transfer workload more efficiently, potentially improving performance in distributed environments.

Highlights

Distributed KVCache Transfer: Introduced an environment variable SGLANG_DISAGGREGATION_CP_ALL_RANKS_TRANSFER to enable all Control Plane (CP) ranks to participate in KVCache transfer, rather than only CP rank 0.
KVCache Page Distribution: Implemented logic to distribute KVCache pages among CP ranks, ensuring each rank processes only its assigned subset of pages during transfer.
Backend Integration: Applied the new KVCache distribution logic to Mooncake, Mori, and Nixl disaggregation backends.
Refined CP Rank Handling: Updated the is_dummy_cp_rank determination and assertions related to decode CP size to align with the new all-ranks transfer capability.

🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console.

Changelog

python/sglang/srt/disaggregation/common/conn.py
- Added enable_cp_all_ranks_transfer attribute to CommonKVSender based on a new environment variable.
- Modified the condition for is_dummy_cp_rank to incorporate the enable_cp_all_ranks_transfer flag.
- Included a debug log for CommonKVSender initialization.
- Updated assertions and target_cp_ranks logic to support all CP ranks in KVCache transfer.
python/sglang/srt/disaggregation/mooncake/conn.py
- Imported the page_indices_to_cp_rank_page_indices utility function.
- Replaced local_rank with prefill_unique_rank for consistent identification of prefill senders.
- Integrated page index filtering within the send method to ensure each CP rank handles its specific portion of KV pages.
python/sglang/srt/disaggregation/mori/conn.py
- Imported the page_indices_to_cp_rank_page_indices utility function.
- Implemented page index filtering within the send method for distributing KV pages among CP ranks.
python/sglang/srt/disaggregation/nixl/conn.py
- Imported the page_indices_to_cp_rank_page_indices utility function.
- Applied page index filtering within the send method to manage KV page distribution across CP ranks.
python/sglang/srt/disaggregation/utils.py
- Introduced a new function page_indices_to_cp_rank_page_indices to logically partition global KV page indices among multiple CP ranks.
python/sglang/srt/environ.py
- Defined a new environment variable SGLANG_DISAGGREGATION_CP_ALL_RANKS_TRANSFER with a default value of True.

Activity

The pull request was initiated by ShangmingCai to enable all Control Plane (CP) ranks for KVCache transfer, moving beyond the previous limitation of only CP rank 0.
The description outlines the motivation and modifications, indicating a focus on improving distributed KVCache handling.
Standard contribution checklists are present, suggesting the PR is in an early review stage or awaiting further checks.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	`/gemini review`	Performs a code review for the current pull request in its current state.
Pull Request Summary	`/gemini summary`	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	`/gemini help`	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

gemini-code-assist

Code Review

This pull request introduces a feature to enable all Context Parallelism (CP) ranks for KVCache transfer, controlled by a new environment variable. The changes are well-implemented and correctly enable the intended functionality. My main feedback is regarding code duplication across several backend-specific connection files (mooncake/conn.py, mori/conn.py, and nixl/conn.py). I've suggested a refactoring to improve maintainability.

Signed-off-by: Shangming Cai <csmthu@gmail.com>

ShangmingCai · 2026-03-03T09:25:09Z

/rerun-stage stage-c-test-8-gpu-h20

github-actions · 2026-03-03T09:25:28Z

✅ Triggered stage-c-test-8-gpu-h20 to run independently (skipping dependencies).

github-actions · 2026-03-03T09:25:34Z

🔗 View workflow run

ShangmingCai · 2026-03-03T09:27:02Z

/rerun-stage stage-b-test-large-2-gpu

github-actions · 2026-03-03T09:27:22Z

✅ Triggered stage-b-test-large-2-gpu to run independently (skipping dependencies).

github-actions · 2026-03-03T09:27:28Z

🔗 View workflow run

Signed-off-by: Shangming Cai <csmthu@gmail.com>

ShangmingCai · 2026-03-03T09:40:28Z

/rerun-stage stage-c-test-8-gpu-h20

github-actions · 2026-03-03T09:40:51Z

✅ Triggered stage-c-test-8-gpu-h20 to run independently (skipping dependencies).

github-actions · 2026-03-03T09:40:57Z

🔗 View workflow run

ShangmingCai · 2026-03-03T10:02:20Z

/rerun-stage stage-b-test-large-2-gpu

github-actions · 2026-03-03T10:02:42Z

✅ Triggered stage-b-test-large-2-gpu to run independently (skipping dependencies).

github-actions · 2026-03-03T10:02:48Z

🔗 View workflow run

Signed-off-by: Shangming Cai <csmthu@gmail.com>

ShangmingCai · 2026-03-03T11:11:13Z

/rerun-stage stage-c-test-8-gpu-h20

github-actions · 2026-03-03T11:11:37Z

✅ Triggered stage-c-test-8-gpu-h20 to run independently (skipping dependencies).

github-actions · 2026-03-03T11:11:43Z

🔗 View workflow run

ShangmingCai · 2026-03-03T11:32:50Z

Related CI has passed.

Signed-off-by: Shangming Cai <csmthu@gmail.com>

ShangmingCai added 14 commits March 2, 2026 19:29

[PD] Enable all CP ranks for KVCache transfer

b593401

Signed-off-by: Shangming Cai <csmthu@gmail.com>

fix

2e28b68

Signed-off-by: Shangming Cai <csmthu@gmail.com>

fix

062d049

Signed-off-by: Shangming Cai <csmthu@gmail.com>

fix

6c04cf1

Signed-off-by: Shangming Cai <csmthu@gmail.com>

fix

7a52cd2

Signed-off-by: Shangming Cai <csmthu@gmail.com>

add log

080f1f1

Signed-off-by: Shangming Cai <csmthu@gmail.com>

add debug log

0a68fad

Signed-off-by: Shangming Cai <csmthu@gmail.com>

fix

594534d

Signed-off-by: Shangming Cai <csmthu@gmail.com>

add log

3d61fe1

Signed-off-by: Shangming Cai <csmthu@gmail.com>

fix

4854b3e

Signed-off-by: Shangming Cai <csmthu@gmail.com>

fix

e3324f0

Signed-off-by: Shangming Cai <csmthu@gmail.com>

fix

e364eb2

Signed-off-by: Shangming Cai <csmthu@gmail.com>

add support for nixl

3cd7af5

Signed-off-by: Shangming Cai <csmthu@gmail.com>

add support for mori

0803cd2

Signed-off-by: Shangming Cai <csmthu@gmail.com>

ShangmingCai requested review from ByronHsu and hnyls2002 as code owners March 3, 2026 09:10

gemini-code-assist Bot reviewed Mar 3, 2026

View reviewed changes

Comment thread python/sglang/srt/disaggregation/mooncake/conn.py Outdated

optimize name

63004c6

Signed-off-by: Shangming Cai <csmthu@gmail.com>

add missing change

5b4d2ae

Signed-off-by: Shangming Cai <csmthu@gmail.com>

fix comment

bd48922

Signed-off-by: Shangming Cai <csmthu@gmail.com>

ShangmingCai merged commit facde4c into main Mar 3, 2026
82 of 90 checks passed

ShangmingCai deleted the support_cp_transfer branch March 3, 2026 11:35

whybeyoung pushed a commit to whybeyoung/sglang that referenced this pull request Mar 3, 2026

[PD] Enable all CP ranks for KVCache transfer (sgl-project#19765)

aec5ccd

Signed-off-by: Shangming Cai <csmthu@gmail.com>

Kangyan-Zhou pushed a commit to Kangyan-Zhou/sglang that referenced this pull request Mar 4, 2026

[PD] Enable all CP ranks for KVCache transfer (sgl-project#19765)

9e9d567

Signed-off-by: Shangming Cai <csmthu@gmail.com>

magicYang1573 pushed a commit to magicYang1573/sglang that referenced this pull request Mar 9, 2026

[PD] Enable all CP ranks for KVCache transfer (sgl-project#19765)

a842523

Signed-off-by: Shangming Cai <csmthu@gmail.com>

Wangzheee pushed a commit to Wangzheee/sglang that referenced this pull request Mar 21, 2026

[PD] Enable all CP ranks for KVCache transfer (sgl-project#19765)

608b56a

Signed-off-by: Shangming Cai <csmthu@gmail.com>

ShangmingCai mentioned this pull request Mar 30, 2026

[Roadmap] Prefill-Decode Disaggregation Roadmap (2026 Q2) #21703

Open

14 tasks

JustinTong0323 pushed a commit to JustinTong0323/sglang that referenced this pull request Apr 7, 2026

[PD] Enable all CP ranks for KVCache transfer (sgl-project#19765)

9eaa7dd

Signed-off-by: Shangming Cai <csmthu@gmail.com>

yhyang201 pushed a commit to yhyang201/sglang that referenced this pull request Apr 22, 2026

[PD] Enable all CP ranks for KVCache transfer (sgl-project#19765)

fa11f21

Signed-off-by: Shangming Cai <csmthu@gmail.com>

Conversation

ShangmingCai commented Mar 3, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Motivation

Modifications

Accuracy Tests

Benchmarking and Profiling

Checklist

Review Process

Uh oh!

gemini-code-assist Bot commented Mar 3, 2026

Summary of Changes

Highlights

Footnotes

Uh oh!

gemini-code-assist Bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

Uh oh!

ShangmingCai commented Mar 3, 2026

Uh oh!

github-actions Bot commented Mar 3, 2026

Uh oh!

github-actions Bot commented Mar 3, 2026

Uh oh!

ShangmingCai commented Mar 3, 2026

Uh oh!

github-actions Bot commented Mar 3, 2026

Uh oh!

github-actions Bot commented Mar 3, 2026

Uh oh!

ShangmingCai commented Mar 3, 2026

Uh oh!

github-actions Bot commented Mar 3, 2026

Uh oh!

github-actions Bot commented Mar 3, 2026

Uh oh!

ShangmingCai commented Mar 3, 2026

Uh oh!

github-actions Bot commented Mar 3, 2026

Uh oh!

github-actions Bot commented Mar 3, 2026

Uh oh!

ShangmingCai commented Mar 3, 2026

Uh oh!

github-actions Bot commented Mar 3, 2026

Uh oh!

github-actions Bot commented Mar 3, 2026

Uh oh!

ShangmingCai commented Mar 3, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

ShangmingCai commented Mar 3, 2026 •

edited

Loading