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

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