[HiCache] Fix deadlock when creating new group (#15805)

Signed-off-by: Xuchun Shang <xuchun.shang@gmail.com>

python/sglang/srt/distributed/parallel_state.py

@@ -1717,6 +1717,55 @@ def initialize_model_parallel(
     )
 
 
+def create_custom_parallel_group(
+    group_ranks: List[int], backend: str = "gloo"
+) -> Optional[torch.distributed.ProcessGroup]:
+    """
+    Create a custom parallel group based on the provided ranks.
+
+    Args:
+        group_ranks: The list of ranks that the CURRENT process wants to join.
+                     (e.g., Rank 0 passes [0...7], Rank 8 passes [8...15])
+        backend: The communication backend (default: "gloo").
+
+    Returns:
+        The ProcessGroup if the current rank is in group_ranks, else None.
+    """
+    assert torch.distributed.is_initialized()
+
+    world_size = torch.distributed.get_world_size()
+    rank = torch.distributed.get_rank()
+
+    local_config = sorted(list(set(group_ranks)))
+    gathered_configs = [None for _ in range(world_size)]
+
+    torch.distributed.all_gather_object(gathered_configs, local_config)
+
+    unique_groups = []
+    seen_signatures = set()
+
+    for config in gathered_configs:
+        config_tuple = tuple(config)
+        if config_tuple not in seen_signatures:
+            seen_signatures.add(config_tuple)
+            unique_groups.append(list(config_tuple))
+
+    unique_groups.sort(key=lambda x: x[0])
+
+    my_new_group = None
+
+    for g_ranks in unique_groups:
+        group = torch.distributed.new_group(ranks=g_ranks, backend=backend)
+
+        if set(g_ranks) == set(local_config):
+            my_new_group = group
+            logger.debug(
+                f"Rank {rank} successfully created/joined custom group: {g_ranks}"
+            )
+
+    return my_new_group
+
+
 def ensure_model_parallel_initialized(
     tensor_model_parallel_size: int,
     expert_model_parallel_size: int,

python/sglang/srt/managers/cache_controller.py

@@ -309,9 +309,13 @@ class HiCacheController:
             # create a new communication group for synchronizing storage operations across TP workers
             self.tp_world_size = torch.distributed.get_world_size(group=tp_group)
             if self.tp_world_size > 1:
+                from sglang.srt.distributed.parallel_state import (
+                    create_custom_parallel_group,
+                )
+
                 group_ranks = torch.distributed.get_process_group_ranks(tp_group)
-                self.prefetch_tp_group = torch.distributed.new_group(
-                    group_ranks, backend="gloo"
+                self.prefetch_tp_group = create_custom_parallel_group(
+                    group_ranks=group_ranks, backend="gloo"
                 )
 
             # Select the get and set functions