huggingface/candle · error

`q` and `k` last dims must match

Error message

`q` and `k` last dims must match

What it means

In the Metal SDPA kernel, the query and key tensors must agree in their last (head/embedding) dimension so Q·K^T is a valid matmul per head. The op checks q_l.dim(D::Minus1) == k_l.dim(D::Minus1) and bails otherwise.

Source

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

        use candle::backend::BackendStorage;
        use candle_metal_kernels::SdpaDType;

        let device = q.device();

        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;

View on GitHub (pinned to d5fee525bf)

Solutions

  1. Make q and k projections produce the same head_dim (adjust linear layer out_features)
  2. Fix the reshape/split so both q and k end with head_dim as the last dimension
  3. Add a projection layer (linear) to map one side's head dim to the other's
  4. Validate q.dims()[..].last() == k.dims()[..].last() before calling SDPA

Example fix

// before
let q = q_proj.forward(&x)?;   // head_dim 128
let k = k_proj_small.forward(&x)?; // head_dim 64
let out = sdpa(&q, &k, &v, ...)?;
// after
assert_eq!(q.dim(D::Minus1)?, k.dim(D::Minus1)?);
let k = k_proj.forward(&x)?; // head_dim 128
let out = sdpa(&q, &k, &v, ...)?;
Defensive patterns

Strategy: validation

Validate before calling

fn check_qk_last_dim(q: &Tensor, k: &Tensor) -> candle::Result<()> {
    if q.dim(candle::D::Minus1)? != k.dim(candle::D::Minus1)? {
        candle::bail!("q head_dim {:?} != k head_dim {:?}", q.shape(), k.shape());
    }
    Ok(())
}

Type guard

fn qk_dims_ok(q: &Tensor, k: &Tensor) -> bool {
    q.dim(candle::D::Minus1) == k.dim(candle::D::Minus1) && q.dim(candle::D::Minus1).is_ok()
}

Try / catch

match sdpa(&q, &k, &v, &mask, false, Some(scale)) {
    Ok(y) => y,
    Err(e) if e.to_string().contains("last dims must match") => Err(candle::Error::msg("q/k head_dim config bug").bt()),
    Err(e) => Err(e),
}

Prevention

When it happens

Trigger: Calling the SDPA op on Metal with q and k tensors whose last dims differ, e.g. mismatched head_dim between query and key projections, or k carrying a different embedding size.

Common situations: Misconfigured GQA/model projections where q_proj outputs a different head dim than k_proj; reshaping errors that flatten wrong dims; feeding cross-attention with mismatched encoder/decoder dims without a projection.

Related errors


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