From a9bfa08e6692cf5ef691651a58b382415cbcb07b Mon Sep 17 00:00:00 2001
From: hnyls2002 <lsyincs@gmail.com>
Date: Fri, 23 Oct 2026 21:14:25 -0700
Subject: [PATCH] revert megagate routing in prefill; prefill graph capture
 fails with it

---
 python/sglang/srt/layers/moe/mega_gate.py     | 10 ++++----
 .../layers/moe/test_mega_moe_deepgemm_api.py  | 25 +++++++++++--------
 2 files changed, 20 insertions(+), 15 deletions(-)

diff --git a/python/sglang/srt/layers/moe/mega_gate.py b/python/sglang/srt/layers/moe/mega_gate.py
index 3eb31f941f93..0376a93b358e 100644
--- a/python/sglang/srt/layers/moe/mega_gate.py
+++ b/python/sglang/srt/layers/moe/mega_gate.py
@@ -19,6 +19,8 @@
 from sglang.srt.runtime_context import get_exec, get_forward, get_parallel
 from sglang.srt.state_capturer.routed_experts import get_global_experts_capturer
 
+# DeepGEMM bf16_mega_gate scores one token block per launch.
+_MEGA_GATE_MAX_ROWS = 2048
 # Per-forward cache of the gate's row masks, installed by mega_gate_forward_scope.
 _ROW_MASKS_ATTR = "_mega_gate_row_masks"
 
@@ -82,14 +84,12 @@ def mega_gate_supported(moe) -> bool:
 def should_use_mega_gate(moe, hidden_states: torch.Tensor, forward_batch) -> bool:
     if not envs.SGLANG_OPT_DEEPGEMM_MEGA_GATE.get() or forward_batch is None:
         return False
-    # Prefill, decode and verify; draft extend keeps the regular router.
+    # Prefill graph capture clears is_extend_in_batch, so gate on the mode itself.
     mode = forward_batch.forward_mode
-    # The routing and the Mega-mHC input go to MegaMoE's buffer, sized to its token cap.
-    max_rows = envs.SGLANG_OPT_DEEPGEMM_MEGA_MOE_NUM_MAX_TOKENS_PER_RANK.get()
     return (
-        (mode.is_decode() or mode.is_extend())
+        (mode.is_decode() or mode.is_target_verify())
         and moe.mega_gate_supported
-        and 0 < hidden_states.shape[0] <= max_rows
+        and 0 < hidden_states.shape[0] <= _MEGA_GATE_MAX_ROWS
         and not get_forward().sp_active
         and hidden_states.dtype == torch.bfloat16
         and hidden_states.is_contiguous()
diff --git a/test/registered/unit/layers/moe/test_mega_moe_deepgemm_api.py b/test/registered/unit/layers/moe/test_mega_moe_deepgemm_api.py
index 4b5be67b73b3..7364d8ffddcb 100644
--- a/test/registered/unit/layers/moe/test_mega_moe_deepgemm_api.py
+++ b/test/registered/unit/layers/moe/test_mega_moe_deepgemm_api.py
@@ -30,6 +30,7 @@
     get_context,
     get_exec,
     get_flags,
+    get_forward,
     get_model,
     get_parallel,
 )
@@ -554,22 +555,26 @@ def test_mega_gate_admission_falls_back_for_unsupported_routers(self):
                 self.assertFalse(use(moe, hidden, decode))
             with envs.SGLANG_OPT_DEEPGEMM_MEGA_GATE.override(True):
                 self.assertTrue(use(moe, hidden, decode))
-                for mode in (ForwardMode.TARGET_VERIFY, ForwardMode.EXTEND):
-                    self.assertTrue(
-                        use(moe, hidden, SimpleNamespace(forward_mode=mode))
-                    )
-                self.assertFalse(
+                self.assertTrue(
                     use(
                         moe,
                         hidden,
-                        SimpleNamespace(forward_mode=ForwardMode.DRAFT_EXTEND_V2),
+                        SimpleNamespace(forward_mode=ForwardMode.TARGET_VERIFY),
                     )
                 )
                 self.assertFalse(use(moe, hidden.float(), decode))
-                with envs.SGLANG_OPT_DEEPGEMM_MEGA_MOE_NUM_MAX_TOKENS_PER_RANK.override(
-                    16
-                ):
-                    self.assertFalse(use(moe, hidden, decode))
+                self.assertFalse(
+                    use(moe, torch.zeros((4096, 256), dtype=torch.bfloat16), decode)
+                )
+                # Prefill graph capture runs with is_extend_in_batch cleared.
+                with get_forward().scoped(is_extend_in_batch=False):
+                    self.assertFalse(
+                        use(
+                            moe,
+                            hidden,
+                            SimpleNamespace(forward_mode=ForwardMode.EXTEND),
+                        )
+                    )
                 self.assertFalse(use(moe, hidden, None))
 
     def test_mega_gate_routes_vision_hash_and_padding(self):