Fix uniform_random routing simulation to sample without replacement (#43680)

Signed-off-by: Elvir Crncevic <elvircrn@gmail.com>
Signed-off-by: Elvir Crnčević <elvircrn@gmail.com>
Signed-off-by: Tyler Michael Smith <tlrmchlsmth@gmail.com>
Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
Co-authored-by: Tyler Michael Smith <tlrmchlsmth@gmail.com>
This commit is contained in:
Elvir Crnčević
2026-08-10 15:28:57 +00:00
committed by GitHub
co-authored by Claude Opus 4.6 gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> Tyler Michael Smith
parent fac808b36f
commit 8fea1d306f
@@ -131,14 +131,10 @@ class DistributionBasedRouting(RoutingStrategy):
"""Sample expert IDs based on the specified distribution."""
if self.distribution == "uniform":
# Uniform random sampling
return torch.randint(
low=0,
high=num_experts,
size=(num_tokens, top_k),
dtype=indices_type,
device=device,
)
# Generate random scores, and take the top-k to avoid duplicate topk_ids
scores = torch.rand(num_tokens, num_experts, device=device)
_, topk_ids = torch.topk(scores, top_k, dim=-1)
return topk_ids.to(indices_type)
elif self.distribution == "normal":
# For normal distribution, sample continuous values and map to