{"record":{"id":"cf1acc3aef7af55f","repo":"jax-ml/jax","slug":"k-pages-and-v-pages-must-have-the-same-shape-got-cf1acc","errorCode":null,"errorMessage":"k_pages and v_pages must have the same shape. Got {k_pages.shape} and {v_pages.shape}","messagePattern":"k_pages and v_pages must have the same shape\\. Got (.+?) and (.+?)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/experimental/pallas/ops/tpu/paged_attention/paged_attention_kernel.py","lineNumber":441,"sourceCode":"        k_scales_pages, (*k_scales_pages.shape[:-1], k_pages.shape[-1])\n    )\n  else:\n    k_scales_pages = None\n  if isinstance(v_pages, quantization_utils.QuantizedTensor):\n    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    )","sourceCodeStart":423,"sourceCodeEnd":459,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/pallas/ops/tpu/paged_attention/paged_attention_kernel.py#L423-L459","documentation":"The TPU paged-attention kernel requires K and V page tables to be identical in shape (num_kv_heads, num_pages, page_size, head_dim) because they are indexed in lockstep. If k_pages.shape != v_pages.shape the kernel raises immediately, since page-based lookup would read out of bounds or mix layouts.","triggerScenarios":"Calling paged_attention(q, k_pages, v_pages, ...) where k_pages and v_pages come from different KV-cache allocations, were sliced differently, or one was transposed/padded relative to the other.","commonSituations":"Building a KV cache where V was allocated with a different page count; copying from a HF-style cache that stores K and V separately with padding applied to only one; batched decode loops that trim pages on only one tensor.","solutions":["Allocate K and V caches from the same shape spec and never resize them independently","Re-pad or slice the mismatched tensor to match, e.g. v_pages = v_pages[:, :k_pages.shape[1]]","Verify shapes right before the call: assert k_pages.shape == v_pages.shape"],"exampleFix":"// before\nout = paged_attention(q, k_pages, v_pages[:,:,:-1], ...)\n// after\nout = paged_attention(q, k_pages, v_pages, ...)  # keep caches identical","handlingStrategy":"validation","validationCode":"assert k_pages.shape == v_pages.shape, (k_pages.shape, v_pages.shape)","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Allocate K and V caches together from one shape tuple","Add an invariant check in the cache class that both halves stay identical"],"tags":["jax","pallas","tpu","paged-attention","shape-validation"],"backgroundTag":"shape-validation-failed","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}