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

test/registered/unit/mem_cache/test_dsa_layer_split_broadcast.py

@@ -20,7 +20,7 @@ import torch
 import torch.multiprocessing as mp
 
 from sglang.test.ci.ci_register import register_cuda_ci
-from sglang.test.test_utils import CustomTestCase
+from sglang.test.test_utils import CustomTestCase, publish_build_topology
 
 register_cuda_ci(est_time=28, stage="base-c", runner_config="4-gpu-b200")
 
@@ -54,10 +54,8 @@ def _run(rank: int, world: int, port: int):
         distributed_init_method=f"tcp://127.0.0.1:{port}",
         backend="nccl",
     )
-    initialize_model_parallel(
-        tensor_model_parallel_size=world,
-        attention_context_model_parallel_size=world,
-    )
+    publish_build_topology(tp_size=world, attn_cp_size=world, world_rank=rank)
+    initialize_model_parallel()
 
     from sglang.srt.mem_cache.dsa_cache_layer_split import (
         LayerSplitDSATokenToKVPool,

test/registered/unit/mem_cache/test_page_interleave_shard.py

@@ -84,7 +84,7 @@ from sglang.srt.runtime_context import get_parallel, publish
 from sglang.srt.server_args import ServerArgs
 from sglang.srt.utils import ceil_div
 from sglang.test.ci.ci_register import register_cpu_ci
-from sglang.test.test_utils import CustomTestCase
+from sglang.test.test_utils import CustomTestCase, publish_build_topology
 
 register_cpu_ci(est_time=30, suite="base-a-test-cpu")
 
@@ -1215,10 +1215,8 @@ def _dist_init(rank, world, port, attn_cp_size):
         ServerArgs(model_path="dummy", tp_size=world, attn_cp_size=attn_cp_size),
         role="scheduler",
     )
-    initialize_model_parallel(
-        tensor_model_parallel_size=world,
-        attention_context_model_parallel_size=attn_cp_size,
-    )
+    publish_build_topology(tp_size=world, attn_cp_size=attn_cp_size, world_rank=rank)
+    initialize_model_parallel()
 
 
 def _gather_make_spec(shard_rank, max_prefix_groups=16, chunk_groups=4):