{"record":{"id":"42955700af18eef8","repo":"jax-ml/jax","slug":"the-number-of-query-heads-must-be-a-multiple-of-ke","errorCode":null,"errorMessage":"The number of query heads must be a multiple of key/value heads, but got {query_arr.shape[-2]} vs {K}","messagePattern":"The number of query heads must be a multiple of key/value heads, but got (.+?) vs (.+?)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/nn/functions.py","lineNumber":1216,"sourceCode":"    if t.ndim != len(shape):\n      raise ValueError(f\"{name} ndim should be {len(shape)}, but got {t.ndim}\")\n    if dtype is not None and t.dtype != dtype:\n      raise ValueError(f\"{name} dtype should be {dtype}, but got {t.dtype}\")\n    for i in range(t.ndim):\n      if shape[i] != -1 and t.shape[i] != shape[i]:\n        raise ValueError(f\"{name} shape should be {shape}: but got {t.shape}\")\n\n  B, S, K, H = key_arr.shape\n  _check_shape_and_dtype(value_arr, [B, S, K, H], key_arr.dtype, 'value')\n  _check_shape_and_dtype(query_arr, [B, -1, -1, H], key_arr.dtype, 'query')\n  _check_shape_and_dtype(mask, [-1] * 4, np.dtype(bool), 'mask')\n  _check_shape_and_dtype(bias, [-1] * 4, None, 'bias')\n  _check_shape_and_dtype(query_seq_lengths, [B], np.dtype('int32'),\n                         'query_seq_lengths')\n  _check_shape_and_dtype(key_value_seq_lengths, [B], np.dtype('int32'),\n                         'key_value_seq_lengths')\n  if query_arr.shape[-2] % K != 0:\n    raise ValueError(f\"The number of query heads must be a multiple of \"\n                     f\"key/value heads, but got {query_arr.shape[-2]} vs {K}\")\n\n  scale_val = (1.0 / np.sqrt(H)) if scale is None else scale\n\n  match implementation:\n    case 'xla':\n      out = _dot_product_attention_xla(\n          query_arr, key_arr, value_arr, bias, mask, is_causal=is_causal,\n          scale=scale_val, q_seqlen=query_seq_lengths,\n          kv_seqlen=key_value_seq_lengths,\n          local_window_size=local_window_size,\n          return_residual=return_residual,\n      )\n    case 'cudnn':\n      use_padding = (\n           query_seq_lengths is not None or key_value_seq_lengths is not None\n      )\n      if use_padding:","sourceCodeStart":1198,"sourceCodeEnd":1234,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/nn/functions.py#L1198-L1234","documentation":"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).","triggerScenarios":"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.","commonSituations":"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.","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"],"exampleFix":"// before\nn_q_heads, n_kv_heads = 30, 8  # 30 % 8 != 0\n\n// after\nn_q_heads, n_kv_heads = 32, 8  # 32 % 8 == 0 (GQA factor 4)","handlingStrategy":"validation","validationCode":"assert q.shape[-2] % k.shape[-2] == 0, (\n    f'query heads {q.shape[-2]} not a multiple of kv heads {k.shape[-2]}')","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Validate head counts in model config: num_q_heads % num_kv_heads == 0","Add a config sanity check function for transformer hyperparams"],"tags":["jax","nn","attention","gqa","shape-validation"],"backgroundTag":"shape-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}