Skip to content

[PD] Add kv_cache_dtype consistency check for PD Disaggregation#19407

Merged
ShangmingCai merged 2 commits intomainfrom
notify_pd_diff_kvtype
Feb 26, 2026
Merged

[PD] Add kv_cache_dtype consistency check for PD Disaggregation#19407
ShangmingCai merged 2 commits intomainfrom
notify_pd_diff_kvtype

Conversation

@ShangmingCai
Copy link
Copy Markdown
Collaborator

@ShangmingCai ShangmingCai commented Feb 26, 2026

Motivation

A more robust solution to achieve #19398

Modifications

Accuracy Tests

Benchmarking and Profiling

Checklist

Review Process

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

@ShangmingCai
[PD] Add kv_cache_dtype consistency check for PD Disaggregation
660246e 
Signed-off-by: Shangming Cai <csmthu@gmail.com>
@gemini-code-assist
Copy link
Copy Markdown
Contributor

gemini-code-assist Bot commented Feb 26, 2026

Summary of Changes

Hello @ShangmingCai, I'm Gemini Code Assist1! 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 robustness of the PD Disaggregation system by implementing a critical consistency check for the KV cache data type. It ensures that all participating servers, specifically prefill and decode servers, operate with the same kv_cache_dtype, preventing potential runtime issues arising from mismatched configurations. The changes involve propagating this data type information through the system's communication channels and validating it at key integration points.

Highlights

  • KV Cache Dtype Consistency: Introduced a consistency check for kv_cache_dtype between prefill and decode servers in the disaggregation architecture to prevent mismatches.
  • Data Propagation: Extended the PrefillServerInfo and BootstrapServer to include kv_cache_dtype, ensuring this information is properly communicated and stored across the distributed components.
  • Robustness: Added runtime error handling to explicitly raise an error if kv_cache_dtype values do not match between servers, improving system stability and debugging.

🧠 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 kv_cache_dtype as an optional string attribute to the PrefillServerInfo class and ensured its proper type conversion.
    • Implemented a runtime check in ensure_parallel_info to verify that the kv_cache_dtype of prefill and decode servers are identical, raising an error on mismatch.
    • Included kv_cache_dtype in the information registered to the bootstrap server.
    • Added kv_cache_dtype as an attribute to the BootstrapServer class.
    • Modified _handle_route_put to extract and store kv_cache_dtype from incoming requests.
    • Updated _handle_route_get to include kv_cache_dtype when constructing PrefillServerInfo for responses.
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.

Footnotes

  1. 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. ↩

@ShangmingCai
Copy link
Copy Markdown
Collaborator Author

ShangmingCai commented Feb 26, 2026

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

@github-actions
Copy link
Copy Markdown
Contributor

github-actions Bot commented Feb 26, 2026

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

@github-actions
Copy link
Copy Markdown
Contributor

github-actions Bot commented Feb 26, 2026

🔗 View workflow run

gemini-code-assist[bot]
gemini-code-assist Bot reviewed Feb 26, 2026
View reviewed changes
Copy link
Copy Markdown
Contributor

@gemini-code-assist gemini-code-assist Bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Code Review

This pull request adds a consistency check for kv_cache_dtype in a prefill/decode disaggregated setup. It ensures that both prefill and decode servers are configured with the same KV cache data type. The changes involve propagating the kv_cache_dtype from the prefill server to the bootstrap server, and then having the decode server fetch and validate it. The implementation is straightforward and correctly addresses the need for configuration consistency. I have one suggestion to further improve the robustness of the system by adding a consistency check on the bootstrap server itself.

Comment thread python/sglang/srt/disaggregation/common/conn.py
@ShangmingCai
fix
8ca1bc4 
Signed-off-by: Shangming Cai <csmthu@gmail.com>
@ShangmingCai
Copy link
Copy Markdown
Collaborator Author

ShangmingCai commented Feb 26, 2026

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

@github-actions
Copy link
Copy Markdown
Contributor

github-actions Bot commented Feb 26, 2026

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

@github-actions
Copy link
Copy Markdown
Contributor

github-actions Bot commented Feb 26, 2026

🔗 View workflow run

@ShangmingCai ShangmingCai merged commit 27fd014 into main Feb 26, 2026
80 of 88 checks passed
@ShangmingCai ShangmingCai deleted the notify_pd_diff_kvtype branch February 26, 2026 09:16
magicYang1573 pushed a commit to magicYang1573/sglang that referenced this pull request Mar 9, 2026
@ShangmingCai
[PD] Add kv_cache_dtype consistency check for PD Disaggregation (sgl-…
340b971 
…project#19407)

Signed-off-by: Shangming Cai <csmthu@gmail.com>
Wangzheee pushed a commit to Wangzheee/sglang that referenced this pull request Mar 21, 2026
@ShangmingCai @Wangzheee
[PD] Add kv_cache_dtype consistency check for PD Disaggregation (sgl-…
18a05d8 
…project#19407)

Signed-off-by: Shangming Cai <csmthu@gmail.com>
JustinTong0323 pushed a commit to JustinTong0323/sglang that referenced this pull request Apr 7, 2026
@ShangmingCai
[PD] Add kv_cache_dtype consistency check for PD Disaggregation (sgl-…
74e2b2d 
…project#19407)

Signed-off-by: Shangming Cai <csmthu@gmail.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@ByronHsu ByronHsu Awaiting requested review from ByronHsu ByronHsu is a code owner

@hnyls2002 hnyls2002 Awaiting requested review from hnyls2002 hnyls2002 is a code owner

1 more reviewer

@gemini-code-assist gemini-code-assist[bot] gemini-code-assist[bot] left review comments

Reviewers whose approvals may not affect merge requirements

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

1 participant

@ShangmingCai