[Fix] dsv41: Revert MegaGate routing in prefill

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()

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):