[S2V] Add sequence parallelism support to WanS2VTransformer3DModel (#4276)

Signed-off-by: Chendi Xue <chendi.xue@intel.com>

docker/Dockerfile.xpu

@@ -101,7 +101,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \
     source /opt/intel/oneapi/ccl/2021.15/env/vars.sh --force && \
     export CMAKE_PREFIX_PATH="$(python3 -c 'import site; print(site.getsitepackages()[0])'):${CMAKE_PREFIX_PATH}"
 
-ENV LD_LIBRARY_PATH="$LD_LIBRARY_PATH:/usr/local/lib/"
+ENV LD_LIBRARY_PATH="/opt/intel/oneapi/ccl/2021.15/lib:$LD_LIBRARY_PATH:/usr/local/lib/"
 
 ENV VLLM_TARGET_DEVICE=xpu
 ENV VLLM_WORKER_MULTIPROC_METHOD=spawn

docs/models/supported_models.md

@@ -37,7 +37,7 @@ th {
 | `WanPipeline` | Wan2.1-T2V, Wan2.2-T2V, Wan2.2-TI2V | `Wan-AI/Wan2.1-T2V-1.3B-Diffusers`, `Wan-AI/Wan2.1-T2V-14B-Diffusers`, `Wan-AI/Wan2.2-T2V-A14B-Diffusers`, `Wan-AI/Wan2.2-TI2V-5B-Diffusers` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
 | `WanImageToVideoPipeline` | Wan2.2-I2V | `Wan-AI/Wan2.2-I2V-A14B-Diffusers` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
 | `Cosmos3OmniDiffusersPipeline` | Cosmos3 T2I, T2V, I2V, V2V, T2V with sound, action policy | `nvidia/Cosmos3-Nano`, `nvidia/Cosmos3-Super` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
-| `WanSpeechToVideoPipeline` | Wan2.2-S2V | `Wan-AI/Wan2.2-S2V-14B` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
+| `Wan22S2VPipeline` | Wan2.2-S2V | `Wan-AI/Wan2.2-S2V-14B` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
 | `Wan22VACEPipeline` | Wan2.1-VACE | `Wan-AI/Wan2.1-VACE-1.3B-diffusers`, `Wan-AI/Wan2.1-VACE-14B-diffusers` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
 | `LTX2Pipeline` | LTX-2-T2V | `Lightricks/LTX-2` | ✅︎ | ✅︎ | | |
 | `LTX2ImageToVideoPipeline` | LTX-2-I2V | `Lightricks/LTX-2` | ✅︎ | ✅︎ | | |

docs/user_guide/diffusion_features.md

@@ -146,7 +146,7 @@ The following tables show which models support each feature:
 | Model                        | ⚡TeaCache | ⚡Cache-DiT | 🔀SP (Ulysses & Ring) | 🔀CFG-Parallel | 🔀Tensor-Parallel | Pipeline-Parallel | 🔀HSDP | 💾CPU Offload (Layerwise) | 💾VAE-Patch-Parallel | 💾Quantization | 🔄Step Execution |
 |------------------------------|:---------:|:----------:|:---------------------:|:--------------:|:-----------------:|:-----------------:|:------:|:-------------------------:|:--------------------:|:--------------:|:----------------:|
 | **Wan2.2**                   |     ❌     |     ✅      |           ✅           |       ✅        |         ✅         |         ✅         |   ✅    |             ✅             |  ✅ (encode/decode)   |       ❌        |        ❌         |
-| **Wan2.2-S2V**               |     ❌     |     ✅      |           ❌           |       ❌        |         ✅         |         ❌         |   ❌    |             ✅             |  ✅ (encode/decode)   |       ❌        |        ❌         |
+| **Wan2.2-S2V**               |     ❌     |     ✅      |           ✅           |       ✅        |         ✅         |         ❌         |   ❌    |             ✅             |  ✅ (encode/decode)   |       ❌        |        ❌         |
 | **Wan2.1-VACE**              |     ❌     |     ✅      |           ✅           |       ✅        |         ✅         |         ❌         |   ✅    |             ✅             |      ✅ (decode)      |       ❌        |        ❌         |
 | **LTX-2**                    |     ❌     |     ✅      |           ✅           |       ✅        |         ✅         |         ❌         |   ✅    |             ✅             |          ❌           |       ❌        |        ❌         |
 | **LTX-2.3**                  |     ❌     |     ✅      |           ✅           |       ✅        |         ✅         |         ❌         |   ❌    |             ✅             |      ✅ (decode)      |       ❌        |        ❌         |

vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_s2v.py

@@ -606,13 +606,17 @@ class Wan22S2VPipeline(
 
         t5_checkpoint = os.path.join(model_path, "models_t5_umt5-xxl-enc-bf16.pth")
         self.text_encoder = UMT5EncoderModel(_WAN_UMT5_CONFIG)
-        self.text_encoder = _load_wan_t5_as_umt5(self.text_encoder, t5_checkpoint, dtype=dtype).to(self.device)
+        _cpu_offload = self.od_config.enable_cpu_offload or self.od_config.enable_layerwise_offload
+        self.text_encoder = _load_wan_t5_as_umt5(self.text_encoder, t5_checkpoint, dtype=dtype)
+        if not _cpu_offload:
+            self.text_encoder = self.text_encoder.to(self.device)
 
         # -- VAE (original Wan2.1 VAE, loaded via diffusers from_single_file) --
         vae_pth = os.path.join(model_path, "Wan2.1_VAE.pth")
         self.vae = DistributedAutoencoderKLWan.from_single_file(vae_pth, torch_dtype=dtype)
         self.vae.init_distributed()
-        self.vae = self.vae.to(self.device)
+        if not _cpu_offload:
+            self.vae = self.vae.to(self.device)
 
         # -- Audio encoder (wav2vec2) --
         wav2vec_path = os.path.join(model_path, "wav2vec2-large-xlsr-53-english")

vllm_omni/diffusion/models/wan2_2/wan2_2_s2v_transformer.py

@@ -31,6 +31,7 @@ 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.distributed.sp_plan import SequenceParallelInput, SequenceParallelOutput
 from vllm_omni.diffusion.layers.rope import RotaryEmbeddingS2VGrid, RotaryEmbeddingWanS2V, WanS2VRotaryPosEmbed
 from vllm_omni.platforms import current_omni_platform
 
@@ -1058,6 +1059,32 @@ def _sample_indices(num_frames, stride, expand_ratio, c):
     return indices
 
 
+# ---------------------------------------------------------------------------
+# SP helper modules
+# ---------------------------------------------------------------------------
+
+
+class WanS2VBlocksPrepare(nn.Module):
+    """Pass-through module providing SP split hook point before transformer blocks.
+
+    Returns (hidden_states, freqs) unchanged. The SP hook shards both outputs
+    via split_output=True along the sequence dimension.
+    """
+
+    def forward(self, hidden_states: torch.Tensor, freqs: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        return hidden_states, freqs
+
+
+class WanS2VBlocksOutputGather(nn.Module):
+    """Pass-through module providing SP gather hook point after transformer blocks.
+
+    Returns hidden_states unchanged. The SP gather hook collects from all ranks.
+    """
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        return hidden_states
+
+
 # ---------------------------------------------------------------------------
 # Main model
 # ---------------------------------------------------------------------------
@@ -1087,6 +1114,17 @@ class WanS2VTransformer3DModel(nn.Module):
 
     _hsdp_shard_conditions = [_is_transformer_block]
 
+    _sp_plan = {
+        # Shard both hidden_states and freqs via blocks_prepare split_output
+        # hidden_states: [B, seq, dim], freqs: [B, seq, 1, head_dim//2] (complex)
+        "blocks_prepare": {
+            0: SequenceParallelInput(split_dim=1, expected_dims=3, split_output=True, auto_pad=True),
+            1: SequenceParallelInput(split_dim=1, expected_dims=4, split_output=True, auto_pad=True),
+        },
+        # Gather hidden_states after transformer blocks
+        "blocks_output_gather": SequenceParallelOutput(gather_dim=1, expected_dims=3),
+    }
+
     def __init__(
         self,
         cond_dim=0,
@@ -1157,6 +1195,10 @@ class WanS2VTransformer3DModel(nn.Module):
             ]
         )
 
+        # SP hook modules (pass-through; hooks shard/gather at these boundaries)
+        self.blocks_prepare = WanS2VBlocksPrepare()
+        self.blocks_output_gather = WanS2VBlocksOutputGather()
+
         # Output head
         self.norm_out = FP32LayerNorm(dim, eps, elementwise_affine=False)
         self.proj_out = nn.Linear(dim, math.prod(patch_size) * out_dim)
@@ -1396,8 +1438,21 @@ class WanS2VTransformer3DModel(nn.Module):
             ]
         return x, seq_lens, rope_embs, mask_input
 
+    def _is_sp_active(self) -> bool:
+        from vllm_omni.diffusion.forward_context import get_forward_context, is_forward_context_available
+
+        if not is_forward_context_available():
+            return False
+        return get_forward_context().sp_active
+
     def after_transformer_block(self, block_idx, hidden_states):
         if block_idx in self.audio_injector.injected_block_id.keys():
+            sp_active = self._is_sp_active()
+            if sp_active:
+                from vllm_omni.diffusion.distributed.sp_sharding import sp_gather, sp_shard
+
+                hidden_states = sp_gather(hidden_states, dim=1, validate=False)
+
             audio_attn_id = self.audio_injector.injected_block_id[block_idx]
             audio_emb = self.merged_audio_emb
             num_frames = audio_emb.shape[1]
@@ -1424,6 +1479,9 @@ class WanS2VTransformer3DModel(nn.Module):
             residual_out = rearrange(residual_out, "(b t) n c -> b (t n) c", t=num_frames)
             hidden_states[:, : self.original_seq_len].add_(residual_out)
 
+            if sp_active:
+                hidden_states = sp_shard(hidden_states, dim=1, validate=False)
+
         return hidden_states
 
     # ------------------------------------------------------------------
@@ -1574,6 +1632,21 @@ class WanS2VTransformer3DModel(nn.Module):
 
         context_lens = None
 
+        # SP: shard hidden_states and freqs before entering blocks
+        x, self.pre_compute_freqs = self.blocks_prepare(x, self.pre_compute_freqs)
+
+        # SP: adjust segment boundary for the local shard
+        if self._is_sp_active():
+            from vllm_omni.diffusion.distributed.parallel_state import (
+                get_sequence_parallel_rank,
+            )
+
+            sp_rank = get_sequence_parallel_rank()
+            chunk_size = x.shape[1]
+            global_start = sp_rank * chunk_size
+            local_seg_boundary = max(0, min(e0[1] - global_start, chunk_size))
+            e0 = [e0[0], local_seg_boundary]
+
         # Transformer blocks
         kwargs = dict(
             e=e0,
@@ -1588,6 +1661,9 @@ class WanS2VTransformer3DModel(nn.Module):
             x = block(x, **kwargs)
             x = self.after_transformer_block(idx, x)
 
+        # SP: gather hidden_states from all ranks
+        x = self.blocks_output_gather(x)
+
         # Output norm, projection & unpatchify
         x = x[:, : self.original_seq_len]
         shift, scale = self.output_scale_shift_prepare(e)