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

test/registered/layers/mamba/test_mamba2_mixer.py

@@ -17,6 +17,7 @@ from sglang.srt.distributed.parallel_state import (
 from sglang.srt.runtime_context import get_parallel
 from sglang.srt.utils import get_device, get_device_count
 from sglang.test.ci.ci_register import register_cuda_ci, register_xpu_ci
+from sglang.test.test_utils import publish_build_topology
 
 register_cuda_ci(est_time=30, stage="base-b", runner_config="2-gpu-large")
 register_xpu_ci(est_time=60, suite="nightly-xpu-2-gpu", nightly=True)
@@ -105,7 +106,8 @@ def mixer2_gated_norm_tensor_parallel(
         local_rank=local_rank,
         backend=get_default_distributed_backend(device.type),
     )
-    initialize_model_parallel(tensor_model_parallel_size=world_size)
+    publish_build_topology(tp_size=world_size, world_rank=local_rank)
+    initialize_model_parallel()
 
     # create random weights an inputs
     weight = torch.rand((hidden_size,), dtype=dtype, device=device)