revert sdpa_mask_without_vmap changes

Author: echarlaix

optimum/exporters/openvino/patching_utils.py

@@ -186,21 +186,23 @@ def _orig_sdpa_mask_without_vmap(
 # Compatibility wrapper for sdpa_mask_without_vmap from optimum.
 # The installed optimum version expects (batch_size, cache_position: Tensor, kv_length, ...),
 # but transformers >= 5.5 passes (batch_size, q_length: int, kv_length: int, q_offset: int, ...).
-def sdpa_mask_without_vmap(**kwargs):
-    kwargs.pop("use_vmap", None)
-    if is_transformers_version("<", "5"):
-        return _orig_sdpa_mask_without_vmap(**kwargs)
-    elif (
-        is_transformers_version(">=", "5.4")
-        and is_transformers_version("<", "5.9")
-        and isinstance(kwargs.get("q_length", None), torch.Tensor)
-    ):
-        q_length = kwargs.pop("q_length", None)
-        q_offset = kwargs.pop("q_offset", 0)
-        cache_position = torch.arange(q_offset, q_offset + q_length, device=q_length.device)
-        return sdpa_mask(q_length=cache_position, use_vmap=False, **kwargs)
+def sdpa_mask_without_vmap(batch_size, q_length=None, kv_length=None, q_offset=0, kv_offset=0, **kwargs):
+    import inspect
+
+    sig = inspect.signature(_orig_sdpa_mask_without_vmap)
+    if is_transformers_version(">=", "5.5") and "cache_position" in sig.parameters and q_length is not None:
+        # Old optimum signature: (batch_size, cache_position, kv_length, kv_offset, ...)
+        cache_position = torch.arange(q_length, dtype=torch.long) + q_offset
+        kwargs.pop("q_offset", None)
+        kwargs.pop("allow_is_bidirectional_skip", None)
+        kwargs.pop("allow_torch_fix", None)
+        kwargs.pop("use_vmap", None)
+        kwargs.pop("device", None)
+        return _orig_sdpa_mask_without_vmap(batch_size, cache_position, kv_length, kv_offset=kv_offset, **kwargs)
     else:
-        return sdpa_mask(use_vmap=False, **kwargs)
+        return _orig_sdpa_mask_without_vmap(
+            batch_size, q_length=q_length, kv_length=kv_length, q_offset=q_offset, kv_offset=kv_offset, **kwargs
+        )
 
 
 # Adapted from https://github.com/huggingface/transformers/blob/v4.53.0/src/transformers/masking_utils.py#L433