huggingface/candle · error

`k` and `v` head dims must match

Error message

`k` and `v` head dims must match

What it means

For grouped-query/multi-query attention the key and value tensors must have the same number of KV heads (third dim from the end), since V is expanded against K's head count inside the Metal kernel. The op checks v_l.dim(D::Minus(3)) == k_l.dim(D::Minus(3)) and bails otherwise.

Source

Thrown at candle-nn/src/ops.rs:1065

        let out_dims = vec![q_l.dim(0)?, q_l.dim(1)?, q_l.dim(2)?, v_l.dim(3)?];
        let elem_count: usize = out_dims.iter().product();
        let out_shape = Shape::from_dims(&out_dims);
        let out_layout = Layout::contiguous(out_shape.clone());

        let output = device
            .new_buffer_builder()
            .with_size_for(elem_count, q.dtype())
            .with_label("sdpa_o")
            .build()?;

        // q,k must have matching emb dim
        if q_l.dim(D::Minus1)? != k_l.dim(D::Minus1)? {
            candle::bail!("`q` and `k` last dims must match");
        }

        // k,v must have matching n kv heads
        if v_l.dim(D::Minus(3))? != k_l.dim(D::Minus(3))? {
            candle::bail!("`k` and `v` head dims must match");
        }

        // n_heads % n_kv_heads == 0; n_heads >= 1, n_kv_heads >= 1.
        if q_l.dim(D::Minus(3))? % k_l.dim(D::Minus(3))? != 0 {
            candle::bail!("query `n_heads` must be a multiple of `n_kv_heads`");
        }

        let k_head = k_l.dim(D::Minus1)?;
        let q_head = q_l.dim(D::Minus1)?;
        let q_seq = q_l.dim(2)?;
        let k_seq = k_l.dim(2)?;

        let mut implementation_supports_use_case = q_head == k_head;
        let supported_head_dim = q_head == 32
            || q_head == 64
            || q_head == 72
            || q_head == 80
            || q_head == 96

View on GitHub (pinned to d5fee525bf)

Solutions

1. Ensure k and v are reshaped to the same (b, n_kv_heads, seq, head_dim) layout
2. Fix the kv projection so both k and v produce n_kv_heads heads
3. Check the KV cache stores k and v with identical head dimensions
4. Validate k.dim(1) == v.dim(1) (or dim(D::Minus(3))) before calling SDPA

Example fix

// before
let k = k_proj.forward(&x)?.reshape((b, 8, s, hd))?;
let v = v_proj.forward(&x)?.reshape((b, 4, s, hd))?; // mismatch
// after
let n_kv = 4;
let k = k_proj.forward(&x)?.reshape((b, n_kv, s, hd))?;
let v = v_proj.forward(&x)?.reshape((b, n_kv, s, hd))?;

fn check_kv_heads(k: &Tensor, v: &Tensor) -> candle::Result<()> {
    if k.dim(candle::D::Minus(3))? != v.dim(candle::D::Minus(3))? {
        candle::bail!("k heads {} != v heads {}", k.dim(1)?, v.dim(1)?);
    }
    Ok(())
}

fn kv_heads_ok(k: &Tensor, v: &Tensor) -> bool {
    k.dim(candle::D::Minus(3)).ok() == v.dim(candle::D::Minus(3)).ok()
}

match sdpa(&q, &k, &v, &mask, false, Some(scale)) {
    Ok(y) => y,
    Err(e) if e.to_string().contains("head dims must match") => Err(candle::Error::msg("k/v head count bug in kv cache or projection").bt()),
    Err(e) => Err(e),
}

When it happens

Trigger: Calling SDPA on Metal where k and v have different head counts, e.g. after a wrong reshape/split of the kv projection output or mismatched k/v projections in a custom attention layer.

Common situations: GQA implementations where k and v are split with different head counts by mistake; cache bug where cached k and freshly computed v have different shapes; transposed or wrongly permuted kv tensors.

Related errors

• `q` and `k` last dims must match
• query `n_heads` must be a multiple of `n_kv_heads`
• convtr1d: shape mismatch on c_in {:?} {:?}
• input rank ({}) must be >= weight rank ({})
• Meta SDPA does not support q head dim {q_head}: q dims {:?},

AI-assisted analysis of huggingface/candle@d5fee525bf (2026-09-02). Data as JSON: /api/errors/4b6d86ff8c9d1586. Report an issue: GitHub.