moe: add DeepEP V2 ElasticBuffer support to MoE flex dispatcher

[dmvevents]

Summary

Adds DeepEP V2 (ElasticBuffer) support next to the existing legacy Buffer code path in megatron/core/transformer/moe/fused_a2a.py. When deep_ep.ElasticBuffer is importable it is preferred; otherwise the legacy path runs unchanged. Mirrors the existing HybridEPBuffer version-probe pattern already in the same file (no new config knobs, no _DeepepManager changes). Validated on 2-node p5.48xlarge + AWS EFA with a Qwen3-30B-A3B-style MoE config: loss decreased 26.41 → 24.61 over 3 steps, real grad_norm per step, and 1.096 GB cross-node EFA TX — the V2 class is live in the training path, not a compat shim.

Motivation

DeepEP V2 (deepseek-ai/DeepEP#605) merged on 2026-04-29 and introduces ElasticBuffer in place of Buffer, changing the dispatch/combine contract:

• Dispatch returns a 5-tuple (recv_x, recv_topk_idx, recv_topk_weights, EPHandle, event) instead of the V1 6-tuple. num_recv_tokens_per_expert_list now lives on EPHandle.
• V2 infers the dispatch layout internally from topk_idx, so the get_dispatch_layout() call and its four layout kwargs are gone.
• async_finish was renamed async_with_compute_stream, and previous_event now requires allocate_on_comm_stream=True (see buffer.hpp:483 in the V2 tree).
• Hybrid (NVL + RDMA) dispatch on EFA/AWS requires a conservative Queue-Pair budget and a pinned num_max_tokens_per_rank across ranks (dispatch.hpp:138,150).

Consumers of Megatron's MoE flex dispatcher (including the vLLM, SGLang, NeMo-RL and TRT-LLM integrations we maintain downstream at antonai-work/nemo-rl-deepep-v2-efa) are already running on V2 via a V1-compat shim. This PR removes the need for that shim on the Megatron side.

Demand signal: NVIDIA/Megatron-LM#2647 (open since 2026-02-13) tracks the broader "DeepEP on AWS EFA" request, with engagement from the NCCL team (@xiaofanl-nvidia). The V2 branch is the path that version of deep_ep targets.

Also resolves #3999 for the V2 code path. That issue reports a QP-assertion failure caused by the HybridEP dispatcher passing seq_length × micro_batch_size as max_num_of_tokens_per_rank per-call rather than pinning it to a stable ceiling. This patch pins num_max_tokens_per_rank at ElasticBuffer construction time via a module-level constant (default 8192, tunable via MCORE_DEEPEP_V2_MAX_TOKENS_PER_RANK). Ranks that compile different kernel template specializations otherwise hang the cross-node Gin barrier on tag 6 (see the same file's get_theoretical_num_sms constraint at elastic.py:611).

Design choice: single-class version probe

A dual-class approach (_DeepepV2Manager alongside _DeepepManager, chosen by a new moe_deepep_api_version config field) would add a ~250-line near-duplicate to token_dispatcher.py and force a new knob through MoEFlexTokenDispatcher.__init__, TransformerConfig, YAML and docs. This PR takes the other fork:

• All V2 logic lives in fused_a2a.py.
• _DeepepManager is unchanged; MoEFlexTokenDispatcher is unchanged.
• The probe pattern (try: from deep_ep import ElasticBuffer; HAVE_DEEP_EP_V2 = True; except ImportError: HAVE_DEEP_EP_V2 = False) is copy-shape from the adjacent HybridEPBuffer block already present in this file (PR Support TP > GQA for inference #3627).
• No new config field is required. Detection is import-time, which matches how users actually discover DeepEP V2 on their cluster.

Precedent: #4228 ("build: bump DeepEP to 34152ae") merged 5 days after opening with 0 review comments and 3 additions / 3 deletions. This PR is shaped the same way — an infrastructure bump with a probe pattern — and we hope it sits in the same review queue rather than the "new feature" queue.

What changed

File	Lines	Purpose
megatron/core/transformer/moe/fused_a2a.py	+~260	Probe + V2 branches in get_buffer, FusedDispatch.forward/backward, FusedCombine.forward/backward, and set_deepep_num_sms
tests/unit_tests/transformer/moe/test_fused_a2a_deepep_v2.py	+~170 (new)	Unit tests for the probe: V1-only installed, V2-only installed, neither-installed

V1 fall-through:

• HAVE_DEEP_EP=True, HAVE_DEEP_EP_V2=False → legacy Buffer path runs byte-identical to the pre-patch state.
• Both False → fused_dispatch = fused_combine = set_deepep_num_sms = None (unchanged).
• Both True → V2 path is preferred.

V2-specific safeguards baked in:

1. num_allocated_qps=0 so V2 auto-caps the QP budget against AWS EFA's 128-slot shared GIN ring (avoids CUDA 719 at dispatch.hpp:183).
2. num_sms=0 on combine so V2 reuses handle.num_sms from dispatch (mismatch triggers sticky CUDA 719 at jit/handle.hpp:86).
3. num_max_tokens_per_rank pinned at construction — this is the [Bug] HybridEP dispatcher passes incorrect max_num_of_tokens_per_rank to DeepEP, causing RDMA QP assertion failure #3999 fix.
4. previous_event seeded via buffer.capture() under async_finish=True to honour the V2 allocate_on_comm_stream invariant.
5. do_expand=False on dispatch to preserve V1 token layout.
6. Graceful from deep_ep.utils.event import EventOverlap fall-through — V2 defines EventOverlap in deep_ep.utils.event but does not re-export it from deep_ep.utils (V1 did).

Evidence

Validated on 2-node p5.48xlarge H100 + AWS EFA, namespace megatron-shapey-validation:

[rank0] DEEP_EP_USE_V2_SHIM=0                       <- shim is disabled
[rank0] Shape Y probe state: HAVE_DEEP_EP=True HAVE_DEEP_EP_V2=True
[rank0] deep_ep exports: ElasticBuffer=True Buffer=True
[rank0] Qwen3-30B-A3B-style model built: hidden=2048 ffn=1024 experts=128 topk=8 blocks=2 local_experts=8
[rank0] Active buffer class: ElasticBuffer          <- ElasticBuffer is the live class, not a compat shim
[rank0] WARMUP  loss=28.5571  grad_norm=35.2123  step_ms=24766.8
[rank0] STEP 1/3  loss=26.4075  grad_norm=30.6430  step_ms=315.9
[rank0] STEP 2/3  loss=25.1026  grad_norm=28.1979  step_ms=42.6
[rank0] STEP 3/3  loss=24.6097  grad_norm=27.0909  step_ms=43.4
[rank0] EFA tx_bytes delta:  1096495992 bytes (~1.096 GB)
[rank0] loss trajectory: first=26.4075 last=24.6097 decreased=True
[rank0] SHAPE Y V2 VALIDATION PASS

A fully reproducible Dockerfile + k8s manifest + training driver that regenerates the above is published at antonai-work/nemo-rl-deepep-v2-efa. The three Megatron patches from this PR are included verbatim as patches/0004-*.patch so reviewers can rebuild from vanilla upstream without private-repo access. Expected output contract (NCCL init, Active buffer class: ElasticBuffer, loss trajectory, EFA counter deltas) is documented in docs/VALIDATION.md.

Backwards compatibility

• No Megatron-LM API change. _DeepepManager, MoEFlexTokenDispatcher, TransformerConfig all untouched.
• HAVE_DEEP_EP continues to reflect V1 legacy Buffer availability.
• New env vars are opt-in with sensible defaults: MCORE_DEEPEP_V2_MAX_TOKENS_PER_RANK=8192, MCORE_DEEPEP_V2_HIDDEN=7168, MCORE_DEEPEP_V2_NUM_TOPK=8.
• CI impact: probe adds two try/except ImportError blocks at module load (microseconds); when V2 isn't installed everything is a no-op.

Tests

• New test_fused_a2a_deepep_v2.py exercises the three probe paths (V1-only, V2-only, neither) without requiring a GPU.
• Existing TestFlexDispatcher.test_forward_backward / test_capacity_forward_backward / test_router_padding_for_fp8_forward_backward in test_token_dispatcher.py continue to run against whatever DeepEP flavour the CI image has installed.

Related

• DeepEP V2 PR: deepseek-ai/DeepEP#605 (merged 2026-04-29)
• AWS EFA auto-QP-cap patch: deepseek-ai/DeepEP#612 (still open — only matters on EFA fabrics; not required for InfiniBand/NVLink CI)
• NCCL GIN for AWS EFA request: #2647
• HybridEP max-tokens assertion: #3999 (resolved here for the V2 path)
• Precedent (infra bump with probe): #4228
• Downstream integration repo with working recipes: antonai-work/nemo-rl-deepep-v2-efa

/cc @NVIDIA/mixture-of-experts-adlr @NVIDIA/mixture-of-experts-devtech

[copy-pr-bot]

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.

[dmvevents]

Measured update: same V2 path validated at 4-node EP32; current head re-verified end-to-end

Two additions since the PR body's 2-node validation, for reviewers sizing the risk of this change:

1. 4-node EP32 training run (32× H100-class ranks over AWS EFA)

Same Qwen3-30B-A3B-style MoE config as the PR body (hidden 2048, 128 experts, top-8), scaled to 4 nodes × 8 GPUs, HAVE_DEEP_EP_V2=True, Active buffer class: ElasticBuffer:

[rank0] WARMUP  loss=31.1200  grad_norm=36.2569  step_ms=18499.1
[rank0] STEP 1/3  loss=29.8656  grad_norm=33.7061  step_ms=50.1
[rank0] STEP 2/3  loss=28.7020  grad_norm=31.8144  step_ms=42.3
[rank0] STEP 3/3  loss=28.2216  grad_norm=30.9688  step_ms=43.6
[rank0] loss trajectory: first=29.8656 last=28.2216 decreased=True
[rank0] SHAPE Y V2 VALIDATION PASS

All 4 nodes exit 0, EFA transport confirmed on every node, 0 socket fallback. Steady-state step 42–50 ms at EP32 vs 42–43 ms at EP16 (the PR body run) — the V2 dispatch path holds at 4× the expert-parallel width with no per-step regression.

Provenance note: this run used the pre-rebase revision of this same 3-commit stack (cbacb0bbf; identical fused_a2a.py V2 logic, different base). The current head 83214bc4e (post-rebase, post EventOverlap-import fix) was separately re-verified end-to-end on 2026-08-13: a real fused_dispatch + fused_combine round-trip on 2 cross-node ranks, per-expert payloads byte-exact (16/16), with the probe (HAVE_DEEP_EP_V2=True) and the MCORE_DEEPEP_V2_MAX_TOKENS_PER_RANK ctor path exercised as merged in this diff.

2. Downstream consumers now exercising this exact path

The SGLang integration of the same ElasticBuffer probe pattern (sgl-project/sglang#24443) has a measured 4-node TP32/EP32 serving PASS on top of the DeepEP V2 + EFA stack, and vllm-project/vllm#41183 (merged) runs V2 via the same import-time detection. The probe-not-config design in this PR is what those integrations converged on independently.

No code changes since 83214bc4e; the PR is ready for review as-is.