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

benchmark/kernels/all_reduce/benchmark_all_reduce.py

@@ -21,6 +21,7 @@ from sglang.srt.distributed.parallel_state import (
     init_distributed_environment,
     initialize_model_parallel,
 )
+from sglang.test.test_utils import publish_build_topology
 
 
 def parse_args():
@@ -85,7 +86,8 @@ def init_dist(backend: str):
         distributed_init_method=distributed_init_method,
         local_rank=rank,
     )
-    initialize_model_parallel(tensor_model_parallel_size=world_size)
+    publish_build_topology(world_rank=rank, tp_size=world_size)
+    initialize_model_parallel()
     return dist.group.WORLD