{"record":{"id":"89792dc9305b2d42","repo":"jax-ml/jax","slug":"head-dim-of-q-must-be-the-same-as-that-of-k-v-got-89792d","errorCode":null,"errorMessage":"head_dim of Q must be the same as that of K/V. Got {head_dim} and {head_dim_k}.","messagePattern":"head_dim of Q must be the same as that of K/V\\. Got (.+?) and (.+?)\\.","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/experimental/pallas/ops/tpu/paged_attention/paged_attention_kernel.py","lineNumber":451,"sourceCode":"  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:\n    raise ValueError(\n        f\"The dtype of `lengths` must be int32. Got {lengths.dtype}\"\n    )\n\n  # TODO(dinghua): get the actual cores per chip once there's an official API.","sourceCodeStart":433,"sourceCodeEnd":469,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/pallas/ops/tpu/paged_attention/paged_attention_kernel.py#L433-L469","documentation":"Every query head is dotted against K/V with the same head_dim, so q.shape[-1] must equal the head_dim baked into k_pages/v_pages. When the last dimensions differ, the dot product inside the kernel is shape-incompatible and the kernel raises this error before launch.","triggerScenarios":"Calling paged_attention with q head_dim 128 but a KV cache allocated with page tensors of head_dim 64 (or 256); mixing a model checkpoint's head_dim with a cache built for another model.","commonSituations":"Swapping LoRA adapters or models that share tokenizer but not head dims; reusing a persisted KV cache from a different architecture; config drift between rope_theta/head_dim settings.","solutions":["Make q's last dim match the cache: project q to head_dim_k, or rebuild the cache with head_dim == q.shape[-1]","Clear and re-allocate the KV cache whenever the model's head_dim changes","Assert equality before inference: assert q.shape[-1] == k_pages.shape[-1]"],"exampleFix":"// before\nq = proj(x)              # head_dim 128\ncache head_dim = 64\n// after\nq = proj(x) @ w_rescale   # or rebuild cache with head_dim 128","handlingStrategy":"validation","validationCode":"assert q.shape[-1] == k_pages.shape[-1] == v_pages.shape[-1]","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Derive cache allocation from the model's head_dim config","Rebuild caches whenever head_dim-affecting config changes"],"tags":["jax","pallas","tpu","paged-attention","head-dim-validation"],"backgroundTag":"shape-validation-failed","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}