[Bugfix] Keep a shared MAX_LEN prefill CUDA graph bucket when the graph captures a DP gather (MegaMoE sparse-DP hang) (#37933)

Co-authored-by: shyeh25 <206795756+shyeh25@users.noreply.github.com>
Co-authored-by: Po-Han Huang (NVIDIA) <53919306+nvpohanh@users.noreply.github.com>

test/registered/unit/model_executor/runner/test_prefill_cuda_graph_padding.py

@@ -2,10 +2,13 @@ import unittest
 from types import SimpleNamespace
 from unittest import mock
 
+from sglang.srt.layers.moe.utils import MoeA2ABackend
+from sglang.srt.model_executor import forward_batch_info
 from sglang.srt.model_executor.cuda_graph_config import Backend
 from sglang.srt.model_executor.forward_batch_info import (
     CaptureHiddenMode,
     ForwardMode,
+    prefill_graph_tolerates_sum_len,
 )
 from sglang.srt.model_executor.runner.prefill_cuda_graph_runner import (
     PrefillCudaGraphRunner,
@@ -72,6 +75,68 @@ class TestPrefillCudaGraphPadding(CustomTestCase):
             forward_batch, num_qo_tokens=16
         )
 
+    def _megamoe_no_prefill_cp(self, graph_has_dp_gather=False):
+        return (
+            mock.patch.object(
+                forward_batch_info,
+                "get_flags",
+                return_value=SimpleNamespace(
+                    dp=SimpleNamespace(prefill_graph_has_dp_gather=graph_has_dp_gather)
+                ),
+            ),
+            mock.patch(
+                "sglang.srt.layers.moe.utils.get_moe_a2a_backend",
+                return_value=MoeA2ABackend.MEGAMOE,
+            ),
+            mock.patch(
+                "sglang.srt.layers.attention.dsa.utils.is_dsa_enable_prefill_cp",
+                return_value=False,
+            ),
+            mock.patch(
+                "sglang.srt.layers.cp.utils.is_mla_cp_enabled",
+                return_value=False,
+            ),
+        )
+
+    def test_megamoe_idle_rank_without_graph_gather_keeps_sum_len(self):
+        # Without a DP gather in the graph, an eager idle rank matches the
+        # replaying peers, so the sparse batch keeps per-rank buckets.
+        runner = self._make_runner()
+        flags, a2a, dsa_cp, mla_cp = self._megamoe_no_prefill_cp()
+        with flags, a2a, dsa_cp, mla_cp:
+            self.assertTrue(prefill_graph_tolerates_sum_len())
+            self.assertFalse(
+                runner._has_inactive_dp_rank(
+                    SimpleNamespace(global_num_tokens_cpu=[8, 0])
+                )
+            )
+
+    def test_megamoe_all_ranks_busy_keeps_per_rank_buckets(self):
+        runner = self._make_runner()
+        flags, a2a, dsa_cp, mla_cp = self._megamoe_no_prefill_cp()
+        with flags, a2a, dsa_cp, mla_cp:
+            self.assertTrue(prefill_graph_tolerates_sum_len())
+            self.assertFalse(
+                runner._has_inactive_dp_rank(
+                    SimpleNamespace(global_num_tokens_cpu=[8, 16])
+                )
+            )
+
+    def test_megamoe_graph_with_dp_gather_forces_shared_bucket(self):
+        # A DP gather captured in the graph has fixed MAX_LEN geometry; per-rank
+        # buckets or an eager idle rank would deadlock its all_gather.
+        runner = self._make_runner()
+        flags, a2a, dsa_cp, mla_cp = self._megamoe_no_prefill_cp(
+            graph_has_dp_gather=True
+        )
+        with flags, a2a, dsa_cp, mla_cp:
+            self.assertFalse(prefill_graph_tolerates_sum_len())
+            self.assertTrue(
+                runner._has_inactive_dp_rank(
+                    SimpleNamespace(global_num_tokens_cpu=[8, 0])
+                )
+            )
+
 
 if __name__ == "__main__":
     unittest.main()