jax-ml/jax · error · ValueError

The number of query heads must be a multiple of key/value he

Error message

The number of query heads must be a multiple of key/value heads, but got {query_arr.shape[-2]} vs {K}

What it means

For grouped-query attention (GQA), the number of query heads must be an integer multiple of the number of key/value heads. dot_product_attention raises this when query_arr.shape[-2] % K != 0 (K = key/value head count).

Source

Thrown at jax/_src/nn/functions.py:1216

    if t.ndim != len(shape):
      raise ValueError(f"{name} ndim should be {len(shape)}, but got {t.ndim}")
    if dtype is not None and t.dtype != dtype:
      raise ValueError(f"{name} dtype should be {dtype}, but got {t.dtype}")
    for i in range(t.ndim):
      if shape[i] != -1 and t.shape[i] != shape[i]:
        raise ValueError(f"{name} shape should be {shape}: but got {t.shape}")

  B, S, K, H = key_arr.shape
  _check_shape_and_dtype(value_arr, [B, S, K, H], key_arr.dtype, 'value')
  _check_shape_and_dtype(query_arr, [B, -1, -1, H], key_arr.dtype, 'query')
  _check_shape_and_dtype(mask, [-1] * 4, np.dtype(bool), 'mask')
  _check_shape_and_dtype(bias, [-1] * 4, None, 'bias')
  _check_shape_and_dtype(query_seq_lengths, [B], np.dtype('int32'),
                         'query_seq_lengths')
  _check_shape_and_dtype(key_value_seq_lengths, [B], np.dtype('int32'),
                         'key_value_seq_lengths')
  if query_arr.shape[-2] % K != 0:
    raise ValueError(f"The number of query heads must be a multiple of "
                     f"key/value heads, but got {query_arr.shape[-2]} vs {K}")

  scale_val = (1.0 / np.sqrt(H)) if scale is None else scale

  match implementation:
    case 'xla':
      out = _dot_product_attention_xla(
          query_arr, key_arr, value_arr, bias, mask, is_causal=is_causal,
          scale=scale_val, q_seqlen=query_seq_lengths,
          kv_seqlen=key_value_seq_lengths,
          local_window_size=local_window_size,
          return_residual=return_residual,
      )
    case 'cudnn':
      use_padding = (
           query_seq_lengths is not None or key_value_seq_lengths is not None
      )
      if use_padding:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Set num_kv_heads to a divisor of num_query_heads (e.g. 32 query heads with 8 or 4 KV heads)
If you didn't intend GQA, make query and key/value head counts equal
Check config plumbing from your config object down to the projection layer shapes

Example fix

// before
n_q_heads, n_kv_heads = 30, 8  # 30 % 8 != 0

// after
n_q_heads, n_kv_heads = 32, 8  # 32 % 8 == 0 (GQA factor 4)

Defensive patterns

Strategy: validation

Validate before calling

assert q.shape[-2] % k.shape[-2] == 0, (
    f'query heads {q.shape[-2]} not a multiple of kv heads {k.shape[-2]}')

Prevention

Validate head counts in model config: num_q_heads % num_kv_heads == 0
Add a config sanity check function for transformer hyperparams

When it happens

Trigger: query with 12 heads and key/value with 8 heads; accidentally passing num_key_heads that doesn't divide num_query_heads when building GQA/MHA modules.

Common situations: Configuring GQA (e.g. 32 query heads, 8 KV heads works; 30 vs 8 fails); editing transformer config YAML where kv_heads is set independently; porting models with non-divisible head layouts.

Related errors

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