jax-ml/jax · error · ValueError

pages_per_compute_block must be divisible by pages per seque

Error message

pages_per_compute_block must be divisible by pages per sequence. Got {pages_per_compute_block} and {pages_per_sequence}.

What it means

The paged-attention kernel processes KV pages in fixed blocks of pages_per_compute_block; pages_per_sequence (from page_indices.shape[1]) must be divisible by it so each compute block covers whole pages across all sequences. Note the check is written as pages_per_sequence % pages_per_compute_block, so in practice pages_per_compute_block must be 1 or a divisor of pages_per_seq.

Source

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

  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.
  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'"
      )

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Use pages_per_compute_block=1 (always valid)
Make pages_per_sequence a multiple of pages_per_compute_block (e.g. allocate 8 pages/seq for block=4)
Recompute page_indices padding so shape[1] divides evenly

Example fix

// before
paged_attention(q, k, v, page_indices, lens, pages_per_compute_block=4)  # pages_per_seq=6
// after
paged_attention(q, k, v, page_indices, lens, pages_per_compute_block=2)  # 6 % 2 == 0

Defensive patterns

Strategy: validation

Validate before calling

pages_per_seq = page_indices.shape[1]
assert pages_per_seq % pages_per_compute_block == 0, (pages_per_seq, pages_per_compute_block)

Prevention

Default to pages_per_compute_block=1; only raise it after fixing cache geometry
Keep pages_per_seq a power-of-two multiple of the compute block

When it happens

Trigger: Calling paged_attention with pages_per_compute_block=2 while page_indices has an odd pages_per_sequence (e.g. 7); passing a larger compute block than the per-sequence page budget.

Common situations: Tuning pages_per_compute_block for performance without checking cache geometry; variable sequence budgets where pages_per_seq is not a multiple of the block size.

Related errors

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