jax-ml/jax · error · ValueError

The dtype of `kv_lens`, `page_indices`, and `cu_q_lens` must

Error message

The dtype of `kv_lens`, `page_indices`, and `cu_q_lens` must be int32. Got {kv_lens.dtype=}, {page_indices.dtype=}, {cu_q_lens.dtype=}.

What it means

The ragged paged attention kernel's metadata tensors (kv_lens, page_indices, cu_q_lens) are fed directly into the Pallas TPU kernel, which only supports int32 scalars/indices. static_validate_inputs rejects any other dtype (e.g. int64 or uint32).

Source

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

    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=}."
    )
  if num_q_heads % num_kv_heads != 0:
    raise ValueError(f"{num_q_heads=} must be divisible by {num_kv_heads=}")
  if sliding_window is not None and sliding_window <= 0:
    raise ValueError(f"{sliding_window=} must be positive.")
  if soft_cap is not None and soft_cap == 0.0:
    raise ValueError(f"{soft_cap=} must not be 0.0.")
  if (
      num_kv_pages_per_block is not None
      and not 0 < num_kv_pages_per_block <= pages_per_seq
  ):
    raise ValueError(
        f"{num_kv_pages_per_block=} must be in range (0, {pages_per_seq}]."
    )
  if num_queries_per_block is not None and num_queries_per_block <= 0:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Cast all three tensors to jnp.int32 before calling the function
Use explicit dtype when constructing: jnp.arange(..., dtype=jnp.int32)
If sourcing from NumPy, use np.asarray(..., dtype=np.int32)

Example fix

// before
page_indices = np.arange(num_seqs * pages_per_seq).reshape(num_seqs, -1)
// after
page_indices = jnp.asarray(np.arange(num_seqs * pages_per_seq, dtype=np.int32).reshape(num_seqs, -1))

Defensive patterns

Strategy: validation

Validate before calling

kv_lens = jnp.asarray(kv_lens, jnp.int32)
page_indices = jnp.asarray(page_indices, jnp.int32)
cu_q_lens = jnp.asarray(cu_q_lens, jnp.int32)
assert kv_lens.dtype == page_indices.dtype == cu_q_lens.dtype == jnp.int32

Prevention

Normalize metadata to int32 at the boundary of your serving/inference code
Avoid passing raw numpy int64 arrays; always jnp.asarray with dtype=jnp.int32

When it happens

Trigger: Calling ragged_paged_attention with metadata produced by default NumPy ops (often int64 on Linux) or by jnp.arange without an explicit dtype, e.g. page_indices = np.arange(...).

Common situations: Building the page table or length arrays with numpy on a platform where the default integer type is int64, or deserializing metadata from a serving framework that stores int64.

Related errors

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