{"record":{"id":"3468d119efd0bd6e","repo":"jax-ml/jax","slug":"pages-per-compute-block-must-be-divisible-by-pages","errorCode":null,"errorMessage":"pages_per_compute_block must be divisible by pages per sequence. Got {pages_per_compute_block} and {pages_per_sequence}.","messagePattern":"pages_per_compute_block must be divisible by pages per sequence\\. Got (.+?) and (.+?)\\.","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/experimental/pallas/ops/tpu/paged_attention/paged_attention_kernel.py","lineNumber":456,"sourceCode":"  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.\n  if megacore_mode == \"kv_head\":\n    if num_kv_heads % 2 != 0:\n      raise ValueError(\n          \"number of KV heads must be even when megacore_mode is 'kv_head'\"\n      )","sourceCodeStart":438,"sourceCodeEnd":474,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/pallas/ops/tpu/paged_attention/paged_attention_kernel.py#L438-L474","documentation":"The paged-attention kernel processes KV pages in fixed blocks of pages_per_compute_block; pages_per_sequence (from page_indices.shape[1]) must be divisible by it so each compute block covers whole pages across all sequences. Note the check is written as pages_per_sequence % pages_per_compute_block, so in practice pages_per_compute_block must be 1 or a divisor of pages_per_seq.","triggerScenarios":"Calling paged_attention with pages_per_compute_block=2 while page_indices has an odd pages_per_sequence (e.g. 7); passing a larger compute block than the per-sequence page budget.","commonSituations":"Tuning pages_per_compute_block for performance without checking cache geometry; variable sequence budgets where pages_per_seq is not a multiple of the block size.","solutions":["Use pages_per_compute_block=1 (always valid)","Make pages_per_sequence a multiple of pages_per_compute_block (e.g. allocate 8 pages/seq for block=4)","Recompute page_indices padding so shape[1] divides evenly"],"exampleFix":"// before\npaged_attention(q, k, v, page_indices, lens, pages_per_compute_block=4)  # pages_per_seq=6\n// after\npaged_attention(q, k, v, page_indices, lens, pages_per_compute_block=2)  # 6 % 2 == 0","handlingStrategy":"validation","validationCode":"pages_per_seq = page_indices.shape[1]\nassert pages_per_seq % pages_per_compute_block == 0, (pages_per_seq, pages_per_compute_block)","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Default to pages_per_compute_block=1; only raise it after fixing cache geometry","Keep pages_per_seq a power-of-two multiple of the compute block"],"tags":["jax","pallas","tpu","paged-attention","layout-validation"],"backgroundTag":"invalid-kernel-configuration","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}