Check the topology identities where the layout is written, and build at the published widths (#40340)

test/manual/ep/test_flashinfer_dispatcher.py

@@ -14,7 +14,7 @@ from sglang.srt.layers.moe.token_dispatcher.flashinfer import FlashinferDispatch
 from sglang.srt.layers.moe.utils import initialize_moe_config
 from sglang.srt.runtime_context import get_context, publish
 from sglang.srt.server_args import ServerArgs, set_global_server_args_for_scheduler
-from sglang.test.test_utils import CustomTestCase
+from sglang.test.test_utils import CustomTestCase, publish_build_topology
 
 
 class TestFlashinferDispatcher(CustomTestCase):
@@ -44,9 +44,8 @@ class TestFlashinferDispatcher(CustomTestCase):
         publish(server_args, role="scheduler")
         initialize_moe_config()
 
-        initialize_model_parallel(
-            tensor_model_parallel_size=world_size, expert_model_parallel_size=world_size
-        )
+        publish_build_topology(tp_size=world_size, ep_size=world_size, world_rank=rank)
+        initialize_model_parallel()
 
     @classmethod
     def tearDownClass(cls):

test/manual/test_custom_allreduce.py

@@ -18,7 +18,7 @@ from sglang.srt.distributed.parallel_state import (
     initialize_model_parallel,
 )
 from sglang.srt.server_args import ServerArgs, set_global_server_args_for_scheduler
-from sglang.test.test_utils import CustomTestCase
+from sglang.test.test_utils import CustomTestCase, publish_build_topology
 
 
 def get_open_port() -> int:
@@ -98,7 +98,8 @@ class TestCustomAllReduce(CustomTestCase):
             distributed_init_method=distributed_init_method,
             local_rank=rank,
         )
-        initialize_model_parallel(tensor_model_parallel_size=world_size)
+        publish_build_topology(tp_size=world_size, world_rank=rank)
+        initialize_model_parallel()
         group = get_tensor_model_parallel_group().device_group
 
         # Set global server args to avoid "Global server args is not set yet!" error
@@ -161,7 +162,8 @@ class TestCustomAllReduce(CustomTestCase):
             distributed_init_method=distributed_init_method,
             local_rank=rank,
         )
-        initialize_model_parallel(tensor_model_parallel_size=world_size)
+        publish_build_topology(tp_size=world_size, world_rank=rank)
+        initialize_model_parallel()
         group = get_tensor_model_parallel_group().device_group
 
         # Set global server args to avoid "Global server args is not set yet!" error

test/manual/test_quick_allreduce.py

@@ -23,7 +23,7 @@ from sglang.srt.distributed.parallel_state import (
     graph_capture,
     initialize_model_parallel,
 )
-from sglang.test.test_utils import CustomTestCase
+from sglang.test.test_utils import CustomTestCase, publish_build_topology
 
 torch.manual_seed(42)
 random.seed(44)  # keep the deterministic seed
@@ -117,7 +117,8 @@ class TestQuickAllReduce(CustomTestCase):
             distributed_init_method=distributed_init_method,
             local_rank=rank,
         )
-        initialize_model_parallel(tensor_model_parallel_size=world_size)
+        publish_build_topology(tp_size=world_size, world_rank=rank)
+        initialize_model_parallel()
         group = get_tensor_model_parallel_group().device_group
 
         # A small all_reduce for warmup.
@@ -186,7 +187,8 @@ class TestQuickAllReduce(CustomTestCase):
             distributed_init_method=distributed_init_method,
             local_rank=rank,
         )
-        initialize_model_parallel(tensor_model_parallel_size=world_size)
+        publish_build_topology(tp_size=world_size, world_rank=rank)
+        initialize_model_parallel()
         group = get_tensor_model_parallel_group().device_group
 
         for sz in self.TEST_SIZES:

test/manual/test_weight_loader_v2_equiv.py

@@ -24,6 +24,7 @@ import unittest
 import torch
 
 from sglang.srt.environ import envs
+from sglang.test.test_utils import publish_build_topology
 
 MODEL = "Qwen/Qwen2-0.5B"
 
@@ -43,7 +44,8 @@ def _init_model_parallel() -> None:
             local_rank=0,
             distributed_init_method="tcp://127.0.0.1:29634",
         )
-        initialize_model_parallel(tensor_model_parallel_size=1)
+        publish_build_topology(tp_size=1)
+        initialize_model_parallel()
         monkey_patch_vllm_parallel_state()
     except AssertionError:
         pass