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

benchmark/kernels/all_reduce/benchmark_fused_ar_rms_amd.py

@@ -41,6 +41,7 @@ from sglang.srt.distributed.parallel_state import (
     initialize_model_parallel,
     set_custom_all_reduce,
 )
+from sglang.test.test_utils import publish_build_topology
 
 Shape = Tuple[int, int]
 
@@ -381,7 +382,8 @@ def main():
         distributed_init_method="env://",
         backend="nccl",
     )
-    initialize_model_parallel(tensor_model_parallel_size=world_size)
+    publish_build_topology(world_rank=rank, tp_size=world_size)
+    initialize_model_parallel()
 
     prefill_shapes = parse_shapes(args.prefill_shapes)
     decode_shapes = parse_shapes(args.decode_shapes)