DeepSpeed

11b24d21 - AutoEP: one program per token in the fused weighted restore forward (#8678)

AutoEP: one program per token in the fused weighted restore forward (#8678) ## Summary - Keep `combine_impl="fused_weighted_sum"` opt-in and change only its forward scheduling: one Triton program handles a whole token, walking the hidden dimension in up to 1024-column chunks instead of launching a program per token and 256-column tile. The backward kernel and its row/score gradients are unchanged. - Preserve the existing routing, expert reorder, collectives, grouped GEMM and top-k semantics. The default combine path is unaffected. ## Numerical contract - Each row value and routing score are multiplied in FP32, the top-k products accumulate in FP32 in a fixed slot order, and the output is cast to the row dtype once. FP contraction is disabled so the product rounds before addition. Hidden-tail loads/stores are masked; non-power-of-two top-k requires no padded loads; token and source-row address arithmetic uses int64. - **Not bitwise unchanged versus eager:** eager's top-k reduction may sum FP32 products in a different order. The existing operator and model parity tolerances are unchanged; repeated fused calls are deterministic on fixed inputs. The backward implementation is unchanged. - Measured step-0 loss was 13.237928 → 13.237917 at 8k and 13.276934 → 13.276920 at 16k. These are **step-0 observations, not upper bounds on differences at later training steps**. No long-run convergence claim is made. ## Testing Done - Previously validated on NVIDIA H100: **241 passed**, including fused restore and AutoEP unit coverage. The operator tests compare forward, row gradients and routing-score gradients against the eager weighted sum. Existing separate model parity tests cover gradients through the router and experts; they were not rerun here. The wide-hidden test covers 2048 columns and a 1030-column tail with top-k 6, plus repeated-run determinism. - New focused FP32 cancellation case: **1 passed, 0 skipped** on one NVIDIA H100 80GB HBM3 (torch 2.10, CUDA 13.0, Triton 3.6). It checks that the product is rounded before it is added, including a non-power-of-two top-k and a 1030-column hidden tail. As a negative control, allowing FP contraction makes this same case fail. The negative-control change is not in this PR; existing tolerances were not relaxed. - Changed-files pre-commit checks passed. ## Performance and attribution Within the same NVIDIA H100 allocation, at top-k 8 and hidden 2048, the restore forward fell from about **0.60 → 0.105 ms** per call at 8k tokens, and **1.20 → 0.20 ms** at 16k. The complete AutoEP native-path old/new comparison (two repeats per arm, identical setup apart from this forward change) is the primary E2E evidence: | Sequence | AutoEP before | AutoEP after | Difference | | --- | ---: | ---: | ---: | | 8192 | 2547.53 ms | 2500.44 ms | −47.1 ms (−1.85%) | | 16384 | 5244.91 ms | 5145.86 ms | −99.1 ms (−1.89%) | Peak allocated and reserved memory did not increase. On this allocation, the remaining AutoEP/Megatron step-time gaps were about 0.43%/0.61% at 8k/16k; those numbers are context, **not** evidence that Transformer Engine's unpermute has identical semantics. Its unpermute does not apply the routing weights, so no direct kernel-to-kernel speedup is claimed. This PR's code is independent of #8671 and changes different files. The **reported E2E baseline already includes #8671's gather-based reorder**, so those timings should not be compared with runs that omit it. No new full E2E or convergence run was required for this focused forward review. 🤖 Generated with [GitHub Copilot CLI](https://docs.github.com/copilot/github-copilot-cli) --------- Signed-off-by: yh0903 <helloyu0903@gmail.com> Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>