diff --git a/vllm/model_executor/layers/fused_moe/router/routing_simulator_router.py b/vllm/model_executor/layers/fused_moe/router/routing_simulator_router.py index 650a87f24c..853e6e6758 100644 --- a/vllm/model_executor/layers/fused_moe/router/routing_simulator_router.py +++ b/vllm/model_executor/layers/fused_moe/router/routing_simulator_router.py @@ -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