diff --git a/python/sglang/srt/distributed/parallel_state.py b/python/sglang/srt/distributed/parallel_state.py index 8f0cda98d..8e3429dec 100644 --- a/python/sglang/srt/distributed/parallel_state.py +++ b/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, diff --git a/python/sglang/srt/managers/cache_controller.py b/python/sglang/srt/managers/cache_controller.py index cfb069fc9..d90c76b94 100644 --- a/python/sglang/srt/managers/cache_controller.py +++ b/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