[Feature] Add AllGather-KV sequence-parallel attention str… (#4968)

Signed-off-by: zhuhh97 <zhuhonghan@huawei.com>
Signed-off-by: dph97 <nenbaying@163.com>
Signed-off-by: zhuhh97 <zhh@unknown>
Co-authored-by: Claude Opus 4.7 <noreply@anthropic.com>
Co-authored-by: dengyunyang <584797741@qq.com>

vllm_omni/diffusion/distributed/group_coordinator.py

@@ -199,9 +199,15 @@ class GroupCoordinator:
         return input_
 
     def all_gather(
-        self, input_: torch.Tensor, dim: int = 0, separate_tensors: bool = False
+        self,
+        input_: torch.Tensor,
+        dim: int = 0,
+        separate_tensors: bool = False,
+        group: ProcessGroup | None = None,
     ) -> torch.Tensor | list[torch.Tensor]:
-        world_size = self.world_size
+        if group is None:
+            group = self.device_group
+        world_size = torch.distributed.get_world_size(group)
         # Bypass the function if we are using only 1 GPU.
         if world_size == 1:
             return input_
@@ -214,7 +220,7 @@ class GroupCoordinator:
         input_size[0] *= world_size
         output_tensor = torch.empty(input_size, dtype=input_.dtype, device=input_.device)
         # All-gather.
-        torch.distributed.all_gather_into_tensor(output_tensor, input_.contiguous(), group=self.device_group)
+        torch.distributed.all_gather_into_tensor(output_tensor, input_.contiguous(), group=group)
         if dim != 0:
             input_size[0] //= world_size
             output_tensor = output_tensor.reshape(
@@ -1002,6 +1008,7 @@ class SequenceParallelGroupCoordinator(GroupCoordinator):
 
         ulysses_group = kwargs.get("ulysses_group", None)
         ring_group = kwargs.get("ring_group", None)
+        allgather_group = kwargs.get("allgather_group", None)
         if ulysses_group is None:
             raise RuntimeError(
                 "Please pass argument 'ulysses_group' when calling init func of SequenceParallelGroupCoordinator"
@@ -1010,10 +1017,17 @@ class SequenceParallelGroupCoordinator(GroupCoordinator):
             raise RuntimeError(
                 "Please pass argument 'ring_group' when calling init func of SequenceParallelGroupCoordinator"
             )
+        if allgather_group is None:
+            raise RuntimeError(
+                "Please pass argument 'allgather_group' when calling init func of SequenceParallelGroupCoordinator"
+            )
         self.ulysses_group = ulysses_group
         self.ring_group = ring_group
+        self.allgather_group = allgather_group
 
         self.ulysses_world_size = torch.distributed.get_world_size(self.ulysses_group)
         self.ulysses_rank = torch.distributed.get_rank(self.ulysses_group)
         self.ring_world_size = torch.distributed.get_world_size(self.ring_group)
         self.ring_rank = torch.distributed.get_rank(self.ring_group)
+        self.allgather_world_size = torch.distributed.get_world_size(self.allgather_group)
+        self.allgather_rank = torch.distributed.get_rank(self.allgather_group)

vllm_omni/diffusion/distributed/parallel_state.py

@@ -303,6 +303,14 @@ def get_ring_parallel_rank():
     return get_sp_group().ring_rank
 
 
+def get_allgather_parallel_world_size():
+    return get_sp_group().allgather_world_size
+
+
+def get_allgather_parallel_rank():
+    return get_sp_group().allgather_rank
+
+
 def get_expert_parallel_group_ranks() -> list[list[int]]:
     assert vllm_parallel_state._EP is not None, "expert parallel group is not initialized"
     assert _EXPERT_PARALLEL_GROUP_RANKS is not None, "expert parallel group ranks are not initialized"
@@ -548,13 +556,12 @@ def set_seq_parallel_pg(
     world_size: int,
     use_ulysses_low: bool = True,
     sp_group_ranks: list[list[int]] | None = None,
-) -> tuple[torch.distributed.ProcessGroup, torch.distributed.ProcessGroup]:
+    sp_allgather_degree: int = 1,
+) -> tuple[torch.distributed.ProcessGroup, torch.distributed.ProcessGroup, torch.distributed.ProcessGroup]:
     """
-    Initialize sequence-parallel Ulysses and Ring process groups.
+    Initialize Ulysses, Ring, and AllGather-KV process groups.
 
-    This builds sequence-parallel (SP) subgroups inside each data-parallel (DP)
-    slice. The SP group size is sp_ulysses_degree * sp_ring_degree, and
-    world_size must be divisible by that size.
+    AllGather-KV is mutually exclusive with Ulysses and Ring.
 
     Args:
         sp_ulysses_degree: Size of each Ulysses subgroup.
@@ -569,7 +576,7 @@ def set_seq_parallel_pg(
             ranges.
 
     Returns:
-        ulyssess_pg (torch.distributed.ProcessGroup): The Ulysses process group
+        ulysses_pg (torch.distributed.ProcessGroup): The Ulysses process group
             for this rank.
         ring_pg (torch.distributed.ProcessGroup): The Ring process group for
             this rank.
@@ -586,6 +593,31 @@ def set_seq_parallel_pg(
         - If sp_group_ranks is None, groups are auto-generated within each DP
           slice using offsets of size sp_size.
     """
+    if sp_allgather_degree > 1:
+        if sp_ulysses_degree > 1 or sp_ring_degree > 1:
+            raise ValueError("AllGather-KV is mutually exclusive with Ulysses and Ring")
+        sp_size = sp_allgather_degree
+        if sp_group_ranks is None:
+            sp_group_ranks = [list(range(offset, offset + sp_size)) for offset in range(0, world_size, sp_size)]
+        if len(sp_group_ranks) * sp_size != world_size or any(len(ranks) != sp_size for ranks in sp_group_ranks):
+            raise ValueError(f"Invalid sp_group_ranks: expected {world_size // sp_size} groups of size {sp_size}.")
+
+        ulysses_pg = ring_pg = allgather_pg = None
+        for group_ranks in sp_group_ranks:
+            group = torch.distributed.new_group(group_ranks)
+            if rank in group_ranks:
+                allgather_pg = group
+        for singleton_rank in range(world_size):
+            group = torch.distributed.new_group([singleton_rank])
+            if rank == singleton_rank:
+                ulysses_pg = group
+        for singleton_rank in range(world_size):
+            group = torch.distributed.new_group([singleton_rank])
+            if rank == singleton_rank:
+                ring_pg = group
+        assert ulysses_pg is not None and ring_pg is not None and allgather_pg is not None
+        return ulysses_pg, ring_pg, allgather_pg
+
     sp_size = sp_ring_degree * sp_ulysses_degree
     dp_size = world_size // sp_size
 
@@ -619,7 +651,7 @@ def set_seq_parallel_pg(
                     ulysses_ranks = group_ranks[i * sp_ulysses_degree : (i + 1) * sp_ulysses_degree]
                     group = torch.distributed.new_group(ulysses_ranks)
                     if rank in ulysses_ranks:
-                        ulyssess_pg = group
+                        ulysses_pg = group
                         local_ulysses = list(ulysses_ranks)
                 for i in range(num_ring_pgs):
                     ring_ranks = group_ranks[i::num_ring_pgs]
@@ -639,7 +671,7 @@ def set_seq_parallel_pg(
                     ulysses_ranks = group_ranks[i::num_ulysses_pgs]
                     group = torch.distributed.new_group(ulysses_ranks)
                     if rank in ulysses_ranks:
-                        ulyssess_pg = group
+                        ulysses_pg = group
                         local_ulysses = list(ulysses_ranks)
         if local_sp_group is not None:
             logger.info(
@@ -662,7 +694,7 @@ def set_seq_parallel_pg(
                     )
                     group = torch.distributed.new_group(ulysses_ranks)
                     if rank in ulysses_ranks:
-                        ulyssess_pg = group
+                        ulysses_pg = group
 
                 for i in range(num_ring_pgs):
                     ring_ranks = list(range(i + offset, sp_size + offset, num_ring_pgs))
@@ -683,9 +715,15 @@ def set_seq_parallel_pg(
                     ulysses_ranks = list(range(i + offset, sp_size + offset, num_ulysses_pgs))
                     group = torch.distributed.new_group(ulysses_ranks)
                     if rank in ulysses_ranks:
-                        ulyssess_pg = group
+                        ulysses_pg = group
 
-    return ulyssess_pg, ring_pg
+    allgather_pg = None
+    for singleton_rank in range(world_size):
+        group = torch.distributed.new_group([singleton_rank])
+        if rank == singleton_rank:
+            allgather_pg = group
+    assert allgather_pg is not None
+    return ulysses_pg, ring_pg, allgather_pg
 
 
 def initialize_model_parallel(
@@ -694,6 +732,7 @@ def initialize_model_parallel(
     sequence_parallel_size: int | None = None,
     ulysses_degree: int = 1,
     ring_degree: int = 1,
+    allgather_degree: int = 1,
     tensor_parallel_size: int = 1,
     pipeline_parallel_size: int = 1,
     fully_shard_degree: int = 1,
@@ -710,9 +749,12 @@ def initialize_model_parallel(
         data_parallel_size: number of data parallelism groups.
         cfg_parallel_size: number of GPUs used for Classifier Free Guidance (CFG) parallelism.
         sequence_parallel_size: number of GPUs used for sequence parallelism.
-            sequence_parallel_size = ulysses_degree * ring_degree
+            Uses allgather_degree when AllGather-KV is enabled, otherwise
+            ulysses_degree * ring_degree.
         ulysses_degree: number of GPUs used for ulysses sequence parallelism.
         ring_degree: number of GPUs used for ring sequence parallelism.
+        allgather_degree: number of GPUs used for AllGather-KV sequence parallelism
+            (causal=False only). Mutually exclusive with ulysses/ring in v1.
         tensor_parallel_size: number of GPUs used for tensor parallelism.
         pipeline_parallel_size: number of GPUs used for pipeline parallelism.
         fully_shard_degree: number of GPUs used for fully sharded data parallelism (HSDP shard dimension).
@@ -749,16 +791,22 @@ def initialize_model_parallel(
     world_size: int = torch.distributed.get_world_size()
     backend = backend or torch.distributed.get_backend(get_world_group().device_group)
 
-    if sequence_parallel_size is None:
-        sequence_parallel_size = ring_degree * ulysses_degree
-        logger.info(
-            f"sequence_parallel_size is not provided, using ring_degree * ulysses_degree = {sequence_parallel_size}"
+    if allgather_degree > 1:
+        assert ulysses_degree == 1 and ring_degree == 1, (
+            "AllGather-KV (allgather_degree>1) is mutually exclusive with Ulysses/Ring in v1. "
+            f"Got ulysses_degree={ulysses_degree}, ring_degree={ring_degree}, "
+            f"allgather_degree={allgather_degree}."
         )
 
-    if sequence_parallel_size != ring_degree * ulysses_degree:
+    expected_sequence_parallel_size = allgather_degree if allgather_degree > 1 else ring_degree * ulysses_degree
+    if sequence_parallel_size is None:
+        sequence_parallel_size = expected_sequence_parallel_size
+        logger.info("sequence_parallel_size is not provided, using %d", sequence_parallel_size)
+
+    if sequence_parallel_size != expected_sequence_parallel_size:
         raise ValueError(
-            "sequence_parallel_size is not equal to ring_degree * ulysses_degree,"
-            f" but got {sequence_parallel_size} != {ring_degree} * {ulysses_degree}"
+            f"sequence_parallel_size must be {expected_sequence_parallel_size} for the configured SP mode, "
+            f"but got {sequence_parallel_size}"
         )
 
     dit_parallel_size = (
@@ -835,9 +883,10 @@ def initialize_model_parallel(
 
     global _SP
     assert _SP is None, "sequence parallel group is already initialized"
-    ulysses_pg, ring_pg = set_seq_parallel_pg(
+    ulysses_pg, ring_pg, allgather_pg = set_seq_parallel_pg(
         sp_ulysses_degree=ulysses_degree,
         sp_ring_degree=ring_degree,
+        sp_allgather_degree=allgather_degree,
         rank=get_world_group().rank_in_group,
         world_size=dit_parallel_size,
         sp_group_ranks=sp_group_ranks,
@@ -849,6 +898,7 @@ def initialize_model_parallel(
         parallel_mode="sequence",
         ulysses_group=ulysses_pg,
         ring_group=ring_pg,
+        allgather_group=allgather_pg,
     )
     if use_moe_parallel_mapping:
         # Diffusion normally uses its own SP group. Map it to vLLM PCP only for

vllm_omni/diffusion/distributed/sp_plan.py

@@ -77,6 +77,7 @@ class SequenceParallelConfig:
 
     ulysses_degree: int = 1
     ring_degree: int = 1
+    allgather_degree: int = 1
     convert_to_fp32: bool = True
 
     # Internal state - populated by setup()
@@ -85,18 +86,18 @@ class SequenceParallelConfig:
     _device: torch.device | None = None
 
     def __post_init__(self) -> None:
-        if self.ulysses_degree < 1 or self.ring_degree < 1:
-            raise ValueError("`ulysses_degree` and `ring_degree` must be >= 1.")
+        if self.ulysses_degree < 1 or self.ring_degree < 1 or self.allgather_degree < 1:
+            raise ValueError("SP degrees must be >= 1.")
 
-        if self.ulysses_degree == 1 and self.ring_degree == 1:
-            raise ValueError(
-                "At least one of `ulysses_degree` or `ring_degree` must be > 1 to use sequence parallelism."
-            )
+        if self.allgather_degree > 1 and (self.ulysses_degree > 1 or self.ring_degree > 1):
+            raise ValueError("AllGather-KV is mutually exclusive with Ulysses and Ring.")
+        if self.ulysses_degree == self.ring_degree == self.allgather_degree == 1:
+            raise ValueError("At least one SP degree must be > 1.")
 
     @property
     def sequence_parallel_size(self) -> int:
         """Total sequence parallel world size."""
-        return self.ulysses_degree * self.ring_degree
+        return self.allgather_degree if self.allgather_degree > 1 else self.ulysses_degree * self.ring_degree
 
     def get_world_size(self) -> int:
         """Get the sequence parallel world size from parallel state.
@@ -178,13 +179,12 @@ class SequenceParallelConfig:
         self._world_size = world_size
         self._device = device
 
-        expected_sp_size = self.ulysses_degree * self.ring_degree
+        expected_sp_size = self.sequence_parallel_size
         actual_sp_size = self.get_world_size()
 
         if expected_sp_size != actual_sp_size:
             raise ValueError(
-                f"Configuration mismatch: ulysses_degree ({self.ulysses_degree}) * "
-                f"ring_degree ({self.ring_degree}) = {expected_sp_size}, but "
+                f"Configuration mismatch: expected SP size {expected_sp_size}, but "
                 f"actual sequence parallel world size is {actual_sp_size}."
             )