scaled_dot_product_attention with alibi bias issue on long sequence length

[nlpcat]

🐛 Describe the bug

A100 40GB.

scaled_dot_product_attention + alibi_attention_mask + enable_mem_efficient + seq_length 8192: throw an error bias_4d_view.stride(0) overflows

scaled_dot_product_attention + alibi_attention_mask + enable_math + seq_length 8192: no error

scaled_dot_product_attention + alibi_attention_mask + enable_mem_efficient + seq_length 4096: no error

import torch
dtype = torch.bfloat16
seq_length = 8192
num_heads = 32
query_layer = torch.rand([1, num_heads, seq_length, 64], dtype=dtype, device="cuda:0")
key_layer = torch.rand([1, num_heads, seq_length, 64], dtype=dtype, device="cuda:0")
value_layer = torch.rand([1, num_heads, seq_length, 64], dtype=dtype, device="cuda:0")
alibi = torch.rand([1, num_heads, seq_length, seq_length], dtype=dtype, device="cuda:0")

# throw an error
with torch.backends.cuda.sdp_kernel(enable_flash=False, enable_math=False, enable_mem_efficient=True):
     context_layer = torch.nn.functional.scaled_dot_product_attention(
         query_layer, key_layer, value_layer, attn_mask=alibi, dropout_p=0.0
     )

# no error
with torch.backends.cuda.sdp_kernel(enable_flash=False, enable_math=True, enable_mem_efficient=False):
     context_layer = torch.nn.functional.scaled_dot_product_attention(
         query_layer, key_layer, value_layer, attn_mask=alibi, dropout_p=0.0
     )

Versions

nightly version
torch==2.1.0-dev20230825+cu118

cc @jbschlosser @bhosmer @cpuhrsch @erichan1 @drisspg

[nlpcat]

follow up issues on this #96099 and #104310

[nlpcat]

@drisspg

[drisspg]

Ill take a look