jax-ml/jax · error · ValueError

batch size must be even when megacore_mode is 'batch'

Error message

batch size must be even when megacore_mode is 'batch'

What it means

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.

Source

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

  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":
      q_block_spec = pl.BlockSpec(
          (None, num_groups, None, head_dim),
          lambda core_index, b, h, *_: (b, h * num_cores + core_index, 0, 0),
      )
    elif megacore_mode == "batch":
      q_block_spec = pl.BlockSpec(

View on GitHub (pinned to 1e1c6a8fc0)

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

Example fix

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

Defensive patterns

Strategy: validation

Validate before calling

if megacore_mode == 'batch':
    assert q.shape[0] % 2 == 0, 'batch must be even'

Prevention

Switch to megacore_mode=None for batch=1 decode
Pad odd batches with zero-length dummy sequences when throughput matters

When it happens

Trigger: Calling paged_attention with megacore_mode='batch' and q.shape[0] odd (e.g. batch=1 decode step, or 3 sequences in continuous batching).

Common situations: Autoregressive decode with batch=1 while batch megacore is enabled; last ragged batch of a generation loop; benchmarking with odd batch sizes.

Related errors

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