jax-ml/jax · error · ValueError

{vmem_limit_bytes=} must be positive.

Error message

{vmem_limit_bytes=} must be positive.

What it means

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.

Source

Thrown at jax/experimental/pallas/ops/tpu/ragged_paged_attention/kernel.py:285

        f" {cu_q_lens.dtype=}."
    )
  if num_q_heads % num_kv_heads != 0:
    raise ValueError(f"{num_q_heads=} must be divisible by {num_kv_heads=}")
  if sliding_window is not None and sliding_window <= 0:
    raise ValueError(f"{sliding_window=} must be positive.")
  if soft_cap is not None and soft_cap == 0.0:
    raise ValueError(f"{soft_cap=} must not be 0.0.")
  if (
      num_kv_pages_per_block is not None
      and not 0 < num_kv_pages_per_block <= pages_per_seq
  ):
    raise ValueError(
        f"{num_kv_pages_per_block=} must be in range (0, {pages_per_seq}]."
    )
  if num_queries_per_block is not None and num_queries_per_block <= 0:
    raise ValueError(f"{num_queries_per_block=} must be positive.")
  if vmem_limit_bytes is not None and vmem_limit_bytes <= 0:
    raise ValueError(f"{vmem_limit_bytes=} must be positive.")
  del sm_scale  # No constraints on sm_scale.
  del mask_value  # No consstraints on mask_value.


def ragged_paged_attention_kernel(
    # Prefetch
    kv_lens_ref,  # [max_num_seqs]
    page_indices_ref,  # [max_num_seqs, pages_per_seq]
    cu_q_lens_ref,  # [max_num_seqs + 1]
    seq_buf_idx_ref,
    # TODO(jevinjiang): if OOM in SMEM, consider pack to other scalar refs.
    num_seqs_ref,
    # Input
    q_ref,  # [num_q_per_blk, num_q_heads_per_blk, head_dim]
    kv_pages_hbm_ref,  # [total_num_pages, page_size, num_combined_kv_heads, head_dim]
    # Output
    o_ref,  # [num_q_per_blk, num_q_heads_per_blk, head_dim]
    # Scratch

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Pass None instead of 0 to use the default VMEM limit
2. Clamp computed budgets: vmem_limit_bytes = budget if budget > 0 else None
3. Re-check the VMEM size constant you subtract from (per TPU generation it differs)

Example fix

// before
attn(..., vmem_limit_bytes=vmem_total - reserved)  # <= 0
// after
budget = vmem_total - reserved
attn(..., vmem_limit_bytes=budget if budget > 0 else None)

vmem_limit_bytes = vmem_limit_bytes if (vmem_limit_bytes is not None and vmem_limit_bytes > 0) else None

When it happens

Trigger: Passing vmem_limit_bytes=0 or a negative value, often as a sentinel meaning 'no limit' or from subtractive arithmetic computing a budget.

Common situations: 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'.

Understand the failure class

Background: "Must be a positive integer", "Invalid value", "Unsupported": the invalid-argument-value error family, when a library rejects the value you pass — this error's family across 35 libraries.

Related errors

• Invalid memory space: {memory_space!r}
• {num_queries_per_block=} must be positive.
• {bkv_compute=} should not be greater than {bkv=}
• Compiler params for platform {platform} cannot be used for {
• Memory space {self.memory_space} is not supported by mesh {s

AI-assisted analysis of jax-ml/jax@1e1c6a8fc0 (2026-08-27). Data as JSON: /api/errors/8b3a47bb57b76889. Report an issue: GitHub.