From c086554d1fd02749bdc2bcad37636e90ade01fec Mon Sep 17 00:00:00 2001 From: Honghan Zhu Date: Tue, 8 Sep 2026 15:53:20 +0800 Subject: [PATCH] [Perf]Boogu-Image Skip dense attention masks (#6871) Signed-off-by: zhuhh97 --- .../test_boogu_image_attention_mask.py | 30 +++++++++++++++++++ .../boogu_image/boogu_image_transformer.py | 11 ++----- vllm_omni/diffusion/models/utils.py | 30 ++++++++++++++++++- 3 files changed, 62 insertions(+), 9 deletions(-) create mode 100644 tests/diffusion/models/boogu_image/test_boogu_image_attention_mask.py diff --git a/tests/diffusion/models/boogu_image/test_boogu_image_attention_mask.py b/tests/diffusion/models/boogu_image/test_boogu_image_attention_mask.py new file mode 100644 index 000000000..dbaf3b71a --- /dev/null +++ b/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) diff --git a/vllm_omni/diffusion/models/boogu_image/boogu_image_transformer.py b/vllm_omni/diffusion/models/boogu_image/boogu_image_transformer.py index 2025859ed..66031ca3a 100644 --- a/vllm_omni/diffusion/models/boogu_image/boogu_image_transformer.py +++ b/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( diff --git a/vllm_omni/diffusion/models/utils.py b/vllm_omni/diffusion/models/utils.py index 122646219..2d783d907 100644 --- a/vllm_omni/diffusion/models/utils.py +++ b/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",