jax-ml/jax · error · ValueError

megacore_mode must be one of ['kv_head', 'batch', None]

Error message

megacore_mode must be one of ['kv_head', 'batch', None]

What it means

megacore_mode selects how work is partitioned across the TPU Matmul cores and only accepts the literal strings 'kv_head', 'batch', or None. Any other value (including case variants or typo'd strings) falls through to the final else branch and raises this error.

Source

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

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

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Use exactly one of 'kv_head', 'batch', or None (Python None, not 'none')
2. Normalize config input: mode = None if mode in ('', 'none', 'None') else mode
3. Add a unit test asserting megacore_mode membership in the allowed set

Example fix

// before
paged_attention(..., megacore_mode='none')
// after
paged_attention(..., megacore_mode=None)

ALLOWED = ('kv_head', 'batch', None)
megacore_mode = None if megacore_mode in ('', 'none', 'None') else megacore_mode
assert megacore_mode in ALLOWED

def is_valid_megacore_mode(m): return m in ('kv_head', 'batch', None)

When it happens

Trigger: Passing megacore_mode='KV_HEAD', 'batch_head', 'none' (string), 'kv', or an unset config variable defaulting to '' instead of None.

Common situations: Config-driven megacore selection from YAML/env vars where None is serialized as the string 'none'; typos in copied sample code; API changes from older versions that accepted other modes.

Understand the failure class

Background: "Must be a positive integer", "Invalid value", "Unsupported": the invalid-argument-value error family, when a library rejects the value you pass — this error's family across 35 libraries.

Related errors

• number of KV heads must be even when megacore_mode is 'kv_he
• batch size must be even when megacore_mode is 'batch'
• k_pages and v_pages must have the same shape. Got {k_pages.s
• Number of Q heads must be divisible by number of KV heads. G
• head_dim of Q must be the same as that of K/V. Got {head_dim

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