jax-ml/jax · error · ValueError

number of KV heads must be even when megacore_mode is 'kv_he

Error message

number of KV heads must be even when megacore_mode is 'kv_head'

What it means

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.

Source

Thrown at jax/experimental/pallas/ops/tpu/paged_attention/paged_attention_kernel.py:472

    )
  if pages_per_sequence % pages_per_compute_block != 0:
    raise ValueError(
        "pages_per_compute_block must be divisible by pages per sequence. Got"
        f" {pages_per_compute_block} and {pages_per_sequence}."
    )
  if lengths.shape != (batch_size,):
    raise ValueError("`lengths` and `q` must have the same batch size")
  if batch_size_paged_indices != batch_size:
    raise ValueError("`page_indices` and `q` must have the same batch size")
  if lengths.dtype != jnp.int32:
    raise ValueError(
        f"The dtype of `lengths` must be int32. Got {lengths.dtype}"
    )

  # TODO(dinghua): get the actual cores per chip once there's an official API.
  if megacore_mode == "kv_head":
    if num_kv_heads % 2 != 0:
      raise ValueError(
          "number of KV heads must be even when megacore_mode is 'kv_head'"
      )
    num_cores = 2
  elif megacore_mode == "batch":
    if batch_size % 2 != 0:
      raise ValueError("batch size must be even when megacore_mode is 'batch'")
    num_cores = 2
  elif megacore_mode is None:
    num_cores = 1
  else:
    raise ValueError("megacore_mode must be one of ['kv_head', 'batch', None]")

  num_groups = num_q_heads // num_kv_heads
  if (num_groups) % 8 != 0:
    # Reshape q to hint XLA to pick a <1x128> layout otherwise it will pick a
    # <8x128> layout for a <1x128> memref inside the kernel and error out.
    q = q.reshape(batch_size, num_q_heads, 1, head_dim)
    if megacore_mode == "kv_head":

View on GitHub (pinned to 1e1c6a8fc0)

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

Example fix

// before
paged_attention(q, k, v, idx, lens, megacore_mode='kv_head')  # num_kv_heads=1
// after
paged_attention(q, k, v, idx, lens, megacore_mode='batch')  # or None

Defensive patterns

Strategy: validation

Validate before calling

if megacore_mode == 'kv_head':
    assert k_pages.shape[0] % 2 == 0, 'num_kv_heads must be even'

Prevention

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

When it happens

Trigger: 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).

Common situations: 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.

Related errors

AI-assisted analysis of jax-ml/jax@1e1c6a8fc0 (2026-08-27). Data as JSON: /api/errors/fb8f68eada0c84c7. Report an issue: GitHub.