huggingface / huggingface/candle
Metal SDPA panics during quantized LLaMA GQA decode
- Dominant language
- Rust
- Stars
- 21k
- Forks
- 1.8k
- Avg merge
- 16h 42m
- Merged PRs (30d)
- 25
Description
## Summary
`candle-transformers` quantized LLaMA inference can panic on macOS/Metal during the single-token decode path for GQA models.
Observed panic:
```text
index out of bounds: the len is 2 but the index is 2
```
This is related to, but distinct from, #3388. That issue covers Metal SDPA with non-square masks + GQA and notes that standard decode with `seq_len == 1` is unaffected. In this case the panic happens after prompt prefill, when the decode loop calls `ModelWeights::forward()` with a single next-token tensor and `index_pos == prompt_len`.
## Environment
- OS: macOS 25.4.0
- Backend selected by app: `Device::metal_if_available(0)`
- Candle version: `0.10.2`
- Model path: GGUF through `candle-transformers::models::quantized_llama::ModelWeights`
- Model family: TinyLlama / LLaMA GQA GGUF (`llama.attention.head_count_kv != llama.attention.head_count`)
## Path that appears to fail
In `candle-transformers/src/models/quantized_llama.rs`, `LayerWeights::forward_attn()` routes Metal single-token decode through SDPA:
```rust
let y = if q.device().is_metal() && seq_len == 1 {
candle_nn::ops::sdpa(
&q,
&k,
&v,
None,
false,
1. / (self.head_dim as f32).sqrt(),
1.,
)?
} else {
let k = crate::utils::repeat_kv(k, self.n_head / self.n_kv_head)?;
let v = crate::utils::repeat_kv(v, self.n_head / self.n_kv_head)?;
// manual attention path
}
```
For GQA, `self.n_head != self.n_kv_head`, so the Metal SDPA branch receives Q heads and KV heads with different counts. The manual branch expands KV via `repeat_kv` first and does not hit this panic.
## Workaround downstream
In our app we had to skip Metal for LLaMA GQA GGUFs before model load:
```rust
head_count_kv > 0 && head_count_kv != head_count
```
When that condition is true, we select CPU instead of Metal. This avoids the panic, but gives up Metal acceleration for TinyLlama-style models.
## Possible fix
A minimal fix may be to avoid the Metal SDPA fast path for GQA decode and use the existing manual attention path instead:
```rust
let y = if q.device().is_metal() && seq_len == 1 && self.n_head == self.n_kv_head {
candle_nn::ops::sdpa(...)
} else {
let k = crate::utils::repeat_kv(k, self.n_head / self.n_kv_head)?;
let v = crate::utils::repeat_kv(v, self.n_head / self.n_kv_head)?;
// existing manual attention path
}
```
That should keep Metal for non-GQA decode, while avoiding the panic for GQA models.
Contributor guide
No contributing guide indexed for this repository
Assessment
This issue has not been assessed yet.