[BUG] Shared expert overlap delays independent expert wgrad until input-gradient merge

Author: fkunerCreated Sep 16, 2026Updated Sep 18, 2026
Labelscommunity-requestwaiting-on-maintainers

Describe the bug

When moe_shared_expert_overlap=True and overlap_dispatch_backward_with_experts_wgrad=True are enabled together, the Flex dispatcher's autograd graph makes routed-expert weight-gradient (wgrad) submission depend on the shared/routed input-gradient merge. This introduces a dependency on shared-expert backward before wgrad can be submitted, undermining the intended dispatch-backward/wgrad overlap.

The dependency is visible in dev commit bb5dfd08f09ce06c5925af453fef06b3129f199d:

  1. MoELayer.dispatch inserts _RegisterDelayedWgradForExperts before calling the token dispatcher.
  2. Flex token_dispatch feeds that same tensor into dispatch and the shared-expert branch.
  3. The registration node's backward therefore waits for the gradients from both branches before calling backward_dw.
mermaid
flowchart LR
  D["Dispatch backward"] --> M["Shared + routed input-gradient merge"]
  S["Shared FC1 backward"] --> M
  M --> W["Submit routed expert wgrad"]

This dependency is in the common Flex graph structure; it is not inherently specific to DeepEP V2.

Steps/Code to reproduce bug

On the dev revision above, use a supported Flex backend and a MoE layer with shared experts, enabling:

python
moe_token_dispatcher_type = "flex"
moe_shared_expert_overlap = True
overlap_dispatch_backward_with_experts_wgrad = True

Run a forward/backward pass with gradients enabled. To inspect the scheduling dependency:

  • Register an autograd node post-hook on the dispatched hidden states.
  • Register a tensor gradient hook on the input received by shared_experts.pre_forward_comm; this observes the shared/routed input-gradient merge.
  • Record entry into MoELayer.backward_dw.

The graph requires the merge before wgrad submission, even though routed-expert dgrad has its own completion event.

Expected behavior

Submit routed-expert wgrad once expert dgrad is ready and dispatch backward has submitted communication, without waiting for the shared-expert input-gradient merge. Preserve the completion wait and gradient-processing hooks at their existing boundary.

Additional context

This report concerns an avoidable scheduling dependency. Whether removing it improves latency depends on workload size and the available communication-overlap window.

cc @NVIDIA/mcore-oncall