{"record":{"id":"2aaf2f83b589363b","repo":"jax-ml/jax","slug":"lengths-and-q-must-have-the-same-batch-size-2aaf2f","errorCode":null,"errorMessage":"`lengths` and `q` must have the same batch size","messagePattern":"`lengths` 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":461,"sourceCode":"        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      )\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","sourceCodeStart":443,"sourceCodeEnd":479,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/pallas/ops/tpu/paged_attention/paged_attention_kernel.py#L443-L479","documentation":"The kernel needs one sequence length per query batch element to mask padded positions. It derives batch_size from q.shape[0] and requires lengths to have exactly that shape; otherwise the masking metadata does not line up with the queries and the kernel raises this error.","triggerScenarios":"Passing lengths of shape (batch*beam,) while q has batch dims after reshaping for beam search, or a scalar/padded lengths array from a dataloader with a different batch size than q.","commonSituations":"Beam search where q is reshaped but lengths is not; last partial batch in a generation loop where lengths buffer was preallocated to max batch; multi-host sharding where batch sizes diverge.","solutions":["Slice or pad lengths to exactly q.shape[0]: lengths = lengths[:q.shape[0]]","Keep lengths and q produced from the same batch slicing step","For ragged/variable batch counts, use the ragged_paged_attention kernel instead"],"exampleFix":"// before\nout = paged_attention(q, lengths=all_lengths, ...)  # all_lengths longer than batch\n// after\nout = paged_attention(q, lengths=all_lengths[:q.shape[0]], ...)","handlingStrategy":"validation","validationCode":"assert lengths.shape == (q.shape[0],), (lengths.shape, q.shape)","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Generate q, lengths, page_indices from a single batch-prep function","Slice all batch arrays together at the end of the generation loop"],"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"}