| FazBrowse GitHub Viewer | Trending | | Home |
| Tools: [Download Repo ZIP] [Original HTTPS Page] |
Sorry, something went wrong.
Signed-off-by: Hongwei Chen <hongweichen@microsoft.com>
Signed-off-by: Hongwei Chen <hongweichen@microsoft.com>
Signed-off-by: Hongwei Chen <hongweichen@microsoft.com>
|
@hwchen2017 thanks for the PR, I like the part that use scatter_add to replace bincount, very neat! |
Sorry, something went wrong.
| Back | FazBrowse Home | New Git URL |
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.