sdpaex : fix backward when attn_mask is not provided

[kshitij12345]

Repro

import torch
import thunder
from thunder.executors.sdpaex import sdpa_ex

batch = 10
seq_len = 128
num_heads = 4
dim_per_head = 32

requires_grad=True
query = torch.randn([batch, seq_len, num_heads, dim_per_head], device="cuda", requires_grad=requires_grad)
key = torch.randn([batch, seq_len, num_heads, dim_per_head], device="cuda", requires_grad=requires_grad)
value = torch.randn([batch, seq_len, num_heads, dim_per_head], device="cuda", requires_grad=requires_grad)

def fn(query, key, value):
    return torch.nn.functional.scaled_dot_product_attention(query, key, value)

cfn = thunder.jit(fn, executors=[sdpa_ex])

thunder_result = cfn(query, key, value)

grad_output = torch.rand_like(thunder_result)
grads = torch.autograd.grad(thunder_result, (query, key, value), grad_outputs=grad_output)

Error

File "thunder.backward_fn_2", line 19, in backward_fn
  File "/home/kkalambarkar/lightning-thunder/thunder/executors/sdpaex.py", line 325, in _scaled_dot_product_efficient_attention_backward_impl
    if not utils.same_shape(grad_attn_mask.shape, attn_mask.shape):
AttributeError: 'NoneType' object has no attribute 'shape'

Problem -

lightning-thunder/thunder/executors/sdpaex.py

Lines 324 to 327 in 35ca2e9

	if not utils.same_shape(grad_attn_mask.shape, attn_mask.shape):
	# Needs to sum over the number of heads dimension in grad_attn_mask
	# if the number of heads in attention mask is expanded in _attention_mask_memory_efficient_helper.
	grad_attn_mask = torch.sum(grad_attn_mask, dim=1, keepdim=True)

It is possible that here grad_attn_mask and attn_mask are both None, so calling .shape leads to error.

[mruberry]