[Diffusion][Feature] Add SP for Omnigen2 (#3206)

Signed-off-by: Zhang Jian <jianmusings@gmail.com>

docs/user_guide/diffusion_features.md

@@ -126,7 +126,7 @@ The following tables show which models support each feature:
 | **MagiHuman**            |     ❌     |     ❌      |           ❌           |       ❓        |         ✅         |          ❌          |   ❌    |             ✅             |          ❌           |       ❌        |        ❌         |
 | **MammothModa2(T2I)**    |     ❌     |     ❌      |           ❌           |       ❌        |         ❌         |          ❌          |   ❌    |             ❌             |          ❌           |       ❌        |        ❌         |
 | **Nextstep_1(T2I)**      |     ❓     |     ❓      |           ❌           |       ✅        |         ✅         |          ❌          |   ❌    |             ✅             |          ❌           |       ❌        |        ❌         |
-| **OmniGen2**             |     ❌     |     ✅      |           ❌           |       ❌        |         ✅         |          ❌          |   ❌    |             ❌             |          ❌           |       ❌        |        ❌         |
+| **OmniGen2**             |     ❌     |     ✅      |           ✅           |       ❌        |         ✅         |          ❌          |   ❌    |             ❌             |          ❌           |       ❌        |        ❌         |
 | **Ovis-Image**           |     ❌     |     ✅      |           ❌           |       ✅        |         ❌         |          ❌          |   ❌    |             ✅             |          ❌           |       ❌        |        ❌         |
 | **Qwen-Image**           |     ✅     |     ✅      |           ✅           |       ✅        |         ✅         |          ❌          |   ✅    |             ✅             |      ✅ (decode)      |       ✅        |        ✅         |
 | **Qwen-Image-2512**      |     ✅     |     ✅      |           ✅           |       ✅        |         ✅         |          ❌          |   ✅    |             ✅             |      ✅ (decode)      |       ✅        |        ✅         |

vllm_omni/diffusion/models/omnigen2/omnigen2_transformer.py

@@ -22,6 +22,10 @@ from vllm.model_executor.layers.quantization import QuantizationConfig
 from vllm.model_executor.model_loader.weight_utils import default_weight_loader
 
 from vllm_omni.diffusion.attention.layer import Attention
+from vllm_omni.diffusion.distributed.sp_plan import (
+    SequenceParallelInput,
+    SequenceParallelOutput,
+)
 from vllm_omni.platforms import current_omni_platform
 
 logger = logging.getLogger(__name__)
@@ -808,6 +812,36 @@ class OmniGen2TransformerBlock(nn.Module):
         return hidden_states
 
 
+class _OmniGen2SPInputBoundary(nn.Module):
+    """SP input-sharding boundary marker for ``OmniGen2Transformer2DModel._sp_plan``.
+
+    Resets ``sp_original_seq_len`` on every call: OmniGen2 runs the transformer
+    once per CFG branch and the unconditional / conditional branches have
+    different joint sequence lengths (text prefix differs). The SP framework's
+    auto-pad records the length only on the first call, so without this reset
+    a stale length from an earlier branch truncates the gather of later ones.
+    """
+
+    def forward(self, hidden_states: torch.Tensor, rotary_emb: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        from vllm_omni.diffusion.forward_context import (
+            get_forward_context,
+            is_forward_context_available,
+        )
+
+        if is_forward_context_available():
+            ctx = get_forward_context()
+            ctx.sp_original_seq_len = None
+            ctx.sp_padding_size = 0
+        return hidden_states, rotary_emb
+
+
+class _OmniGen2SPOutputBoundary(nn.Module):
+    """SP output-gathering boundary marker for ``OmniGen2Transformer2DModel._sp_plan``."""
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        return hidden_states
+
+
 class OmniGen2Transformer2DModel(nn.Module):
     """
     OmniGen2 Transformer 2D Model.
@@ -836,6 +870,17 @@ class OmniGen2Transformer2DModel(nn.Module):
         timestep_scale: Scale factor for timestep embeddings
     """
 
+    # ``attention_mask`` is intentionally omitted from the plan: OmniGen2Attention
+    # accepts but never reads it. ``auto_pad=True`` requires Ulysses SP — Ring SP
+    # cannot use attention masks.
+    _sp_plan = {
+        "sp_input_boundary": {
+            0: SequenceParallelInput(split_dim=1, expected_dims=3, split_output=True, auto_pad=True),
+            1: SequenceParallelInput(split_dim=1, expected_dims=3, split_output=True, auto_pad=True),
+        },
+        "sp_output_boundary": SequenceParallelOutput(gather_dim=1, expected_dims=3),
+    }
+
     def __init__(
         self,
         patch_size: int = 2,
@@ -983,6 +1028,9 @@ class OmniGen2Transformer2DModel(nn.Module):
         # Add learnable embeddings to distinguish different images
         self.image_index_embedding = nn.Parameter(torch.randn(5, hidden_size))  # support max 5 ref images
 
+        self.sp_input_boundary = _OmniGen2SPInputBoundary()
+        self.sp_output_boundary = _OmniGen2SPOutputBoundary()
+
     def img_patch_embed_and_refine(
         self,
         hidden_states,
@@ -1251,9 +1299,13 @@ class OmniGen2Transformer2DModel(nn.Module):
 
         hidden_states = joint_hidden_states
 
+        hidden_states, rotary_emb = self.sp_input_boundary(hidden_states, rotary_emb)
+
         for layer in self.layers:
             hidden_states = layer(hidden_states, attention_mask, rotary_emb, temb)
 
+        hidden_states = self.sp_output_boundary(hidden_states)
+
         # 4. Output norm & projection
         hidden_states = self.norm_out(hidden_states, temb)