Add TRANSFORMERS_SDPA_FORCE_GQA envvar in sdpa_attention_forward (#49459)

Add `TRANSFORMERS_SDPA_FORCE_GQA` env var

Author: jakobknauer

docs/source/en/reference/environment_variables.md

@@ -53,6 +53,7 @@ Optimized kernels are downloaded from the Hub when [kernels](https://github.com/
 | `USE_HUB_KERNELS` | boolean, enabled by default | Set to a false value to fall back to the reference implementations. |
 | `TRANSFORMERS_DISABLE_DEEPGEMM_LINEAR` | `1` | Forces the Triton fallback instead of DeepGEMM for FP8 linear layers in finegrained FP8 quantization. |
 | `TRANSFORMERS_DISABLE_TORCH_CHECK` | `1` | Skips the minimum PyTorch version check at import. |
+| `TRANSFORMERS_SDPA_FORCE_GQA` | `1` | Forces PyTorch SDPA to use `enable_gqa=True` for grouped-query attention. Otherwise, Transformers selects between setting `enable_gqa=True` and repeating key/value heads. |
 
 > [!WARNING]
 > `TRANSFORMERS_DISABLE_TORCH_CHECK` removes a guard rather than changing behavior. On an unsupported PyTorch version you get import errors or silently wrong results instead of a clear message.

src/transformers/integrations/sdpa_attention.py

@@ -1,3 +1,5 @@
+import os
+
 import torch
 
 from ..generation.continuous_batching.cache import PagedAttentionCache
@@ -28,6 +30,8 @@ def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
 
 
 def use_gqa_in_sdpa(attention_mask: torch.Tensor | None, key: torch.Tensor, value: torch.Tensor) -> bool:
+    if os.environ.get("TRANSFORMERS_SDPA_FORCE_GQA", "0") == "1":
+        return True
     # XPU and MPS support masked GQA under more recent torch versions
     if _is_torch_xpu_available:
         return _is_torch_greater_or_equal_than_2_8

tests/utils/test_modeling_utils.py

@@ -3290,6 +3290,42 @@ def test_position_bias_in_sdpa_uses_a_fused_kernel(self, use_mask, is_causal):
             )
         self.assertEqual(attn_output.shape, (1, 16, 8, 64))
 
+    @parameterized.expand([("0", False), ("1", True)])
+    def test_transformers_sdpa_force_gqa_env_var_overrides_use_gqa_in_sdpa_attention_forward(
+        self, env_var_value, expected_enable_gqa
+    ):
+        from transformers.integrations.sdpa_attention import sdpa_attention_forward
+
+        module = torch.nn.Module()
+        module.num_key_value_groups = 2
+        query = torch.randn(1, 4, 16, 1, dtype=torch.float16)
+        key = torch.randn(1, 2, 16, 1, dtype=torch.float16)
+        value = torch.randn(1, 2, 16, 1, dtype=torch.float16)
+
+        # use an attention mask and set _is_torch_{xpu,mps}_available to False, so that use_gqa_in_sdpa
+        # falls back to returning False when the env var is not set to 1
+        attention_mask = torch.zeros(1, 1, 16, 16, dtype=torch.float16)
+
+        with (
+            patch("transformers.integrations.sdpa_attention._is_torch_xpu_available", False),
+            patch("transformers.integrations.sdpa_attention._is_torch_mps_available", False),
+            patch(
+                "torch.nn.functional.scaled_dot_product_attention",
+                wraps=torch.nn.functional.scaled_dot_product_attention,
+            ) as mock_sdpa,
+            patch.dict(os.environ, {"TRANSFORMERS_SDPA_FORCE_GQA": env_var_value}),
+        ):
+            sdpa_attention_forward(module, query, key, value, attention_mask=attention_mask)
+
+        mock_sdpa.assert_called_once()
+        self.assertEqual(mock_sdpa.call_args.kwargs.get("enable_gqa", False), expected_enable_gqa)
+        self.assertEqual(
+            mock_sdpa.call_args.args[1].shape[1], 2 if expected_enable_gqa else 4
+        )  # check if key was repeated
+        self.assertEqual(
+            mock_sdpa.call_args.args[2].shape[1], 2 if expected_enable_gqa else 4
+        )  # check if value was repeated
+
     def test_flash_attn_available_no_keyerror_when_missing_from_distribution_map(self):
         # Regression test for https://github.com/huggingface/transformers/issues/45520.
         # When flash_attn is importable but not present in PACKAGE_DISTRIBUTION_MAPPING