diff --git a/tests/model_executor/test_phi4mm_attention.py b/tests/model_executor/test_phi4mm_attention.py new file mode 100644 index 0000000000..51710f84bc --- /dev/null +++ b/tests/model_executor/test_phi4mm_attention.py @@ -0,0 +1,28 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import torch + +from vllm.model_executor.models.phi4mm_utils import MultiHeadedAttention + + +def test_eager_attention_does_not_retain_probability_matrix(): + module = MultiHeadedAttention( + n_head=2, + n_feat=16, + dropout_rate=0.0, + use_pt_scaled_dot_product_attention=False, + ) + module.eval() + batch, time, feat = 2, 8, 16 + x = torch.randn(batch, time, feat) + with torch.no_grad(): + out = module(x, x, x, None, None, None) + + assert out.shape == (batch, time, feat) + assert getattr(module, "attn", None) is None + + +def test_multiheaded_attention_defaults_to_scaled_dot_product(): + module = MultiHeadedAttention(n_head=2, n_feat=16, dropout_rate=0.0) + assert module.use_pt_scaled_dot_product_attention is True diff --git a/vllm/model_executor/models/phi4mm_audio.py b/vllm/model_executor/models/phi4mm_audio.py index 2a16b4ca75..9daedcf38e 100644 --- a/vllm/model_executor/models/phi4mm_audio.py +++ b/vllm/model_executor/models/phi4mm_audio.py @@ -129,7 +129,7 @@ class ConformerEncoderLayer(nn.Module): default False. use_pt_scaled_dot_product_attention: bool, optional if set to True, use pytorch's scaled dot product attention - implementation in training. + implementation in training. Default: True. attn_group_sizes: int, optional the number of groups to use for attention, default 1 (Multi-Head Attention), @@ -163,7 +163,7 @@ class ConformerEncoderLayer(nn.Module): attention_glu_type: str = "swish", activation_checkpointing: str = "", export: bool = False, - use_pt_scaled_dot_product_attention: bool = False, + use_pt_scaled_dot_product_attention: bool = True, attn_group_sizes: int = 1, ) -> None: super().__init__() @@ -808,7 +808,7 @@ class ConformerEncoder(TransformerEncoderBase): default 4 use_pt_scaled_dot_product_attention: whether to use pytorch scaled dot product attention in training. - Default: False + Default: True nemo_conv_settings: dict, optional A dictionary of settings for NeMo Subsampling. default: None @@ -873,7 +873,7 @@ class ConformerEncoder(TransformerEncoderBase): activation_checkpointing: str = "", relative_attention_bias_args: dict[str, Any] | None = None, time_reduction: int = 4, - use_pt_scaled_dot_product_attention: bool = False, + use_pt_scaled_dot_product_attention: bool = True, nemo_conv_settings: dict[str, Any] | None = None, conv2d_extra_padding: Literal["feat", "feat_time", "none", True] = "none", replication_pad_for_subsample_embedding: bool = False, diff --git a/vllm/model_executor/models/phi4mm_utils.py b/vllm/model_executor/models/phi4mm_utils.py index af54defd3f..db1b83f343 100644 --- a/vllm/model_executor/models/phi4mm_utils.py +++ b/vllm/model_executor/models/phi4mm_utils.py @@ -1642,7 +1642,7 @@ class MultiHeadedAttention(nn.Module): NOTE: this will NOT be used in ONNX decoding due to a lack of support. In that case, we use the original attention implementation, which shows no regression. - default: False. + default: True. n_value: int, optional if set to values other than -1, use a different dimension for value. With the default value (i.e. -1), it is backward compatible. @@ -1666,7 +1666,7 @@ class MultiHeadedAttention(nn.Module): attention_inner_dim: int = -1, glu_type: str = "swish", bias_in_glu: bool = True, - use_pt_scaled_dot_product_attention: bool = False, + use_pt_scaled_dot_product_attention: bool = True, n_value: int = -1, group_size: int = 1, ) -> None: @@ -1690,7 +1690,6 @@ class MultiHeadedAttention(nn.Module): self.linear_v = nn.Linear(n_value, attention_inner_dim // group_size) self.linear_out = nn.Linear(attention_inner_dim // group_size, n_value) - self.attn = torch.jit.Attribute(None, Tensor | None) self.dropout = nn.Dropout(p=dropout_rate) self.dropout_rate = dropout_rate self.use_pt_scaled_dot_product_attention = use_pt_scaled_dot_product_attention @@ -1799,8 +1798,6 @@ class MultiHeadedAttention(nn.Module): attn = masked_softmax(scores, mask) # (batch, head, time1, time2) - self.attn = attn - p_attn = self.dropout(attn) x = torch.matmul(p_attn.to(v.dtype), v) # (batch, head, time1, d_k) if pos_v is not None: