diff --git a/tests/model_executor/models/minicpmo_4_5/test_vision_flash_attention.py b/tests/model_executor/models/minicpmo_4_5/test_vision_flash_attention.py new file mode 100644 index 000000000..3139dbeac --- /dev/null +++ b/tests/model_executor/models/minicpmo_4_5/test_vision_flash_attention.py @@ -0,0 +1,353 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Regression coverage for MiniCPM-o 4.5 SigLIP FlashAttention.""" + +from __future__ import annotations + +from collections.abc import Callable + +import pytest +import torch +from pytest_mock import MockerFixture +from transformers.utils import is_flash_attn_2_available + +from tests.helpers.mark import hardware_test +from vllm_omni.model_executor.models.minicpmo_4_5 import minicpmo_4_5_omni_llm as minicpmo_model +from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni_llm import ( + SiglipAttention, + SiglipEncoder, + SiglipFlashAttention2, + SiglipVisionConfig, +) + +_VISION_NUM_HEADS = 16 +_VISION_HEAD_DIM = 72 +_VISION_HIDDEN_SIZE = _VISION_NUM_HEADS * _VISION_HEAD_DIM +_DTYPE = torch.bfloat16 + + +def _attention_config(*, hidden_size: int, num_heads: int) -> SiglipVisionConfig: + return SiglipVisionConfig( + hidden_size=hidden_size, + num_attention_heads=num_heads, + attention_dropout=0.0, + ) + + +def _encoder_config(*, num_layers: int = 3) -> SiglipVisionConfig: + config = SiglipVisionConfig( + hidden_size=32, + intermediate_size=64, + num_hidden_layers=num_layers, + num_attention_heads=4, + attention_dropout=0.0, + ) + config._attn_implementation = "flash_attention_2" + return config + + +def _capture_output( + outputs: dict[str, torch.Tensor], name: str +) -> Callable[[torch.nn.Module, tuple[torch.Tensor, ...], torch.Tensor], None]: + def hook( + _module: torch.nn.Module, + _inputs: tuple[torch.Tensor, ...], + output: torch.Tensor, + ) -> None: + outputs[name] = output + + return hook + + +@pytest.mark.core_model +@pytest.mark.cpu +def test_qkv_bshd_layout_reuses_projection_storage(mocker: MockerFixture) -> None: + batch_size = 2 + seq_len = 7 + hidden_size = 32 + num_heads = 4 + head_dim = hidden_size // num_heads + attention = SiglipFlashAttention2(_attention_config(hidden_size=hidden_size, num_heads=num_heads)).eval() + + projected: dict[str, torch.Tensor] = {} + attention.q_proj.register_forward_hook(_capture_output(projected, "query")) + attention.k_proj.register_forward_hook(_capture_output(projected, "key")) + attention.v_proj.register_forward_hook(_capture_output(projected, "value")) + + flash_inputs: dict[str, torch.Tensor] = {} + + def fake_flash_attention( + query_states: torch.Tensor, + key_states: torch.Tensor, + value_states: torch.Tensor, + _attention_mask: torch.Tensor | None, + _query_length: int, + *, + dropout: float = 0.0, + softmax_scale: float | None = None, + unpadding_metadata: minicpmo_model.FlashAttentionUnpaddingMetadata | None = None, + ) -> torch.Tensor: + del dropout, softmax_scale, unpadding_metadata + assert _attention_mask is None + assert _query_length == seq_len + flash_inputs.update( + query=query_states, + key=key_states, + value=value_states, + ) + return query_states + + flash_forward = mocker.patch.object( + attention, + "_flash_attention_forward", + side_effect=fake_flash_attention, + ) + + hidden_states = torch.randn(batch_size, seq_len, hidden_size) + with torch.inference_mode(): + output, _ = attention(hidden_states) + + assert output.shape == hidden_states.shape + flash_forward.assert_called_once() + + expected_shape = (batch_size, seq_len, num_heads, head_dim) + expected_stride = ( + seq_len * num_heads * head_dim, + num_heads * head_dim, + head_dim, + 1, + ) + for name in ("query", "key", "value"): + projected_tensor = projected[name] + flash_tensor = flash_inputs[name] + assert flash_tensor.shape == expected_shape + assert flash_tensor.stride() == expected_stride + assert flash_tensor.is_contiguous() + assert flash_tensor.storage_offset() == projected_tensor.storage_offset() + assert flash_tensor.untyped_storage().data_ptr() == projected_tensor.untyped_storage().data_ptr() + + +@pytest.mark.core_model +@pytest.mark.cpu +@pytest.mark.parametrize( + ("attention_mask", "expected_calls"), + [ + pytest.param(torch.tensor([[True, True, False], [True, False, False]]), 1, id="padded"), + pytest.param(None, 0, id="dense"), + ], +) +def test_encoder_computes_unpadding_metadata_once_per_forward( + mocker: MockerFixture, + attention_mask: torch.Tensor | None, + expected_calls: int, +) -> None: + encoder = SiglipEncoder(_encoder_config()).eval() + metadata_spy = mocker.spy(minicpmo_model, "_get_unpad_data") + received_metadata: list[minicpmo_model.FlashAttentionUnpaddingMetadata | None] = [] + + def fake_flash_attention( + query_states: torch.Tensor, + key_states: torch.Tensor, + value_states: torch.Tensor, + _attention_mask: torch.Tensor | None, + _query_length: int, + *, + dropout: float = 0.0, + softmax_scale: float | None = None, + unpadding_metadata: minicpmo_model.FlashAttentionUnpaddingMetadata | None = None, + ) -> torch.Tensor: + del key_states, value_states, _attention_mask, _query_length, dropout, softmax_scale + received_metadata.append(unpadding_metadata) + return query_states + + for layer in encoder.layers: + mocker.patch.object(layer.self_attn, "_flash_attention_forward", side_effect=fake_flash_attention) + + hidden_states = torch.randn(2, 3, 32) + with torch.inference_mode(): + encoder( + inputs_embeds=hidden_states, + attention_mask=attention_mask, + output_attentions=False, + output_hidden_states=False, + return_dict=False, + ) + + assert metadata_spy.call_count == expected_calls + assert len(received_metadata) == len(encoder.layers) + if attention_mask is None: + assert all(metadata is None for metadata in received_metadata) + else: + assert received_metadata[0] is not None + assert all(metadata is received_metadata[0] for metadata in received_metadata) + + +@pytest.mark.core_model +@pytest.mark.cpu +def test_encoder_recomputes_unpadding_metadata_for_new_mask(mocker: MockerFixture) -> None: + encoder = SiglipEncoder(_encoder_config(num_layers=1)).eval() + metadata_spy = mocker.spy(minicpmo_model, "_get_unpad_data") + received_metadata: list[minicpmo_model.FlashAttentionUnpaddingMetadata | None] = [] + + def fake_flash_attention( + query_states: torch.Tensor, + key_states: torch.Tensor, + value_states: torch.Tensor, + _attention_mask: torch.Tensor | None, + _query_length: int, + *, + dropout: float = 0.0, + softmax_scale: float | None = None, + unpadding_metadata: minicpmo_model.FlashAttentionUnpaddingMetadata | None = None, + ) -> torch.Tensor: + del key_states, value_states, _attention_mask, _query_length, dropout, softmax_scale + received_metadata.append(unpadding_metadata) + return query_states + + mocker.patch.object( + encoder.layers[0].self_attn, + "_flash_attention_forward", + side_effect=fake_flash_attention, + ) + hidden_states = torch.randn(2, 4, 32) + masks = ( + torch.tensor([[True, True, True, False], [True, True, False, False]]), + torch.tensor([[True, True, False, False], [True, False, False, False]]), + ) + + with torch.inference_mode(): + for attention_mask in masks: + encoder( + inputs_embeds=hidden_states, + attention_mask=attention_mask, + output_attentions=False, + output_hidden_states=False, + return_dict=False, + ) + + assert metadata_spy.call_count == len(masks) + first_metadata, second_metadata = received_metadata + assert first_metadata is not None + assert second_metadata is not None + assert first_metadata is not second_metadata + assert not torch.equal(first_metadata[0], second_metadata[0]) + + +def _require_flash_attention() -> None: + if not torch.cuda.is_available(): + pytest.skip("CUDA is required for FlashAttention execution") + if not is_flash_attn_2_available(): + pytest.skip("Transformers reports that FlashAttention 2 is unavailable") + + +def _make_cuda_attention_pair() -> tuple[SiglipAttention, SiglipFlashAttention2]: + config = _attention_config( + hidden_size=_VISION_HIDDEN_SIZE, + num_heads=_VISION_NUM_HEADS, + ) + torch.manual_seed(0) + eager_attention = SiglipAttention(config).eval().to(device="cuda", dtype=_DTYPE) + flash_attention = SiglipFlashAttention2(config).eval().to(device="cuda", dtype=_DTYPE) + flash_attention.load_state_dict(eager_attention.state_dict()) + return eager_attention, flash_attention + + +@pytest.mark.core_model +@hardware_test(res={"cuda": "L4"}) +def test_encoder_metadata_reuse_matches_per_layer_recomputation() -> None: + _require_flash_attention() + encoder = SiglipEncoder(_encoder_config()).eval().to(device="cuda", dtype=_DTYPE) + torch.manual_seed(3) + hidden_states = torch.randn(2, 64, 32, device="cuda", dtype=_DTYPE) + attention_mask = torch.arange(64, device="cuda").unsqueeze(0) < torch.tensor( + (64, 43), + device="cuda", + ).unsqueeze(1) + + with torch.inference_mode(): + expected = hidden_states + for encoder_layer in encoder.layers: + expected = encoder_layer( + expected, + attention_mask, + output_attentions=False, + unpadding_metadata=None, + )[0] + actual = encoder( + inputs_embeds=hidden_states, + attention_mask=attention_mask, + output_attentions=False, + output_hidden_states=False, + return_dict=False, + )[0] + + torch.testing.assert_close(actual, expected, atol=0.0, rtol=0.0) + + +@pytest.mark.core_model +@hardware_test(res={"cuda": "L4"}) +def test_flash_attention_unpadded_matches_eager_reference() -> None: + _require_flash_attention() + eager_attention, flash_attention = _make_cuda_attention_pair() + torch.manual_seed(1) + hidden_states = torch.randn( + 2, + 64, + _VISION_HIDDEN_SIZE, + device="cuda", + dtype=_DTYPE, + ) + + with torch.inference_mode(): + expected, _ = eager_attention(hidden_states) + actual, _ = flash_attention(hidden_states) + + torch.testing.assert_close(actual, expected, atol=2e-2, rtol=2e-2) + + +@pytest.mark.core_model +@hardware_test(res={"cuda": "L4"}) +@pytest.mark.parametrize( + "valid_lengths", + [ + pytest.param((64, 43), id="moderate_padding"), + pytest.param((64, 17), id="high_padding"), + ], +) +def test_flash_attention_padded_matches_per_sample_unpadded( + valid_lengths: tuple[int, int], +) -> None: + _require_flash_attention() + _, flash_attention = _make_cuda_attention_pair() + batch_size = len(valid_lengths) + seq_len = max(valid_lengths) + torch.manual_seed(2) + hidden_states = torch.randn( + batch_size, + seq_len, + _VISION_HIDDEN_SIZE, + device="cuda", + dtype=_DTYPE, + ) + attention_mask = torch.arange(seq_len, device="cuda").unsqueeze(0) < torch.tensor( + valid_lengths, + device="cuda", + ).unsqueeze(1) + + with torch.inference_mode(): + padded_output, _ = flash_attention( + hidden_states, + attention_mask=attention_mask, + ) + unpadded_outputs = [ + flash_attention(hidden_states[index : index + 1, :valid_length])[0] + for index, valid_length in enumerate(valid_lengths) + ] + + for index, valid_length in enumerate(valid_lengths): + torch.testing.assert_close( + padded_output[index : index + 1, :valid_length], + unpadded_outputs[index], + atol=2e-2, + rtol=2e-2, + ) diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_llm.py b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_llm.py index ce00f21ed..a182848cf 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_llm.py +++ b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_llm.py @@ -770,8 +770,11 @@ AutoImageProcessor.register(MiniCPMOConfig, MiniCPMVImageProcessor, exist_ok=Tru # ============== SigLIP Vision Transformer Classes ============== +FlashAttentionUnpaddingMetadata: TypeAlias = tuple[torch.Tensor, torch.Tensor, int] + + # Copied from transformers.models.llama.modeling_llama._get_unpad_data -def _get_unpad_data(attention_mask): +def _get_unpad_data(attention_mask: torch.Tensor) -> FlashAttentionUnpaddingMetadata: seqlens_in_batch = attention_mask.sum(dim=-1, dtype=torch.int32) indices = torch.nonzero(attention_mask.flatten(), as_tuple=False).flatten() max_seqlen_in_batch = seqlens_in_batch.max().item() @@ -1116,6 +1119,7 @@ class SiglipFlashAttention2(SiglipAttention): past_key_value: tuple[torch.Tensor] | None = None, output_attentions: bool = False, use_cache: bool = False, + unpadding_metadata: FlashAttentionUnpaddingMetadata | None = None, **kwargs, ) -> tuple[torch.Tensor, torch.Tensor | None, tuple[torch.Tensor] | None]: output_attentions = False @@ -1178,7 +1182,13 @@ class SiglipFlashAttention2(SiglipAttention): value_states = value_states.to(target_dtype) attn_output = self._flash_attention_forward( - query_states, key_states, value_states, attention_mask, q_len, dropout=dropout_rate + query_states, + key_states, + value_states, + attention_mask, + q_len, + dropout=dropout_rate, + unpadding_metadata=unpadding_metadata, ) attn_output = attn_output.reshape(bsz, q_len, self.embed_dim).contiguous() @@ -1190,7 +1200,15 @@ class SiglipFlashAttention2(SiglipAttention): return attn_output, attn_weights def _flash_attention_forward( - self, query_states, key_states, value_states, attention_mask, query_length, dropout=0.0, softmax_scale=None + self, + query_states, + key_states, + value_states, + attention_mask, + query_length, + dropout=0.0, + softmax_scale=None, + unpadding_metadata: FlashAttentionUnpaddingMetadata | None = None, ): """ Calls the forward method of Flash Attention - if the input hidden states contain at least one padding token @@ -1218,7 +1236,12 @@ class SiglipFlashAttention2(SiglipAttention): if attention_mask is not None: batch_size = query_states.shape[0] query_states, key_states, value_states, indices_q, cu_seq_lens, max_seq_lens = self._upad_input( - query_states, key_states, value_states, attention_mask, query_length + query_states, + key_states, + value_states, + attention_mask, + query_length, + unpadding_metadata=unpadding_metadata, ) cu_seqlens_q, cu_seqlens_k = cu_seq_lens @@ -1245,8 +1268,18 @@ class SiglipFlashAttention2(SiglipAttention): return attn_output - def _upad_input(self, query_layer, key_layer, value_layer, attention_mask, query_length): - indices_k, cu_seqlens_k, max_seqlen_in_batch_k = _get_unpad_data(attention_mask) + def _upad_input( + self, + query_layer, + key_layer, + value_layer, + attention_mask, + query_length, + unpadding_metadata: FlashAttentionUnpaddingMetadata | None = None, + ): + if unpadding_metadata is None: + unpadding_metadata = _get_unpad_data(attention_mask) + indices_k, cu_seqlens_k, max_seqlen_in_batch_k = unpadding_metadata batch_size, kv_seq_len, num_key_value_heads, head_dim = key_layer.shape key_layer = index_first_axis( @@ -1316,6 +1349,7 @@ class SiglipEncoderLayer(nn.Module): hidden_states: torch.Tensor, attention_mask: torch.Tensor, output_attentions: bool | None = False, + unpadding_metadata: FlashAttentionUnpaddingMetadata | None = None, ) -> tuple[torch.FloatTensor]: """ Args: @@ -1330,11 +1364,19 @@ class SiglipEncoderLayer(nn.Module): residual = hidden_states hidden_states = self.layer_norm1(hidden_states) - hidden_states, attn_weights = self.self_attn( - hidden_states=hidden_states, - attention_mask=attention_mask, - output_attentions=output_attentions, - ) + if self._use_flash_attention_2: + hidden_states, attn_weights = self.self_attn( + hidden_states=hidden_states, + attention_mask=attention_mask, + output_attentions=output_attentions, + unpadding_metadata=unpadding_metadata, + ) + else: + hidden_states, attn_weights = self.self_attn( + hidden_states=hidden_states, + attention_mask=attention_mask, + output_attentions=output_attentions, + ) hidden_states = residual + hidden_states residual = hidden_states @@ -1435,6 +1477,7 @@ class SiglipEncoder(nn.Module): self.config = config self.layers = nn.ModuleList([SiglipEncoderLayer(config) for _ in range(config.num_hidden_layers)]) self.gradient_checkpointing = False + self._use_flash_attention_2 = config._attn_implementation == "flash_attention_2" # Ignore copy def forward( @@ -1475,6 +1518,9 @@ class SiglipEncoder(nn.Module): all_attentions = () if output_attentions else None hidden_states = inputs_embeds + unpadding_metadata = ( + _get_unpad_data(attention_mask) if self._use_flash_attention_2 and attention_mask is not None else None + ) for encoder_layer in self.layers: if output_hidden_states: encoder_states = encoder_states + (hidden_states,) @@ -1484,12 +1530,14 @@ class SiglipEncoder(nn.Module): hidden_states, attention_mask, output_attentions, + unpadding_metadata, ) else: layer_outputs = encoder_layer( hidden_states, attention_mask, output_attentions=output_attentions, + unpadding_metadata=unpadding_metadata, ) hidden_states = layer_outputs[0]