{"record":{"id":"073fd08f0a61f444","repo":"jax-ml/jax","slug":"number-of-q-heads-must-be-divisible-by-number-of-k-073fd0","errorCode":null,"errorMessage":"Number of Q heads must be divisible by number of KV heads. Got {num_q_heads} and {num_kv_heads}.","messagePattern":"Number of Q heads must be divisible by number of KV heads\\. Got (.+?) and (.+?)\\.","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/experimental/pallas/ops/tpu/paged_attention/paged_attention_kernel.py","lineNumber":446,"sourceCode":"    v_pages, v_scales_pages = v_pages.weight, v_pages.scales\n    assert isinstance(v_scales_pages, jax.Array)  # For typing.\n    v_scales_pages = jnp.broadcast_to(\n        v_scales_pages, (*v_scales_pages.shape[:-1], v_pages.shape[-1])\n    )\n  else:\n    v_scales_pages = None\n\n  batch_size, num_q_heads, head_dim = q.shape\n  num_kv_heads, _, page_size, head_dim_k = k_pages.shape\n  batch_size_paged_indices, pages_per_sequence = page_indices.shape\n\n  if k_pages.shape != v_pages.shape:\n    raise ValueError(\n        f\"k_pages and v_pages must have the same shape. Got {k_pages.shape} and\"\n        f\" {v_pages.shape}\"\n    )\n  if num_q_heads % num_kv_heads != 0:\n    raise ValueError(\n        \"Number of Q heads must be divisible by number of KV heads. Got\"\n        f\" {num_q_heads} and {num_kv_heads}.\"\n    )\n  if head_dim_k != head_dim:\n    raise ValueError(\n        \"head_dim of Q must be the same as that of K/V. Got\"\n        f\" {head_dim} and {head_dim_k}.\"\n    )\n  if pages_per_sequence % pages_per_compute_block != 0:\n    raise ValueError(\n        \"pages_per_compute_block must be divisible by pages per sequence. Got\"\n        f\" {pages_per_compute_block} and {pages_per_sequence}.\"\n    )\n  if lengths.shape != (batch_size,):\n    raise ValueError(\"`lengths` and `q` must have the same batch size\")\n  if batch_size_paged_indices != batch_size:\n    raise ValueError(\"`page_indices` and `q` must have the same batch size\")\n  if lengths.dtype != jnp.int32:","sourceCodeStart":428,"sourceCodeEnd":464,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/pallas/ops/tpu/paged_attention/paged_attention_kernel.py#L428-L464","documentation":"Paged attention on TPU implements grouped-query attention (GQA) by mapping each KV head to num_q_heads // num_kv_heads query heads. That mapping requires num_q_heads to be an exact multiple of num_kv_heads; otherwise the head-grouping in the kernel is undefined and this ValueError is raised.","triggerScenarios":"Calling paged_attention with q of shape (batch, 8, dim) and k_pages with num_kv_heads=3, or any combination where num_q_heads % num_kv_heads != 0 (e.g. 8 Q heads with 6 KV heads).","commonSituations":"Switching a model from MHA to GQA/MQA with incompatible head counts; loading checkpoints where the KV head replication factor is fractional; typoing head counts in config.","solutions":["Set num_kv_heads to a divisor of num_q_heads (e.g. 1 for MQA, num_q_heads for MHA, or num_q_heads/rep factor)","Fix the query projection so q reshapes to (batch, num_kv_heads * group, dim)","Re-check model config: num_attention_heads and num_key_value_heads must satisfy divisibility"],"exampleFix":"// before\nq shape (b, 8, d); k_pages with num_kv_heads=6\n// after\nq shape (b, 8, d); k_pages with num_kv_heads=2  # 8 % 2 == 0","handlingStrategy":"validation","validationCode":"num_q_heads, num_kv_heads = q.shape[1], k_pages.shape[0]\nassert num_q_heads % num_kv_heads == 0, (num_q_heads, num_kv_heads)","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Validate head counts once at model-build time","Encode num_kv_heads = num_q_heads // group_size in config instead of an independent value"],"tags":["jax","pallas","tpu","paged-attention","gqa","head-dim-validation"],"backgroundTag":"invalid-model-configuration","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}