{"record":{"id":"cd2a00ea9bb93560","repo":"jax-ml/jax","slug":"expected-cu-q-lens-shape-to-be-max-num-seqs","errorCode":null,"errorMessage":"Expected {cu_q_lens.shape=} to be ({max_num_seqs + 1},)  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":255,"sourceCode":"  _, _, 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    )\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:","sourceCodeStart":237,"sourceCodeEnd":273,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/pallas/ops/tpu/ragged_paged_attention/kernel.py#L237-L273","documentation":"static_validate_inputs requires cu_q_lens (cumulative query lengths, the ragged-layout offset array) to have exactly max_num_seqs + 1 entries, where max_num_seqs = page_indices.shape[0]. It is the standard prefix-sum layout: entry i is the start of sequence i and the last entry is the total query count.","triggerScenarios":"Passing a cu_q_lens that is too short/long for the batch, or forgetting the terminal total element (passing per-sequence q_lens instead of cumulative sums), when calling ragged_paged_attention.","commonSituations":"Converting variable-length batched queries into ragged format and forgetting the exclusive-scan final element, or reusing kv_lens as cu_q_lens.","solutions":["Compute cu_q_lens = jnp.concatenate([[0], jnp.cumsum(q_lens)]) so its length is num_seqs + 1","Verify cu_q_lens.shape == (page_indices.shape[0] + 1,) before the call","Do not pass per-sequence lengths where cumulative offsets are expected"],"exampleFix":"// before\ncu_q_lens = q_lens  # wrong: per-seq lengths\n// after\ncu_q_lens = jnp.concatenate([jnp.zeros(1, jnp.int32), jnp.cumsum(q_lens, dtype=jnp.int32)])","handlingStrategy":"validation","validationCode":"cu_q_lens = jnp.concatenate([jnp.zeros(1, jnp.int32), jnp.cumsum(q_lens, dtype=jnp.int32)])\nassert cu_q_lens.shape == (page_indices.shape[0] + 1,)","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Always derive cu_q_lens via concatenate([0], cumsum(q_lens)) instead of hand-building","Check length == num_seqs + 1 before the call"],"tags":["jax","pallas","tpu","ragged-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"}