jax-ml/jax · error · ValueError

Q head_dim {head_dim} must be the same as that of K/V {head_

Error message

Q head_dim {head_dim} must be the same as that of K/V {head_dim_k}.

What it means

Static validation of ragged paged attention checks that q's head_dim equals the head_dim of the combined K/V pages tensor (which packs K and V along the head axis). Since each query is dotted with K/V of the same dimensionality, a mismatch is a hard compile-time error.

Source

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

    mask_value: float | None = None,
    k_scale: float | None = None,
    v_scale: float | None = None,
    # Kernel tuning params.
    num_kv_pages_per_block: int | None = None,
    num_queries_per_block: int | None = None,
    vmem_limit_bytes: int | None = None,
):
  _, num_q_heads, head_dim = q.shape
  _, _, num_combined_kv_heads, head_dim_k = kv_pages.shape
  assert num_combined_kv_heads % 2 == 0
  assert isinstance(k_scale, float) or k_scale is None
  assert isinstance(v_scale, float) or v_scale is None
  num_kv_heads = num_combined_kv_heads // 2
  max_num_seqs, pages_per_seq = page_indices.shape
  if num_seqs.shape != (1,):
    raise ValueError(f"{num_seqs.shape=} must be (1,)")
  if head_dim_k != head_dim:
    raise ValueError(
        f"Q head_dim {head_dim} must be the same as that of K/V {head_dim_k}."
    )
  if kv_lens.shape != (max_num_seqs,):
    raise ValueError(
        f"Expected {kv_lens.shape=} to be ({max_num_seqs},) where"
        " `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(

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Rebuild kv_pages with head_dim equal to q.shape[-1]
2. Project q to the cache's head_dim before calling
3. Always allocate the cache from the model config's head_dim, not a hardcoded value

Example fix

// before
kv_pages = stack_cache(k_heads_64, v_heads_64)  # q head_dim 128
// after
kv_pages = stack_cache(k_heads_128, v_heads_128)  # matches q

head_dim_k = kv_pages.shape[-1]  # combined K/V pages tensor
assert q.shape[-1] == head_dim_k

When it happens

Trigger: Calling ragged_paged_attention with q.shape[-1]=128 but kv_pages built with head_dim=64; loading a KV cache saved from a model with a different head dimension; mis-stacking K and V so the combined tensor's last dim no longer matches.

Common situations: Model swaps reusing a persisted paged KV cache; GQA refactor changing head_dim; building kv_pages by stacking K/V with an incorrect axis order.

Related errors

• 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
• {num_seqs[0]=} must be less or equal to {max_num_seqs=}
• {pages_per_seq=} must be greater or equal to {min_pages_per_
• Total q tokens {cu_q_lens[num_seqs[0]]} must be less or equa

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