[refactor] Collect MoE and DP-attention runtime state into typed flag groups (#30347)

test/registered/ops/test_aiter_allreduce_fusion_amd.py

@@ -420,9 +420,9 @@ class TestAiterAllreduceFusionGate(CustomTestCase):
             stack.enter_context(
                 mock.patch.object(comm, "get_global_server_args", lambda: server_args)
             )
-            stack.enter_context(
-                mock.patch.object(comm, "is_dp_attention_enabled", lambda: dp_attention)
-            )
+            from sglang.srt.runtime_context import get_flags
+
+            stack.enter_context(get_flags().dp.override(enabled=dp_attention))
             stack.enter_context(
                 mock.patch.object(comm, "get_moe_a2a_backend", lambda: a2a_backend)
             )

test/registered/ops/test_aiter_greedy_sample_amd.py

@@ -38,7 +38,9 @@ def _mock_global_server_args(backend="pytorch"):
         device_group = None
 
     sampler_mod.get_tp_group = lambda: _DummyTPGroup()
-    sampler_mod.is_dp_attention_enabled = lambda: False
+    from sglang.srt.runtime_context import get_flags
+
+    get_flags().dp.enabled = False
 
 
 def _make_sampling_info(batch_size, vocab_size, device="cuda"):

test/registered/unit/test_module_state_ratchet.py

@@ -0,0 +1,72 @@
+"""Ratchet guard: module-level runtime state in the flag-owning layers may
+only shrink.
+
+Runtime flags live on ``get_flags()`` groups (``moe`` / ``dp`` / ``capture``),
+where they get lifecycle reset, typo-safe writes, and the transactional
+test-override primitive. A new module-level global written through a
+``global`` statement in these modules recreates the pattern this replaced:
+state with ad-hoc lifecycle that leaks across unit-test teardowns and cannot
+be overridden scoped.
+
+The pin lists the survivors by name: the DP-attention topology values (owned
+by the parallel vertical) and the TBO comm stream (a resource, owned by the
+resources vertical). Migrating one of them must shrink its pin; adding a name
+fails the ratchet.
+"""
+
+from sglang.test.ci.ci_register import register_cpu_ci
+
+register_cpu_ci(est_time=5, suite="base-a-test-cpu")
+
+import ast
+import unittest
+from pathlib import Path
+
+import sglang.srt
+from sglang.test.test_utils import CustomTestCase
+
+_SRT_ROOT = Path(next(iter(sglang.srt.__path__)))
+
+_PINNED_GLOBALS = {
+    "layers/moe/utils.py": frozenset(),
+    "layers/dp_attention.py": frozenset(
+        {
+            # DP-attention topology (parallel vertical scope).
+            "_ATTN_DP_RANK",
+            "_ATTN_DP_SIZE",
+            "_LOCAL_ATTN_DP_SIZE",
+            "_LOCAL_ATTN_DP_RANK",
+            # Comm stream resource (resources vertical scope).
+            "_DP_TBO_COMM_STREAM",
+        }
+    ),
+}
+
+
+class TestModuleStateRatchet(CustomTestCase):
+    def test_global_statements_match_the_pins(self):
+        for rel, pinned in _PINNED_GLOBALS.items():
+            tree = ast.parse((_SRT_ROOT / rel).read_text())
+            declared = {
+                name
+                for node in ast.walk(tree)
+                if isinstance(node, ast.Global)
+                for name in node.names
+            }
+            grown = declared - pinned
+            self.assertFalse(
+                grown,
+                f"{rel} declares new module-level runtime state {sorted(grown)}; "
+                "put runtime flags on a get_flags() group instead "
+                "(see runtime_context.MoeFlags / DpFlags).",
+            )
+            shrunk = pinned - declared
+            self.assertFalse(
+                shrunk,
+                f"{rel} no longer declares {sorted(shrunk)}; "
+                "shrink the pin in this file to lock in the progress.",
+            )
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/unit/test_runtime_context.py

@@ -265,6 +265,108 @@ class _FakeResolvedArgs:
     _resolved_overrides: list = dataclasses.field(default_factory=list)
 
 
+class TestMoeFlagsGroup(_IsolatedServerArgs):
+    """flags.moe: materialized by initialize_moe_config; the ACTIVE backends
+    swap under the speculative contexts and restore on exit."""
+
+    def _init(self, **kw):
+        from types import SimpleNamespace
+
+        from sglang.srt.layers.moe.utils import initialize_moe_config
+
+        defaults = dict(
+            moe_a2a_backend="none",
+            moe_runner_backend="auto",
+            speculative_moe_runner_backend=None,
+            speculative_moe_a2a_backend=None,
+            deepep_mode="auto",
+            deepep_config=None,
+            enable_two_batch_overlap=False,
+            enable_single_batch_overlap=False,
+            tbo_token_distribution_threshold=0.48,
+            disable_flashinfer_cutlass_moe_fp4_allgather=False,
+            quantization=None,
+        )
+        defaults.update(kw)
+        initialize_moe_config(SimpleNamespace(**defaults))
+
+    def test_lazy_defaults_before_initialize(self):
+        from sglang.srt.layers.moe.utils import (
+            get_moe_a2a_backend,
+            get_moe_runner_backend,
+            is_tbo_enabled,
+        )
+
+        reset_context()
+        self.assertTrue(get_moe_a2a_backend().is_none())
+        self.assertEqual(get_moe_runner_backend().name, "AUTO")
+        self.assertFalse(is_tbo_enabled())
+
+    def test_initialize_materializes_group(self):
+        from sglang.srt.layers.moe.utils import get_moe_a2a_backend, is_tbo_enabled
+
+        self._init(moe_a2a_backend="deepep", enable_two_batch_overlap=True)
+        self.assertTrue(get_moe_a2a_backend().is_deepep())
+        self.assertTrue(is_tbo_enabled())
+        self.assertEqual(get_flags().moe.deepep_config, "")
+
+    def test_speculative_swap_and_restore(self):
+        from sglang.srt.layers.moe.utils import (
+            get_moe_a2a_backend,
+            get_moe_runner_backend,
+            speculative_moe_a2a_backend_context,
+            speculative_moe_backend_context,
+        )
+
+        self._init(
+            moe_a2a_backend="deepep",
+            moe_runner_backend="triton",
+            speculative_moe_runner_backend="auto",
+            speculative_moe_a2a_backend="none",
+        )
+        with speculative_moe_backend_context(), speculative_moe_a2a_backend_context():
+            self.assertEqual(get_moe_runner_backend().name, "AUTO")
+            self.assertTrue(get_moe_a2a_backend().is_none())
+            # MTP layers are unquantized: fp4 allgather is forced off
+            self.assertTrue(get_flags().moe.disable_fp4_allgather)
+        self.assertEqual(get_moe_runner_backend().name, "TRITON")
+        self.assertTrue(get_moe_a2a_backend().is_deepep())
+        self.assertFalse(get_flags().moe.disable_fp4_allgather)
+
+    def test_swap_restores_on_exception(self):
+        from sglang.srt.layers.moe.utils import (
+            get_moe_runner_backend,
+            speculative_moe_backend_context,
+        )
+
+        self._init(moe_runner_backend="triton", speculative_moe_runner_backend="auto")
+        with self.assertRaises(RuntimeError):
+            with speculative_moe_backend_context():
+                raise RuntimeError("boom")
+        self.assertEqual(get_moe_runner_backend().name, "TRITON")
+
+
+class TestDpFlagsGroup(_IsolatedServerArgs):
+    """flags.dp: the DP-attention runtime flags; is_dp_attention_enabled is a
+    thin shim over the group leaf."""
+
+    def test_shim_reads_the_leaf(self):
+        from sglang.srt.layers.dp_attention import is_dp_attention_enabled
+
+        reset_context()
+        self.assertFalse(is_dp_attention_enabled())
+        get_flags().dp.enabled = True
+        self.assertTrue(is_dp_attention_enabled())
+
+    def test_scoped_override_forces_the_predicate(self):
+        from sglang.srt.layers.dp_attention import is_dp_attention_enabled
+
+        reset_context()
+        with get_flags().dp.override(enabled=True):
+            self.assertTrue(is_dp_attention_enabled())
+        self.assertFalse(is_dp_attention_enabled())
+
+
 class TestPublishLifecycle(_IsolatedServerArgs):
     """Publish installs the resolved server_args and seeds the capture tier."""