[Diffusion] Non-Intrusive Sequence Parallelism (SP) Model Support Abstraction for vLLM-Omni Framework (#779)

Signed-off-by: mxuax <mxuax@connect.ust.hk>
Signed-off-by: XU Mingshi <91017482+mxuax@users.noreply.github.com>

docs/user_guide/diffusion/parallelism_acceleration.md

@@ -27,7 +27,7 @@ The following table shows which models are currently supported by parallelism me
 | **Qwen-Image-Edit** | `Qwen/Qwen-Image-Edit` | ✅ | ✅ | ✅ | ❌ |
 | **Qwen-Image-Edit-2509** | `Qwen/Qwen-Image-Edit-2509` | ✅ | ✅ | ✅ | ❌ |
 | **Qwen-Image-Layered** | `Qwen/Qwen-Image-Layered` | ✅ | ✅ | ✅ | ❌ |
-| **Z-Image** | `Tongyi-MAI/Z-Image-Turbo` | ❌ | ❌ | ❌ | ✅ (TP=2 only) |
+| **Z-Image** | `Tongyi-MAI/Z-Image-Turbo` | ✅ | ✅ | ❌ | ✅ (TP=2 only) |
 | **Stable-Diffusion3.5** | `stabilityai/stable-diffusion-3.5` | ❌ | ❌ | ❌ | ❌ |
 
 
@@ -220,100 +220,139 @@ To measure the parallelism methods, we run benchmarks with **Qwen/Qwen-Image** m
 
 ##### How to parallelize a new model
 
-If a diffusion model has been deployed in vLLM-Omni and supports single-card inference, you can refer to the following instructions to parallelize it with [Ulysses-SP](https://arxiv.org/pdf/2309.14509).
+NOTE: "Terminology: SP vs CP"
+    Our "Sequence Parallelism" (SP) corresponds to "Context Parallelism" (CP) in the [diffusers library](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/_modeling_parallel.py).
+    We use "Sequence Parallelism" to align with vLLM-Omni's terminology.
 
+---
 
-This section uses **Qwen-Image** (`QwenImageTransformer2DModel`) as the reference implementation. Qwen-Image is a **dual-stream** transformer (text + image) that performs **joint attention** across the concatenated sequences. Because of that, when enabling sequence parallel you typically:
+###### Non-intrusive `_sp_plan` (Recommended)
 
-- Chunk **image tokens** (`hidden_states`) across SP ranks along the **sequence dimension**.
-- Keep **text embeddings** (`encoder_hidden_states`) **replicated** on all SP ranks for correctness (unless you implement full joint-SP semantics and explicitly split text too).
-- Chunk **image RoPE freqs** to match the chunked image tokens.
-- `all_gather` the final output back along the sequence dimension.
+The `_sp_plan` mechanism allows SP without modifying `forward()` logic. The framework automatically registers hooks to shard inputs and gather outputs at module boundaries.
 
-First, add the sequence-parallel helpers and (for Qwen-Image) the forward-context flag. Then, in the transformer's `forward()`:
+**Requirements for `forward()` function:**
 
-- Chunk `hidden_states` by SP world size.
-- Set `get_forward_context().split_text_embed_in_sp = False` because Qwen-Image uses joint attention and we would use a full text embedding in each rank.
-- After `pos_embed`, chunk `img_freqs` on `dim=0` (token axis) to match the chunked image tokens; only chunk `txt_freqs` if you explicitly split text embeddings in SP.
+- Tensor operations that need sharding/gathering must happen at **`nn.Module` boundaries** (not inline Python operations)
+- If your `forward()` contains inline tensor operations (e.g., `torch.cat`, `pad_sequence`) that need sharding, **extract them into a submodule**
 
-Taking `vllm_omni/diffusion/models/qwen_image/qwen_image_transformer.py` as an example:
+**When to create a submodule:**
 
-```diff
-from vllm_omni.diffusion.distributed.parallel_state import (
-    get_sequence_parallel_rank,
-    get_sequence_parallel_world_size,
-    get_sp_group,
+```python
+# ❌ BAD: Inline operations - hooks cannot intercept
+def forward(self, x, cap_feats):
+    unified = torch.cat([x, cap_feats], dim=1)  # Cannot be sharded via _sp_plan
+    ...
+
+# ✅ GOOD: Extract into a submodule
+class UnifiedPrepare(nn.Module):
+    def forward(self, x, cap_feats):
+        return torch.cat([x, cap_feats], dim=1)  # Now can be sharded via _sp_plan
+
+class MyModel(nn.Module):
+    def __init__(self):
+        self.unified_prepare = UnifiedPrepare()  # Submodule
+
+    def forward(self, x, cap_feats):
+        unified = self.unified_prepare(x, cap_feats)  # Hook can intercept here
+```
+
+---
+
+###### Defining `_sp_plan`
+
+**Type definitions** (see [diffusers `_modeling_parallel.py`](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/_modeling_parallel.py) for reference):
+
+```python
+from vllm_omni.diffusion.distributed.sp_plan import (
+    SequenceParallelInput,   # Corresponds to diffusers' ContextParallelInput
+    SequenceParallelOutput,  # Corresponds to diffusers' ContextParallelOutput
 )
-from vllm_omni.diffusion.forward_context import get_forward_context
-
-class QwenImageTransformer2DModel(...):
-    def forward(
-        self,
-        hidden_states: torch.Tensor,
-        encoder_hidden_states: torch.Tensor = None,
-        encoder_hidden_states_mask: torch.Tensor = None,
-        timestep: torch.LongTensor = None,
-        img_shapes: list[tuple[int, int, int]] | None = None,
-        txt_seq_lens: list[int] | None = None,
-        ...
-    ):
-+       if self.parallel_config.sequence_parallel_size > 1:
-+           # Chunk image tokens along sequence dimension.
-+           hidden_states = torch.chunk(
-+               hidden_states,
-+               get_sequence_parallel_world_size(),
-+               dim=-2,
-+           )[get_sequence_parallel_rank()]
-+
-+           # Qwen-Image uses *dual-stream* (text + image) and runs *joint attention*.
-+           # Text embeddings should be replicated across SP ranks for correctness.
-+           get_forward_context().split_text_embed_in_sp = False
-
-        hidden_states = self.img_in(hidden_states)
-
-        ...
-        image_rotary_emb = self.pos_embed(img_shapes, txt_seq_lens, device=hidden_states.device)
-
-+       def get_rotary_emb_chunk(freqs):
-+           return torch.chunk(freqs, get_sequence_parallel_world_size(), dim=0)[
-+               get_sequence_parallel_rank()
-+           ]
-+
-+       if self.parallel_config.sequence_parallel_size > 1:
-+           img_freqs, txt_freqs = image_rotary_emb
-+           img_freqs = get_rotary_emb_chunk(img_freqs)
-+           if get_forward_context().split_text_embed_in_sp:
-+               txt_freqs = get_rotary_emb_chunk(txt_freqs)
-+           image_rotary_emb = (img_freqs, txt_freqs)
 ```
 
-Next, at the end of the `forward()` function, call `get_sp_group().all_gather` to gather the chunked outputs across devices and concatenate them along the sequence dimension (matching the earlier `dim=-2` chunking):
+| Parameter | Description |
+|-----------|-------------|
+| `split_dim` | Dimension to split/gather (usually `1` for sequence) |
+| `expected_dims` | Expected tensor rank for validation (optional) |
+| `split_output` | `False`: shard **input** parameters; `True`: shard **output** tensors |
+| `auto_pad` | Auto-pad if sequence not divisible by world_size (Ulysses only) |
 
-```diff
-class QwenImageTransformer2DModel(...):
-    def forward(...):
-        ...
-        hidden_states = self.norm_out(hidden_states, temb)
-        output = self.proj_out(hidden_states)
+**Key naming convention:**
 
-+       if self.parallel_config.sequence_parallel_size > 1:
-+           output = get_sp_group().all_gather(output, dim=-2)
-        return Transformer2DModelOutput(sample=output)
+| Key | Meaning | Python equivalent |
+|-----|---------|-------------------|
+| `""` | Root model | `model` |
+| `"blocks.0"` | First element of ModuleList | `model.blocks[0]` |
+| `"blocks.*"` | All elements of ModuleList | `for b in model.blocks` |
+| `"outputs.main"` | ModuleDict entry | `model.outputs["main"]` |
+
+**Dictionary key types:**
+
+| Key type | `split_output` | Description |
+|----------|----------------|-------------|
+| `"param_name"` (str) | `False` | Shard **input parameter** by name |
+| `0`, `1` (int) | `True` | Shard **output tuple** by index |
+
+**Example** (similar to [diffusers `transformer_wan.py`](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/transformers/transformer_wan.py)):
+
+```python
+class MyTransformer(nn.Module):
+    _sp_plan = {
+        # Shard rope module OUTPUTS (returns tuple)
+        "rope": {
+            0: SequenceParallelInput(split_dim=1, expected_dims=4, split_output=True),  # cos
+            1: SequenceParallelInput(split_dim=1, expected_dims=4, split_output=True),  # sin
+        },
+        # Shard transformer block INPUT parameter
+        "blocks.0": {
+            "hidden_states": SequenceParallelInput(split_dim=1, expected_dims=3),
+        },
+        # Gather at final projection
+        "proj_out": SequenceParallelOutput(gather_dim=1, expected_dims=3),
+    }
 ```
 
-Finally, you can set the parallel configuration and pass it to `Omni` and start parallel inference with:
-```diff
-from vllm_omni import Omni
-+from vllm_omni.diffusion.data import DiffusionParallelConfig
-ulysses_degree = 2
+---
 
-omni = Omni(
-    model="Qwen/Qwen-Image",
-+    parallel_config=DiffusionParallelConfig(ulysses_degree=2)
-)
+###### Hook flow
 
-outputs = omni.generate(prompt="A cat sitting on a windowsill", num_inference_steps=50)
 ```
+Input → [SequenceParallelSplitHook: pre_forward] → Module.forward() → [post_forward] → ...
+                                                                              ↓
+... → [SequenceParallelGatherHook: post_forward] → Output
+```
+
+1. **SplitHook** shards tensors before/after the target module
+2. **Attention layers** handle Ulysses/Ring communication internally
+3. **GatherHook** collects sharded outputs
+
+The framework automatically applies these hooks when `sequence_parallel_size > 1`.
+
+---
+
+###### Method 2: Intrusive modification (For complex cases)
+
+For models with dynamic sharding logic that cannot be expressed via `_sp_plan`:
+
+```python
+from vllm_omni.diffusion.distributed.sp_sharding import sp_shard, sp_gather
+
+def forward(self, hidden_states, ...):
+    if self.parallel_config.sequence_parallel_size > 1:
+        hidden_states = sp_shard(hidden_states, dim=1)
+        # ... computation ...
+        output = sp_gather(output, dim=1)
+    return output
+```
+
+---
+
+###### Choosing the right approach
+
+| Scenario | Approach |
+|----------|----------|
+| Standard transformer | `_sp_plan` |
+| Inline tensor ops need sharding | Extract to submodule + `_sp_plan` |
+| Dynamic/conditional sharding | Intrusive modification |
 
 
 ### CFG-Parallel

tests/e2e/offline_inference/test_sequence_parallel.py

@@ -9,6 +9,7 @@ correctly with diffusion models. It uses minimal settings to keep test time
 short for CI.
 """
 
+import gc
 import os
 import sys
 import time
@@ -21,7 +22,7 @@ import torch.distributed as dist
 from PIL import Image
 
 # ruff: noqa: E402
-REPO_ROOT = Path(__file__).resolve().parents[2]
+REPO_ROOT = Path(__file__).resolve().parents[3]
 if str(REPO_ROOT) not in sys.path:
     sys.path.insert(0, str(REPO_ROOT))
 
@@ -30,7 +31,6 @@ from vllm_omni.diffusion.data import DiffusionParallelConfig
 from vllm_omni.diffusion.distributed.parallel_state import device_count
 from vllm_omni.diffusion.envs import get_device_name
 
-# Use random weights model for testing
 models = ["riverclouds/qwen_image_random"]
 
 PROMPT = "a photo of a cat sitting on a laptop keyboard"
@@ -51,44 +51,145 @@ def _diff_metrics(a: Image.Image, b: Image.Image) -> tuple[float, float]:
     return abs_diff.mean().item(), abs_diff.max().item()
 
 
-def _get_images(output):
-    """Extract images from output, handling both dict and SimpleNamespace types.
+def _cleanup_distributed():
+    """Clean up distributed environment and GPU resources."""
+    if dist.is_initialized():
+        dist.destroy_process_group()
 
-    The output structure varies depending on serialization path:
-    - Direct memory: SimpleNamespace with .images attribute
-    - SHM serialization: dict with "images" key (dataclass converted via asdict)
-    - Wrapped output: SimpleNamespace(output=...) which needs unwrapping
-    """
-    # Check if output has direct images attribute (diffusion mode)
-    if hasattr(output, "images") and output.images:
-        return output.images
+    for key in ["MASTER_ADDR", "MASTER_PORT", "RANK", "WORLD_SIZE", "LOCAL_RANK"]:
+        os.environ.pop(key, None)
 
-    # Check request_output for pipeline mode
-    if output.request_output is None:
-        return None
+    gc.collect()
 
-    if isinstance(output.request_output, list) and len(output.request_output) == 0:
-        return None
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+        torch.cuda.synchronize()
 
-    item = output.request_output[0]
+    time.sleep(5)
 
-    # Handle wrapped SimpleNamespace (e.g. from omni_stage.py)
-    # Some items are wrapped as SimpleNamespace(request_id=..., output=...)
-    while hasattr(item, "output") and not hasattr(item, "images"):
-        item = item.output
-        if item is None:
-            return None
 
-    # Handle both dict (from SHM serialization) and object (direct) types
-    if isinstance(item, dict):
-        return item.get("images")
-    return getattr(item, "images", None)
+def _run_baseline(model_name: str, dtype: torch.dtype, attn_backend: str, height: int, width: int, seed: int):
+    """Run baseline inference (no SP)."""
+    baseline_parallel_config = DiffusionParallelConfig(ulysses_degree=1, ring_degree=1)
+    baseline = Omni(
+        model=model_name,
+        parallel_config=baseline_parallel_config,
+        dtype=dtype,
+        attention_backend=attn_backend,
+    )
+
+    try:
+        outputs = baseline.generate(
+            PROMPT,
+            height=height,
+            width=width,
+            num_inference_steps=4,
+            guidance_scale=0.0,
+            generator=torch.Generator(get_device_name()).manual_seed(seed),
+            num_outputs_per_prompt=1,
+        )
+        return outputs[0].request_output[0].images
+    finally:
+        baseline.close()
+        _cleanup_distributed()
+
+
+def _run_sp(
+    model_name: str,
+    dtype: torch.dtype,
+    attn_backend: str,
+    height: int,
+    width: int,
+    seed: int,
+    ulysses_degree: int,
+    ring_degree: int,
+):
+    """Run SP inference."""
+    sp_parallel_config = DiffusionParallelConfig(ulysses_degree=ulysses_degree, ring_degree=ring_degree)
+    sp = Omni(
+        model=model_name,
+        parallel_config=sp_parallel_config,
+        dtype=dtype,
+        attention_backend=attn_backend,
+    )
+
+    try:
+        outputs = sp.generate(
+            PROMPT,
+            height=height,
+            width=width,
+            num_inference_steps=4,
+            guidance_scale=0.0,
+            generator=torch.Generator(get_device_name()).manual_seed(seed),
+            num_outputs_per_prompt=1,
+        )
+        return outputs[0].request_output[0].images
+    finally:
+        sp.close()
+        _cleanup_distributed()
+
+
+@pytest.mark.parametrize("model_name", models)
+@pytest.mark.parametrize("dtype", [torch.bfloat16])
+@pytest.mark.parametrize("attn_backend", ["sdpa"])
+def test_baseline_only(model_name: str, dtype: torch.dtype, attn_backend: str):
+    """Test baseline inference only (no SP)."""
+    height = 256
+    width = 256
+    seed = 42
+
+    images = _run_baseline(model_name, dtype, attn_backend, height, width, seed)
+
+    assert images is not None
+    assert len(images) == 1
+    assert images[0].width == width
+    assert images[0].height == height
+
+
+@pytest.mark.parametrize("model_name", models)
+@pytest.mark.parametrize("dtype", [torch.bfloat16])
+@pytest.mark.parametrize("attn_backend", ["sdpa"])
+def test_sp_ulysses2_only(model_name: str, dtype: torch.dtype, attn_backend: str):
+    """Test SP inference only (ulysses=2)."""
+    if device_count() < 2:
+        pytest.skip(f"Test requires 2 GPUs but only {device_count()} available")
+
+    height = 256
+    width = 256
+    seed = 42
+
+    images = _run_sp(model_name, dtype, attn_backend, height, width, seed, ulysses_degree=2, ring_degree=1)
+
+    assert images is not None
+    assert len(images) == 1
+    assert images[0].width == width
+    assert images[0].height == height
+
+
+@pytest.mark.parametrize("model_name", models)
+@pytest.mark.parametrize("dtype", [torch.bfloat16])
+@pytest.mark.parametrize("attn_backend", ["sdpa"])
+def test_sp_ring2_only(model_name: str, dtype: torch.dtype, attn_backend: str):
+    """Test SP inference only (ring=2)."""
+    if device_count() < 2:
+        pytest.skip(f"Test requires 2 GPUs but only {device_count()} available")
+
+    height = 256
+    width = 256
+    seed = 42
+
+    images = _run_sp(model_name, dtype, attn_backend, height, width, seed, ulysses_degree=1, ring_degree=2)
+
+    assert images is not None
+    assert len(images) == 1
+    assert images[0].width == width
+    assert images[0].height == height
 
 
 @pytest.mark.parametrize("model_name", models)
 @pytest.mark.parametrize("ulysses_degree", [1, 2])
 @pytest.mark.parametrize("ring_degree", [1, 2])
-@pytest.mark.parametrize("dtype", [torch.bfloat16])  # Only test bfloat16 to reduce CI time
+@pytest.mark.parametrize("dtype", [torch.bfloat16])
 @pytest.mark.parametrize("attn_backend", ["sdpa"])
 def test_sequence_parallel(
     model_name: str,
@@ -100,214 +201,60 @@ def test_sequence_parallel(
     """Compare baseline (ulysses_degree=1) vs SP (ulysses_degree>1) outputs."""
     if ulysses_degree <= 1 and ring_degree <= 1:
         pytest.skip(
-            "This test compares ulysses_degree * ring_degree = 1 vs ulysses_degree * ring_degree > 1; provide ulysses_degree or ring_degree>1."
+            "This test compares ulysses_degree * ring_degree = 1 vs ulysses_degree * ring_degree > 1; "
+            "provide ulysses_degree or ring_degree>1."
         )
 
-    # Skip if not enough GPUs available for SP run
-    if device_count() < ulysses_degree * ring_degree:
-        pytest.skip(f"Test requires {ulysses_degree * ring_degree} GPUs but only {device_count()} available")
+    sp_size = ulysses_degree * ring_degree
+    if device_count() < sp_size:
+        pytest.skip(f"Test requires {sp_size} GPUs but only {device_count()} available")
 
-    # Use minimal settings for fast testing
     height = 256
     width = 256
-    num_inference_steps = 4  # Minimal steps for fast test
     seed = 42
 
-    # Step 1: Baseline (no Ulysses sequence parallel)
-    baseline_parallel_config = DiffusionParallelConfig(ulysses_degree=1, ring_degree=1)
-    baseline = Omni(
-        model=model_name,
-        parallel_config=baseline_parallel_config,
-        dtype=dtype,
-        attention_backend=attn_backend,
-    )
-    try:
-        outputs = baseline.generate(
-            PROMPT,
-            height=height,
-            width=width,
-            num_inference_steps=num_inference_steps,
-            guidance_scale=0.0,
-            generator=torch.Generator(get_device_name()).manual_seed(seed),
-            num_outputs_per_prompt=1,
-        )
-        baseline_images = outputs[0].request_output[0].images
-    finally:
-        baseline.close()
-        if dist.is_initialized():
-            dist.destroy_process_group()
-        for key in ["MASTER_ADDR", "MASTER_PORT", "RANK", "WORLD_SIZE", "LOCAL_RANK"]:
-            os.environ.pop(key, None)
-        time.sleep(2)  # Wait for resources to release
+    baseline_images = _run_baseline(model_name, dtype, attn_backend, height, width, seed)
+    assert baseline_images is not None and len(baseline_images) == 1
 
-    assert baseline_images is not None
-    assert len(baseline_images) == 1
-    assert baseline_images[0].width == width
-    assert baseline_images[0].height == height
+    sp_images = _run_sp(model_name, dtype, attn_backend, height, width, seed, ulysses_degree, ring_degree)
+    assert sp_images is not None and len(sp_images) == 1
 
-    # Step 2: SP (Ulysses-SP + Ring-SP)
-    sp_parallel_config = DiffusionParallelConfig(ulysses_degree=ulysses_degree, ring_degree=ring_degree)
-    sp = Omni(
-        model=model_name,
-        parallel_config=sp_parallel_config,
-        dtype=dtype,
-        attention_backend=attn_backend,
-    )
-    try:
-        outputs = sp.generate(
-            PROMPT,
-            height=height,
-            width=width,
-            num_inference_steps=num_inference_steps,
-            guidance_scale=0.0,
-            generator=torch.Generator(get_device_name()).manual_seed(seed),
-            num_outputs_per_prompt=1,
-        )
-        sp_images = outputs[0].request_output[0].images
-    finally:
-        sp.close()
-        if dist.is_initialized():
-            dist.destroy_process_group()
-        for key in ["MASTER_ADDR", "MASTER_PORT", "RANK", "WORLD_SIZE", "LOCAL_RANK"]:
-            os.environ.pop(key, None)
-        time.sleep(2)
-
-    assert sp_images is not None
-    assert len(sp_images) == 1
-    assert sp_images[0].width == width
-    assert sp_images[0].height == height
-
-    # Step 3: Compare outputs
     mean_abs_diff, max_abs_diff = _diff_metrics(baseline_images[0], sp_images[0])
 
-    # FP16/BF16 may differ slightly due to different computation order under parallelism.
-    if dtype in (torch.float16, torch.bfloat16):
-        mean_threshold = 2e-2
-        max_threshold = 2e-1
-    else:
-        mean_threshold = 1e-2
-        max_threshold = 1e-1
-
-    print(
-        "Image diff stats (baseline ulysses_degree*ring_degree=1 vs SP): "
-        f"mean_abs_diff={mean_abs_diff:.6e}, max_abs_diff={max_abs_diff:.6e}; "
-        f"thresholds: mean<={mean_threshold:.6e}, max<={max_threshold:.6e}; "
-        f"ulysses_degree={ulysses_degree}, ring_degree={ring_degree}, dtype={dtype}"
-    )
+    mean_threshold = 2e-2
+    max_threshold = 2e-1
 
     assert mean_abs_diff <= mean_threshold and max_abs_diff <= max_threshold, (
-        f"Image diff exceeded threshold: mean_abs_diff={mean_abs_diff:.6e}, max_abs_diff={max_abs_diff:.6e} "
-        f"(thresholds: mean<={mean_threshold:.6e}, max<={max_threshold:.6e}); "
-        f"ulysses_degree={ulysses_degree}, ring_degree={ring_degree}, dtype={dtype}"
+        f"Image diff exceeded threshold: mean={mean_abs_diff:.6e}, max={max_abs_diff:.6e}"
     )
 
 
 @pytest.mark.parametrize("model_name", models)
 @pytest.mark.parametrize("dtype", [torch.bfloat16])
 @pytest.mark.parametrize("attn_backend", ["sdpa"])
-def test_sequence_parallel_ulysses_sp_only(
-    model_name: str,
-    dtype: torch.dtype,
-    attn_backend: str,
-):
-    """Test sequence parallel with ulysses_degree=4, ring_degree=1, and the image size (332x332) where the sequence length is NOT divisible by sp_size."""
+def test_sequence_parallel_ulysses4(model_name: str, dtype: torch.dtype, attn_backend: str):
+    """Test SP with ulysses_degree=4."""
     ulysses_degree = 4
     ring_degree = 1
 
-    # Skip if not enough GPUs available for SP run
     if device_count() < ulysses_degree * ring_degree:
         pytest.skip(f"Test requires {ulysses_degree * ring_degree} GPUs but only {device_count()} available")
 
-    # (272/8) * (272/8) = 17 * 17 = 289 Not divisible by sp_size=4
     height = 272
     width = 272
-    num_inference_steps = 4  # Minimal steps for fast test
     seed = 42
 
-    # Step 1: Baseline (no Ulysses sequence parallel)
-    baseline_parallel_config = DiffusionParallelConfig(ulysses_degree=1, ring_degree=1)
-    baseline = Omni(
-        model=model_name,
-        parallel_config=baseline_parallel_config,
-        dtype=dtype,
-        attention_backend=attn_backend,
-    )
-    try:
-        outputs = baseline.generate(
-            PROMPT,
-            height=height,
-            width=width,
-            num_inference_steps=num_inference_steps,
-            guidance_scale=0.0,
-            generator=torch.Generator(get_device_name()).manual_seed(seed),
-            num_outputs_per_prompt=1,
-        )
-        baseline_images = outputs[0].request_output[0].images
-    finally:
-        baseline.close()
-        if dist.is_initialized():
-            dist.destroy_process_group()
-        for key in ["MASTER_ADDR", "MASTER_PORT", "RANK", "WORLD_SIZE", "LOCAL_RANK"]:
-            os.environ.pop(key, None)
-        time.sleep(2)  # Wait for resources to release
+    baseline_images = _run_baseline(model_name, dtype, attn_backend, height, width, seed)
+    assert baseline_images is not None and len(baseline_images) == 1
 
-    assert baseline_images is not None
-    assert len(baseline_images) == 1
-    assert baseline_images[0].width == width
-    assert baseline_images[0].height == height
+    sp_images = _run_sp(model_name, dtype, attn_backend, height, width, seed, ulysses_degree, ring_degree)
+    assert sp_images is not None and len(sp_images) == 1
 
-    # Step 2: SP (Ulysses-SP + Ring-SP)
-    sp_parallel_config = DiffusionParallelConfig(ulysses_degree=ulysses_degree, ring_degree=ring_degree)
-    sp = Omni(
-        model=model_name,
-        parallel_config=sp_parallel_config,
-        dtype=dtype,
-        attention_backend=attn_backend,
-    )
-    try:
-        outputs = sp.generate(
-            PROMPT,
-            height=height,
-            width=width,
-            num_inference_steps=num_inference_steps,
-            guidance_scale=0.0,
-            generator=torch.Generator(get_device_name()).manual_seed(seed),
-            num_outputs_per_prompt=1,
-        )
-        sp_images = outputs[0].request_output[0].images
-    finally:
-        sp.close()
-        if dist.is_initialized():
-            dist.destroy_process_group()
-        for key in ["MASTER_ADDR", "MASTER_PORT", "RANK", "WORLD_SIZE", "LOCAL_RANK"]:
-            os.environ.pop(key, None)
-        time.sleep(2)
-
-    assert sp_images is not None
-    assert len(sp_images) == 1
-    assert sp_images[0].width == width
-    assert sp_images[0].height == height
-
-    # Step 3: Compare outputs
     mean_abs_diff, max_abs_diff = _diff_metrics(baseline_images[0], sp_images[0])
 
-    # FP16/BF16 may differ slightly due to different computation order under parallelism.
-    if dtype in (torch.float16, torch.bfloat16):
-        mean_threshold = 2e-2
-        max_threshold = 2e-1
-    else:
-        mean_threshold = 1e-2
-        max_threshold = 1e-1
-
-    print(
-        "Image diff stats (baseline ulysses_degree*ring_degree=1 vs SP): "
-        f"mean_abs_diff={mean_abs_diff:.6e}, max_abs_diff={max_abs_diff:.6e}; "
-        f"thresholds: mean<={mean_threshold:.6e}, max<={max_threshold:.6e}; "
-        f"ulysses_degree={ulysses_degree}, ring_degree={ring_degree}, dtype={dtype}"
-    )
+    mean_threshold = 2e-2
+    max_threshold = 2e-1
 
     assert mean_abs_diff <= mean_threshold and max_abs_diff <= max_threshold, (
-        f"Image diff exceeded threshold: mean_abs_diff={mean_abs_diff:.6e}, max_abs_diff={max_abs_diff:.6e} "
-        f"(thresholds: mean<={mean_threshold:.6e}, max<={max_threshold:.6e}); "
-        f"ulysses_degree={ulysses_degree}, ring_degree={ring_degree}, dtype={dtype}"
+        f"Image diff exceeded threshold: mean={mean_abs_diff:.6e}, max={max_abs_diff:.6e}"
     )

tools/pre_commit/check_pickle_imports.py

@@ -20,7 +20,7 @@ ALLOWED_FILES = {
     "tests/e2e/offline_inference/utils.py",
     "tests/utils.py",
     "vllm_omni/diffusion/distributed/group_coordinator.py",
-    "tests/diffusion/attention/test_sequence_parallel.py",
+    "tests/diffusion/attention/test_attention_sp.py",
 }
 
 PICKLE_RE = re.compile(

vllm_omni/diffusion/distributed/__init__.py

@@ -0,0 +1,38 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Distributed utilities for vLLM-Omni diffusion models."""
+
+from vllm_omni.diffusion.distributed.sp_plan import (
+    SequenceParallelConfig,
+    SequenceParallelInput,
+    SequenceParallelModelPlan,
+    SequenceParallelOutput,
+    SequenceParallelPartialInput,
+    get_sp_plan_from_model,
+    validate_sp_plan,
+)
+from vllm_omni.diffusion.distributed.sp_sharding import (
+    ShardingValidator,
+    get_sharding_validator,
+    sp_gather,
+    sp_shard,
+    sp_shard_with_padding,
+)
+
+__all__ = [
+    # Config
+    "SequenceParallelConfig",
+    # Plan types
+    "SequenceParallelInput",
+    "SequenceParallelOutput",
+    "SequenceParallelPartialInput",
+    "SequenceParallelModelPlan",
+    "validate_sp_plan",
+    "get_sp_plan_from_model",
+    # Sharding utilities
+    "sp_shard",
+    "sp_gather",
+    "sp_shard_with_padding",
+    "ShardingValidator",
+    "get_sharding_validator",
+]

vllm_omni/diffusion/distributed/sp_plan.py

@@ -0,0 +1,457 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM and The HuggingFace Team
+# Type definitions in this module are adapted from HuggingFace diffusers library:
+#   diffusers/src/diffusers/models/_modeling_parallel.py
+"""Sequence Parallelism configuration and plan type definitions.
+
+This module provides:
+1. SequenceParallelConfig: Configuration for SP (ulysses_degree, ring_degree)
+2. SequenceParallelInput/Output: Type definitions for _sp_plan declarations
+3. Validation utilities for _sp_plan
+
+A _sp_plan is a dictionary that specifies how to shard/gather tensors at
+different points in a model's forward pass. This allows automatic handling
+of sequence parallelism without modifying the model's forward() method.
+
+NOTE: Our "Sequence Parallelism" (SP) corresponds to "Context Parallelism" (CP)
+in diffusers. We use "Sequence Parallelism" to align with vLLM-Omni terminology.
+
+Example:
+    class MyTransformer(nn.Module):
+        _sp_plan = {
+            # Split inputs before model forward
+            "": {
+                "hidden_states": SequenceParallelInput(split_dim=1, expected_dims=3),
+                "encoder_hidden_states": SequenceParallelInput(split_dim=1, expected_dims=3),
+            },
+            # Split RoPE embeddings after pos_embed layer
+            "pos_embed": {
+                0: SequenceParallelInput(split_dim=0, expected_dims=2, split_output=True),
+            },
+            # Gather output after proj_out layer
+            "proj_out": SequenceParallelOutput(gather_dim=1, expected_dims=3),
+        }
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    import torch
+    import torch.nn as nn
+
+
+# =============================================================================
+# Sequence Parallel Configuration
+# =============================================================================
+
+
+@dataclass
+class SequenceParallelConfig:
+    """Configuration for Sequence Parallelism using vLLM-Omni's parallel state.
+
+    This class provides a unified interface for SP configuration that integrates
+    with vLLM-Omni's existing SequenceParallelGroupCoordinator. Unlike diffusers'
+    DeviceMesh-based approach (ContextParallelConfig), this uses the existing
+    parallel state management.
+
+    Note: This corresponds to `ContextParallelConfig` in diffusers library.
+
+    Args:
+        ulysses_degree: Number of devices for Ulysses (All-to-All) attention.
+            Sequence is split across devices, with Q/K/V redistributed via
+            All-to-All communication. Best for moderate sequences with good
+            interconnect bandwidth.
+        ring_degree: Number of devices for Ring attention. Sequence is split
+            across devices, with K/V passed in a ring topology. Best for long
+            sequences with limited memory/bandwidth.
+        convert_to_fp32: Whether to convert output and LSE to float32 for
+            numerical stability in ring attention.
+
+    Note:
+        ulysses_degree * ring_degree = sequence_parallel_size
+        vLLM-Omni supports hybrid Ulysses-Ring attention (both > 1).
+    """
+
+    ulysses_degree: int = 1
+    ring_degree: int = 1
+    convert_to_fp32: bool = True
+
+    # Internal state - populated by setup()
+    _rank: int | None = None
+    _world_size: int | None = None
+    _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 and self.ring_degree == 1:
+            raise ValueError(
+                "At least one of `ulysses_degree` or `ring_degree` must be > 1 to use sequence parallelism."
+            )
+
+    @property
+    def sequence_parallel_size(self) -> int:
+        """Total sequence parallel world size."""
+        return self.ulysses_degree * self.ring_degree
+
+    def get_world_size(self) -> int:
+        """Get the sequence parallel world size from parallel state.
+
+        Returns:
+            The world size for sequence parallelism.
+
+        Raises:
+            RuntimeError: If parallel state is not initialized.
+        """
+        from vllm_omni.diffusion.distributed.parallel_state import get_sequence_parallel_world_size
+
+        return get_sequence_parallel_world_size()
+
+    def get_rank(self) -> int:
+        """Get the current rank in the sequence parallel group.
+
+        Returns:
+            The rank within the sequence parallel group.
+
+        Raises:
+            RuntimeError: If parallel state is not initialized.
+        """
+        from vllm_omni.diffusion.distributed.parallel_state import get_sequence_parallel_rank
+
+        return get_sequence_parallel_rank()
+
+    def get_ulysses_world_size(self) -> int:
+        """Get the Ulysses parallel world size.
+
+        Returns:
+            The world size for Ulysses (All-to-All) parallelism.
+        """
+        from vllm_omni.diffusion.distributed.parallel_state import get_ulysses_parallel_world_size
+
+        return get_ulysses_parallel_world_size()
+
+    def get_ulysses_rank(self) -> int:
+        """Get the current rank in the Ulysses parallel group.
+
+        Returns:
+            The rank within the Ulysses parallel group.
+        """
+        from vllm_omni.diffusion.distributed.parallel_state import get_ulysses_parallel_rank
+
+        return get_ulysses_parallel_rank()
+
+    def get_ring_world_size(self) -> int:
+        """Get the Ring parallel world size.
+
+        Returns:
+            The world size for Ring attention parallelism.
+        """
+        from vllm_omni.diffusion.distributed.parallel_state import get_ring_parallel_world_size
+
+        return get_ring_parallel_world_size()
+
+    def get_ring_rank(self) -> int:
+        """Get the current rank in the Ring parallel group.
+
+        Returns:
+            The rank within the Ring parallel group.
+        """
+        from vllm_omni.diffusion.distributed.parallel_state import get_ring_parallel_rank
+
+        return get_ring_parallel_rank()
+
+    def setup(self, rank: int, world_size: int, device: torch.device) -> None:
+        """Initialize the config with runtime parallel state.
+
+        This is called automatically when sequence parallelism is enabled.
+
+        Args:
+            rank: The global rank of this process.
+            world_size: Total world size.
+            device: The device for this rank.
+        """
+        self._rank = rank
+        self._world_size = world_size
+        self._device = device
+
+        expected_sp_size = self.ulysses_degree * self.ring_degree
+        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"actual sequence parallel world size is {actual_sp_size}."
+            )
+
+    def is_initialized(self) -> bool:
+        """Check if the config has been initialized with runtime state.
+
+        Returns:
+            True if setup() has been called, False otherwise.
+        """
+        return self._rank is not None
+
+
+# =============================================================================
+# Sequence Parallel Plan Type Definitions
+# =============================================================================
+
+
+@dataclass(frozen=True)
+class SequenceParallelInput:
+    """Configuration for splitting an input tensor across sequence parallel ranks.
+
+    This specifies how to shard a tensor in the pre-forward or post-forward hook
+    of a layer. The tensor will be split along the specified dimension.
+
+    Note: This corresponds to `ContextParallelInput` in diffusers library.
+
+    Args:
+        split_dim: The dimension along which to split the tensor.
+        expected_dims: Expected number of dimensions. If provided, validates that
+            the tensor has this many dimensions before splitting. If the tensor
+            has a different number of dimensions, splitting is skipped with a warning.
+        split_output: If True, split the output of the layer instead of the input.
+            This is useful for layers whose outputs should be split after preprocessing
+            (e.g., RoPE embeddings).
+        auto_pad: If True, automatically pad the tensor if its size along split_dim
+            is not divisible by world_size. Creates an attention mask to indicate
+            valid vs padding positions. The mask is stored in ForwardContext.
+            Note: Ring attention does not support attention mask, so auto_pad
+            should only be used with Ulysses SP.
+
+    Example:
+        # Split hidden_states along sequence dimension (dim 1)
+        SequenceParallelInput(split_dim=1, expected_dims=3)
+
+        # Split RoPE output along sequence dimension (dim 0)
+        SequenceParallelInput(split_dim=0, expected_dims=2, split_output=True)
+
+        # Split with auto-padding for variable-length sequences
+        SequenceParallelInput(split_dim=1, expected_dims=3, auto_pad=True)
+    """
+
+    split_dim: int
+    expected_dims: int | None = None
+    split_output: bool = False
+    auto_pad: bool = False
+
+    def __repr__(self) -> str:
+        return (
+            f"SequenceParallelInput(split_dim={self.split_dim}, "
+            f"expected_dims={self.expected_dims}, split_output={self.split_output}, "
+            f"auto_pad={self.auto_pad})"
+        )
+
+
+@dataclass(frozen=True)
+class SequenceParallelOutput:
+    """Configuration for gathering an output tensor across sequence parallel ranks.
+
+    This specifies how to gather a tensor in the post-forward hook of a layer.
+    The tensor will be gathered along the specified dimension from all ranks.
+
+    Note: This corresponds to `ContextParallelOutput` in diffusers library.
+
+    Args:
+        gather_dim: The dimension along which to gather the tensor.
+        expected_dims: Expected number of dimensions. If provided, validates that
+            the tensor has this many dimensions before gathering.
+
+    Example:
+        # Gather output along sequence dimension (dim 1)
+        SequenceParallelOutput(gather_dim=1, expected_dims=3)
+    """
+
+    gather_dim: int
+    expected_dims: int | None = None
+
+    def __repr__(self) -> str:
+        return f"SequenceParallelOutput(gather_dim={self.gather_dim}, expected_dims={self.expected_dims})"
+
+
+@dataclass(frozen=True)
+class SequenceParallelPartialInput:
+    """Configuration for partially splitting a tensor (e.g., split image part, keep text part).
+
+    This is designed for models like LongCat/Qwen where RoPE embeddings need special handling:
+    - Text portion: kept full across all ranks (for joint attention)
+    - Image portion: split across ranks
+
+    The tensor is assumed to be concatenated as [text_part, image_part] along split_dim.
+
+    Note: This is an extension beyond diffusers' standard ContextParallelInput,
+    designed for vLLM-Omni's dual-stream attention models.
+
+    Args:
+        split_dim: The dimension along which to split the image portion.
+        text_len_source: How to determine text length:
+            - str: Name of a forward parameter that contains text length
+            - int: Fixed text length value
+        expected_dims: Expected number of dimensions for validation.
+        split_output: If True, split the output instead of input.
+
+    Example:
+        # Split RoPE: text portion (from txt_ids.shape[0]) kept full, image portion split
+        SequenceParallelPartialInput(
+            split_dim=0,
+            text_len_source="txt_ids",  # Get text length from txt_ids.shape[0]
+            expected_dims=2,
+            split_output=True,
+        )
+
+        # Or with fixed text length
+        SequenceParallelPartialInput(
+            split_dim=0,
+            text_len_source=512,  # Fixed text length
+            expected_dims=2,
+            split_output=True,
+        )
+    """
+
+    split_dim: int
+    text_len_source: str | int
+    expected_dims: int | None = None
+    split_output: bool = False
+
+    def __repr__(self) -> str:
+        return (
+            f"SequenceParallelPartialInput(split_dim={self.split_dim}, "
+            f"text_len_source={self.text_len_source!r}, expected_dims={self.expected_dims}, "
+            f"split_output={self.split_output})"
+        )
+
+
+# =============================================================================
+# Type Aliases for _sp_plan Structure
+# =============================================================================
+
+# Any input config type
+AnySequenceParallelInput = SequenceParallelInput | SequenceParallelPartialInput
+
+# Input specification: maps parameter names (str) or output indices (int) to split config
+SequenceParallelInputType = dict[
+    str | int,
+    AnySequenceParallelInput | list[AnySequenceParallelInput] | tuple[AnySequenceParallelInput, ...],
+]
+
+# Output specification: single or multiple gather configs
+SequenceParallelOutputType = SequenceParallelOutput | list[SequenceParallelOutput] | tuple[SequenceParallelOutput, ...]
+
+# Full model plan: maps module names to input/output specifications
+# - Key "" refers to the model itself (root level)
+# - Key "module_name" refers to a submodule
+# - Key "module_name.*" refers to all children of a ModuleList
+#
+# Example of a complete _sp_plan:
+#
+#     _sp_plan = {
+#         # Root level: split model inputs before any submodule
+#         "": {
+#             "hidden_states": SequenceParallelInput(split_dim=1, expected_dims=3),
+#         },
+#         # Submodule: split outputs of pos_embed (RoPE) layer
+#         "pos_embed": {
+#             0: SequenceParallelInput(split_dim=1, expected_dims=4, split_output=True),  # cos
+#             1: SequenceParallelInput(split_dim=1, expected_dims=4, split_output=True),  # sin
+#         },
+#         # Submodule: gather outputs of proj_out layer
+#         "proj_out": SequenceParallelOutput(gather_dim=1, expected_dims=3),
+#     }
+#
+SequenceParallelModelPlan = dict[str, SequenceParallelInputType | SequenceParallelOutputType]
+
+
+# =============================================================================
+# Validation Utilities
+# =============================================================================
+
+
+def _is_valid_input_config(value: object) -> bool:
+    """Check if a value is a valid input configuration type."""
+    return isinstance(value, (SequenceParallelInput, SequenceParallelPartialInput))
+
+
+def _is_valid_input_config_list(value: object) -> bool:
+    """Check if a value is a list/tuple of valid input configurations."""
+    if not isinstance(value, (list, tuple)):
+        return False
+    return all(_is_valid_input_config(x) for x in value)
+
+
+def validate_sp_plan(plan: SequenceParallelModelPlan) -> None:
+    """Validate a _sp_plan dictionary for correctness.
+
+    Args:
+        plan: The _sp_plan dictionary to validate.
+
+    Raises:
+        ValueError: If the plan is invalid.
+    """
+    if not isinstance(plan, dict):
+        raise ValueError(f"_sp_plan must be a dict, got {type(plan).__name__}")
+
+    for module_id, module_plan in plan.items():
+        if not isinstance(module_id, str):
+            raise ValueError(f"_sp_plan keys must be strings, got {type(module_id).__name__}")
+
+        # Check if it's an output specification (SequenceParallelOutput or list/tuple thereof)
+        if isinstance(module_plan, SequenceParallelOutput):
+            continue
+        if isinstance(module_plan, (list, tuple)):
+            if all(isinstance(x, SequenceParallelOutput) for x in module_plan):
+                continue
+            if _is_valid_input_config_list(module_plan):
+                # List of inputs for a specific parameter (when output is tuple)
+                continue
+
+        # Otherwise, should be an input specification dict
+        if isinstance(module_plan, dict):
+            for key, value in module_plan.items():
+                if not isinstance(key, (str, int)):
+                    raise ValueError(
+                        f"Input spec keys must be str or int, got {type(key).__name__} for module '{module_id}'"
+                    )
+                if isinstance(key, int) and not _is_valid_input_config(value):
+                    raise ValueError(
+                        f"Integer keys (output indices) must map to SequenceParallelInput/PartialInput, "
+                        f"got {type(value).__name__} for module '{module_id}'[{key}]"
+                    )
+                if _is_valid_input_config(value):
+                    if isinstance(key, int) and not value.split_output:
+                        raise ValueError(
+                            f"Integer keys (output indices) require split_output=True, "
+                            f"got split_output=False for module '{module_id}'[{key}]"
+                        )
+                elif _is_valid_input_config_list(value):
+                    pass  # Valid list of input configs
+                else:
+                    raise ValueError(
+                        f"Input spec values must be SequenceParallelInput/PartialInput or list thereof, "
+                        f"got {type(value).__name__} for module '{module_id}'['{key}']"
+                    )
+        else:
+            raise ValueError(
+                f"_sp_plan values must be dict (input spec) or SequenceParallelOutput, "
+                f"got {type(module_plan).__name__} for module '{module_id}'"
+            )
+
+
+def get_sp_plan_from_model(model: nn.Module) -> SequenceParallelModelPlan | None:
+    """Get the _sp_plan from a model if it exists.
+
+    Args:
+        model: The model to get the plan from.
+
+    Returns:
+        The _sp_plan dictionary, or None if not defined.
+    """
+    plan = getattr(model, "_sp_plan", None)
+    if plan is not None:
+        validate_sp_plan(plan)
+    return plan

vllm_omni/diffusion/distributed/sp_sharding.py

@@ -0,0 +1,268 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project and The HuggingFace Team
+"""Sequence Parallelism sharding utilities.
+
+This module provides low-level sharding and gathering functions for
+Sequence Parallelism. These can be used directly in model forward methods
+for semi-intrusive SP support, or internally by the SP hooks.
+"""
+
+from __future__ import annotations
+
+from contextlib import contextmanager
+from dataclasses import dataclass, field
+
+import torch
+from vllm.logger import init_logger
+
+from vllm_omni.diffusion.distributed.parallel_state import (
+    get_sequence_parallel_rank,
+    get_sequence_parallel_world_size,
+    get_sp_group,
+)
+
+logger = init_logger(__name__)
+
+
+def sp_shard(
+    tensor: torch.Tensor,
+    dim: int,
+    validate: bool = True,
+) -> torch.Tensor:
+    """Shard a tensor along the specified dimension for sequence parallelism.
+
+    The tensor is split into world_size chunks along dim, and this rank
+    receives its corresponding chunk.
+
+    Args:
+        tensor: The tensor to shard.
+        dim: The dimension along which to split.
+        validate: If True, validate that the tensor size is divisible by world_size.
+
+    Returns:
+        The shard for this rank.
+
+    Raises:
+        ValueError: If validate=True and tensor size is not divisible by world_size.
+
+    Example:
+        # In model forward:
+        hidden_states = sp_shard(hidden_states, dim=1)
+    """
+    world_size = get_sequence_parallel_world_size()
+
+    if world_size == 1:
+        return tensor
+
+    rank = get_sequence_parallel_rank()
+    size = tensor.size(dim)
+
+    if validate and size % world_size != 0:
+        raise ValueError(
+            f"Tensor size along dim {dim} ({size}) must be divisible by "
+            f"world_size ({world_size}) for sequence parallel sharding."
+        )
+
+    return tensor.chunk(world_size, dim=dim)[rank]
+
+
+def sp_gather(
+    tensor: torch.Tensor,
+    dim: int,
+    validate: bool = True,
+) -> torch.Tensor:
+    """Gather a tensor along the specified dimension from all sequence parallel ranks.
+
+    The sharded tensors from all ranks are concatenated along dim.
+
+    Args:
+        tensor: The local shard to gather.
+        dim: The dimension along which to gather.
+        validate: If True, validate tensor consistency (currently unused).
+
+    Returns:
+        The full tensor gathered from all ranks.
+
+    Example:
+        # At end of model forward:
+        output = sp_gather(output, dim=1)
+    """
+    world_size = get_sequence_parallel_world_size()
+
+    if world_size == 1:
+        return tensor
+
+    sp_group = get_sp_group()
+    return sp_group.all_gather(tensor, dim=dim)
+
+
+def sp_shard_with_padding(
+    tensor: torch.Tensor,
+    dim: int,
+    pad_value: float = 0.0,
+) -> tuple[torch.Tensor, int]:
+    """Shard a tensor with automatic padding if not divisible by world_size.
+
+    This is useful for variable-length sequences where padding may be needed.
+
+    Args:
+        tensor: The tensor to shard.
+        dim: The dimension along which to split.
+        pad_value: Value to use for padding.
+
+    Returns:
+        Tuple of (sharded_tensor, padding_size). The padding_size indicates
+        how much padding was added to the original tensor before sharding.
+
+    Example:
+        sharded, pad_size = sp_shard_with_padding(hidden_states, dim=1)
+        # ... process ...
+        output = sp_gather(output, dim=1)
+        if pad_size > 0:
+            output = output[..., :-pad_size]  # Remove padding
+    """
+    world_size = get_sequence_parallel_world_size()
+
+    if world_size == 1:
+        return tensor, 0
+
+    size = tensor.size(dim)
+    remainder = size % world_size
+
+    if remainder == 0:
+        return sp_shard(tensor, dim, validate=False), 0
+
+    # Pad to make divisible
+    pad_size = world_size - remainder
+    pad_shape = list(tensor.shape)
+    pad_shape[dim] = pad_size
+    padding = torch.full(pad_shape, pad_value, dtype=tensor.dtype, device=tensor.device)
+    tensor = torch.cat([tensor, padding], dim=dim)
+
+    return sp_shard(tensor, dim, validate=False), pad_size
+
+
+# NOTE: This class is a vLLM-Omni extension for
+# debugging intrusive SP implementations.
+# Purpose:
+# - Help developers detect bugs when implementing intrusive SP
+# - Verify that every sharded tensor is properly gathered
+# - Warn about common mistakes (double shard, gather without shard)
+#
+# When to use:
+# - During development/debugging of intrusive SP code
+# - In tests to verify shard/gather correctness
+@dataclass
+class ShardingValidator:
+    """Validator for tracking and verifying sharding operations.
+
+    This class helps ensure that sharding and gathering operations are
+    correctly paired in model forward passes. It tracks which tensors
+    have been sharded and verifies that they are properly gathered.
+
+    Usage:
+        validator = ShardingValidator()
+        with validator.track():
+            hidden_states = validator.shard(hidden_states, "hidden_states", dim=1)
+            # ... model computation ...
+            output = validator.gather(output, "hidden_states", dim=1)
+        validator.validate()  # Raises if any shard was not gathered
+
+    Attributes:
+        _sharded: Set of tensor names that have been sharded.
+        _gathered: Set of tensor names that have been gathered.
+        _enabled: Whether tracking is currently enabled.
+    """
+
+    _sharded: set[str] = field(default_factory=set)
+    _gathered: set[str] = field(default_factory=set)
+    _enabled: bool = False
+
+    def reset(self) -> None:
+        """Reset the validator state for a new forward pass."""
+        self._sharded.clear()
+        self._gathered.clear()
+
+    @contextmanager
+    def track(self):
+        """Context manager to enable tracking for a forward pass."""
+        self._enabled = True
+        self.reset()
+        try:
+            yield
+        finally:
+            self._enabled = False
+
+    def shard(
+        self,
+        tensor: torch.Tensor,
+        name: str,
+        dim: int,
+        validate_divisible: bool = True,
+    ) -> torch.Tensor:
+        """Shard a tensor and track the operation.
+
+        Args:
+            tensor: The tensor to shard.
+            name: A name to identify this tensor for validation.
+            dim: The dimension along which to split.
+            validate_divisible: If True, validate divisibility.
+
+        Returns:
+            The sharded tensor.
+        """
+        if self._enabled:
+            if name in self._sharded:
+                logger.warning(f"Tensor '{name}' sharded multiple times")
+            self._sharded.add(name)
+
+        return sp_shard(tensor, dim, validate=validate_divisible)
+
+    def gather(
+        self,
+        tensor: torch.Tensor,
+        name: str,
+        dim: int,
+    ) -> torch.Tensor:
+        """Gather a tensor and track the operation.
+
+        Args:
+            tensor: The local shard to gather.
+            name: The name used when sharding (for validation).
+            dim: The dimension along which to gather.
+
+        Returns:
+            The gathered tensor.
+        """
+        if self._enabled:
+            if name not in self._sharded:
+                logger.warning(f"Tensor '{name}' gathered without being sharded")
+            self._gathered.add(name)
+
+        return sp_gather(tensor, dim)
+
+    def validate(self) -> None:
+        """Validate that all sharded tensors were gathered.
+
+        Raises:
+            ValueError: If any sharded tensor was not gathered.
+        """
+        unmatched = self._sharded - self._gathered
+        if unmatched:
+            raise ValueError(
+                f"The following tensors were sharded but not gathered: {unmatched}. "
+                f"This may indicate a bug in the model's SP implementation."
+            )
+
+
+# Global validator instance for convenience
+_global_validator = ShardingValidator()
+
+
+def get_sharding_validator() -> ShardingValidator:
+    """Get the global sharding validator instance.
+
+    Returns:
+        The global ShardingValidator.
+    """
+    return _global_validator

vllm_omni/diffusion/forward_context.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 from contextlib import contextmanager
 from dataclasses import dataclass
 
@@ -21,6 +23,14 @@ class ForwardContext:
     split_text_embed_in_sp: bool = False
     # whether to split the text embed in sequence parallel, if True, the text embed will be split in sequence parallel
 
+    # Sequence Parallel padding support
+    # When sequence length is not divisible by SP world size, padding is added
+    # These values are used by SequenceParallelGatherHook to remove padding,
+    # and by attention layers to create attention masks dynamically
+    sp_padding_size: int = 0
+    # Original sequence length before padding (for removing padding in gather)
+    sp_original_seq_len: int | None = None
+
     def __post_init__(self):
         pass
 

vllm_omni/diffusion/hooks.py

@@ -1,102 +0,0 @@
-from __future__ import annotations
-
-from collections.abc import Callable
-from dataclasses import dataclass
-from typing import Any
-
-import torch.nn as nn
-
-
-class BaseState:
-    """Base class for hook state containers."""
-
-    def reset(self) -> None:  # pragma: no cover - default is no-op
-        pass
-
-
-class StateManager:
-    """Manage per-context hook state instances."""
-
-    def __init__(self, state_cls: Callable[[], BaseState]):
-        self._state_cls = state_cls
-        self._states: dict[str, BaseState] = {}
-        self._context: str = "default"
-
-    def set_context(self, name: str) -> None:
-        self._context = name or "default"
-
-    def get_state(self) -> BaseState:
-        if self._context not in self._states:
-            self._states[self._context] = self._state_cls()
-        return self._states[self._context]
-
-    def reset(self) -> None:
-        self._states.clear()
-
-
-class ModelHook:
-    """Base class for model hooks that can override a module's forward."""
-
-    def initialize_hook(self, module: nn.Module) -> nn.Module:
-        raise NotImplementedError
-
-    def new_forward(self, module: nn.Module, *args: Any, **kwargs: Any):
-        raise NotImplementedError
-
-    def reset_state(self, module: nn.Module) -> nn.Module:
-        raise NotImplementedError
-
-
-@dataclass
-class _WrappedForward:
-    module: nn.Module
-
-    def __call__(self, *args: Any, **kwargs: Any):
-        registry: HookRegistry | None = getattr(self.module, "_hook_registry", None)
-        if registry is None or not registry._hooks:
-            return self.module._original_forward(*args, **kwargs)
-        return registry.dispatch(*args, **kwargs)
-
-
-class HookRegistry:
-    """Registry of hooks attached to a module."""
-
-    def __init__(self, module: nn.Module):
-        self.module = module
-        self._hooks: dict[str, ModelHook] = {}
-
-    @classmethod
-    def get_or_create(cls, module: nn.Module) -> HookRegistry:
-        registry: HookRegistry | None = getattr(module, "_hook_registry", None)
-        if registry is None:
-            registry = cls(module)
-            setattr(module, "_hook_registry", registry)
-
-            # Wrap module.forward once so hooks can intercept calls.
-            if not hasattr(module, "_original_forward"):
-                module._original_forward = module.forward  # type: ignore[attr-defined]
-                module.forward = _WrappedForward(module)  # type: ignore[assignment]
-
-        return registry
-
-    def register_hook(self, name: str, hook: ModelHook) -> None:
-        hook.initialize_hook(self.module)
-        self._hooks[name] = hook
-
-    def get_hook(self, name: str) -> ModelHook | None:
-        return self._hooks.get(name)
-
-    def dispatch(self, *args: Any, **kwargs: Any):
-        # For now we support a single active hook and call it directly.
-        # This can be extended to a chain if needed.
-        if not self._hooks:
-            return self.module._original_forward(*args, **kwargs)  # type: ignore[attr-defined]
-        # Deterministic order: sort by name.
-        name = sorted(self._hooks.keys())[0]
-        hook = self._hooks[name]
-        return hook.new_forward(self.module, *args, **kwargs)
-
-    def reset_hook(self, name: str) -> None:
-        hook = self._hooks.get(name)
-        if hook is not None:
-            hook.reset_state(self.module)

vllm_omni/diffusion/hooks/__init__.py

@@ -0,0 +1,33 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Hook mechanism for model forward interception."""
+
+from vllm_omni.diffusion.hooks.base import (
+    BaseState,
+    HookRegistry,
+    ModelHook,
+    StateManager,
+)
+from vllm_omni.diffusion.hooks.sequence_parallel import (
+    SequenceParallelGatherHook,
+    SequenceParallelSplitHook,
+    apply_sequence_parallel,
+    disable_sequence_parallel_for_model,
+    enable_sequence_parallel_for_model,
+    remove_sequence_parallel,
+)
+
+__all__ = [
+    # Base hooks
+    "BaseState",
+    "StateManager",
+    "ModelHook",
+    "HookRegistry",
+    # Sequence parallel hooks (corresponds to diffusers' context_parallel)
+    "SequenceParallelSplitHook",
+    "SequenceParallelGatherHook",
+    "apply_sequence_parallel",
+    "remove_sequence_parallel",
+    "enable_sequence_parallel_for_model",
+    "disable_sequence_parallel_for_model",
+]

vllm_omni/diffusion/hooks/base.py

@@ -0,0 +1,247 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Base hook classes for model forward interception.
+
+This module provides the foundational hook mechanism that allows intercepting
+and modifying model forward passes without invasive changes to model code.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+from dataclasses import dataclass
+from typing import Any
+
+import torch.nn as nn
+
+
+class BaseState:
+    """Base class for hook state containers."""
+
+    def reset(self) -> None:  # pragma: no cover - default is no-op
+        pass
+
+
+class StateManager:
+    """Manage per-context hook state instances."""
+
+    def __init__(self, state_cls: Callable[[], BaseState]):
+        self._state_cls = state_cls
+        self._states: dict[str, BaseState] = {}
+        self._context: str = "default"
+
+    def set_context(self, name: str) -> None:
+        self._context = name or "default"
+
+    def get_state(self) -> BaseState:
+        if self._context not in self._states:
+            self._states[self._context] = self._state_cls()
+        return self._states[self._context]
+
+    def reset(self) -> None:
+        self._states.clear()
+
+
+class ModelHook:
+    """Base class for model hooks that can override a module's forward.
+
+    Hooks can intercept the forward pass at two points:
+    - pre_forward: Called before the original forward, can modify args/kwargs
+    - post_forward: Called after the original forward, can modify output
+
+    Subclasses can override either or both methods. The default implementations
+    pass through args/kwargs/output unchanged.
+
+    For more complex behavior, override new_forward to completely replace
+    the forward logic.
+    """
+
+    def initialize_hook(self, module: nn.Module) -> nn.Module:
+        """Initialize the hook when it's registered to a module.
+
+        Args:
+            module: The module this hook is being attached to.
+
+        Returns:
+            The module (possibly modified).
+        """
+        return module
+
+    def pre_forward(self, module: nn.Module, *args: Any, **kwargs: Any) -> tuple[tuple, dict]:
+        """Called before the module's forward pass.
+
+        Args:
+            module: The module being called.
+            *args: Positional arguments to forward.
+            **kwargs: Keyword arguments to forward.
+
+        Returns:
+            Tuple of (args, kwargs) to pass to the forward method.
+        """
+        return args, kwargs
+
+    def post_forward(self, module: nn.Module, output: Any) -> Any:
+        """Called after the module's forward pass.
+
+        Args:
+            module: The module that was called.
+            output: The output from the forward method.
+
+        Returns:
+            The (possibly modified) output.
+        """
+        return output
+
+    def new_forward(self, module: nn.Module, *args: Any, **kwargs: Any) -> Any:
+        """Override the module's forward pass completely.
+
+        The default implementation calls pre_forward, then the original forward,
+        then post_forward. Override this method for more complex behavior.
+
+        Args:
+            module: The module being called.
+            *args: Positional arguments to forward.
+            **kwargs: Keyword arguments to forward.
+
+        Returns:
+            The output of the forward pass.
+        """
+        args, kwargs = self.pre_forward(module, *args, **kwargs)
+        output = module._original_forward(*args, **kwargs)  # type: ignore[attr-defined]
+        return self.post_forward(module, output)
+
+    def reset_state(self, module: nn.Module) -> nn.Module:
+        """Reset any state associated with this hook.
+
+        Args:
+            module: The module this hook is attached to.
+
+        Returns:
+            The module.
+        """
+        return module
+
+
+@dataclass
+class _WrappedForward:
+    """Wrapper that intercepts forward calls and dispatches to hooks."""
+
+    module: nn.Module
+
+    def __call__(self, *args: Any, **kwargs: Any):
+        registry: HookRegistry | None = getattr(self.module, "_hook_registry", None)
+        if registry is None or not registry._hooks:
+            return self.module._original_forward(*args, **kwargs)
+        return registry.dispatch(*args, **kwargs)
+
+
+class HookRegistry:
+    """Registry of hooks attached to a module.
+
+    Manages multiple hooks that can intercept a module's forward pass.
+    Hooks are called in sorted order by name for determinism.
+    """
+
+    def __init__(self, module: nn.Module):
+        self.module = module
+        self._hooks: dict[str, ModelHook] = {}
+
+    @classmethod
+    def get_or_create(cls, module: nn.Module) -> HookRegistry:
+        """Get existing registry or create a new one for the module.
+
+        Args:
+            module: The module to get/create a registry for.
+
+        Returns:
+            The HookRegistry for this module.
+        """
+        registry: HookRegistry | None = getattr(module, "_hook_registry", None)
+        if registry is None:
+            registry = cls(module)
+            setattr(module, "_hook_registry", registry)
+
+            # Wrap module.forward once so hooks can intercept calls.
+            if not hasattr(module, "_original_forward"):
+                module._original_forward = module.forward  # type: ignore[attr-defined]
+                module.forward = _WrappedForward(module)  # type: ignore[assignment]
+
+        return registry
+
+    def register_hook(self, name: str, hook: ModelHook) -> None:
+        """Register a hook with the given name.
+
+        Args:
+            name: Unique name for this hook.
+            hook: The hook instance to register.
+        """
+        hook.initialize_hook(self.module)
+        self._hooks[name] = hook
+
+    def remove_hook(self, name: str) -> None:
+        """Remove a hook by name.
+
+        Args:
+            name: The name of the hook to remove.
+        """
+        if name in self._hooks:
+            del self._hooks[name]
+
+    def get_hook(self, name: str) -> ModelHook | None:
+        """Get a hook by name.
+
+        Args:
+            name: The name of the hook.
+
+        Returns:
+            The hook if found, None otherwise.
+        """
+        return self._hooks.get(name)
+
+    def dispatch(self, *args: Any, **kwargs: Any) -> Any:
+        """Dispatch a forward call through registered hooks.
+
+        Currently supports a single active hook. Multiple hooks are called
+        in sorted order by name, with each hook's output passed to the next.
+
+        Args:
+            *args: Positional arguments to forward.
+            **kwargs: Keyword arguments to forward.
+
+        Returns:
+            The output of the forward pass.
+        """
+        if not self._hooks:
+            return self.module._original_forward(*args, **kwargs)  # type: ignore[attr-defined]
+
+        # For single hook case, call directly
+        if len(self._hooks) == 1:
+            hook = next(iter(self._hooks.values()))
+            return hook.new_forward(self.module, *args, **kwargs)
+
+        # For multiple hooks, chain them in sorted order
+        # Each hook can modify args/kwargs via pre_forward
+        sorted_hooks = sorted(self._hooks.items(), key=lambda x: x[0])
+
+        # Apply all pre_forward hooks
+        for _, hook in sorted_hooks:
+            args, kwargs = hook.pre_forward(self.module, *args, **kwargs)
+
+        # Call original forward
+        output = self.module._original_forward(*args, **kwargs)  # type: ignore[attr-defined]
+
+        # Apply all post_forward hooks in reverse order
+        for _, hook in reversed(sorted_hooks):
+            output = hook.post_forward(self.module, output)
+
+        return output
+
+    def reset_hook(self, name: str) -> None:
+        """Reset a hook's state by name.
+
+        Args:
+            name: The name of the hook to reset.
+        """
+        hook = self._hooks.get(name)
+        if hook is not None:
+            hook.reset_state(self.module)

vllm_omni/diffusion/hooks/sequence_parallel.py

@@ -0,0 +1,703 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project and the HuggingFace Team.
+# All rights reserved.
+#
+# This module is adapted from HuggingFace diffusers library:
+#   diffusers/src/diffusers/hooks/context_parallel.py
+#
+# NOTE: Our "Sequence Parallelism" (SP) corresponds to "Context Parallelism" (CP) in diffusers.
+# We use the term "Sequence Parallelism" to align with vLLM-Omni's existing terminology.
+#
+# Key adaptations from diffusers:
+#   - ModuleForwardMetadata: parameter lookup logic (adapted)
+#   - SequenceParallelSplitHook/GatherHook: hook structure (adapted from ContextParallel*)
+#   - apply_sequence_parallel: registration logic (adapted from apply_context_parallel)
+#
+# Key differences from diffusers:
+#   - Uses vLLM-Omni's SequenceParallelGroupCoordinator instead of DeviceMesh
+#   - Uses sp_shard/sp_gather from sp_sharding.py instead of funcol operations
+#   - Supports Ulysses + Ring hybrid parallelism
+#
+"""Sequence Parallelism hooks for non-intrusive SP support.
+
+This module implements the hook-based mechanism for applying sequence parallelism
+to models without modifying their forward() methods.
+
+Usage:
+    1. Define _sp_plan on your model class (corresponds to diffusers' _cp_plan)
+    2. Call apply_sequence_parallel(model, config, plan) to enable SP
+    3. Call remove_sequence_parallel(model, plan) to disable SP
+
+The hooks automatically shard inputs before forward and gather outputs after,
+based on the plan specification.
+"""
+
+from __future__ import annotations
+
+import inspect
+from dataclasses import dataclass
+from typing import Any
+
+import torch
+import torch.nn as nn
+from vllm.logger import init_logger
+
+from vllm_omni.diffusion.distributed.sp_plan import (
+    AnySequenceParallelInput,
+    SequenceParallelConfig,
+    SequenceParallelInput,
+    SequenceParallelModelPlan,
+    SequenceParallelOutput,
+    SequenceParallelPartialInput,
+)
+from vllm_omni.diffusion.distributed.sp_sharding import sp_gather, sp_shard
+from vllm_omni.diffusion.hooks.base import HookRegistry, ModelHook
+
+logger = init_logger(__name__)
+
+# Hook name templates for identifying SP hooks
+_SP_INPUT_HOOK_TEMPLATE = "sp_input---{}"
+_SP_OUTPUT_HOOK_TEMPLATE = "sp_output---{}"
+
+
+@dataclass
+class ModuleForwardMetadata:
+    """Metadata for mapping forward() parameter names to args/kwargs positions.
+
+    This caches the inspection of a module's forward signature to efficiently
+    locate parameters by name in subsequent calls.
+    """
+
+    cached_parameter_indices: dict[str, int] | None = None
+    _cls: type | None = None
+
+    def _get_parameter_from_args_kwargs(
+        self,
+        identifier: str,
+        args: tuple = (),
+        kwargs: dict | None = None,
+    ) -> tuple[Any, bool, int | None]:
+        """Get a parameter value from args or kwargs by name.
+
+        Args:
+            identifier: The parameter name to look up.
+            args: Positional arguments passed to forward.
+            kwargs: Keyword arguments passed to forward.
+
+        Returns:
+            Tuple of (value, is_kwarg, index).
+            - value: The parameter value (or None if not found)
+            - is_kwarg: True if found in kwargs
+            - index: Position in args if found there
+
+        Raises:
+            ValueError: If parameter not found in signature.
+        """
+        kwargs = kwargs or {}
+
+        # First check kwargs
+        if identifier in kwargs:
+            return kwargs[identifier], True, None
+
+        # Check cached indices
+        if self.cached_parameter_indices is not None:
+            index = self.cached_parameter_indices.get(identifier, None)
+            if index is None:
+                raise ValueError(f"Parameter '{identifier}' not found in cached indices.")
+            if index < len(args):
+                return args[index], False, index
+            return None, False, index
+
+        # Build cache from forward signature
+        if self._cls is None:
+            raise ValueError("Model class is not set for metadata.")
+
+        parameters = list(inspect.signature(self._cls.forward).parameters.keys())
+        parameters = parameters[1:]  # Skip `self`
+        self.cached_parameter_indices = {param: i for i, param in enumerate(parameters)}
+
+        if identifier not in self.cached_parameter_indices:
+            raise ValueError(f"Parameter '{identifier}' not found in function signature.")
+
+        index = self.cached_parameter_indices[identifier]
+
+        if index >= len(args):
+            return None, False, index
+
+        return args[index], False, index
+
+
+def _unwrap_module(module: nn.Module) -> nn.Module:
+    """Unwrap a module from any wrappers to get the original class.
+
+    Args:
+        module: Potentially wrapped module.
+
+    Returns:
+        The unwrapped module.
+    """
+    # Handle common wrappers
+    while hasattr(module, "_modules") and len(module._modules) == 1:
+        inner = next(iter(module._modules.values()))
+        if inner is not None:
+            module = inner
+        else:
+            break
+    return module
+
+
+class SequenceParallelSplitHook(ModelHook):
+    """Hook for splitting inputs before a module's forward pass.
+
+    This hook is registered to modules that need their inputs sharded
+    across sequence parallel ranks. It intercepts the forward call,
+    shards specified inputs according to the plan, and passes the
+    sharded inputs to the original forward.
+
+    For split_output=True inputs, it shards the output instead.
+
+    Supports both SequenceParallelInput (full split) and SequenceParallelPartialInput
+    (partial split for text/image separation).
+
+    Note: This corresponds to `ContextParallelSplitHook` in diffusers.
+    """
+
+    def __init__(
+        self,
+        metadata: dict[str | int, AnySequenceParallelInput | list[AnySequenceParallelInput]],
+        config: SequenceParallelConfig,
+    ) -> None:
+        super().__init__()
+        self.metadata = metadata
+        self.config = config
+        self.module_forward_metadata: ModuleForwardMetadata | None = None
+        # Cache for text lengths resolved from kwargs
+        self._text_len_cache: dict[str, int] = {}
+
+    def initialize_hook(self, module: nn.Module) -> nn.Module:
+        cls = _unwrap_module(module).__class__
+        self.module_forward_metadata = ModuleForwardMetadata(_cls=cls)
+        return module
+
+    def pre_forward(self, module: nn.Module, *args: Any, **kwargs: Any) -> tuple[tuple, dict]:
+        """Shard inputs before forward."""
+        args_list = list(args)
+        # Clear text length cache for this forward pass
+        self._text_len_cache.clear()
+
+        for name, spm in self.metadata.items():
+            # Skip if this is a split_output entry (handled in post_forward)
+            if isinstance(spm, (SequenceParallelInput, SequenceParallelPartialInput)) and spm.split_output:
+                continue
+
+            # Get the parameter value
+            input_val, is_kwarg, index = self.module_forward_metadata._get_parameter_from_args_kwargs(
+                name, args_list, kwargs
+            )
+
+            if input_val is None:
+                continue
+
+            # Shard the input
+            if isinstance(input_val, torch.Tensor):
+                input_val = self._prepare_sp_input(input_val, spm, args_list, kwargs)
+            elif isinstance(input_val, (list, tuple)):
+                # Handle list/tuple of tensors with per-element config
+                if not isinstance(spm, (list, tuple)):
+                    raise ValueError(
+                        f"Expected list/tuple of SequenceParallelInput for parameter '{name}' "
+                        f"which is a list/tuple, but got {type(spm).__name__}"
+                    )
+                if len(input_val) != len(spm):
+                    raise ValueError(f"Expected {len(spm)} elements for parameter '{name}', got {len(input_val)}")
+                sharded_input_val = []
+                for i, x in enumerate(input_val):
+                    if torch.is_tensor(x) and not spm[i].split_output:
+                        x = self._prepare_sp_input(x, spm[i], args_list, kwargs)
+                    sharded_input_val.append(x)
+                input_val = type(input_val)(sharded_input_val)
+            else:
+                raise ValueError(f"Unsupported input type for sharding: {type(input_val).__name__}")
+
+            # Update args or kwargs
+            if is_kwarg:
+                kwargs[name] = input_val
+            elif index is not None and index < len(args_list):
+                args_list[index] = input_val
+            else:
+                raise ValueError(f"Failed to update parameter '{name}' after sharding.")
+
+        # Store kwargs for post_forward to resolve text lengths
+        self._last_kwargs = kwargs
+        self._last_args = tuple(args_list)
+
+        return tuple(args_list), kwargs
+
+    def post_forward(self, module: nn.Module, output: Any) -> Any:
+        """Shard outputs for split_output=True entries."""
+        is_tensor = isinstance(output, torch.Tensor)
+        is_tensor_list = isinstance(output, (list, tuple)) and all(isinstance(x, torch.Tensor) for x in output)
+
+        if not is_tensor and not is_tensor_list:
+            # No tensor outputs to shard
+            return output
+
+        output_list = [output] if is_tensor else list(output)
+
+        for index, spm in self.metadata.items():
+            if not isinstance(index, int):
+                continue
+            if not isinstance(spm, (SequenceParallelInput, SequenceParallelPartialInput)) or not spm.split_output:
+                continue
+            if index >= len(output_list):
+                raise ValueError(f"Index {index} out of bounds for output of length {len(output_list)}.")
+
+            output_list[index] = self._prepare_sp_input(output_list[index], spm, self._last_args, self._last_kwargs)
+
+        return output_list[0] if is_tensor else type(output)(output_list)
+
+    def _resolve_text_len(
+        self,
+        sp_input: SequenceParallelPartialInput,
+        args: tuple,
+        kwargs: dict,
+    ) -> int:
+        """Resolve text length from the source specification."""
+        source = sp_input.text_len_source
+
+        if isinstance(source, int):
+            return source
+
+        # String source - look up from kwargs or args
+        if source in self._text_len_cache:
+            return self._text_len_cache[source]
+
+        # Try to get from kwargs/args
+        try:
+            val, _, _ = self.module_forward_metadata._get_parameter_from_args_kwargs(source, args, kwargs)
+            if val is None:
+                raise ValueError(f"Parameter '{source}' is None, cannot determine text length.")
+            if isinstance(val, torch.Tensor):
+                # TODO: Currently assumes batch_size=1, where shape[0] is sequence length.
+                # For batch inference support, this should be updated to handle
+                # shape (batch_size, seq_len, ...) where text_len varies per sample.
+                text_len = val.shape[0]
+            elif isinstance(val, int):
+                text_len = val
+            else:
+                raise ValueError(f"Cannot determine text length from '{source}' of type {type(val).__name__}")
+            self._text_len_cache[source] = text_len
+            return text_len
+        except ValueError as e:
+            raise ValueError(f"Failed to resolve text_len_source '{source}': {e}") from e
+
+    def _prepare_sp_input(
+        self,
+        x: torch.Tensor,
+        sp_input: AnySequenceParallelInput,
+        args: tuple = (),
+        kwargs: dict | None = None,
+    ) -> torch.Tensor:
+        """Shard a tensor according to the input specification."""
+        kwargs = kwargs or {}
+
+        if sp_input.expected_dims is not None and x.dim() != sp_input.expected_dims:
+            logger.warning_once(f"Expected tensor with {sp_input.expected_dims} dims, got {x.dim()}. Skipping split.")
+            return x
+
+        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 sp_shard(x, sp_input.split_dim, validate=False)
+        elif isinstance(sp_input, SequenceParallelPartialInput):
+            # Partial split: keep text portion, split image portion
+            text_len = self._resolve_text_len(sp_input, args, kwargs)
+            dim = sp_input.split_dim
+
+            # Split tensor into text and image portions
+            text_part = x.narrow(dim, 0, text_len)
+            image_part = x.narrow(dim, text_len, x.size(dim) - text_len)
+
+            # Only shard the image portion
+            image_part_sharded = sp_shard(image_part, dim, validate=False)
+
+            # Concatenate back: [text_full, image_sharded]
+            return torch.cat([text_part, image_part_sharded], dim=dim)
+        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.
+
+        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
+        """
+        from vllm_omni.diffusion.attention.selector import (
+            _BACKENDS_SUPPORT_ATTENTION_MASK,
+            get_attn_backend,
+        )
+        from vllm_omni.diffusion.distributed.parallel_state import (
+            get_ring_parallel_world_size,
+            get_sequence_parallel_rank,
+            get_sequence_parallel_world_size,
+        )
+        from vllm_omni.diffusion.forward_context import get_forward_context, is_forward_context_available
+
+        world_size = get_sequence_parallel_world_size()
+        if world_size == 1:
+            return x
+
+        seq_len = x.size(dim)
+        remainder = seq_len % world_size
+
+        if remainder == 0:
+            # No padding needed
+            return sp_shard(x, dim, validate=False)
+
+        # Check backend compatibility
+        attn_backend = get_attn_backend(-1)
+        if attn_backend.get_name() not in _BACKENDS_SUPPORT_ATTENTION_MASK:
+            raise ValueError(
+                f"Sequence length ({seq_len}) is not divisible by SP world size ({world_size}). "
+                f"Cannot use {attn_backend.get_name()} which does not support attention_mask. "
+                f"Please switch to SDPA or Ascend attention backend."
+            )
+
+        # Ring attention does not support attention_mask
+        if get_ring_parallel_world_size() > 1:
+            raise ValueError(
+                f"Sequence length ({seq_len}) is not divisible by SP world size ({world_size}). "
+                f"Cannot use Ring attention which does not support attention_mask. "
+                f"Please switch to Ulysses SP only."
+            )
+
+        # Calculate padding
+        pad_size = world_size - remainder
+        padded_seq_len = seq_len + pad_size
+
+        # Pad the tensor
+        pad_shape = list(x.shape)
+        pad_shape[dim] = pad_size
+        padding = torch.zeros(pad_shape, dtype=x.dtype, device=x.device)
+        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():
+            ctx = get_forward_context()
+            # Only set if not already set (first auto_pad tensor wins)
+            if ctx.sp_original_seq_len is None:
+                ctx.sp_padding_size = pad_size
+                ctx.sp_original_seq_len = seq_len
+                logger.debug(
+                    f"Auto-padded sequence from {seq_len} to {padded_seq_len} "
+                    f"(pad_size={pad_size}, world_size={world_size}, dim={dim})"
+                )
+
+        # Shard the padded tensor
+        rank = get_sequence_parallel_rank()
+        return x_padded.chunk(world_size, dim=dim)[rank]
+
+
+class SequenceParallelGatherHook(ModelHook):
+    """Hook for gathering outputs after a module's forward pass.
+
+    This hook is registered to modules that need their outputs gathered
+    from all sequence parallel ranks. It intercepts the output and gathers
+    it according to the plan specification.
+
+    Note: This corresponds to `ContextParallelGatherHook` in diffusers.
+    """
+
+    def __init__(
+        self,
+        metadata: SequenceParallelOutput | list[SequenceParallelOutput],
+        config: SequenceParallelConfig,
+    ) -> None:
+        super().__init__()
+        if isinstance(metadata, SequenceParallelOutput):
+            metadata = [metadata]
+        self.metadata = metadata
+        self.config = config
+
+    def initialize_hook(self, module: nn.Module) -> nn.Module:
+        return module
+
+    def post_forward(self, module: nn.Module, output: Any) -> Any:
+        """Gather outputs after forward and remove padding if applied."""
+        from vllm_omni.diffusion.forward_context import get_forward_context, is_forward_context_available
+
+        is_tensor = isinstance(output, torch.Tensor)
+
+        if is_tensor:
+            output = [output]
+        elif not (isinstance(output, (list, tuple)) and all(isinstance(x, torch.Tensor) for x in output)):
+            raise ValueError(f"Expected tensor or list/tuple of tensors, got {type(output).__name__}")
+
+        output = list(output)
+
+        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
+
+        for i, spm in enumerate(self.metadata):
+            if spm is None:
+                continue
+
+            x = output[i]
+            if spm.expected_dims is not None and x.dim() != spm.expected_dims:
+                logger.warning_once(
+                    f"Expected output tensor with {spm.expected_dims} dims, got {x.dim()}. Skipping gather."
+                )
+                continue
+
+            # Gather from all ranks
+            gathered = sp_gather(x, spm.gather_dim, validate=False)
+
+            # Remove padding if it was applied
+            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})")
+
+            output[i] = gathered
+
+        return output[0] if is_tensor else type(output)(output)
+
+
+def _get_submodule_by_name(model: nn.Module, name: str) -> nn.Module | list[nn.Module]:
+    """Get a submodule by dotted name, supporting wildcards.
+
+    Args:
+        model: The root module.
+        name: Dotted path to submodule. Use "*" to match all children
+            of a ModuleList.
+
+    Returns:
+        The submodule or list of submodules if wildcard used.
+
+    Raises:
+        ValueError: If the path is invalid or module not found.
+    """
+    if name.count("*") > 1:
+        raise ValueError("Wildcard '*' can only be used once in the name")
+    return _find_submodule_by_name(model, name)
+
+
+def _find_submodule_by_name(model: nn.Module, name: str) -> nn.Module | list[nn.Module]:
+    """Recursive helper for _get_submodule_by_name."""
+    if name == "":
+        return model
+
+    first_atom, remaining_name = name.split(".", 1) if "." in name else (name, "")
+
+    if first_atom == "*":
+        if not isinstance(model, nn.ModuleList):
+            raise ValueError("Wildcard '*' can only be used with ModuleList")
+        submodules = []
+        for submodule in model:
+            subsubmodules = _find_submodule_by_name(submodule, remaining_name)
+            if not isinstance(subsubmodules, list):
+                subsubmodules = [subsubmodules]
+            submodules.extend(subsubmodules)
+        return submodules
+    else:
+        if hasattr(model, first_atom):
+            submodule = getattr(model, first_atom)
+            return _find_submodule_by_name(submodule, remaining_name)
+        else:
+            raise ValueError(f"'{first_atom}' is not a submodule of '{model.__class__.__name__}'")
+
+
+def apply_sequence_parallel(
+    module: nn.Module,
+    config: SequenceParallelConfig,
+    plan: SequenceParallelModelPlan,
+) -> None:
+    """Apply sequence parallel hooks to a model according to the plan.
+
+    This function registers hooks on the specified submodules to automatically
+    shard inputs and gather outputs for sequence parallelism.
+
+    Note: This corresponds to `apply_context_parallel` in diffusers.
+
+    The complete SP flow is:
+    1. Input sharding (SequenceParallelSplitHook): Split sequence across SP ranks
+    2. Attention parallelism (handled by vLLM-Omni's Attention layer):
+       - Ulysses: All-to-All over Q/K/V heads
+       - Ring: K/V circulation in ring topology
+       - Hybrid: Both (Ulysses handles head redistribution, Ring handles K/V)
+    3. Output gathering (SequenceParallelGatherHook): Gather sequence from SP ranks
+
+    Args:
+        module: The model to apply SP to.
+        config: The sequence parallel configuration.
+        plan: Dictionary mapping module names to input/output specifications.
+
+    Example:
+        config = SequenceParallelConfig(ulysses_degree=2)
+        plan = {
+            "": {"hidden_states": SequenceParallelInput(split_dim=1, expected_dims=3)},
+            "proj_out": SequenceParallelOutput(gather_dim=1, expected_dims=3),
+        }
+        apply_sequence_parallel(model, config, plan)
+
+    Note:
+        vLLM-Omni's Attention layer automatically handles the internal
+        parallelism (Ulysses All-to-All or Ring attention) based on the
+        forward_context configuration. This function only handles input/output
+        sharding for the model as a whole.
+    """
+    logger.debug(
+        f"Applying sequence parallel with config: ulysses={config.ulysses_degree}, "
+        f"ring={config.ring_degree}, plan keys: {list(plan.keys())}"
+    )
+
+    for module_id, sp_model_plan in plan.items():
+        submodule = _get_submodule_by_name(module, module_id)
+        if not isinstance(submodule, list):
+            submodule = [submodule]
+
+        logger.debug(f"Applying SP hooks to '{module_id}' ({len(submodule)} module(s))")
+
+        for m in submodule:
+            if isinstance(sp_model_plan, dict):
+                # Input specification
+                hook = SequenceParallelSplitHook(sp_model_plan, config)
+                hook_name = _SP_INPUT_HOOK_TEMPLATE.format(module_id)
+            elif isinstance(sp_model_plan, (SequenceParallelOutput, list, tuple)):
+                # Output specification
+                if isinstance(sp_model_plan, SequenceParallelOutput):
+                    sp_model_plan = [sp_model_plan]
+                if not all(isinstance(x, SequenceParallelOutput) or x is None for x in sp_model_plan):
+                    raise ValueError(f"Expected SequenceParallelOutput elements, got {sp_model_plan}")
+                hook = SequenceParallelGatherHook(sp_model_plan, config)
+                hook_name = _SP_OUTPUT_HOOK_TEMPLATE.format(module_id)
+            else:
+                raise ValueError(f"Unsupported plan type: {type(sp_model_plan).__name__}")
+
+            registry = HookRegistry.get_or_create(m)
+            registry.register_hook(hook_name, hook)
+
+
+def remove_sequence_parallel(
+    module: nn.Module,
+    plan: SequenceParallelModelPlan,
+) -> None:
+    """Remove sequence parallel hooks from a model.
+
+    Note: This corresponds to `remove_context_parallel` in diffusers.
+
+    Args:
+        module: The model to remove SP from.
+        plan: The same plan used when applying SP.
+    """
+    for module_id, sp_model_plan in plan.items():
+        submodule = _get_submodule_by_name(module, module_id)
+        if not isinstance(submodule, list):
+            submodule = [submodule]
+
+        for m in submodule:
+            registry = getattr(m, "_hook_registry", None)
+            if registry is None:
+                continue
+
+            if isinstance(sp_model_plan, dict):
+                hook_name = _SP_INPUT_HOOK_TEMPLATE.format(module_id)
+            elif isinstance(sp_model_plan, (SequenceParallelOutput, list, tuple)):
+                hook_name = _SP_OUTPUT_HOOK_TEMPLATE.format(module_id)
+            else:
+                continue
+
+            registry.remove_hook(hook_name)
+
+
+def enable_sequence_parallel_for_model(
+    model: nn.Module,
+    config: SequenceParallelConfig | None = None,
+) -> None:
+    """Enable sequence parallelism for a model using its _sp_plan.
+
+    This is a convenience function that reads the model's _sp_plan attribute
+    and applies sequence parallelism automatically.
+
+    Note: This corresponds to `enable_context_parallel_for_model` in diffusers,
+    but uses vLLM-Omni's _sp_plan instead of diffusers' _cp_plan.
+
+    The function performs two main tasks:
+    1. Applies _sp_plan hooks to shard inputs and gather outputs
+    2. Ensures Attention layers are configured for the correct parallel mode
+       (handled automatically by vLLM-Omni's forward_context mechanism)
+
+    Args:
+        model: The model to enable SP for. Must have a _sp_plan attribute.
+        config: Optional config. If None, uses default based on current
+            parallel state.
+
+    Raises:
+        ValueError: If model has no _sp_plan defined.
+
+    Note:
+        vLLM-Omni supports Ulysses + Ring hybrid parallelism:
+        - ulysses_degree > 1: Uses All-to-All communication over Q/K/V heads
+        - ring_degree > 1: Uses Ring attention with K/V passing
+        - Both > 1: Hybrid mode (Ulysses handles head redistribution,
+          Ring handles K/V circulation)
+    """
+    from vllm_omni.diffusion.distributed.parallel_state import (
+        get_ring_parallel_world_size,
+        get_ulysses_parallel_world_size,
+    )
+    from vllm_omni.diffusion.distributed.sp_plan import get_sp_plan_from_model
+
+    plan = get_sp_plan_from_model(model)
+    if plan is None:
+        raise ValueError(
+            f"Model {model.__class__.__name__} has no _sp_plan defined. "
+            f"Define _sp_plan as a class attribute or pass a plan explicitly."
+        )
+
+    if config is None:
+        # Create config from current parallel state
+        ulysses_degree = get_ulysses_parallel_world_size()
+        ring_degree = get_ring_parallel_world_size()
+        config = SequenceParallelConfig(
+            ulysses_degree=ulysses_degree,
+            ring_degree=ring_degree,
+        )
+        if ulysses_degree > 1 and ring_degree > 1:
+            mode = "hybrid"
+        elif ulysses_degree > 1:
+            mode = "ulysses"
+        else:
+            mode = "ring"
+        logger.info(
+            f"Created SP config from parallel state: "
+            f"ulysses_degree={ulysses_degree}, ring_degree={ring_degree}, "
+            f"mode={mode}"
+        )
+
+    apply_sequence_parallel(model, config, plan)
+    logger.info(f"Enabled sequence parallelism for {model.__class__.__name__}")
+
+
+def disable_sequence_parallel_for_model(model: nn.Module) -> None:
+    """Disable sequence parallelism for a model.
+
+    Note: This corresponds to `disable_context_parallel_for_model` in diffusers.
+
+    Args:
+        model: The model to disable SP for.
+    """
+    from vllm_omni.diffusion.distributed.sp_plan import get_sp_plan_from_model
+
+    plan = get_sp_plan_from_model(model)
+    if plan is not None:
+        remove_sequence_parallel(model, plan)

vllm_omni/diffusion/models/qwen_image/qwen_image_transformer.py

@@ -27,14 +27,11 @@ from vllm_omni.diffusion.attention.backends.abstract import (
     AttentionMetadata,
 )
 from vllm_omni.diffusion.attention.layer import Attention
-from vllm_omni.diffusion.attention.selector import _BACKENDS_SUPPORT_ATTENTION_MASK, get_attn_backend
 from vllm_omni.diffusion.cache.base import CachedTransformer
 from vllm_omni.diffusion.data import OmniDiffusionConfig
-from vllm_omni.diffusion.distributed.parallel_state import (
-    get_ring_parallel_world_size,
-    get_sequence_parallel_rank,
-    get_sequence_parallel_world_size,
-    get_sp_group,
+from vllm_omni.diffusion.distributed.sp_plan import (
+    SequenceParallelInput,
+    SequenceParallelOutput,
 )
 from vllm_omni.diffusion.forward_context import get_forward_context
 from vllm_omni.diffusion.layers.adalayernorm import AdaLayerNorm
@@ -43,6 +40,55 @@ from vllm_omni.diffusion.layers.rope import RotaryEmbedding
 logger = init_logger(__name__)
 
 
+class ImageRopePrepare(nn.Module):
+    """Prepares image hidden_states and RoPE embeddings for sequence parallel.
+
+    This module encapsulates the input linear projection and RoPE computation.
+    Similar to Z-Image's UnifiedPrepare, this creates a module boundary where
+    _sp_plan can shard outputs via split_output=True.
+
+    The key insight is that hidden_states and vid_freqs must be sharded together
+    to maintain dimension alignment for RoPE computation in attention layers.
+
+    Note: Our _sp_plan corresponds to diffusers' _cp_plan (Context Parallelism).
+    """
+
+    def __init__(self, img_in: nn.Linear, pos_embed: nn.Module):
+        super().__init__()
+        self.img_in = img_in
+        self.pos_embed = pos_embed
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        img_shapes: list[tuple[int, int, int]],
+        txt_seq_lens: list[int],
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        """Prepare hidden_states and RoPE for SP.
+
+        Args:
+            hidden_states: [batch, img_seq_len, channels]
+            img_shapes: List of (frame, height, width) tuples
+            txt_seq_lens: List of text sequence lengths
+
+        Returns:
+            hidden_states: Processed hidden states [batch, img_seq_len, dim]
+            vid_freqs: Image RoPE frequencies [img_seq_len, rope_dim]
+            txt_freqs: Text RoPE frequencies [txt_seq_len, rope_dim]
+
+        Note: _sp_plan will shard hidden_states and vid_freqs via split_output=True
+              txt_freqs is kept replicated for dual-stream attention
+        """
+        # Apply input projection
+        hidden_states = self.img_in(hidden_states)
+
+        # Compute RoPE embeddings
+        image_rotary_emb = self.pos_embed(img_shapes, txt_seq_lens, device=hidden_states.device)
+        vid_freqs, txt_freqs = image_rotary_emb
+
+        return hidden_states, vid_freqs, txt_freqs
+
+
 class QwenTimestepProjEmbeddings(nn.Module):
     def __init__(self, embedding_dim, use_additional_t_cond=False):
         super().__init__()
@@ -436,6 +482,25 @@ class QwenImageCrossAttention(nn.Module):
         hidden_states_mask: torch.Tensor | None = None,
         encoder_hidden_states_mask: torch.Tensor | None = None,
     ) -> tuple[torch.Tensor, torch.Tensor]:
+        # Check for SP auto_pad: create attention mask dynamically if padding was applied
+        # In Ulysses mode, attention is computed on the FULL sequence (after All-to-All)
+        if (
+            self.parallel_config is not None
+            and self.parallel_config.sequence_parallel_size > 1
+            and hidden_states_mask is None
+        ):
+            ctx = get_forward_context()
+            if ctx.sp_original_seq_len is not None and ctx.sp_padding_size > 0:
+                # Create mask for the full (padded) sequence
+                # valid positions = True, padding positions = False
+                batch_size = hidden_states.shape[0]
+                padded_seq_len = ctx.sp_original_seq_len + ctx.sp_padding_size
+                hidden_states_mask = torch.ones(
+                    batch_size, padded_seq_len, dtype=torch.bool, device=hidden_states.device
+                )
+                hidden_states_mask[:, ctx.sp_original_seq_len :] = False
+
+        # if mask is all true, set it to None
         if hidden_states_mask is not None and hidden_states_mask.all():
             hidden_states_mask = None
         if encoder_hidden_states_mask is not None and encoder_hidden_states_mask.all():
@@ -718,6 +783,31 @@ class QwenImageTransformer2DModel(CachedTransformer):
     # used for torch compile optimizations
     _repeated_blocks = ["QwenImageTransformerBlock"]
 
+    # Sequence Parallelism plan (following diffusers' _cp_plan pattern)
+    # Similar to Z-Image's UnifiedPrepare, we use ImageRopePrepare to create
+    # a module boundary where _sp_plan can shard hidden_states and vid_freqs together.
+    #
+    # Key insight: hidden_states and vid_freqs MUST be sharded together to maintain
+    # dimension alignment for RoPE computation in attention layers.
+    #
+    # auto_pad=True enables automatic padding when sequence length is not divisible
+    # by SP world size. This creates an attention mask stored in ForwardContext
+    # that attention layers can use to ignore padding positions.
+    #
+    # Note: _sp_plan corresponds to diffusers' _cp_plan (Context Parallelism)
+    _sp_plan = {
+        # Shard ImageRopePrepare outputs (hidden_states and vid_freqs must be sharded together)
+        "image_rope_prepare": {
+            # hidden_states: auto_pad=True for variable sequence length support
+            0: SequenceParallelInput(split_dim=1, expected_dims=3, split_output=True, auto_pad=True),
+            # vid_freqs: auto_pad=True to match hidden_states padding
+            1: SequenceParallelInput(split_dim=0, expected_dims=2, split_output=True, auto_pad=True),
+            # txt_freqs (index 2) is NOT sharded - kept replicated for dual-stream attention
+        },
+        # Gather output at proj_out
+        "proj_out": SequenceParallelOutput(gather_dim=1, expected_dims=3),
+    }
+
     def __init__(
         self,
         od_config: OmniDiffusionConfig,
@@ -773,6 +863,10 @@ class QwenImageTransformer2DModel(CachedTransformer):
         self.gradient_checkpointing = False
         self.zero_cond_t = zero_cond_t
 
+        # ImageRopePrepare module for _sp_plan to shard hidden_states and vid_freqs together
+        # This ensures RoPE dimensions align with hidden_states after sharding
+        self.image_rope_prepare = ImageRopePrepare(self.img_in, self.pos_embed)
+
     def forward(
         self,
         hidden_states: torch.Tensor,
@@ -816,49 +910,17 @@ class QwenImageTransformer2DModel(CachedTransformer):
         # else:
         #     lora_scale = 1.0
 
-        original_seq_len = None
-        seq_padding = 0
-        hidden_states_mask = None
-
+        # Set split_text_embed_in_sp = False for dual-stream attention
+        # QwenImage uses *dual-stream* (text + image) and runs a *joint attention*.
+        # Text embeddings must be replicated across SP ranks for correctness.
         if self.parallel_config.sequence_parallel_size > 1:
-            batch_size, seq_len, channels = hidden_states.shape
-            sp_size = get_sequence_parallel_world_size()
-
-            if seq_len % sp_size != 0:
-                #  ring_attn, sage_attn do not support attention_mask
-                if get_attn_backend(-1).get_name() not in _BACKENDS_SUPPORT_ATTENTION_MASK:
-                    raise ValueError(
-                        f"When generating image shape that the sequence length is NOT divisible by sp_size={sp_size},"
-                        f"cannot use {get_attn_backend(-1).get_name()} which does not support attention_mask."
-                        f"Please switch to SDPA or Ascend attention backend."
-                    )
-                # ring attention does not support attention_mask
-                if get_ring_parallel_world_size() > 1:
-                    raise ValueError(
-                        f"When generating image shape that the sequence length is NOT divisible by sp_size={sp_size},"
-                        f"cannot use ring attention which does not support attention_mask."
-                        f"Please switch to Ulysses SP only."
-                    )
-
-                seq_padding = sp_size - (seq_len % sp_size)
-                original_seq_len = seq_len
-
-                hidden_states_mask = torch.ones(
-                    batch_size, seq_len + seq_padding, dtype=torch.bool, device=hidden_states.device
-                )
-                hidden_states_mask[:, seq_len:] = False
-                padding_tensor = torch.zeros(
-                    batch_size, seq_padding, channels, dtype=hidden_states.dtype, device=hidden_states.device
-                )
-                hidden_states = torch.cat([hidden_states, padding_tensor], dim=1)
-
-            hidden_states = torch.chunk(hidden_states, sp_size, dim=-2)[get_sequence_parallel_rank()]
-            # NOTE:
-            # QwenImage uses *dual-stream* (text + image) and runs a *joint attention*.
-            # text embeddings to be replicated across SP ranks for correctness.
             get_forward_context().split_text_embed_in_sp = False
 
-        hidden_states = self.img_in(hidden_states)
+        # Prepare hidden_states and RoPE via ImageRopePrepare module
+        # _sp_plan will shard hidden_states and vid_freqs together via split_output=True
+        # txt_freqs is kept replicated for dual-stream attention
+        hidden_states, vid_freqs, txt_freqs = self.image_rope_prepare(hidden_states, img_shapes, txt_seq_lens)
+        image_rotary_emb = (vid_freqs, txt_freqs)
 
         # Ensure timestep tensor is on the same device and dtype as hidden_states
         timestep = timestep.to(device=hidden_states.device, dtype=hidden_states.dtype)
@@ -885,23 +947,6 @@ class QwenImageTransformer2DModel(CachedTransformer):
             else self.time_text_embed(timestep, guidance, hidden_states, additional_t_cond)
         )
 
-        image_rotary_emb = self.pos_embed(img_shapes, txt_seq_lens, device=hidden_states.device)
-
-        def get_rotary_emb_chunk(freqs, padding=0):
-            # Pad rotary embeddings if needed
-            if padding > 0:
-                padding_tensor = torch.zeros(padding, freqs.shape[-1], dtype=freqs.dtype, device=freqs.device)
-                freqs = torch.cat([freqs, padding_tensor], dim=0)
-            freqs = torch.chunk(freqs, get_sequence_parallel_world_size(), dim=0)[get_sequence_parallel_rank()]
-            return freqs
-
-        if self.parallel_config.sequence_parallel_size > 1:
-            img_freqs, txt_freqs = image_rotary_emb
-            img_freqs = get_rotary_emb_chunk(img_freqs, seq_padding)
-            if get_forward_context().split_text_embed_in_sp:
-                txt_freqs = get_rotary_emb_chunk(txt_freqs)
-            image_rotary_emb = (img_freqs, txt_freqs)
-
         for index_block, block in enumerate(self.transformer_blocks):
             encoder_hidden_states, hidden_states = block(
                 hidden_states=hidden_states,
@@ -911,7 +956,6 @@ class QwenImageTransformer2DModel(CachedTransformer):
                 image_rotary_emb=image_rotary_emb,
                 joint_attention_kwargs=attention_kwargs,
                 modulate_index=modulate_index,
-                hidden_states_mask=hidden_states_mask,
             )
 
         if self.zero_cond_t:
@@ -920,14 +964,8 @@ class QwenImageTransformer2DModel(CachedTransformer):
         hidden_states = self.norm_out(hidden_states, temb)
         output = self.proj_out(hidden_states)
 
-        if self.parallel_config.sequence_parallel_size > 1:
-            output = get_sp_group().all_gather(output, dim=-2)
-
-            # Remove padding if it was added
-            if original_seq_len is not None:
-                output = output[:, :original_seq_len, :]
-
-        torch.cuda.empty_cache()
+        # Note: SP gather is handled automatically by _sp_plan's SequenceParallelGatherHook
+        # on proj_out output. No manual all_gather needed here.
 
         return Transformer2DModelOutput(sample=output)
 

vllm_omni/diffusion/models/z_image/pipeline_z_image.py

@@ -176,6 +176,9 @@ class ZImagePipeline(nn.Module):
         self.transformer = ZImageTransformer2DModel()
         self.tokenizer = AutoTokenizer.from_pretrained(model, subfolder="tokenizer", local_files_only=local_files_only)
 
+        # Note: Context parallelism is applied centrally in registry.initialize_model()
+        # following diffusers' pattern of enable_parallelism() at model loading time
+
         self.vae_scale_factor = (
             2 ** (len(self.vae.config.block_out_channels) - 1) if hasattr(self, "vae") and self.vae is not None else 8
         )

vllm_omni/diffusion/models/z_image/z_image_transformer.py

@@ -1,5 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# _sp_plan definition adapted from HuggingFace diffusers library (_cp_plan)
 
 # Copyright 2025 Alibaba Z-Image Team and The HuggingFace Team. All rights reserved.
 #
@@ -32,7 +33,14 @@ from vllm.model_executor.layers.linear import (
 from vllm.model_executor.model_loader.weight_utils import default_weight_loader
 
 from vllm_omni.diffusion.attention.layer import Attention
-from vllm_omni.diffusion.forward_context import get_forward_context, is_forward_context_available
+from vllm_omni.diffusion.distributed.sp_plan import (
+    SequenceParallelInput,
+    SequenceParallelOutput,
+)
+from vllm_omni.diffusion.forward_context import (
+    get_forward_context,
+    is_forward_context_available,
+)
 from vllm_omni.diffusion.layers.rope import RotaryEmbedding
 
 ADALN_EMBED_DIM = 256
@@ -41,6 +49,63 @@ SEQ_MULTI_OF = 32
 logger = init_logger(__name__)
 
 
+class UnifiedPrepare(nn.Module):
+    """Prepares unified tensors for transformer blocks.
+
+    This module encapsulates the unification of x and cap tensors into unified
+    sequences. Similar to how Wan's `rope` module outputs rotary embeddings,
+    this module outputs unified tensors that can be sharded via _sp_plan's
+    split_output=True mechanism.
+
+    This follows the diffusers pattern where tensor preparation happens in
+    a dedicated submodule, enabling _sp_plan hooks to work at module boundaries.
+    """
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        x_cos: torch.Tensor,
+        x_sin: torch.Tensor,
+        cap_feats: torch.Tensor,
+        cap_cos: torch.Tensor,
+        cap_sin: torch.Tensor,
+        x_item_seqlens: list[int],
+        cap_item_seqlens: list[int],
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+        """Combine x and cap tensors into unified sequences.
+
+        Returns:
+            unified: Combined hidden states [batch, seq_len, dim]
+            unified_cos: Combined RoPE cos [batch, seq_len, rope_dim]
+            unified_sin: Combined RoPE sin [batch, seq_len, rope_dim]
+            unified_attn_mask: Combined attention mask [batch, seq_len]
+        """
+        bsz = x.shape[0]
+        device = x.device
+
+        unified = []
+        unified_cos = []
+        unified_sin = []
+        for i in range(bsz):
+            x_len = x_item_seqlens[i]
+            cap_len = cap_item_seqlens[i]
+            unified.append(torch.cat([x[i][:x_len], cap_feats[i][:cap_len]]))
+            unified_cos.append(torch.cat([x_cos[i][:x_len], cap_cos[i][:cap_len]]))
+            unified_sin.append(torch.cat([x_sin[i][:x_len], cap_sin[i][:cap_len]]))
+
+        unified_item_seqlens = [a + b for a, b in zip(cap_item_seqlens, x_item_seqlens)]
+        unified_max_item_seqlen = max(unified_item_seqlens)
+
+        unified = pad_sequence(unified, batch_first=True, padding_value=0.0)
+        unified_cos = pad_sequence(unified_cos, batch_first=True, padding_value=0.0)
+        unified_sin = pad_sequence(unified_sin, batch_first=True, padding_value=0.0)
+        unified_attn_mask = torch.zeros((bsz, unified_max_item_seqlen), dtype=torch.bool, device=device)
+        for i, seq_len in enumerate(unified_item_seqlens):
+            unified_attn_mask[i, :seq_len] = 1
+
+        return unified, unified_cos, unified_sin, unified_attn_mask
+
+
 def _positive_divisors(n: int) -> set[int]:
     if n <= 0:
         return set()
@@ -456,8 +521,52 @@ class RopeEmbedder:
 
 
 class ZImageTransformer2DModel(nn.Module):
+    """Z-Image Transformer model for image generation.
+
+    Sequence Parallelism:
+        This model supports non-intrusive SP via _sp_plan. The plan specifies:
+        - Input splitting at first main transformer block (unified sequence)
+        - RoPE (cos/sin) splitting along sequence dimension
+        - Attention mask splitting along sequence dimension
+        - Output gathering at final_layer
+
+        The SP is applied to the main `layers` transformer blocks where the
+        unified image+caption sequence is processed jointly.
+
+        Note: noise_refiner and context_refiner are NOT parallelized as they
+        process image and caption separately before unification.
+
+        Important: The default _sp_plan assumes patch_size=2 and f_patch_size=1.
+        If using different patch configurations, update _sp_plan accordingly.
+
+        Note: Our "Sequence Parallelism" (SP) corresponds to "Context Parallelism" (CP) in diffusers.
+    """
+
     _repeated_blocks = ["ZImageTransformerBlock"]
 
+    # Sequence Parallelism for Z-Image (following diffusers' _cp_plan pattern)
+    # Similar to how Wan uses `rope` module's split_output to shard rotary embeddings,
+    # Z-Image uses `unified_prepare` module's split_output to shard unified tensors.
+    #
+    # The _sp_plan specifies sharding/gathering at module boundaries:
+    # - unified_prepare: Split all 4 outputs (unified, cos, sin, attn_mask) via split_output=True
+    # - layers.0: hidden_states input is already sharded from unified_prepare output
+    # - all_final_layer.2-1: Gather outputs after the final layer
+    #
+    # Note: _sp_plan corresponds to diffusers' _cp_plan (Context Parallelism)
+    _sp_plan = {
+        # Shard unified_prepare outputs (similar to Wan's rope module)
+        # This shards all 4 return values: unified, unified_cos, unified_sin, unified_attn_mask
+        "unified_prepare": {
+            0: SequenceParallelInput(split_dim=1, expected_dims=3, split_output=True),  # unified
+            1: SequenceParallelInput(split_dim=1, expected_dims=3, split_output=True),  # unified_cos
+            2: SequenceParallelInput(split_dim=1, expected_dims=3, split_output=True),  # unified_sin
+            3: SequenceParallelInput(split_dim=1, expected_dims=2, split_output=True),  # unified_attn_mask
+        },
+        # Gather output at final_layer (default: patch_size=2, f_patch_size=1)
+        "all_final_layer.2-1": SequenceParallelOutput(gather_dim=1, expected_dims=3),
+    }
+
     def __init__(
         self,
         all_patch_size=(2,),
@@ -572,6 +681,11 @@ class ZImageTransformer2DModel(nn.Module):
 
         self.rope_embedder = RopeEmbedder(theta=rope_theta, axes_dims=axes_dims, axes_lens=axes_lens)
 
+        # UnifiedPrepare module for combining x and cap tensors
+        # This enables _cp_plan to shard outputs via split_output=True
+        # Similar to how Wan's rope module enables rotary embedding sharding
+        self.unified_prepare = UnifiedPrepare()
+
     def unpatchify(self, x: list[torch.Tensor], size: list[tuple], patch_size, f_patch_size) -> list[torch.Tensor]:
         pH = pW = patch_size
         pF = f_patch_size
@@ -772,31 +886,19 @@ class ZImageTransformer2DModel(nn.Module):
         for layer in self.context_refiner:
             cap_feats = layer(cap_feats, cap_attn_mask, cap_cos, cap_sin)
 
-        # unified
-        unified = []
-        unified_cos = []
-        unified_sin = []
-        for i in range(bsz):
-            x_len = x_item_seqlens[i]
-            cap_len = cap_item_seqlens[i]
-            unified.append(torch.cat([x[i][:x_len], cap_feats[i][:cap_len]]))
-            unified_cos.append(torch.cat([x_cos[i][:x_len], cap_cos[i][:cap_len]]))
-            unified_sin.append(torch.cat([x_sin[i][:x_len], cap_sin[i][:cap_len]]))
-        unified_item_seqlens = [a + b for a, b in zip(cap_item_seqlens, x_item_seqlens)]
-        assert unified_item_seqlens == [len(_) for _ in unified]
-        unified_max_item_seqlen = max(unified_item_seqlens)
-
-        unified = pad_sequence(unified, batch_first=True, padding_value=0.0)
-        unified_cos = pad_sequence(unified_cos, batch_first=True, padding_value=0.0)
-        unified_sin = pad_sequence(unified_sin, batch_first=True, padding_value=0.0)
-        unified_attn_mask = torch.zeros((bsz, unified_max_item_seqlen), dtype=torch.bool, device=device)
-        for i, seq_len in enumerate(unified_item_seqlens):
-            unified_attn_mask[i, :seq_len] = 1
+        # Prepare unified tensors via UnifiedPrepare module
+        # This enables _cp_plan to shard outputs via split_output=True
+        unified, unified_cos, unified_sin, unified_attn_mask = self.unified_prepare(
+            x, x_cos, x_sin, cap_feats, cap_cos, cap_sin, x_item_seqlens, cap_item_seqlens
+        )
 
+        # Main transformer blocks
         for layer in self.layers:
             unified = layer(unified, unified_attn_mask, unified_cos, unified_sin, adaln_input)
 
+        # Final layer
         unified = self.all_final_layer[f"{patch_size}-{f_patch_size}"](unified, adaln_input)
+
         unified = list(unified.unbind(dim=0))
         x = self.unpatchify(unified, x_size, patch_size, f_patch_size)
 

vllm_omni/diffusion/registry.py

@@ -3,9 +3,15 @@
 
 import importlib
 
+import torch.nn as nn
+from vllm.logger import init_logger
 from vllm.model_executor.models.registry import _LazyRegisteredModel, _ModelRegistry
 
 from vllm_omni.diffusion.data import OmniDiffusionConfig
+from vllm_omni.diffusion.distributed.sp_plan import SequenceParallelConfig, get_sp_plan_from_model
+from vllm_omni.diffusion.hooks.sequence_parallel import apply_sequence_parallel
+
+logger = init_logger(__name__)
 
 _DIFFUSION_MODELS = {
     # arch:(mod_folder, mod_relname, cls_name)
@@ -100,7 +106,24 @@ DiffusionModelRegistry = _ModelRegistry(
 
 def initialize_model(
     od_config: OmniDiffusionConfig,
-):
+) -> nn.Module:
+    """Initialize a diffusion model from the registry.
+
+    This function:
+    1. Loads the model class from the registry
+    2. Instantiates the model with the config
+    3. Configures VAE optimization settings
+    4. Applies sequence parallelism if enabled (similar to diffusers' enable_parallelism)
+
+    Args:
+        od_config: The OmniDiffusion configuration.
+
+    Returns:
+        The initialized pipeline model.
+
+    Raises:
+        ValueError: If the model class is not found in the registry.
+    """
     model_class = DiffusionModelRegistry._try_load_model_cls(od_config.model_class_name)
     if model_class is not None:
         model = model_class(od_config=od_config)
@@ -110,11 +133,73 @@ def initialize_model(
         if hasattr(model.vae, "use_tiling"):
             model.vae.use_tiling = od_config.vae_use_tiling
 
+        # Apply sequence parallelism if enabled
+        # This follows diffusers' pattern where enable_parallelism() is called
+        # at model loading time, not inside individual model files
+        _apply_sequence_parallel_if_enabled(model, od_config)
+
         return model
     else:
         raise ValueError(f"Model class {od_config.model_class_name} not found in diffusion model registry.")
 
 
+def _apply_sequence_parallel_if_enabled(model, od_config: OmniDiffusionConfig) -> None:
+    """Apply sequence parallelism hooks if SP is enabled.
+
+    This is the centralized location for enabling SP, similar to diffusers'
+    ModelMixin.enable_parallelism() method. It applies _sp_plan hooks to
+    transformer models that define them.
+
+    Note: Our "Sequence Parallelism" (SP) corresponds to "Context Parallelism" (CP) in diffusers.
+    We use _sp_plan instead of diffusers' _cp_plan.
+
+    Args:
+        model: The pipeline model (e.g., ZImagePipeline).
+        od_config: The OmniDiffusion configuration.
+    """
+
+    try:
+        sp_size = od_config.parallel_config.sequence_parallel_size
+        if sp_size <= 1:
+            return
+
+        # Find transformer model(s) in the pipeline that have _sp_plan
+        transformer_attrs = ["transformer", "dit", "unet"]
+        for attr in transformer_attrs:
+            if not hasattr(model, attr):
+                continue
+
+            transformer = getattr(model, attr)
+            if transformer is None:
+                continue
+
+            plan = get_sp_plan_from_model(transformer)
+            if plan is None:
+                continue
+
+            # Create SP config
+            sp_config = SequenceParallelConfig(
+                ulysses_degree=od_config.parallel_config.ulysses_degree,
+                ring_degree=od_config.parallel_config.ring_degree,
+            )
+
+            # Apply hooks according to the plan
+            mode = (
+                "hybrid"
+                if sp_config.ulysses_degree > 1 and sp_config.ring_degree > 1
+                else ("ulysses" if sp_config.ulysses_degree > 1 else "ring")
+            )
+            logger.info(
+                f"Applying sequence parallelism to {transformer.__class__.__name__} "
+                f"(sp_size={sp_size}, mode={mode}, ulysses={sp_config.ulysses_degree}, ring={sp_config.ring_degree})"
+            )
+            apply_sequence_parallel(transformer, sp_config, plan)
+            return  # Only apply to first transformer found
+
+    except Exception as e:
+        logger.warning(f"Failed to apply sequence parallelism: {e}. Continuing without SP hooks.")
+
+
 _DIFFUSION_POST_PROCESS_FUNCS = {
     # arch: post_process_func
     # `post_process_func` function must be placed in {mod_folder}/{mod_relname}.py,