AutoEP: support combine_impl="fused_weighted_sum" with AutoEP+AutoTP folding

[pengdurice]

AutoEP: support combine_impl="fused_weighted_sum" with AutoEP+AutoTP folding

Part of: #8672 (Q4 2026 roadmap), "Mixture-of-Expert (MoE) Training Scaling and Kernel Development": "AutoEP Performance Optimization & Parallel Folding (AutoEP + AutoTP)" (PoC @tohtana @pengdurice). Follows #8670.

Problem

combine_impl="fused_weighted_sum" was refused whenever tensor_parallel.autotp_size > 1:

• deepspeed/module_inject/auto_ep_config.py (validate_autoep_config) raised ValueError.
• deepspeed/module_inject/auto_ep_layer.py (set_deepspeed_parallelism) raised the same error.

The reason was that the folded path restores tokens with its own function, restore_combined in
deepspeed/moe/ep_tp_dispatch.py, and the fused kernel was only wired into the unfolded path.

restore_combined works like this:

1. Each TP peer sends only part of the routed assignments to the experts (partition_assignments).
2. After the experts, it all-gathers the rows, their token index, their top-k slot and their routing weight over the
   TP group.
3. It casts the routing weights to the activation dtype, multiplies, and adds the products with one index_add_ per
   top-k slot, all in the activation dtype.

So under folding, users could not get the memory saving and the FP32 accumulation of the fused kernel, and the
folded restore did its arithmetic in bf16.

Change

• New restore_combined_fused in deepspeed/moe/ep_tp_dispatch.py:
  • It does the same differentiable all-gathers as restore_combined for the rows and the routing weights, so the
    tp_size gradient factor that the folded AVERAGE and EXPERT_TP_CANCEL strategies cancel is unchanged.
  • It all-gathers one flat index (token * top_k + slot) instead of the token index and the slot separately.
  • It keeps the routing weights in the router dtype, scatters them into [tokens, top_k], and calls the existing
    fused_weighted_restore. The kernel is unchanged.
  • It raises a clear RuntimeError if the gathered rows do not cover every (token, slot). The folded payload never
    drops or pads assignments today, so this cannot trigger now; it protects the kernel from reading an unset source
    row if capacity dropping is added later.
  • validate_coverage=True runs the same coverage check as restore_combined.
• AutoEPMoELayer._forward calls restore_combined_fused in the folded branch when
  combine_impl == "fused_weighted_sum". The default ("auto" → "weighted_sum") still calls restore_combined.
• The two autotp_size > 1 refusals are removed. The other fused refusals stay: expert_tensor_parallel_size > 1,
  DeepEP, score_apply other than "post", and unsupported device or dtype.
• Docs: docs/code-docs/source/autoep.rst describes the folded fused restore and drops the refusal from the list;
  docs/_pages/config-json.md drops the autotp_size=1 requirement.

Tests

All GPU runs: one node of 8 NVIDIA H200 GPUs. torch 2.8.0+cu126 with Triton 3.4.0, and torch 2.9.1+cu128 with Triton 3.5.1.

Test	What it checks	GPUs	torch 2.8.0	torch 2.9.1
test_autoep_unit.py::TestAutoEPConfig::test_fused_combine_accepts_folded_tensor_parallelism (was ..._rejects_...)	the config is accepted with tp_size=2	0	passed	not run
test_autoep_autotp_dispatch.py::test_restore_combined_fused_rejects_a_missing_assignment	a missing (token, slot) raises instead of reading an unset row	0	passed	passed
test_autoep_autotp_dispatch.py::TestRestoreCombinedFusedTensorParallel (fp32 and bf16 rows)	2 TP peers with different row counts, each back-propagating a different gradient: the output is bitwise equal to the unpartitioned fused_weighted_restore, and the row and score gradients, summed over the peers, match the unpartitioned gradients	2	2 passed	2 passed
test_autoep_fused_parity.py::TestAutoEPFusedParityFoldedTensorParallel (ZeRO 0 and 2, with and without activation checkpointing)	one training step with TP 2 × EP 2: the loss, output, input gradient, every parameter gradient and every parameter update of the fused run match the default folded run	4	4 passed	4 passed
Rest of test_autoep_autotp_dispatch.py	no change to restore_combined	0	9 passed	9 passed
test_autoep_fused_token_ops.py	the kernel is unchanged	1	63 passed	not run
test_autoep_autotp_grad_parity.py with DEEPSPEED_RUN_H100_TESTS=1	no change to the default folded path	up to 8	22 passed	not run
Whole test_autoep_unit.py	no other config change	up to 8	198 passed (master: 198 passed)	not run

Rerun on master d371f33cc with this change (torch 2.8.0+cu126): test_autoep_autotp_dispatch.py 12
passed; test_autoep_fused_parity.py 4 passed (the 4 new folded cases) and 3 failed; test_autoep_unit.py 198 passed.
The 3 failures are existing cases (TestAutoEPFusedParityExpertParallel x 2, TestAutoEPFusedParityLocalExperts) and
fail the same on master without this change: "Expected mat_a to be BFloat16 matrix got Half", torch 2.8's
_grouped_mm takes bf16 only and these cases run in fp16.

The end-to-end folded test runs in FP32, on MockMoEOnlyTransformer, without grouped GEMM:

• FP32: the default folded restore rounds each product and each top-k addition to the activation dtype. In fp16 the
  two restores differ by 1 to 1.5 units in the last place of the MoE output (about 1,600 in this model), which is more
  than PARITY_TOLERANCE allows after the LM head. In FP32 the two forward passes are bitwise equal.
• No attention: the key bias of nn.MultiheadAttention has a true gradient of zero, so its computed gradient is
  rounding noise, and Adam's first step turns that noise into ±lr. The parameter-update check then fails by 2 × lr no
  matter which restore is used.
• No grouped GEMM: on torch 2.8, torch._grouped_mm accepts only bf16, so the fp32 run would fail before reaching the
  restore.

Existing tests in test_autoep_fused_parity.py that fail on master c6af3faf2 in this environment, before this change
(not changed here):

• torch 2.8.0: all 3 fail with RuntimeError: Expected mat_a to be BFloat16 matrix got Half. The tests use fp16 and
  the default grouped GEMM.
• torch 2.9.1: TestAutoEPFusedParityExpertParallel (2 cases) fails on the input gradient: 5 of 1,024 elements, with
  the largest relative difference 2.1% against a 1% limit. Under fp16 the input gradient still carries the loss scale
  of 256. TestAutoEPFusedParityLocalExperts passes.

On master 5f266e2 (2026-10-15; 1 node of 8 H200, driver 580.126.09, torch 2.8.0+cu126): 214 passed, 3 failed: the 3 fp16 cases of test_autoep_fused_parity.py (TestAutoEPFusedParityExpertParallel x2, TestAutoEPFusedParityLocalExperts: torch 2.8 _grouped_mm takes bf16 only), which fail the same way on master 5f266e2 without this change.
Failures that also fail on master are rerun on master in the same job.

Benchmark

One TP group of 2 H200 GPUs on one node. bf16 rows, FP32 routing weights, random routing with distinct experts per
token. The hidden size, top-k and expert count are the MiniMax-M3 and GLM-5.2 values. Memory is the peak above the
inputs. Time is one forward and backward, averaged over 10 runs after 3 warm-up runs, and includes the identical
partition step. Errors are against an FP64 reference.

Shape (tokens per TP group, top-k, experts, hidden)	Forward peak, GiB (default → fused)	Backward peak, GiB	Time, ms	Output error, relative L2	Score gradient error, relative L2
8,192, 4, 128, 6144 (MiniMax-M3)	1.22 → 0.75	1.59 → 0.84	13.3 → 7.4	3.58e-3 → 1.66e-3	2.35e-3 → 8.4e-8
32,768, 4, 128, 6144 (MiniMax-M3)	4.88 → 3.00	6.38 → 3.38	41.1 → 21.0	3.57e-3 → 1.66e-3	2.34e-3 → 8.3e-8
8,192, 8, 256, 6144 (GLM-5.2)	2.35 → 1.50	3.09 → 1.60	26.7 → 12.0	4.36e-3 → 1.66e-3	2.34e-3 → 8.6e-8
32,768, 8, 256, 6144 (GLM-5.2)	9.39 → 6.01	12.38 → 6.38	91.9 → 38.8	4.36e-3 → 1.66e-3	2.35e-3 → 8.7e-8

The fused output error of 1.66e-3 is the rounding of the bf16 output itself. The score gradient is the router
gradient; the default folded path computes it through bf16 weights, the fused path in FP32.

Script: bench_folded_restore.py (run with torchrun --nproc_per_node=2).

Notes

• The default folded path and the fused path give close but not bitwise-equal results, because the default path
  multiplies and adds in the activation dtype.
• AutoEP+AutoTP folding with ZeRO-3 is still refused on master; this change does not alter that.

Duplicate check

Searched open and recent pull requests on 2026-10-15 ("fused_weighted_sum", "combine", "folding", "AutoEP + AutoTP"): none enables the fused combine under folding. #8829 (optimizer offload under folding) and #8823 (ZeRO-1/2 clipping under folding) change other parts of folding; #8745 and #8739 change the DeepEP combine.

Proposal

New functionality: per CONTRIBUTING.md it needs an accepted proposal issue first. Draft:
descriptions/r3/peng-autoep-fused-combine-folded-tp-v2-proposal-issue.md. TODO(user): open the issue, wait for acceptance, then link it here.

AI assistance

This change, its tests and this description were prepared with AI assistance (Claude).