diff --git a/tests/diffusion/distributed/test_sequence_parallel.py b/tests/diffusion/distributed/test_sequence_parallel.py index 9f76b3b75..c02667b93 100644 --- a/tests/diffusion/distributed/test_sequence_parallel.py +++ b/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) diff --git a/vllm_omni/diffusion/distributed/utils.py b/vllm_omni/diffusion/distributed/utils.py index 851328e62..fd4310e3f 100644 --- a/vllm_omni/diffusion/distributed/utils.py +++ b/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) diff --git a/vllm_omni/diffusion/forward_context.py b/vllm_omni/diffusion/forward_context.py index f2b6927d2..ea230ed9b 100644 --- a/vllm_omni/diffusion/forward_context.py +++ b/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. diff --git a/vllm_omni/diffusion/models/wan2_2/wan2_2_transformer.py b/vllm_omni/diffusion/models/wan2_2/wan2_2_transformer.py index 3f4d3ec9a..d61d1c82b 100644 --- a/vllm_omni/diffusion/models/wan2_2/wan2_2_transformer.py +++ b/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 (