{"record":{"id":"c60e4296e9e9e62b","repo":"jax-ml/jax","slug":"expected-kv-lens-shape-to-be-max-num-seqs","errorCode":null,"errorMessage":"Expected {kv_lens.shape=} to be ({max_num_seqs},) where `max_num_seqs` is `page_indices.shape[0]`.","messagePattern":"Expected (.+?) to be \\((.+?),\\) where `max_num_seqs` is `page_indices\\.shape\\[0\\]`\\.","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/experimental/pallas/ops/tpu/ragged_paged_attention/kernel.py","lineNumber":250,"sourceCode":"    num_kv_pages_per_block: int | None = None,\n    num_queries_per_block: int | None = None,\n    vmem_limit_bytes: int | None = None,\n):\n  _, num_q_heads, head_dim = q.shape\n  _, _, num_combined_kv_heads, head_dim_k = kv_pages.shape\n  assert num_combined_kv_heads % 2 == 0\n  assert isinstance(k_scale, float) or k_scale is None\n  assert isinstance(v_scale, float) or v_scale is None\n  num_kv_heads = num_combined_kv_heads // 2\n  max_num_seqs, pages_per_seq = page_indices.shape\n  if num_seqs.shape != (1,):\n    raise ValueError(f\"{num_seqs.shape=} must be (1,)\")\n  if head_dim_k != head_dim:\n    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    )","sourceCodeStart":232,"sourceCodeEnd":268,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/pallas/ops/tpu/ragged_paged_attention/kernel.py#L232-L268","documentation":"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.","triggerScenarios":"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,)).","commonSituations":"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.","solutions":["Make kv_lens exactly 1-D with length equal to page_indices.shape[0]","Ensure both tensors are trimmed to the same number of sequences (e.g. kv_lens = kv_lens[:page_indices.shape[0]])","Check you did not accidentally pass q_lens or cumulative lens instead of per-sequence kv lengths"],"exampleFix":"// before\nout = ragged_paged_attention(q, k, v, page_indices, kv_lens, cu_q_lens)  # kv_lens.shape == (17,), page_indices.shape == (16, ...)\n// after\nkv_lens = kv_lens[:page_indices.shape[0]]\nout = ragged_paged_attention(q, k, v, page_indices, kv_lens, cu_q_lens)","handlingStrategy":"validation","validationCode":"assert page_indices.ndim == 2\nmax_num_seqs = page_indices.shape[0]\nassert kv_lens.shape == (max_num_seqs,), (kv_lens.shape, max_num_seqs)","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Build kv_lens and page_indices from the same num_seqs variable","Assert metadata shapes against page_indices.shape[0] before calling the kernel"],"tags":["jax","pallas","tpu","paged-attention","shape-validation"],"backgroundTag":"tensor-shape-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}