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

benchmark/hf3fs/bench_zerocopy.py

@@ -5,10 +5,10 @@ import torch
 from tqdm import tqdm
 
 from sglang.srt.distributed import (
-    get_world_group,
     init_distributed_environment,
     initialize_model_parallel,
 )
+from sglang.srt.distributed.parallel_state import get_world_group
 from sglang.srt.managers.cache_controller import (
     HiCacheController,
     PrefetchOperation,
@@ -17,6 +17,7 @@ from sglang.srt.managers.cache_controller import (
 from sglang.srt.mem_cache.allocator import TokenToKVPoolAllocator
 from sglang.srt.mem_cache.memory_pool import MHATokenToKVPool
 from sglang.srt.mem_cache.pool_host.mha import MHATokenToKVPoolHost
+from sglang.test.test_utils import publish_build_topology
 
 init_distributed_environment(
     world_size=1,
@@ -26,10 +27,8 @@ init_distributed_environment(
     backend="gloo",
 )
 
-initialize_model_parallel(
-    tensor_model_parallel_size=1,
-    pipeline_model_parallel_size=1,
-)
+publish_build_topology()
+initialize_model_parallel()
 
 group = get_world_group().cpu_group
 

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
 
 

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)

benchmark/kernels/all_reduce/benchmark_fused_ar_rms_quant_amd.py

@@ -47,6 +47,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]
 FP8_DTYPE = torch.float8_e4m3fnuz
@@ -400,7 +401,8 @@ def main() -> None:
         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()
 
     if rank == 0:
         print(

benchmark/kernels/all_reduce/benchmark_mscclpp.py

@@ -30,6 +30,7 @@ from sglang.srt.distributed.parallel_state import (
     initialize_model_parallel,
     set_mscclpp_all_reduce,
 )
+from sglang.test.test_utils import publish_build_topology
 
 
 def torch_allreduce(torch_input: torch.Tensor, group: ProcessGroup) -> torch.Tensor:
@@ -173,7 +174,8 @@ if __name__ == "__main__":
         rank=rank,
         local_rank=rank % 8,
     )
-    initialize_model_parallel(tensor_model_parallel_size=world_size)
+    publish_build_topology(world_rank=rank, tp_size=world_size)
+    initialize_model_parallel()
     group = get_tensor_model_parallel_group().device_group
     cpu_group = get_tensor_model_parallel_group().cpu_group
     pynccl_comm = get_tensor_model_parallel_group().pynccl_comm

benchmark/kernels/all_reduce/benchmark_torch_symm_mem.py

@@ -44,6 +44,7 @@ from sglang.srt.distributed.parallel_state import (
     initialize_model_parallel,
     set_torch_symm_mem_all_reduce,
 )
+from sglang.test.test_utils import publish_build_topology
 from sglang.utils import is_in_ci
 
 IS_CI = is_in_ci()
@@ -188,7 +189,8 @@ if __name__ == "__main__":
         rank=rank,
         local_rank=rank % 8,
     )
-    initialize_model_parallel(tensor_model_parallel_size=world_size)
+    publish_build_topology(world_rank=rank, tp_size=world_size)
+    initialize_model_parallel()
     group = get_tensor_model_parallel_group().device_group
     cpu_group = get_tensor_model_parallel_group().cpu_group
     pynccl_comm = get_tensor_model_parallel_group().pynccl_comm

benchmark/kernels/flashinfer_allreduce_fusion/benchmark_fused_collective.py

@@ -30,14 +30,16 @@ import torch.distributed as dist  # type: ignore
 
 from sglang.kernels.ops.quantization.fp8_kernel import fp8_dtype as SGLANG_FP8_DTYPE
 from sglang.kernels.ops.quantization.fp8_kernel import static_quant_fp8
-from sglang.srt.distributed import get_tp_group, tensor_model_parallel_all_reduce
+from sglang.srt.distributed import tensor_model_parallel_all_reduce
 from sglang.srt.distributed.parallel_state import (
     cleanup_dist_env_and_memory,
+    get_tp_group,
     graph_capture,
     init_distributed_environment,
     initialize_model_parallel,
 )
 from sglang.srt.layers.layernorm import RMSNorm  # noqa
+from sglang.test.test_utils import publish_build_topology
 
 try:
     from sgl_kernel import fused_add_rmsnorm as SGL_FUSED_ADD_RMS_NORM
@@ -1178,7 +1180,8 @@ def main():
         local_rank=rank,
         backend="nccl",
     )
-    initialize_model_parallel(tensor_model_parallel_size=world_size)
+    publish_build_topology(world_rank=rank, tp_size=world_size)
+    initialize_model_parallel()
 
     # Validate world size (must be > 1 for collective operations)
     if world_size <= 1:

benchmark/kernels/fused_moe_triton/benchmark_sglang_fused_moe_triton.py

@@ -26,6 +26,7 @@ from sglang.srt.layers.moe.topk import (
     select_experts,
 )
 from sglang.srt.server_args import ServerArgs, set_global_server_args_for_scheduler
+from sglang.test.test_utils import publish_build_topology
 
 
 def fused_moe_triton_api(
@@ -227,10 +228,8 @@ def main():
             backend="nccl" if torch.cuda.is_available() else "gloo",
         )
 
-        initialize_model_parallel(
-            tensor_model_parallel_size=1,
-            expert_model_parallel_size=1,
-        )
+        publish_build_topology()
+        initialize_model_parallel()
 
         model_config = get_model_config(args.model, args.tp_size, args.ep_size)
         benchmark.run(

benchmark/kernels/fused_moe_triton/benchmark_vllm_vs_sglang_fused_moe_triton.py

@@ -15,6 +15,7 @@ from sglang.srt.distributed.parallel_state import (
 from sglang.srt.layers.moe.moe_runner.triton_utils.fused_moe import (
     fused_moe as fused_moe_sglang,
 )
+from sglang.test.test_utils import publish_build_topology
 
 from .common_utils import get_model_config
 
@@ -243,10 +244,8 @@ def main():
             backend="nccl" if torch.cuda.is_available() else "gloo",
         )
 
-        initialize_model_parallel(
-            tensor_model_parallel_size=1,
-            pipeline_model_parallel_size=1,
-        )
+        publish_build_topology()
+        initialize_model_parallel()
 
         shape_configs = get_model_config(args.model, args.tp_size, args.ep_size)
         benchmark.run(