[Feature][Diffusion] Add BOOGU Image sequence parallel support (#5718)

Signed-off-by: Xenoryn <timberh2000@outlook.com>
Co-authored-by: Claude <noreply@anthropic.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>

vllm_omni/diffusion/hooks/sequence_parallel.py

@@ -182,6 +182,12 @@ class SequenceParallelSplitHook(ModelHook):
         self.module_forward_metadata: ModuleForwardMetadata | None = None
         # Cache for text lengths resolved from kwargs
         self._text_len_cache: dict[str, int] = {}
+        self._shard_groups = {
+            sp_input.shard_group
+            for value in metadata.values()
+            for sp_input in (value if isinstance(value, (list, tuple)) else (value,))
+            if isinstance(sp_input, SequenceParallelInput) and sp_input.shard_group is not None
+        }
 
     def initialize_hook(self, module: nn.Module) -> nn.Module:
         cls = _unwrap_module(module).__class__
@@ -193,6 +199,16 @@ class SequenceParallelSplitHook(ModelHook):
         args_list = list(args)
         # Clear text length cache for this forward pass
         self._text_len_cache.clear()
+        if self._shard_groups:
+            from vllm_omni.diffusion.forward_context import (
+                get_forward_context,
+                is_forward_context_available,
+            )
+
+            if is_forward_context_available():
+                ctx = get_forward_context()
+                for shard_group in self._shard_groups:
+                    ctx.sp_shard_metadata.pop(shard_group, None)
 
         for name, spm in self.metadata.items():
             # Skip if this is a split_output entry (handled in post_forward)
@@ -376,7 +392,11 @@ class SequenceParallelSplitHook(ModelHook):
         if isinstance(sp_input, SequenceParallelInput):
             # Full split with optional auto-padding
             if sp_input.auto_pad:
-                return self._shard_with_auto_pad(x, sp_input.split_dim)
+                return self._shard_with_auto_pad(
+                    x,
+                    sp_input.split_dim,
+                    sp_input.shard_group,
+                )
             _maybe_validate_strict_divisibility(dim=sp_input.split_dim, seq_len=x.size(sp_input.split_dim))
             return sp_shard(x, sp_input.split_dim, validate=False)
         elif isinstance(sp_input, SequenceParallelPartialInput):
@@ -397,13 +417,17 @@ class SequenceParallelSplitHook(ModelHook):
         else:
             raise ValueError(f"Unsupported input config type: {type(sp_input).__name__}")
 
-    def _shard_with_auto_pad(self, x: torch.Tensor, dim: int) -> torch.Tensor:
-        """Shard tensor with automatic padding and attention mask creation.
+    def _shard_with_auto_pad(
+        self,
+        x: torch.Tensor,
+        dim: int,
+        shard_group: str | None,
+    ) -> torch.Tensor:
+        """Pad and shard a tensor while recording its global sequence length.
 
-        When sequence length is not divisible by SP world size, this method:
-        1. Pads the tensor to make it divisible
-        2. Creates an attention mask indicating valid vs padding positions
-        3. Stores the mask and padding info in ForwardContext
+        Models with several independent sequences read the recorded length back
+        via `shard_group` to build masks. Legacy singleton padding fields are
+        also populated for single-sequence models.
         """
         from vllm_omni.diffusion.attention.selector import get_attn_backend_for_capability
         from vllm_omni.diffusion.distributed.parallel_state import (
@@ -418,10 +442,20 @@ class SequenceParallelSplitHook(ModelHook):
             return x
 
         seq_len = x.size(dim)
+        rank = get_sequence_parallel_rank()
         remainder = seq_len % world_size
 
+        if shard_group is not None and is_forward_context_available():
+            ctx = get_forward_context()
+            existing = ctx.sp_shard_metadata.setdefault(shard_group, seq_len)
+            if existing != seq_len:
+                raise ValueError(
+                    f"All tensors in shard_group={shard_group!r} must have the "
+                    f"same global sequence length, but got {existing} and {seq_len}."
+                )
+
         if remainder == 0:
-            # No padding needed
+            # Keyed groups record their length even when no padding is needed.
             return sp_shard(x, dim, validate=False)
 
         # Check backend compatibility
@@ -464,8 +498,11 @@ class SequenceParallelSplitHook(ModelHook):
         x_padded = torch.cat([x, padding], dim=dim)
 
         # Store padding info in forward context (only once, for primary tensor)
-        # Attention layers will create masks dynamically using this info
-        if is_forward_context_available():
+        # Attention layers will create masks dynamically using this info.
+        # Keyed groups record their length in sp_shard_metadata above and must
+        # not clobber the legacy singleton fields, whose semantics assume a
+        # single sequence per boundary.
+        if shard_group is None and is_forward_context_available():
             ctx = get_forward_context()
             # Only set if not already set (first auto_pad tensor wins)
             if ctx.sp_original_seq_len is None:
@@ -477,7 +514,6 @@ class SequenceParallelSplitHook(ModelHook):
                 )
 
         # Shard the padded tensor
-        rank = get_sequence_parallel_rank()
         return x_padded.chunk(world_size, dim=dim)[rank]
 
 
@@ -523,12 +559,6 @@ class SequenceParallelGatherHook(ModelHook):
         if len(output) != len(self.metadata):
             raise ValueError(f"Expected {len(self.metadata)} outputs, got {len(output)}.")
 
-        # Check if padding was applied during split
-        original_seq_len = None
-        if is_forward_context_available():
-            ctx = get_forward_context()
-            original_seq_len = ctx.sp_original_seq_len
-
         actually_gathered = False
 
         for i, spm in enumerate(self.metadata):
@@ -546,6 +576,15 @@ class SequenceParallelGatherHook(ModelHook):
             gathered = sp_gather(x, spm.gather_dim, validate=False)
 
             # Remove padding if it was applied
+            original_seq_len = None
+            if is_forward_context_available():
+                ctx = get_forward_context()
+                if spm.shard_group is None:
+                    original_seq_len = ctx.sp_original_seq_len
+                else:
+                    original_seq_len = ctx.sp_shard_metadata.get(spm.shard_group)
+                    if original_seq_len is None:
+                        raise RuntimeError(f"No SP shard metadata was registered for shard_group={spm.shard_group!r}.")
             if original_seq_len is not None and gathered.size(spm.gather_dim) > original_seq_len:
                 gathered = gathered.narrow(spm.gather_dim, 0, original_seq_len)
                 logger.debug(f"Removed padding: gathered shape {gathered.shape} (original_seq_len={original_seq_len})")
@@ -559,6 +598,9 @@ class SequenceParallelGatherHook(ModelHook):
             ctx._sp_shard_depth = max(0, ctx._sp_shard_depth - 1)
             if ctx._sp_equal_pad_stack:
                 ctx._sp_equal_pad_stack.pop()
+            for spm in self.metadata:
+                if spm is not None and spm.shard_group is not None:
+                    ctx.sp_shard_metadata.pop(spm.shard_group, None)
 
         return output[0] if is_tensor else type(output)(output)