Comfy-Org/ComfyUI · error · ValueError

Key/value head count mismatch for GQA: {key_heads} != {value

Error message

Key/value head count mismatch for GQA: {key_heads} != {value_heads}

What it means

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.

Source

Thrown at comfy/ops.py:42

from comfy.cli_args import args, PerformanceFeature
import comfy.float
import json
import comfy.memory_management
import comfy.pinned_memory
import comfy.utils

import comfy_aimdo.model_vbar
import comfy_aimdo.torch

def run_every_op():
    if torch.compiler.is_compiling():
        return

    comfy.model_management.throw_exception_if_processing_interrupted()

def gqa_repeat_factor(query_heads, key_heads, value_heads):
    if key_heads != value_heads:
        raise ValueError(f"Key/value head count mismatch for GQA: {key_heads} != {value_heads}")
    if query_heads == key_heads:
        return 1
    if query_heads % key_heads != 0:
        raise ValueError(f"Query heads must be divisible by key/value heads for GQA: {query_heads} vs {key_heads}")
    return query_heads // key_heads

def repeat_kv_for_gqa(k, v, query_heads, head_dim):
    n_rep = gqa_repeat_factor(query_heads, k.shape[head_dim], v.shape[head_dim])
    if n_rep > 1:
        k = k.repeat_interleave(n_rep, dim=head_dim)
        v = v.repeat_interleave(n_rep, dim=head_dim)
    return k, v

def scaled_dot_product_attention(q, k, v, *args, **kwargs):
    attn_mask = args[0] if len(args) > 0 else kwargs.get("attn_mask")
    if kwargs.get("enable_gqa", False) and attn_mask is not None:
        k, v = repeat_kv_for_gqa(k, v, q.shape[-3], -3)
        kwargs["enable_gqa"] = False

View on GitHub (pinned to 1c6d8d45b3)

Solutions

1. Check the K and V projection shapes in the checkpoint and confirm they use the same head count.
2. Verify head_dim indexes the head axis in your tensor layout (commonly 2 for (B,S,H,D)).
3. If writing a custom attention, reshape K/V to the layout repeat_kv_for_gqa expects before calling.

Example fix

# before
k, v = repeat_kv_for_gqa(k, v, q_heads, head_dim=3)  # layout is (B,S,H,D)
# after
k, v = repeat_kv_for_gqa(k, v, q_heads, head_dim=2)

k_heads, v_heads = k.shape[head_dim], v.shape[head_dim]
if k_heads != v_heads:
    raise ValueError(f"K/V head mismatch ({k_heads} vs {v_heads}); check projections or head_dim axis")
k2, v2 = repeat_kv_for_gqa(k, v, q_heads, head_dim)

def kv_heads_match(k, v, head_dim) -> bool:
    return k.shape[head_dim] == v.shape[head_dim]

When it happens

Trigger: 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).

Common situations: 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.

Related errors

• {name} inner dimension {inner_dim} is not divisible by head
• Query heads must be divisible by key/value heads for GQA: {q
• `only_cross_attention` can only be set to True if `added_kv_
• Hidden size {params.hidden_size} must be divisible by num_he
• Hidden size {params.hidden_size} must be divisible by num_he

AI-assisted analysis of Comfy-Org/ComfyUI@1c6d8d45b3 (2026-08-14). Data as JSON: /api/errors/dc1a6fc6f6262712. Report an issue: GitHub.