[Diffusion][Feature] Non-Intrusive Sequence Parallelism (SP) Support for Wan2.2 (#966)

Signed-off-by: mxuax <mxuax@connect.ust.hk>
Signed-off-by: XU Mingshi <91017482+mxuax@users.noreply.github.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>

docs/user_guide/diffusion/parallelism_acceleration.md

@@ -48,7 +48,7 @@ The following table shows which models are currently supported by parallelism me
 
 | Model | Model Identifier | Ulysses-SP | Ring-SP | Tensor-Parallel |
 |-------|------------------|------------|---------|--------------------------|
-| **Wan2.2** | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | ❌ | ❌ | ❌ |
+| **Wan2.2** | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | ✅ | ✅ | ❌ |
 
 ### Tensor Parallelism
 

examples/offline_inference/text_to_video/text_to_video.py

@@ -3,11 +3,13 @@
 
 import argparse
 import os
+import time
 from pathlib import Path
 
 import numpy as np
 import torch
 
+from vllm_omni.diffusion.data import DiffusionParallelConfig
 from vllm_omni.entrypoints.omni import Omni
 from vllm_omni.outputs import OmniRequestOutput
 from vllm_omni.utils.platform_utils import detect_device_type, is_npu
@@ -40,6 +42,23 @@ def parse_args() -> argparse.Namespace:
         action="store_true",
         help="Enable CPU offloading for diffusion models.",
     )
+    parser.add_argument(
+        "--ulysses_degree",
+        type=int,
+        default=1,
+        help="Number of GPUs used for ulysses sequence parallelism.",
+    )
+    parser.add_argument(
+        "--ring_degree",
+        type=int,
+        default=1,
+        help="Number of GPUs used for ring sequence parallelism.",
+    )
+    parser.add_argument(
+        "--enforce_eager",
+        action="store_true",
+        help="Disable torch.compile and force eager execution.",
+    )
     return parser.parse_args()
 
 
@@ -52,6 +71,13 @@ def main():
     vae_use_slicing = is_npu()
     vae_use_tiling = is_npu()
 
+    # Configure parallel settings (only SP is supported for Wan)
+    # Note: cfg_parallel and tensor_parallel are not implemented for Wan models
+    parallel_config = DiffusionParallelConfig(
+        ulysses_degree=args.ulysses_degree,
+        ring_degree=args.ring_degree,
+    )
+
     # Check if profiling is requested via environment variable
     profiler_enabled = bool(os.getenv("VLLM_TORCH_PROFILER_DIR"))
 
@@ -62,12 +88,25 @@ def main():
         boundary_ratio=args.boundary_ratio,
         flow_shift=args.flow_shift,
         enable_cpu_offload=args.enable_cpu_offload,
+        parallel_config=parallel_config,
+        enforce_eager=args.enforce_eager,
     )
 
     if profiler_enabled:
         print("[Profiler] Starting profiling...")
         omni.start_profile()
 
+    # Print generation configuration
+    print(f"\n{'=' * 60}")
+    print("Generation Configuration:")
+    print(f"  Model: {args.model}")
+    print(f"  Inference steps: {args.num_inference_steps}")
+    print(f"  Frames: {args.num_frames}")
+    print(f"  Parallel configuration: ulysses_degree={args.ulysses_degree}, ring_degree={args.ring_degree}")
+    print(f"  Video size: {args.width}x{args.height}")
+    print(f"{'=' * 60}\n")
+
+    generation_start = time.perf_counter()
     frames = omni.generate(
         args.prompt,
         negative_prompt=args.negative_prompt,
@@ -79,6 +118,11 @@ def main():
         num_inference_steps=args.num_inference_steps,
         num_frames=args.num_frames,
     )
+    generation_end = time.perf_counter()
+    generation_time = generation_end - generation_start
+
+    # Print profiling results
+    print(f"Total generation time: {generation_time:.4f} seconds ({generation_time * 1000:.2f} ms)")
 
     # Extract video frames from OmniRequestOutput
     if isinstance(frames, list) and len(frames) > 0:

vllm_omni/diffusion/models/wan2_2/wan2_2_transformer.py

@@ -17,6 +17,10 @@ from vllm.model_executor.layers.linear import QKVParallelLinear, ReplicatedLinea
 from vllm.model_executor.model_loader.weight_utils import default_weight_loader
 
 from vllm_omni.diffusion.attention.layer import Attention
+from vllm_omni.diffusion.distributed.sp_plan import (
+    SequenceParallelInput,
+    SequenceParallelOutput,
+)
 
 logger = init_logger(__name__)
 
@@ -497,6 +501,16 @@ class WanTransformer3DModel(nn.Module):
     This is an optimized version of the diffusers WanTransformer3DModel that uses
     vLLM's efficient QKVParallelLinear and RMSNorm implementations.
 
+    Sequence Parallelism:
+        This model supports non-intrusive SP via _sp_plan. The plan specifies:
+        - RoPE (cos/sin) splitting via rope module's split_output
+        - hidden_states splitting at first transformer block input
+        - Output gathering at proj_out layer
+
+        The video sequence (flattened patches) is parallelized across GPUs.
+
+        Note: Our "Sequence Parallelism" (SP) corresponds to "Context Parallelism" (CP) in diffusers.
+
     Args:
         patch_size: 3D patch dimensions for video embedding (t_patch, h_patch, w_patch)
         num_attention_heads: Number of attention heads
@@ -520,6 +534,29 @@ class WanTransformer3DModel(nn.Module):
         "to_qkv": ["to_q", "to_k", "to_v"],
     }
 
+    # Sequence Parallelism for Wan (following diffusers' _cp_plan pattern)
+    #
+    # The _sp_plan specifies sharding/gathering at module boundaries:
+    # - rope: Split both RoPE outputs (freqs_cos, freqs_sin) via split_output=True
+    # - blocks.0: Split hidden_states input at the first transformer block
+    # - proj_out: Gather outputs after the final projection layer
+    #
+    # Note: _sp_plan corresponds to diffusers' _cp_plan (Context Parallelism)
+    _sp_plan = {
+        # Shard RoPE embeddings after rope module computes them
+        "rope": {
+            0: SequenceParallelInput(split_dim=1, expected_dims=4, split_output=True),  # freqs_cos [1, seq, 1, dim]
+            1: SequenceParallelInput(split_dim=1, expected_dims=4, split_output=True),  # freqs_sin [1, seq, 1, dim]
+        },
+        # Shard hidden_states at first transformer block input
+        # (after patch_embedding + flatten + transpose)
+        "blocks.0": {
+            "hidden_states": SequenceParallelInput(split_dim=1, expected_dims=3),  # [B, seq, dim]
+        },
+        # Gather at proj_out (final linear projection before unpatchify)
+        "proj_out": SequenceParallelOutput(gather_dim=1, expected_dims=3),
+    }
+
     def __init__(
         self,
         patch_size: tuple[int, int, int] = (1, 2, 2),
@@ -613,8 +650,11 @@ class WanTransformer3DModel(nn.Module):
         post_patch_height = height // p_h
         post_patch_width = width // p_w
 
+        # Compute RoPE embeddings (sharded by _sp_plan via split_output=True)
         rotary_emb = self.rope(hidden_states)
 
+        # Patch embedding and flatten to sequence
+        # (hidden_states is sharded at blocks.0 input by _sp_plan)
         hidden_states = self.patch_embedding(hidden_states)
         hidden_states = hidden_states.flatten(2).transpose(1, 2)
 

vllm_omni/diffusion/registry.py

@@ -164,7 +164,10 @@ def _apply_sequence_parallel_if_enabled(model, od_config: OmniDiffusionConfig) -
             return
 
         # Find transformer model(s) in the pipeline that have _sp_plan
-        transformer_attrs = ["transformer", "dit", "unet"]
+        # Include transformer_2 for two-stage models (e.g., Wan MoE)
+        transformer_attrs = ["transformer", "transformer_2", "dit", "unet"]
+        applied_count = 0
+
         for attr in transformer_attrs:
             if not hasattr(model, attr):
                 continue
@@ -190,11 +193,17 @@ def _apply_sequence_parallel_if_enabled(model, od_config: OmniDiffusionConfig) -
                 else ("ulysses" if sp_config.ulysses_degree > 1 else "ring")
             )
             logger.info(
-                f"Applying sequence parallelism to {transformer.__class__.__name__} "
+                f"Applying sequence parallelism to {transformer.__class__.__name__} ({attr}) "
                 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
+            applied_count += 1
+
+        if applied_count == 0:
+            logger.warning(
+                f"Sequence parallelism is enabled (sp_size={sp_size}) but no transformer with _sp_plan found. "
+                "SP hooks not applied. Consider adding _sp_plan to your transformer model."
+            )
 
     except Exception as e:
         logger.warning(f"Failed to apply sequence parallelism: {e}. Continuing without SP hooks.")