jax-ml/jax · error · ValueError

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

Error message

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

What it means

Every query head is dotted against K/V with the same head_dim, so q.shape[-1] must equal the head_dim baked into k_pages/v_pages. When the last dimensions differ, the dot product inside the kernel is shape-incompatible and the kernel raises this error before launch.

Source

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

  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:
    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.

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Make q's last dim match the cache: project q to head_dim_k, or rebuild the cache with head_dim == q.shape[-1]
Clear and re-allocate the KV cache whenever the model's head_dim changes
Assert equality before inference: assert q.shape[-1] == k_pages.shape[-1]

Example fix

// before
q = proj(x)              # head_dim 128
cache head_dim = 64
// after
q = proj(x) @ w_rescale   # or rebuild cache with head_dim 128

Defensive patterns

Strategy: validation

Validate before calling

assert q.shape[-1] == k_pages.shape[-1] == v_pages.shape[-1]

Prevention

Derive cache allocation from the model's head_dim config
Rebuild caches whenever head_dim-affecting config changes

When it happens

Trigger: Calling paged_attention with q head_dim 128 but a KV cache allocated with page tensors of head_dim 64 (or 256); mixing a model checkpoint's head_dim with a cache built for another model.

Common situations: Swapping LoRA adapters or models that share tokenizer but not head dims; reusing a persisted KV cache from a different architecture; config drift between rope_theta/head_dim settings.

Related errors

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