Wan2.2: Fix the bug of using cache-dit with ulysses in t2v and i2v (#3927)

Signed-off-by: Mengkejiergeli Ba <mengkejiergeli.ba@intel.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>

tests/diffusion/distributed/test_sequence_parallel.py

@@ -23,6 +23,7 @@ from PIL import Image
 from tests.helpers.mark import hardware_test
 from tests.helpers.runtime import OmniRunner
 from vllm_omni.diffusion.data import DiffusionParallelConfig
+from vllm_omni.diffusion.distributed.utils import build_local_sp_padding_mask
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.platforms import current_omni_platform
 
@@ -412,3 +413,85 @@ def test_sp_correctness_advanced(model_name: str):
             f"{baseline_str:<12} {r['sp_ms']:.0f}ms{'':<7} {speedup_str:<10} {status}"
         )
     print("=" * 70)
+
+
+@pytest.mark.skipif(
+    not (current_omni_platform.is_cuda() or current_omni_platform.is_xpu()),
+    reason="Only tested on CUDA and XPU",
+)
+@pytest.mark.diffusion
+@pytest.mark.parallel
+@pytest.mark.core_model
+def test_local_sp_padding_mask(monkeypatch: pytest.MonkeyPatch) -> None:
+    """A partially padded SP shard must receive a local-length mask."""
+    mask = build_local_sp_padding_mask(
+        batch_size=2,
+        local_seq_len=4,
+        sp_original_seq_len=5,
+        sp_padding_size=3,
+        sequence_parallel_rank=1,
+        device=torch.device(current_omni_platform.device_type),
+    )
+
+    expected = torch.tensor(
+        [
+            [True, False, False, False],
+            [True, False, False, False],
+        ],
+        dtype=torch.bool,
+        device=mask.device,
+    )
+    assert mask is not None
+    assert mask.shape == (2, 4)
+    assert torch.equal(mask, expected)
+
+
+@pytest.mark.skipif(
+    not (current_omni_platform.is_cuda() or current_omni_platform.is_xpu()),
+    reason="Only tested on CUDA and XPU",
+)
+@pytest.mark.diffusion
+@pytest.mark.parallel
+@pytest.mark.core_model
+def test_local_sp_padding_mask_no_padding(monkeypatch: pytest.MonkeyPatch) -> None:
+    """A rank whose local shard contains no padding should not get a mask."""
+    mask = build_local_sp_padding_mask(
+        batch_size=2,
+        local_seq_len=4,
+        sp_original_seq_len=5,
+        sp_padding_size=3,
+        sequence_parallel_rank=0,
+        device=torch.device(current_omni_platform.device_type),
+    )
+
+    assert mask is None
+
+
+@pytest.mark.skipif(
+    not (current_omni_platform.is_cuda() or current_omni_platform.is_xpu()),
+    reason="Only tested on CUDA and XPU",
+)
+@pytest.mark.diffusion
+@pytest.mark.parallel
+@pytest.mark.core_model
+def test_wan_sp_plan() -> None:
+    """Wan2.2 must shard hidden states before CacheDiT-wrapped transformer blocks."""
+    try:
+        from vllm_omni.diffusion.distributed.sp_plan import validate_sp_plan
+        from vllm_omni.diffusion.models.wan2_2.wan2_2_transformer import WanTransformer3DModel
+    except ImportError as exc:
+        pytest.skip(f"WanTransformer3DModel not available: {exc}")
+
+    plan = getattr(WanTransformer3DModel, "_sp_plan", None)
+
+    assert plan is not None
+    assert "_sp_shard_point" in plan
+    assert "blocks.0" not in plan
+
+    shard_plan = plan["_sp_shard_point"]
+    assert 0 in shard_plan
+    assert shard_plan[0].split_dim == 1
+    assert shard_plan[0].expected_dims == 3
+    assert shard_plan[0].split_output is True
+
+    validate_sp_plan(plan)

vllm_omni/diffusion/distributed/utils.py

@@ -14,3 +14,26 @@ def get_local_device() -> torch.device:
         return current_omni_platform.get_torch_device(current_omni_platform.current_device())
     local_rank = int(os.environ.get("LOCAL_RANK", 0))
     return current_omni_platform.get_torch_device(local_rank)
+
+
+def build_local_sp_padding_mask(
+    batch_size: int,
+    local_seq_len: int,
+    sp_original_seq_len: int | None,
+    sp_padding_size: int,
+    sequence_parallel_rank: int,
+    device,
+):
+    """Build a per-rank SP padding mask that matches the local shard shape."""
+    if sp_original_seq_len is None or sp_padding_size <= 0:
+        return None
+
+    shard_start = sequence_parallel_rank * local_seq_len
+    shard_end = shard_start + local_seq_len
+    valid_tokens = max(0, min(sp_original_seq_len, shard_end) - shard_start)
+
+    if valid_tokens >= local_seq_len:
+        return None
+
+    token_mask = torch.arange(local_seq_len, device=device) < valid_tokens
+    return token_mask.unsqueeze(0).expand(batch_size, -1)

vllm_omni/diffusion/forward_context.py

@@ -4,6 +4,7 @@ from contextlib import contextmanager
 from dataclasses import dataclass
 from typing import TYPE_CHECKING
 
+import torch
 import vllm.ir
 from vllm.config import VllmConfig
 
@@ -88,6 +89,40 @@ def is_forward_context_available() -> bool:
     return _forward_context is not None
 
 
+def build_local_sp_padding_mask(
+    batch_size: int,
+    local_seq_len: int,
+    device,
+):
+    """Build a per-rank SP padding mask that matches the local shard shape.
+
+    Auto-padding is applied before sequence-parallel sharding, so attention on each
+    rank must receive a mask for its local shard, not for the global padded sequence.
+    """
+    if not is_forward_context_available():
+        return None
+
+    ctx = get_forward_context()
+    if ctx.sp_original_seq_len is None or ctx.sp_padding_size <= 0:
+        return None
+
+    from vllm_omni.diffusion.distributed.parallel_state import (
+        get_sequence_parallel_rank,
+    )
+    from vllm_omni.diffusion.distributed.utils import (
+        build_local_sp_padding_mask as build_local_sp_padding_mask_for_rank,
+    )
+
+    return build_local_sp_padding_mask_for_rank(
+        batch_size=batch_size,
+        local_seq_len=local_seq_len,
+        sp_original_seq_len=ctx.sp_original_seq_len,
+        sp_padding_size=ctx.sp_padding_size,
+        sequence_parallel_rank=get_sequence_parallel_rank(),
+        device=device,
+    )
+
+
 def get_ulysses_mode(*, default: str = "strict") -> str:
     """Resolve the Ulysses-SP mode from the current ForwardContext.
 

vllm_omni/diffusion/models/wan2_2/wan2_2_transformer.py

@@ -40,7 +40,7 @@ from vllm_omni.diffusion.distributed.sp_plan import (
     SequenceParallelInput,
     SequenceParallelOutput,
 )
-from vllm_omni.diffusion.forward_context import get_forward_context
+from vllm_omni.diffusion.forward_context import build_local_sp_padding_mask, get_forward_context
 from vllm_omni.diffusion.layers.adalayernorm import AdaLayerNorm
 from vllm_omni.diffusion.layers.norm import LayerNorm, RMSNorm
 from vllm_omni.diffusion.layers.rope import RotaryEmbeddingWan
@@ -799,7 +799,7 @@ class WanTransformer3DModel(nn.Module):
     # The _sp_plan specifies sharding/gathering at module boundaries:
     # - rope: Split both RoPE outputs (freqs_cos, freqs_sin) via split_output=True
     # - timestep_proj_prepare: Split timestep_proj for TI2V models (4D tensor)
-    # - blocks.0: Split hidden_states input at the first transformer block
+    # - _sp_shard_point: Split hidden_states before CacheDiT-wrapped transformer blocks
     # - proj_out: Gather outputs after the final projection layer
     #
     # Note: _sp_plan corresponds to diffusers' _cp_plan (Context Parallelism)
@@ -822,10 +822,10 @@ class WanTransformer3DModel(nn.Module):
                 split_dim=1, expected_dims=4, split_output=True, auto_pad=True
             ),  # [B, seq, 6, 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, auto_pad=True),  # [B, seq, dim]
+        # Shard hidden_states before entering transformer blocks. This keeps
+        # CacheDiT block wrappers aligned with the local SP shard.
+        "_sp_shard_point": {
+            0: SequenceParallelInput(split_dim=1, expected_dims=3, split_output=True, auto_pad=True),
         },
         # Shard output scale/shift for TI2V (3D); T2V outputs 2D and skips sharding
         "output_scale_shift_prepare": {
@@ -856,7 +856,6 @@ class WanTransformer3DModel(nn.Module):
         quant_config: QuantizationConfig | None = None,
     ):
         super().__init__()
-
         # Store config for compatibility
         self.config = type(
             "Config",
@@ -938,6 +937,7 @@ class WanTransformer3DModel(nn.Module):
 
         # SP helper modules
         self.timestep_proj_prepare = TimestepProjPrepare()
+        self._sp_shard_point = nn.Identity()
         if is_pipeline_last_stage():
             self.output_scale_shift_prepare = OutputScaleShiftPrepare(inner_dim)
         else:
@@ -981,10 +981,11 @@ class WanTransformer3DModel(nn.Module):
             self._cached_rope_emb = rotary_emb
 
         if is_pipeline_first_stage():
-            # Patch embedding and flatten to sequence
-            # (hidden_states is sharded at blocks.0 input by _sp_plan)
+            # Patch embedding and flatten to sequence. SP sharding happens at
+            # _sp_shard_point so downstream block wrappers see local tensors.
             hidden_states = self.patch_embedding(hidden_states)
             hidden_states = hidden_states.flatten(2).transpose(1, 2)
+            hidden_states = self._sp_shard_point(hidden_states)
         else:
             if intermediate_tensors is None:
                 raise RuntimeError("intermediate_tensors must be provided for non-first PP stages")
@@ -1017,18 +1018,12 @@ class WanTransformer3DModel(nn.Module):
             parallel_config is not None
             and parallel_config.sequence_parallel_size > 1
             and parallel_config.mask_sp_padding
-            and ctx.sp_original_seq_len is not None
-            and ctx.sp_padding_size > 0
         ):
-            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 = build_local_sp_padding_mask(
+                hidden_states.shape[0],
+                hidden_states.shape[1],
+                hidden_states.device,
             )
-            hidden_states_mask[:, ctx.sp_original_seq_len :] = False
             if hidden_states_mask.all():
                 hidden_states_mask = None
         elif (