[sanketpurandare]

Enqueue independent shared experts before router/dispatch/routed compute on
a separate lower-priority CUDA stream and join both branches before combination.
Reuse the common MoE boundary, including Kimi's latent routed projection,
without duplicating router/sharding/TP-add logic. Dist-MoE's BF16 and MXFP8
transforms opt in automatically; no-shared-expert modules are unaffected.

The caller may already have lowest priority, so cache one routed priority-1
and shared priority0 stream per indexed device rather than changing global
stream policy. Initialize them before graph capture. Explicit waits and
record_stream retain input/output allocator ownership in train and no-grad eval.
Autograd uses each branch's forward stream during backward.

Moving branch readiness can change BF16 input-gradient rounding. A storage-free
input fork combines shared+routed+router contributions in the original order,
including separate TP-sharded input edges and unused gradients. No parameter
gradient cache, extra tensor copy, precision change or owned output is added.
GraphTrainer's own parallel-branch lowering remains downstream work.

Related DeepEP-specific #5263 by @yuweih205 uses a different dispatch callback;
this is generic MoE/Dist-MoE integration and does not copy that implementation.

Verification

Owning-commit CPU checks: 82 passed in 7.77s.

Strict relevant Sphinx builds and all non-type repository hooks passed.
Three unchanged main type findings remain attributed, not reported as a clean
full type run. External reviewer and CI acceptance are separate, still pending.

Complete-stack integration (not this PR in isolation): real four-GPU PP2/DP2/
EP2 and PP4/DP1/EP1 ran ten updates with exact model/optimizer tensor parity,
optimizer graphs and sequential validation. A fresh native PP2/PP4 checkpoint
matrix matched resumed state exactly. These are debug-model contract checks,
not 671B throughput or convergence results.

The final 207-case GPU matrix covers in-place gradients, dense attention, PP/FSDP train/eval graphs, shared-stream backward/capture and fused numerical contracts.

[tianyu-l]

I got a little confused: DistMoE has routed experts comm-comp overlap. Why'd we still need shared expert overlap?

[sanketpurandare]

Dist-MoE's communication/GEMM overlap is inside the routed branch; it doesn't enqueue the separate shared FFN. This launches that independent branch before routed metadata/barrier work and joins it before the outputs are combined.

[tianyu-l]

@claude review

[claude]

Claude finished @tianyu-l's task in 4m 29s —— View job

---

Review: request changes

The design holds up and the gradient-order handling is careful. Two things block it. Turning this on unconditionally from DistMoeTransform crashes --create_seed_checkpoint for every Dist-MoE model with shared experts. It also changes the "verified" 671B recipes without any 671B perf data.

Blocking

1. Seed-checkpoint creation now raises. torchtitan/models/common/moe/moe.py:1116-1119
For create_seed_checkpoint, training_engine.py:347-349 uses init_device="cpu" and buffer_device=None. That makes _init_self_buffers resolve buffer_device to router.tokens_per_expert_E.device, which is CPU. Then _expert_compute_streams(cpu) raises ValueError("... requires a CUDA device"). DistMoeRoutedExperts._init_self_buffers does nothing, so this path worked before the PR. Fix: only prime the streams when buffer_device.type == "cuda". The forward path already validates the device when the overlap actually runs.
Fix this →

2. The opt-in lives on the wrong switch and is always on. torchtitan/config/transform/dist_moe.py:174-175
Shared-expert overlap is a generic MoE feature; the GPU tests run it with stock RoutedExperts. Even so, DistMoeTransform is the only way to turn it on, and it can't be turned off there. That silently changes:

• deepseek_v3_671b_dist_moe_bf16 / _mxfp8, which are documented as "verified" recipes
• the Muse Spark Dist-MoE recipes
• the RL generator in torchtitan_recipes/rl/dapo_math.py, which captures FULL CUDA graphs in vLLM

The PR body says the evidence is "debug-model contract checks, not 671B throughput or convergence results". Per "protect battle-tested code paths", please pick one:

• add overlap_shared_experts: bool = False to DistMoeTransform, or let recipes set MoE.Config.overlap_shared_experts directly; or
• attach a 671B trace or step-time comparison that shows the gain.

Either would also answer @tianyu-l's question concretely. The README explains where overlap could happen, not that it does.

Correctness question

3. record_stream only covers shared_x_TD. moe.py:916-917

• routed_x_TD and routed_padding_mask_T come from _maybe_shard_routed_branch_inputs_across_tp (spmd.redistribute I->S(0) and R->S(0)). Those run on the caller stream.
• Both tensors are then consumed on the routed stream, and the router gate and routed experts save them for backward.

If the local redistribute allocates new memory instead of returning a dim-0 view of x_TD, the block goes back to the caller-stream pool after the routed-stream backward with no event recorded. The caller stream could then reuse it while routed kernels are still running.

Please either record_stream(routed) those tensors when they don't alias x_TD, or add a comment explaining why they always alias. Does the 207-case matrix include EP+TP with dense SP off?

Non-blocking

4. Generic tests are skipped without an optional package. TestSharedExpertStreams sits in tests/unit_tests/gpu/models/common/test_dist_moe.py, which calls pytest.importorskip on dist_moe at module level (line 22). The tests only exercise common MoE with stock routed experts, so any CI without that package skips them. Suggest moving them to their own file, e.g. test_moe_shared_expert_streams.py.
5. The one configuration that enables the overlap is untested. No unit test runs it with Dist-MoE routed experts; the CPU test only checks that the flag gets set.
6. README:
   • Drop the "[MoE] Opt-in shared-expert overlap with deferred DeepEP dispatch #5263 ... not a copy of that dispatch callback" sentence. Where the code came from belongs in the PR description, not user docs.
   • priority0 / priority-1 should be priority 0 / priority -1.

Looks good

• Gradient order is preserved. _ExpertInputs reproduces the original (shared + routed) + router sum. In the TP case, where shared and routed take different inputs, the shared gradient still reaches x_TD before the redistribute backward, which matches the sequential order. Both cases are covered by bit-exact CPU tests.
• Kimi refactor is behavior-neutral. Moving Kimi onto _forward_routed_experts removes its duplicated TP and shared-add code. aux_loss_denominator already reached the router through **router_kwargs, so nothing changes there. MuseSparkMoE rejects shared experts, so it is unaffected.
• Eval is safe. Joining both streams before the add, plus record_stream on shared_TD and out_TD, covers no-grad eval, where autograd keeps no intermediates.