2D/3D Mask Attention issue

[Cemberk]

Issue seen during Huggingface Transformers, Cohere2 model tests during the test test_generation_beyond_sliding_window_0_flash_attention_2

Minimal Reproducer:

import torch
import torch.nn.functional as F

query = torch.randn(2, 32, 128, 64, device='cuda', dtype=torch.float16)
key = torch.randn(2, 32, 128, 64, device='cuda', dtype=torch.float16) 
value = torch.randn(2, 32, 128, 64, device='cuda', dtype=torch.float16)

mask = torch.ones(2, 128, device='cuda', dtype=torch.bool)

output = F.scaled_dot_product_attention(query, key, value, attn_mask=mask, is_causal=True)

[rocm-repo-management-api]

Jenkins build for 70f58b228561dd635e742485616d0665d5e0f403 commit finished as FAILURE
Links: Blue Ocean view / Build artifacts