{"record":{"id":"4b6d86ff8c9d1586","repo":"huggingface/candle","slug":"k-and-v-head-dims-must-match","errorCode":null,"errorMessage":"`k` and `v` head dims must match","messagePattern":"`k` and `v` head dims must match","errorType":"validation","errorClass":null,"httpStatus":null,"severity":"error","filePath":"candle-nn/src/ops.rs","lineNumber":1065,"sourceCode":"        let out_dims = vec![q_l.dim(0)?, q_l.dim(1)?, q_l.dim(2)?, v_l.dim(3)?];\n        let elem_count: usize = out_dims.iter().product();\n        let out_shape = Shape::from_dims(&out_dims);\n        let out_layout = Layout::contiguous(out_shape.clone());\n\n        let output = device\n            .new_buffer_builder()\n            .with_size_for(elem_count, q.dtype())\n            .with_label(\"sdpa_o\")\n            .build()?;\n\n        // q,k must have matching emb dim\n        if q_l.dim(D::Minus1)? != k_l.dim(D::Minus1)? {\n            candle::bail!(\"`q` and `k` last dims must match\");\n        }\n\n        // k,v must have matching n kv heads\n        if v_l.dim(D::Minus(3))? != k_l.dim(D::Minus(3))? {\n            candle::bail!(\"`k` and `v` head dims must match\");\n        }\n\n        // n_heads % n_kv_heads == 0; n_heads >= 1, n_kv_heads >= 1.\n        if q_l.dim(D::Minus(3))? % k_l.dim(D::Minus(3))? != 0 {\n            candle::bail!(\"query `n_heads` must be a multiple of `n_kv_heads`\");\n        }\n\n        let k_head = k_l.dim(D::Minus1)?;\n        let q_head = q_l.dim(D::Minus1)?;\n        let q_seq = q_l.dim(2)?;\n        let k_seq = k_l.dim(2)?;\n\n        let mut implementation_supports_use_case = q_head == k_head;\n        let supported_head_dim = q_head == 32\n            || q_head == 64\n            || q_head == 72\n            || q_head == 80\n            || q_head == 96","sourceCodeStart":1047,"sourceCodeEnd":1083,"githubUrl":"https://github.com/huggingface/candle/blob/d5fee525bfde3273eb7c9b75fd2bc4937be867ca/candle-nn/src/ops.rs#L1047-L1083","documentation":"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.","triggerScenarios":"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.","commonSituations":"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.","solutions":["Ensure k and v are reshaped to the same (b, n_kv_heads, seq, head_dim) layout","Fix the kv projection so both k and v produce n_kv_heads heads","Check the KV cache stores k and v with identical head dimensions","Validate k.dim(1) == v.dim(1) (or dim(D::Minus(3))) before calling SDPA"],"exampleFix":"// before\nlet k = k_proj.forward(&x)?.reshape((b, 8, s, hd))?;\nlet v = v_proj.forward(&x)?.reshape((b, 4, s, hd))?; // mismatch\n// after\nlet n_kv = 4;\nlet k = k_proj.forward(&x)?.reshape((b, n_kv, s, hd))?;\nlet v = v_proj.forward(&x)?.reshape((b, n_kv, s, hd))?;","handlingStrategy":"validation","validationCode":"fn check_kv_heads(k: &Tensor, v: &Tensor) -> candle::Result<()> {\n    if k.dim(candle::D::Minus(3))? != v.dim(candle::D::Minus(3))? {\n        candle::bail!(\"k heads {} != v heads {}\", k.dim(1)?, v.dim(1)?);\n    }\n    Ok(())\n}","typeGuard":"fn kv_heads_ok(k: &Tensor, v: &Tensor) -> bool {\n    k.dim(candle::D::Minus(3)).ok() == v.dim(candle::D::Minus(3)).ok()\n}","tryCatchPattern":"match sdpa(&q, &k, &v, &mask, false, Some(scale)) {\n    Ok(y) => y,\n    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()),\n    Err(e) => Err(e),\n}","preventionTips":["Reshape k and v with the same n_kv_heads constant","Verify KV cache stores k/v with identical shapes","Share one n_kv_heads config value for both projections","Assert k.shape() == v.shape() (minus seq-len growth) when writing cache"],"tags":["sdpa","shape-mismatch","gqa","metal"],"backgroundTag":"attention-shape-mismatch","analyzedSha":"d5fee525bfde3273eb7c9b75fd2bc4937be867ca","analyzedAt":"2026-09-02T00:15:47.023Z","contentChangedAt":null,"schemaVersion":2},"datasetVersion":"2026-09-09T06:17:21.866Z"}