jax-ml/jax · error · ValueError

{num_queries_per_block=} must be positive.

Error message

{num_queries_per_block=} must be positive.

What it means

num_queries_per_block optionally sets how many query rows the ragged paged attention kernel processes per block; it must be a positive integer. Zero or negative values are rejected by static_validate_inputs.

Source

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

        "The dtype of `kv_lens`, `page_indices`, and `cu_q_lens` must be"
        f" int32. Got {kv_lens.dtype=}, {page_indices.dtype=},"
        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

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Pass None to let autotuning choose
2. Ensure the computed value is >= 1; guard with max(1, value) if derived from arithmetic
3. Validate tuning configs before passing them to the kernel

Example fix

// before
attn(..., num_queries_per_block=num_q // blk)  # 0 when num_q < blk
// after
attn(..., num_queries_per_block=max(1, num_q // blk) if blk <= num_q else None)

num_queries_per_block = None if num_queries_per_block is None else max(1, int(num_queries_per_block))
if num_queries_per_block is not None:
    assert num_queries_per_block > 0

When it happens

Trigger: Passing num_queries_per_block=0 (intending autotune/disabled) or a negative tuning value from a config sweep.

Common situations: Hyperparameter search code writing 0 as 'unset', or arithmetic that underflows (e.g. max_seq_len // something == 0 for tiny inputs).

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_kv_pages_per_block=} must be in range (0, {pages_per_se
• {vmem_limit_bytes=} must be positive.
• block_kv must be a multiple of {NUM_LANES}
• block_q must be a multiple of {NUM_LANES}

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