[Bugfix] Solve Ulysses-SP sequence length not divisible by SP degree (using padding and attention mask) (#672)

Signed-off-by: Didan Deng <[email protected]>

Author: wtomin

tests/e2e/offline_inference/test_sequence_parallel.py

@@ -202,3 +202,114 @@ def test_sequence_parallel(
         f"(thresholds: mean<={mean_threshold:.6e}, max<={max_threshold:.6e}); "
         f"ulysses_degree={ulysses_degree}, ring_degree={ring_degree}, dtype={dtype}"
     )
+
+
+@pytest.mark.parametrize("model_name", models)
+@pytest.mark.parametrize("dtype", [torch.bfloat16])
+@pytest.mark.parametrize("attn_backend", ["sdpa"])
+def test_sequence_parallel_ulysses_sp_only(
+    model_name: str,
+    dtype: torch.dtype,
+    attn_backend: str,
+):
+    """Test sequence parallel with ulysses_degree=4, ring_degree=1, and the image size (332x332) where the sequence length is NOT divisible by sp_size."""
+    ulysses_degree = 4
+    ring_degree = 1
+
+    # Skip if not enough GPUs available for SP run
+    if device_count() < ulysses_degree * ring_degree:
+        pytest.skip(f"Test requires {ulysses_degree * ring_degree} GPUs but only {device_count()} available")
+
+    # (272/8) * (272/8) = 17 * 17 = 289 Not divisible by sp_size=4
+    height = 272
+    width = 272
+    num_inference_steps = 4  # Minimal steps for fast test
+    seed = 42
+
+    # Step 1: Baseline (no Ulysses sequence parallel)
+    baseline_parallel_config = DiffusionParallelConfig(ulysses_degree=1, ring_degree=1)
+    baseline = Omni(
+        model=model_name,
+        parallel_config=baseline_parallel_config,
+        dtype=dtype,
+        attention_backend=attn_backend,
+    )
+    try:
+        outputs = baseline.generate(
+            PROMPT,
+            height=height,
+            width=width,
+            num_inference_steps=num_inference_steps,
+            guidance_scale=0.0,
+            generator=torch.Generator(get_device_name()).manual_seed(seed),
+            num_outputs_per_prompt=1,
+        )
+        baseline_images = outputs[0].request_output[0].images
+    finally:
+        baseline.close()
+        if dist.is_initialized():
+            dist.destroy_process_group()
+        for key in ["MASTER_ADDR", "MASTER_PORT", "RANK", "WORLD_SIZE", "LOCAL_RANK"]:
+            os.environ.pop(key, None)
+        time.sleep(2)  # Wait for resources to release
+
+    assert baseline_images is not None
+    assert len(baseline_images) == 1
+    assert baseline_images[0].width == width
+    assert baseline_images[0].height == height
+
+    # Step 2: SP (Ulysses-SP + Ring-SP)
+    sp_parallel_config = DiffusionParallelConfig(ulysses_degree=ulysses_degree, ring_degree=ring_degree)
+    sp = Omni(
+        model=model_name,
+        parallel_config=sp_parallel_config,
+        dtype=dtype,
+        attention_backend=attn_backend,
+    )
+    try:
+        outputs = sp.generate(
+            PROMPT,
+            height=height,
+            width=width,
+            num_inference_steps=num_inference_steps,
+            guidance_scale=0.0,
+            generator=torch.Generator(get_device_name()).manual_seed(seed),
+            num_outputs_per_prompt=1,
+        )
+        sp_images = outputs[0].request_output[0].images
+    finally:
+        sp.close()
+        if dist.is_initialized():
+            dist.destroy_process_group()
+        for key in ["MASTER_ADDR", "MASTER_PORT", "RANK", "WORLD_SIZE", "LOCAL_RANK"]:
+            os.environ.pop(key, None)
+        time.sleep(2)
+
+    assert sp_images is not None
+    assert len(sp_images) == 1
+    assert sp_images[0].width == width
+    assert sp_images[0].height == height
+
+    # Step 3: Compare outputs
+    mean_abs_diff, max_abs_diff = _diff_metrics(baseline_images[0], sp_images[0])
+
+    # FP16/BF16 may differ slightly due to different computation order under parallelism.
+    if dtype in (torch.float16, torch.bfloat16):
+        mean_threshold = 2e-2
+        max_threshold = 2e-1
+    else:
+        mean_threshold = 1e-2
+        max_threshold = 1e-1
+
+    print(
+        "Image diff stats (baseline ulysses_degree*ring_degree=1 vs SP): "
+        f"mean_abs_diff={mean_abs_diff:.6e}, max_abs_diff={max_abs_diff:.6e}; "
+        f"thresholds: mean<={mean_threshold:.6e}, max<={max_threshold:.6e}; "
+        f"ulysses_degree={ulysses_degree}, ring_degree={ring_degree}, dtype={dtype}"
+    )
+
+    assert mean_abs_diff <= mean_threshold and max_abs_diff <= max_threshold, (
+        f"Image diff exceeded threshold: mean_abs_diff={mean_abs_diff:.6e}, max_abs_diff={max_abs_diff:.6e} "
+        f"(thresholds: mean<={mean_threshold:.6e}, max<={max_threshold:.6e}); "
+        f"ulysses_degree={ulysses_degree}, ring_degree={ring_degree}, dtype={dtype}"
+    )

vllm_omni/diffusion/attention/backends/abstract.py

@@ -48,6 +48,8 @@ def supports_head_size(cls, head_size: int) -> bool:
 @dataclass
 class AttentionMetadata:
     attn_mask: torch.Tensor | None = None
+    joint_attn_mask: torch.Tensor | None = None
+    # a joint mask for the joint query, key, and value, depends the joint_strategy
     joint_query: torch.Tensor | None = None
     # a replicated tensor among processes appended to the front or rear of query, depends the joint_strategy
     joint_key: torch.Tensor | None = None

vllm_omni/diffusion/attention/parallel/ulysses.py

@@ -163,6 +163,36 @@ def pre_attention(
             joint_len=joint_len,
             joint_strategy=joint_strategy,
         )
+
+        if attn_metadata is not None:
+            if is_joint:
+                if attn_metadata.joint_attn_mask is None and attn_metadata.attn_mask is None:
+                    attn_metadata.attn_mask = None
+                else:
+                    if attn_metadata.attn_mask is None:
+                        attn_metadata.attn_mask = torch.ones(
+                            [query.shape[0], query.shape[1] - attn_metadata.joint_attn_mask.shape[1]],
+                            dtype=torch.bool,
+                            device=query.device,
+                        )
+                    elif attn_metadata.joint_attn_mask is None:
+                        attn_metadata.joint_attn_mask = torch.ones(
+                            [query.shape[0], query.shape[1] - attn_metadata.attn_mask.shape[1]],
+                            dtype=torch.bool,
+                            device=query.device,
+                        )
+                    attn_metadata.attn_mask = (
+                        torch.cat([attn_metadata.joint_attn_mask, attn_metadata.attn_mask], dim=1)
+                        if joint_strategy == "front"
+                        else torch.cat([attn_metadata.attn_mask, attn_metadata.joint_attn_mask], dim=1)
+                    )
+
+            if attn_metadata.attn_mask is not None:
+                # the final attn_mask is ready, the length should be aligedn with query length
+                assert attn_metadata.attn_mask.shape[1] == query.shape[1], (
+                    f"attn_mask length: {attn_metadata.attn_mask.shape[1]} != query length: {query.shape[1]}"
+                )
+                attn_metadata.attn_mask = attn_metadata.attn_mask.bool().contiguous()
         return query, key, value, attn_metadata, ctx
 
     def post_attention(self, attn_output: torch.Tensor, ctx: ParallelAttentionContext | None) -> torch.Tensor:

vllm_omni/diffusion/models/qwen_image/qwen_image_transformer.py

@@ -23,9 +23,11 @@
     AttentionMetadata,
 )
 from vllm_omni.diffusion.attention.layer import Attention
+from vllm_omni.diffusion.attention.selector import get_attn_backend
 from vllm_omni.diffusion.cache.base import CachedTransformer
 from vllm_omni.diffusion.data import OmniDiffusionConfig
 from vllm_omni.diffusion.distributed.parallel_state import (
+    get_ring_parallel_world_size,
     get_sequence_parallel_rank,
     get_sequence_parallel_world_size,
     get_sp_group,
@@ -373,7 +375,14 @@ def forward(
         encoder_hidden_states: torch.Tensor,
         vid_freqs: torch.Tensor,
         txt_freqs: torch.Tensor,
+        hidden_states_mask: torch.Tensor | None = None,
+        encoder_hidden_states_mask: torch.Tensor | None = None,
     ):
+        # 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
+        if encoder_hidden_states_mask is not None and encoder_hidden_states_mask.all():
+            encoder_hidden_states_mask = None
         seq_len_txt = encoder_hidden_states.shape[1]
 
         # Compute QKV for image stream (sample projections)
@@ -416,30 +425,63 @@ def forward(
         joint_value = torch.cat([txt_value, img_value], dim=1)
 
         # Compute joint attention
-
         if (
             self.parallel_config is not None
             and self.parallel_config.sequence_parallel_size > 1
             and not get_forward_context().split_text_embed_in_sp
         ):
             # if using sequence parallel, but not splitting text embed,
             #  we need to pass text embedding to attention layer as joint qkv
+            attn_metadata = AttentionMetadata(
+                joint_query=txt_query,
+                joint_key=txt_key,
+                joint_value=txt_value,
+                joint_strategy="front",
+            )
+            if hidden_states_mask is not None:
+                attn_metadata.attn_mask = hidden_states_mask
+            if encoder_hidden_states_mask is not None:
+                attn_metadata.joint_attn_mask = encoder_hidden_states_mask
+
             joint_hidden_states = self.attn(
                 img_query,
                 img_key,
                 img_value,
-                AttentionMetadata(
-                    joint_query=txt_query,
-                    joint_key=txt_key,
-                    joint_value=txt_value,
-                    joint_strategy="front",
-                ),
+                attn_metadata,
             )
         else:
+            attn_metadata = None
+            if hidden_states_mask is not None or encoder_hidden_states_mask is not None:
+                mask_list = []
+                if encoder_hidden_states_mask is not None:
+                    mask_list.append(encoder_hidden_states_mask)
+                else:
+                    mask_list.append(
+                        torch.ones(
+                            [encoder_hidden_states.shape[0], encoder_hidden_states.shape[1]],
+                            dtype=torch.bool,
+                            device=encoder_hidden_states.device,
+                        )
+                    )
+                if hidden_states_mask is not None:
+                    mask_list.append(hidden_states_mask)
+                else:
+                    mask_list.append(
+                        torch.ones(
+                            [hidden_states.shape[0], hidden_states.shape[1]],
+                            dtype=torch.bool,
+                            device=hidden_states.device,
+                        )
+                    )
+                joint_mask = (
+                    None if len(mask_list) == 0 else torch.cat(mask_list, dim=1) if len(mask_list) > 1 else mask_list[0]
+                )
+                attn_metadata = AttentionMetadata(attn_mask=joint_mask)
             joint_hidden_states = self.attn(
                 joint_query,
                 joint_key,
                 joint_value,
+                attn_metadata,
             )
         joint_hidden_states = joint_hidden_states.flatten(2, 3)
         joint_hidden_states = joint_hidden_states.to(joint_query.dtype)
@@ -547,6 +589,7 @@ def forward(
         image_rotary_emb: tuple[torch.Tensor, torch.Tensor],
         joint_attention_kwargs: dict[str, Any] | None = None,
         modulate_index: list[int] | None = None,
+        hidden_states_mask: torch.Tensor | None = None,
     ) -> tuple[torch.Tensor, torch.Tensor]:
         # Get modulation parameters for both streams
         img_mod_params = self.img_mod(temb)  # [B, 6*dim]
@@ -577,6 +620,8 @@ def forward(
             encoder_hidden_states=txt_modulated,  # Text stream (will be processed as "context")
             vid_freqs=image_rotary_emb[0],
             txt_freqs=image_rotary_emb[1],
+            hidden_states_mask=hidden_states_mask,
+            encoder_hidden_states_mask=encoder_hidden_states_mask,
         )
 
         # QwenAttnProcessor2_0 returns (img_output, txt_output) when encoder_hidden_states is provided
@@ -732,14 +777,48 @@ def forward(
         # else:
         #     lora_scale = 1.0
 
+        original_seq_len = None
+        seq_padding = 0
+        hidden_states_mask = None
+
         if self.parallel_config.sequence_parallel_size > 1:
-            hidden_states = torch.chunk(hidden_states, get_sequence_parallel_world_size(), dim=-2)[
-                get_sequence_parallel_rank()
-            ]
+            batch_size, seq_len, channels = hidden_states.shape
+            sp_size = get_sequence_parallel_world_size()
+
+            if seq_len % sp_size != 0:
+                #  flash_attn, ring_attn, sage_attn do not support attention_mask
+                if get_attn_backend(-1).get_name() != "SDPA" and get_attn_backend(-1).get_name() != "ASCEND":
+                    raise ValueError(
+                        f"When generating image shape that the sequence length is NOT divisible by sp_size={sp_size},"
+                        f"cannot use {get_attn_backend(-1).get_name()} which does not support attention_mask."
+                        f"Please switch to SDPA or Ascend attention backend."
+                    )
+                # ring attention does not support attention_mask
+                if get_ring_parallel_world_size() > 1:
+                    raise ValueError(
+                        f"When generating image shape that the sequence length is NOT divisible by sp_size={sp_size},"
+                        f"cannot use ring attention which does not support attention_mask."
+                        f"Please switch to Ulysses SP only."
+                    )
+
+                seq_padding = sp_size - (seq_len % sp_size)
+                original_seq_len = seq_len
+
+                hidden_states_mask = torch.ones(
+                    batch_size, seq_len + seq_padding, dtype=torch.bool, device=hidden_states.device
+                )
+                hidden_states_mask[:, seq_len:] = False
+                padding_tensor = torch.zeros(
+                    batch_size, seq_padding, channels, dtype=hidden_states.dtype, device=hidden_states.device
+                )
+                hidden_states = torch.cat([hidden_states, padding_tensor], dim=1)
+
+            hidden_states = torch.chunk(hidden_states, sp_size, dim=-2)[get_sequence_parallel_rank()]
             # NOTE:
             # QwenImage uses *dual-stream* (text + image) and runs a *joint attention*.
             # text embeddings to be replicated across SP ranks for correctness.
             get_forward_context().split_text_embed_in_sp = False
+
         hidden_states = self.img_in(hidden_states)
 
         # Ensure timestep tensor is on the same device and dtype as hidden_states
@@ -769,13 +848,17 @@ def forward(
 
         image_rotary_emb = self.pos_embed(img_shapes, txt_seq_lens, device=hidden_states.device)
 
-        def get_rotary_emb_chunk(freqs):
+        def get_rotary_emb_chunk(freqs, padding=0):
+            # Pad rotary embeddings if needed
+            if padding > 0:
+                padding_tensor = torch.zeros(padding, freqs.shape[-1], dtype=freqs.dtype, device=freqs.device)
+                freqs = torch.cat([freqs, padding_tensor], dim=0)
             freqs = torch.chunk(freqs, get_sequence_parallel_world_size(), dim=0)[get_sequence_parallel_rank()]
             return freqs
 
         if self.parallel_config.sequence_parallel_size > 1:
             img_freqs, txt_freqs = image_rotary_emb
-            img_freqs = get_rotary_emb_chunk(img_freqs)
+            img_freqs = get_rotary_emb_chunk(img_freqs, seq_padding)
             if get_forward_context().split_text_embed_in_sp:
                 txt_freqs = get_rotary_emb_chunk(txt_freqs)
             image_rotary_emb = (img_freqs, txt_freqs)
@@ -789,6 +872,7 @@ def get_rotary_emb_chunk(freqs):
                 image_rotary_emb=image_rotary_emb,
                 joint_attention_kwargs=attention_kwargs,
                 modulate_index=modulate_index,
+                hidden_states_mask=hidden_states_mask,
             )
 
         if self.zero_cond_t:
@@ -799,6 +883,11 @@ def get_rotary_emb_chunk(freqs):
 
         if self.parallel_config.sequence_parallel_size > 1:
             output = get_sp_group().all_gather(output, dim=-2)
+
+            # Remove padding if it was added
+            if original_seq_len is not None:
+                output = output[:, :original_seq_len, :]
+
         return Transformer2DModelOutput(sample=output)
 
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: