Reduce redundant work in AutoEP token routing - #8209
Queued
hwchen2017 wants to merge 4 commits into
Queued
Conversation
Signed-off-by: Hongwei Chen <hongweichen@microsoft.com>
Signed-off-by: Hongwei Chen <hongweichen@microsoft.com>
Signed-off-by: Hongwei Chen <hongweichen@microsoft.com>
delock
approved these changes
Aug 4, 2026
delock
enabled auto-merge
August 4, 2026 07:04
Collaborator
|
@hwchen2017 thanks for the PR, I like the part that use |
delock
added this pull request to the merge queue
Aug 4, 2026
Any commits made after this event will not be merged.
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
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.
Motivations
count_tokens_per_expertfunction was called three times in each forward pass, and thetorch.bincountinside it will introduce cpu-gpu sync. But the results of the first call could be reused.Changes
• Reuse the router's histogram. The router already computes num_tokens_per_expert; reuse it through compute_split_plan and the ep_size == 1 path instead of recomputing it in AutoEPMoELayer.forward .
• Faster count_tokens_per_expert. Replace torch.bincount with a pre-sized zeros(num_experts, int32) + scatter_add_ , avoiding the device-to-host sync that bincount needs . The helper now always returns an int32 histogram; the unused out_dtype / deterministic_safe params and padding logic are removed.
• Remove deterministic_safe path in
count_tokens_per_expert. The histogram of integers is inherently deterministic. The op just sums 1 per bucket. Integer addition is associative and commutative, so the atomic accumulation order has zero effect on the result — every run produces identical counts• Remove the TokenReorderer module. Its logic (argsort by expert + score gather) is a two-liner, now inlined directly in the layer forward.
Performance
The time below is measured from the moe gate kernel to the last kernel before first all-to-all communication.
A100: 2.3ms -> 1.7ms.
H200: 0.89ms -> 0.53ms.