{"record":{"id":"0cb44620c8775179","repo":"jax-ml/jax","slug":"total-q-tokens-cu-q-lens-num-seqs-0-must-be-le","errorCode":null,"errorMessage":"Total q tokens {cu_q_lens[num_seqs[0]]} must be less or equal to {max_num_batched_tokens=}.","messagePattern":"Total q tokens (.+?) must be less or equal to (.+?)\\.","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/experimental/pallas/ops/tpu/ragged_paged_attention/kernel.py","lineNumber":202,"sourceCode":"      v_scale=v_scale,\n      num_kv_pages_per_block=num_kv_pages_per_block,\n      num_queries_per_block=num_queries_per_block,\n      vmem_limit_bytes=vmem_limit_bytes,\n  )\n  max_num_batched_tokens = q.shape[0]\n  page_size = kv_pages.shape[1]\n  max_num_seqs, pages_per_seq = page_indices.shape\n  if num_seqs[0] > max_num_seqs:\n    raise ValueError(f\"{num_seqs[0]=} must be less or equal to {max_num_seqs=}\")\n  max_kv_len = jnp.max(kv_lens)\n  min_pages_per_seq = pl.cdiv(max_kv_len, page_size)\n  if pages_per_seq < min_pages_per_seq:\n    raise ValueError(\n        f\"{pages_per_seq=} must be greater or equal to\"\n        f\" {min_pages_per_seq=} given {max_kv_len=} and {page_size=}.\"\n    )\n  if cu_q_lens[num_seqs[0]] > max_num_batched_tokens:\n    raise ValueError(\n        f\"Total q tokens {cu_q_lens[num_seqs[0]]} must be less or equal to\"\n        f\" {max_num_batched_tokens=}.\"\n    )\n  for i in range(num_seqs[0]):\n    q_len = cu_q_lens[i + 1] - cu_q_lens[i]\n    kv_len = kv_lens[i]\n    if q_len > kv_len:\n      raise ValueError(\n          f\"{q_len=} must be less or equal to {kv_len=} at sequence {i}.\"\n      )\n\n\n# Expect to run these checks during compile time.\ndef static_validate_inputs(\n    q: jax.Array,  # [max_num_batched_tokens, num_q_heads, head_dim]\n    kv_pages: jax.Array,  # [total_num_pages, page_size, num_combined_kv_heads, head_dim]\n    kv_lens: jax.Array,  # i32[max_num_seqs]\n    page_indices: jax.Array,  # i32[max_num_seqs, pages_per_seq]","sourceCodeStart":184,"sourceCodeEnd":220,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/pallas/ops/tpu/ragged_paged_attention/kernel.py#L184-L220","documentation":"cu_q_lens is the cumulative query-token count (like cu_seqlens in flash-attention) and its final relevant entry cu_q_lens[num_seqs[0]] gives total queries. q is padded to max_num_batched_tokens rows, so the total must fit within q.shape[0]; otherwise the kernel would gather query rows beyond the buffer.","triggerScenarios":"Calling ragged_paged_attention where q was sliced to fewer rows than sum of per-sequence query lengths, or cu_q_lens was built against a different token budget than the actual q allocation.","commonSituations":"Chunked prefill schedulers that batch tokens up to a budget but build cu_q_lens against a larger one; off-by-one in cumulative sums; padding q to the wrong axis length.","solutions":["Allocate q with q.shape[0] >= int(cu_q_lens[num_seqs]) and pad remaining rows with zeros","Rebuild cu_q_lens from the actual packed batch: cu = jnp.concatenate([jnp.array([0]), jnp.cumsum(q_lens)])","Validate cu_q_lens[-1 relevant] <= q.shape[0] before launch"],"exampleFix":"// before\nq = q_tokens[:1024]  # but cu_q_lens[num_seqs] == 1200\n// after\nq = jnp.pad(q_tokens[:1200], ((0, 88), (0,0), (0,0)))  # fit budget","handlingStrategy":"validation","validationCode":"total_q = int(cu_q_lens[int(num_seqs[0])])\nassert total_q <= q.shape[0], (total_q, q.shape[0])","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Build cu_q_lens with jnp.cumsum from the actual packed lengths","Pad q to the token budget and mask with lengths=0 rows"],"tags":["jax","pallas","tpu","ragged-attention","cumulative-length-validation"],"backgroundTag":"shape-validation-failed","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}