jax-ml/jax · error · ValueError

{num_q_heads=} must be divisible by {num_kv_heads=}

Error message

{num_q_heads=} must be divisible by {num_kv_heads=}

What it means

This kernel only implements grouped-query attention (GQA) / multi-query attention where the number of query heads is an integer multiple of the number of KV heads. static_validate_inputs enforces num_q_heads % num_kv_heads == 0 so head-to-head mapping inside the kernel is well defined.

Source

Thrown at jax/experimental/pallas/ops/tpu/ragged_paged_attention/kernel.py:270

        " `max_num_seqs` is `page_indices.shape[0]`."
    )
  if cu_q_lens.shape != (max_num_seqs + 1,):
    raise ValueError(
        f"Expected {cu_q_lens.shape=} to be ({max_num_seqs + 1},)  where"
        " `max_num_seqs` is `page_indices.shape[0]`."
    )
  if (
      kv_lens.dtype != jnp.int32
      or page_indices.dtype != jnp.int32
      or cu_q_lens.dtype != jnp.int32
  ):
    raise ValueError(
        "The dtype of `kv_lens`, `page_indices`, and `cu_q_lens` must be"
        f" int32. Got {kv_lens.dtype=}, {page_indices.dtype=},"
        f" {cu_q_lens.dtype=}."
    )
  if num_q_heads % num_kv_heads != 0:
    raise ValueError(f"{num_q_heads=} must be divisible by {num_kv_heads=}")
  if sliding_window is not None and sliding_window <= 0:
    raise ValueError(f"{sliding_window=} must be positive.")
  if soft_cap is not None and soft_cap == 0.0:
    raise ValueError(f"{soft_cap=} must not be 0.0.")
  if (
      num_kv_pages_per_block is not None
      and not 0 < num_kv_pages_per_block <= pages_per_seq
  ):
    raise ValueError(
        f"{num_kv_pages_per_block=} must be in range (0, {pages_per_seq}]."
    )
  if num_queries_per_block is not None and num_queries_per_block <= 0:
    raise ValueError(f"{num_queries_per_block=} must be positive.")
  if vmem_limit_bytes is not None and vmem_limit_bytes <= 0:
    raise ValueError(f"{vmem_limit_bytes=} must be positive.")
  del sm_scale  # No constraints on sm_scale.
  del mask_value  # No consstraints on mask_value.

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Fix the head configuration so num_q_heads is a multiple of num_kv_heads (e.g. 32 q heads with 8 kv heads)
2. Repeat/pad KV heads to a divisor of num_q_heads if the model truly has an odd ratio (interleaved repeat like jnp.repeat(kv, ratio, axis=1))
3. Double-check that q,k,v were not transposed so the head axis is actually axis 1

Example fix

// before
q: (seq, 12, d); k/v: (seq, 8, d)  # 12 % 8 != 0
// after
k = jnp.repeat(k, 12 // 8 if 12 % 8 == 0 else 1, axis=1)  # better: choose configs like 12 q / 6 kv or 16 q / 8 kv
# preferred: use num_q_heads=16, num_kv_heads=8

n_q, n_kv = q.shape[1], k.shape[1]
assert n_q % n_kv == 0, f'{n_q=} not divisible by {n_kv=}'

When it happens

Trigger: Passing q with a head count not divisible by k/v's head count, e.g. 12 query heads with 8 KV heads.

Common situations: Porting a model config to the TPU ragged attention path where head ratios like 8:1 or 7:1 (GQA) are expected; typos in num_heads vs num_kv_heads config; using MHA weights with a partial KV head set.

Related errors

• Number of Q heads must be divisible by number of KV heads. G
• number of KV heads must be even when megacore_mode is 'kv_he
• {soft_cap=} must not be 0.0.
• Not implemented: {num_combined_kv_heads=} can not be XLA ful
• In MHA, expected number of 'key' heads ({num_kv_heads}) to b

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