[NVIDIA][MegaMoE] Fix the default spec decoding moe backend when the target model uses flashinfer megamoe

python/sglang/srt/arg_groups/fields/spec.py

@@ -172,15 +172,15 @@ class Spec(msgspec.Struct):
     speculative_moe_runner_backend: A[
         Optional[str],
         Arg(
-            help="Choose the runner backend for MoE in speculative decoding.",
+            help="Choose the runner backend for MoE in speculative decoding. Defaults to --moe-runner-backend, except flashinfer_megamoe defaults the draft to flashinfer_trtllm (draft layers are typically unquantized, which flashinfer_megamoe's global-expert-id dispatch doesn't support).",
             choices=MOE_RUNNER_BACKEND_CHOICES,
             resolvable=True,
         ),
     ] = None
     speculative_moe_a2a_backend: A[
         Optional[str],
         Arg(
-            help="Choose the backend for MoE A2A in speculative decoding",
+            help="Choose the backend for MoE A2A in speculative decoding. Defaults to --moe-a2a-backend, except flashinfer_megamoe defaults the draft to none (draft layers are typically unquantized, which flashinfer_megamoe's global-expert-id dispatch doesn't support).",
             choices=[
                 "none",
                 "deepep",

python/sglang/srt/arg_groups/overrides.py

@@ -1584,9 +1584,22 @@ def _pipeline_parallel_overlap_disable(view: Any) -> dict:
 def _speculative_moe_runner_default(view: Any) -> dict:
     """Default the speculative (draft) MoE runner backend to the resolved
     target-model backend. Invoked at the head of the speculative-decoding
-    hook, after the MoE kernel chain has resolved."""
+    hook, after the MoE kernel chain has resolved.
+
+    flashinfer_megamoe routes by global expert id inside the mega kernel and
+    only supports the quantized MoE layers it prepared weights for. Draft
+    (MTP/NextN) layers are typically unquantized and run through a regular
+    runner, so inheriting it would mismatch the draft's runner and corrupt
+    expert-id dispatch. Default the draft to flashinfer_trtllm instead; users
+    with a quantized draft can set --speculative-moe-runner-backend explicitly.
+    (DeepGEMM megamoe is selected via the a2a backend and falls back per-layer
+    when its weights are not prepared, so it needs no override.)
+    """
     if view.speculative_moe_runner_backend is None:
-        return {"speculative_moe_runner_backend": view.moe_runner_backend}
+        default = view.moe_runner_backend
+        if default == "flashinfer_megamoe":
+            default = "flashinfer_trtllm"
+        return {"speculative_moe_runner_backend": default}
     return {}
 
 

python/sglang/srt/layers/moe/utils.py

@@ -443,6 +443,27 @@ def get_deepep_v2_fp8_scale_format() -> DeepEPv2Fp8ScaleFormat:
     )
 
 
+# Only flashinfer_megamoe is special-cased: its global-expert-id dispatch
+# assumes quantized layers, which draft (MTP/NextN) layers typically are not.
+# DeepGEMM megamoe (a2a "megamoe") falls back per-layer when its weights
+# are not prepared, so drafts can inherit it. A quantized draft that wants
+# flashinfer_megamoe must set the --speculative-moe-* flags explicitly.
+def _default_speculative_a2a_backend(
+    target_a2a_backend: MoeA2ABackend,
+) -> MoeA2ABackend:
+    if target_a2a_backend.is_flashinfer_megamoe():
+        return MoeA2ABackend.NONE
+    return target_a2a_backend
+
+
+def _default_speculative_moe_runner_backend(
+    target_runner_backend: MoeRunnerBackendLike,
+) -> MoeRunnerBackendLike:
+    if target_runner_backend.is_flashinfer_megamoe():
+        return MoeRunnerBackend.FLASHINFER_TRTLLM
+    return target_runner_backend
+
+
 def initialize_moe_config():
     """Seed the MoE runtime flags from the published configuration.
 
@@ -460,12 +481,12 @@ def initialize_moe_config():
     moe.speculative_runner_backend = (
         resolve_moe_runner_backend(spec.speculative_moe_runner_backend)
         if spec.speculative_moe_runner_backend is not None
-        else moe.runner_backend
+        else _default_speculative_moe_runner_backend(moe.runner_backend)
     )
     moe.speculative_a2a_backend = (
         MoeA2ABackend(spec.speculative_moe_a2a_backend)
         if spec.speculative_moe_a2a_backend is not None
-        else moe.a2a_backend
+        else _default_speculative_a2a_backend(moe.a2a_backend)
     )
     moe.deepep_mode = DeepEPMode(exec_moe.deepep_mode)
     moe.deepep_config = exec_moe.deepep_config or ""

test/registered/unit/test_model_overrides.py

@@ -2318,6 +2318,17 @@ def test_speculative_moe_runner_default_pass(self):
             ),
             {"speculative_moe_runner_backend": "triton"},
         )
+        self.assertEqual(
+            _speculative_moe_runner_default(
+                ResolvedView(
+                    SimpleNamespace(
+                        speculative_moe_runner_backend=None,
+                        moe_runner_backend="flashinfer_megamoe",
+                    )
+                )
+            ),
+            {"speculative_moe_runner_backend": "flashinfer_trtllm"},
+        )
         # user-set draft backend survives
         self.assertEqual(
             _speculative_moe_runner_default(
@@ -2735,7 +2746,7 @@ def _view(**kw):
                 self.assertEqual(
                     _deepseek_spec_moe_resolution(
                         _view(
-                            speculative_moe_runner_backend="triton",
+                            speculative_moe_runner_backend="flashinfer_trtllm",
                             speculative_moe_a2a_backend="none",
                         )
                     ),

test/registered/unit/test_runtime_context.py

@@ -1161,6 +1161,26 @@ def test_speculative_swap_and_restore(self):
         self.assertFalse(get_flags().moe.disable_fp4_allgather)
         self.assertFalse(get_flags().moe.speculative_context)
 
+    def test_flashinfer_megamoe_speculative_defaults(self):
+        self._init(
+            moe_a2a_backend="flashinfer_megamoe",
+            moe_runner_backend="flashinfer_megamoe",
+        )
+        moe = get_flags().moe
+        self.assertEqual(moe.speculative_runner_backend.value, "flashinfer_trtllm")
+        self.assertTrue(moe.speculative_a2a_backend.is_none())
+
+    def test_flashinfer_megamoe_speculative_explicit_wins(self):
+        self._init(
+            moe_a2a_backend="flashinfer_megamoe",
+            moe_runner_backend="flashinfer_megamoe",
+            speculative_moe_runner_backend="triton",
+            speculative_moe_a2a_backend="deepep",
+        )
+        moe = get_flags().moe
+        self.assertEqual(moe.speculative_runner_backend.value, "triton")
+        self.assertTrue(moe.speculative_a2a_backend.is_deepep())
+
     def test_swap_restores_on_exception(self):
         from sglang.srt.layers.moe.utils import (
             get_moe_runner_backend,