From d96e593fd0176c7245935aab2411913494adc5da Mon Sep 17 00:00:00 2001 From: Byron Hsu Date: Mon, 18 May 2026 09:16:12 -0700 Subject: [PATCH] [Benchmark] Add SGLANG_SIMULATE_UNIFORM_EXPERTS for balanced expert routing with dummy weights (#25571) Co-authored-by: Byron Hsu Co-authored-by: Cursor --- python/sglang/srt/environ.py | 1 + python/sglang/srt/layers/moe/topk.py | 12 ++++++++++++ 2 files changed, 13 insertions(+) diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index be06b3d02..9f791751a 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -205,6 +205,7 @@ class Envs: SGLANG_DISABLE_TP_MEMORY_INBALANCE_CHECK = EnvBool(False) SGLANG_SIMULATE_ACC_LEN = EnvFloat(-1) SGLANG_SIMULATE_ACC_METHOD = EnvStr("match-expected") + SGLANG_SIMULATE_UNIFORM_EXPERTS = EnvBool(False) SGLANG_TORCH_PROFILER_DIR = EnvStr("/tmp") SGLANG_OTLP_EXPORTER_SCHEDULE_DELAY_MILLIS = EnvInt(500) SGLANG_OTLP_EXPORTER_MAX_EXPORT_BATCH_SIZE = EnvInt(64) diff --git a/python/sglang/srt/layers/moe/topk.py b/python/sglang/srt/layers/moe/topk.py index 62a22e572..fb1d32123 100644 --- a/python/sglang/srt/layers/moe/topk.py +++ b/python/sglang/srt/layers/moe/topk.py @@ -1465,6 +1465,18 @@ def select_experts( renormalize=renormalize, ) + if envs.SGLANG_SIMULATE_UNIFORM_EXPERTS.get(): + # Benchmark-only: override gating with uniform round-robin expert assignment + # to avoid expert imbalance from dummy/random weights. Do NOT use in production. + num_tokens, k = topk_ids.shape + num_experts = router_logits.shape[1] + offsets = torch.randint(0, num_experts, (num_tokens, 1), device=topk_ids.device) + steps = torch.arange(k, device=topk_ids.device).unsqueeze(0) + topk_ids = ((offsets + steps * (num_experts // k)) % num_experts).to( + topk_ids.dtype + ) + topk_weights = torch.ones_like(topk_weights) / k + topk_ids, topk_weights = _post_process_topk_ids( topk_ids=topk_ids, topk_weights=topk_weights,