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>

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)