[main](moe):Refine A2A overlap under CUDA_DEVICE_MAX_CONNECTIONS=1

[Baidu-AIAK]

Problem Description

The current overlap_moe_expert_parallel_comm in Megatron-LM requires CUDA_DEVICE_MAX_CONNECTIONS to be set to a relatively large value (#2180 (comment)).

However, in certain scenarios, ensuring the correctness of communication and computation scheduling requires setting CUDA_DEVICE_MAX_CONNECTIONS=1. This constraint can weaken the parallelism of All-to-All (A2A) overlap in practice, leading to a noticeable gap between the achieved optimization benefits and the theoretical expectations.

Root Cause Analysis

The A2A overlap optimization works by splitting modules and overlapping the computation of one micro-batch with the communication of another, thereby hiding EP communication latency. The current scheduling logic can be summarized as follows:

Stream	Phase 1	Phase 2	Phase 3
comm_stream	combine_bwd	dispatch_fwd → dispatch_bwd	combine_fwd
comp_stream	attn_fwd → post_attn_fwd	mlp_bwd → mlp_bwd_dw → mlp_fwd	post_attn_bwd → attn_bwd

The theoretical execution timeline of the overlap optimization is illustrated below:

However, in memory-constrained scenarios where both TP and CP are greater than 1, we are forced to set CUDA_DEVICE_MAX_CONNECTIONS=1 to ensure the correctness of synchronization. Under this constraint, the effective execution timeline of the overlap becomes as follows:

As can be observed, the original overlap logic is completely disrupted. This issue is also discussed in #2630 (comment) . We attribute this behavior to setting CUDA_DEVICE_MAX_CONNECTIONS=1, which enforces a serialized kernel submission model. Under this model, the launch order of computation and communication kernels is forced to be consistent across all devices, effectively eliminating the intended concurrency between computation and communication.

• We initially observed that when CUDA_DEVICE_MAX_CONNECTIONS=1 is set, achieving overlap between computation and communication requires launching the communication first, followed by the computation.

  • In the current implementation, the launch of mlp_bwd occurs before dispatch_fwd. As a result, under CUDA_DEVICE_MAX_CONNECTIONS=1, it is impossible to overlap mlp_bwd with dispatch_fwd. Although this overlap can be enabled by reordering the launch sequence, doing so degrades the overlap behavior of several subsequent modules. This issue is also discussed in [QUESTION] MoE communication & computation can only overlap partially #2180 (comment) .

  • As illustrated in the figure above, this launch-order constraint causes subsequent A2A communication modules to overlap with the next computation module instead, leading to a misaligned (shifted) overlap pattern.

• For combine_fwd and post_attn_bwd → attn_bwd, regardless of which is launched first, overlap cannot be achieved when CUDA_DEVICE_MAX_CONNECTIONS=1. This behavior is also reported in [QUESTION] MoE communication & computation can only overlap partially #2180 (comment)

  • We believe this is due to inherent data dependencies: the computation of mlp_bda must occur after mlp.combine, which forces combine_fwd, post_attn_bwd → attn_bwd, and PP_fwd to execute serially.

  • In the current implementation, combine_fwd is launched before post_attn_bwd → attn_bwd. However, because combine_fwd ends with a bda operation, it cannot overlap with post_attn_bwd → attn_bwd.

Furthermore, since the launch of PP_fwd occurs after post_attn_bwd → attn_bwd, PP_fwd cannot be overlapped either.

Our Solution

Our solution can be summarized in two steps. First, we split the original combine module into two separate stages, combine and post_combine, thereby decoupling communication from computation within the combine phase. Second, by further adjusting the scheduling logic, we are able to achieve overlap between all A2A communications and PP communications, even when CUDA_DEVICE_MAX_CONNECTIONS=1 is enforced.

The newly designed scheduling logic can be summarized as follows:

Stream	Phase 1	Phase 2	Phase 3	Phase 4	Phase 5	Phase 6
comm_stream	combine_bwd	dispatch_fwd	dispatch_bwd	combine_fwd	PP_fwd	PP_bwd
comp_stream	post_combine_bwd → attn_fwd → post_attn_fwd	mlp_bwd	mlp_fwd	mlp_bwd_dw → post_attn_bwd → post_combine_fwd	attn_bwd	attn_bwd_dw

The theoretical execution timeline of the overlap optimization is illustrated below:

Evaluation

Finally, we evaluated the performance of the DeepSeek-V3.1 model on 64 Hopper-architecture GPUs using this optimization. With DeepEP optimization and A2A overlap enabled, and with CUDA_DEVICE_MAX_CONNECTIONS=1 set, our approach achieves more than a 10% performance improvement in our target scenario compared to the current A2A overlap implementation in Megatron-LM.

Nsight Systems Comparison

Before Optimization (Nsight Systems):

As shown above, the actual Nsight Systems trace matches the previously illustrated behavior of the current implementation. The overlap logic is misaligned, resulting in exposed (non-overlapped) communication phases for both combine_fwd and PP_fwd.

After Optimization (Nsight Systems):

With our optimization applied, all A2A communications and PP communications are successfully overlapped.

Summary

Overall, our optimization provides a more effective All-to-All (A2A) overlap solution for scenarios in which CUDA_DEVICE_MAX_CONNECTIONS=1 must be enforced. In other words, our approach makes A2A overlap optimization no longer heavily dependent on setting CUDA_DEVICE_MAX_CONNECTIONS=32.

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