huggingface / huggingface/candle

Metal SDPA panics during quantized LLaMA GQA decode

Open
#3,522 0 comments 0 reactions 0 assignees View on GitHub
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.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.