feat(moe): DeepEP v2 dispatcher, Torch 2.13 PP, and selectable EP variants

[hemildesai]

What does this PR do?

This PR upgrades the MoE stack to DeepEP v2 while retaining selectable DeepEP v1
and HybridEP builds. It adds the DeepEP v2 ElasticBuffer dispatcher, Torch 2.13
pipeline/FSDP synchronization fixes, tuned communication defaults, benchmark and
finetuning recipes, HybridEP JIT-cache reuse, and a selectable CUDA development
image.

The branch includes main through a664ea4a242883a394c812232ed46b4e426c2472.
All 26 commits are signed off.

This replaces the earlier stacked PRs #2850, #2782, #2851, #2853, #2854, and
#2855. Those PRs are closed; this is a standalone PR against main.

Changelog

• build(moe): add conflicting uv extras for DeepEP v1, DeepEP v2, and
  HybridEP (deepep_v1, deepep_v2, hybridep; moe selects v2 by default),
  with updated dependency pins and lockfiles.
• feat(moe): add a DeepEP v2 ElasticBuffer dispatcher with non-reentrant
  activation-checkpointing support and SM/QP configuration propagated through
  dispatch and backward combine.
• fix(moe): initialize the ElasticBuffer during MoE parallelization with the
  actual EP process group. Non-PP fixed and packed recipes infer capacity from
  local batch size and sequence length during distributed setup; PP uses
  microbatch size and sequence length. Unknown dynamic shapes retain a one-time
  first-dispatch fallback, and the training loop has no dispatcher-specific
  buffer lifecycle.
• perf(moe): share the dispatch-owned ElasticBuffer handle checkpoint state
  with the paired combine operation. This prevents non-reentrant checkpoint
  unpack from replaying forward combine communication.
• fix(pp): adapt pipeline execution and FSDP MoE gradient synchronization to
  Torch 2.13, including final-microbatch reduction without duplicate reduction.
• perf(moe): preserve Qwen throughput after the rebase by disabling the outer
  FSDP wrapper for PP1 Qwen3 30B.
• test/perf(moe): add DeepEP v2 dispatcher tests and Qwen3 MoE benchmark and
  finetuning recipes.
• perf(moe): reuse HybridEP JIT kernels through a stable cache bridge.
• build(docker): add a selectable EP CUDA development image with
  hardlink-deduplicated per-backend virtual environments.

Performance

DeepEP v2 pin

The final stack uses Torch 2.13.0, source-built Transformer Engine 2.17.0, and
DeepEP v2 pinned to af9a0403
(af9a0403188392824fc3057452822235873e0612), the direct parent of upstream
099d5f2. The A/B used a DeepEP-only image overlay, leaving Torch, TE, CUDA,
NCCL, NVSHMEM, and AutoModel source unchanged.

Cached dispatch-plus-combine forward/backward probes show that the pre-099d5f2
pin recovers communication bandwidth:

Exact shape	SM / QP	DeepEP dd758caf	DeepEP af9a0403	Delta
Qwen3 MoE 235B, EP32	40 / 48	90.54 GB/s	93.96 GB/s	+3.8%
DeepSeek V3, EP64	6 / 33	79.88 GB/s	86.34 GB/s	+8.1%
DeepSeek V3, EP64	10 / 33	81.82 GB/s	87.48 GB/s	+6.9%

Pinned end-to-end controls:

Model	DeepEP dd758caf	DeepEP af9a0403	Delta	Pinned W&B
Qwen3 MoE 30B, EP8	31.729% MFU	31.765% MFU	+0.036 pp	yhrdlrwb
Qwen3 MoE 235B, PP4 / EP32	27.864% MFU	27.673% MFU	-0.191 pp	nj349lea
DeepSeek V3, PP4 / EP64	28.953% MFU	28.948% MFU	-0.005 pp	qtmhe2q1

The full-model Qwen3 235B controls use 33 QPs. An explicit 48-QP run fell to
7.7-9.3% steady-state MFU even though 48 QPs helped the isolated EP group; that
setting oversubscribes the full PP4/DP2 topology. Earlier local benchmark records
labeled 48/65 QPs logged the override but did not append it to the recipe
arguments, so those benchmark runs used the 33-QP backend default.

The Torch 2.13 image source-builds TE against its packaged cuDNN 9.20 and stages
the matching uv-installed cuDNN ahead of the base-image copy. The prebuilt TE
wheel requested an unavailable cuDNN descriptor in fused-attention backward.

Dispatcher comparison

All runs use H100s and non-reentrant activation checkpointing. MFU excludes the
first two warmup iterations.

Model	Dispatcher	Avg iteration	Avg MFU	W&B
Qwen3 MoE 30B, PP1 / EP8	DeepEP v2 pinned, 20 SM	19.197 s	31.765%	yhrdlrwb
Qwen3 MoE 30B, PP1 / EP8	DeepEP v1	19.786 s	30.820%	hy0g1bjn
Qwen3 MoE 30B, PP1 / EP8	HybridEP	18.834 s	32.377%	b1q5e72l
Qwen3 MoE 235B, PP4 / EP32	DeepEP v2 pinned, 48 SM / 33 QP	17.724 s	27.673%	nj349lea
Qwen3 MoE 235B, PP4 / EP32	DeepEP v1, 20 SM	17.305 s	28.343%	l3zaloye
Qwen3 MoE 235B, PP4 / EP32	HybridEP	17.282 s	28.380%	uih3sven
DeepSeek V3, PP4 / EP64	DeepEP v2 pinned, 6 SM / 33 QP	7.148 s	28.948%	qtmhe2q1
DeepSeek V3, PP4 / EP64	DeepEP v1, 20 SM	7.092 s	29.178%	0x5kvcaz
DeepSeek V3, PP4 / EP64	HybridEP	7.149 s	28.943%	21pxafgk

The original post-rebase Qwen3 235B v2 result was 18.759 s / 26.147% at
20 SM / 33 QP. Besides needing more communication resources, dispatch and
combine independently registered the same mutable handle with non-reentrant
checkpoint saved-tensor hooks. During combine backward, the second registration
forced recomputation through forward combine, adding an extra EP communication
replay per checkpointed layer. The paired handle now reuses the dispatch autograd
context, so recomputation stops after dispatch as it does for the legacy path.

An exact EP32 probe measured 4.7950 ms for the duplicate-save checkpoint path and
3.9740 ms for the paired path, a 17.1% reduction. With the final pin, a 40-SM /
33-QP control reached 17.776 s / 27.592%
(pq5739dj);
48 SM reached 17.724 s / 27.673%. The 0.081-point difference is not material.

Convergence

Qwen3 MoE 30B

HellaSwag, 100 finetuning steps:

Dispatcher	Loss (step 0 -> 99)	Grad norm (step 0 -> 99)	Avg TPS, steps 10-99	W&B
DeepEP v2	3.0198 -> 1.8865	15.7949 -> 1.9008	24,059.75	q7oczj2r
DeepEP v1	3.0203 -> 1.8864	15.7982 -> 1.8733	24,421.85	dbfy047a
HybridEP	3.0197 -> 1.8882	15.7967 -> 1.7678	26,701.75	jp9osl8a

Qwen3 MoE 235B

HellaSwag, 100 finetuning steps, PP4 / EP32, LR 1e-5, torch_mm experts, and
GC every 10 steps:

Dispatcher	Loss (step 0 -> 99)	Grad norm (step 0 -> 99)	Avg TPS, steps 10-99	W&B
DeepEP v2, paired handle	3.0099 -> 2.8120	535.8504 -> 108.9569	10,514.00	6xho3eke
DeepEP v1	3.0100 -> 2.7815	535.3629 -> 126.4499	9,070.51	bzrtojyv
HybridEP	3.0108 -> 2.7630	535.5029 -> 108.3242	11,768.04	d85a69be

The DeepEP v2 run maintained a stable 30.45 GiB peak through step 99 while its
loss and grad-norm trajectory matched the v1 and HybridEP baselines. A longer
Tulu-3 SFT validation also completed 1000 steps with DeepEP v2, PP4/EP32, and TE
FusedAdam fp32 master weights: loss 0.77 -> approximately 0.50 and grad norm near
0.35 (1mdcfrss).

Numerical parity probes after the backward/checkpoint-replay fix report maximum
differences of 7.8e-3 for dispatcher output, 4.8e-7 for hidden gradients, and
6.1e-5 for router-probability gradients.

Validation

• Rebased onto main at a664ea4a2.
• Pinned image probe: Torch 2.13.0, TE 2.17.0, NCCL 2.30.7, NVSHMEM 3.7.2,
  DeepEP af9a0403, and ElasticBuffer import all pass.
• Final affected suite: 782 passed, 30 warnings.
• Pinned affected suite: 119 passed, 28 warnings.
• Broader affected suite: 218 passed, 45 skipped; the vision AC file passes
  10/10 in isolation.
• uv lock --check, pre-commit hooks, Ruff, YAML recipe lint, and
  git diff --check pass.
• Exact-shape EP32/EP64 probes and pinned Qwen 30B, Qwen 235B, and DeepSeek V3
  controls completed successfully.
• All 26 commits include DCO sign-off.

Before your PR is "Ready for review"

Pre checks:

• Followed contributor guidelines.
• Added or updated dispatcher, PP, recipe, and cache tests.
• DCO sign-off on all commits.

[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.

[hemildesai]

/ok to test b43f131

[hemildesai]

/ok to test 5c89be7

[hemildesai]

Rebased onto latest main (includes #3022 which fixes the claude-review policy test). /ok to test 0df6c8e

[hemildesai]

/ok to test 0df6c8e

[akoumpa]

/ok to test 0d45a92