Reduce redundant work in AutoEP token routing (#8209)
**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>
Co-authored-by: Ma, Guokai <guokai.ma@gmail.com>
