jax-ml/jax · error · ValueError

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

Error message

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

What it means

The kernel needs one sequence length per query batch element to mask padded positions. It derives batch_size from q.shape[0] and requires lengths to have exactly that shape; otherwise the masking metadata does not line up with the queries and the kernel raises this error.

Source

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

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

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Slice or pad lengths to exactly q.shape[0]: lengths = lengths[:q.shape[0]]
2. Keep lengths and q produced from the same batch slicing step
3. For ragged/variable batch counts, use the ragged_paged_attention kernel instead

Example fix

// before
out = paged_attention(q, lengths=all_lengths, ...)  # all_lengths longer than batch
// after
out = paged_attention(q, lengths=all_lengths[:q.shape[0]], ...)

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

When it happens

Trigger: Passing lengths of shape (batch*beam,) while q has batch dims after reshaping for beam search, or a scalar/padded lengths array from a dataloader with a different batch size than q.

Common situations: Beam search where q is reshaped but lengths is not; last partial batch in a generation loop where lengths buffer was preallocated to max batch; multi-host sharding where batch sizes diverge.

Related errors

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