jax-ml/jax · error · ValueError

In MHA, expected number of 'key' heads ({num_kv_heads}) to b

Error message

In MHA, expected number of 'key' heads ({num_kv_heads}) to be a multiple of the number of 'query' heads ({num_q_heads})

What it means

For multi-head attention (not MQA), the number of query heads must be divisible by the number of KV heads (GQA). Note the message text says 'key heads to be a multiple of query heads' but the check is num_q_heads % num_kv_heads != 0 — i.e. q_heads must be a multiple of kv_heads; otherwise per-KV-head grouping of query heads is impossible.

Source

Thrown at jax/experimental/pallas/ops/tpu/splash_attention/splash_attention_kernel.py:956

    raise ValueError(
        "partial_mask_blocks must be of type np.bool_ but got"
        f" {partial_mask_blocks.dtype}"
    )

  if len(k.shape) != expected_kv_rank:
    raise ValueError(
        f"Expected {expected_kv_rank}-dim 'key' tensor for MQA. Instead got a"
        f" {len(k.shape)}-dim one."
    )

  if k.shape[kv_head_dimension] != head_dim_qk:
    raise ValueError(
        f"Expected 'key' head dimension to be: {head_dim_qk}. Instead got:"
        f" {k.shape[kv_head_dimension]}."
    )

  if not is_mqa and num_q_heads % num_kv_heads != 0:
    raise ValueError(
        f"In MHA, expected number of 'key' heads ({num_kv_heads}) to be a"
        f" multiple of the number of 'query' heads ({num_q_heads})"
    )

  if k.shape[:-1] != v.shape[:-1]:
    raise ValueError(
        f"Expected 'key' {k.shape} and 'value' {v.shape} to have the same "
        "leading dimensions."
    )

  assert bkv_compute is not None
  if bkv % bkv_compute:
    raise ValueError(f"{bkv=} must be a multiple of {bkv_compute=}.")
  if bkv_compute % NUM_LANES:
    raise ValueError(f"{bkv_compute=} must be a multiple of {NUM_LANES}.")

  kv_seq_len = k.shape[kv_seq_len_dimension]

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Make num_q_heads a multiple of num_kv_heads (e.g. 32 Q / 8 KV)
If you truly need unmatched heads, repeat/expand K,V heads: jnp.repeat(k, num_q_heads//num_kv_heads, axis=0) after choosing an integer ratio, or run MHA (equal heads)
Double-check the head axis order — you may have swapped num_heads and num_kv_heads when reshaping

Example fix

// before
q: [12, seq, hd], k/v: [8, seq, hd]
// after
q: [16, seq, hd], k/v: [8, seq, hd]  # or repeat k,v to 16 heads for MHA

Defensive patterns

Strategy: validation

Validate before calling

assert q.shape[0] % k.shape[0] == 0, 'num_q_heads must be divisible by num_kv_heads'

Type guard

def gqa_ok(q_heads: int, kv_heads: int) -> bool:
    return q_heads % kv_heads == 0

Prevention

Sanity-check GQA ratios in config validation
Print head counts once at model init

When it happens

Trigger: Passing q with 12 heads and k/v with 8 heads; any head configuration where num_q_heads % num_kv_heads != 0 (e.g. 48 Q heads, 32 KV heads).

Common situations: Using GQA ratios from Llama-style models incorrectly (e.g. 8 KV heads is fine for 32 Q heads, but 7 KV heads is not); mixing configs between model and kernel.

Related errors

AI-assisted analysis of jax-ml/jax@1e1c6a8fc0 (2026-08-27). Data as JSON: /api/errors/3980d98554f92269. Report an issue: GitHub.