Metal SDPA kernel produces NaN with non-square mask + GQA

Hey there! Working with Opus 4.6, I stumbled upon what might be a limitation in candle and thought I would share it with you. Here is the summary Opus came up with:

## Summary

`candle_nn::ops::sdpa` on Metal produces NaN outputs when called with:
- A **non-square attention mask** (Q sequence length != K sequence length)
- **Grouped Query Attention** (different number of Q and K heads)

Each of these works individually. The bug only manifests when both conditions are true simultaneously.

## Reproduction

```rust
use candle_core::{Device, DType, Tensor};

let device = Device::new_metal(0)?;
let dtype = DType::F16;

let b = 1;
let q_heads = 32;
let kv_heads = 8;
let q_len = 98; // non-square: q_len != kv_len
let kv_len = 1748;
let head_dim = 128;

let q = Tensor::randn(0f32, 1.0, (b, q_heads, q_len, head_dim), &device)?.to_dtype(dtype)?;
let k = Tensor::randn(0f32, 1.0, (b, kv_heads, kv_len, head_dim), &device)?.to_dtype(dtype)?;
let v = Tensor::randn(0f32, 1.0, (b, kv_heads, kv_len, head_dim), &device)?.to_dtype(dtype)?;

// Non-square mask: (q_len, kv_len) with -inf for masked positions
let mask_data: Vec = (0..q_len)
.flat_map(|i| (0..kv_len).map(move |j| if j <= i + (kv_len - q_len) { 0.0 } else { f32::NEG_INFINITY }))
.collect();
let mask = Tensor::from_vec(mask_data, (q_len, kv_len), &device)?
.to_dtype(dtype)?
.broadcast_as((b, q_heads, q_len, kv_len))?;

let scale = 1.0 / (head_dim as f32).sqrt();
let out = candle_nn::ops::sdpa(&q, &k, &v, Some(&mask), false, scale, 1.0)?;

// Check for NaN
let nan_count = out.to_dtype(DType::F32)?
.ne(&out.to_dtype(DType::F32)?)? // NaN != NaN
.to_dtype(DType::F32)?
.sum_all()?
.to_scalar::()?;

assert_eq!(nan_count, 0.0, "SDPA output contains {nan_count} NaN values");
```

## Conditions

| Q heads | K heads | Mask shape | Result |
|---------|---------|------------|--------|
| 32 | 32 | square (N, N) | OK |
| 32 | 8 | square (N, N) | OK |
| 32 | 32 | non-square (M, N) | OK (unverified) |
| **32** | **8** | **non-square (M, N)** | **NaN** |

- Tested on Apple M-series (Metal backend)
- DType: F16 (typical for GGUF inference)
- The output is entirely NaN (~200K NaN values for the dimensions above)
- Inputs (Q, K, V, mask) contain no NaN and have reasonable magnitudes

## Context

This was discovered while implementing KV cache reuse for position-independent caching (PIC). During cache-reuse prefill, only a subset of tokens (Q) attend to the full cached KV, producing a non-square mask. The standard prefill path (square mask) and decode path (no mask, seq_len=1) are unaffected.

## Workaround

Bypass Metal SDPA for the non-square + GQA case and use manual attention:

```rust
let k = repeat_kv(k, n_kv_groups)?; // expand 8 -> 32 heads
let v = repeat_kv(v, n_kv_groups)?;
let scale = (head_dim as f64).sqrt();
let att = (q.matmul(&k.transpose(2, 3)?)? / scale)?;
let att = att.broadcast_add(&mask)?;
let att = candle_nn::ops::softmax_last_dim(&att)?;
att.matmul(&v)?
```

This stays on GPU but loses the kernel fusion and GQA-native handling of the Metal SDPA kernel.

Start at the candle_nn::ops::sdpa entry point and trace the Metal SDPA path. Reproduce the supplied F16 case, then compare non-square masks with equal and differing Q/K head counts. Done means the non-square plus GQA case produces no NaNs while the other listed combinations remain working.

Written by the indexing model from the issue text.