[Perf][Wan2.2] Skip attention mask for zero-padded SP sequences to avoid varlen path (#3763)

Signed-off-by: jjuvonen-amd <joni.juvonen@amd.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>

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)

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."""
 

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]:

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