{"record":{"id":"ba6987378c873a09","repo":"jax-ml/jax","slug":"the-dtype-of-kv-lens-page-indices-and-cu-q","errorCode":null,"errorMessage":"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=}.","messagePattern":"The dtype of `kv_lens`, `page_indices`, and `cu_q_lens` must be int32\\. Got (.+?), (.+?), (.+?)\\.","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/experimental/pallas/ops/tpu/ragged_paged_attention/kernel.py","lineNumber":264,"sourceCode":"    raise ValueError(\n        f\"Q head_dim {head_dim} must be the same as that of K/V {head_dim_k}.\"\n    )\n  if kv_lens.shape != (max_num_seqs,):\n    raise ValueError(\n        f\"Expected {kv_lens.shape=} to be ({max_num_seqs},) where\"\n        \" `max_num_seqs` is `page_indices.shape[0]`.\"\n    )\n  if cu_q_lens.shape != (max_num_seqs + 1,):\n    raise ValueError(\n        f\"Expected {cu_q_lens.shape=} to be ({max_num_seqs + 1},)  where\"\n        \" `max_num_seqs` is `page_indices.shape[0]`.\"\n    )\n  if (\n      kv_lens.dtype != jnp.int32\n      or page_indices.dtype != jnp.int32\n      or cu_q_lens.dtype != jnp.int32\n  ):\n    raise ValueError(\n        \"The dtype of `kv_lens`, `page_indices`, and `cu_q_lens` must be\"\n        f\" int32. Got {kv_lens.dtype=}, {page_indices.dtype=},\"\n        f\" {cu_q_lens.dtype=}.\"\n    )\n  if num_q_heads % num_kv_heads != 0:\n    raise ValueError(f\"{num_q_heads=} must be divisible by {num_kv_heads=}\")\n  if sliding_window is not None and sliding_window <= 0:\n    raise ValueError(f\"{sliding_window=} must be positive.\")\n  if soft_cap is not None and soft_cap == 0.0:\n    raise ValueError(f\"{soft_cap=} must not be 0.0.\")\n  if (\n      num_kv_pages_per_block is not None\n      and not 0 < num_kv_pages_per_block <= pages_per_seq\n  ):\n    raise ValueError(\n        f\"{num_kv_pages_per_block=} must be in range (0, {pages_per_seq}].\"\n    )\n  if num_queries_per_block is not None and num_queries_per_block <= 0:","sourceCodeStart":246,"sourceCodeEnd":282,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/pallas/ops/tpu/ragged_paged_attention/kernel.py#L246-L282","documentation":"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).","triggerScenarios":"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(...).","commonSituations":"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.","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)"],"exampleFix":"// before\npage_indices = np.arange(num_seqs * pages_per_seq).reshape(num_seqs, -1)\n// after\npage_indices = jnp.asarray(np.arange(num_seqs * pages_per_seq, dtype=np.int32).reshape(num_seqs, -1))","handlingStrategy":"validation","validationCode":"kv_lens = jnp.asarray(kv_lens, jnp.int32)\npage_indices = jnp.asarray(page_indices, jnp.int32)\ncu_q_lens = jnp.asarray(cu_q_lens, jnp.int32)\nassert kv_lens.dtype == page_indices.dtype == cu_q_lens.dtype == jnp.int32","typeGuard":null,"tryCatchPattern":null,"preventionTips":["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"],"tags":["jax","pallas","tpu","dtype-validation","paged-attention"],"backgroundTag":"tensor-dtype-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}