{"record":{"id":"770c49b21ed96738","repo":"jax-ml/jax","slug":"page-indices-and-q-must-have-the-same-batch-si","errorCode":null,"errorMessage":"`page_indices` and `q` must have the same batch size","messagePattern":"`page_indices` and `q` must have the same batch size","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/experimental/pallas/ops/tpu/paged_attention/paged_attention_kernel.py","lineNumber":463,"sourceCode":"  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      )\n    num_cores = 2\n  elif megacore_mode == \"batch\":\n    if batch_size % 2 != 0:\n      raise ValueError(\"batch size must be even when megacore_mode is 'batch'\")\n    num_cores = 2\n  elif megacore_mode is None:\n    num_cores = 1","sourceCodeStart":445,"sourceCodeEnd":481,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/pallas/ops/tpu/paged_attention/paged_attention_kernel.py#L445-L481","documentation":"page_indices maps each sequence in the batch to its KV pages, so its leading dimension must equal q's batch dimension. If page_indices.shape[0] != q.shape[0] the kernel cannot associate queries with their page tables and raises immediately.","triggerScenarios":"Passing a page_indices table sized for a different batch (e.g. max_num_seqs from vLLM-style scheduling) while q contains only the currently-scheduled sequences.","commonSituations":"Porting vLLM-style continuous batching where page_indices covers capacity, not the current batch; multi-step decode loops that shrink the batch as sequences finish without slicing page_indices.","solutions":["Slice page_indices to the live batch: page_indices = page_indices[:q.shape[0]]","Regenerate page_indices from the scheduler each step alongside q","Keep q, lengths, and page_indices derived from one batch spec object"],"exampleFix":"// before\npaged_attention(q, k, v, page_indices_full, lengths, ...)\n// after\npaged_attention(q, k, v, page_indices_full[:q.shape[0]], lengths, ...)","handlingStrategy":"validation","validationCode":"assert page_indices.shape[0] == q.shape[0], (page_indices.shape[0], q.shape[0])","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Regenerate page_indices from the scheduler each step rather than reusing a fixed table","Keep batch-indexed tensors in one dict so they are always sliced in sync"],"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"}