jax-ml/jax · error · ValueError

k_pages and v_pages must have the same shape. Got {k_pages.s

Error message

k_pages and v_pages must have the same shape. Got {k_pages.shape} and {v_pages.shape}

What it means

The TPU paged-attention kernel requires K and V page tables to be identical in shape (num_kv_heads, num_pages, page_size, head_dim) because they are indexed in lockstep. If k_pages.shape != v_pages.shape the kernel raises immediately, since page-based lookup would read out of bounds or mix layouts.

Source

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

        k_scales_pages, (*k_scales_pages.shape[:-1], k_pages.shape[-1])
    )
  else:
    k_scales_pages = None
  if isinstance(v_pages, quantization_utils.QuantizedTensor):
    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}."
    )

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Allocate K and V caches from the same shape spec and never resize them independently
2. Re-pad or slice the mismatched tensor to match, e.g. v_pages = v_pages[:, :k_pages.shape[1]]
3. Verify shapes right before the call: assert k_pages.shape == v_pages.shape

Example fix

// before
out = paged_attention(q, k_pages, v_pages[:,:,:-1], ...)
// after
out = paged_attention(q, k_pages, v_pages, ...)  # keep caches identical

assert k_pages.shape == v_pages.shape, (k_pages.shape, v_pages.shape)

When it happens

Trigger: Calling paged_attention(q, k_pages, v_pages, ...) where k_pages and v_pages come from different KV-cache allocations, were sliced differently, or one was transposed/padded relative to the other.

Common situations: Building a KV cache where V was allocated with a different page count; copying from a HF-style cache that stores K and V separately with padding applied to only one; batched decode loops that trim pages on only one tensor.

Related errors

• `lengths` and `q` must have the same batch size
• `page_indices` and `q` must have the same batch size
• Expected {kv_lens.shape=} to be ({max_num_seqs},) where `max
• Seed key_data must be 1D.
• Leading dimension of seed key_data must be 1.

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