{"record":{"id":"fb8f68eada0c84c7","repo":"jax-ml/jax","slug":"number-of-kv-heads-must-be-even-when-megacore-mode","errorCode":null,"errorMessage":"number of KV heads must be even when megacore_mode is 'kv_head'","messagePattern":"number of KV heads must be even when megacore_mode is 'kv_head'","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/experimental/pallas/ops/tpu/paged_attention/paged_attention_kernel.py","lineNumber":472,"sourceCode":"    )\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\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\":","sourceCodeStart":454,"sourceCodeEnd":490,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/pallas/ops/tpu/paged_attention/paged_attention_kernel.py#L454-L490","documentation":"With megacore_mode='kv_head' the kernel splits work across the TPU's two Matmul cores by splitting KV heads in half, which requires an even number of KV heads. An odd num_kv_heads cannot be split evenly so the kernel raises this error at configuration time.","triggerScenarios":"Calling paged_attention with megacore_mode='kv_head' and k_pages whose num_kv_heads is odd (e.g. GQA with 3 KV heads, or MQA with 1 KV head).","commonSituations":"Enabling megacore for throughput on models with odd KV head counts; switching a GQA model from 4 to 3 KV heads; using MQA (num_kv_heads=1) with the default kv_head mode.","solutions":["Use megacore_mode=None or 'batch' when num_kv_heads is odd","Choose an even num_kv_heads in the model config (e.g. 2, 4, 8)","If batch is even, megacore_mode='batch' splits across batch instead"],"exampleFix":"// before\npaged_attention(q, k, v, idx, lens, megacore_mode='kv_head')  # num_kv_heads=1\n// after\npaged_attention(q, k, v, idx, lens, megacore_mode='batch')  # or None","handlingStrategy":"validation","validationCode":"if megacore_mode == 'kv_head':\n    assert k_pages.shape[0] % 2 == 0, 'num_kv_heads must be even'","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Choose megacore mode from parity: kv_head if num_kv_heads%2==0, batch if batch%2==0, else None","Prefer even KV head counts in GQA configs"],"tags":["jax","pallas","tpu","paged-attention","megacore","gqa"],"backgroundTag":"invalid-kernel-configuration","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}