Skip to content

[PD-Disagg] Unify prefill info data transition flow, all with PrefillServerInfo#19195

Merged
hnyls2002 merged 5 commits intomainfrom
lsyin/introduce-prefill-info-dataclass
Feb 24, 2026
Merged

[PD-Disagg] Unify prefill info data transition flow, all with PrefillServerInfo#19195
hnyls2002 merged 5 commits intomainfrom
lsyin/introduce-prefill-info-dataclass

Conversation

@hnyls2002
Copy link
Copy Markdown
Collaborator

@hnyls2002 hnyls2002 commented Feb 23, 2026

Summary

  • Replace 5 separate per-bootstrap-addr Dict tables in CommonKVManager with a single PrefillServerInfo dataclass and prefill_info_table
  • Use dataclasses.asdict() / PrefillServerInfo(**data) for serialization end-to-end between bootstrap server and decode fetcher
  • Move _fetch_prefill_server_info from CommonKVReceiver to CommonKVManager
  • Fix missing cleanup of prefill_page_size_table and follow_bootstrap_room_table on node failure in mooncake/nixl

@gemini-code-assist
Copy link
Copy Markdown
Contributor

Summary of Changes

Hello @hnyls2002, 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 refactors the handling of prefill parallelization information within the disaggregation module. By consolidating disparate data points into a single, well-defined dataclass, the change significantly improves the organization and maintainability of the codebase. This unification simplifies how parallelization parameters are passed, stored, and accessed throughout the system, leading to a more robust and easier-to-understand data flow for disaggregated prefill operations.

Highlights

  • Data Structure Unification: Introduced a new PrefillParallelInfo dataclass to encapsulate all prefill parallelization-related information (attention TP size, DP size, PP size, page size, and follow_bootstrap_room flag).
  • Refactored Data Storage: Replaced multiple individual dictionaries (e.g., prefill_attn_tp_size_table, prefill_dp_size_table) with a single prefill_info_table that stores instances of the PrefillParallelInfo dataclass, improving data cohesion and simplifying access.
  • Streamlined Data Flow: Updated the logic for fetching, caching, and utilizing prefill parallel information across various components to consistently use the new PrefillParallelInfo object, reducing redundancy and potential for inconsistencies.

🧠 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
    • Imported the dataclasses module.
    • Defined the PrefillParallelInfo dataclass to centralize prefill parallelization parameters.
    • Replaced individual prefill info tables with a single prefill_info_table storing PrefillParallelInfo objects.
    • Updated ensure_parallel_info method to fetch and store PrefillParallelInfo objects.
    • Modified _fetch_prefill_parallel_info to return an Optional[PrefillParallelInfo] instead of a tuple.
    • Adjusted _handle_route_get to serialize PrefillParallelInfo objects to JSON using dataclasses.asdict.
  • python/sglang/srt/disaggregation/decode.py
    • Updated _resolve_dp_rank to retrieve prefill parallel information from the prefill_info_table using the new PrefillParallelInfo object.
  • python/sglang/srt/disaggregation/mooncake/conn.py
    • Modified heartbeat_checker to iterate over keys in prefill_info_table.
    • Updated _handle_node_failure to remove entries from prefill_info_table directly upon node failure.
  • python/sglang/srt/disaggregation/nixl/conn.py
    • Modified heartbeat_checker to iterate over keys in prefill_info_table.
    • Updated _handle_node_failure to remove entries from prefill_info_table directly upon node failure.
Activity
  • The pull request was opened with a 'WIP' (Work In Progress) status, indicating ongoing development.
  • No further human activity, such as comments or reviews, has been recorded since the initial submission.
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.

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

The pull request unifies the prefill information data transition flow by introducing the PrefillParallelInfo dataclass and replacing multiple individual tables with a single prefill_info_table. This refactoring improves code maintainability and readability. However, the current implementation of data fetching is less robust than the previous one, as it lacks explicit type enforcement and backward compatibility for JSON keys, which could lead to runtime errors during rolling updates or when encountering uninitialized values.

@hnyls2002 hnyls2002 changed the title [PD-Disagg] Unify prefill info data transition flow, all with PrefillParallelInfo [PD-Disagg] Unify prefill info data transition flow, all with PrefillServerInfo Feb 23, 2026
@hnyls2002
Copy link
Copy Markdown
Collaborator Author

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

@github-actions
Copy link
Copy Markdown
Contributor

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

@github-actions
Copy link
Copy Markdown
Contributor

🔗 View workflow run

@hnyls2002
Copy link
Copy Markdown
Collaborator Author

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

@github-actions
Copy link
Copy Markdown
Contributor

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

@github-actions
Copy link
Copy Markdown
Contributor

🔗 View workflow run

@hnyls2002 hnyls2002 merged commit d160c5b into main Feb 24, 2026
111 of 119 checks passed
@hnyls2002 hnyls2002 deleted the lsyin/introduce-prefill-info-dataclass branch February 24, 2026 00:15
zhuxinjie-nz pushed a commit to zhuxinjie-nz/sglang that referenced this pull request Feb 24, 2026
xiaobaicxy added a commit to xiaobaicxy/sglang that referenced this pull request Feb 24, 2026
…o xverse_moe

* 'xverse_moe' of https://github.com/xiaobaicxy/sglang: (275 commits)
  fix: add missing blank line after docstring in serving_transcription.py (sgl-project#19206)
  Whisper model support & `/v1/audio/transcriptions` endpoint & benchmark (sgl-project#16983)
  fix: patch docker image fixes (sgl-project#19100)
  [PD-Disagg] Unify prefill info data transition flow, all with `PrefillServerInfo` (sgl-project#19195)
  [CI] Tiny enhance the dp attention load blance benchmark (sgl-project#19194)
  add new ci user (sgl-project#19133)
  [CI] fix the teardown output of disaggregation test (sgl-project#19193)
  [PD-Disagg] Support query dp rank from bootstrap server. (sgl-project#19168)
  [Kernel Slimming] Migrate AWQ marlin repack kernel to JIT (sgl-project#18949)
  [Diffusion] Match rotary_embedding module name style (sgl-project#19179)
  [Refactor] Split rotary_embedding.py into a modular package (sgl-project#19144)
  [NPU] bump sgl-kernel-npu to 2026.02.01.post2 (sgl-project#19178)
  Use single mma warp group for short q_len in FA to optimize decoding performance (sgl-project#18985)
  Reorganize topk logic to clean up code and expose logical experts (sgl-project#16945)
  [ROCm] Use unreg path for custom all-reduce during CUDA graph capture (sgl-project#19162)
  [diffusion] feat: detect Flux2 custom VAE path from component_paths (sgl-project#19170)
  [AMD] ENV flags tuning and cleanup (sgl-project#19176)
  Fix bench_one_batch_server by moving the print statements (sgl-project#19175)
  Update rocm7.2 Dockerfile to install amdsmi for QuickReduce Initialization (sgl-project#19091)
  Revert "Refactor graph input buffers (sgl-project#18991)" (sgl-project#19173)
  ...
magicYang1573 pushed a commit to magicYang1573/sglang that referenced this pull request Mar 9, 2026
Wangzheee pushed a commit to Wangzheee/sglang that referenced this pull request Mar 21, 2026
JustinTong0323 pushed a commit to JustinTong0323/sglang that referenced this pull request Apr 7, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant