Avoid collective token preparation for AutoEP DeepEP

[yh0903]

Summary

• Route AutoEP's EP>1 DeepEP path directly from router/utilization accounting into _deepep_route, before collective-backend token preparation.
• Avoid DeepEP-unused stable argsort, score/expert gathers, [T*K,H] routed-input expansion, score preparation, split-count all-to-all, and D2H split materialization.
• Share the output finalization tail across DeepEP, standard communication, and EP1 so output shape, shared experts, router logits, and cache clearing retain their existing contracts.

This is an opt-in backend cleanup: the default standard communication path and EP1 path are unchanged.

Compatibility:

• DeepEP continues to bypass fused weighted restore because its combine already restores and reduces token rows.
• DeepEP now bypasses split-plan construction entirely, including async split planning.
• The existing DeepEP incompatibility with autoep_non_moe compile is unchanged.
• Folded tensor parallelism and standard communication retain their existing paths.

Testing Done

• Local code review completed
• Unit tests added/updated
• Integration tests pass
• Manual testing performed

Correctness:

• Repository pre-commit hooks pass for all three changed files.
• Rebased H100 targeted suite passed after upstream AutoEP score-correction bias changes.
• CPU/mock contracts verify DeepEP does not call argsort, standard score application, or split-plan construction, while standard communication and EP1 still use their existing preparation.
• H100 cleanup OFF/ON parity covers output, loss, exact routes, input gradient, routing-score and router-parameter gradients, all expert gradients, optimizer deltas, activation checkpointing on/off, skewed routing, and an empty expert.

Performance, fixed-routing Qwen3-30B-A3B, EP16, TP1, BF16, activation checkpointing on, no profiler:

• L8 smoke: 115.01 ms -> 109.05 ms, a 5.18% improvement.
• L48 fresh ABBA + BAAB blocks: 682.86 ms -> 665.42 ms, a 2.55% improvement.
• Four paired L48 deltas were all positive: 42.50, 25.32, 9.56, and 18.56 ms.
• Paired mean delta was 23.98 ms, with 95% CI [1.82, 46.14] ms.
• Maximum loss difference was 0.00682; p95 did not regress.
• Peak allocated/reserved changed by approximately +2.2 MiB / +22 MiB, within the no-regression gate.

Profiling:

• Sparse CUDA-event observer overhead was -1.69%, within run-to-run noise.
• The early-route path emits no standard split-plan preparation.
• One steady-state lean Nsight capture was collected only for explanation; traced wall time is not used as the performance headline.

Matched DeepEP V2 context (same-allocation paired blocks; not a merge gate):

• Megatron uses Add DeepEP v2 flex dispatcher backend NVIDIA/Megatron-LM#5153 head eb688c4a... with the same DeepEP 01dc3aaa..., ElasticBuffer V2, NCCL 2.30.4, 12 SMs, 16 QPs, fixed routing, data, and checkpointing semantics.
• Two fresh allocations each discarded one full AutoEP arm and one full Megatron arm before ABBA/BAAB measurement.
• All four paired L48 deltas favored AutoEP: 38.62, 67.21, 80.91, and 62.18 ms.
• Paired mean delta was 62.23 ms, with 95% CI [34.20, 90.26] ms; pooled medians were 649.76 ms for AutoEP and 714.45 ms for Megatron, an AutoEP speedup of 9.06%.
• Measured-window loss differed by at most 0.00593.
• AutoEP used about 0.63 GiB more allocated and 2.77 GiB more reserved memory.
• The median result does not represent average elapsed time: across all 80 recorded measured steps, AutoEP averaged 766.67 ms versus 716.22 ms for Megatron, so AutoEP was about 7.0% slower by arithmetic mean.
• AutoEP had 32/80 steps over 800 ms versus 0/80 for Megatron; only 3/80 crossed one second. The repeated slow steps dominate full-run throughput and remain unexplained.

Follow-up fixed-routing instrumentation localizes those repeated slow steps:

• Two additional AutoEP arms again had fast medians (633.17 and 638.05 ms) but slower means (754.56 and 747.22 ms), with 16/40 measured steps over 800 ms.
• Slow versus fast median inflation was concentrated in forward (+281 ms and +248 ms in the two arms); backward and optimizer medians were effectively unchanged.
• A detailed follow-up arm attributed essentially all of the slow-forward increase to the DeepEP dispatch call: +382.07 ms dispatch versus +1.12 ms expert compute, -3.43 ms combine, and +3.94 ms router.
• Cross-rank inspection shows that dispatch is the synchronization surface, not yet the root cause: at every hotspot, 14-15 ranks wait about 233-357 ms, while one late-arriving rank spends only about 0.55-0.74 ms in dispatch. The late rank's preceding MoE/combine work is normally only 1.5-3.0 ms, placing most of the originating delay in the uninstrumented non-MoE forward region between MoE layers; one observed case accumulated the delay in the router call itself.
• The synchronized measured-window wallclock was 16893.88 ms, versus 16894.84 ms from summing the 20 recorded critical-step times. The step measurements therefore account for the full measured training window; the mean/median reversal is not an omitted gap between steps.
• A second detailed run split the inter-MoE region into decoder, attention, and normalization calls. The delayed call site moved between input RMSNorm, self-attention, router, and otherwise uninstrumented Python gaps on different ranks and layers. In each case one rank paused for roughly 260-435 ms, after which the remaining ranks waited in the next dispatch. This pattern is inconsistent with a specific DeepEP or transformer kernel regression.
• A causal run with identical fixed routing and instrumentation but Python cyclic GC disabled after initialization removed the tail completely: mean 933.08 -> 557.31 ms, p95 1740.38 -> 581.60 ms, and steps over 800 ms 11/20 -> 0/20. The measured-window mean was 557.24 ms, matching the recorded 557.31 ms. This comparison used separate allocations, so a same-allocation paired confirmation is still required before treating the magnitude as final.

The same-allocation AutoEP-only confirmation is now complete. A dual-warm default -> managed -> managed -> default block used full measured-window timing:

• Default automatic GC: 688.57 ms mean, 535.13 ms median, 1072.52 ms median p95, 11/40 steps over 800 ms, and 4/40 over one second.
• python_gc_policy="disable_during_training": 526.42 ms mean, 522.89 ms median, 551.12 ms median p95, and no steps over 800 ms.
• Paired measured-window savings were 166.37 and 157.92 ms; the paired mean was 162.14 ms, with 95% CI [108.45, 215.84] ms.
• Peak allocated/reserved memory was identical, observed routing was identical, and maximum paired loss differences were 0.00266 and 0.01439.

The opt-in engine-managed policy is isolated in Draft PR #8451. It is independent of this PR's DeepEP local-preparation cleanup.

Rewriting or overlapping DeepEP dispatch would optimize the waiting point rather than the source of the tail.

Natural-routing same-allocation context:

• All four paired L48 deltas also favored AutoEP: 67.03, 55.87, 80.45, and 101.26 ms.
• Paired mean delta was 76.15 ms, with 95% CI [45.10, 107.21] ms; pooled medians were 686.49 ms for AutoEP and 760.23 ms for Megatron, an AutoEP speedup of 9.70%.
• Measured-window loss differed by at most 0.00970.
• Across all 80 measured steps, AutoEP averaged 793.65 ms versus 801.08 ms for Megatron, only a 0.93% average-time advantage despite the larger median signal.
• AutoEP had 28/80 steps over 800 ms versus 14/80 for Megatron, and 9/80 over one second versus 2/80; only one AutoEP step exceeded twice its arm median.
• A separate tokens-per-expert audit showed that the natural-routing workload is not fully matched. Per-layer sorted expert-load total variation had a 9.84% median, but indexed expert variation had a 64.72% median and rank-receive variation had a 46.46% median. The two frameworks therefore see similar load-shape distributions assigned to different expert/rank identities.
• The natural-routing speedup is reported as end-to-end context, not as a pure framework execution gap. The fixed-routing paired comparison remains the controlled cross-framework result.

🤖 Generated with GitHub Copilot CLI

[chatgpt-codex-connector]

Codex Review Summary

This comment shows the latest Codex review activity on this pull request.

Review	Status	Commit	Review trigger
📝 Code Review	✅ Completed 2026-09-06T08:30:25.367949Z	91cf688	Draft marked ready

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review" or "@codex security review".

Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings.

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 91cf688f8c

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

[chatgpt-codex-connector]

Add the mandatory sign-off trailer

This is a one-parent, non-merge commit, but its commit message contains no Signed-off-by: trailer. The repository requires every non-merge commit to carry one, so the commit does not satisfy the project’s contribution requirements and must be recreated with --signoff.

AGENTS.md reference: AGENTS.md:L6-L8

Useful? React with 👍 / 👎.

[tohtana]

Thank you @yh0903!
Let me note one small concern in the parity test. Given the lr and atol, it seems the test might not catch the error. (The correct amount of update is -1e-4 and atol=5e-4). Maybe we can increase lr?

[yh0903]

Good catch, thank you @tohtana. Confirmed the math: at step 1, Adam's bias correction makes every updated parameter's delta equal to +/-lr regardless of the gradient's magnitude, so with the shared config's default lr=1e-4 the correct delta (1e-4) is smaller than atol=5e-4 on the parameter_deltas check. that check could not have told a correct update apart from a missing or wrong-signed one.

Fixed in the latest commit by overriding lr=1e-2 for this test only (the update is applied once, after every value under comparison is already captured, so nothing else in the test depends on the old magnitude).

[tohtana]

Thank you for the update! Looks good to me.