{"record":{"id":"6c2e599c2a706e4f","repo":"jax-ml/jax","slug":"batch-size-must-be-even-when-megacore-mode-is-bat","errorCode":null,"errorMessage":"batch size must be even when megacore_mode is 'batch'","messagePattern":"batch size must be even when megacore_mode is 'batch'","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/experimental/pallas/ops/tpu/paged_attention/paged_attention_kernel.py","lineNumber":478,"sourceCode":"  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\n  else:\n    raise ValueError(\"megacore_mode must be one of ['kv_head', 'batch', None]\")\n\n  num_groups = num_q_heads // num_kv_heads\n  if (num_groups) % 8 != 0:\n    # Reshape q to hint XLA to pick a <1x128> layout otherwise it will pick a\n    # <8x128> layout for a <1x128> memref inside the kernel and error out.\n    q = q.reshape(batch_size, num_q_heads, 1, head_dim)\n    if megacore_mode == \"kv_head\":\n      q_block_spec = pl.BlockSpec(\n          (None, num_groups, None, head_dim),\n          lambda core_index, b, h, *_: (b, h * num_cores + core_index, 0, 0),\n      )\n    elif megacore_mode == \"batch\":\n      q_block_spec = pl.BlockSpec(","sourceCodeStart":460,"sourceCodeEnd":496,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/pallas/ops/tpu/paged_attention/paged_attention_kernel.py#L460-L496","documentation":"With megacore_mode='batch' the kernel splits the batch dimension across the two Matmul cores, requiring an even batch size. An odd batch leaves one core idle/broken so the kernel rejects the configuration up front.","triggerScenarios":"Calling paged_attention with megacore_mode='batch' and q.shape[0] odd (e.g. batch=1 decode step, or 3 sequences in continuous batching).","commonSituations":"Autoregressive decode with batch=1 while batch megacore is enabled; last ragged batch of a generation loop; benchmarking with odd batch sizes.","solutions":["Use megacore_mode=None or 'kv_head' (if KV heads are even) for odd batches","Pad the batch to an even size with dummy sequences (mask via lengths=0)","Pick megacore mode dynamically based on parity of batch and num_kv_heads"],"exampleFix":"// before\npaged_attention(q, k, v, idx, lens, megacore_mode='batch')  # batch=1\n// after\npaged_attention(q, k, v, idx, lens, megacore_mode=None)  # batch=1","handlingStrategy":"validation","validationCode":"if megacore_mode == 'batch':\n    assert q.shape[0] % 2 == 0, 'batch must be even'","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Switch to megacore_mode=None for batch=1 decode","Pad odd batches with zero-length dummy sequences when throughput matters"],"tags":["jax","pallas","tpu","paged-attention","megacore","batch-size"],"backgroundTag":"invalid-kernel-configuration","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}