Integrate DeepEPv2 - #4783
Open
irexyc wants to merge 8 commits into
Open
Conversation
Contributor
There was a problem hiding this comment.
Pull request overview
This PR integrates a DeepEPv2-based MoE all-to-all (“deepep”) backend into TurboMind for expert-parallel (EP) execution (including multi-node), wiring it through communicator topology, MoE FFN dispatch/combine, build system integration, and Python/CLI configuration.
Changes:
- Adds a DeepEP-backed token dispatcher and MoE A2A implementation path for EP>1, plus supporting CUDA kernels/utilities.
- Extends the fused Allreduce+RMSNorm communicator API to accept
local_token_nums_countand to support separate reduce/all-gather groups. - Exposes
moe_a2a_backendvia engine config + CLI and updates build/packaging/install steps to ship required NCCL compatibility pieces.
Reviewed changes
Copilot reviewed 31 out of 31 changed files in this pull request and generated 6 comments.
Show a summary per file
| File | Description |
|---|---|
| src/turbomind/turbomind.cc | Adds EP/MLP communicator splits for the deepep MoE A2A layout. |
| src/turbomind/python/CMakeLists.txt | Adjusts RPATH handling for the Python extension build/install paths. |
| src/turbomind/models/llama/unified_decoder.h | Threads a new local_token_nums_count parameter and stores mlp_group_. |
| src/turbomind/models/llama/unified_decoder.cc | Uses d_mlp_group and passes token-count metadata through Allreduce+RMSNorm; plumbs MoE token counts into shard selection. |
| src/turbomind/models/llama/moe_ffn_layer.h | Extends ForwardParam to carry local token counts and updates shard-input API. |
| src/turbomind/models/llama/moe_ffn_layer.cc | Adds DeepEP-backed MoE A2A implementation (dispatch/combine) and selects implementation by backend. |
| src/turbomind/models/llama/context.h | Adds h_ep_group and d_mlp_group to the model communicators bundle. |
| src/turbomind/engine/engine_config.h | Introduces moe_a2a_backend engine config field (default “auto”). |
| src/turbomind/comm/test_comm.cu | Updates tests to pass local_token_nums_count into Allreduce+RMSNorm. |
| src/turbomind/comm/nccl/nccl.cu | Updates NCCL communicator implementation for the new Allreduce+RMSNorm grouping/counting behavior. |
| src/turbomind/comm/nccl/deepep/utils.h | Adds runtime discovery helpers used by DeepEP integration. |
| src/turbomind/comm/nccl/deepep/token_dispatcher.h | Declares DeepEP token dispatcher wrapper used by the MoE A2A path. |
| src/turbomind/comm/nccl/deepep/token_dispatcher.cc | Implements DeepEP token dispatch/combine wrapper and JIT initialization. |
| src/turbomind/comm/nccl/deepep/symm_ctx.h | Adds NCCL symmetric memory context wrapper for DeepEP runtime. |
| src/turbomind/comm/nccl/deepep/patches/deepep-cpp-wrapper.patch | Patches upstream DeepEP sources for TurboMind integration needs (thread-local caches, stream API changes, etc.). |
| src/turbomind/comm/nccl/deepep/patches/apply_patch.cmake | Adds a CMake script to apply/verify the DeepEP patch during fetch. |
| src/turbomind/comm/nccl/deepep/nccl_stub.cc | Adds an NCCL symbol compatibility stub to enforce minimum NCCL runtime requirements. |
| src/turbomind/comm/nccl/deepep/moe_a2a_utils.h | Declares DeepEP MoE A2A helper APIs and CUDA kernels. |
| src/turbomind/comm/nccl/deepep/moe_a2a_utils.cu | Implements DeepEP MoE routing/mapping/shared-combine CUDA kernels. |
| src/turbomind/comm/nccl/deepep/CMakeLists.txt | Fetches/builds DeepEP, applies patch, builds token_dispatcher + NCCL stub, and checks NCCL version/headers. |
| src/turbomind/comm/nccl/CMakeLists.txt | Links the NCCL communicator with the DeepEP token dispatcher subproject. |
| src/turbomind/comm/device_comm.h | Extends the device-comm interface with local_token_nums_count. |
| src/turbomind/comm/cuda_ipc/fused_allreduce_ex.cu | Extends CUDA-IPC fused allreduce implementation to match the new signature. |
| src/turbomind/comm/cuda_ipc/cuda_ipc_comm.h | Updates CUDA-IPC communicator interface to include local_token_nums_count. |
| setup.py | Adds CUDA-version-based NCCL dependency selection and forces a NCCL minimum via setup requirements. |
| lmdeploy/turbomind/turbomind.py | Adds “auto/default/deepep” backend selection logic and passes backend into the C++ engine config. |
| lmdeploy/messages.py | Documents and validates new moe_a2a_backend config field. |
| lmdeploy/cli/utils.py | Adds --moe-a2a-backend CLI argument. |
| lmdeploy/cli/serve.py | Wires --moe-a2a-backend into server engine config. |
| lmdeploy/cli/cli.py | Adds EP and MoE A2A backend flags to the chat CLI surface. |
| CMakeLists.txt | Installs turbomind_nccl_stub alongside Python bindings when present. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Comment on lines
+1
to
+20
|
|
||
| #include "csrc/jit/no_ref.hpp" | ||
| #include "csrc/kernels/backend/symmetric.hpp" | ||
| #include "deep_ep/include/deep_ep/common/compiled.cuh" | ||
| #include "src/turbomind/comm/env.h" | ||
| #include "src/turbomind/comm/nccl/deepep/utils.h" | ||
|
|
||
| #include <nccl.h> | ||
| #include <nccl_device.h> | ||
| #include <nccl_device/core.h> | ||
|
|
||
| #include <memory> | ||
| #include <vector> | ||
|
|
||
| #define NCCLCHECK(e) \ | ||
| if (auto ec = e; ec != ncclSuccess) { \ | ||
| auto msg = fmt::format("NCCL error {}:{} '{}'", __FILE__, __LINE__, ncclGetErrorString(ec)); \ | ||
| throw std::runtime_error(msg.c_str()); \ | ||
| } | ||
|
|
Comment on lines
+5
to
+11
| #include <nccl.h> | ||
|
|
||
| #include <cstdio> | ||
| #include <dlfcn.h> | ||
| #include <filesystem> | ||
| #include <regex> | ||
| #include <string> |
| const int inner_tp = std::min(tp0, tp1); | ||
|
|
||
| TM_CHECK(tp0 % inner_tp == 0 && tp1 % inner_tp == 0); | ||
| const int inner_tp = global_n_ranks_ / local_token_nums_count; |
Comment on lines
+12
to
+13
| add_subdirectory(deepep) | ||
| target_link_libraries(nccl_comm PRIVATE token_dispatcher) |
Comment on lines
+262
to
+269
| if (p.ep_size > 1) { | ||
| c.h_ep_group = c.h_comm->Split(p.mlp_tp_rank, 0); | ||
| // For moe-args, we actually use `AllreduceResidualRMSnorm` to do the dispatch / combine | ||
| // For moe-a2a, we omit the all-gather before dispatch and omit the all-reduce after combine. | ||
| if (p.moe_a2a_backend == "deepep") { | ||
| TM_CHECK_EQ(communicator_type_, "nccl"); | ||
| c.d_mlp_group = c.d_comm->Split(p.ep_rank, 0, 0); | ||
| } |
Comment on lines
+162
to
+164
| cuda_version = None if os.name == 'nt' else get_cuda_version() | ||
| extra_deps = [] if os.name == 'nt' else get_turbomind_deps(cuda_version) | ||
| extra_setup_deps = [] if os.name == 'nt' else get_extra_setup_deps(cuda_version) |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Thanks for your contribution and we appreciate it a lot. The following instructions would make your pull request more healthy and more easily receiving feedbacks. If you do not understand some items, don't worry, just make the pull request and seek help from maintainers.
Motivation
Please describe the motivation of this PR and the goal you want to achieve through this PR.
Modification
Please briefly describe what modification is made in this PR.
BC-breaking (Optional)
Does the modification introduce changes that break the backward-compatibility of the downstream repositories?
If so, please describe how it breaks the compatibility and how the downstream projects should modify their code to keep compatibility with this PR.
Use cases (Optional)
If this PR introduces a new feature, it is better to list some use cases here, and update the documentation.
Checklist