Fix FlashInfer A2A IMA by DP-synchronizing the decode graph bucket (#30242) (#30450)

python/sglang/srt/layers/moe/token_dispatcher/flashinfer.py

@@ -7,7 +7,10 @@ import torch
 
 from sglang.kernel_api_logging import debug_kernel_api
 from sglang.srt.environ import envs
-from sglang.srt.layers.dp_attention import get_dp_global_num_tokens
+from sglang.srt.layers.dp_attention import (
+    get_dp_global_num_tokens,
+    is_dp_attention_enabled,
+)
 from sglang.srt.layers.moe.token_dispatcher import (
     BaseDispatcher,
     CombineInput,
@@ -20,11 +23,9 @@ from sglang.srt.layers.moe.token_dispatcher.flashinfer_utils import (
 )
 from sglang.srt.layers.moe.topk import StandardTopKOutput, TopKOutput
 from sglang.srt.layers.moe.utils import get_moe_runner_backend
-from sglang.srt.model_executor.runner_utils.capture_mode import get_is_capture_mode
 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
@@ -200,45 +201,55 @@ class FlashinferDispatcher(BaseDispatcher):
         # runtime_max_tokens_per_rank selection
         # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
         # MoeAlltoAll uses fixed-geometry buffers shaped
-        # [ep_size, runtime_max_tokens_per_rank, ...], so every EP rank
+        # [ep_size, runtime_max_tokens_per_rank, ...], so every EP rank must pass
-        # must pass the same value.  Three cases:
+        # the SAME value. This code (Python) runs during eager forwards and during
+        # CUDA-graph *capture*; on *replay* dispatch() is not re-executed and the
+        # value baked at capture is reused. Two cases, both rank-invariant:
         #
-        # Case 1 — max(dp_global):
+        # Case 1 — max(dp_global): DP attention feeding EP. The scheduler
-        #   DP attention with require_mlp_tp_gather=True.  The scheduler
+        #   all-gathers per-DP-rank token counts into dp_global (length dp_size,
-        #   all-gathered per-DP-rank token counts into dp_global (a list
+        #   identical on every rank), which differ across ranks, so we must take
-        #   of length dp_size); max() is uniform across all ranks and
+        #   the max. FlashInfer A2A forces require_mlp_tp_gather=True (see
-        #   sizes the workspace for the fattest rank.
+        #   require_mlp_tp_gather()), so: eager reads the live list; capture sees
+        #   [num_tokens] * dp_size (uniform capture bs) and bakes max() == the
+        #   bucket; replay reuses that baked value and every rank replays the same
+        #   bucket because the decode graph runner sizes it from the cross-rank
+        #   max. Without this, per-rank buckets could diverge -> geometry mismatch
+        #   -> illegal memory access (issue #30242).
         #
-        # Case 2 — self.max_num_tokens (static capacity):
+        # Case 2 — x.shape[0]: no per-rank DP list (dp_global absent or scalar).
-        #   EP>1 during live (non-capture) inference with
+        #   This is SP attention feeding EP (tokens are sequence-parallel scattered
-        #   require_mlp_tp_gather=False.  The scheduler only stored the
+        #   uniformly, so x.shape[0] is already identical on every EP rank), a
-        #   local token count, so x.shape[0] can differ across EP ranks
+        #   single EP rank, or CUDA-graph capture of those. x.shape[0] is
-        #   that span different DP groups.  The static workspace capacity
+        #   rank-invariant here, so it is both correct and right-sized.
-        #   is the same on every rank, so it is always safe.
-        #
-        # Case 3 — x.shape[0] (actual tensor size):
-        #   Everything else: EP=1, sequence-parallel (post-scatter), or
-        #   CUDA graph capture.  In these situations x.shape[0] is the
-        #   same on every EP rank.  During CUDA graph capture
-        #   (get_is_capture_mode()=True) the graph runner ensures all
-        #   ranks capture with the same batch size, so we skip Case 2
-        #   and land here — using x.shape[0] avoids baking the
-        #   (potentially much larger) static max into the captured graph.
         dp_global = get_dp_global_num_tokens()
         if dp_global is not None and len(dp_global) > 1:
             # Case 1
             self.runtime_max_tokens_per_rank = max(dp_global)
-        elif (
-            self.ep_size > 1
-            and not get_is_capture_mode()
-            and not require_mlp_tp_gather(get_global_server_args())
-        ):
-            # Case 2
-            self.runtime_max_tokens_per_rank = self.max_num_tokens
         else:
-            # Case 3
+            # Case 2. Guard against the #30242 failure mode: DP attention must
+            # never land here with ep_size > 1, because there x.shape[0] differs
+            # across ranks and is NOT a safe fixed geometry. DP attention is
+            # routed to Case 1 via require_mlp_tp_gather=True; reaching here with
+            # DP attention on and ep_size > 1 means the DP all-gather was skipped
+            # (e.g. SGLANG_SCHEDULER_SKIP_ALL_GATHER, unsupported) -> fail fast.
+            assert not is_dp_attention_enabled() or self.ep_size == 1, (
+                "FlashInfer A2A: DP attention reached the x.shape[0] fallback "
+                f"with ep_size={self.ep_size} > 1 (dp_global={dp_global}); "
+                "runtime_max_tokens_per_rank would not be rank-invariant."
+            )
             self.runtime_max_tokens_per_rank = x.shape[0]
 
+        # The recv buffer reserves runtime_max_tokens_per_rank slots for THIS
+        # rank, so it must cover this rank's own tokens. This holds in both cases
+        # (Case 1: max(dp_global) >= the local count; Case 2: exactly x.shape[0]),
+        # so a violation signals a sizing/plumbing bug (e.g. an un-adjusted spec
+        # count) rather than a benign case.
+        assert self.runtime_max_tokens_per_rank >= x.shape[0], (
+            f"runtime_max_tokens_per_rank={self.runtime_max_tokens_per_rank} < "
+            f"x.shape[0]={x.shape[0]}: MoeAlltoAll recv buffer would overflow."
+        )
+
         # Passing topk_ids + invalid_token_expert_id triggers the sanitize step
         # inside moe_a2a. The recv buffer has shape
         # [ep_size, max_tokens_per_rank, ...], so any rank below max leaves

python/sglang/srt/utils/common.py

@@ -3460,6 +3460,16 @@ def require_mlp_tp_gather(server_args: ServerArgs):
             return True
         elif get_moe_a2a_backend().is_none():
             return True
+        elif get_moe_a2a_backend().is_flashinfer():
+            # FlashInfer MoE A2A needs a rank-invariant, DP-synchronized per-rank
+            # token count: MoeAlltoAll uses fixed-geometry buffers and the decode
+            # cuda-graph bucket must be identical across EP ranks, otherwise ranks
+            # replay different-sized graphs -> geometry mismatch -> illegal memory
+            # access (issue #30242). No literal MLP TP-gather happens here -- the
+            # MoE stays SCATTERED and the a2a op owns dispatch/combine -- but we
+            # reuse this flag's DP-sync bookkeeping (uniform global_num_tokens +
+            # max-based graph bucket). See #30432 re: the misleading flag name.
+            return True
         else:
             return (
                 server_args.moe_dense_tp_size

test/registered/unit/test_legacy_global_ratchet.py

@@ -25,7 +25,7 @@ _SRT_ROOT = Path(next(iter(sglang.srt.__path__)))
 # Baselines counted over python/sglang/srt/**/*.py, including each function's
 # own def line. Ratchet: decrease-only.
 _RATCHETS = [
-    ("get_global_server_args", r"\bget_global_server_args\s*\(", 280),
+    ("get_global_server_args", r"\bget_global_server_args\s*\(", 279),
     (
         "set_global_server_args_for_*",
         r"\bset_global_server_args_for_(?:scheduler|tokenizer)\s*\(",