jax-ml/jax · error · ValueError

Expected {kv_lens.shape=} to be ({max_num_seqs},) where `max

Error message

Expected {kv_lens.shape=} to be ({max_num_seqs},) where `max_num_seqs` is `page_indices.shape[0]`.

What it means

Raised by static_validate_inputs in JAX's TPU ragged paged attention kernel when the kv_lens tensor's shape does not match the number of sequences implied by page_indices. kv_lens must be a 1-D int32 array of length max_num_seqs = page_indices.shape[0]. This is an input-contract violation caught eagerly before the Pallas kernel launches.

Source

Thrown at jax/experimental/pallas/ops/tpu/ragged_paged_attention/kernel.py:250

    num_kv_pages_per_block: int | None = None,
    num_queries_per_block: int | None = None,
    vmem_limit_bytes: int | None = None,
):
  _, num_q_heads, head_dim = q.shape
  _, _, num_combined_kv_heads, head_dim_k = kv_pages.shape
  assert num_combined_kv_heads % 2 == 0
  assert isinstance(k_scale, float) or k_scale is None
  assert isinstance(v_scale, float) or v_scale is None
  num_kv_heads = num_combined_kv_heads // 2
  max_num_seqs, pages_per_seq = page_indices.shape
  if num_seqs.shape != (1,):
    raise ValueError(f"{num_seqs.shape=} must be (1,)")
  if head_dim_k != head_dim:
    raise ValueError(
        f"Q head_dim {head_dim} must be the same as that of K/V {head_dim_k}."
    )
  if kv_lens.shape != (max_num_seqs,):
    raise ValueError(
        f"Expected {kv_lens.shape=} to be ({max_num_seqs},) where"
        " `max_num_seqs` is `page_indices.shape[0]`."
    )
  if cu_q_lens.shape != (max_num_seqs + 1,):
    raise ValueError(
        f"Expected {cu_q_lens.shape=} to be ({max_num_seqs + 1},)  where"
        " `max_num_seqs` is `page_indices.shape[0]`."
    )
  if (
      kv_lens.dtype != jnp.int32
      or page_indices.dtype != jnp.int32
      or cu_q_lens.dtype != jnp.int32
  ):
    raise ValueError(
        "The dtype of `kv_lens`, `page_indices`, and `cu_q_lens` must be"
        f" int32. Got {kv_lens.dtype=}, {page_indices.dtype=},"
        f" {cu_q_lens.dtype=}."
    )

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Make kv_lens exactly 1-D with length equal to page_indices.shape[0]
2. Ensure both tensors are trimmed to the same number of sequences (e.g. kv_lens = kv_lens[:page_indices.shape[0]])
3. Check you did not accidentally pass q_lens or cumulative lens instead of per-sequence kv lengths

Example fix

// before
out = ragged_paged_attention(q, k, v, page_indices, kv_lens, cu_q_lens)  # kv_lens.shape == (17,), page_indices.shape == (16, ...)
// after
kv_lens = kv_lens[:page_indices.shape[0]]
out = ragged_paged_attention(q, k, v, page_indices, kv_lens, cu_q_lens)

assert page_indices.ndim == 2
max_num_seqs = page_indices.shape[0]
assert kv_lens.shape == (max_num_seqs,), (kv_lens.shape, max_num_seqs)

When it happens

Trigger: Calling ragged_paged_attention (or ref_ragged_paged_attention / dynamic_validate_inputs) with a kv_lens array whose length differs from page_indices.shape[0], or with extra/missing dimensions (e.g. shape (1, N) instead of (N,)).

Common situations: Building the paged-attention metadata by hand from a serving stack where the batch size was padded differently from the page table, or after slicing/trimming kv_lens for leftover sequences without also 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
• `page_indices` and `q` must have the same batch size
• 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/c60e4296e9e9e62b. Report an issue: GitHub.