Skip to content

perf(moe): upgrade DeepEP to V2 ElasticBuffer - #3942

Draft
elfiegg wants to merge 1 commit into
NVIDIA-NeMo:mainfrom
elfiegg:elfieg/perf/deepep-v2
Draft

elfiegg wants to merge 1 commit into
NVIDIA-NeMo:mainfrom
elfiegg:elfieg/perf/deepep-v2

Conversation

@elfiegg

@elfiegg elfiegg commented Sep 19, 2026

Copy link
Copy Markdown

What does this PR do

Upgrades AutoModel from the legacy DeepEP Buffer API to DeepEP V2 ElasticBuffer and adds a sync-free expanded-dispatch path for grouped-matmul MoE experts.

  • Uses device-resident per-expert token counts, removing the previous item/permutation synchronization.
  • Preserves expanded layouts through combine backward and masks unused grouped-matmul rows.
  • Keeps the existing hybridep dispatcher label as a compatibility alias for V2 hierarchical mode.
  • Pins DeepEP 2.1.0 and NCCL 2.30.7.

Validation

Well tested on the Lyris GB200 cluster:

  • AutoModel integration: 1 GB200 node, 2 GPUs / 2 NCCL ranks actively used. Forward and backward numerical parity passed on both ranks.
  • Torch profiler: no aten::item, aten::_local_scalar_dense, or cudaDeviceSynchronize inside the measured DeepEP V2 workload.
  • Unit tests: 25 passed.
  • Upstream DeepEP V2 reference: 1 GB200 node, 2 GPUs / 2 ranks; sync-free direct dispatch/combine passed.
  • Ruff 0.12.12 formatting/checks and git diff --check passed.

Multi-node hierarchical/GIN mode is not yet validated, so this PR is intentionally a draft.

Changelog

  • Upgrade DeepEP to V2 ElasticBuffer.
  • Add zero-per-dispatch-CPU-sync execution for torch_mm grouped experts.
  • Add V2 compatibility for the existing hybridep label.
  • Add unit and one-node distributed profiler coverage.

Before your PR is "Ready for review"

  • Validated on Lyris GB200 hardware.
  • Forward/backward numerical parity passed.
  • Measured steady-state workload has no CPU-GPU synchronization.
  • Unit tests and Ruff checks passed.
  • Validate hierarchical/GIN mode across multiple GB200 nodes.
  • Refresh lock files through the repository workflow.

Signed-off-by: Elfie Guo <elfieg@nvidia.com>
@copy-pr-bot

copy-pr-bot Bot commented Sep 19, 2026

Copy link
Copy Markdown

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.

This branch was successfully deployed

3 active deployments
public 5879f1b3 Deployed Sep 19, 2026 by copy-pr-bot[bot] via release / finalize / notify #4701
test 5879f1b3 Deployed Sep 19, 2026 by copy-pr-bot[bot] via cicd-wait-in-queue #10870
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant