config: delete the redundant full stamp in initialize_model_parallel (#39202)

test/manual/ep/test_flashinfer_dispatcher.py

@@ -20,14 +20,9 @@ from sglang.test.test_utils import CustomTestCase
 class TestFlashinferDispatcher(CustomTestCase):
     @classmethod
     def setUpClass(cls):
-        server_args = ServerArgs(model_path="dummy")
-        server_args.moe_runner_backend = "flashinfer_cutlass"
-        server_args.moe_a2a_backend = "flashinfer"
-        cls.server_args = server_args
-        set_global_server_args_for_scheduler(server_args)
-        publish(server_args, role="scheduler")
-        initialize_moe_config()
-
+        # Dist-init first: world_size (and so the tp/ep width ServerArgs must
+        # carry) is only known after it, and init_distributed_environment
+        # itself reads no published config.
         init_distributed_environment(
             world_size=-1,  # Auto-detect from environment
             rank=-1,  # Auto-detect from environment
@@ -38,6 +33,17 @@ class TestFlashinferDispatcher(CustomTestCase):
         rank = torch.distributed.get_rank()
         device = torch.device(f"cuda:{rank % torch.cuda.device_count()}")
         torch.cuda.set_device(device)
+
+        server_args = ServerArgs(
+            model_path="dummy", tp_size=world_size, ep_size=world_size
+        )
+        server_args.moe_runner_backend = "flashinfer_cutlass"
+        server_args.moe_a2a_backend = "flashinfer"
+        cls.server_args = server_args
+        set_global_server_args_for_scheduler(server_args)
+        publish(server_args, role="scheduler")
+        initialize_moe_config()
+
         initialize_model_parallel(
             tensor_model_parallel_size=world_size, expert_model_parallel_size=world_size
         )

test/registered/eplb/test_lplb_distributed.py

@@ -212,6 +212,13 @@ def _worker_main(local_rank: int, world_size: int):
     set_global_server_args_for_scheduler(
         ServerArgs(
             model_path="dummy",
+            # Match the tp/ep width initialize_model_parallel is about to
+            # build below -- get_parallel()'s derived widths (attn_tp_size,
+            # moe_ep_size, ...) are projected from this at publish time, and
+            # nothing here should leave that projection reflecting a width
+            # this process never actually runs at.
+            tp_size=world_size,
+            ep_size=world_size,
         )
     )
 

test/registered/unit/distributed/test_mmencoder_forwards_attn_cp_size.py

@@ -0,0 +1,67 @@
+"""MMEncoder must forward attn_cp_size to initialize_model_parallel, not just
+tp_size.
+
+Real 2-GPU hardware confirmed a live mismatch this test guards against
+statically: with `tp_size=2, attn_cp_size=2` published, calling
+`initialize_model_parallel(tensor_model_parallel_size=2)` alone builds the
+live attention-TP group at width 2, while `get_parallel().attn_tp_size`
+(derived from the published config) answers 1 -- `VisionAttention`
+(`layers/attention/vision.py`) reads that derived value as its own
+weight-sharding width, so the mismatch is a real, silent wrong-sharding bug,
+not just a reporting discrepancy. Forwarding
+`attention_context_model_parallel_size=get_parallel().attn_cp_size` too
+makes the two agree (confirmed on the same hardware).
+
+A full `MMEncoder` instantiation needs real weights and a live process
+group, so this checks the one line that matters statically: the call passes
+`attention_context_model_parallel_size` as well as
+`tensor_model_parallel_size`. Not a substitute for testing `MMEncoder`
+end-to-end under `--attn-cp-size > 1` on real hardware, but cheap enough to
+run everywhere and catches the specific regression class (a future edit
+that reverts to the tp-only call).
+"""
+
+import ast
+import os
+
+import sglang.srt.disaggregation.encoder.server as encoder_server_module
+from sglang.test.ci.ci_register import register_cpu_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_cpu_ci(est_time=5, suite="base-a-test-cpu")
+
+
+class TestMMEncoderForwardsAttnCpSize(CustomTestCase):
+    def test_initialize_model_parallel_call_forwards_attn_cp_size(self):
+        path = encoder_server_module.__file__
+        tree = ast.parse(open(path).read())
+        calls = [
+            node
+            for node in ast.walk(tree)
+            if isinstance(node, ast.Call)
+            and isinstance(node.func, ast.Name)
+            and node.func.id == "initialize_model_parallel"
+        ]
+        self.assertEqual(
+            len(calls),
+            1,
+            f"expected exactly one initialize_model_parallel(...) call in "
+            f"{os.path.basename(path)}, found {len(calls)} -- update this "
+            "test if that's now intentional",
+        )
+        kwarg_names = {kw.arg for kw in calls[0].keywords}
+        self.assertIn(
+            "attention_context_model_parallel_size",
+            kwarg_names,
+            "MMEncoder's initialize_model_parallel(...) call must forward "
+            "attention_context_model_parallel_size (not just "
+            "tensor_model_parallel_size), or get_parallel().attn_tp_size "
+            "silently disagrees with the group actually built whenever "
+            "--attn-cp-size > 1 -- confirmed on real 2-GPU hardware",
+        )
+
+
+if __name__ == "__main__":
+    import unittest
+
+    unittest.main()

test/registered/unit/test_runtime_context.py

@@ -1375,7 +1375,8 @@ class TestParallelLeafReads(_IsolatedServerArgs):
 
 
 class TestDerivedWidths(_IsolatedOverrides):
-    """The widths no flag sets are computed from the leaves and stamped.
+    """The widths no flag sets are computed from the leaves and permanently
+    overridable.
 
     `attn_tp_size` and its siblings used to be read back off the group
     coordinator that was built from them, which made the answer depend on
@@ -1390,7 +1391,7 @@ class TestDerivedWidths(_IsolatedOverrides):
         self.addCleanup(
             lambda: (
                 parallel.clear_derived_widths(),
-                parallel.stamp_derived_widths(**self._saved_derived),
+                parallel.override_permanently(**self._saved_derived),
             )
         )
 
@@ -1439,13 +1440,14 @@ class TestDerivedWidths(_IsolatedOverrides):
             get_parallel().attn_tp_size
         self.assertIn("not available", str(caught.exception))
 
-    def test_a_stamp_and_a_live_group_both_win_over_the_leaves(self):
-        """Order is stamp, then live group, then the leaves. Where a group
-        exists it is the truth -- elastic scale-up moves the group without
-        restamping -- so the leaf derivation only answers where there is none.
+    def test_a_permanent_override_and_a_live_group_both_win_over_the_leaves(self):
+        """Order is permanent override, then live group, then the leaves.
+        Where a group exists it is the truth -- elastic scale-up moves the
+        group without a fresh override -- so the leaf derivation only
+        answers where there is none.
         """
         parallel = get_parallel()
-        parallel.stamp_derived_widths(attn_tp_size=7)
+        parallel.override_permanently(attn_tp_size=7)
         self.addCleanup(parallel.clear_derived_widths)
         with parallel.override(tp_size=8, attn_dp_size=2):
             self.assertEqual(parallel.attn_tp_size, 7)
@@ -1464,9 +1466,9 @@ class TestDerivedWidths(_IsolatedOverrides):
         self.assertEqual(widths["moe_tp_size"], 8 // 4 // 2)
         self.assertEqual(widths["attn_dcp_size"], 1)
 
-    def test_the_world_size_is_not_stamped(self):
+    def test_the_world_size_is_not_permanently_overridden(self):
         """It is not a quotient, and the live getter is right at every moment.
-        A stamp taken when the groups are built would answer with the launch
+        A value fixed when the groups are built would answer with the launch
         count after `try_admit_scale_ranks` expands WORLD, and with the joining
         cohort's own width on a scale-joiner, which lays its groups out at
         `tp * pp` while WORLD spans `ep_join_rank_offset + tp * pp`."""
@@ -1481,31 +1483,30 @@ class TestDerivedWidths(_IsolatedOverrides):
         )
         self.assertNotIn("world_size", widths)
         parallel = get_parallel()
-        parallel.stamp_derived_widths(attn_tp_size=4)
+        parallel.override_permanently(attn_tp_size=4)
         with patch(f"{_PS}.get_world_size", return_value=9):
             self.assertEqual(parallel.world_size, 9)
 
-    def test_a_stamped_width_is_what_the_reader_answers_with(self):
+    def test_a_permanently_overridden_width_is_what_the_reader_answers_with(self):
         parallel = get_parallel()
-        parallel.stamp_derived_widths(attn_tp_size=4, moe_tp_size=1)
+        parallel.override_permanently(attn_tp_size=4, moe_tp_size=1)
         with patch(
             f"{_PS}.get_attn_tensor_model_parallel_world_size",
             side_effect=AssertionError("the group must not be asked"),
         ):
             self.assertEqual(parallel.attn_tp_size, 4)
 
-    def test_an_override_still_wins_over_the_stamp(self):
+    def test_a_scoped_override_still_wins_over_the_permanent_one(self):
         parallel = get_parallel()
-        parallel.stamp_derived_widths(attn_tp_size=4)
+        parallel.override_permanently(attn_tp_size=4)
         with parallel.override(attn_tp_size=1):
             self.assertEqual(parallel.attn_tp_size, 1)
         self.assertEqual(parallel.attn_tp_size, 4)
 
     def test_the_group_is_never_consulted(self):
-        """There is no third source. A quotient comes from an override, a stamp
-        or the published leaf -- never from a group coordinator, which could
-        only ever agree, since `initialize_model_parallel` stamps as its last
-        statement."""
+        """There is no third source. A quotient comes from a scoped override, a
+        permanent override, or the published leaf -- never from a group
+        coordinator."""
         reset_context()
         self.addCleanup(reset_context)
         with patch(
@@ -1528,50 +1529,51 @@ class TestDerivedWidths(_IsolatedOverrides):
             with self.assertRaisesRegex(RuntimeError, r"derived parallel width"):
                 get_parallel().attn_tp_size
 
-    def test_a_temporary_disable_beats_the_stamp(self):
+    def test_a_temporary_disable_beats_the_permanent_override(self):
         """`disable_dp_size()` runs a draft scope without DP attention. It moves
         the module global the legacy getter reads, so it has to move the derived
-        width too -- the stamp wins over the live group, and a scope that left
-        it alone would answer with the target model's width for its duration."""
+        width too -- the scoped override wins over the permanent one, and a
+        scope that left it alone would answer with the target model's width
+        for its duration."""
         from sglang.srt.layers import dp_attention
 
         parallel = get_parallel()
-        parallel.stamp_derived_widths(attn_dp_size=4)
+        parallel.override_permanently(attn_dp_size=4)
         with patch.object(dp_attention, "_ATTN_DP_SIZE", 4):
             with dp_attention.disable_dp_size():
                 self.assertEqual(dp_attention.get_attention_dp_size(), 1)
                 self.assertEqual(parallel.attn_dp_size, 1)
             self.assertEqual(parallel.attn_dp_size, 4)
 
-    def test_the_stamp_is_cleared_and_restamped(self):
+    def test_the_permanent_override_is_cleared_and_reset(self):
         parallel = get_parallel()
-        parallel.stamp_derived_widths(attn_dp_size=2)
+        parallel.override_permanently(attn_dp_size=2)
         self.assertEqual(parallel.attn_dp_size, 2)
-        # Elastic scaling restamps where it updates the live width.
-        parallel.stamp_derived_widths(attn_dp_size=4)
+        # Elastic scaling overrides again where it updates the live width.
+        parallel.override_permanently(attn_dp_size=4)
         self.assertEqual(parallel.attn_dp_size, 4)
         parallel.clear_derived_widths()
         with parallel.override(tp_size=8, attn_dp_size=1):
             self.assertEqual(parallel.attn_dp_size, 1)
 
-    def test_reset_context_drops_the_stamp(self):
-        """The stamp belongs to the lifecycle that made it.
+    def test_reset_context_drops_the_permanent_override(self):
+        """The permanent override belongs to the lifecycle that made it.
 
-        `_derived_width` prefers the stamp over the published leaf, so a stamp
-        that outlived `reset_context()` would let the next test read the
-        previous topology.
+        `_derived_width` prefers it over the published leaf, so one that
+        outlived `reset_context()` would let the next test read the previous
+        topology.
         """
         parallel = get_parallel()
-        parallel.stamp_derived_widths(attn_tp_size=4)
+        parallel.override_permanently(attn_tp_size=4)
         self.assertEqual(parallel.attn_tp_size, 4)
         reset_context()
         self.addCleanup(reset_context)
         publish(ServerArgs(model_path="dummy", tp_size=1), role="test")
         self.assertEqual(get_parallel().attn_tp_size, 1)
 
-    def test_the_rank_helper_agrees_with_the_stamp(self):
+    def test_the_rank_helper_agrees_with_the_override(self):
         """`compute_dp_attention_world_info` keeps the ranks and takes the
-        widths from the same derivation the stamp uses."""
+        widths from the same derivation `override_permanently`'s callers use."""
         from sglang.srt.layers.dp_attention import compute_dp_attention_world_info
 
         for tp_size, dp_size, attn_cp_size in ((8, 2, 1), (8, 2, 2), (16, 4, 2)):
@@ -1590,6 +1592,126 @@ class TestDerivedWidths(_IsolatedOverrides):
             self.assertEqual(attn_tp_size, widths["attn_tp_size"])
             self.assertEqual(attn_dp_size, widths["attn_dp_size"])
 
+    def test_recomputing_from_published_leaves_matches_the_publish_bag(self):
+        """`initialize_model_parallel` no longer overrides anything -- see
+        `test_initialize_model_parallel_no_longer_touches_the_bag` below --
+        which makes this the load-bearing half of 16-field-registry-design.md
+        §6e: every real caller must forward leaves that already match its own
+        published config, because nothing corrects a mismatch anymore.
+        `scheduler.py`'s `ps.attn_dp_size`/`ps.moe_ep_size`/etc, and the
+        weight-cache daemon's own already-published config, both do -- this
+        pins that the formula they'd recompute from those leaves
+        (`derive_attention_widths`, `derive_parallel_widths`, the same ones
+        `publish` itself used) agrees with what's already in the bag, across
+        the widths `test_the_rank_helper_agrees_with_the_override` does not
+        vary -- moe_ep_size, moe_dp_size, and dcp_size -- using real
+        `publish()`.
+
+        A caller that does NOT keep the two in sync is a bug in that caller,
+        not something this framework silently corrects: two real ones existed
+        (`test/registered/eplb/test_lplb_distributed.py` and
+        `test/manual/ep/test_flashinfer_dispatcher.py`, both publishing a
+        placeholder config and then building real groups at a width it never
+        reflected) and were fixed by publishing the actual width instead of
+        relying on a correction to paper over the mismatch.
+        """
+        shapes = (
+            dict(tp_size=8),
+            dict(tp_size=8, dp_size=2, enable_dp_attention=True),
+            dict(tp_size=8, ep_size=4, moe_dp_size=2),
+            dict(tp_size=8, dcp_size=8),
+        )
+        for shape in shapes:
+            with self.subTest(shape=shape):
+                reset_context()
+                self.addCleanup(reset_context)
+                publish(ServerArgs(model_path="dummy", **shape), role="test")
+                parallel = get_parallel()
+                published = {
+                    "attn_tp_size": parallel.attn_tp_size,
+                    "attn_dp_size": parallel.attn_dp_size,
+                    "moe_ep_size": parallel.moe_ep_size,
+                    "moe_tp_size": parallel.moe_tp_size,
+                    "dcp_enabled": parallel.dcp_enabled,
+                    "attn_dcp_size": parallel.attn_dcp_size,
+                }
+                # What every real `initialize_model_parallel` caller forwards:
+                # its own already-published leaves, through the same two
+                # functions the bag was projected with.
+                recomputed = derive_parallel_widths(
+                    tp_size=parallel.tp_size,
+                    attn_cp_size=parallel.attn_cp_size,
+                    attn_dp_size=(
+                        parallel.dp_size if parallel.enable_dp_attention else 1
+                    ),
+                    moe_ep_size=parallel.ep_size,
+                    moe_dp_size=parallel.moe_dp_size,
+                    dcp_size=parallel.dcp_size,
+                    dcp_enabled=parallel.dcp_size > 1,
+                )
+                self.assertEqual(published, recomputed)
+
+    def test_initialize_model_parallel_no_longer_touches_the_bag(self):
+        """§6e, landed: `initialize_model_parallel` used to recompute and
+        permanently override the six derived widths on `get_parallel()`
+        after building its groups; that call is gone. Publish a placeholder
+        config (tp_size defaults to 1), then build real groups at a
+        different width -- the published leaf must now stay exactly what it
+        was, because nothing corrects it. This is the behavior a caller
+        relies on being told about, loudly, the first time it publishes and
+        builds inconsistently -- see
+        `test_recomputing_from_published_leaves_matches_the_publish_bag`
+        for why every real caller must not do that.
+        """
+        from unittest.mock import Mock
+
+        from sglang.srt.distributed import parallel_state
+
+        reset_context()
+        self.addCleanup(reset_context)
+        publish(ServerArgs(model_path="dummy"), role="test")
+        self.assertEqual(get_parallel().attn_tp_size, 1)
+        self.assertEqual(get_parallel().moe_ep_size, 1)
+
+        world_size = 8
+        with (
+            patch.object(parallel_state, "_WORLD", None),
+            patch.object(parallel_state, "_TP", None),
+            patch.object(parallel_state, "_DCP", None),
+            patch.object(parallel_state, "_ATTN_CP", None),
+            patch.object(parallel_state, "_ATTN_TP", None),
+            patch.object(parallel_state, "_MOE_DP", None),
+            patch.object(parallel_state, "_MOE_EP", None),
+            patch.object(parallel_state, "_MOE_TP", None),
+            patch.object(parallel_state, "_PP", None),
+            patch.object(parallel_state, "_SELF_PP", None),
+            patch("torch.distributed.is_initialized", return_value=True),
+            patch("torch.distributed.get_world_size", return_value=world_size),
+            patch("torch.distributed.get_rank", return_value=0),
+            patch("torch.distributed.get_backend", return_value="nccl"),
+            patch.object(
+                parallel_state,
+                "init_model_parallel_group",
+                return_value=Mock(device_group=Mock()),
+            ),
+            patch.object(parallel_state, "get_world_group") as mock_world_group,
+        ):
+            mock_world_group.return_value = Mock(device_group=Mock(), local_rank=0)
+            parallel_state.initialize_model_parallel(
+                tensor_model_parallel_size=world_size,
+                expert_model_parallel_size=world_size,
+            )
+        self.addCleanup(parallel_state.destroy_model_parallel)
+
+        self.assertEqual(
+            get_parallel().attn_tp_size,
+            1,
+            "initialize_model_parallel must not touch the published leaf -- "
+            "a caller that needs it corrected must publish a config that "
+            "already matches the width it is about to build",
+        )
+        self.assertEqual(get_parallel().moe_ep_size, 1)
+
 
 class TestTheDerivedHalfIsDeclared(CustomTestCase):
     """The quotients are declared beside the leaves, in the same class.