Repository: pytorch
Issue: [FlexAttention] Support the number of shared query heads in GQA to not be the power of 2
Labels: triaged,oncall: pt2,module: higher order operators,module: pt2-dispatcher,module: flex attention

### 🚀 The feature, motivation and pitch

I've implemented FlexAttention for SmolLM, but since HF team uses 3 query heads per KV-head, the compilation fails for `seq_length=5`:

```
LoweringException: ValueError: Number of shared query heads sharing the same KV head must be power of 2. 
  target: flex_attention
  args[0]: TensorBox(StorageBox(
    InputBuffer(name='primals_1', layout=FixedLayout('cuda:0', torch.float32, size=[1, 9, 5, 64], stride=[2880, 320, 64, 1]))
  ))
  args[1]: TensorBox(StorageBox(
    InputBuffer(name='primals_2', layout=FixedLayout('cuda:0', torch.float32, size=[1, 3, 5, 64], stride=[960, 320, 64, 1]))
  ))
  args[2]: TensorBox(StorageBox(
    InputBuffer(name='primals_3', layout=FixedLayout('cuda:0', torch.float32, size=[1, 3, 5, 64], stride=[960, 320, 64, 1]))
  ))
  args[3]: Subgraph(name='sdpa_score0', graph_module=<lambda>(), graph=None)
  args[4]: (5, 5, TensorBox(StorageBox(
    InputBuffer(name='primals_5', layout=FixedLayout('cuda:0', torch.int32, size=[1, 1, 1], stride=[1, 1, 1]))
  )), TensorBox(StorageBox(
    InputBuffer(name='primals_4', layout=FixedLayout('cuda:0', torch.int32, size=[1, 1, 1, 1], stride=[1, 1, 1, 1]))
  )), TensorBox(StorageBox(
    InputBuffer(name='primals_9', layout=FixedLayout('cuda:0', torch.int32, size=[1, 1, 1], stride=[1, 1, 1]))
  )), TensorBox(StorageBox(
    InputBuffer(name='primals_10', layout=FixedLayout('cuda:0', torch.int32, size=[1, 1, 1, 1], stride=[1, 1, 1, 1]))
  )), TensorBox(StorageBox(
    InputBuffer(name='primals_11', layout=FixedLayout('cuda:0', torch.int32, size=[1, 1, 1], stride=[1, 1, 1]))
  )), TensorBox(StorageBox(
    InputBuffer(name='primals_12', layout=FixedLayout('cuda:0', torch.int32, size=[1, 1, 1, 1], stride=[1, 1, 1, 1]))
  )), TensorBox(StorageBox(
    InputBuffer(name='primals_13', layout=FixedLayout('cuda:0', torch.int32, size=[1, 1, 1], stride=[1, 1, 1]))
  )), TensorBox(StorageBox(
    InputBuffer(name='primals_14', layout=FixedLayout('cuda:0', torch.int32, size=[1, 1, 1, 1], stride=[1, 1, 1, 1]))
  )), 128, 128, Subgraph(name='sdpa_mask0', graph_module=<lambda>(), graph=None))
  args[5]: 0.125
  args[6]: {'PRESCALE_QK': False, 'ROWS_GUARANTEED_SAFE': False, 'BLOCKS_ARE_CONTIGUOUS': False, 'OUTPUT_LOGSUMEXP': True}
  args[7]: ()
  args[8]: (TensorBox(StorageBox(
    InputBuffer(name='primals_6', layout=FixedLayout('cuda:0', torch.int64, size=[5], stride=[1]))
  )), TensorBox(StorageBox(
    InputBuffer(name='primals_7', layout=FixedLayout('cuda:0', torch.bool, size=[5], stride=[1]))
  )), TensorBox(StorageBox(
    InputBuffer(name='primals_8', layout=FixedLayout('cuda:0', torch.int64, size=[5], stride=[1]))
  )))

Set TORCH_LOGS="+dynamo" and TORCHDYNAMO_VERBOSE=1 for more information


You can suppress this exception and fall back to eager by setting:
    import torch._dynamo
    torch._dynamo.config.suppress_errors = True
```

Weird thing is that the compilation works for `seq_length>128`!!!

### Alternatives

I believe I can pad the head dimension with zeros while I wait for the fix. Please correct me if I am wrong.

### Additional context

I use nightly build: `2.6.0.dev20241212+cu124`

Minimal reproduction code:
```python
from torch.nn.attention.flex_attention import flex_attention, create_block_mask
import torch

seq_len = 128  # 5

# smollest SmolLM 
torch.compile(flex_attention, fullgraph=True, dynamic=False, mode="max-autotune-no-cudagraphs", disable=False)(
    torch.randn(1, 9, seq_len, 64).to('cuda'),
    torch.randn(1, 3, seq_len, 64).to('cuda'),
    torch.randn(1, 3, seq_len, 64).to('cuda'),
    block_mask=create_block_mask(
        lambda b, h, q_idx, kv_idx: q_idx >= kv_idx, 
        None, None, seq_len, seq_len, 
        _compile=True
    ),
    enable_gqa=True
)
```

Change `seq_len` to 5 to reproduce the error.

@drisspg  you've been a great help with FlexAttention, could you have a look here? Thank you!

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