Skip to content

Reduce redundant work in AutoEP token routing - #8209

Queued
hwchen2017 wants to merge 4 commits into
masterfrom
hongwei/moe_routing
Queued

Reduce redundant work in AutoEP token routing#8209
hwchen2017 wants to merge 4 commits into
masterfrom
hongwei/moe_routing

Conversation

@hwchen2017

@hwchen2017 hwchen2017 commented Aug 3, 2026

Copy link
Copy Markdown
Contributor

Motivations
count_tokens_per_expert function was called three times in each forward pass, and the torch.bincount inside 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.

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
delock enabled auto-merge August 4, 2026 07:04
@delock

delock commented Aug 4, 2026

Copy link
Copy Markdown
Collaborator

@hwchen2017 thanks for the PR, I like the part that use scatter_add to replace bincount, very neat!

@delock
delock added this pull request to the merge queue Aug 4, 2026
Any commits made after this event will not be merged.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants