Check correctness for score_mod implementations

[oraluben]

No description provided.

[oraluben]

Hi @drisspg , I'd like to propose a change in the main torch repo in the torch.nn.attention.flex_attention.create_mask[1] function, let it to returns the real bias, instead of the mask.

This allows to finish this PR with minimal efforts and would help other people to validate their flex attn implementations IMO.
If the change has been adopted, the current behaviour can be simply restored with torch.where(torch.isneginf(out), False, True), while it's not easy to get the bias from the mask, if the user do not want to just ctrl-c-v create_mask.

The change might look like this with backward compatibility:

diff --git a/torch/nn/attention/flex_attention.py b/torch/nn/attention/flex_attention.py
index 7b0a3354642..ebd17dde5c4 100644
--- a/torch/nn/attention/flex_attention.py
+++ b/torch/nn/attention/flex_attention.py
@@ -711,6 +711,7 @@ def create_mask(
     Q_LEN: int,
     KV_LEN: int,
     device: str = "cuda",
+    _return_score: bool = False,
     _compile: bool = False,
 ) -> Tensor:
     r"""This function creates a mask tensor from a mod_fn function.
@@ -747,7 +748,8 @@ def create_mask(
             score_mod = mod_fn
             score_mod = _vmap_for_bhqkv(score_mod, prefix=(0,))  # first input is score
             out = score_mod(torch.zeros(B, H, Q_LEN, KV_LEN, device=device), b, h, m, n)
-            mask = torch.where(torch.isneginf(out), False, True)
+            if not _return_score:
+                mask = torch.where(torch.isneginf(out), False, True)
             return mask
         elif mod_type == _ModificationType.MASK_MOD:
             mask_mod = mod_fn

Looking forward to know if it make sense to you, thanks!

[1] https://github.com/pytorch/pytorch/blob/87843ee9ab50778a98eda62fd7498d44c69488bd/torch/nn/attention/flex_attention.py#L810-L815