Add an opt-in AutoEP non-MoE regional compile mode

[yh0903]

Summary

• add an opt-in engine.compile(compile_mode="autoep_non_moe") mode that regionally compiles decoder blocks containing AutoEP layers
• keep AutoEPMoELayer.forward as an explicit compiler-disabled graph break, so routing, token movement, expert compute, and AllToAll collectives remain eager
• discover decoder regions from the injected AutoEP module hierarchy instead of relying on Qwen-specific model.layers traversal
• fail fast for unsupported configurations instead of silently changing compile behavior

The default engine.compile() behavior and the default eager AutoEP path are unchanged.

Support boundary

This first version uses vanilla torch.compile/Inductor with the standard AutoEP comm backend, sequence and pipeline parallel sizes of one, and ZeRO stages 0, 1, and 2. Distributed performance and parity validation currently target ZeRO stage 1.

DeepEP, DeepCompile, AutoEP+AutoTP folding, sequence or pipeline parallelism, ZeRO stage 3, optimizer or parameter offload, compiled autograd, DeepCompile schedules, and any fullgraph or dynamic value other than False are rejected.

Performance

Qwen3-30B-A3B, 48 layers, EP16, sequence length 1024, BF16, activation checkpointing enabled, fixed routing, 2x8 H100. Each allocation discarded one fixed warm arm and used two runs per variant.

Allocation	Implementation / order	Eager median	Compiled median	Throughput gain	Eager p95 median	Compiled p95 median
1	pre-production canary, ECCE	936.66 ms	885.82 ms	+5.74%	1627.42 ms	1534.37 ms
2	production engine API, CEEC	932.61 ms	875.10 ms	+6.57%	1537.12 ms	1446.36 ms

Both allocations reduced peak allocated memory by 1.60 GiB and peak reserved memory by 2.01 GiB. The maximum paired loss differences were 0.00617 and 0.00194. All 16 ranks reported zero Dynamo counter changes in the measured window and exactly 2,880 eager AutoEP calls per compiled arm (30 steps x 48 layers x forward/replay).

An exact same-stack Nsight census explained the clean E2E improvement:

Metric	Eager	Compiled	Delta
GPU kernels	22,340	15,236	-7,104 (-31.8%)
NCCL kernels	390	390	unchanged
non-NCCL kernel union	233.73 ms	191.42 ms	-42.31 ms
launch API time	111.76 ms	83.07 ms	-28.69 ms

Compiled execution also removed 1.59 GiB of D2D traffic per captured step. These performance runs used the PR1+PR2+PR3 benchmark stack (#8326, #8331, and #8359) to isolate the remaining Transformer/runtime fragmentation. This PR is based directly on master and has no code dependency on those changes.

Project rebaseline

This is context for the broader AutoEP optimization effort, not a merge gate for this PR. A same-allocation fixed-routing discard + ABBA + BAAB comparison of the full #8326 + #8331 + #8359 + #8380 stack against Megatron Core measured:

Framework	Median step	Median of per-arm p95	Peak allocated	Peak reserved
AutoEP PR1-4	887.21 ms	1244.49 ms	40.37 GiB	45.29 GiB
Megatron Core	814.58 ms	821.30 ms	41.68 GiB	45.80 GiB

The current residual is 72.64 ms / AutoEP +8.9%. All four paired deltas favored Megatron and stayed within 66.55–77.19 ms; AutoEP used 1.30 GiB less allocated and 0.51 GiB less reserved memory.

A separate low-overhead outer CUDA-event trace placed the typical residual primarily in:

• forward: AutoEP +39.14 ms
• backward including ordinary gradient synchronization: AutoEP +46.27 ms
• optimizer: AutoEP -9.00 ms

The AutoEP tail is not a measured-window recompilation. In both traced AutoEP arms, step 27's forward increased from a typical 266–271 ms to 678 ms while backward and optimizer remained near their medians. With only ten measured steps per arm, the reported p95 is the maximum sample; this deterministic forward-only spike is a follow-up investigation rather than evidence that regional compile regressed the typical path.

Testing Done

• Local code review completed
• pre-commit run --files passed for all five changed files
• exact latest-master targeted unit and 2-rank H100 parity suite passed after readiness hardening
• warmup explicitly captured nonzero Dynamo graphs/calls before the zero-recompilation measured-window assertion
• parity covers output, loss, input gradient, router gradients, expert gradients, non-MoE gradients, optimizer parameter deltas, exact routing assignments, and activation checkpointing on/off
• parity tolerances were tightened for input gradients, parameter gradients, and SGD optimizer deltas
• fail-fast coverage includes dynamic=None, fullgraph=None, DeepEP, sequence parallelism, and pipeline parallelism
• full Qwen3-30B-A3B production-API smoke passed with activation checkpointing enabled
• two independent L48 paired allocations completed with ECCE/CEEC order
• exact eager/compiled Nsight kernel census completed
• GitHub full CI needs a maintainer rerun: the selected suite reached 85% with no observed test failure, then the Modal controller cancelled it at the 75-minute limit

🤖 Generated with GitHub Copilot CLI