Add DeepEP v2 flex dispatcher backend by Autumn1998 · Pull Request #5153 · NVIDIA/Megatron-LM

Autumn1998 · 2026-06-04T08:15:39Z

What does this PR do ?

The DeepEP V2 support as a backend of flex dispatcher.

PR on dev: #4793

⚠️ For major changes (either in lines of code or in its impact), please make sure to first share a design doc with the team. If you're unsure what's the best way to do so, contact @NVIDIA/mcore-oncall.

Issue tracking

For PRs from open-source community contributors:

New features: a linked issue is required. Please open a feature request and reference it here before submitting the PR.
Small updates (bug fixes, minor improvements): a linked issue is recommended and will accelerate the PR review process.

Linked issue:

Contribution process

Pre-checks

I have added relevant unit tests
I have added relevant functional tests
I have added proper typing to my code Typing guidelines
I have added relevant documentation
I have run the autoformatter.sh on my PR

Code review

Feel free to message or comment @NVIDIA/mcore-oncall to help accelerate your merge into main. The less complex your PR is, the faster it will be approved and merged!

All PRs start as draft. If you open a non-draft PR, it will be automatically converted to draft.

Step 1: Mark PR as "Ready for Review"

When your PR is ready, click Ready for Review.
An oncall reviewer is auto-assigned and expert reviewers are notified based on your changes.
- Some PRs may jump straight to step 2. This is determined by .github/CODEOWNERS.

⚠️ Only mark as ready once merge-conflicts are resolved and the CI is passing.
Final Review might get declined if these requirements are not fulfilled.

Step 2: Final Review

For PRs that change megatron/core, once all expert reviewers have approved, the Final Review label is applied automatically and final reviewers are assigned.

For PRs outside megatron/core, this step is skipped.

Step 3: Approved

Once all required reviewers have approved, the Approved label is applied automatically.

Merge

Any member of mcore-engineers will be able to merge your PR.

copy-pr-bot · 2026-06-04T08:15:43Z

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

github-actions · 2026-06-04T08:15:51Z

This PR has been automatically converted to draft because all PRs must start as drafts.

When you are ready for review, click Ready for Review to begin the review process. This will:

Add the oncall reviewer (optional reviewer)
Add required review teams based on your changes

See the contribution guide for more details.

Autumn1998 · 2026-06-04T08:18:27Z

/ok to test 646cceb

Autumn1998 · 2026-06-08T07:31:16Z

/ok to test c0a73a0

zhongbozhu

Please add test coverage to this file: https://github.com/NVIDIA/Megatron-LM/blob/main/tests/unit_tests/transformer/moe/test_grouped_tensor_dispatcher_numerics.py

zhongbozhu

Please also consider updating this:

Megatron-LM/megatron/core/transformer/moe/moe_utils.py

Line 1505 in dff2f2e

    
           def _deepep_permute_pads_grouped_tensor_input(config: TransformerConfig) -> bool:

zhongbozhu · 2026-09-22T03:17:28Z

    HAVE_DEEP_EP = False

+try:
+    from deep_ep import ElasticBuffer


Is there anyway we can leave some instructions for people who want to use deepepv2 about how to update their container and its potential impact like HybridEP might need to be uninstalled?

Added the installation instructions and the note about potentially uninstalling HybridEP here in the missing-DeepEP-v2 error message
(

Megatron-LM/megatron/core/transformer/moe/token_dispatcher.py

Lines 1511 to 1518 in 4e37802

if deepepv2_dispatch is None:

raise ImportError(

"DeepEP v2 (deep_ep.ElasticBuffer) is unavailable. Install or upgrade DeepEP "

"in your container following https://github.com/deepseek-ai/DeepEP#installation "

"and ensure its CUDA, PyTorch, and NCCL requirements are met. "

"HybridEP also uses the deep_ep Python module and may need to be uninstalled "

"first; replacing it with DeepEP v2 can make the hybridep backend unavailable."

)

).

guihong-nv

The NCCL warmup and input-lifetime fixes address my earlier comments. I left one new comment about a race in backward. The uneven-token-count issue in the existing thread still needs to be handled. I couldn't run the multi-GPU tests here.

guihong-nv · 2026-09-22T03:21:01Z

+        previous_event = (
+            ctx.buffer.capture() if ctx.async_finish and ctx.allocate_on_comm_stream else None
+        )
+        grad_x, _, _, _, event = ctx.buffer.dispatch(
+            grad_output.contiguous(),


Can we make grad_output contiguous before capturing previous_event, in both backward methods? With the default async settings, a non-contiguous gradient queues its copy after the event, but DeepEP only waits for that event before reading the copied tensor. The communication stream can then read unfinished data and produce wrong gradients. Moving the input conversions before capture would close this race.

Signed-off-by: tongliu <tongliu@nvidia.com>

alpha-baby · 2026-09-23T10:35:32Z

could support deepep v2.5：deepseek-ai/DeepEP#763

fanshiqing · 2026-09-23T23:06:51Z

/claude review

fanshiqing · 2026-09-23T23:11:32Z

            raise ValueError(
                f"Invalid backend: {self.config.moe_flex_dispatcher_backend}"
-                "Please set --moe-flex-dispatcher-backend to deepep, hybridep, or ncclep"
+                "Please set --moe-flex-dispatcher-backend to deepep, deepepv2, hybridep, or ncclep"


Hi @Autumn1998 , do we have perf comparison between deepepv2 and our ncclep on any popular models?

It would be nice if we have some shared public comparison across these different dispatcher-backend. So that it's easier for us to make a choice between them

fanshiqing

LGTM!

fanshiqing · 2026-09-23T23:13:05Z

/ok to test 4e37802

claude · 2026-09-23T23:27:28Z

Light review — `Add DeepEP v2 flex dispatcher backend`

The backend wiring is consistent: every == "deepep" predicate that gates padding ownership or activation-freeing (should_free_input, build_transformer_layer_callables, _deepep_permute_pads_grouped_tensor_input) was widened to in ("deepep", "deepepv2"), and _DeepepV2Manager deliberately skips _DeepepManager.__init__ so a v2-only image without the v1 Buffer symbol still constructs. Tightening the deprecated moe_enable_deepep guard from == "hybridep" to != "deepep" is a real fix — it used to silently overwrite an explicit ncclep selection.

One blocking item and one small inaccuracy.

1. The determinism-coverage gate will fail the `linting` job

tools/check_kernel_determinism_coverage.py runs on every PR push. Both changed core files are registered kernel sources in tests/unit_tests/determinism/kernels/manifest.py:

changed source	manifest entry	required test
`megatron/core/transformer/moe/fused_a2a.py`	`moe_fused_a2a` (`kind="external-lib"`)	`tests/unit_tests/determinism/kernels/test_moe_kernels.py`
`megatron/core/transformer/moe/token_dispatcher.py`	`moe_token_dispatchers` (`kind="torch-op"`)	`tests/unit_tests/determinism/kernels/test_moe_kernels.py`

test_moe_kernels.py is not in this PR's diff, and the PR carries only complexity: medium / Approved — no determinism-exempt. So check() emits a violation per source ("kernel … changed but none of its determinism tests did").

This is also substantively warranted, not just a gate: ElasticBuffer.dispatch / .combine is a brand-new external-library reduction over the EP group, exactly the \b\w*buffer\.(low_latency_)?(dispatch|combine)\s*\( pattern the registry tracks. The existing flex-deepep-ep2 cell covers v1 only.

Adding a sibling cell in test_moe_kernels.py::test_moe_layer_replays, immediately after the existing flex-deepep-ep2 entry in its @pytest.mark.parametrize("dispatcher,ep,extra", [...]) list:

from megatron.core.transformer.moe.fused_a2a import HAVE_DEEP_EP, HAVE_DEEP_EP_V2

            pytest.param(
                "flex",
                2,
                {"moe_flex_dispatcher_backend": "deepepv2"},
                id="flex-deepepv2-ep2",
                marks=pytest.mark.skipif(not HAVE_DEEP_EP_V2, reason="DeepEP v2 not installed"),
            ),

and refreshing the two manifest notes so they no longer read v1-only:

    KernelEntry(
        name="moe_token_dispatchers",
        sources=(
            "megatron/core/transformer/moe/token_dispatcher.py",
            "megatron/core/transformer/moe/moe_layer.py",
        ),
        tests=(K + "test_moe_kernels.py",),
        kind="torch-op",
        notes="MoELayer replay through the allgather / alltoall dispatchers (EP=1, EP=2) and "
        "flex+DeepEP v1/v2 when available.",
    ),
    KernelEntry(
        name="moe_fused_a2a",
        sources=("megatron/core/transformer/moe/fused_a2a.py",),
        tests=(K + "test_moe_kernels.py",),
        kind="external-lib",
        notes="DeepEP v1 and v2 flex dispatcher cells (skipped without deep_ep / >=2 GPUs). "
        "HybridEP and NCCL-EP backends are not yet replayed bit-exactly here.",
    ),

The cell skips cleanly wherever deep_ep.ElasticBuffer is absent, matching how the v1 cell already behaves — so it satisfies the gate without making CI depend on a v2 container.

2. Stale error text now reachable via `deepepv2`

megatron/core/transformer/transformer_config.py:2055 — the predicate above it was widened to in ("deepep", "deepepv2"), but the message still names only deepep, so a deepepv2 user gets an error that does not mention their backend:

                raise ValueError(
                    "Flex token dispatcher with deepep/deepepv2 backend does not support "
                    "moe_pad_expert_input_to_capacity"
                )

Not blocking

tests/unit_tests/transformer/moe/test_token_dispatcher_capacity.py still parametrizes ["deepep", "hybridep"]. Extending it is optional here — moe_expert_capacity_factor reaches _DeepepV2Manager.setup_metadata through the inherited v1 implementation, so the masking path is shared code rather than anything new in this PR.

zhongbozhu · 2026-10-01T23:22:57Z

/claude fix

svcnvidia-nemo-ci · 2026-10-01T23:30:59Z

❌ Claude fix stopped because a workflow step failed. Inspect the run.

Connor-XY · 2026-10-01T23:48:57Z

/claude fix

svcnvidia-nemo-ci · 2026-10-01T23:58:59Z

❌ Claude fix stopped because a workflow step failed. Inspect the run.

Autumn1998 requested review from a team as code owners June 4, 2026 08:15

svcnvidia-nemo-ci marked this pull request as draft June 4, 2026 08:15

Autumn1998 changed the title ~~Tongliu/deepepv2 flex dispatcher main~~ Add DeepEP v2 flex dispatcher backend Jun 4, 2026

Autumn1998 mentioned this pull request Jun 4, 2026

[Dev] Add DeepEP v2 flex dispatcher backend #4793

Merged

5 tasks

copy-pr-bot Bot temporarily deployed to public June 4, 2026 08:19 Inactive

copy-pr-bot Bot temporarily deployed to test June 4, 2026 08:19 Inactive

copy-pr-bot Bot temporarily deployed to public June 4, 2026 08:22 Inactive

copy-pr-bot Bot temporarily deployed to public June 4, 2026 08:28 Inactive

Autumn1998 marked this pull request as ready for review June 8, 2026 07:30

copy-pr-bot Bot temporarily deployed to public June 8, 2026 07:31 Inactive

copy-pr-bot Bot temporarily deployed to test June 8, 2026 07:32 Inactive

svcnvidia-nemo-ci added the complexity: medium label Jun 8, 2026

copy-pr-bot Bot temporarily deployed to public June 8, 2026 07:35 Inactive

Autumn1998 requested a review from guihong-nv September 22, 2026 02:31

Merge branch 'main' into tongliu/deepepv2-flex-dispatcher-main

6377af2

zhongbozhu requested changes Sep 22, 2026

View reviewed changes

zhongbozhu reviewed Sep 22, 2026

View reviewed changes

guihong-nv reviewed Sep 22, 2026

View reviewed changes

Autumn1998 added 3 commits September 22, 2026 00:13

Fix DeepEP v2 backward input event ordering

e2f2392

Signed-off-by: tongliu <tongliu@nvidia.com>

Handle DeepEP v2 grouped-tensor padding

266ab43

Signed-off-by: tongliu <tongliu@nvidia.com>

Add DeepEP v2 installation guidance

4e37802

Signed-off-by: tongliu <tongliu@nvidia.com>

Autumn1998 requested a review from zhongbozhu September 22, 2026 07:20

Autumn1998 force-pushed the tongliu/deepepv2-flex-dispatcher-main branch from 64b8723 to 4e37802 Compare September 22, 2026 09:37

Autumn1998 requested a review from guihong-nv September 22, 2026 10:16

zhongbozhu approved these changes Sep 22, 2026

View reviewed changes

jaredcasper approved these changes Sep 23, 2026

View reviewed changes

ericharper approved these changes Sep 23, 2026

View reviewed changes

fanshiqing reviewed Sep 23, 2026

View reviewed changes

fanshiqing approved these changes Sep 23, 2026

View reviewed changes

svcnvidia-nemo-ci added the Approved All necessary approvals have been made label Sep 23, 2026

copy-pr-bot Bot deployed to public September 23, 2026 23:16 Active

NVSkySC assigned Autumn1998 Sep 29, 2026

	if deepepv2_dispatch is None:
	raise ImportError(
	"DeepEP v2 (deep_ep.ElasticBuffer) is unavailable. Install or upgrade DeepEP "
	"in your container following https://github.com/deepseek-ai/DeepEP#installation "
	"and ensure its CUDA, PyTorch, and NCCL requirements are met. "
	"HybridEP also uses the deep_ep Python module and may need to be uninstalled "
	"first; replacing it with DeepEP v2 can make the hybridep backend unavailable."
	)

Conversation

Autumn1998 commented Jun 4, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What does this PR do ?

Issue tracking

Contribution process

Pre-checks

Code review

Step 1: Mark PR as "Ready for Review"

Step 2: Final Review

Step 3: Approved

Merge

Uh oh!

copy-pr-bot Bot commented Jun 4, 2026

Uh oh!

github-actions Bot commented Jun 4, 2026

Uh oh!

Autumn1998 commented Jun 4, 2026

Uh oh!

Autumn1998 commented Jun 8, 2026

Uh oh!

zhongbozhu left a comment

Choose a reason for hiding this comment

Uh oh!

zhongbozhu left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

zhongbozhu Sep 22, 2026

Choose a reason for hiding this comment

Uh oh!

Autumn1998 Sep 22, 2026

Choose a reason for hiding this comment

Uh oh!

guihong-nv left a comment

Choose a reason for hiding this comment

Uh oh!

guihong-nv Sep 22, 2026

Choose a reason for hiding this comment

Uh oh!

Autumn1998 Sep 22, 2026

Choose a reason for hiding this comment

Uh oh!

alpha-baby commented Sep 23, 2026

Uh oh!

fanshiqing commented Sep 23, 2026

Uh oh!

fanshiqing Sep 23, 2026

Choose a reason for hiding this comment

Uh oh!

fanshiqing Sep 23, 2026

Choose a reason for hiding this comment

Uh oh!

fanshiqing left a comment

Choose a reason for hiding this comment

Uh oh!

fanshiqing commented Sep 23, 2026

Uh oh!

claude Bot commented Sep 23, 2026

Light review — Add DeepEP v2 flex dispatcher backend

1. The determinism-coverage gate will fail the linting job

2. Stale error text now reachable via deepepv2

Not blocking

Uh oh!

zhongbozhu commented Oct 1, 2026

Uh oh!

svcnvidia-nemo-ci commented Oct 1, 2026

Uh oh!

Connor-XY commented Oct 1, 2026

Uh oh!

svcnvidia-nemo-ci commented Oct 1, 2026

Uh oh!

This branch was successfully deployed

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

Autumn1998 commented Jun 4, 2026 •

edited

Loading

zhongbozhu left a comment •

edited

Loading

Light review — `Add DeepEP v2 flex dispatcher backend`

1. The determinism-coverage gate will fail the `linting` job

2. Stale error text now reachable via `deepepv2`