Refactoring Mooncake TE as a shared distributed component#17810
Refactoring Mooncake TE as a shared distributed component#17810ShangmingCai merged 10 commits intomainfrom
Conversation
Signed-off-by: Shangming Cai <csmthu@gmail.com>
ShangmingCai
requested review from
ByronHsu,
Fridge003,
Ying1123,
ch-wan,
hanming-lu,
hnyls2002,
ispobock,
merrymercy,
xiezhq-hermann and
yizhang2077
as code owners
Summary of ChangesHello @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 significantly refactors the Mooncake Transfer Engine, transitioning it from an isolated component within the disaggregation module to a shared, globally accessible entity within the distributed module. This change is designed to allow various advanced features, such as PD, HiCache, and EPD, to leverage a single, consistent instance of the transfer engine, thereby enhancing resource efficiency and simplifying its integration across the system. The modifications involve centralizing the engine's initialization and access mechanisms, ensuring that all dependent modules retrieve the same instance rather than creating their own. Highlights
🧠 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. Using Gemini Code AssistThe 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
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 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
|
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request refactors the Mooncake Transfer Engine into a shared distributed component, allowing various advanced features like PD, HiCache, and EPD to reuse a single instance. This change effectively centralizes the management and initialization of the transfer engine, which should lead to better resource utilization and a more consistent approach across modules. The introduction of init_mooncake_transfer_engine and get_mooncake_transfer_engine functions in parallel_state.py and mooncake_transfer_engine.py is a good design choice for managing this shared resource. The changes are well-aligned with the stated motivation.
python/sglang/srt/distributed/device_communicators/mooncake_transfer_engine.py
Show resolved
Hide resolved
python/sglang/srt/mem_cache/storage/mooncake_store/mooncake_store.py
Outdated
Show resolved
Hide resolved
|
/tag-and-rerun-ci |
|
CC: @mickqian |
Signed-off-by: Shangming Cai <csmthu@gmail.com>
ShangmingCai
requested review from
iforgetmyname and
ping1jing2
as code owners
Signed-off-by: Shangming Cai <csmthu@gmail.com>
Signed-off-by: Shangming Cai <csmthu@gmail.com>
|
/rerun-failed-ci |
|
The only failed test is not related: |
|
@yizhang2077 Do you have time to help me review this PR? |
…t#17810) Signed-off-by: Shangming Cai <csmthu@gmail.com>
* www/pr/ks: (265 commits) [BugFix][PD]Fix metadata_buffer_index leak when aborted in PD (sgl-project#17483) Refactoring Mooncake TE as a shared distributed component (sgl-project#17810) [ModelOPT] Support Qwen 3 Next Coder NVFP4 (sgl-project#18224) Update author information in pyproject.toml (sgl-project#18453) [Kimi-K2.5] Fix missing `quant_config` in `KimiK25` (sgl-project#18440) Add tensor parallelism support to LFM2 ShortConv layers (sgl-project#17777) [diffusion] chore: revise process title (sgl-project#18446) Fix TRT-LLM MLA backend applying k_scale to BF16 KV cache in BMM1 (sgl-project#18396) [diffusion] refactor: group component loaders under the component_loaders/ directory (sgl-project#18438) [ModelOpt] Fix broken Qwen3-235B-A22B-Instruct-2507-NVFP4 launch (sgl-project#18189) [diffusion] feat: support efficient sequence shard (sgl-project#18161) [CI] fix: notebook ci may not working (sgl-project#18417) fix: sync server_args.kv_cache_dtype when detecting FP8 KV cache (sgl-project#18394) [Fix] Fix backend selection after flashinfer version update (sgl-project#18364) [diffusion] platform: support WAN/FLUX/Qwen-Image/Qwen-Image-edit on Ascend (sgl-project#13662) fix: fix NVFP4 Kimi-K2.5 weight mapping and exclude list (sgl-project#18370) [diffusion] feat: support saving videos directly on the server to avoid the overhead of tensor transfer (sgl-project#18253) [diffusion] fix: respect dist_timeout option (sgl-project#18386) [Doc] Fix outdated `--fp4-gemm-backend` documentation (sgl-project#18350) [diffusion] fix: remove unnecessary norm_type argument from GLM-Image dits (sgl-project#18382) ...
3 participants
Motivation and Modifications
By migrating the Mooncake transfer engine from the disaggregation module to the distributed module, various advanced features such as PD, HiCache, and EPD can reuse the same transfer engine instance, instead of initializing and using different instances in different modules, which would otherwise waste resources.
TODO
Accuracy Tests
Benchmarking and Profiling
Checklist
Review Process
/tag-run-ci-label,/rerun-failed-ci,/tag-and-rerun-ci