Add an opt-in fused weighted restore for AutoEP - #8326
Open
yh0903 wants to merge 8 commits into
Open
Conversation
The AutoEP path moves routed rows around the two expert all-to-alls with
general-purpose tensor ops. It materializes a padded copy of the token matrix,
gathers through an advanced index, scatters expert outputs into a zero-filled
buffer, and builds a [tokens, top_k, hidden] FP32 intermediate only to apply
routing weights and reduce over top-k. Every one of those steps costs a full
pass over the routed activations, and they repeat in every MoE layer.
Add expert_parallel.local_token_backend. It defaults to "eager", which leaves
the current implementation untouched. Setting it to "fused" swaps the reorder
and the weighted restore for Triton kernels that touch each row once.
Writing the reorder out as four passes shows that all of them are the same row
gather, where perm maps an expert-major slot to its source row, inv is its
inverse, and a negative index reads as zero:
permute forward out[i] = tokens[perm[i]] gather by perm
permute backward dtok[j] = dout[inv[j]] gather by inv
unpermute forward out[j] = expert_out[inv[j]] gather by inv
unpermute backward dexp[i] = dout[perm[i]] gather by perm
So one kernel serves all four. perm is injective on real rows, so no pass needs
atomics, and carrying inv costs 4 bytes per row against the 2 * hidden bytes per
row that the padded copy and the zero-filled scatter buffer cost today.
The restore goes straight from [tokens * top_k, hidden] to [batch, seq, hidden],
weighting each row and reducing over top-k in registers. It keeps the eager
dtype discipline: FP32 product and accumulation, one cast on the way out. Its
backward produces both the expert-output gradient and the routing-score
gradient, reducing the score gradient over the hidden dimension per token so
that no cross-program atomic is needed.
The collectives, the router and the grouped GEMM are unchanged, so a measured
difference between the two backends belongs to the local token engine alone.
"fused" is rejected rather than quietly ignored wherever it would have nothing
to replace or would change semantics: folded tensor parallelism, which restores
combined tokens from assignment metadata; an explicit combine_impl="legacy_bmm";
a resolved score_apply other than "post"; and non-CUDA, non-Triton or
non-bf16/fp16 execution, which is checked once before any collective so that
ranks fail together instead of stalling. A run that asked for "fused" and
silently got "eager" would otherwise report the difference between a backend and
itself.
The alignment and index generation that both backends share moves to
generate_local_expert_permute_indices, so the two cannot drift apart.
Signed-off-by: yh0903 <helloyu0903@gmail.com>
The kernel tests cover the reorder and the weighted restore in isolation. They cannot show that a step taken through the fused backend trains the same model, which is the property that decides whether the backend is safe to select. Add a parity test that runs one step through each backend from the same initial state and the same batch, then compares the loss, the block output, the input gradient, every trainable gradient by name, and the parameter delta the optimizer produced. Router and expert gradients reach the comparison by different routes through the restore, so they are asserted to be present rather than left to a bulk comparison that would pass on an empty set. The benchmarked configuration recomputes each MoE block in backward, so the expert-parallel case runs with and without activation checkpointing. A local case with autoep_size=1 covers the branch that skips the all-to-alls. Also fix the fail-fast assertion in the kernel tests: it matched on the required score_apply rather than the rejected one, so it failed against a correct message. Validated on an H100 node: 139 passed, none skipped, for the kernel and config tests, and all three parity cases passed. Signed-off-by: yh0903 <helloyu0903@gmail.com>
Timing the two fused ops against the eager ones they replaced, on an H100 at the
canonical shape, separates them cleanly:
weighted restore forward 2.98x fwd+bwd 1.87x saves 0.215 ms per layer
expert reorder forward 0.89x fwd+bwd 1.01x saves 0.004 ms per layer
The reorder is a wash, and its forward is slower than the eager one. Three
attempts to fix that -- a wider hidden tile, more warps, and several rows per
program -- all landed within noise of the same number. PyTorch's advanced-index
gather and scatter are already close to optimal for this access pattern, and the
fused version additionally has to build an inverse index the eager one does not
need. So it is removed: it carried a kernel, an autograd pair, an index buffer
and a shared-helper refactor, and returned nothing.
An earlier measurement had the reorder at 1.23x. That was an artifact of timing
one iteration at a time and synchronising after each, which charges every
iteration the launch latency a training step hides by queueing work ahead of the
GPU. Timing a batch of iterations between one pair of events measures the
regime these ops actually run in, and the reorder's apparent win disappeared.
What remains is the weighted restore, so it is spelled as what it is: another
implementation of the combine, selected by combine_impl="fused_weighted_sum"
alongside the existing weighted_sum and legacy_bmm, rather than by a second
config key describing a "local token backend" that now moves no tokens. This
also dissolves the question of what a fused backend should do when someone asks
for legacy_bmm: they are alternatives in one enum and cannot both be chosen.
It is still rejected rather than quietly ignored where it has nothing to
replace: folded tensor parallelism, a resolved score_apply other than "post",
and non-CUDA, non-Triton or non-bf16/fp16 execution, checked once before any
collective so ranks fail together instead of stalling.
Validated on an H100 node: 121 kernel and config tests pass, and all three
eager-versus-fused parity cases pass, comparing loss, output, input gradient,
every named gradient, and the optimizer's parameter delta.
Signed-off-by: yh0903 <helloyu0903@gmail.com>
Validate AutoTP folding and expert tensor parallelism independently so neither configuration can mask the other before process-group setup. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: yh0903 <helloyu0903@gmail.com>
Reject malformed tensor contracts before launching Triton, keep malformed permutations deterministic, and use a differentiable PyTorch backward when create_graph requests higher-order autograd. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: yh0903 <helloyu0903@gmail.com>
Keep the optional fused restore unavailable on HIP without importing pytorch-triton-rocm, matching DeepSpeed import-time device compatibility. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: yh0903 <helloyu0903@gmail.com>
Keep comments focused on numerical behavior, fail-fast ordering, and non-obvious kernel choices while removing repeated implementation narration. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: yh0903 <helloyu0903@gmail.com>
yh0903
marked this pull request as ready for review
August 26, 2026 23:23
yh0903
requested review from
hwchen2017,
loadams,
tjruwase and
tohtana
as code owners
August 26, 2026 23:23
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: c02ea19b5f
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
combine_impl="fused_weighted_sum"path while keeping the existing eager weighted reduction as the default.[tokens * top_k, hidden]expert rows directly to[batch, sequence, hidden]in one Triton pass with FP32 weighting and accumulation, including expert-row and routing-score gradients.Performance
[-1.12%, +3.02%]; this is not yet a statistically conclusive end-to-end win.Testing Done
Local code review completed
Unit tests added/updated
Integration tests pass
Manual testing performed
Repository pre-commit hooks passed for all nine changed files.
H100 fused token-op suite: 23 passed, including top-k 2/4/6/8, non-power-of-two hidden sizes, input-contract validation, forward/backward parity, and double backward.
H100 configuration tests: 5 passed for standard, folded AutoTP, expert tensor parallelism, and score-application validation.
Full-step eager/fused parity: 3 passed, covering loss, output, input gradients, router/expert gradients, optimizer parameter deltas, activation checkpointing on/off, EP2, and local experts.
Earlier final restore-only sweep: 121 kernel/config tests passed before the additional review-driven safety cases were added.