Add opt-in fused AutoEP DeepEP row weighting (#8644)
## Summary
- Add experimental `expert_parallel.row_weighting_impl="fused"` for the
AutoEP DeepEP route at both `score_apply="pre"` and `"post"` boundaries.
The default `"auto"` still resolves to the existing eager expression;
unsupported configurations fail explicitly rather than falling back.
- The Triton operator multiplies each received row by its FP32 routing
weight, rounds at the same BF16/FP16 boundary, and leaves routing,
expert GEMM, DeepEP combine, and the top-k reduction unchanged.
- In a prior within-allocation H200 A/B screen, fused row weighting
reduced median step time by 2.10% (four repeats, 95% CI 0.075%-3.45%).
Peak allocated memory differed by 192 KiB and peak reserved memory was
510 MiB higher; the latter is an observation, not an explained or
accepted memory improvement.
## Numerical semantics
- The forward and row gradients use eager's FP32 products and rounding.
The routing-weight gradient sums the same FP32 products in another
order, so bitwise equality with eager is not expected. Against an FP64
reference at 4096 x 2048, fused/eager maximum absolute errors were
1.77e-5/1.49e-5 and maximum relative errors were 8.5e-5/2.9e-4; neither
order is consistently more accurate.
- The model regression snapshots retained routing-score, router-gate,
full-parameter, and input gradients **after `engine.backward()` and
before `engine.step()` clips gradients**, using `||fused - eager||_2 /
||eager||_2 <= 2e-2` for each checked tensor/group. Same-seed
eager/eager controls were bit-exact; the largest previously measured
native fused/eager model-gradient relative L2 was 5.6e-3 (reentrant
checkpointing), leaving about 3.6x headroom. Near-zero gradient sign
flips can cause learning-rate-scale differences in Adam's first
parameter update, so coordinate-wise update equality is not used as the
native-path gate.
## Testing Done
- [x] Local code review completed.
- [x] Unit and GPU-gated operator tests cover default selection,
fail-fast inputs, score boundaries, forward/row-gradient parity, and the
unmodified fused weight reduction against FP64 with an FP32 tree-sum
error bound. The previously run GPU operator checks are reused; they are
not a substitute for the model-level test.
- [x] Native-backward distributed BF16 model tests: **2 passed, 0
skipped** (`pre` and `post`) on 4 ranks of an 8-GPU NVIDIA H100 80GB
HBM3 node. Neither case replaces `d_weights`; both check exact routes
and normalized errors through routing scores, router gate, all model
parameters, and model input. The separate summation-order-matched
attribution controls remain in place.
- [x] Reviewer follow-up: only explicit eager/fused comparisons use
`lr=1e-2`; the original default/cleanup/legacy cases retain `lr=1e-4`. A
targeted run on 4 ranks of an 8-GPU NVIDIA H100 80GB HBM3 node passed
**7/7, 0 skipped**: three config-scope cases, both cleanup/legacy cases,
and eager/fused comparisons at both score boundaries.
- [x] Changed-files pre-commit checks passed for all ten files touched
by this PR.
- A prior broad run accidentally used a stale validation bundle and did
not exercise the two native-backward cases; it reported 72 passed and
one legacy non-reentrant checkpoint/optimizer-delta mismatch. That
legacy case passed in the targeted rerun after restoring its original
learning rate. Full-suite pass is **not** claimed from the broad run.
## Dependencies and scope
- The cached-layout repair (#8636) and independent MoE gradient-norm fix
(#8635) have merged into upstream `master`. The GPU tests used an
earlier pinned source with both corrections applied separately; neither
is included in this PR. The norm fix applies to `FP16_Optimizer`, which
is also used by some BF16 configurations. The new model-level gradient
gate reads gradients **before** this wrapper's norm/clipping in
`engine.step()`.
- The integrated DeepEP route currently accepts BF16 rows. The operator
also supports FP16 rows, but is CUDA/Triton-only, requires contiguous
`[N, H]` rows and FP32 `[N, 1]` weights, and does not support ROCm or
non-DeepEP/EP=1 routes. DeepEP v2 requires a compatible NCCL build with
GIN support.
---------
Signed-off-by: yh0903 <helloyu0903@gmail.com>
Co-authored-by: pengdurice <pengduhit@gmail.com>
