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>
Elvir Crnčević committed
Aug 10, 2026 at 17:28 UTC
8fea1d306fdfe9b7f47a8f45114b2bb23e8c0d9b
1 file changed
+4
-8
vllm/model_executor/layers/fused_moe/router/routing_simulator_router.py
+4
-8
@@ -131,14 +131,10 @@ class DistributionBasedRouting(RoutingStrategy):
131
"""Sample expert IDs based on the specified distribution."""
132
133
if self.distribution == "uniform":
134
- # Uniform random sampling
135
- return torch.randint(
136
- low=0,
137
- high=num_experts,
138
- size=(num_tokens, top_k),
139
- dtype=indices_type,
140
- device=device,
141
- )
134
+ # Generate random scores, and take the top-k to avoid duplicate topk_ids
135
+ scores = torch.rand(num_tokens, num_experts, device=device)
136
+ _, topk_ids = torch.topk(scores, top_k, dim=-1)
137
+ return topk_ids.to(indices_type)
138
139
elif self.distribution == "normal":
140
# For normal distribution, sample continuous values and map to