[Perf]Boogu-Image Skip dense attention masks (#6871)

Signed-off-by: zhuhh97 <zhuhonghan@huawei.com>

tests/diffusion/models/boogu_image/test_boogu_image_attention_mask.py

@@ -0,0 +1,30 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+import pytest
+import torch
+
+from vllm_omni.diffusion.models.utils import make_attention_mask
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
+
+
+@pytest.mark.parametrize("seq_lengths", [[16], [16, 16]])
+def test_attention_mask_is_skipped_for_dense_batch(seq_lengths):
+    hidden_states = torch.empty(len(seq_lengths), 16, 64)
+
+    assert make_attention_mask(hidden_states, seq_lengths) is None
+
+
+def test_attention_mask_preserves_variable_length_padding():
+    hidden_states = torch.empty(2, 16, 64)
+
+    attention_mask = make_attention_mask(hidden_states, [16, 12])
+
+    expected = torch.tensor(
+        [
+            [True] * 16,
+            [True] * 12 + [False] * 4,
+        ]
+    )
+    torch.testing.assert_close(attention_mask, expected)

vllm_omni/diffusion/models/boogu_image/boogu_image_transformer.py

@@ -35,6 +35,7 @@ from vllm.model_executor.model_loader.weight_utils import default_weight_loader
 from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
 from vllm_omni.diffusion.attention.layer import Attention
 from vllm_omni.diffusion.data import OmniDiffusionConfig
+from vllm_omni.diffusion.models.utils import make_attention_mask
 from vllm_omni.platforms import current_omni_platform
 
 if TYPE_CHECKING:
@@ -1554,17 +1555,11 @@ class BooguImageTransformer2DModel(nn.Module):
         img_hidden_states = combined_img_hidden_states
 
         # Joint mask for [instruct + image].
-        max_seq_len = max(seq_lengths)
-        joint_attention_mask = hidden_states.new_zeros(batch_size, max_seq_len, dtype=torch.bool)
-        for i, seq_len in enumerate(seq_lengths):
-            joint_attention_mask[i, :seq_len] = True
+        joint_attention_mask = make_attention_mask(hidden_states, seq_lengths)
 
         # Dual-stream (double-stream) stage.
         if self.num_double_stream_layers > 0:
-            max_img_len = max(combined_img_seq_lengths)
-            img_attention_mask = hidden_states.new_zeros(batch_size, max_img_len, dtype=torch.bool)
-            for i, img_seq_len in enumerate(combined_img_seq_lengths):
-                img_attention_mask[i, :img_seq_len] = True
+            img_attention_mask = make_attention_mask(hidden_states, combined_img_seq_lengths)
 
             for layer in self.double_stream_layers:
                 img_hidden_states, instruct_hidden_states = layer(

vllm_omni/diffusion/models/utils.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 from __future__ import annotations
 
@@ -31,6 +31,34 @@ if TYPE_CHECKING:
 Style = Literal["colwise", "colwise_rep", "rowwise", "rowwise_rep", "replicate"]
 
 
+def make_attention_mask(
+    hidden_states: torch.Tensor,
+    seq_lengths: list[int],
+) -> torch.Tensor | None:
+    """Build a padding mask only when the batch contains padding.
+
+    ``seq_lengths`` must be non-empty and contain valid sequence lengths.
+    Returns a boolean mask of shape ``(len(seq_lengths), max(seq_lengths))``
+    on the device of ``hidden_states``, with True marking valid tokens,
+    or None when all sequence lengths are equal.
+
+    Passing an all-true CUDA mask makes attention backends inspect a device
+    scalar to select the dense path, which introduces a host-device sync.
+    """
+    max_seq_len = max(seq_lengths)
+    if all(seq_len == max_seq_len for seq_len in seq_lengths):
+        return None
+
+    attention_mask = hidden_states.new_zeros(
+        len(seq_lengths),
+        max_seq_len,
+        dtype=torch.bool,
+    )
+    for i, seq_len in enumerate(seq_lengths):
+        attention_mask[i, :seq_len] = True
+    return attention_mask
+
+
 def replace_linear_class(
     linear: nn.Linear,
     style: Style = "replicate",