{"record":{"id":"8b3a47bb57b76889","repo":"jax-ml/jax","slug":"vmem-limit-bytes-must-be-positive","errorCode":null,"errorMessage":"{vmem_limit_bytes=} must be positive.","messagePattern":"(.+?) must be positive\\.","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/experimental/pallas/ops/tpu/ragged_paged_attention/kernel.py","lineNumber":285,"sourceCode":"        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:\n    raise ValueError(f\"{num_queries_per_block=} must be positive.\")\n  if vmem_limit_bytes is not None and vmem_limit_bytes <= 0:\n    raise ValueError(f\"{vmem_limit_bytes=} must be positive.\")\n  del sm_scale  # No constraints on sm_scale.\n  del mask_value  # No consstraints on mask_value.\n\n\ndef ragged_paged_attention_kernel(\n    # Prefetch\n    kv_lens_ref,  # [max_num_seqs]\n    page_indices_ref,  # [max_num_seqs, pages_per_seq]\n    cu_q_lens_ref,  # [max_num_seqs + 1]\n    seq_buf_idx_ref,\n    # TODO(jevinjiang): if OOM in SMEM, consider pack to other scalar refs.\n    num_seqs_ref,\n    # Input\n    q_ref,  # [num_q_per_blk, num_q_heads_per_blk, head_dim]\n    kv_pages_hbm_ref,  # [total_num_pages, page_size, num_combined_kv_heads, head_dim]\n    # Output\n    o_ref,  # [num_q_per_blk, num_q_heads_per_blk, head_dim]\n    # Scratch","sourceCodeStart":267,"sourceCodeEnd":303,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/pallas/ops/tpu/ragged_paged_attention/kernel.py#L267-L303","documentation":"vmem_limit_bytes optionally caps TPU vector memory (VMEM) usage for the ragged paged attention kernel and must be positive. Zero or negative byte limits are meaningless and rejected by static_validate_inputs.","triggerScenarios":"Passing vmem_limit_bytes=0 or a negative value, often as a sentinel meaning 'no limit' or from subtractive arithmetic computing a budget.","commonSituations":"Memory-budget heuristics that compute vmem_limit_bytes = total - reserved and underflow to <= 0 when the reservation exceeds the total; porting configs where 0 meant 'auto'.","solutions":["Pass None instead of 0 to use the default VMEM limit","Clamp computed budgets: vmem_limit_bytes = budget if budget > 0 else None","Re-check the VMEM size constant you subtract from (per TPU generation it differs)"],"exampleFix":"// before\nattn(..., vmem_limit_bytes=vmem_total - reserved)  # <= 0\n// after\nbudget = vmem_total - reserved\nattn(..., vmem_limit_bytes=budget if budget > 0 else None)","handlingStrategy":"validation","validationCode":"vmem_limit_bytes = vmem_limit_bytes if (vmem_limit_bytes is not None and vmem_limit_bytes > 0) else None","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Use None for the default VMEM limit","Sanity-check computed memory budgets for underflow before passing them"],"tags":["jax","pallas","tpu","memory-limit","validation"],"backgroundTag":"invalid-argument-value","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}