[Benchmark] Add SGLANG_SIMULATE_UNIFORM_EXPERTS for balanced expert routing with dummy weights (#25571)

Co-authored-by: Byron Hsu <byronhsu@Byrons-MacBook-Pro.local>
Co-authored-by: Cursor <cursoragent@cursor.com>

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)

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,