[DeepSeek V4] Add W4A4 MegaMoE server flag (#35918)

test/registered/cp/test_deepseek_v4_flash_fp4_b200_cp.py

@@ -41,8 +41,6 @@ _DEEPEP_ENV = {
 _MEGAMOE_ENV = {
     "SGLANG_ENABLE_CP_V2": "1",
     "SGLANG_OPT_DEEPGEMM_MEGA_MOE_NUM_MAX_TOKENS_PER_RANK": "8320",
-    "SGLANG_OPT_DEEPGEMM_MEGA_MOE_USE_FP4_ACTS": "1",
-    "SGLANG_OPT_DEEPGEMM_MEGA_MOE_USE_MXF4_KIND": "1",
 }
 
 
@@ -126,6 +124,7 @@ class TestDSV4FlashFP4B200Balanced_CP_Megamoe(
                 "--enable-dp-attention",
                 "--moe-a2a-backend",
                 "megamoe",
+                "--enable-w4a4-mxfp4-megamoe",
                 "--speculative-algorithm",
                 "EAGLE",
                 "--speculative-num-steps",

test/registered/models_e2e/test_deepseek_v4_flash_fp4_megamoe_b200.py

@@ -34,8 +34,6 @@ _W4A8_MEGAMOE_ENV = {
 
 _W4A4_MEGAMOE_ENV = {
     "SGLANG_OPT_DEEPGEMM_MEGA_MOE_NUM_MAX_TOKENS_PER_RANK": "4096",
-    "SGLANG_OPT_DEEPGEMM_MEGA_MOE_USE_FP4_ACTS": "1",
-    "SGLANG_OPT_DEEPGEMM_MEGA_MOE_USE_MXF4_KIND": "1",
 }
 
 
@@ -115,6 +113,7 @@ class TestDSV4FlashFP4B200W4A4MegaMoE(
                 "--enable-dp-attention",
                 "--moe-a2a-backend",
                 "megamoe",
+                "--enable-w4a4-mxfp4-megamoe",
                 "--speculative-algorithm",
                 "EAGLE",
                 "--speculative-num-steps",

test/registered/unit/server_args/test_server_args.py

@@ -43,6 +43,38 @@ _mock_device.start()
 
 
 class TestPrepareServerArgs(CustomTestCase):
+    def test_enable_w4a4_mxfp4_megamoe_sets_deepgemm_env(self):
+        deepgemm_env = {
+            "DG_USE_FP4_ACTS": "0",
+            "DG_USE_MXF4_KIND": "0",
+        }
+        with patch.dict(os.environ, deepgemm_env, clear=False):
+            try:
+                args = prepare_server_args(
+                    ["--model-path", "dummy", "--enable-w4a4-mxfp4-megamoe"]
+                )
+            except SystemExit as exc:
+                self.fail(
+                    "--enable-w4a4-mxfp4-megamoe must be accepted by the CLI "
+                    f"parser, got SystemExit({exc.code})"
+                )
+
+            self.assertTrue(args.enable_w4a4_mxfp4_megamoe)
+            self.assertEqual(os.environ["DG_USE_FP4_ACTS"], "1")
+            self.assertEqual(os.environ["DG_USE_MXF4_KIND"], "1")
+
+    def test_w4a4_mxfp4_megamoe_disabled_preserves_deepgemm_env(self):
+        deepgemm_env = {
+            "DG_USE_FP4_ACTS": "0",
+            "DG_USE_MXF4_KIND": "0",
+        }
+        with patch.dict(os.environ, deepgemm_env, clear=False):
+            args = prepare_server_args(["--model-path", "dummy"])
+
+            self.assertFalse(args.enable_w4a4_mxfp4_megamoe)
+            self.assertEqual(os.environ["DG_USE_FP4_ACTS"], "0")
+            self.assertEqual(os.environ["DG_USE_MXF4_KIND"], "0")
+
     def test_prefill_decode_interval(self):
         args = ServerArgs(model_path="dummy", prefill_decode_interval=16)
         self.assertEqual(args.prefill_decode_interval, 16)

test/registered/unit/test_environ.py

@@ -114,6 +114,21 @@ class TestDeprecatedEnvRegistry(unittest.TestCase):
         caught = self._apply(old_name, _DeprecatedEnv())
         self.assertIn(f"{old_name} is deprecated", str(caught[0].message))
 
+    def test_w4a4_mxfp4_megamoe_envs_warn_to_use_cli_flag(self):
+        old_names = (
+            "SGLANG_OPT_DEEPGEMM_MEGA_MOE_USE_FP4_ACTS",
+            "SGLANG_OPT_DEEPGEMM_MEGA_MOE_USE_MXF4_KIND",
+        )
+        for old_name in old_names:
+            with self.subTest(old_name=old_name):
+                os.environ[old_name] = "1"
+                self.addCleanup(os.environ.pop, old_name, None)
+
+                caught = self._apply(old_name, _DEPRECATED_ENVS[old_name])
+
+                self.assertIn("--enable-w4a4-mxfp4-megamoe", str(caught[0].message))
+                self.assertIsNone(_DEPRECATED_ENVS[old_name].replacement)
+
     def test_renamed_env_forwards_value(self):
         old_name, new_name = "SGLANG_TEST_OLD_ENV", "SGLANG_TEST_NEW_ENV"
         os.environ[old_name] = "abc"