Repository: pytorch
Issue: Some FlexAttention learned bias bugs/limitations
Labels: triaged,oncall: pt2,module: higher order operators,module: pt2-dispatcher,module: flex attention

### 🐛 Describe the bug

## Ex 1
```Python
import torch
from torch.nn.attention.flex_attention import flex_attention, create_block_mask, create_mask
torch.set_default_device('cuda')

flex_attention = torch.compile(flex_attention, dynamic=False)

result = torch.randn((), requires_grad=True)
def score_mod(score, b, h, q, kv):
    return score * result

S = 8192
torch.manual_seed(0)
q, k, v = [torch.randn(1, 1, S, 64, dtype=torch.float16, requires_grad=True) for _ in range(3)]
flex_attention(q, k, v, score_mod=score_mod).sum().backward()
```

```Shell
  File "/home/chilli/local/pytorch/torch/fx/interpreter.py", line 230, in run_node
    return getattr(self, n.op)(n.target, args, kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/chilli/local/pytorch/torch/_inductor/graph.py", line 1147, in call_function
    raise LoweringException(e, target, args, kwargs).with_traceback(
  File "/home/chilli/local/pytorch/torch/_inductor/graph.py", line 1137, in call_function
    out = lowerings[target](*args, **kwargs)  # type: ignore[index]
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/chilli/local/pytorch/torch/_inductor/lowering.py", line 452, in wrapped
    out = decomp_fn(*args, **kwargs)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/chilli/local/pytorch/torch/_inductor/kernel/flex_attention.py", line 2226, in flex_attention_backward
    joint_outputs = process_joint_outputs(
                    ^^^^^^^^^^^^^^^^^^^^^^
  File "/home/chilli/local/pytorch/torch/_inductor/kernel/flex_attention.py", line 2103, in process_joint_outputs
    grads_out = [get_out(x) for x in other_grads]
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/chilli/local/pytorch/torch/_inductor/kernel/flex_attention.py", line 2103, in <listcomp>
    grads_out = [get_out(x) for x in other_grads]
                 ^^^^^^^^^^
  File "/home/chilli/local/pytorch/torch/_inductor/kernel/flex_attention.py", line 2100, in get_out
    assert buf.name is not None
           ^^^^^^^^^^^^^^^^^^^^
torch._inductor.exc.LoweringException: AssertionError: 
  target: flex_attention_backward
  args[0]: TensorBox(StorageBox(
    InputBuffer(name='primals_1', layout=FixedLayout('cuda:0', torch.float16, size=[1, 1, 8192, 64], stride=[524288, 524288, 64, 1]))
  ))
  args[1]: TensorBox(StorageBox(
    InputBuffer(name='primals_2', layout=FixedLayout('cuda:0', torch.float16, size=[1, 1, 8192, 64], stride=[524288, 524288, 64, 1]))
  ))
  args[2]: TensorBox(StorageBox(
```

## Ex 2
```Python
import torch
from torch.nn.attention.flex_attention import flex_attention, create_block_mask, create_mask
torch.set_default_device('cuda')

flex_attention = torch.compile(flex_attention, dynamic=False)

result = torch.randn((1,), requires_grad=True)
def score_mod(score, b, h, q, kv):
    return score * result[score.new_zeros((), dtype=torch.int)]

S = 8192
torch.manual_seed(0)
q, k, v = [torch.randn(1, 1, S, 64, dtype=torch.float16, requires_grad=True) for _ in range(3)]
flex_attention(q, k, v, score_mod=score_mod).sum().backward()
```

```Shell
Traceback (most recent call last):
  File "/home/chilli/.conda/envs/py311/lib/python3.11/site-packages/triton/language/core.py", line 35, in wrapper
    return fn(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^
  File "/home/chilli/.conda/envs/py311/lib/python3.11/site-packages/triton/language/core.py", line 1268, in broadcast_to
    return semantic.broadcast_impl_shape(input, shape, _builder)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/chilli/.conda/envs/py311/lib/python3.11/site-packages/triton/language/semantic.py", line 732, in broadcast_impl_shape
    raise ValueError(f"Cannot broadcast, rank mismatch: {src_shape}, {shape}")
ValueError: Cannot broadcast, rank mismatch: [1], [64, 64]

The above exception was the direct cause of the following exception:

triton.compiler.errors.CompilationError: at 96:33:
    if CHECK_BLOCK_BOUNDARY:
        grad_scores = tl.where(offs_n2[None, :] < KV_LEN, grad_scores, 0.0)

    # ~~~~~~~~~~~~~~~~~~~ Apply other buffer grad writes ~~~~~~~~~~~~~
    if WRITE_DQ:
        scatter_mask = offs_m2[:, None] < Q_LEN and offs_n2[None, :] < KV_LEN
        tmp12 = tl.full([1], 0, tl.int32)
        tmp13 = (ds)
        tmp14 = (pre_mod_scores)
        tmp15 = tmp13 * tmp14
        tmp16 = tmp15.to(tl.float32)
        tl.atomic_add(in_ptr17 + tl.broadcast_to(tmp12, tmp16.shape), tmp16, scatter_mask, sem='relaxed')
                                 ^

The above exception was the direct cause of the following exception:

triton.compiler.errors.CompilationError: at 79:17:
            dq = bwd_dq_block_mn(
                arg_Q, arg_K, arg_V, arg_LSE, arg_DELTA, arg_DO, arg_DQ, arg_DV, arg_KV_NUM_BLKS, arg_KV_IDX, arg_Q_NUM_BLKS, arg_Q_IDX, arg_FULL_KV_NUM_BLKS, arg_FULL_KV_IDX, arg_FULL_Q_NUM_BLKS, arg_FULL_Q_IDX, in_ptr16, in_ptr17, out_ptr0,
                dq, q, kT_ptrs, vT_ptrs, do, Di, lse, Q_LEN, KV_LEN,
                off_z, off_hq, offs_m2, offs_n2,
                stride_kn, stride_kd, stride_vn, stride_vd,
                kv_indices, sparse_kv_num_blocks,
                MATMUL_PRECISION, RCP_LN2,
                IS_FULL_BLOCKS, CHECK_BLOCK_BOUNDARY=True,
            )
```

### Versions

N/A

cc @chauhang @penguinwu @zou3519 @ydwu4 @bdhirsh @yf225 @drisspg @yanboliang @BoyuanFeng