[Security] Do not retain Phi-4 MM audio attention matrices (#60614)

Signed-off-by: Juan Pérez de Algaba <jperezde@redhat.com>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>

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

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,

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: