From 8fea1d306fdfe9b7f47a8f45114b2bb23e8c0d9b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Elvir=20Crn=C4=8Devi=C4=87?= Date: Mon, 10 Aug 2026 17:28:57 +0200 Subject: [PATCH] Fix uniform_random routing simulation to sample without replacement (#43680) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Elvir Crncevic Signed-off-by: Elvir Crnčević Signed-off-by: Tyler Michael Smith Co-authored-by: Claude Opus 4.6 Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> Co-authored-by: Tyler Michael Smith --- .../fused_moe/router/routing_simulator_router.py | 12 ++++-------- 1 file changed, 4 insertions(+), 8 deletions(-) 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