jax-ml/jax · error · ValueError

`page_indices` and `q` must have the same batch size

Error message

`page_indices` and `q` must have the same batch size

What it means

page_indices maps each sequence in the batch to its KV pages, so its leading dimension must equal q's batch dimension. If page_indices.shape[0] != q.shape[0] the kernel cannot associate queries with their page tables and raises immediately.

Source

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

  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.
  if megacore_mode == "kv_head":
    if num_kv_heads % 2 != 0:
      raise ValueError(
          "number of KV heads must be even when megacore_mode is 'kv_head'"
      )
    num_cores = 2
  elif megacore_mode == "batch":
    if batch_size % 2 != 0:
      raise ValueError("batch size must be even when megacore_mode is 'batch'")
    num_cores = 2
  elif megacore_mode is None:
    num_cores = 1

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Slice page_indices to the live batch: page_indices = page_indices[:q.shape[0]]
2. Regenerate page_indices from the scheduler each step alongside q
3. Keep q, lengths, and page_indices derived from one batch spec object

Example fix

// before
paged_attention(q, k, v, page_indices_full, lengths, ...)
// after
paged_attention(q, k, v, page_indices_full[:q.shape[0]], lengths, ...)

assert page_indices.shape[0] == q.shape[0], (page_indices.shape[0], q.shape[0])

When it happens

Trigger: Passing a page_indices table sized for a different batch (e.g. max_num_seqs from vLLM-style scheduling) while q contains only the currently-scheduled sequences.

Common situations: Porting vLLM-style continuous batching where page_indices covers capacity, not the current batch; multi-step decode loops that shrink the batch as sequences finish without slicing page_indices.

Related errors

• k_pages and v_pages must have the same shape. Got {k_pages.s
• `lengths` 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/770c49b21ed96738. Report an issue: GitHub.