FazBrowse GitHub Viewer | Trending |
URL:
| Home
Tools: [Download Repo ZIP]   [Original HTTPS Page]

Reduce redundant work in AutoEP token routing by hwchen2017 · Pull Request #8209 · deepspeedai/DeepSpeed · GitHub

Reduce redundant work in AutoEP token routing - #8209

Merged
delock merged 4 commits into
masterfrom
hongwei/moe_routing
Aug 4, 2026
Merged

Reduce redundant work in AutoEP token routing#8209
delock merged 4 commits into
masterfrom
hongwei/moe_routing

Conversation

hwchen2017 commented Aug 3, 2026
edited
Loading

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

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 added this pull request to the merge queue Aug 4, 2026
Merged via the queue into master with commit 58aab42 Aug 4, 2026
12 checks passed
delock deleted the hongwei/moe_routing branch August 4, 2026 08:08
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

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants


Back | FazBrowse Home | New Git URL