jax-ml/jax · error · ValueError

Number of Q heads must be divisible by number of KV heads. G

Error message

Number of Q heads must be divisible by number of KV heads. Got {num_q_heads} and {num_kv_heads}.

What it means

Paged attention on TPU implements grouped-query attention (GQA) by mapping each KV head to num_q_heads // num_kv_heads query heads. That mapping requires num_q_heads to be an exact multiple of num_kv_heads; otherwise the head-grouping in the kernel is undefined and this ValueError is raised.

Source

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

    v_pages, v_scales_pages = v_pages.weight, v_pages.scales
    assert isinstance(v_scales_pages, jax.Array)  # For typing.
    v_scales_pages = jnp.broadcast_to(
        v_scales_pages, (*v_scales_pages.shape[:-1], v_pages.shape[-1])
    )
  else:
    v_scales_pages = None

  batch_size, num_q_heads, head_dim = q.shape
  num_kv_heads, _, page_size, head_dim_k = k_pages.shape
  batch_size_paged_indices, pages_per_sequence = page_indices.shape

  if k_pages.shape != v_pages.shape:
    raise ValueError(
        f"k_pages and v_pages must have the same shape. Got {k_pages.shape} and"
        f" {v_pages.shape}"
    )
  if num_q_heads % num_kv_heads != 0:
    raise ValueError(
        "Number of Q heads must be divisible by number of KV heads. Got"
        f" {num_q_heads} and {num_kv_heads}."
    )
  if head_dim_k != head_dim:
    raise ValueError(
        "head_dim of Q must be the same as that of K/V. Got"
        f" {head_dim} and {head_dim_k}."
    )
  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:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Set num_kv_heads to a divisor of num_q_heads (e.g. 1 for MQA, num_q_heads for MHA, or num_q_heads/rep factor)
Fix the query projection so q reshapes to (batch, num_kv_heads * group, dim)
Re-check model config: num_attention_heads and num_key_value_heads must satisfy divisibility

Example fix

// before
q shape (b, 8, d); k_pages with num_kv_heads=6
// after
q shape (b, 8, d); k_pages with num_kv_heads=2  # 8 % 2 == 0

Defensive patterns

Strategy: validation

Validate before calling

num_q_heads, num_kv_heads = q.shape[1], k_pages.shape[0]
assert num_q_heads % num_kv_heads == 0, (num_q_heads, num_kv_heads)

Prevention

Validate head counts once at model-build time
Encode num_kv_heads = num_q_heads // group_size in config instead of an independent value

When it happens

Trigger: Calling paged_attention with q of shape (batch, 8, dim) and k_pages with num_kv_heads=3, or any combination where num_q_heads % num_kv_heads != 0 (e.g. 8 Q heads with 6 KV heads).

Common situations: Switching a model from MHA to GQA/MQA with incompatible head counts; loading checkpoints where the KV head replication factor is fractional; typoing head counts in config.

Related errors

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