diff --git a/tests/test_diffusion_config_propagation.py b/tests/test_diffusion_config_propagation.py index 389dc9c98..0991c6601 100644 --- a/tests/test_diffusion_config_propagation.py +++ b/tests/test_diffusion_config_propagation.py @@ -62,6 +62,16 @@ class TestParallelConfigPropagation: assert od.parallel_config.ulysses_degree == 2 assert od.parallel_config.ring_degree == 1 + def test_mask_sp_padding_roundtrip(self): + pc = DiffusionParallelConfig(ulysses_degree=2, mask_sp_padding=True) + od = _roundtrip_diffusion_config(model="x", parallel_config=pc) + assert od.parallel_config.mask_sp_padding is True + + def test_mask_sp_padding_defaults_false(self): + pc = DiffusionParallelConfig(ulysses_degree=2) + od = _roundtrip_diffusion_config(model="x", parallel_config=pc) + assert od.parallel_config.mask_sp_padding is False + def test_cfg_parallel_roundtrip(self): pc = DiffusionParallelConfig(cfg_parallel_size=2) od = _roundtrip_diffusion_config(model="x", parallel_config=pc) diff --git a/vllm_omni/diffusion/data.py b/vllm_omni/diffusion/data.py index 4017ac984..00f3b0b5a 100644 --- a/vllm_omni/diffusion/data.py +++ b/vllm_omni/diffusion/data.py @@ -130,6 +130,14 @@ class DiffusionParallelConfig: use_hsdp: bool = False """Enable Hybrid Sharded Data Parallel (HSDP) for model weight sharding.""" + mask_sp_padding: bool = False + """If True, generate a boolean attention mask for zero-padded SP tokens + when sequence length is not divisible by the SP world size. The mask + routes attention through the varlen path (unpad→kernel→repad), which is + correct but carries additional overhead. When False (default), padding + tokens are left unmasked; since _shard_with_auto_pad always pads with + zeros, their contribution to attention output is negligible.""" + hsdp_shard_size: int = -1 """Number of GPUs to shard weights across within each replica group. -1 means auto-calculate.""" 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 8edf7ec19..cda01e083 100644 --- a/vllm_omni/diffusion/models/wan2_2/wan2_2_transformer.py +++ b/vllm_omni/diffusion/models/wan2_2/wan2_2_transformer.py @@ -1009,28 +1009,43 @@ class WanTransformer3DModel(nn.Module): if encoder_hidden_states_image is not None: encoder_hidden_states = torch.concat([encoder_hidden_states_image, encoder_hidden_states], dim=1) - # Check for SP auto_pad: create attention mask dynamically if padding was applied - hidden_states_mask = None # default - config = get_forward_context().omni_diffusion_config - parallel_config = config.parallel_config - if parallel_config is not None and parallel_config.sequence_parallel_size > 1: - ctx = get_forward_context() - if ctx.sp_original_seq_len is not None and ctx.sp_padding_size > 0: - # Create mask for the full (padded) sequence - # valid positions = True, padding positions = False - 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[:, ctx.sp_original_seq_len :] = False - - # if mask is all true, set it to None - if hidden_states_mask is not None and hidden_states_mask.all(): - hidden_states_mask = None + hidden_states_mask = None + ctx = get_forward_context() + parallel_config = ctx.omni_diffusion_config.parallel_config + if ( + 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[:, ctx.sp_original_seq_len :] = False + if hidden_states_mask.all(): + hidden_states_mask = None + elif ( + parallel_config is not None + and parallel_config.sequence_parallel_size > 1 + and not parallel_config.mask_sp_padding + and ctx.sp_original_seq_len is not None + and ctx.sp_padding_size > 0 + ): + logger.warning_once( + "SP auto-padding applied %d token(s) (seq_len=%d, ulysses_degree=%d). " + "Padding tokens are not masked from attention (mask_sp_padding=False), " + "which avoids the varlen attention path but may produce minor numerical differences. " + "Set parallel_config.mask_sp_padding=True to restore strict masking.", + ctx.sp_padding_size, + ctx.sp_original_seq_len, + parallel_config.sequence_parallel_size, + ) # Transformer blocks for block in self.blocks[self.start_layer : self.end_layer]: diff --git a/vllm_omni/diffusion/models/wan2_2/wan2_2_vace_transformer.py b/vllm_omni/diffusion/models/wan2_2/wan2_2_vace_transformer.py index 3f66a0d1a..871f0df86 100644 --- a/vllm_omni/diffusion/models/wan2_2/wan2_2_vace_transformer.py +++ b/vllm_omni/diffusion/models/wan2_2/wan2_2_vace_transformer.py @@ -10,6 +10,7 @@ from typing import Any import torch import torch.nn as nn import torch.nn.functional as F +from vllm.logger import init_logger from vllm.model_executor.layers.quantization.base_config import QuantizationConfig from vllm_omni.diffusion.distributed.sp_plan import SequenceParallelInput @@ -21,6 +22,8 @@ from vllm_omni.diffusion.models.wan2_2.wan2_2_transformer import ( WanTransformerBlock, ) +logger = init_logger(__name__) + class VaceWanTransformerBlock(WanTransformerBlock): """VACE variant of WanTransformerBlock with proj_in/proj_out for skip connections.""" @@ -214,12 +217,16 @@ class WanVACETransformer3DModel(WanTransformer3DModel): # Shard hidden_states via _sp_plan hook (before VACE, not at blocks.0) hidden_states = self._sp_shard_point(hidden_states) - # SP state and attention mask for padding hidden_states_mask = None ctx = get_forward_context() parallel_config = ctx.omni_diffusion_config.parallel_config sp_size = parallel_config.sequence_parallel_size if parallel_config is not None else 1 - if ctx.sp_original_seq_len is not None and ctx.sp_padding_size > 0: + if ( + parallel_config is not None + and parallel_config.mask_sp_padding + and ctx.sp_original_seq_len is not None + and ctx.sp_padding_size > 0 + ): padded_seq_len = ctx.sp_original_seq_len + ctx.sp_padding_size hidden_states_mask = torch.ones( batch_size, @@ -228,6 +235,21 @@ class WanVACETransformer3DModel(WanTransformer3DModel): device=hidden_states.device, ) hidden_states_mask[:, ctx.sp_original_seq_len :] = False + elif ( + parallel_config is not None + and not parallel_config.mask_sp_padding + and ctx.sp_original_seq_len is not None + and ctx.sp_padding_size > 0 + ): + logger.warning_once( + "SP auto-padding applied %d token(s) (seq_len=%d, ulysses_degree=%d). " + "Padding tokens are not masked from attention (mask_sp_padding=False), " + "which avoids the varlen attention path but may produce minor numerical differences. " + "Set parallel_config.mask_sp_padding=True to restore strict masking.", + ctx.sp_padding_size, + ctx.sp_original_seq_len, + sp_size, + ) # VACE: embed context and run conditioning blocks vace_hints = None