From f23f48df98f702496c85d34e67069f6b8d94979e Mon Sep 17 00:00:00 2001
From: YAMY <74099316+YAMY1234@users.noreply.github.com>
Date: Fri, 12 Jun 2026 16:11:44 -0700
Subject: [PATCH] fix(moe): make FlashInfer A2A robust to collapsed
 global_num_tokens (moe_dense_tp_size NaN) (#27945)

---
 .../sglang/srt/layers/moe/token_dispatcher/flashinfer.py   | 7 +++++++
 1 file changed, 7 insertions(+)

diff --git a/python/sglang/srt/layers/moe/token_dispatcher/flashinfer.py b/python/sglang/srt/layers/moe/token_dispatcher/flashinfer.py
index bbdadc7c8..f66c83137 100644
--- a/python/sglang/srt/layers/moe/token_dispatcher/flashinfer.py
+++ b/python/sglang/srt/layers/moe/token_dispatcher/flashinfer.py
@@ -23,6 +23,7 @@ from sglang.srt.layers.moe.utils import get_moe_runner_backend
 from sglang.srt.server_args import get_global_server_args
 from sglang.srt.speculative.spec_info import SpeculativeAlgorithm
 from sglang.srt.utils import get_int_env_var
+from sglang.srt.utils.common import require_mlp_tp_gather
 
 try:
     from flashinfer import nvfp4_block_scale_interleave
@@ -210,6 +211,12 @@ class FlashinferDispatcher(BaseDispatcher):
             # DP attention: multiple DP ranks with different token counts.
             # Use the max across ranks so the A2A workspace fits the fattest.
             self.runtime_max_tokens_per_rank = max(dp_global)
+        elif self.ep_size > 1 and not require_mlp_tp_gather(get_global_server_args()):
+            # require_mlp_tp_gather is False, so the scheduler collapsed
+            # global_num_tokens to the local count; x.shape[0] then differs
+            # across EP ranks and breaks the fixed-geometry MoeAlltoAll. Use
+            # the static all-rank capacity instead.
+            self.runtime_max_tokens_per_rank = self.max_num_tokens
         else:
             # dp_size=1 or SP: use the actual input tensor size (post-scatter
             # in SP mode, full batch otherwise).  Avoids the pre-scatter