{"record":{"id":"dc1a6fc6f6262712","repo":"Comfy-Org/ComfyUI","slug":"key-value-head-count-mismatch-for-gqa-key-heads","errorCode":null,"errorMessage":"Key/value head count mismatch for GQA: {key_heads} != {value_heads}","messagePattern":"Key/value head count mismatch for GQA: (.+?) != (.+?)","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"comfy/ops.py","lineNumber":42,"sourceCode":"from comfy.cli_args import args, PerformanceFeature\nimport comfy.float\nimport json\nimport comfy.memory_management\nimport comfy.pinned_memory\nimport comfy.utils\n\nimport comfy_aimdo.model_vbar\nimport comfy_aimdo.torch\n\ndef run_every_op():\n    if torch.compiler.is_compiling():\n        return\n\n    comfy.model_management.throw_exception_if_processing_interrupted()\n\ndef gqa_repeat_factor(query_heads, key_heads, value_heads):\n    if key_heads != value_heads:\n        raise ValueError(f\"Key/value head count mismatch for GQA: {key_heads} != {value_heads}\")\n    if query_heads == key_heads:\n        return 1\n    if query_heads % key_heads != 0:\n        raise ValueError(f\"Query heads must be divisible by key/value heads for GQA: {query_heads} vs {key_heads}\")\n    return query_heads // key_heads\n\ndef repeat_kv_for_gqa(k, v, query_heads, head_dim):\n    n_rep = gqa_repeat_factor(query_heads, k.shape[head_dim], v.shape[head_dim])\n    if n_rep > 1:\n        k = k.repeat_interleave(n_rep, dim=head_dim)\n        v = v.repeat_interleave(n_rep, dim=head_dim)\n    return k, v\n\ndef scaled_dot_product_attention(q, k, v, *args, **kwargs):\n    attn_mask = args[0] if len(args) > 0 else kwargs.get(\"attn_mask\")\n    if kwargs.get(\"enable_gqa\", False) and attn_mask is not None:\n        k, v = repeat_kv_for_gqa(k, v, q.shape[-3], -3)\n        kwargs[\"enable_gqa\"] = False","sourceCodeStart":24,"sourceCodeEnd":60,"githubUrl":"https://github.com/Comfy-Org/ComfyUI/blob/1c6d8d45b3693bfbb32385b410d813a7fd6be216/comfy/ops.py#L24-L60","documentation":"gqa_repeat_factor() implements grouped-query attention by repeating K/V heads to match query heads; this only works when the key and value head counts are equal. A K/V mismatch means the attention weights are malformed or the head-dim argument points at the wrong axis, so it fails before any repeat arithmetic.","triggerScenarios":"repeat_kv_for_gqa(k, v, q_heads, head_dim) where k.shape[head_dim] != v.shape[head_dim] — e.g. an attention module whose KV projections have diverging head counts, or head_dim passed as the wrong axis (e.g. 2 vs 3 for different layouts).","commonSituations":"Loading a checkpoint with mismatched kv head configs; permuting K/V tensors to a different layout (BSHD vs BHSD) without updating head_dim; buggy fused-attention shims.","solutions":["Check the K and V projection shapes in the checkpoint and confirm they use the same head count.","Verify head_dim indexes the head axis in your tensor layout (commonly 2 for (B,S,H,D)).","If writing a custom attention, reshape K/V to the layout repeat_kv_for_gqa expects before calling."],"exampleFix":"# before\nk, v = repeat_kv_for_gqa(k, v, q_heads, head_dim=3)  # layout is (B,S,H,D)\n# after\nk, v = repeat_kv_for_gqa(k, v, q_heads, head_dim=2)","handlingStrategy":"validation","validationCode":"k_heads, v_heads = k.shape[head_dim], v.shape[head_dim]\nif k_heads != v_heads:\n    raise ValueError(f\"K/V head mismatch ({k_heads} vs {v_heads}); check projections or head_dim axis\")\nk2, v2 = repeat_kv_for_gqa(k, v, q_heads, head_dim)","typeGuard":"def kv_heads_match(k, v, head_dim) -> bool:\n    return k.shape[head_dim] == v.shape[head_dim]","tryCatchPattern":null,"preventionTips":["Confirm K and V projections use the same head count in the config.","Double-check head_dim against your (B,S,H,D) vs (B,H,S,D) layout."],"tags":["attention","gqa","shape-mismatch","heads"],"backgroundTag":null,"analyzedSha":"1c6d8d45b3693bfbb32385b410d813a7fd6be216","analyzedAt":"2026-08-14T19:37:18.893Z","schemaVersion":2},"datasetVersion":"2026-08-15T22:17:37.221Z"}